From 23f9760dae99c5d1a5e7b620263ae10ba91f6838 Mon Sep 17 00:00:00 2001
From: mrmr1993 <mr_1993@hotmail.co.uk>
Date: Mon, 4 Dec 2023 22:52:23 +0000
Subject: [PATCH 1/3] Add Column

---
 optimism/src/mips/column.rs | 4 ++++
 optimism/src/mips/mod.rs    | 1 +
 2 files changed, 5 insertions(+)
 create mode 100644 optimism/src/mips/column.rs

diff --git a/optimism/src/mips/column.rs b/optimism/src/mips/column.rs
new file mode 100644
index 0000000000..3d2c234747
--- /dev/null
+++ b/optimism/src/mips/column.rs
@@ -0,0 +1,4 @@
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum Column {
+    ScratchState(usize),
+}
diff --git a/optimism/src/mips/mod.rs b/optimism/src/mips/mod.rs
index e4b4a77241..8394c1f946 100644
--- a/optimism/src/mips/mod.rs
+++ b/optimism/src/mips/mod.rs
@@ -1,3 +1,4 @@
+pub mod column;
 pub mod interpreter;
 pub mod registers;
 pub mod witness;

From bb4076b829637c3fbc0d7ef4abd74c2c80b496e1 Mon Sep 17 00:00:00 2001
From: mrmr1993 <mr_1993@hotmail.co.uk>
Date: Mon, 4 Dec 2023 22:54:44 +0000
Subject: [PATCH 2/3] Add write_column helper to environment

---
 optimism/src/mips/witness.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/optimism/src/mips/witness.rs b/optimism/src/mips/witness.rs
index 547b0c806d..3f489e34a7 100644
--- a/optimism/src/mips/witness.rs
+++ b/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,
@@ -211,6 +212,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;

From 228b1b6b8a6328d8e84c98c64953786ac456524d Mon Sep 17 00:00:00 2001
From: mrmr1993 <mr_1993@hotmail.co.uk>
Date: Mon, 4 Dec 2023 22:57:40 +0000
Subject: [PATCH 3/3] Import alloc_scratch

---
 optimism/src/mips/interpreter.rs | 4 ++++
 optimism/src/mips/witness.rs     | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/optimism/src/mips/interpreter.rs b/optimism/src/mips/interpreter.rs
index 9224e286af..56ebb2bf1d 100644
--- a/optimism/src/mips/interpreter.rs
+++ b/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>
diff --git a/optimism/src/mips/witness.rs b/optimism/src/mips/witness.rs
index 3f489e34a7..f25504c8f9 100644
--- a/optimism/src/mips/witness.rs
+++ b/optimism/src/mips/witness.rs
@@ -99,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(