Merge pull request #1381 from o1-labs/feature/mips-columns

Add support for scratch_state to MIPS VM

optimism/src/mips/column.rs

@@ -0,0 +1,4 @@
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum Column {
+    ScratchState(usize),
+}

optimism/src/mips/interpreter.rs

@@ -194,6 +194,10 @@ pub enum ITypeInstruction {
 }
 
 pub trait InterpreterEnv {
+    type Position;
+
+    fn alloc_scratch(&mut self) -> Self::Position;
+
     type Variable: Clone
         + std::ops::Add<Self::Variable, Output = Self::Variable>
         + std::ops::Mul<Self::Variable, Output = Self::Variable>

optimism/src/mips/mod.rs

@@ -1,3 +1,4 @@
+pub mod column;
 pub mod interpreter;
 pub mod registers;
 pub mod witness;

optimism/src/mips/witness.rs

@@ -4,6 +4,7 @@ use crate::{
         PAGE_SIZE,
     },
     mips::{
+        column::Column,
         interpreter::{
             self, ITypeInstruction, Instruction, InstructionPart, InstructionParts, InterpreterEnv,
             JTypeInstruction, RTypeInstruction,
@@ -98,6 +99,14 @@ fn memory_size(total: usize) -> String {
 }
 
 impl<Fp: Field> InterpreterEnv for Env<Fp> {
+    type Position = Column;
+
+    fn alloc_scratch(&mut self) -> Self::Position {
+        let scratch_idx = self.scratch_state_idx;
+        self.scratch_state_idx += 1;
+        Column::ScratchState(scratch_idx)
+    }
+
     type Variable = u32;
 
     fn overwrite_register_checked(
@@ -211,6 +220,12 @@ impl<Fp: Field> Env<Fp> {
         }
     }
 
+    pub fn write_column(&mut self, column: Column, value: u64) {
+        match column {
+            Column::ScratchState(idx) => self.scratch_state[idx] = value.into(),
+        }
+    }
+
     pub fn get_memory_direct(&self, addr: u32) -> u8 {
         let page = addr >> PAGE_ADDRESS_SIZE;
         let page_address = (addr & PAGE_ADDRESS_MASK) as usize;