Merge branch 'kimchi/generic-cols' into zkvm/keccak/generic-cols

o1-labs · Nov 1, 2023 · 0d19643 · 0d19643
2 parents 01312b7 + bb0a7f0
commit 0d19643
Show file tree

Hide file tree

Showing 16 changed files with 135 additions and 60 deletions.
diff --git a/kimchi/src/bench.rs b/kimchi/src/bench.rs
@@ -77,12 +77,7 @@ impl BenchmarkCtx {
     }
 
     /// Produces a proof
-    pub fn create_proof(
-        &self,
-    ) -> (
-        ProverProof<Vesta, OpeningProof<Vesta>, KIMCHI_COLS>,
-        Vec<Fp>,
-    ) {
+    pub fn create_proof(&self) -> (ProverProof<Vesta, OpeningProof<Vesta>>, Vec<Fp>) {
         // create witness
         let witness: [Vec<Fp>; KIMCHI_COLS] = array::from_fn(|_| vec![1u32.into(); self.num_gates]);
 
@@ -102,13 +97,7 @@ impl BenchmarkCtx {
     }
 
     #[allow(clippy::type_complexity)]
-    pub fn batch_verification(
-        &self,
-        batch: &[(
-            ProverProof<Vesta, OpeningProof<Vesta>, KIMCHI_COLS>,
-            Vec<Fp>,
-        )],
-    ) {
+    pub fn batch_verification(&self, batch: &[(ProverProof<Vesta, OpeningProof<Vesta>>, Vec<Fp>)]) {
         // verify the proof
         let batch: Vec<_> = batch
             .iter()

diff --git a/kimchi/src/circuits/argument.rs b/kimchi/src/circuits/argument.rs
@@ -13,6 +13,7 @@ use serde::{Deserialize, Serialize};
 use super::{
     expr::{constraints::ExprOps, Cache, ConstantExpr, Constants},
     gate::{CurrOrNext, GateType},
+    wires::KIMCHI_COLS,
 };
 use CurrOrNext::{Curr, Next};
 
@@ -36,8 +37,8 @@ pub enum ArgumentType {
 /// created with ArgumentData and F = Field or F = PrimeField, then the constraints
 /// are built as expressions of real field elements and can be evaluated directly on
 /// the witness without using the prover.
-pub struct ArgumentEnv<F: 'static, T> {
-    data: Option<ArgumentData<F>>,
+pub struct ArgumentEnv<F: 'static, T, const COLUMNS: usize = KIMCHI_COLS> {
+    data: Option<ArgumentData<F, COLUMNS>>,
     phantom_data: PhantomData<T>,
 }
 
@@ -51,10 +52,14 @@ impl<F, T> Default for ArgumentEnv<F, T> {
     }
 }
 
-impl<F: Field, T: ExprOps<F>> ArgumentEnv<F, T> {
+impl<F: Field, T: ExprOps<F>, const COLUMNS: usize> ArgumentEnv<F, T, COLUMNS> {
     /// Initialize the environment for creating constraints of real field elements that can be
     /// evaluated directly over the witness without the prover/verifier
-    pub fn create(witness: ArgumentWitness<F>, coeffs: Vec<F>, constants: Constants<F>) -> Self {
+    pub fn create(
+        witness: ArgumentWitness<F, COLUMNS>,
+        coeffs: Vec<F>,
+        constants: Constants<F>,
+    ) -> Self {
         ArgumentEnv {
             data: Some(ArgumentData {
                 witness,
@@ -129,24 +134,24 @@ impl<F: Field, T: ExprOps<F>> ArgumentEnv<F, T> {
 }
 
 /// Argument environment data for constraints of field elements
-pub struct ArgumentData<F: 'static> {
+pub struct ArgumentData<F: 'static, const COLUMNS: usize = KIMCHI_COLS> {
     /// Witness rows
-    pub witness: ArgumentWitness<F>,
+    pub witness: ArgumentWitness<F, COLUMNS>,
     /// Gate coefficients
     pub coeffs: Vec<F>,
     /// Constants
     pub constants: Constants<F>,
 }
 
 /// Witness data for a argument
-pub struct ArgumentWitness<T> {
+pub struct ArgumentWitness<T, const COLUMNS: usize = KIMCHI_COLS> {
     /// Witness for current row
-    pub curr: Vec<T>,
+    pub curr: [T; COLUMNS],
     /// Witness for next row
-    pub next: Vec<T>,
+    pub next: [T; COLUMNS],
 }
 
-impl<T> std::ops::Index<(CurrOrNext, usize)> for ArgumentWitness<T> {
+impl<T, const COLUMNS: usize> std::ops::Index<(CurrOrNext, usize)> for ArgumentWitness<T, COLUMNS> {
     type Output = T;
 
     fn index(&self, idx: (CurrOrNext, usize)) -> &T {
@@ -168,7 +173,10 @@ pub trait Argument<F: PrimeField> {
     const CONSTRAINTS: u32;
 
     /// Constraints for this argument
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, cache: &mut Cache) -> Vec<T>;
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        cache: &mut Cache,
+    ) -> Vec<T>;
 
     /// Returns the set of constraints required to prove this argument.
     fn constraints(cache: &mut Cache) -> Vec<E<F>> {

diff --git a/kimchi/src/circuits/expr.rs b/kimchi/src/circuits/expr.rs
@@ -2660,13 +2660,20 @@ pub mod constraints {
         fn literal(x: F) -> Self;
 
         // Witness variable
-        fn witness(row: CurrOrNext, col: usize, env: Option<&ArgumentData<F>>) -> Self;
+        fn witness<const COLUMNS: usize>(
+            row: CurrOrNext,
+            col: usize,
+            env: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self;
 
         /// Coefficient
-        fn coeff(col: usize, env: Option<&ArgumentData<F>>) -> Self;
+        fn coeff<const COLUMNS: usize>(col: usize, env: Option<&ArgumentData<F, COLUMNS>>) -> Self;
 
         /// Create a constant
-        fn constant(expr: ConstantExpr<F>, env: Option<&ArgumentData<F>>) -> Self;
+        fn constant<const COLUMNS: usize>(
+            expr: ConstantExpr<F>,
+            env: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self;
 
         /// Cache item
         fn cache(&self, cache: &mut Cache) -> Self;
@@ -2726,15 +2733,22 @@ pub mod constraints {
             Expr::Constant(ConstantExpr::Literal(x))
         }
 
-        fn witness(row: CurrOrNext, col: usize, _: Option<&ArgumentData<F>>) -> Self {
+        fn witness<const COLUMNS: usize>(
+            row: CurrOrNext,
+            col: usize,
+            _: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self {
             witness(col, row)
         }
 
-        fn coeff(col: usize, _: Option<&ArgumentData<F>>) -> Self {
+        fn coeff<const COLUMNS: usize>(col: usize, _: Option<&ArgumentData<F, COLUMNS>>) -> Self {
             coeff(col)
         }
 
-        fn constant(expr: ConstantExpr<F>, _: Option<&ArgumentData<F>>) -> Self {
+        fn constant<const COLUMNS: usize>(
+            expr: ConstantExpr<F>,
+            _: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self {
             Expr::Constant(expr)
         }
 
@@ -2784,21 +2798,28 @@ pub mod constraints {
             x
         }
 
-        fn witness(row: CurrOrNext, col: usize, env: Option<&ArgumentData<F>>) -> Self {
+        fn witness<const COLUMNS: usize>(
+            row: CurrOrNext,
+            col: usize,
+            env: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self {
             match env {
                 Some(data) => data.witness[(row, col)],
                 None => panic!("Missing witness"),
             }
         }
 
-        fn coeff(col: usize, env: Option<&ArgumentData<F>>) -> Self {
+        fn coeff<const COLUMNS: usize>(col: usize, env: Option<&ArgumentData<F, COLUMNS>>) -> Self {
             match env {
                 Some(data) => data.coeffs[col],
                 None => panic!("Missing coefficients"),
             }
         }
 
-        fn constant(expr: ConstantExpr<F>, env: Option<&ArgumentData<F>>) -> Self {
+        fn constant<const COLUMNS: usize>(
+            expr: ConstantExpr<F>,
+            env: Option<&ArgumentData<F, COLUMNS>>,
+        ) -> Self {
             match env {
                 Some(data) => expr.value(&data.constants),
                 None => panic!("Missing constants"),

diff --git a/kimchi/src/circuits/gate.rs b/kimchi/src/circuits/gate.rs
@@ -264,7 +264,8 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
             zk_rows: cs.zk_rows,
         };
         // Create the argument environment for the constraints over field elements
-        let env = ArgumentEnv::<F, F>::create(argument_witness, self.coeffs.clone(), constants);
+        let env =
+            ArgumentEnv::<F, F, COLUMNS>::create(argument_witness, self.coeffs.clone(), constants);
 
         // Check the wiring (i.e. copy constraints) for this gate
         // Note: Gates can operated on row Curr or Curr and Next.
@@ -299,7 +300,9 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
                 // TODO: implement the verification for the generic gate
                 vec![]
             }
-            GateType::Poseidon => poseidon::Poseidon::constraint_checks(&env, &mut cache),
+            GateType::Poseidon => {
+                poseidon::Poseidon::constraint_checks::<F, COLUMNS>(&env, &mut cache)
+            }
             GateType::CompleteAdd => complete_add::CompleteAdd::constraint_checks(&env, &mut cache),
             GateType::VarBaseMul => varbasemul::VarbaseMul::constraint_checks(&env, &mut cache),
             GateType::EndoMul => endosclmul::EndosclMul::constraint_checks(&env, &mut cache),
@@ -358,7 +361,7 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
         &self,
         row: usize,
         witness: &[Vec<F>; COLUMNS],
-    ) -> CircuitGateResult<ArgumentWitness<F>> {
+    ) -> CircuitGateResult<ArgumentWitness<F, COLUMNS>> {
         // Get the part of the witness relevant to this gate
         let witness_curr: [F; COLUMNS] = (0..witness.len())
             .map(|col| witness[col][row])
@@ -375,9 +378,9 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
             [F::zero(); COLUMNS]
         };
 
-        Ok(ArgumentWitness::<F> {
-            curr: witness_curr.to_vec(),
-            next: witness_next.to_vec(),
+        Ok(ArgumentWitness::<F, COLUMNS> {
+            curr: witness_curr,
+            next: witness_next,
         })
     }
 }

diff --git a/kimchi/src/circuits/polynomials/complete_add.rs b/kimchi/src/circuits/polynomials/complete_add.rs
@@ -97,7 +97,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::CompleteAdd);
     const CONSTRAINTS: u32 = 7;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        cache: &mut Cache,
+    ) -> Vec<T> {
         // This function makes 2 + 1 + 1 + 1 + 2 = 7 constraints
         let x1 = env.witness_curr(0);
         let y1 = env.witness_curr(1);

diff --git a/kimchi/src/circuits/polynomials/endomul_scalar.rs b/kimchi/src/circuits/polynomials/endomul_scalar.rs
@@ -167,7 +167,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::EndoMulScalar);
     const CONSTRAINTS: u32 = 11;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        cache: &mut Cache,
+    ) -> Vec<T> {
         let n0 = env.witness_curr(0);
         let n8 = env.witness_curr(1);
         let a0 = env.witness_curr(2);

diff --git a/kimchi/src/circuits/polynomials/endosclmul.rs b/kimchi/src/circuits/polynomials/endosclmul.rs
@@ -186,7 +186,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::EndoMul);
     const CONSTRAINTS: u32 = 11;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        cache: &mut Cache,
+    ) -> Vec<T> {
         let b1 = env.witness_curr(11);
         let b2 = env.witness_curr(12);
         let b3 = env.witness_curr(13);

diff --git a/kimchi/src/circuits/polynomials/foreign_field_add/circuitgates.rs b/kimchi/src/circuits/polynomials/foreign_field_add/circuitgates.rs
@@ -142,7 +142,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::ForeignFieldAdd);
     const CONSTRAINTS: u32 = 4;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         let foreign_modulus: [T; LIMB_COUNT] = array::from_fn(|i| env.coeff(i));
 
         // stored as coefficient for better correspondance with the relation being proved

diff --git a/kimchi/src/circuits/polynomials/foreign_field_mul/circuitgates.rs b/kimchi/src/circuits/polynomials/foreign_field_mul/circuitgates.rs
@@ -193,7 +193,10 @@ where
     const CONSTRAINTS: u32 = 11;
     // DEGREE is 4
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         let mut constraints = vec![];
 
         //

diff --git a/kimchi/src/circuits/polynomials/generic.rs b/kimchi/src/circuits/polynomials/generic.rs
@@ -74,7 +74,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::Generic);
     const CONSTRAINTS: u32 = 2;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         // First generic gate
         let left_coeff1 = env.coeff(0);
         let right_coeff1 = env.coeff(1);

diff --git a/kimchi/src/circuits/polynomials/poseidon.rs b/kimchi/src/circuits/polynomials/poseidon.rs
@@ -337,7 +337,10 @@ where
     const ARGUMENT_TYPE: ArgumentType = ArgumentType::Gate(GateType::Poseidon);
     const CONSTRAINTS: u32 = 15;
 
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        cache: &mut Cache,
+    ) -> Vec<T> {
         let mut res = vec![];
 
         let mut idx = 0;

diff --git a/kimchi/src/circuits/polynomials/range_check/circuitgates.rs b/kimchi/src/circuits/polynomials/range_check/circuitgates.rs
@@ -178,7 +178,10 @@ where
     //   * Operates on Curr row
     //   * Range constrain all limbs except vp0 and vp1 (barring plookup constraints, which are done elsewhere)
     //   * Constrain that combining all limbs equals the limb stored in column 0
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         // 1) Apply range constraints on the limbs
         //    * Columns 1-2 are 12-bit copy constraints
         //        * They are copied 3 rows ahead (to the final row) and are constrained by lookups
@@ -279,7 +282,10 @@ where
     //   * Operates on Curr and Next row
     //   * Range constrain all limbs (barring plookup constraints, which are done elsewhere)
     //   * Constrain that combining all limbs equals the value v2 stored in row Curr, column 0
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         // 1) Apply range constraints on limbs for Curr row
         //    * Column 2 is a 2-bit crumb
         let mut constraints = vec![crumb(&env.witness_curr(2))];

diff --git a/kimchi/src/circuits/polynomials/rot.rs b/kimchi/src/circuits/polynomials/rot.rs
@@ -213,7 +213,10 @@ where
     // (stored in coefficient as a power-of-two form)
     //   * Operates on Curr row
     //   * Shifts the words by `rot` bits and then adds the excess to obtain the rotated word.
-    fn constraint_checks<T: ExprOps<F>>(env: &ArgumentEnv<F, T>, _cache: &mut Cache) -> Vec<T> {
+    fn constraint_checks<T: ExprOps<F>, const COLUMNS: usize>(
+        env: &ArgumentEnv<F, T, COLUMNS>,
+        _cache: &mut Cache,
+    ) -> Vec<T> {
         // Check that the last 8 columns are 2-bit crumbs
         // C1..C8: x * (x - 1) * (x - 2) * (x - 3) = 0
         let mut constraints = (7..KIMCHI_COLS)