Logup for o1vm

kimchi/src/circuits/lookup/lookups.rs

@@ -300,7 +300,7 @@ impl<F: Copy> SingleLookup<F> {
 }
 
 /// The table ID associated with a particular lookup
-#[derive(Clone, Serialize, Deserialize, Debug)]
+#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum LookupTableID {
     /// Look up the value from the given fixed table ID
     Constant(i32),

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/column.rs

@@ -1,9 +1,13 @@
 use crate::interpreters::mips::Instruction::{self, IType, JType, RType};
+use kimchi::circuits::lookup::lookups::LookupTableID;
 use kimchi_msm::{
     columns::{Column, ColumnIndexer},
     witness::Witness,
 };
-use std::ops::{Index, IndexMut};
+use std::{
+    collections::HashMap,
+    ops::{Index, IndexMut},
+};
 use strum::EnumCount;
 
 use super::{ITypeInstruction, JTypeInstruction, RTypeInstruction};
@@ -52,7 +56,7 @@ pub const N_MIPS_COLS: usize = N_MIPS_REL_COLS + N_MIPS_SEL_COLS;
 
 /// Abstract columns (or variables of our multi-variate polynomials) that will
 /// be used to describe our constraints.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum ColumnAlias {
     // Can be seen as the abstract indexed variable X_{i}
     ScratchState(usize),
@@ -61,26 +65,41 @@ pub enum ColumnAlias {
     ScratchStateInverse(usize),
     InstructionCounter,
     Selector(usize),
+    Lookup(LookupTableID, usize),
+    // TODO add
+    LookupPartialSum((u32, usize)),
+    /// Multiplicities, indexed. This corresponds to the `m_i`. First
+    /// indexed by table ID, then internal index.
+    LookupMultiplicity((u32, usize)),
+    /// The lookup aggregation, i.e. `phi`
+    LookupAggregation,
+    /// The fixed tables. The parameter is considered to the indexed table.
+    LookupFixedTable(u32),
 }
 
 /// The columns used by the MIPS circuit. The MIPS circuit is split into three
 /// main opcodes: RType, JType, IType. The columns are shared between different
 /// instruction types. (the total number of columns refers to the maximum of
 /// columns used by each mode)
-impl From<ColumnAlias> for usize {
-    fn from(alias: ColumnAlias) -> usize {
+impl TryFrom<ColumnAlias> for usize {
+    type Error = String;
+    fn try_from(alias: ColumnAlias) -> Result<usize, String> {
         // Note that SCRATCH_SIZE + 1 is for the error
         match alias {
             ColumnAlias::ScratchState(i) => {
                 assert!(i < SCRATCH_SIZE);
-                i
+                Ok(i)
             }
             ColumnAlias::ScratchStateInverse(i) => {
                 assert!(i < SCRATCH_SIZE_INVERSE);
-                SCRATCH_SIZE + i
+                Ok(SCRATCH_SIZE + i)
             }
-            ColumnAlias::InstructionCounter => SCRATCH_SIZE + SCRATCH_SIZE_INVERSE,
-            ColumnAlias::Selector(s) => SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 + s,
+            ColumnAlias::InstructionCounter => Ok(SCRATCH_SIZE + SCRATCH_SIZE_INVERSE),
+            ColumnAlias::Selector(s) => Ok(SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 + s),
+            ColumnAlias::Lookup(table_id, idx) => Err(format!(
+                "Could not convert a lookup in table {:?} at index {} to a usize.",
+                table_id, idx
+            )),
         }
     }
 }
@@ -124,7 +143,11 @@ impl From<Instruction> for usize {
 /// - the (at most) 4 bytes of the preimage key that are currently being
 ///   processed
 /// - 4 helpers to check if at least n bytes were read in the current row
-pub type MIPSWitness<T> = Witness<N_MIPS_COLS, T>;
+pub struct MIPSWitness<T> {
+    lookup_witness: HashMap<LookupTableID, Vec<T>>,
+    // TODO: Name better
+    column_witness: Witness<N_MIPS_COLS, T>,
+}
 
 // IMPLEMENTATIONS FOR COLUMN ALIAS
 
@@ -133,13 +156,33 @@ impl<T: Clone> Index<ColumnAlias> for MIPSWitness<T> {
 
     /// Map the column alias to the actual column index.
     fn index(&self, index: ColumnAlias) -> &Self::Output {
-        &self.cols[usize::from(index)]
+        match index {
+            ColumnAlias::ScratchState(_)
+            | ColumnAlias::ScratchStateInverse(_)
+            | ColumnAlias::InstructionCounter
+            | ColumnAlias::Selector(_) => {
+                &self.column_witness.cols[usize::try_from(index).unwrap()]
+            }
+            ColumnAlias::Lookup(table_idx, idx) => {
+                &self.lookup_witness.get(&table_idx).unwrap()[idx]
+            }
+        }
     }
 }
 
 impl<T: Clone> IndexMut<ColumnAlias> for MIPSWitness<T> {
     fn index_mut(&mut self, index: ColumnAlias) -> &mut Self::Output {
-        &mut self.cols[usize::from(index)]
+        match index {
+            ColumnAlias::ScratchState(_)
+            | ColumnAlias::ScratchStateInverse(_)
+            | ColumnAlias::InstructionCounter
+            | ColumnAlias::Selector(_) => {
+                &mut self.column_witness.cols[usize::try_from(index).unwrap()]
+            }
+            ColumnAlias::Lookup(table_idx, idx) => {
+                &mut self.lookup_witness.get_mut(&table_idx).unwrap()[idx]
+            }
+        }
     }
 }
 
@@ -188,13 +231,13 @@ impl<T: Clone> Index<Instruction> for MIPSWitness<T> {
 
     /// Map the column alias to the actual column index.
     fn index(&self, index: Instruction) -> &Self::Output {
-        &self.cols[usize::from(index)]
+        &self.column_witness.cols[usize::from(index)]
     }
 }
 
 impl<T: Clone> IndexMut<Instruction> for MIPSWitness<T> {
     fn index_mut(&mut self, index: Instruction) -> &mut Self::Output {
-        &mut self.cols[usize::from(index)]
+        &mut self.column_witness.cols[usize::from(index)]
     }
 }
 

o1vm/src/interpreters/mips/constraints.rs

@@ -16,17 +16,17 @@ use ark_ff::{Field, One};
 use kimchi::circuits::{
     expr::{ConstantTerm::Literal, Expr, ExprInner, Operations, Variable},
     gate::CurrOrNext,
+    lookup::lookups::LookupTableID,
 };
-use kimchi_msm::columns::ColumnIndexer as _;
-use std::array;
+use std::{array, collections::HashMap};
 use strum::IntoEnumIterator;
-
 use super::column::N_MIPS_SEL_COLS;
 
 /// The environment keeping the constraints between the different polynomials
 pub struct Env<Fp> {
     scratch_state_idx: usize,
     scratch_state_idx_inverse: usize,
+    scratch_lookups: HashMap<LookupTableID, usize>,
     /// A list of constraints, which are multi-variate polynomials over a field,
     /// represented using the expression framework of `kimchi`.
     constraints: Vec<E<Fp>>,
@@ -40,6 +40,7 @@ impl<Fp: Field> Default for Env<Fp> {
         Self {
             scratch_state_idx: 0,
             scratch_state_idx_inverse: 0,
+            scratch_lookups: HashMap::new(),
             constraints: Vec::new(),
             lookups: Vec::new(),
             selector: None,
@@ -65,6 +66,13 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         MIPSColumn::ScratchState(scratch_idx)
     }
 
+    fn alloc_lookup(&mut self, table_id: LookupTableID) -> Self::Position {
+        let idx = self.scratch_lookups.get_mut(&table_id).unwrap();
+        *idx += 1;
+        MIPSColumn::Lookup(table_id, *idx - 1)
+        // TODO
+    }
+
     fn alloc_scratch_inverse(&mut self) -> Self::Position {
         let scratch_idx = self.scratch_state_idx_inverse;
         self.scratch_state_idx_inverse += 1;
@@ -75,7 +83,7 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
 
     fn variable(&self, column: Self::Position) -> Self::Variable {
         Expr::Atom(ExprInner::Cell(Variable {
-            col: column.to_column(),
+            col: column,
             row: CurrOrNext::Curr,
         }))
     }

o1vm/src/interpreters/mips/interpreter.rs

@@ -7,6 +7,7 @@ use crate::{
     lookups::{Lookup, LookupTableIDs},
 };
 use ark_ff::{One, Zero};
+use kimchi::circuits::lookup::lookups::LookupTableID;
 use strum::{EnumCount, IntoEnumIterator};
 use strum_macros::{EnumCount, EnumIter};
 
@@ -172,6 +173,8 @@ pub trait InterpreterEnv {
     /// [crate::interpreters::mips::column::SCRATCH_SIZE]
     fn alloc_scratch(&mut self) -> Self::Position;
 
+    fn alloc_lookup(&mut self, table_id: LookupTableID) -> Self::Position;
+
     fn alloc_scratch_inverse(&mut self) -> Self::Position;
 
     type Variable: Clone

o1vm/src/interpreters/mips/tests_helpers.rs

@@ -8,7 +8,7 @@ use crate::{
     preimage_oracle::PreImageOracleT,
 };
 use rand::{CryptoRng, Rng, RngCore};
-use std::{fs, path::PathBuf};
+use std::{collections::HashMap, fs, path::PathBuf};
 
 // FIXME: we should parametrize the tests with different fields.
 use ark_bn254::Fr as Fp;
@@ -92,6 +92,9 @@ where
         scratch_state_idx_inverse: 0,
         scratch_state: [Fp::from(0); SCRATCH_SIZE],
         scratch_state_inverse: [Fp::from(0); SCRATCH_SIZE_INVERSE],
+        scratch_lookup: HashMap::new(),
+        lookup_fixed_tables: HashMap::new(),
+        lookup_multiplicities: HashMap::new(),
         selector: crate::interpreters::mips::column::N_MIPS_SEL_COLS,
         halt: false,
         // Keccak related