Merge pull request #2951 from o1-labs/feature/logup-values-tracking-a…

…gain

Track lookup values in interpreter and proof

o1vm/src/interpreters/mips/tests_helpers.rs

@@ -93,6 +93,8 @@ where
         scratch_state: [Fp::from(0); SCRATCH_SIZE],
         scratch_state_inverse: [Fp::from(0); SCRATCH_SIZE_INVERSE],
         lookup_multiplicities: LookupMultiplicities::new(),
+        lookup_state_idx: 0,
+        lookup_state: vec![],
         selector: crate::interpreters::mips::column::N_MIPS_SEL_COLS,
         halt: false,
         // Keccak related

o1vm/src/interpreters/mips/witness.rs

@@ -22,6 +22,7 @@ use crate::{
     },
     lookups::{Lookup, LookupTableIDs},
     preimage_oracle::PreImageOracleT,
+    ramlookup::LookupMode,
     utils::memory_size,
 };
 use ark_ff::{Field, PrimeField};
@@ -112,6 +113,8 @@ pub struct Env<Fp, PreImageOracle: PreImageOracleT> {
     pub scratch_state_idx_inverse: usize,
     pub scratch_state: [Fp; SCRATCH_SIZE],
     pub scratch_state_inverse: [Fp; SCRATCH_SIZE_INVERSE],
+    pub lookup_state_idx: usize,
+    pub lookup_state: Vec<Fp>,
     pub halt: bool,
     pub syscall_env: SyscallEnv,
     pub selector: usize,
@@ -175,9 +178,32 @@ impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp,
     }
 
     fn add_lookup(&mut self, lookup: Lookup<Self::Variable>) {
-        let values: Vec<_> = lookup.value.iter().map(|x| Fp::from(*x)).collect();
-        if let Some(idx) = lookup.table_id.ix_by_value(values.as_slice()) {
-            match lookup.table_id {
+        let mut add_value = |x: Fp| {
+            self.lookup_state_idx += 1;
+            self.lookup_state.push(x);
+        };
+        let Lookup {
+            table_id,
+            magnitude: numerator,
+            mode,
+            value: values,
+        } = lookup;
+        let values: Vec<_> = values.into_iter().map(|x| Fp::from(x)).collect();
+
+        // Add lookup numerator
+        match mode {
+            LookupMode::Write => add_value(Fp::from(numerator)),
+            LookupMode::Read => add_value(-Fp::from(numerator)),
+        };
+        // Add lookup table ID
+        add_value(Fp::from(table_id.to_u32()));
+        // Add values
+        for value in values.iter() {
+            add_value(*value);
+        }
+
+        if let Some(idx) = table_id.ix_by_value(values.as_slice()) {
+            match table_id {
                 LookupTableIDs::PadLookup => self.lookup_multiplicities.pad_lookup[idx] += 1,
                 LookupTableIDs::RoundConstantsLookup => {
                     self.lookup_multiplicities.round_constants_lookup[idx] += 1
@@ -935,6 +961,8 @@ impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
             scratch_state_idx_inverse: 0,
             scratch_state: fresh_scratch_state(),
             scratch_state_inverse: fresh_scratch_state(),
+            lookup_state_idx: 0,
+            lookup_state: vec![],
             halt: state.exited,
             syscall_env,
             selector,
@@ -959,6 +987,11 @@ impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
         self.scratch_state_inverse = fresh_scratch_state();
     }
 
+    pub fn reset_lookup_state(&mut self) {
+        self.lookup_state_idx = 0;
+        self.lookup_state = vec![];
+    }
+
     pub fn write_column(&mut self, column: Column, value: u64) {
         self.write_field_column(column, value.into())
     }
@@ -1202,6 +1235,7 @@ impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
     ) -> Instruction {
         self.reset_scratch_state();
         self.reset_scratch_state_inverse();
+        self.reset_lookup_state();
         let (opcode, _instruction) = self.decode_instruction();
 
         self.pp_info(&config.info_at, metadata, start);

o1vm/src/pickles/column_env.rs

@@ -18,6 +18,7 @@ type Evals<F> = Evaluations<F, Radix2EvaluationDomain<F>>;
 pub enum RelationColumnType {
     Scratch(usize),
     ScratchInverse(usize),
+    LookupState(usize),
     InstructionCounter,
     Error,
 }
@@ -43,16 +44,19 @@ pub struct ColumnEnvironment<'a, F: FftField> {
     pub domain: EvaluationDomains<F>,
 }
 
-pub fn get_all_columns() -> Vec<Column<RelationColumnType>> {
+pub fn get_all_columns(num_lookup_columns: usize) -> Vec<Column<RelationColumnType>> {
     let mut cols = Vec::<Column<RelationColumnType>>::with_capacity(
-        SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2 + N_MIPS_SEL_COLS,
+        SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + num_lookup_columns + 2 + N_MIPS_SEL_COLS,
     );
     for i in 0..SCRATCH_SIZE {
         cols.push(Column::Relation(RelationColumnType::Scratch(i)));
     }
     for i in 0..SCRATCH_SIZE_INVERSE {
         cols.push(Column::Relation(RelationColumnType::ScratchInverse(i)));
     }
+    for i in 0..num_lookup_columns {
+        cols.push(Column::Relation(RelationColumnType::LookupState(i)));
+    }
     cols.push(Column::Relation(RelationColumnType::InstructionCounter));
     cols.push(Column::Relation(RelationColumnType::Error));
     for i in 0..N_MIPS_SEL_COLS {
@@ -67,6 +71,7 @@ impl<G> WitnessColumns<G, [G; N_MIPS_SEL_COLS]> {
             Column::Relation(i) => match i {
                 RelationColumnType::Scratch(i) => Some(&self.scratch[i]),
                 RelationColumnType::ScratchInverse(i) => Some(&self.scratch_inverse[i]),
+                RelationColumnType::LookupState(i) => Some(&self.lookup_state[i]),
                 RelationColumnType::InstructionCounter => Some(&self.instruction_counter),
                 RelationColumnType::Error => Some(&self.error),
             },

o1vm/src/pickles/main.rs

@@ -1,4 +1,4 @@
-use ark_ff::UniformRand;
+use ark_ff::{UniformRand, Zero};
 use clap::Parser;
 use kimchi::circuits::domains::EvaluationDomains;
 use log::debug;
@@ -95,6 +95,28 @@ pub fn cannon_main(args: cli::cannon::RunArgs) {
         {
             scratch_chunk.push(*scratch);
         }
+        // Lookup state
+        {
+            let proof_inputs_length = curr_proof_inputs.evaluations.lookup_state.len();
+            let environment_length = mips_wit_env.lookup_state.len();
+            let lookup_state_size = std::cmp::max(proof_inputs_length, environment_length);
+            for idx in 0..lookup_state_size {
+                if idx >= environment_length {
+                    // We pad with 0s for dummy lookups missing from the environment.
+                    curr_proof_inputs.evaluations.lookup_state[idx].push(Fp::zero());
+                } else if idx >= proof_inputs_length {
+                    // We create a new column filled with 0s in the proof inputs.
+                    let mut new_vec =
+                        vec![Fp::zero(); curr_proof_inputs.evaluations.instruction_counter.len()];
+                    new_vec.push(Fp::from(mips_wit_env.lookup_state[idx]));
+                    curr_proof_inputs.evaluations.lookup_state.push(new_vec);
+                } else {
+                    // Push the value to the column.
+                    curr_proof_inputs.evaluations.lookup_state[idx]
+                        .push(Fp::from(mips_wit_env.lookup_state[idx]));
+                }
+            }
+        }
         curr_proof_inputs
             .evaluations
             .instruction_counter

o1vm/src/pickles/proof.rs

@@ -6,6 +6,7 @@ use crate::interpreters::mips::column::{N_MIPS_SEL_COLS, SCRATCH_SIZE, SCRATCH_S
 pub struct WitnessColumns<G, S> {
     pub scratch: [G; SCRATCH_SIZE],
     pub scratch_inverse: [G; SCRATCH_SIZE_INVERSE],
+    pub lookup_state: Vec<G>,
     pub instruction_counter: G,
     pub error: G,
     pub selector: S,
@@ -21,6 +22,7 @@ impl<G: KimchiCurve> ProofInputs<G> {
             evaluations: WitnessColumns {
                 scratch: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
                 scratch_inverse: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
+                lookup_state: vec![],
                 instruction_counter: Vec::with_capacity(domain_size),
                 error: Vec::with_capacity(domain_size),
                 selector: Vec::with_capacity(domain_size),

o1vm/src/pickles/prover.rs

@@ -86,6 +86,7 @@ where
         let WitnessColumns {
             scratch,
             scratch_inverse,
+            lookup_state,
             instruction_counter,
             error,
             selector,
@@ -119,10 +120,15 @@ where
                 eval_col(evals)
             })
             .collect::<Vec<_>>();
+        let lookup_state = lookup_state
+            .into_par_iter()
+            .map(eval_col)
+            .collect::<Vec<_>>();
         let selector = selector.into_par_iter().map(eval_col).collect::<Vec<_>>();
         WitnessColumns {
             scratch: scratch.try_into().unwrap(),
             scratch_inverse: scratch_inverse.try_into().unwrap(),
+            lookup_state,
             instruction_counter: eval_col(instruction_counter),
             error: eval_col(error.clone()),
             selector: selector.try_into().unwrap(),
@@ -134,6 +140,7 @@ where
         let WitnessColumns {
             scratch,
             scratch_inverse,
+            lookup_state,
             instruction_counter,
             error,
             selector,
@@ -151,10 +158,12 @@ where
         // Doing in parallel
         let scratch = scratch.par_iter().map(comm).collect::<Vec<_>>();
         let scratch_inverse = scratch_inverse.par_iter().map(comm).collect::<Vec<_>>();
+        let lookup_state = lookup_state.par_iter().map(comm).collect::<Vec<_>>();
         let selector = selector.par_iter().map(comm).collect::<Vec<_>>();
         WitnessColumns {
             scratch: scratch.try_into().unwrap(),
             scratch_inverse: scratch_inverse.try_into().unwrap(),
+            lookup_state,
             instruction_counter: comm(instruction_counter),
             error: comm(error),
             selector: selector.try_into().unwrap(),
@@ -170,6 +179,7 @@ where
         let WitnessColumns {
             scratch,
             scratch_inverse,
+            lookup_state,
             instruction_counter,
             error,
             selector,
@@ -182,10 +192,15 @@ where
             .into_par_iter()
             .map(eval_d8)
             .collect::<Vec<_>>();
+        let lookup_state = lookup_state
+            .into_par_iter()
+            .map(eval_d8)
+            .collect::<Vec<_>>();
         let selector = selector.into_par_iter().map(eval_d8).collect::<Vec<_>>();
         WitnessColumns {
             scratch: scratch.try_into().unwrap(),
             scratch_inverse: scratch_inverse.try_into().unwrap(),
+            lookup_state,
             instruction_counter: eval_d8(instruction_counter),
             error: eval_d8(error),
             selector: selector.try_into().unwrap(),
@@ -200,6 +215,9 @@ where
     for comm in commitments.scratch_inverse.iter() {
         absorb_commitment(&mut fq_sponge, comm)
     }
+    for comm in commitments.lookup_state.iter() {
+        absorb_commitment(&mut fq_sponge, comm)
+    }
     absorb_commitment(&mut fq_sponge, &commitments.instruction_counter);
     absorb_commitment(&mut fq_sponge, &commitments.error);
     for comm in commitments.selector.iter() {
@@ -314,17 +332,20 @@ where
         let WitnessColumns {
             scratch,
             scratch_inverse,
+            lookup_state,
             instruction_counter,
             error,
             selector,
         } = &polys;
         let eval = |poly: &DensePolynomial<G::ScalarField>| poly.evaluate(point);
         let scratch = scratch.par_iter().map(eval).collect::<Vec<_>>();
         let scratch_inverse = scratch_inverse.par_iter().map(eval).collect::<Vec<_>>();
+        let lookup_state = lookup_state.par_iter().map(eval).collect::<Vec<_>>();
         let selector = selector.par_iter().map(eval).collect::<Vec<_>>();
         WitnessColumns {
             scratch: scratch.try_into().unwrap(),
             scratch_inverse: scratch_inverse.try_into().unwrap(),
+            lookup_state,
             instruction_counter: eval(instruction_counter),
             error: eval(error),
             selector: selector.try_into().unwrap(),
@@ -375,6 +396,15 @@ where
         fr_sponge.absorb(zeta_eval);
         fr_sponge.absorb(zeta_omega_eval);
     }
+
+    for (zeta_eval, zeta_omega_eval) in zeta_evaluations
+        .lookup_state
+        .iter()
+        .zip(zeta_omega_evaluations.lookup_state.iter())
+    {
+        fr_sponge.absorb(zeta_eval);
+        fr_sponge.absorb(zeta_omega_eval);
+    }
     fr_sponge.absorb(&zeta_evaluations.instruction_counter);
     fr_sponge.absorb(&zeta_omega_evaluations.instruction_counter);
     fr_sponge.absorb(&zeta_evaluations.error);
@@ -401,6 +431,7 @@ where
 
     let mut polynomials: Vec<_> = polys.scratch.into_iter().collect();
     polynomials.extend(polys.scratch_inverse);
+    polynomials.extend(polys.lookup_state);
     polynomials.push(polys.instruction_counter);
     polynomials.push(polys.error);
     polynomials.extend(polys.selector);

o1vm/src/pickles/tests.rs

@@ -71,6 +71,7 @@ fn test_small_circuit() {
         evaluations: WitnessColumns {
             scratch: std::array::from_fn(|_| zero_to_n_minus_one(8)),
             scratch_inverse: std::array::from_fn(|_| (0..8).map(|_| Fq::zero()).collect()),
+            lookup_state: vec![],
             instruction_counter: zero_to_n_minus_one(8)
                 .into_iter()
                 .map(|x| x + Fq::one())

o1vm/src/pickles/verifier.rs

@@ -99,6 +99,9 @@ where
     for comm in commitments.scratch_inverse.iter() {
         absorb_commitment(&mut fq_sponge, comm)
     }
+    for comm in commitments.lookup_state.iter() {
+        absorb_commitment(&mut fq_sponge, comm)
+    }
     absorb_commitment(&mut fq_sponge, &commitments.instruction_counter);
     absorb_commitment(&mut fq_sponge, &commitments.error);
     for comm in commitments.selector.iter() {
@@ -148,6 +151,14 @@ where
         fr_sponge.absorb(zeta_eval);
         fr_sponge.absorb(zeta_omega_eval);
     }
+    for (zeta_eval, zeta_omega_eval) in zeta_evaluations
+        .lookup_state
+        .iter()
+        .zip(zeta_omega_evaluations.lookup_state.iter())
+    {
+        fr_sponge.absorb(zeta_eval);
+        fr_sponge.absorb(zeta_omega_eval);
+    }
     fr_sponge.absorb(&zeta_evaluations.instruction_counter);
     fr_sponge.absorb(&zeta_omega_evaluations.instruction_counter);
     fr_sponge.absorb(&zeta_evaluations.error);
@@ -204,7 +215,7 @@ where
     let u_chal = fr_sponge.challenge();
     let u = u_chal.to_field(endo_r);
 
-    let mut evaluations: Vec<_> = get_all_columns()
+    let mut evaluations: Vec<_> = get_all_columns(column_eval.commitment.lookup_state.len())
         .into_iter()
         .map(|column| {
             let commitment = column_eval