Merge pull request #2076 from o1-labs/zkvm/mips/fix-constraint-env-reset

Fix indexing in MIPS constraints environment
o1-labs · Apr 16, 2024 · 4fb098a · 4fb098a
2 parents 046d407 + 083d067
commit 4fb098a
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 9 deletions.
diff --git a/optimism/src/lib.rs b/optimism/src/lib.rs
@@ -8,6 +8,7 @@ pub mod cannon;
 /// A CLI mimicking the Cannon CLI.
 pub mod cannon_cli;
 
+#[cfg(feature = "bn254")]
 /// Integration with folding. Contains common trait implementations to be used by each circuit.
 pub mod folding;
 

diff --git a/optimism/src/main.rs b/optimism/src/main.rs
@@ -19,6 +19,7 @@ use kimchi_optimism::{
     preimage_oracle::PreImageOracle,
     proof, CircuitTrait, DOMAIN_SIZE,
 };
+use log::debug;
 use mina_poseidon::{
     constants::PlonkSpongeConstantsKimchi,
     sponge::{DefaultFqSponge, DefaultFrSponge},
@@ -201,6 +202,7 @@ pub fn main() -> ExitCode {
     {
         // MIPS
         for instr in Instruction::iter().flat_map(|x| x.into_iter()) {
+            debug!("Checking MIPS circuit {:?}", instr);
             let mips_result = prove::<
                 _,
                 OpeningProof,
@@ -218,7 +220,7 @@ pub fn main() -> ExitCode {
                 &mut rng,
             );
             let mips_proof = mips_result.unwrap();
-            eprintln!("Generated a MIPS {:?} proof:\n{:?}", instr, mips_proof);
+            debug!("Generated a MIPS {:?} proof:\n{:?}", instr, mips_proof);
             let mips_verifies = verify::<
                 _,
                 OpeningProof,
@@ -235,9 +237,9 @@ pub fn main() -> ExitCode {
                 Witness::zero_vec(DOMAIN_SIZE),
             );
             if mips_verifies {
-                eprintln!("The MIPS {:?} proof verifies", instr)
+                debug!("The MIPS {:?} proof verifies", instr)
             } else {
-                eprintln!("The MIPS {:?} proof doesn't verify", instr)
+                debug!("The MIPS {:?} proof doesn't verify", instr)
             }
         }
     }
@@ -246,6 +248,7 @@ pub fn main() -> ExitCode {
         // KECCAK
         // FIXME: when folding is applied, the error term will be created to satisfy the folded witness
         for step in keccak::STEPS {
+            debug!("Checking Keccak circuit {:?}", step);
             let keccak_result = prove::<
                 _,
                 OpeningProof,
@@ -263,7 +266,7 @@ pub fn main() -> ExitCode {
                 &mut rng,
             );
             let keccak_proof = keccak_result.unwrap();
-            eprintln!("Generated a Keccak {:?} proof:\n{:?}", step, keccak_proof);
+            debug!("Generated a Keccak {:?} proof:\n{:?}", step, keccak_proof);
             let keccak_verifies = verify::<
                 _,
                 OpeningProof,
@@ -280,9 +283,9 @@ pub fn main() -> ExitCode {
                 Witness::zero_vec(DOMAIN_SIZE),
             );
             if keccak_verifies {
-                eprintln!("The Keccak {:?} proof verifies", step)
+                debug!("The Keccak {:?} proof verifies", step)
             } else {
-                eprintln!("The Keccak {:?} proof doesn't verify", step)
+                debug!("The Keccak {:?} proof doesn't verify", step)
             }
         }
     }

diff --git a/optimism/src/mips/column.rs b/optimism/src/mips/column.rs
@@ -79,7 +79,10 @@ impl<T: Clone> Index<ColumnAlias> for MIPSWitness<T> {
     /// is used by intermediary values when executing the Round step.
     fn index(&self, index: ColumnAlias) -> &Self::Output {
         match index {
-            ColumnAlias::ScratchState(i) => &self.scratch()[i],
+            ColumnAlias::ScratchState(i) => {
+                assert!(i < SCRATCH_SIZE);
+                &self.scratch()[i]
+            }
             ColumnAlias::InstructionCounter => self.instruction_counter(),
         }
     }
@@ -88,7 +91,10 @@ impl<T: Clone> Index<ColumnAlias> for MIPSWitness<T> {
 impl<T: Clone> IndexMut<ColumnAlias> for MIPSWitness<T> {
     fn index_mut(&mut self, index: ColumnAlias) -> &mut Self::Output {
         match index {
-            ColumnAlias::ScratchState(i) => &mut self.cols[i],
+            ColumnAlias::ScratchState(i) => {
+                assert!(i < SCRATCH_SIZE);
+                &mut self.cols[i]
+            }
             ColumnAlias::InstructionCounter => &mut self.cols[SCRATCH_SIZE],
         }
     }
@@ -98,7 +104,10 @@ impl ColumnIndexer for ColumnAlias {
     fn to_column(self) -> Column {
         // TODO: what happens with error? It does not have a corresponding alias
         match self {
-            ColumnAlias::ScratchState(i) => Column::X(i),
+            ColumnAlias::ScratchState(i) => {
+                assert!(i < SCRATCH_SIZE);
+                Column::X(i)
+            }
             ColumnAlias::InstructionCounter => Column::X(SCRATCH_SIZE),
         }
     }

diff --git a/optimism/src/mips/mod.rs b/optimism/src/mips/mod.rs
@@ -65,6 +65,7 @@ impl<F: Field> CircuitTrait<MIPS_COLUMNS, Instruction, F, Env<F>> for MIPSCircui
             interpret_instruction(env, instr);
             circuit.constraints.insert(instr, env.constraints.clone());
             circuit.lookups.insert(instr, env.lookups.clone());
+            env.scratch_state_idx = 0; // Reset the scratch state index for the next instruction
             env.constraints = vec![]; // Clear the constraints for the next instruction
             env.lookups = vec![]; // Clear the lookups for the next instruction
         }