feat: add setup constraints with a for ecdouble

extensions/ecc/circuit/src/weierstrass_chip/double.rs

@@ -1,13 +1,14 @@
-use std::{cell::RefCell, rc::Rc, sync::Arc};
+use std::{cell::RefCell, iter, rc::Rc, sync::Arc};
 
-use itertools::Itertools;
+use itertools::{zip_eq, Itertools};
 use num_bigint_dig::BigUint;
 use num_traits::One;
 use openvm_circuit::arch::{
     AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, MinimalInstruction,
     Result, VmAdapterInterface, VmCoreAir, VmCoreChip,
 };
 use openvm_circuit_primitives::{
+    bigint::utils::big_uint_to_num_limbs,
     var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip},
     SubAir, TraceSubRowGenerator,
 };
@@ -19,9 +20,8 @@ use openvm_mod_circuit_builder::{
 };
 use openvm_stark_backend::{
     interaction::InteractionBuilder,
-    p3_air::BaseAir,
+    p3_air::{AirBuilder, BaseAir},
     p3_field::{AbstractField, Field, PrimeField32},
-    p3_matrix::dense::RowMajorMatrix,
     rap::BaseAirWithPublicValues,
 };
 
@@ -111,7 +111,7 @@ where
         assert_eq!(vars.len(), 3); // x1^2, x3, y3
         assert_eq!(flags.len(), 1); // is_double_flag
 
-        let reads: Vec<AB::Expr> = inputs.concat().iter().map(|x| (*x).into()).collect();
+        let reads: Vec<AB::Expr> = inputs.into_iter().flatten().map(Into::into).collect();
         let writes: Vec<AB::Expr> = self
             .output_indices()
             .iter()
@@ -120,10 +120,26 @@ where
             .collect();
 
         let is_setup = is_valid - flags[0];
+        builder.assert_bool(is_setup.clone());
         let local_opcode_idx = flags[0]
             * AB::Expr::from_canonical_usize(Rv32WeierstrassOpcode::EC_DOUBLE as usize)
-            + is_setup
+            + is_setup.clone()
                 * AB::Expr::from_canonical_usize(Rv32WeierstrassOpcode::SETUP_EC_DOUBLE as usize);
+        // when is_setup, assert `reads` equals `(modulus, a)`
+        for (lhs, &rhs) in zip_eq(
+            &reads,
+            iter::empty()
+                .chain(&self.expr.builder.prime_limbs)
+                .chain(&big_uint_to_num_limbs(
+                    &self.a_biguint,
+                    self.expr.builder.limb_bits,
+                    self.expr.builder.num_limbs,
+                )),
+        ) {
+            builder
+                .when(is_setup.clone())
+                .assert_eq(lhs.clone(), AB::F::from_canonical_usize(rhs));
+        }
 
         let instruction = MinimalInstruction {
             is_valid: is_valid.into(),
@@ -246,6 +262,4 @@ where
     fn air(&self) -> &Self::Air {
         &self.air
     }
-
-    fn finalize(&self, trace: &mut RowMajorMatrix<F>, num_records: usize) {}
 }

extensions/ecc/circuit/src/weierstrass_chip/tests.rs

@@ -189,11 +189,10 @@ fn test_double() {
     assert_eq!(r[2], SampleEcPoints[3].1);
 
     let prime_limbs: [BabyBear; NUM_LIMBS] = prime_limbs(&chip.0.core.air.expr).try_into().unwrap();
-    let mut one_limbs = [BabyBear::ONE; NUM_LIMBS];
-    one_limbs[0] = BabyBear::ONE;
+    let a_limbs = [BabyBear::ZERO; NUM_LIMBS];
     let setup_instruction = rv32_write_heap_default(
         &mut tester,
-        vec![prime_limbs, one_limbs], // inputs[0] = prime, others doesn't matter
+        vec![prime_limbs, a_limbs], // inputs[0] = prime, inputs[1] = a coeff of weierstrass equation
         vec![],
         chip.0.core.air.offset + Rv32WeierstrassOpcode::SETUP_EC_DOUBLE as usize,
     );
@@ -257,7 +256,7 @@ fn test_p256_double() {
         tester.memory_controller(),
         config,
         Rv32WeierstrassOpcode::default_offset(),
-        a,
+        a.clone(),
     );
     assert_eq!(chip.0.core.air.expr.builder.num_variables, 3); // lambda, x3, y3
 
@@ -277,11 +276,11 @@ fn test_p256_double() {
     assert_eq!(r[2], expected_double_y);
 
     let prime_limbs: [BabyBear; NUM_LIMBS] = prime_limbs(&chip.0.core.air.expr).try_into().unwrap();
-    let mut one_limbs = [BabyBear::ONE; NUM_LIMBS];
-    one_limbs[0] = BabyBear::ONE;
+    let a_limbs =
+        biguint_to_limbs::<NUM_LIMBS>(a.clone(), LIMB_BITS).map(BabyBear::from_canonical_u32);
     let setup_instruction = rv32_write_heap_default(
         &mut tester,
-        vec![prime_limbs, one_limbs], // inputs[0] = prime, others doesn't matter
+        vec![prime_limbs, a_limbs], // inputs[0] = prime, inputs[1] = a coeff of weierstrass equation
         vec![],
         chip.0.core.air.offset + Rv32WeierstrassOpcode::SETUP_EC_DOUBLE as usize,
     );

extensions/ecc/sw-setup/src/lib.rs

@@ -513,9 +513,9 @@ pub fn sw_init(input: TokenStream) -> TokenStream {
                 #[cfg(target_os = "zkvm")]
                 {
                     // p1 is (x1, y1), and x1 must be the modulus.
-                    // y1 needs to be non-zero to avoid division by zero in double.
+                    // y1 can be anything for SetupEcAdd, but must equal `a` for SetupEcDouble
                     let modulus_bytes = <#item as openvm_algebra_guest::IntMod>::MODULUS;
-                    let mut one = [0u8; <#item as openvm_algebra_guest::IntMod>::NUM_LIMBS];
+                    let one = [0u8; <#item as openvm_algebra_guest::IntMod>::NUM_LIMBS];
                     one[0] = 1;
                     let p1 = [modulus_bytes.as_ref(), one.as_ref()].concat();
                     // (EcAdd only) p2 is (x2, y2), and x1 - x2 has to be non-zero to avoid division over zero in add.