Merge pull request #2943 from o1-labs/feature/logup-tracking

Add multiplicity tracking for lookups

o1vm/src/interpreters/mips/tests_helpers.rs

@@ -3,7 +3,7 @@ use crate::{
     interpreters::mips::{
         interpreter::{debugging::InstructionParts, InterpreterEnv},
         registers::Registers,
-        witness::{Env as WEnv, SyscallEnv},
+        witness::{Env as WEnv, LookupMultiplicities, SyscallEnv},
     },
     preimage_oracle::PreImageOracleT,
 };
@@ -92,6 +92,7 @@ where
         scratch_state_idx_inverse: 0,
         scratch_state: [Fp::from(0); SCRATCH_SIZE],
         scratch_state_inverse: [Fp::from(0); SCRATCH_SIZE_INVERSE],
+        lookup_multiplicities: LookupMultiplicities::new(),
         selector: crate::interpreters::mips::column::N_MIPS_SEL_COLS,
         halt: false,
         // Keccak related

o1vm/src/interpreters/mips/witness.rs

@@ -20,13 +20,14 @@ use crate::{
             registers::Registers,
         },
     },
-    lookups::Lookup,
+    lookups::{Lookup, LookupTableIDs},
     preimage_oracle::PreImageOracleT,
     utils::memory_size,
 };
-use ark_ff::Field;
+use ark_ff::{Field, PrimeField};
 use core::panic;
 use kimchi::o1_utils::Two;
+use kimchi_msm::LogupTableID;
 use log::{debug, info};
 use std::{
     array,
@@ -64,6 +65,36 @@ impl SyscallEnv {
     }
 }
 
+pub struct LookupMultiplicities {
+    pub pad_lookup: Vec<u64>,
+    pub round_constants_lookup: Vec<u64>,
+    pub at_most_4_lookup: Vec<u64>,
+    pub byte_lookup: Vec<u64>,
+    pub range_check_16_lookup: Vec<u64>,
+    pub sparse_lookup: Vec<u64>,
+    pub reset_lookup: Vec<u64>,
+}
+
+impl LookupMultiplicities {
+    pub fn new() -> Self {
+        LookupMultiplicities {
+            pad_lookup: vec![0; LookupTableIDs::PadLookup.length()],
+            round_constants_lookup: vec![0; LookupTableIDs::RoundConstantsLookup.length()],
+            at_most_4_lookup: vec![0; LookupTableIDs::AtMost4Lookup.length()],
+            byte_lookup: vec![0; LookupTableIDs::ByteLookup.length()],
+            range_check_16_lookup: vec![0; LookupTableIDs::RangeCheck16Lookup.length()],
+            sparse_lookup: vec![0; LookupTableIDs::SparseLookup.length()],
+            reset_lookup: vec![0; LookupTableIDs::ResetLookup.length()],
+        }
+    }
+}
+
+impl Default for LookupMultiplicities {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 /// 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,13 +121,14 @@ pub struct Env<Fp, PreImageOracle: PreImageOracleT> {
     pub preimage_key: Option<[u8; 32]>,
     pub keccak_env: Option<KeccakEnv<Fp>>,
     pub hash_counter: u64,
+    pub lookup_multiplicities: LookupMultiplicities,
 }
 
 fn fresh_scratch_state<Fp: Field, const N: usize>() -> [Fp; N] {
     array::from_fn(|_| Fp::zero())
 }
 
-impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreImageOracle> {
+impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreImageOracle> {
     type Position = Column;
 
     fn alloc_scratch(&mut self) -> Self::Position {
@@ -142,9 +174,30 @@ 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>) {
+        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 {
+                LookupTableIDs::PadLookup => self.lookup_multiplicities.pad_lookup[idx] += 1,
+                LookupTableIDs::RoundConstantsLookup => {
+                    self.lookup_multiplicities.round_constants_lookup[idx] += 1
+                }
+                LookupTableIDs::AtMost4Lookup => {
+                    self.lookup_multiplicities.at_most_4_lookup[idx] += 1
+                }
+                LookupTableIDs::ByteLookup => self.lookup_multiplicities.byte_lookup[idx] += 1,
+                LookupTableIDs::RangeCheck16Lookup => {
+                    self.lookup_multiplicities.range_check_16_lookup[idx] += 1
+                }
+                LookupTableIDs::SparseLookup => self.lookup_multiplicities.sparse_lookup[idx] += 1,
+                LookupTableIDs::ResetLookup => self.lookup_multiplicities.reset_lookup[idx] += 1,
+                // RAM ones, no multiplicities
+                LookupTableIDs::MemoryLookup => (),
+                LookupTableIDs::RegisterLookup => (),
+                LookupTableIDs::SyscallLookup => (),
+                LookupTableIDs::KeccakStepLookup => (),
+            }
+        }
     }
 
     fn instruction_counter(&self) -> Self::Variable {
@@ -819,7 +872,7 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreI
     }
 }
 
-impl<Fp: Field, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
+impl<Fp: PrimeField, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
     pub fn create(page_size: usize, state: State, preimage_oracle: PreImageOracle) -> Self {
         let initial_instruction_pointer = state.pc;
         let next_instruction_pointer = state.next_pc;
@@ -891,6 +944,7 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
             preimage_key: None,
             keccak_env: None,
             hash_counter: 0,
+            lookup_multiplicities: LookupMultiplicities::new(),
         }
     }
 

o1vm/src/lookups.rs

@@ -106,8 +106,36 @@ impl LookupTableID for LookupTableIDs {
         panic!("No runtime tables specified");
     }
 
-    fn ix_by_value<F: PrimeField>(&self, _value: &[F]) -> Option<usize> {
-        todo!()
+    fn ix_by_value<F: PrimeField>(&self, value: &[F]) -> Option<usize> {
+        // Shamelessly copied from below, where it is likely also incorrect.
+        let idx = value[0]
+            .to_bytes()
+            .iter()
+            .rev()
+            .fold(0u64, |acc, &x| acc * 256 + x as u64) as usize;
+        match self {
+            // Fixed tables
+            Self::RoundConstantsLookup
+            | Self::AtMost4Lookup
+            | Self::ByteLookup
+            | Self::RangeCheck16Lookup
+            | Self::ResetLookup => Some(idx),
+            Self::PadLookup => Some(idx - 1),
+            Self::SparseLookup => {
+                // Big yikes. This is copied from below.
+                let res = u64::from_str_radix(&format!("{:x}", idx), 2);
+                if let Ok(ok) = res {
+                    Some(ok as usize)
+                } else {
+                    panic!("Help");
+                }
+            }
+            // Non-fixed tables
+            Self::MemoryLookup
+            | Self::RegisterLookup
+            | Self::SyscallLookup
+            | Self::KeccakStepLookup => None,
+        }
     }
 
     fn all_variants() -> Vec<Self> {

o1vm/tests/test_mips_elf.rs

@@ -1,4 +1,4 @@
-use ark_ff::Field;
+use ark_ff::PrimeField;
 use mina_curves::pasta::Fp;
 use o1vm::{
     cannon::{self, State, VmConfiguration},
@@ -29,7 +29,10 @@ impl MipsTest {
         o1vm::elf_loader::parse_elf(Architecture::Mips, &path).unwrap()
     }
 
-    fn read_word<Fp: Field, T: PreImageOracleT>(env: &mut witness::Env<Fp, T>, addr: u32) -> u32 {
+    fn read_word<Fp: PrimeField, T: PreImageOracleT>(
+        env: &mut witness::Env<Fp, T>,
+        addr: u32,
+    ) -> u32 {
         let bytes: [u8; 4] = [
             env.get_memory_direct(addr),
             env.get_memory_direct(addr + 1),