WIP - MSM: use constant size witness and rename old witness in ProofI…

…nputs

msm/src/proof.rs

@@ -1,122 +1,43 @@
+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 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, 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>>,

msm/src/prover.rs

@@ -17,7 +17,8 @@ use rayon::iter::IntoParallelIterator;
 use rayon::iter::ParallelIterator;
 
 use crate::mvlookup::{self, LookupProof};
-use crate::proof::{Proof, Witness, WitnessColumns};
+use crate::proof::{Proof, ProofInputs};
+use crate::witness::Witness;
 
 pub fn prove<
     G: KimchiCurve,
@@ -26,41 +27,39 @@ pub fn prove<
     EFrSponge: FrSponge<G::ScalarField>,
     Column,
     RNG,
+    const N: usize,
 >(
     domain: EvaluationDomains<G::ScalarField>,
     srs: &OpeningProof::SRS,
-    inputs: Witness<G>,
+    inputs: ProofInputs<N, G>,
     _constraints: Vec<Expr<ConstantExpr<G::ScalarField>, Column>>,
     rng: &mut RNG,
-) -> Proof<G, OpeningProof>
+) -> Proof<N, G, OpeningProof>
 where
     OpeningProof::SRS: Sync,
     RNG: RngCore + CryptoRng,
 {
     // Interpolate all columns on d1, using trait Into.
-    let evaluations: WitnessColumns<Evaluations<G::ScalarField, D<G::ScalarField>>> = inputs
+    let evaluations: Witness<N, Evaluations<G::ScalarField, D<G::ScalarField>>> = inputs
         .evaluations
         .into_par_iter()
         .map(|evals| {
             Evaluations::<G::ScalarField, D<G::ScalarField>>::from_vec_and_domain(evals, domain.d1)
         })
-        .collect::<WitnessColumns<Evaluations<G::ScalarField, D<G::ScalarField>>>>();
+        .collect::<Witness<N, Evaluations<G::ScalarField, D<G::ScalarField>>>>();
 
-    let polys: WitnessColumns<DensePolynomial<G::ScalarField>> = {
+    let polys: Witness<N, DensePolynomial<G::ScalarField>> = {
         let interpolate =
             |evals: Evaluations<G::ScalarField, D<G::ScalarField>>| evals.interpolate();
         evaluations
             .into_par_iter()
             .map(interpolate)
-            .collect::<WitnessColumns<_>>()
+            .collect::<Witness<N, _>>()
     };
 
-    let commitments: WitnessColumns<PolyComm<G>> = {
+    let commitments: Witness<N, PolyComm<G>> = {
         let comm = |poly: &DensePolynomial<G::ScalarField>| srs.commit_non_hiding(poly, 1);
-        (&polys)
-            .into_par_iter()
-            .map(comm)
-            .collect::<WitnessColumns<_>>()
+        (&polys).into_par_iter().map(comm).collect::<Witness<_>>()
     };
 
     let mut fq_sponge = EFqSponge::new(G::other_curve_sponge_params());
@@ -97,10 +96,10 @@ where
     // TODO: Parallelize
     let (zeta_evaluations, zeta_omega_evaluations) = {
         let evals = |point| {
-            let WitnessColumns { x } = &polys;
+            let Witness { cols } = &polys;
             let comm = |poly: &DensePolynomial<G::ScalarField>| poly.evaluate(point);
-            let x = x.iter().map(comm).collect::<Vec<_>>();
-            WitnessColumns { x }
+            let x = cols.iter().map(comm).collect::<Vec<_>>();
+            Witness { cols }
         };
         (evals(&zeta), evals(&zeta_omega))
     };

msm/src/serialization/main.rs

@@ -4,17 +4,19 @@ use kimchi_msm::witness::Witness;
 use poly_commitment::pairing_proof::PairingSRS;
 
 use kimchi_msm::precomputed_srs::get_bn254_srs;
+use kimchi_msm::prover::prove;
 use kimchi_msm::serialization::witness::deserialize_field_element;
-use kimchi_msm::{Fp, BN254, DOMAIN_SIZE, LIMBS_NUM};
+use kimchi_msm::verifier::verify;
+use kimchi_msm::{BaseSponge, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE, LIMBS_NUM};
 
 pub fn main() {
     // FIXME: use a proper RNG
-    let mut _rng = o1_utils::tests::make_test_rng();
+    let mut rng = o1_utils::tests::make_test_rng();
 
     println!("Creating the domain and SRS");
     let domain = EvaluationDomains::<Fp>::create(DOMAIN_SIZE).unwrap();
 
-    let _srs: PairingSRS<BN254> = get_bn254_srs(domain);
+    let srs: PairingSRS<BN254> = get_bn254_srs(domain);
 
     let mut env = witness::Env::<Fp>::create();
     let mut witness: Witness<DOMAIN_SIZE, Vec<Fp>> = Witness {
@@ -36,16 +38,22 @@ pub fn main() {
         }
     }
 
-    // println!("Generating the proof");
-    // let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
-    //     domain,
-    //     &srs,
-    //     witness,
-    //     constraints,
-    //     &mut rng,
-    // );
-
-    // println!("Verifying the proof");
-    // let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
-    // println!("Proof verification result: {verifies}")
+    let _constraints = vec![];
+    let proof_inputs = ProofInputs {
+        evaluations: witness,
+        mvlookups: vec![],
+    };
+
+    println!("Generating the proof");
+    let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
+        domain,
+        &srs,
+        proof_inputs,
+        _constraints,
+        &mut rng,
+    );
+
+    println!("Verifying the proof");
+    let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
+    println!("Proof verification result: {verifies}")
 }

msm/src/verifier.rs

@@ -16,10 +16,11 @@ pub fn verify<
     OpeningProof: OpenProof<G>,
     EFqSponge: Clone + FqSponge<G::BaseField, G, G::ScalarField>,
     EFrSponge: FrSponge<G::ScalarField>,
+    const N: usize,
 >(
     domain: EvaluationDomains<G::ScalarField>,
     srs: &OpeningProof::SRS,
-    proof: &Proof<G, OpeningProof>,
+    proof: &Proof<N, G, OpeningProof>,
 ) -> bool {
     let Proof {
         commitments,