openvm_mod_circuit_builder/
builder.rs

1use std::{cell::RefCell, cmp::min, iter, ops::Deref, rc::Rc};
2
3use itertools::{zip_eq, Itertools};
4use num_bigint::{BigInt, BigUint, Sign};
5use num_traits::{One, Zero};
6use openvm_circuit_primitives::{
7    bigint::{
8        check_carry_mod_to_zero::{CheckCarryModToZeroCols, CheckCarryModToZeroSubAir},
9        check_carry_to_zero::get_carry_max_abs_and_bits,
10        utils::*,
11        OverflowInt,
12    },
13    var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip},
14    ColumnsAir, SubAir, TraceSubRowGenerator,
15};
16use openvm_stark_backend::{
17    interaction::InteractionBuilder,
18    p3_air::{Air, AirBuilder, BaseAir},
19    p3_field::{Field, PrimeCharacteristicRing, PrimeField64},
20    p3_matrix::Matrix,
21    BaseAirWithPublicValues, PartitionedBaseAir,
22};
23
24use super::{FieldVariable, SymbolicExpr};
25
26#[derive(Clone)]
27pub struct ExprBuilderConfig {
28    pub modulus: BigUint,
29    pub num_limbs: usize,
30    pub limb_bits: usize,
31}
32
33impl ExprBuilderConfig {
34    pub fn check_valid(&self) {
35        assert!(self.modulus.bits() <= (self.num_limbs * self.limb_bits) as u64);
36    }
37}
38
39#[derive(Clone)]
40pub struct ExprBuilder {
41    // The prime field.
42    pub prime: BigUint,
43    // Same value, but we need BigInt for computing the quotient.
44    pub prime_bigint: BigInt,
45    pub prime_limbs: Vec<usize>,
46
47    pub num_input: usize,
48    pub num_flags: usize,
49
50    // This should be equal to number of constraints, but declare it to be explicit.
51    pub num_variables: usize,
52
53    pub constants: Vec<(BigUint, Vec<usize>)>, // value and limbs
54
55    /// The number of bits in a canonical representation of a limb.
56    pub limb_bits: usize,
57    /// Number of limbs in canonical representation of the bigint field element.
58    pub num_limbs: usize,
59    proper_max: BigUint,
60    // The max bits that we can range check.
61    pub range_checker_bits: usize,
62    // The max bits that carries are allowed to have.
63    pub max_carry_bits: usize,
64
65    // The number of limbs of the quotient for each constraint.
66    pub q_limbs: Vec<usize>,
67    // The number of limbs of the carries for each constraint.
68    pub carry_limbs: Vec<usize>,
69
70    // The constraints that should be evaluated to zero mod p (doesn't include - p * q part).
71    pub constraints: Vec<SymbolicExpr>,
72
73    // The equations to compute the newly introduced variables. For trace gen only.
74    pub computes: Vec<SymbolicExpr>,
75
76    pub output_indices: Vec<usize>,
77
78    /// flag for debug mode
79    debug: bool,
80
81    /// Whether the builder has been finalized. Only after finalize, we can do generate_subrow and
82    /// eval etc.
83    finalized: bool,
84
85    // Setup opcode is a special op that verifies the modulus is correct.
86    // There are some chips that don't need it because we hardcode the modulus. E.g. the pairing
87    // ones. For those chips need setup, setup is derived: setup = is_valid - sum(all_flags)
88    // Therefore when the chip only supports one opcode, user won't explicitly create a flag for it
89    // and we will create a default flag for it on finalizing.
90    needs_setup: bool,
91}
92
93// Number of bits in BabyBear modulus
94const MODULUS_BITS: usize = 31;
95
96impl ExprBuilder {
97    pub fn new(config: ExprBuilderConfig, range_checker_bits: usize) -> Self {
98        let prime_bigint = BigInt::from_biguint(Sign::Plus, config.modulus.clone());
99        let proper_max = (BigUint::one() << (config.num_limbs * config.limb_bits)) - BigUint::one();
100        // Max carry bits to ensure constraints don't overflow
101        let max_carry_bits = MODULUS_BITS - config.limb_bits - 2;
102        // sanity
103        assert!(config.limb_bits + 2 < MODULUS_BITS);
104        Self {
105            prime: config.modulus.clone(),
106            prime_bigint,
107            prime_limbs: big_uint_to_limbs(&config.modulus, config.limb_bits),
108            num_input: 0,
109            num_flags: 0,
110            limb_bits: config.limb_bits,
111            num_limbs: config.num_limbs,
112            proper_max,
113            range_checker_bits,
114            max_carry_bits: min(max_carry_bits, range_checker_bits),
115            num_variables: 0,
116            constants: vec![],
117            q_limbs: vec![],
118            carry_limbs: vec![],
119            constraints: vec![],
120            computes: vec![],
121            output_indices: vec![],
122            debug: false,
123            finalized: false,
124            needs_setup: false,
125        }
126    }
127
128    // This can be used to debug, when we only want to print something in a specific chip.
129    pub fn set_debug(&mut self) {
130        self.debug = true;
131    }
132
133    #[allow(unused)]
134    fn debug_print(&self, msg: &str) {
135        if self.debug {
136            println!("{msg}");
137        }
138    }
139
140    pub fn is_finalized(&self) -> bool {
141        self.finalized
142    }
143
144    pub fn finalize(&mut self, needs_setup: bool) {
145        self.finalized = true;
146        self.needs_setup = needs_setup;
147
148        // We don't support multi-op chip that doesn't need setup right now.
149        assert!(needs_setup || self.num_flags == 0);
150
151        // setup the default flag if needed
152        if needs_setup && self.num_flags == 0 {
153            self.new_flag();
154        }
155    }
156
157    pub fn new_input(builder: Rc<RefCell<ExprBuilder>>) -> FieldVariable {
158        let mut borrowed = builder.borrow_mut();
159        let num_limbs = borrowed.num_limbs;
160        let limb_bits = borrowed.limb_bits;
161        borrowed.num_input += 1;
162        let (num_input, max_carry_bits) = (borrowed.num_input, borrowed.max_carry_bits);
163        drop(borrowed);
164        FieldVariable {
165            expr: SymbolicExpr::Input(num_input - 1),
166            builder: builder.clone(),
167            limb_max_abs: (1 << limb_bits) - 1,
168            max_overflow_bits: limb_bits,
169            expr_limbs: num_limbs,
170            max_carry_bits,
171        }
172    }
173
174    pub fn new_flag(&mut self) -> usize {
175        self.num_flags += 1;
176        self.num_flags - 1
177    }
178
179    pub fn needs_setup(&self) -> bool {
180        assert!(self.finalized); // Should only be used after finalize.
181        self.needs_setup
182    }
183
184    // Below functions are used when adding variables and constraints manually, need to be careful.
185    // Number of variables, constraints and computes should be consistent,
186    // so there should be same number of calls to the new_var, add_constraint and add_compute.
187    pub fn new_var(&mut self) -> (usize, SymbolicExpr) {
188        self.num_variables += 1;
189        // Allocate space for the new variable, to make sure they are corresponding to the same
190        // variable index.
191        self.constraints.push(SymbolicExpr::Input(0));
192        self.computes.push(SymbolicExpr::Input(0));
193        self.q_limbs.push(0);
194        self.carry_limbs.push(0);
195        (
196            self.num_variables - 1,
197            SymbolicExpr::Var(self.num_variables - 1),
198        )
199    }
200
201    /// Creates a new constant (compile-time known) FieldVariable from `value` where
202    /// the big integer `value` is decomposed into `num_limbs` limbs of `limb_bits` bits,
203    /// with `num_limbs, limb_bits` specified by the builder config.
204    pub fn new_const(builder: Rc<RefCell<ExprBuilder>>, value: BigUint) -> FieldVariable {
205        let mut borrowed = builder.borrow_mut();
206        let index = borrowed.constants.len();
207        let limb_bits = borrowed.limb_bits;
208        let num_limbs = borrowed.num_limbs;
209        let limbs = big_uint_to_num_limbs(&value, limb_bits, num_limbs);
210        let max_carry_bits = borrowed.max_carry_bits;
211        borrowed.constants.push((value.clone(), limbs));
212        drop(borrowed);
213
214        FieldVariable {
215            expr: SymbolicExpr::Const(index, value, num_limbs),
216            builder,
217            limb_max_abs: (1 << limb_bits) - 1,
218            max_overflow_bits: limb_bits,
219            expr_limbs: num_limbs,
220            max_carry_bits,
221        }
222    }
223
224    pub fn set_constraint(&mut self, index: usize, constraint: SymbolicExpr) {
225        let (q_limbs, carry_limbs) = constraint.constraint_limbs(
226            &self.prime,
227            self.limb_bits,
228            self.num_limbs,
229            &self.proper_max,
230        );
231        self.constraints[index] = constraint;
232        self.q_limbs[index] = q_limbs;
233        self.carry_limbs[index] = carry_limbs;
234    }
235
236    pub fn set_compute(&mut self, index: usize, compute: SymbolicExpr) {
237        self.computes[index] = compute;
238    }
239
240    /// Returns `proper_max = 2^{num_limbs * limb_bits} - 1` as a precomputed value.
241    /// Any proper representation of a positive big integer using `num_limbs` limbs with
242    /// `limb_bits` bits each will be `<= proper_max`.
243    pub fn proper_max(&self) -> &BigUint {
244        &self.proper_max
245    }
246}
247
248#[derive(Clone)]
249pub struct FieldExpr {
250    pub builder: ExprBuilder,
251
252    pub check_carry_mod_to_zero: CheckCarryModToZeroSubAir,
253
254    pub range_bus: VariableRangeCheckerBus,
255
256    // any values other than the prime modulus that need to be checked at setup
257    pub setup_values: Vec<BigUint>,
258}
259
260impl FieldExpr {
261    pub fn new(
262        builder: ExprBuilder,
263        range_bus: VariableRangeCheckerBus,
264        needs_setup: bool,
265    ) -> Self {
266        let mut builder = builder;
267        builder.finalize(needs_setup);
268        let subair = CheckCarryModToZeroSubAir::new(
269            builder.prime.clone(),
270            builder.limb_bits,
271            range_bus.inner.index,
272            range_bus.range_max_bits,
273        );
274        FieldExpr {
275            builder,
276            check_carry_mod_to_zero: subair,
277            range_bus,
278            setup_values: vec![],
279        }
280    }
281
282    pub fn new_with_setup_values(
283        builder: ExprBuilder,
284        range_bus: VariableRangeCheckerBus,
285        needs_setup: bool,
286        setup_values: Vec<BigUint>,
287    ) -> Self {
288        let mut ret = Self::new(builder, range_bus, needs_setup);
289        ret.setup_values = setup_values;
290        ret
291    }
292
293    pub fn num_inputs(&self) -> usize {
294        self.builder.num_input
295    }
296
297    pub fn num_vars(&self) -> usize {
298        self.builder.num_variables
299    }
300
301    pub fn num_flags(&self) -> usize {
302        self.builder.num_flags
303    }
304
305    pub fn output_indices(&self) -> &[usize] {
306        &self.builder.output_indices
307    }
308}
309
310impl Deref for FieldExpr {
311    type Target = ExprBuilder;
312
313    fn deref(&self) -> &ExprBuilder {
314        &self.builder
315    }
316}
317
318impl<F: Field> BaseAirWithPublicValues<F> for FieldExpr {}
319impl<F: Field> PartitionedBaseAir<F> for FieldExpr {}
320// No columns provided: width is built dynamically from the expression and `FieldExprCols` uses
321// `Vec<Vec<T>>`.
322impl ColumnsAir for FieldExpr {}
323impl<F: Field> BaseAir<F> for FieldExpr {
324    fn width(&self) -> usize {
325        assert!(self.builder.is_finalized());
326        self.num_limbs * (self.builder.num_input + self.builder.num_variables)
327            + self.builder.q_limbs.iter().sum::<usize>()
328            + self.builder.carry_limbs.iter().sum::<usize>()
329            + self.builder.num_flags
330            + 1 // is_valid
331    }
332}
333
334impl<AB: InteractionBuilder> Air<AB> for FieldExpr {
335    fn eval(&self, builder: &mut AB) {
336        let main = builder.main();
337        let local = main.row_slice(0).expect("window should have two elements");
338        SubAir::eval(self, builder, &local);
339    }
340}
341
342impl<AB: InteractionBuilder> SubAir<AB> for FieldExpr {
343    /// The sub-row slice owned by the expression builder.
344    type AirContext<'a>
345        = &'a [AB::Var]
346    where
347        AB: 'a,
348        AB::Var: 'a,
349        AB::Expr: 'a;
350
351    fn eval<'a>(&'a self, builder: &'a mut AB, local: &'a [AB::Var])
352    where
353        AB::Var: 'a,
354        AB::Expr: 'a,
355    {
356        assert!(self.builder.is_finalized());
357        let FieldExprCols {
358            is_valid,
359            inputs,
360            vars,
361            q_limbs,
362            carry_limbs,
363            flags,
364        } = self.load_vars(local);
365
366        builder.assert_bool(is_valid);
367
368        if self.builder.needs_setup() {
369            let is_setup = flags.iter().fold(is_valid.into(), |acc, &x| acc - x);
370            builder.assert_bool(is_setup.clone());
371            // TODO[jpw]: currently we enforce at the program code level that:
372            // - a valid program must call the correct setup opcodes to be correct
373            // - it would be better if we can constraint this in the circuit, however this has the
374            //   challenge that when the same chip is used across continuation segments, only the
375            //   first segment will have setup called
376
377            let expected = iter::empty()
378                .chain({
379                    let mut prime_limbs = self.builder.prime_limbs.clone();
380                    prime_limbs.resize(self.builder.num_limbs, 0);
381                    prime_limbs
382                })
383                .chain(self.setup_values.iter().flat_map(|x| {
384                    big_uint_to_num_limbs(x, self.builder.limb_bits, self.builder.num_limbs)
385                        .into_iter()
386                }))
387                .collect_vec();
388
389            let reads: Vec<AB::Expr> = inputs
390                .clone()
391                .into_iter()
392                .flatten()
393                .map(Into::into)
394                .take(expected.len())
395                .collect();
396
397            for (lhs, rhs) in zip_eq(&reads, expected) {
398                builder
399                    .when(is_setup.clone())
400                    .assert_eq(lhs.clone(), AB::F::from_usize(rhs));
401            }
402        }
403
404        let inputs = load_overflow::<AB>(inputs, self.limb_bits);
405        let vars = load_overflow::<AB>(vars, self.limb_bits);
406        let constants: Vec<_> = self
407            .constants
408            .iter()
409            .map(|(_, limbs)| {
410                let limbs_expr: Vec<_> = limbs
411                    .iter()
412                    .map(|limb| AB::Expr::from_usize(*limb))
413                    .collect();
414                OverflowInt::from_unsigned_limbs(limbs_expr, self.limb_bits)
415            })
416            .collect();
417
418        for flag in flags.iter() {
419            builder.assert_bool(*flag);
420        }
421        for i in 0..self.constraints.len() {
422            let expr = self.constraints[i]
423                .evaluate_overflow_expr::<AB>(&inputs, &vars, &constants, &flags);
424            self.check_carry_mod_to_zero.eval(
425                builder,
426                (
427                    expr,
428                    CheckCarryModToZeroCols {
429                        carries: carry_limbs[i].clone(),
430                        quotient: q_limbs[i].clone(),
431                    },
432                    is_valid.into(),
433                ),
434            );
435        }
436
437        for var in vars.iter() {
438            for limb in var.limbs().iter() {
439                range_check(
440                    builder,
441                    self.range_bus.inner.index,
442                    self.range_bus.range_max_bits,
443                    self.limb_bits,
444                    limb.clone(),
445                    is_valid,
446                );
447            }
448        }
449    }
450}
451
452type Vecs<T> = Vec<Vec<T>>;
453
454pub struct FieldExprCols<T> {
455    pub is_valid: T,
456    pub inputs: Vecs<T>,
457    pub vars: Vecs<T>,
458    pub q_limbs: Vecs<T>,
459    pub carry_limbs: Vecs<T>,
460    pub flags: Vec<T>,
461}
462
463impl<F: PrimeField64> TraceSubRowGenerator<F> for FieldExpr {
464    type TraceContext<'a> = (&'a VariableRangeCheckerChip, Vec<BigUint>, Vec<bool>);
465    type ColsMut<'a> = &'a mut [F];
466
467    fn generate_subrow<'a>(
468        &'a self,
469        (range_checker, inputs, flags): (&'a VariableRangeCheckerChip, Vec<BigUint>, Vec<bool>),
470        sub_row: &'a mut [F],
471    ) {
472        assert!(self.builder.is_finalized());
473        assert_eq!(inputs.len(), self.num_input);
474        assert_eq!(self.num_variables, self.constraints.len());
475
476        assert_eq!(flags.len(), self.builder.num_flags);
477
478        let limb_bits = self.limb_bits;
479        let mut vars = vec![BigUint::zero(); self.num_variables];
480
481        // BigInt type is required for computing the quotient.
482        let input_bigint = inputs
483            .iter()
484            .map(|x| BigInt::from_biguint(Sign::Plus, x.clone()))
485            .collect::<Vec<BigInt>>();
486        let mut vars_bigint = vec![BigInt::zero(); self.num_variables];
487
488        // OverflowInt type is required for computing the carries.
489        let input_overflow = inputs
490            .iter()
491            .map(|x| OverflowInt::<isize>::from_biguint(x, self.limb_bits, Some(self.num_limbs)))
492            .collect::<Vec<_>>();
493        let zero = OverflowInt::<isize>::from_unsigned_limbs(vec![0], limb_bits);
494        let mut vars_overflow = vec![zero; self.num_variables];
495        // Note: in cases where the prime fits in less limbs than `num_limbs`, we use the smaller
496        // number of limbs.
497        let prime_overflow = OverflowInt::<isize>::from_biguint(&self.prime, self.limb_bits, None);
498
499        let constants: Vec<_> = self
500            .constants
501            .iter()
502            .map(|(_, limbs)| {
503                let limbs_isize: Vec<_> = limbs.iter().map(|i| *i as isize).collect();
504                OverflowInt::from_unsigned_limbs(limbs_isize, self.limb_bits)
505            })
506            .collect();
507
508        let mut all_q = vec![];
509        let mut all_carry = vec![];
510        for i in 0..self.constraints.len() {
511            let r = self.computes[i].compute(&inputs, &vars, &flags, &self.prime);
512            vars[i] = r.clone();
513            vars_bigint[i] = BigInt::from_biguint(Sign::Plus, r);
514            vars_overflow[i] =
515                OverflowInt::<isize>::from_biguint(&vars[i], self.limb_bits, Some(self.num_limbs));
516        }
517        // We need to have all variables computed first because, e.g. constraints[2] might need
518        // variables[3].
519        for i in 0..self.constraints.len() {
520            // expr = q * p
521            let expr_bigint =
522                self.constraints[i].evaluate_bigint(&input_bigint, &vars_bigint, &flags);
523            let q = &expr_bigint / &self.prime_bigint;
524            // If this is not true then the evaluated constraint is not divisible by p.
525            debug_assert_eq!(expr_bigint, &q * &self.prime_bigint);
526            let q_limbs = big_int_to_num_limbs(&q, limb_bits, self.q_limbs[i]);
527            assert_eq!(q_limbs.len(), self.q_limbs[i]); // If this fails, the q_limbs estimate is wrong.
528            for &q in q_limbs.iter() {
529                range_checker.add_count((q + (1 << limb_bits)) as u32, limb_bits + 1);
530            }
531            let q_overflow = OverflowInt::from_signed_limbs(q_limbs.clone(), limb_bits);
532            // compute carries of (expr - q * p)
533            let expr = self.constraints[i].evaluate_overflow_isize(
534                &input_overflow,
535                &vars_overflow,
536                &constants,
537                &flags,
538            );
539            let expr = expr - q_overflow * prime_overflow.clone();
540            let carries = expr.calculate_carries(limb_bits);
541            assert_eq!(carries.len(), self.carry_limbs[i]); // If this fails, the carry limbs estimate is wrong.
542            let max_overflow_bits = expr.max_overflow_bits();
543            let (carry_min_abs, carry_bits) =
544                get_carry_max_abs_and_bits(max_overflow_bits, limb_bits);
545            for &carry in carries.iter() {
546                range_checker.add_count((carry + carry_min_abs as isize) as u32, carry_bits);
547            }
548            all_q.push(vec_isize_to_f::<F>(q_limbs));
549            all_carry.push(vec_isize_to_f::<F>(carries));
550        }
551        for var in vars_overflow.iter() {
552            for limb in var.limbs().iter() {
553                range_checker.add_count(*limb as u32, limb_bits);
554            }
555        }
556
557        let input_limbs = input_overflow
558            .iter()
559            .map(|x| vec_isize_to_f::<F>(x.limbs().to_vec()))
560            .collect::<Vec<_>>();
561        let vars_limbs = vars_overflow
562            .iter()
563            .map(|x| vec_isize_to_f::<F>(x.limbs().to_vec()))
564            .collect::<Vec<_>>();
565
566        sub_row.copy_from_slice(
567            &[
568                vec![F::ONE],
569                input_limbs.concat(),
570                vars_limbs.concat(),
571                all_q.concat(),
572                all_carry.concat(),
573                flags.iter().map(|x| F::from_bool(*x)).collect::<Vec<_>>(),
574            ]
575            .concat(),
576        );
577    }
578}
579
580impl FieldExpr {
581    pub fn canonical_num_limbs(&self) -> usize {
582        self.builder.num_limbs
583    }
584
585    pub fn canonical_limb_bits(&self) -> usize {
586        self.builder.limb_bits
587    }
588
589    pub fn execute(&self, inputs: &[BigUint], flags: &[bool]) -> Vec<BigUint> {
590        assert!(self.builder.is_finalized());
591
592        #[cfg(debug_assertions)]
593        {
594            let is_setup = self.builder.needs_setup() && flags.iter().all(|&x| !x);
595            if is_setup {
596                assert_eq!(inputs[0], self.builder.prime);
597                // Check that inputs.iter().skip(1) has all the setup values as a prefix
598                assert!(inputs.len() > self.setup_values.len());
599                for (expected, actual) in self.setup_values.iter().zip(inputs.iter().skip(1)) {
600                    assert_eq!(expected, actual);
601                }
602            }
603        }
604
605        let mut vars = vec![BigUint::zero(); self.num_variables];
606        for i in 0..self.constraints.len() {
607            let r = self.computes[i].compute(inputs, &vars, flags, &self.prime);
608            vars[i] = r; // r is already owned, no clone needed
609        }
610        vars
611    }
612
613    pub fn execute_with_output(&self, inputs: &[BigUint], flags: &[bool]) -> Vec<BigUint> {
614        let vars = self.execute(inputs, flags);
615        self.builder
616            .output_indices
617            .iter()
618            .map(|i| vars[*i].clone())
619            .collect()
620    }
621
622    pub fn load_vars<T: Clone>(&self, arr: &[T]) -> FieldExprCols<T> {
623        assert!(self.builder.is_finalized());
624        let is_valid = arr[0].clone();
625        let mut idx = 1;
626        let mut inputs = vec![];
627        for _ in 0..self.num_input {
628            inputs.push(arr[idx..idx + self.num_limbs].to_vec());
629            idx += self.num_limbs;
630        }
631        let mut vars = vec![];
632        for _ in 0..self.num_variables {
633            vars.push(arr[idx..idx + self.num_limbs].to_vec());
634            idx += self.num_limbs;
635        }
636        let mut q_limbs = vec![];
637        for q in self.q_limbs.iter() {
638            q_limbs.push(arr[idx..idx + q].to_vec());
639            idx += q;
640        }
641        let mut carry_limbs = vec![];
642        for c in self.carry_limbs.iter() {
643            carry_limbs.push(arr[idx..idx + c].to_vec());
644            idx += c;
645        }
646        let flags = arr[idx..idx + self.num_flags].to_vec();
647        FieldExprCols {
648            is_valid,
649            inputs,
650            vars,
651            q_limbs,
652            carry_limbs,
653            flags,
654        }
655    }
656}
657
658fn load_overflow<AB: AirBuilder>(
659    arr: Vecs<AB::Var>,
660    limb_bits: usize,
661) -> Vec<OverflowInt<AB::Expr>> {
662    let mut result = vec![];
663    for x in arr.into_iter() {
664        let limbs: Vec<AB::Expr> = x.iter().cloned().map(|x| x.into()).collect();
665        result.push(OverflowInt::<AB::Expr>::from_unsigned_limbs(
666            limbs, limb_bits,
667        ));
668    }
669    result
670}
openvm_mod_circuit_builder/builder.rs

openvm_mod_circuit_builder/
builder.rs
