updated mod-builder to support generalized setup and better finalization

Author: Avaneesh-axiom

Cargo.lock

@@ -250,6 +250,7 @@ k256 = { version = "0.13.3", default-features = false }
 elliptic-curve = { version = "0.13.8", default-features = false }
 ecdsa = { version = "0.16.9", default-features = false }
 num-bigint = { version = "0.4.6", default-features = false }
+num-bigint-dig = { version = "0.8.4", default-features = false }
 num-integer = { version = "0.1.46", default-features = false }
 num-traits = { version = "0.2.19", default-features = false }
 ff = { version = "0.13.0", default-features = false }

Cargo.toml

@@ -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::{
@@ -310,18 +311,32 @@ 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();
+
+            if format!("{:?}", flags[0].into()) == "0" {
+                println!("lhs: {:?}", reads);
+                println!("rhs: {:?}", expected);
+            }
+
+            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/builder.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,
@@ -314,14 +315,36 @@ where
             // We will copy over the core part of first 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 first_row = trace.rows().nth(0).unwrap().collect::<Vec<_>>();
-            let first_row_core = first_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 first row, except "is_valid" (the first element of core part) is zero.
-                core_row.copy_from_slice(first_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
+    }
+}

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

@@ -21,6 +21,7 @@ openvm-rv32-adapters = { workspace = true }
 openvm-ecc-transpiler = { workspace = true }
 
 num-bigint = { workspace = true }
+num-bigint-dig = { workspace = true }
 num-traits = { workspace = true }
 strum = { workspace = true }
 derive_more = { workspace = true }