o1-labs · dannywillems · Feb 29, 2024 · Feb 29, 2024
diff --git a/msm/src/constraint.rs b/msm/src/constraint.rs
@@ -9,8 +9,9 @@ use o1_utils::field_helpers::FieldHelpers;
 use o1_utils::foreign_field::ForeignElement;
 
 use crate::columns::{Column, ColumnIndexer, MSMColumnIndexer};
-use crate::proof::{Witness, WitnessColumns};
-use crate::{BN254G1Affine, Ff1, Fp, LIMBS_NUM};
+use crate::proof::ProofInputs;
+use crate::witness::Witness;
+use crate::{BN254G1Affine, Ff1, Fp, LIMBS_NUM, N};
 
 /// Used to represent constraints as multi variate polynomials. The variables
 /// are over the columns.
@@ -83,8 +84,8 @@ impl BuilderEnv<BN254G1Affine> {
     /// Each WitnessColumn stands for both one row and multirow. This
     /// function converts from a vector of one-row instantiation to a
     /// single multi-row form (which is a `Witness`).
-    pub fn get_witness(&self) -> Witness<BN254G1Affine> {
-        let mut x: Vec<Vec<Fp>> = vec![vec![]; 3 * LIMBS_NUM];
+    pub fn get_witness(&self) -> ProofInputs<N, BN254G1Affine> {
+        let mut cols: [Vec<Fp>; N] = std::array::from_fn(|_| vec![]);
 
         for wc in &self.witness_raw {
             let WitnessColumnsIndexer {
@@ -93,14 +94,14 @@ impl BuilderEnv<BN254G1Affine> {
                 c: wc_c,
             } = wc;
             for i in 0..LIMBS_NUM {
-                x[i].push(wc_a[i]);
-                x[LIMBS_NUM + i].push(wc_b[i]);
-                x[2 * LIMBS_NUM + i].push(wc_c[i]);
+                cols[i].push(wc_a[i]);
+                cols[LIMBS_NUM + i].push(wc_b[i]);
+                cols[2 * LIMBS_NUM + i].push(wc_c[i]);
             }
         }
 
-        Witness {
-            evaluations: WitnessColumns { x },
+        ProofInputs {
+            evaluations: Witness { cols },
             mvlookups: vec![],
         }
     }

diff --git a/msm/src/lib.rs b/msm/src/lib.rs
@@ -29,6 +29,11 @@ pub type BN254 = ark_ec::bn::Bn<ark_bn254::Parameters>;
 pub type BN254G1Affine = <BN254 as ark_ec::PairingEngine>::G1Affine;
 pub type BN254G2Affine = <BN254 as ark_ec::PairingEngine>::G2Affine;
 
+/// Number of columns
+/// FIXME: we must move it into the subdirectory of the
+/// foreign field addition circuit
+pub const N: usize = 3 * LIMBS_NUM;
+
 /// The native field we are working with.
 pub type Fp = ark_bn254::Fr;
 
@@ -48,10 +53,13 @@ mod tests {
     use poly_commitment::pairing_proof::PairingSRS;
 
     use crate::{
-        columns::Column, mvlookup::Lookup, proof::Witness, prover::prove, verifier::verify,
+        columns::Column, mvlookup::Lookup, proof::ProofInputs, prover::prove, verifier::verify,
         BaseSponge, Fp, OpeningProof, ScalarSponge, BN254,
     };
 
+    // Number of columns
+    const N: usize = 10;
+
     #[test]
     fn test_completeness() {
         let mut rng = o1_utils::tests::make_test_rng();
@@ -67,20 +75,20 @@ mod tests {
         let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
         srs.full_srs.add_lagrange_basis(domain.d1);
 
-        let witness = Witness::random(domain);
+        let inputs = ProofInputs::random(domain);
         let constraints: Vec<_> = vec![];
 
         // generate the proof
-        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
             domain,
             &srs,
-            witness,
+            inputs,
             constraints,
             &mut rng,
         );
 
         // verify the proof
-        let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
+        let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
         assert!(verifies);
     }
 
@@ -98,22 +106,22 @@ mod tests {
         let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
         srs.full_srs.add_lagrange_basis(domain.d1);
 
-        let witness = Witness::random(domain);
+        let inputs = ProofInputs::random(domain);
         let constraints = vec![];
         // generate the proof
-        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
             domain,
             &srs,
-            witness,
+            inputs,
             constraints.clone(),
             &mut rng,
         );
 
-        let witness_prime = Witness::random(domain);
-        let proof_prime = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+        let inputs_prime = ProofInputs::random(domain);
+        let proof_prime = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
             domain,
             &srs,
-            witness_prime,
+            inputs_prime,
             constraints,
             &mut rng,
         );
@@ -123,7 +131,7 @@ mod tests {
             let mut proof_clone = proof.clone();
             proof_clone.opening_proof = proof_prime.opening_proof;
             let verifies =
-                verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
+                verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
             assert!(!verifies);
         }
 
@@ -134,7 +142,7 @@ mod tests {
             let mut proof_clone = proof.clone();
             proof_clone.commitments = proof_prime.commitments;
             let verifies =
-                verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
+                verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
             assert!(!verifies);
         }
 
@@ -146,7 +154,7 @@ mod tests {
             let mut proof_clone = proof.clone();
             proof_clone.zeta_evaluations = proof_prime.zeta_evaluations;
             let verifies =
-                verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
+                verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
             assert!(!verifies);
         }
     }
@@ -166,27 +174,27 @@ mod tests {
         let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
         srs.full_srs.add_lagrange_basis(domain.d1);
 
-        let mut witness = Witness::random(domain);
+        let mut inputs = ProofInputs::random(domain);
         let constraints = vec![];
         // Take one random f_i (FIXME: taking first one for now)
-        let looked_up_values = witness.mvlookups[0].f[0].clone();
+        let looked_up_values = inputs.mvlookups[0].f[0].clone();
         // We change a random looked up element (FIXME: first one for now)
         let wrong_looked_up_value = Lookup {
             table_id: looked_up_values[0].table_id,
             numerator: looked_up_values[0].numerator,
             value: vec![Fp::rand(&mut rng)],
         };
         // Overwriting the first looked up value
-        witness.mvlookups[0].f[0][0] = wrong_looked_up_value;
+        inputs.mvlookups[0].f[0][0] = wrong_looked_up_value;
         // generate the proof
-        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+        let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
             domain,
             &srs,
-            witness,
+            inputs,
             constraints,
             &mut rng,
         );
-        let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
+        let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
         // FIXME: At the moment, it does verify. It should not. We are missing constraints.
         assert!(!verifies);
     }

diff --git a/msm/src/main.rs b/msm/src/main.rs
@@ -10,7 +10,7 @@ use kimchi_msm::precomputed_srs::get_bn254_srs;
 use kimchi_msm::prover::prove;
 use kimchi_msm::verifier::verify;
 use kimchi_msm::{
-    BN254G1Affine, BaseSponge, Ff1, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE,
+    BN254G1Affine, BaseSponge, Ff1, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE, N,
 };
 
 pub fn generate_random_msm_witness() -> BuilderEnv<BN254G1Affine> {
@@ -43,7 +43,7 @@ pub fn main() {
 
     println!("Generating the proof");
     let constraints = vec![];
-    let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+    let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
         domain,
         &srs,
         witness,
@@ -52,6 +52,6 @@ pub fn main() {
     );
 
     println!("Verifying the proof");
-    let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
+    let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
     println!("Proof verification result: {verifies}")
 }
diff --git a/msm/src/proof.rs b/msm/src/proof.rs
@@ -1,122 +1,42 @@
+use crate::witness::Witness;
 use ark_ff::UniformRand;
 use kimchi::{circuits::domains::EvaluationDomains, curve::KimchiCurve};
 use poly_commitment::{commitment::PolyComm, OpenProof};
-use rand::{prelude::*, thread_rng};
-use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
+use rand::thread_rng;
 
 use crate::mvlookup::{LookupProof, LookupWitness};
 
-/// List all columns of the circuit.
-/// It is parametrized by a type `T` which can be either:
-/// - `Vec<G::ScalarField>` for the evaluations
-/// - `PolyComm<G>` for the commitments
-#[derive(Debug, Clone)]
-pub struct WitnessColumns<T> {
-    pub x: Vec<T>,
-}
-
-impl<'lt, G> IntoIterator for &'lt WitnessColumns<G> {
-    type Item = &'lt G;
-    type IntoIter = std::vec::IntoIter<&'lt G>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        let n = self.x.len();
-        let mut iter_contents = Vec::with_capacity(n);
-        iter_contents.extend(&self.x);
-        iter_contents.into_iter()
-    }
-}
-
-impl<G> IntoParallelIterator for WitnessColumns<G>
-where
-    Vec<G>: IntoParallelIterator,
-{
-    type Iter = <Vec<G> as IntoParallelIterator>::Iter;
-    type Item = <Vec<G> as IntoParallelIterator>::Item;
-
-    fn into_par_iter(self) -> Self::Iter {
-        let n = self.x.len();
-        let mut iter_contents = Vec::with_capacity(n);
-        iter_contents.extend(self.x);
-        iter_contents.into_par_iter()
-    }
-}
-
-impl<G: Send + std::fmt::Debug> FromParallelIterator<G> for WitnessColumns<G> {
-    fn from_par_iter<I>(par_iter: I) -> Self
-    where
-        I: IntoParallelIterator<Item = G>,
-    {
-        let iter_contents = par_iter.into_par_iter().collect::<Vec<_>>();
-        WitnessColumns { x: iter_contents }
-    }
-}
-
-impl<'data, G> IntoParallelIterator for &'data WitnessColumns<G>
-where
-    Vec<&'data G>: IntoParallelIterator,
-{
-    type Iter = <Vec<&'data G> as IntoParallelIterator>::Iter;
-    type Item = <Vec<&'data G> as IntoParallelIterator>::Item;
-
-    fn into_par_iter(self) -> Self::Iter {
-        let n = self.x.len();
-        let mut iter_contents = Vec::with_capacity(n);
-        iter_contents.extend(self.x.iter());
-        iter_contents.into_par_iter()
-    }
-}
-
-impl<'data, G> IntoParallelIterator for &'data mut WitnessColumns<G>
-where
-    Vec<&'data mut G>: IntoParallelIterator,
-{
-    type Iter = <Vec<&'data mut G> as IntoParallelIterator>::Iter;
-    type Item = <Vec<&'data mut G> as IntoParallelIterator>::Item;
-
-    fn into_par_iter(self) -> Self::Iter {
-        let n = self.x.len();
-        let mut iter_contents = Vec::with_capacity(n);
-        iter_contents.extend(self.x.iter_mut());
-        iter_contents.into_par_iter()
-    }
-}
-
 #[derive(Debug)]
-pub struct Witness<G: KimchiCurve> {
-    pub evaluations: WitnessColumns<Vec<G::ScalarField>>,
+pub struct ProofInputs<const N: usize, G: KimchiCurve> {
+    pub evaluations: Witness<N, Vec<G::ScalarField>>,
     pub mvlookups: Vec<LookupWitness<G::ScalarField>>,
 }
 
 // This should be used only for testing purposes.
 // It is not only in the test API because it is used at the moment in the
 // main.rs. It should be moved to the test API when main.rs is replaced with
 // real production code.
-impl<G: KimchiCurve> Witness<G> {
+impl<const N: usize, G: KimchiCurve> ProofInputs<N, G> {
     pub fn random(domain: EvaluationDomains<G::ScalarField>) -> Self {
         let mut rng = thread_rng();
-        let random_n = rng.gen_range(1..1000);
-        Witness {
-            evaluations: WitnessColumns {
-                x: (0..random_n)
-                    .map(|_| {
-                        (0..domain.d1.size as usize)
-                            .map(|_| G::ScalarField::rand(&mut rng))
-                            .collect::<Vec<_>>()
-                    })
-                    .collect::<Vec<_>>(),
-            },
+        let cols: [Vec<G::ScalarField>; N] = std::array::from_fn(|_| {
+            (0..domain.d1.size as usize)
+                .map(|_| G::ScalarField::rand(&mut rng))
+                .collect::<Vec<_>>()
+        });
+        ProofInputs {
+            evaluations: Witness { cols },
             mvlookups: vec![LookupWitness::<G::ScalarField>::random(domain)],
         }
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct Proof<G: KimchiCurve, OpeningProof: OpenProof<G>> {
+pub struct Proof<const N: usize, G: KimchiCurve, OpeningProof: OpenProof<G>> {
     // Columns/PlonK argument
-    pub(crate) commitments: WitnessColumns<PolyComm<G>>,
-    pub(crate) zeta_evaluations: WitnessColumns<G::ScalarField>,
-    pub(crate) zeta_omega_evaluations: WitnessColumns<G::ScalarField>,
+    pub(crate) commitments: Witness<N, PolyComm<G>>,
+    pub(crate) zeta_evaluations: Witness<N, G::ScalarField>,
+    pub(crate) zeta_omega_evaluations: Witness<N, G::ScalarField>,
     // MVLookup argument
     pub(crate) mvlookup_commitments: Option<LookupProof<PolyComm<G>>>,
     pub(crate) mvlookup_zeta_evaluations: Option<LookupProof<G::ScalarField>>,