updated mod-builder

- support checking more constants during setup
- support new finalization (with dummy row and temp range checker)

Author: Avaneesh-axiom

crates/circuits/mod-builder/src/builder.rs

@@ -1,5 +1,6 @@
-use std::{cell::RefCell, ops::Deref, rc::Rc};
+use std::{cell::RefCell, iter, ops::Deref, rc::Rc};
 
+use itertools::{zip_eq, Itertools};
 use num_bigint::{BigInt, BigUint, Sign};
 use num_traits::Zero;
 use openvm_circuit_primitives::{
@@ -209,6 +210,9 @@ pub struct FieldExpr {
     pub check_carry_mod_to_zero: CheckCarryModToZeroSubAir,
 
     pub range_bus: VariableRangeCheckerBus,
+
+    // any values other than the prime modulus that need to be checked at setup
+    pub setup_values: Vec<BigUint>,
 }
 
 impl FieldExpr {
@@ -229,8 +233,20 @@ impl FieldExpr {
             builder,
             check_carry_mod_to_zero: subair,
             range_bus,
+            setup_values: vec![],
         }
     }
+
+    pub fn new_with_setup_values(
+        builder: ExprBuilder,
+        range_bus: VariableRangeCheckerBus,
+        needs_setup: bool,
+        setup_values: Vec<BigUint>,
+    ) -> Self {
+        let mut ret = Self::new(builder, range_bus, needs_setup);
+        ret.setup_values = setup_values;
+        ret
+    }
 }
 
 impl Deref for FieldExpr {
@@ -295,18 +311,27 @@ impl<AB: InteractionBuilder> SubAir<AB> for FieldExpr {
             //   however this has the challenge that when the same chip is used
             //   across continuation segments,
             //   only the first segment will have setup called
-            for i in 0..inputs[0].len().max(self.builder.prime_limbs.len()) {
-                let lhs = if i < inputs[0].len() {
-                    inputs[0][i].into()
-                } else {
-                    AB::Expr::ZERO
-                };
-                let rhs = if i < self.builder.prime_limbs.len() {
-                    AB::Expr::from_canonical_usize(self.builder.prime_limbs[i])
-                } else {
-                    AB::Expr::ZERO
-                };
-                builder.when(is_setup.clone()).assert_eq(lhs, rhs);
+
+            let expected = iter::empty()
+                .chain(self.builder.prime_limbs.clone())
+                .chain(self.setup_values.iter().flat_map(|x| {
+                    big_uint_to_num_limbs(x, self.builder.limb_bits, self.builder.num_limbs)
+                        .into_iter()
+                }))
+                .collect_vec();
+
+            let reads: Vec<AB::Expr> = inputs
+                .clone()
+                .into_iter()
+                .flatten()
+                .map(Into::into)
+                .take(expected.len())
+                .collect();
+
+            for (lhs, rhs) in zip_eq(&reads, expected) {
+                builder
+                    .when(is_setup.clone())
+                    .assert_eq(lhs.clone(), AB::F::from_canonical_usize(rhs));
             }
         }
 

crates/circuits/mod-builder/src/core_chip.rs

@@ -1,5 +1,6 @@
 use itertools::Itertools;
 use num_bigint::BigUint;
+use num_traits::Zero;
 use openvm_circuit::arch::{
     AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, MinimalInstruction,
     Result, VmAdapterInterface, VmCoreAir, VmCoreChip,
@@ -292,20 +293,38 @@ where
         if !self.should_finalize || num_records == 0 {
             return;
         }
-        // We will copy over the core part of last row to padded rows (all rows after num_records).
+
         let core_width = <Self::Air as BaseAir<F>>::width(&self.air);
         let adapter_width = trace.width() - core_width;
-        let last_row = trace
-            .rows()
-            .nth(num_records - 1)
-            .unwrap()
-            .collect::<Vec<_>>();
-        let last_row_core = last_row.split_at(adapter_width).1;
+        let dummy_row = self.generate_dummy_trace_row(adapter_width, core_width);
         for row in trace.rows_mut().skip(num_records) {
-            let core_row = row.split_at_mut(adapter_width).1;
-            // The same as last row, except "is_valid" (the first element of core part) is zero.
-            core_row.copy_from_slice(last_row_core);
-            core_row[0] = F::ZERO;
+            row.copy_from_slice(&dummy_row);
         }
     }
 }
+
+impl FieldExpressionCoreChip {
+    // We will be setting is_valid = 0. That forces all flags be 0 (otherwise setup will be -1).
+    // We generate a dummy row with all flags set to 0, then we set is_valid = 0.
+    fn generate_dummy_trace_row<F: PrimeField32>(
+        &self,
+        adapter_width: usize,
+        core_width: usize,
+    ) -> Vec<F> {
+        let record = FieldExpressionRecord {
+            inputs: vec![BigUint::zero(); self.air.num_inputs()],
+            flags: vec![false; self.air.num_flags()],
+        };
+        let mut row = vec![F::ZERO; adapter_width + core_width];
+        let core_row = &mut row[adapter_width..];
+        // We **do not** want this trace row to update the range checker
+        // so we must create a temporary range checker
+        let tmp_range_checker = SharedVariableRangeCheckerChip::new(self.range_checker.bus());
+        self.air.expr.generate_subrow(
+            (tmp_range_checker.as_ref(), record.inputs, record.flags),
+            core_row,
+        );
+        core_row[0] = F::ZERO; // is_valid = 0
+        row
+    }
+}