Skip to main content

zinc_uair/
degree_counter.rs

1use std::{
2    fmt::{Debug, Display},
3    ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign},
4};
5
6use crypto_primitives::Semiring;
7use num_traits::{CheckedAdd, CheckedMul, CheckedSub};
8
9use crate::{ConstraintBuilder, TraceRow, Uair, ideal::ImpossibleIdeal};
10
11/// Compute the maximum number of multiplicands
12/// in products of witness elements in the UAIR `U`.
13pub fn count_max_degree<U: Uair>() -> usize {
14    count_constraint_degrees::<U>()
15        .into_iter()
16        .max()
17        .unwrap_or(0)
18}
19
20/// Compute the degree of each individual constraint in the UAIR `U`.
21/// Returns a `Vec<usize>` where the i-th element is the degree
22/// of the i-th constraint.
23pub fn count_constraint_degrees<U: Uair>() -> Vec<usize> {
24    let mut dc = ConstraintDegreeCollector::default();
25
26    let sig = U::signature();
27    let (up_dummy, down_dummy) = sig.dummy_rows(DegreeCountingSemiring::var());
28    let up_row = TraceRow::from_slice_with_layout(&up_dummy, sig.total_cols().as_column_layout());
29    let down_row =
30        TraceRow::from_slice_with_layout(&down_dummy, sig.down_cols().as_column_layout());
31
32    U::constrain_general(
33        &mut dc,
34        up_row,
35        down_row,
36        |_| DegreeCountingSemiring::scalar(),
37        |x, _| Some(*x),
38        |_| ImpossibleIdeal,
39    );
40
41    dc.degrees
42}
43
44/// Collects the degree of each constraint in a UAIR by implementing the
45/// `ConstraintBuilder` trait.
46#[derive(Debug, Default)]
47pub(crate) struct ConstraintDegreeCollector {
48    degrees: Vec<usize>,
49}
50
51impl ConstraintBuilder for ConstraintDegreeCollector {
52    type Expr = DegreeCountingSemiring;
53    type Ideal = ImpossibleIdeal;
54
55    fn assert_in_ideal(&mut self, expr: Self::Expr, _ideal: &Self::Ideal) {
56        self.degrees.push(expr.0);
57    }
58
59    fn assert_zero(&mut self, expr: Self::Expr) {
60        self.degrees.push(expr.0);
61    }
62}
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
65pub(crate) struct DegreeCountingSemiring(usize);
66
67impl DegreeCountingSemiring {
68    pub fn var() -> Self {
69        DegreeCountingSemiring(1)
70    }
71
72    pub fn scalar() -> Self {
73        DegreeCountingSemiring(0)
74    }
75}
76
77impl Display for DegreeCountingSemiring {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        Debug::fmt(&self, f)
80    }
81}
82
83macro_rules! impl_binary_additive_op {
84    ($trait:ident, $op:ident) => {
85        impl $trait<&DegreeCountingSemiring> for DegreeCountingSemiring {
86            type Output = Self;
87
88            #[inline(always)]
89            fn $op(self, rhs: &DegreeCountingSemiring) -> Self::Output {
90                DegreeCountingSemiring(std::cmp::max(self.0, rhs.0))
91            }
92        }
93
94        impl $trait<DegreeCountingSemiring> for DegreeCountingSemiring {
95            type Output = Self;
96
97            #[inline(always)]
98            fn $op(self, rhs: DegreeCountingSemiring) -> Self::Output {
99                self.$op(&rhs)
100            }
101        }
102    };
103}
104
105impl_binary_additive_op!(Add, add);
106impl_binary_additive_op!(Sub, sub);
107
108impl Mul<&Self> for DegreeCountingSemiring {
109    type Output = Self;
110
111    #[allow(clippy::arithmetic_side_effects, clippy::suspicious_arithmetic_impl)]
112    #[inline(always)]
113    fn mul(self, rhs: &Self) -> Self::Output {
114        DegreeCountingSemiring(self.0 + rhs.0)
115    }
116}
117
118impl Mul<Self> for DegreeCountingSemiring {
119    type Output = Self;
120
121    #[inline(always)]
122    fn mul(self, rhs: Self) -> Self::Output {
123        self.mul(&rhs)
124    }
125}
126
127macro_rules! impl_additive_op_assign {
128    ($trait:ident, $op:ident) => {
129        impl $trait<&DegreeCountingSemiring> for DegreeCountingSemiring {
130            #[inline(always)]
131            fn $op(&mut self, rhs: &DegreeCountingSemiring) {
132                self.0 = std::cmp::max(self.0, rhs.0);
133            }
134        }
135
136        impl $trait<DegreeCountingSemiring> for DegreeCountingSemiring {
137            #[inline(always)]
138            fn $op(&mut self, rhs: DegreeCountingSemiring) {
139                self.$op(&rhs);
140            }
141        }
142    };
143}
144
145impl_additive_op_assign!(AddAssign, add_assign);
146impl_additive_op_assign!(SubAssign, sub_assign);
147
148impl MulAssign<&Self> for DegreeCountingSemiring {
149    #[allow(clippy::arithmetic_side_effects, clippy::suspicious_op_assign_impl)]
150    #[inline(always)]
151    fn mul_assign(&mut self, rhs: &Self) {
152        self.0 += rhs.0;
153    }
154}
155
156impl MulAssign<Self> for DegreeCountingSemiring {
157    #[inline(always)]
158    fn mul_assign(&mut self, rhs: Self) {
159        self.add_assign(&rhs);
160    }
161}
162
163macro_rules! impl_checked_additive_op {
164    ($trait:ident, $op:ident) => {
165        impl $trait for DegreeCountingSemiring {
166            #[inline(always)]
167            fn $op(&self, rhs: &DegreeCountingSemiring) -> Option<Self::Output> {
168                Some(DegreeCountingSemiring(std::cmp::max(self.0, rhs.0)))
169            }
170        }
171    };
172}
173
174impl_checked_additive_op!(CheckedAdd, checked_add);
175impl_checked_additive_op!(CheckedSub, checked_sub);
176
177impl CheckedMul for DegreeCountingSemiring {
178    #[inline(always)]
179    fn checked_mul(&self, rhs: &Self) -> Option<Self> {
180        Some(DegreeCountingSemiring(self.0.checked_add(rhs.0)?))
181    }
182}
183
184impl Semiring for DegreeCountingSemiring {}