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 crate::{ConstraintBuilder, TraceRow, Uair, ideal::ImpossibleIdeal};
7use crypto_primitives::FixedConfig;
8use num_traits::{CheckedAdd, CheckedMul, CheckedSub, One, Pow, Zero};
9use zinc_utils::add;
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_flattened::<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///
22/// Returns a `Vec<usize>` where the i-th element is the degree of the i-th
23/// emitted constraint, in emission order. Under the
24/// [`crate::ConstraintBuilder::assert_in_fq_ideal`] ordering convention, all
25/// $Q[X]$ degrees appear first, followed by all
26/// $F_{q_i}[X]$ degrees.
27pub fn count_constraint_degrees_flattened<U: Uair>() -> Vec<usize> {
28    let split = count_constraint_degrees::<U>();
29    let mut all = split.q_degrees;
30    all.extend(split.fq_degrees.into_iter().flatten());
31    all
32}
33
34/// Compute the per-family degrees of each constraint in `U`.
35pub fn count_constraint_degrees<U: Uair>() -> ConstraintDegreeCollector {
36    let mut dc = ConstraintDegreeCollector::default();
37
38    let sig = U::signature();
39    let (up_dummy, down_dummy) = sig.dummy_rows(DegreeCountingSemiring::var());
40    let up_row = TraceRow::from_slice_with_layout(&up_dummy, sig.total_cols().as_column_layout());
41    let down_row =
42        TraceRow::from_slice_with_layout(&down_dummy, sig.down_cols().as_column_layout());
43
44    U::constrain_general(
45        &mut dc,
46        &FixedConfig::<DegreeCountingSemiring>::default(),
47        up_row,
48        down_row,
49        |_| DegreeCountingSemiring::scalar(),
50        |x, _| Some(*x),
51        |_| ImpossibleIdeal,
52        |_| ImpossibleIdeal,
53    );
54
55    dc
56}
57
58/// Collects the degree of each constraint in a UAIR by implementing the
59/// `ConstraintBuilder` trait.
60#[derive(Debug, Default)]
61pub struct ConstraintDegreeCollector {
62    pub q_degrees: Vec<usize>,
63    pub fq_degrees: Vec<Vec<usize>>,
64}
65
66impl ConstraintBuilder for ConstraintDegreeCollector {
67    type Expr = DegreeCountingSemiring;
68    type Ideal = ImpossibleIdeal;
69    type FqIdeal = ImpossibleIdeal;
70
71    fn assert_in_ideal(&mut self, expr: Self::Expr, _ideal: &Self::Ideal) {
72        self.q_degrees.push(expr.0);
73    }
74
75    fn assert_zero(&mut self, expr: Self::Expr) {
76        self.q_degrees.push(expr.0);
77    }
78
79    fn assert_in_fq_ideal(&mut self, prime_idx: usize, expr: Self::Expr, _ideal: &Self::FqIdeal) {
80        if self.fq_degrees.len() <= prime_idx {
81            self.fq_degrees.resize(add!(prime_idx, 1), Vec::new());
82        }
83        self.fq_degrees[prime_idx].push(expr.0);
84    }
85}
86
87#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
88pub struct DegreeCountingSemiring(usize);
89
90impl DegreeCountingSemiring {
91    pub fn var() -> Self {
92        DegreeCountingSemiring(1)
93    }
94
95    pub fn scalar() -> Self {
96        DegreeCountingSemiring(0)
97    }
98}
99
100impl Display for DegreeCountingSemiring {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        Debug::fmt(&self, f)
103    }
104}
105
106macro_rules! impl_binary_additive_op {
107    ($trait:ident, $op:ident) => {
108        impl $trait<&DegreeCountingSemiring> for DegreeCountingSemiring {
109            type Output = Self;
110
111            #[inline(always)]
112            fn $op(self, rhs: &DegreeCountingSemiring) -> Self::Output {
113                DegreeCountingSemiring(std::cmp::max(self.0, rhs.0))
114            }
115        }
116
117        impl $trait<DegreeCountingSemiring> for DegreeCountingSemiring {
118            type Output = Self;
119
120            #[inline(always)]
121            fn $op(self, rhs: DegreeCountingSemiring) -> Self::Output {
122                self.$op(&rhs)
123            }
124        }
125    };
126}
127
128impl_binary_additive_op!(Add, add);
129impl_binary_additive_op!(Sub, sub);
130
131impl Mul<&Self> for DegreeCountingSemiring {
132    type Output = Self;
133
134    #[allow(clippy::arithmetic_side_effects, clippy::suspicious_arithmetic_impl)]
135    #[inline(always)]
136    fn mul(self, rhs: &Self) -> Self::Output {
137        DegreeCountingSemiring(self.0 + rhs.0)
138    }
139}
140
141impl Mul<Self> for DegreeCountingSemiring {
142    type Output = Self;
143
144    #[inline(always)]
145    fn mul(self, rhs: Self) -> Self::Output {
146        self.mul(&rhs)
147    }
148}
149
150macro_rules! impl_additive_op_assign {
151    ($trait:ident, $op:ident) => {
152        impl $trait<&DegreeCountingSemiring> for DegreeCountingSemiring {
153            #[inline(always)]
154            fn $op(&mut self, rhs: &DegreeCountingSemiring) {
155                self.0 = std::cmp::max(self.0, rhs.0);
156            }
157        }
158
159        impl $trait<DegreeCountingSemiring> for DegreeCountingSemiring {
160            #[inline(always)]
161            fn $op(&mut self, rhs: DegreeCountingSemiring) {
162                self.$op(&rhs);
163            }
164        }
165    };
166}
167
168impl_additive_op_assign!(AddAssign, add_assign);
169impl_additive_op_assign!(SubAssign, sub_assign);
170
171impl MulAssign<&Self> for DegreeCountingSemiring {
172    #[allow(clippy::arithmetic_side_effects, clippy::suspicious_op_assign_impl)]
173    #[inline(always)]
174    fn mul_assign(&mut self, rhs: &Self) {
175        self.0 += rhs.0;
176    }
177}
178
179impl MulAssign<Self> for DegreeCountingSemiring {
180    #[inline(always)]
181    fn mul_assign(&mut self, rhs: Self) {
182        self.add_assign(&rhs);
183    }
184}
185
186macro_rules! impl_checked_additive_op {
187    ($trait:ident, $op:ident) => {
188        impl $trait for DegreeCountingSemiring {
189            #[inline(always)]
190            fn $op(&self, rhs: &DegreeCountingSemiring) -> Option<Self::Output> {
191                Some(DegreeCountingSemiring(std::cmp::max(self.0, rhs.0)))
192            }
193        }
194    };
195}
196
197impl_checked_additive_op!(CheckedAdd, checked_add);
198impl_checked_additive_op!(CheckedSub, checked_sub);
199
200impl CheckedMul for DegreeCountingSemiring {
201    #[inline(always)]
202    fn checked_mul(&self, rhs: &Self) -> Option<Self> {
203        Some(DegreeCountingSemiring(self.0.checked_add(rhs.0)?))
204    }
205}
206
207impl Zero for DegreeCountingSemiring {
208    #[inline(always)]
209    fn zero() -> Self {
210        Self::scalar()
211    }
212
213    #[inline(always)]
214    fn is_zero(&self) -> bool {
215        self.0 == 0
216    }
217}
218
219impl One for DegreeCountingSemiring {
220    #[inline(always)]
221    fn one() -> Self {
222        Self::scalar()
223    }
224}
225
226impl From<bool> for DegreeCountingSemiring {
227    #[inline(always)]
228    fn from(_value: bool) -> Self {
229        Self::scalar()
230    }
231}
232
233impl Pow<u32> for DegreeCountingSemiring {
234    type Output = Self;
235
236    #[allow(clippy::arithmetic_side_effects)] // degrees are small
237    #[inline(always)]
238    fn pow(self, exp: u32) -> Self {
239        DegreeCountingSemiring(self.0 * exp as usize)
240    }
241}
242
243impl std::iter::Sum for DegreeCountingSemiring {
244    #[allow(clippy::arithmetic_side_effects)] // degrees are small
245    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
246        iter.fold(Self::scalar(), |acc, x| acc + x)
247    }
248}
249
250impl<'a> std::iter::Sum<&'a Self> for DegreeCountingSemiring {
251    #[allow(clippy::arithmetic_side_effects)] // degrees are small
252    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
253        iter.fold(Self::scalar(), |acc, x| acc + x)
254    }
255}
256
257impl std::iter::Product for DegreeCountingSemiring {
258    #[allow(clippy::arithmetic_side_effects)] // degrees are small
259    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
260        iter.fold(Self::scalar(), |acc, x| acc * x)
261    }
262}
263
264impl<'a> std::iter::Product<&'a Self> for DegreeCountingSemiring {
265    #[allow(clippy::arithmetic_side_effects)] // degrees are small
266    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
267        iter.fold(Self::scalar(), |acc, x| acc * x)
268    }
269}