Logup for o1vm

msm/src/logup.rs

@@ -275,13 +275,13 @@ pub struct LogupWitness<F, ID: LookupTableID> {
 #[derive(Debug, Clone)]
 pub struct LookupProof<T, ID> {
     /// The multiplicity polynomials
-    pub(crate) m: BTreeMap<ID, Vec<T>>,
+    pub m: BTreeMap<ID, Vec<T>>,
     /// The polynomial keeping the sum of each row
-    pub(crate) h: BTreeMap<ID, Vec<T>>,
+    pub h: BTreeMap<ID, Vec<T>>,
     /// The "running-sum" over the rows, coined `φ`
-    pub(crate) sum: T,
+    pub sum: T,
     /// All fixed lookup tables values, indexed by their ID
-    pub(crate) fixed_tables: BTreeMap<ID, T>,
+    pub fixed_tables: BTreeMap<ID, T>,
 }
 
 /// Iterator implementation to abstract the content of the structure.
@@ -437,6 +437,7 @@ pub fn constraint_lookups<F: PrimeField, ID: LookupTableID>(
             })
             .collect();
 
+        // isn't it supposed to be not table_id.runtime_create_column()
         if table_id.is_fixed() || table_id.runtime_create_column() {
             let table_lookup = Logup {
                 table_id: *table_id,
@@ -512,6 +513,7 @@ pub mod prover {
     }
 
     /// Represents the environment for the logup argument.
+    #[derive(Clone)]
     pub struct Env<G: KimchiCurve, ID: LookupTableID> {
         /// The polynomial of the multiplicities, indexed by the table ID.
         pub lookup_counters_poly_d1: BTreeMap<ID, Vec<DensePolynomial<G::ScalarField>>>,

o1vm/src/interpreters/mips/witness.rs

@@ -20,16 +20,18 @@ use crate::{
             registers::Registers,
         },
     },
-    lookups::Lookup,
+    lookups::{FixedLookupTables, Lookup, LookupTable, LookupTableIDs},
     preimage_oracle::PreImageOracleT,
     utils::memory_size,
 };
 use ark_ff::Field;
 use core::panic;
 use kimchi::o1_utils::Two;
+use kimchi_msm::LogupTableID as _;
 use log::{debug, info};
 use std::{
     array,
+    collections::HashMap,
     fs::File,
     io::{BufWriter, Write},
 };
@@ -64,6 +66,10 @@ impl SyscallEnv {
     }
 }
 
+// Some type aliases to aid in refactoring to something faster later
+type FixedTableMap<ID, F> = HashMap<ID, LookupTable<F>>;
+type MultiplicityMap<ID> = HashMap<ID, Vec<u32>>;
+
 /// This structure represents the environment the virtual machine state will use
 /// to transition. This environment will be used by the interpreter. The virtual
 /// machine has access to its internal state and some external memory. In
@@ -90,6 +96,8 @@ pub struct Env<Fp, PreImageOracle: PreImageOracleT> {
     pub preimage_key: Option<[u8; 32]>,
     pub keccak_env: Option<KeccakEnv<Fp>>,
     pub hash_counter: u64,
+    pub lookup_fixed_tables: FixedTableMap<LookupTableIDs, Fp>,
+    pub lookup_multiplicities: MultiplicityMap<LookupTableIDs>,
 }
 
 fn fresh_scratch_state<Fp: Field, const N: usize>() -> [Fp; N] {
@@ -142,9 +150,18 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreI
         }
     }
 
-    fn add_lookup(&mut self, _lookup: Lookup<Self::Variable>) {
-        // No-op, constraints only
-        // TODO: keep track of multiplicities of fixed tables here as in Keccak?
+    fn add_lookup(&mut self, lookup: Lookup<Self::Variable>) {
+        if let Some(idx) = LookupTable::is_in_table(&lookup.table_id, lookup.value) {
+            // We found the table, so just add one to the multiplicity.
+            self.lookup_multiplicities
+                .get_mut(&lookup.table_id)
+                .unwrap()[idx] += 1;
+        } else {
+            panic!(
+                "Tried to lookup in non-fixed table: {:?}",
+                LookupTableIDs::from_u32(lookup.table_id)
+            );
+        }
     }
 
     fn instruction_counter(&self) -> Self::Variable {
@@ -867,6 +884,22 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
             }
         };
 
+        let (lookup_fixed_tables, lookup_multiplicities) = {
+            let mut ft = HashMap::new();
+            let mut m = HashMap::new();
+            let fixed_tables = LookupTableIDs::all_variants()
+                .into_iter()
+                .filter(|x| x.is_fixed())
+                .collect::<Vec<_>>();
+
+            for table_id in fixed_tables {
+                ft.insert(table_id, table_id.to_fixed_table::<Fp>());
+                m.insert(table_id, table_id.to_multiplicities_vec());
+            }
+
+            (ft, m)
+        };
+
         Env {
             instruction_counter: state.step,
             memory: initial_memory.clone(),
@@ -891,6 +924,8 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
             preimage_key: None,
             keccak_env: None,
             hash_counter: 0,
+            lookup_fixed_tables,
+            lookup_multiplicities,
         }
     }
 

o1vm/src/lookups.rs

@@ -1,7 +1,12 @@
 //! Instantiation of the lookups for the VM project.
 
+use std::collections::BTreeMap;
+
 use self::LookupTableIDs::*;
-use crate::{interpreters::keccak::pad_blocks, ramlookup::RAMLookup};
+use crate::{
+    interpreters::keccak::pad_blocks,
+    ramlookup::{LookupMode, RAMLookup},
+};
 use ark_ff::{Field, PrimeField};
 use kimchi::{
     circuits::polynomials::keccak::{
@@ -97,7 +102,11 @@ impl LookupTableID for LookupTableIDs {
     }
 
     fn runtime_create_column(&self) -> bool {
-        panic!("No runtime tables specified");
+        if self.is_fixed() {
+            panic!("No runtime tables specified")
+        } else {
+            false
+        }
     }
 
     fn ix_by_value<F: PrimeField>(&self, _value: &[F]) -> Option<usize> {
@@ -121,6 +130,41 @@ impl LookupTableID for LookupTableIDs {
     }
 }
 
+impl LookupTableIDs {
+    /// Panics if the requested table is not fixed!
+    pub fn to_fixed_table<F: Field>(&self) -> LookupTable<F> {
+        match self {
+            Self::PadLookup => LookupTable::<F>::table_pad(),
+            Self::RoundConstantsLookup => LookupTable::<F>::table_round_constants(),
+            Self::AtMost4Lookup => LookupTable::<F>::table_at_most_4(),
+            Self::ByteLookup => LookupTable::<F>::table_byte(),
+            Self::RangeCheck16Lookup => LookupTable::<F>::table_range_check_16(),
+            Self::SparseLookup => LookupTable::<F>::table_sparse(),
+            Self::ResetLookup => LookupTable::<F>::table_reset(),
+            _ => panic!(
+                "Table {:?} is not fixed, so cannot called to_fixed_table on it!",
+                self
+            ),
+        }
+    }
+
+    pub fn to_multiplicities_vec(self) -> Vec<u32> {
+        match self {
+            Self::PadLookup => vec![0; PadLookup.length()],
+            Self::RoundConstantsLookup => vec![0; RoundConstantsLookup.length()],
+            Self::AtMost4Lookup => vec![0; AtMost4Lookup.length()],
+            Self::ByteLookup => vec![0; ByteLookup.length()],
+            Self::RangeCheck16Lookup => vec![0; RangeCheck16Lookup.length()],
+            Self::SparseLookup => vec![0; SparseLookup.length()],
+            Self::ResetLookup => vec![0; ResetLookup.length()],
+            _ => panic!(
+                "Table {:?} is not fixed, so cannot called to_multiplicities_vec on it!",
+                self
+            ),
+        }
+    }
+}
+
 /// Trait that creates all the fixed lookup tables used in the VM
 pub(crate) trait FixedLookupTables<F> {
     /// Checks whether a value is in a table and returns the position if it is or None otherwise.
@@ -276,3 +320,33 @@ impl<F: Field> FixedLookupTables<F> for LookupTable<F> {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct PartitionedLookups<T, ID: LookupTableID> {
+    pub reads: BTreeMap<ID, Vec<Vec<T>>>,
+    pub writes: BTreeMap<ID, Vec<Vec<T>>>,
+}
+
+pub fn partition_lookups<T, ID: LookupTableID>(
+    lookups: Vec<RAMLookup<T, ID>>,
+) -> PartitionedLookups<T, ID> {
+    let mut reads = BTreeMap::new();
+    let mut writes = BTreeMap::new();
+
+    let insert_with_id = |lookup: RAMLookup<T, ID>, table: &mut BTreeMap<ID, Vec<Vec<T>>>| {
+        if let Some(old_vec) = table.get_mut(&lookup.table_id) {
+            old_vec.push(lookup.value)
+        } else {
+            let _ = table.insert(lookup.table_id, vec![lookup.value]);
+        }
+    };
+
+    for lookup in lookups {
+        match lookup.mode {
+            LookupMode::Read => insert_with_id(lookup, &mut reads),
+            LookupMode::Write => insert_with_id(lookup, &mut writes),
+        }
+    }
+
+    PartitionedLookups { reads, writes }
+}

o1vm/src/pickles/column_env.rs

@@ -1,6 +1,6 @@
 use ark_ff::FftField;
 use ark_poly::{Evaluations, Radix2EvaluationDomain};
-use kimchi_msm::columns::Column;
+use kimchi_msm::{columns::Column, logup::prover::QuotientPolynomialEnvironment, LookupTableID};
 
 use crate::{
     interpreters::mips::column::{N_MIPS_SEL_COLS, SCRATCH_SIZE, SCRATCH_SIZE_INVERSE},
@@ -18,10 +18,10 @@ type Evals<F> = Evaluations<F, Radix2EvaluationDomain<F>>;
 /// required to evaluate an expression as a polynomial.
 ///
 /// All are evaluations.
-pub struct ColumnEnvironment<'a, F: FftField> {
+pub struct ColumnEnvironment<'a, F: FftField, ID: LookupTableID> {
     /// The witness column polynomials. Includes relation columns and dynamic
     /// selector columns.
-    pub witness: &'a WitnessColumns<Evals<F>, [Evals<F>; N_MIPS_SEL_COLS]>,
+    pub witness: &'a WitnessColumns<Evals<F>, [Evals<F>; N_MIPS_SEL_COLS], ID>,
     /// The value `prod_{j != 1} (1 - ω^j)`, used for efficiently
     /// computing the evaluations of the unnormalized Lagrange basis
     /// polynomials.
@@ -33,6 +33,9 @@ pub struct ColumnEnvironment<'a, F: FftField> {
     pub challenges: BerkeleyChallenges<F>,
     /// The domains used in the PLONK argument.
     pub domain: EvaluationDomains<F>,
+
+    /// Lookup specific polynomials
+    pub lookup: Option<QuotientPolynomialEnvironment<'a, F, ID>>,
 }
 
 pub fn get_all_columns() -> Vec<Column> {
@@ -47,55 +50,68 @@ pub fn get_all_columns() -> Vec<Column> {
     cols
 }
 
-impl<G> WitnessColumns<G, [G; N_MIPS_SEL_COLS]> {
-    pub fn get_column(&self, col: &Column) -> Option<&G> {
-        match *col {
-            Column::Relation(i) => {
-                if i < SCRATCH_SIZE {
-                    let res = &self.scratch[i];
-                    Some(res)
-                } else if i < SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
-                    let res = &self.scratch_inverse[i - SCRATCH_SIZE];
-                    Some(res)
-                } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
-                    let res = &self.instruction_counter;
-                    Some(res)
-                } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 {
-                    let res = &self.error;
-                    Some(res)
-                } else {
-                    panic!("We should not have that many relation columns. We have {} columns and index {} was given", SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2, i);
-                }
-            }
-            Column::DynamicSelector(i) => {
-                assert!(
-                    i < N_MIPS_SEL_COLS,
-                    "We do not have that many dynamic selector columns. We have {} columns and index {} was given",
-                    N_MIPS_SEL_COLS,
-                    i
-                );
-                let res = &self.selector[i];
+pub fn get_column<'a, F: Clone, ID: LookupTableID>(
+    env: &'a WitnessColumns<F, [F; N_MIPS_SEL_COLS], ID>,
+    col: &Column,
+) -> Option<&'a F> {
+    match *col {
+        Column::Relation(i) => {
+            if i < SCRATCH_SIZE {
+                let res = &env.scratch[i];
                 Some(res)
+            } else if i < SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
+                let res = &env.scratch_inverse[i - SCRATCH_SIZE];
+                Some(res)
+            } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE {
+                let res = &env.instruction_counter;
+                Some(res)
+            } else if i == SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 {
+                let res = &env.error;
+                Some(res)
+            } else {
+                panic!("We should not have that many relation columns. We have {} columns and index {} was given", SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 2, i);
             }
-            _ => {
-                panic!(
-                    "We should not have any other type of columns. The column {:?} was given",
-                    col
-                );
-            }
+        }
+        Column::DynamicSelector(i) => {
+            assert!(
+                i < N_MIPS_SEL_COLS,
+                "We do not have that many dynamic selector columns, given {i} but only have {N_MIPS_SEL_COLS}"
+            );
+            let res = &env.selector[i];
+            Some(res)
+        }
+        Column::LookupPartialSum((table_id, i)) => {
+            let table_id = ID::from_u32(table_id);
+            Some(&env.lookup[&table_id].f[i])
+        }
+        Column::LookupAggregation => Some(&env.lookup_agg),
+        Column::LookupMultiplicity((table_id, i)) => {
+            let table_id = ID::from_u32(table_id);
+            Some(&env.lookup[&table_id].m[i])
+        }
+        Column::LookupFixedTable(table_id) => {
+            let table_id = ID::from_u32(table_id);
+            Some(&env.lookup[&table_id].t)
+        }
+        _ => {
+            panic!(
+                "We should not have any other type of columns. The column {:?} was given",
+                col
+            )
         }
     }
 }
 
-impl<'a, F: FftField> TColumnEnvironment<'a, F, BerkeleyChallengeTerm, BerkeleyChallenges<F>>
-    for ColumnEnvironment<'a, F>
+impl<'a, F: FftField, ID: LookupTableID>
+    TColumnEnvironment<'a, F, BerkeleyChallengeTerm, BerkeleyChallenges<F>>
+    for ColumnEnvironment<'a, F, ID>
 {
     // FIXME: do we change to the MIPS column type?
     // We do not want to keep kimchi_msm/generic prover
     type Column = Column;
 
     fn get_column(&self, col: &Self::Column) -> Option<&'a Evals<F>> {
-        self.witness.get_column(col)
+        get_column(self.witness, col)
     }
 
     fn get_domain(&self, d: Domain) -> Radix2EvaluationDomain<F> {