Merge pull request #2604 from o1-labs/o1vm/column-env-pickles

o1vm/pickles: import column_env

o1vm/src/pickles/column_env.rs

@@ -0,0 +1,102 @@
+use ark_ff::FftField;
+use ark_poly::{Evaluations, Radix2EvaluationDomain};
+
+use crate::interpreters::mips::witness::SCRATCH_SIZE;
+use kimchi::circuits::{
+    domains::EvaluationDomains,
+    expr::{
+        BerkeleyChallengeTerm, BerkeleyChallenges, ColumnEnvironment as TColumnEnvironment,
+        Constants, Domain,
+    },
+};
+
+use super::proof::WitnessColumns;
+
+/// The collection of polynomials (all in evaluation form) and constants
+/// required to evaluate an expression as a polynomial.
+///
+/// All are evaluations.
+pub struct ColumnEnvironment<'a, F: FftField> {
+    /// The witness column polynomials. Includes relation columns and dynamic
+    /// selector columns.
+    pub witness: &'a WitnessColumns<Evaluations<F, Radix2EvaluationDomain<F>>>,
+    /// The value `prod_{j != 1} (1 - ω^j)`, used for efficiently
+    /// computing the evaluations of the unnormalized Lagrange basis
+    /// polynomials.
+    pub l0_1: F,
+    /// Constant values required
+    pub constants: Constants<F>,
+    /// Challenges from the IOP.
+    // FIXME: change for other challenges
+    pub challenges: BerkeleyChallenges<F>,
+    /// The domains used in the PLONK argument.
+    pub domain: EvaluationDomains<F>,
+}
+
+impl<'a, F: FftField> TColumnEnvironment<'a, F, BerkeleyChallengeTerm, BerkeleyChallenges<F>>
+    for ColumnEnvironment<'a, F>
+{
+    // FIXME: do we change to the MIPS column type?
+    // We do not want to keep kimchi_msm/generic prover
+    type Column = kimchi_msm::columns::Column;
+
+    fn get_column(
+        &self,
+        col: &Self::Column,
+    ) -> Option<&'a Evaluations<F, Radix2EvaluationDomain<F>>> {
+        match *col {
+            Self::Column::Relation(i) => {
+                if i < SCRATCH_SIZE {
+                    let res = &self.witness.scratch[i];
+                    Some(res)
+                } else if i == SCRATCH_SIZE {
+                    let res = &self.witness.instruction_counter;
+                    Some(res)
+                } else if i == SCRATCH_SIZE + 1 {
+                    let res = &self.witness.error;
+                    Some(res)
+                } else {
+                    panic!("We should not have that many relation columns");
+                }
+            }
+            Self::Column::DynamicSelector(_i) => {
+                // FIXME: add selectors
+                panic!("Not implemented yet");
+            }
+            _ => {
+                panic!("We should not have any other type of columns")
+            }
+        }
+    }
+
+    fn get_domain(&self, d: Domain) -> Radix2EvaluationDomain<F> {
+        match d {
+            Domain::D1 => self.domain.d1,
+            Domain::D2 => self.domain.d2,
+            Domain::D4 => self.domain.d4,
+            Domain::D8 => self.domain.d8,
+        }
+    }
+
+    fn column_domain(&self, _col: &Self::Column) -> Domain {
+        Domain::D8
+    }
+
+    fn get_constants(&self) -> &Constants<F> {
+        &self.constants
+    }
+
+    fn get_challenges(&self) -> &BerkeleyChallenges<F> {
+        &self.challenges
+    }
+
+    fn vanishes_on_zero_knowledge_and_previous_rows(
+        &self,
+    ) -> &'a Evaluations<F, Radix2EvaluationDomain<F>> {
+        panic!("Not supposed to be used in MIPS. We do not support zero-knowledge for now")
+    }
+
+    fn l0_1(&self) -> F {
+        self.l0_1
+    }
+}

o1vm/src/pickles/main.rs

@@ -79,7 +79,7 @@ pub fn main() -> ExitCode {
         mips_witness::Env::<Fp, PreImageOracle>::create(cannon::PAGE_SIZE as usize, state, po);
 
     // TODO: give this to the prover + verifier
-    let _constraints = {
+    let constraints = {
         let mut mips_con_env = mips_constraints::Env::<Fp>::default();
         let mut constraints = Instruction::iter()
             .flat_map(|instr_typ| instr_typ.into_iter())
@@ -126,12 +126,13 @@ pub fn main() -> ExitCode {
             // FIXME
             let start_iteration = Instant::now();
             debug!("Limit of {DOMAIN_SIZE} reached. We make a proof, verify it (for testing) and start with a new branch new chunk");
-            let _proof: Proof<Vesta> = prover::prove::<
-                Vesta,
-                DefaultFqSponge<VestaParameters, PlonkSpongeConstantsKimchi>,
-                DefaultFrSponge<Fp, PlonkSpongeConstantsKimchi>,
-                _,
-            >(domain_fp, &srs, curr_proof_inputs, &mut rng);
+            let _proof: Proof<Vesta> =
+                prover::prove::<
+                    Vesta,
+                    DefaultFqSponge<VestaParameters, PlonkSpongeConstantsKimchi>,
+                    DefaultFrSponge<Fp, PlonkSpongeConstantsKimchi>,
+                    _,
+                >(domain_fp, &srs, curr_proof_inputs, &constraints, &mut rng);
             debug!(
                 "Proof generated in {elapsed} μs",
                 elapsed = start_iteration.elapsed().as_micros()

o1vm/src/pickles/mod.rs

@@ -12,6 +12,7 @@
 //! O1VM_FLAVOR=pickles bash run-code.sh
 //! ```
 
+pub mod column_env;
 pub mod proof;
 pub mod prover;
 

o1vm/src/pickles/prover.rs

@@ -2,7 +2,12 @@ use ark_ec::AffineRepr;
 use ark_ff::{PrimeField, Zero};
 use ark_poly::{univariate::DensePolynomial, Evaluations, Polynomial, Radix2EvaluationDomain as D};
 use kimchi::{
-    circuits::domains::EvaluationDomains, curve::KimchiCurve, groupmap::GroupMap,
+    circuits::{
+        domains::EvaluationDomains,
+        expr::{l0_1, BerkeleyChallenges, Constants},
+    },
+    curve::KimchiCurve,
+    groupmap::GroupMap,
     plonk_sponge::FrSponge,
 };
 use mina_poseidon::{sponge::ScalarChallenge, FqSponge};
@@ -14,7 +19,11 @@ use poly_commitment::{
 use rand::{CryptoRng, RngCore};
 use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
-use super::proof::{Proof, ProofInputs, WitnessColumns};
+use super::{
+    column_env::ColumnEnvironment,
+    proof::{Proof, ProofInputs, WitnessColumns},
+};
+use crate::E;
 
 /// Make a PlonKish proof for the given circuit. As inputs, we get the execution
 /// trace consisting of evaluations of polynomials over a certain domain
@@ -28,7 +37,6 @@ use super::proof::{Proof, ProofInputs, WitnessColumns};
 ///
 /// The final proof consists of the opening proof, the commitments and the
 /// evaluations at ζ and ζω.
-// FIXME: add constraints (come with the quotient polynomial)
 // TODO: we might need blinders when the evaluation of columns are zeroes.
 pub fn prove<
     G: KimchiCurve,
@@ -39,6 +47,7 @@ pub fn prove<
     domain: EvaluationDomains<G::ScalarField>,
     srs: &SRS<G>,
     inputs: ProofInputs<G>,
+    _constraints: &[E<G::ScalarField>],
     rng: &mut RNG,
 ) -> Proof<G>
 where
@@ -105,7 +114,7 @@ where
     // Based on the regression test
     // `test_regression_constraints_with_selectors`, the highest degree is 6.
     // Therefore, we do evaluate on d8.
-    let _evaluations_d8 = {
+    let evaluations_d8 = {
         let WitnessColumns {
             scratch,
             instruction_counter,
@@ -138,14 +147,42 @@ where
     // Round 2: Creating and committing to the quotient polynomial
     ////////////////////////////////////////////////////////////////////////////
 
+    let (_, endo_r) = G::endos();
+
+    // Constraints combiner
+    let alpha: G::ScalarField = fq_sponge.challenge();
+
+    let zk_rows = 0;
+    let _column_env: ColumnEnvironment<'_, G::ScalarField> = {
+        // FIXME: use a proper Challenge structure
+        let challenges = BerkeleyChallenges {
+            alpha,
+            // No permutation argument for the moment
+            beta: G::ScalarField::zero(),
+            gamma: G::ScalarField::zero(),
+            // No lookup for the moment
+            joint_combiner: G::ScalarField::zero(),
+        };
+        ColumnEnvironment {
+            constants: Constants {
+                endo_coefficient: *endo_r,
+                mds: &G::sponge_params().mds,
+                zk_rows,
+            },
+            challenges,
+            witness: &evaluations_d8,
+            l0_1: l0_1(domain.d1),
+            domain,
+        }
+    };
+
     // FIXME: add quotient polynomial
 
     ////////////////////////////////////////////////////////////////////////////
     // Round 3: Evaluations at ζ and ζω
     ////////////////////////////////////////////////////////////////////////////
 
     let zeta_chal = ScalarChallenge(fq_sponge.challenge());
-    let (_, endo_r) = G::endos();
 
     let zeta = zeta_chal.to_field(endo_r);
     let zeta_omega = zeta * omega;
@@ -209,9 +246,7 @@ where
             (
                 DensePolynomialOrEvaluations::DensePolynomial(poly),
                 // We do not have any blinder, therefore we set to 0.
-                PolyComm {
-                    elems: vec![G::ScalarField::zero()],
-                },
+                PolyComm::new(vec![G::ScalarField::zero()]),
             )
         })
         .collect();