From 00087d6276cec846139c747189551c6caa03d027 Mon Sep 17 00:00:00 2001
From: Danny Willems <be.danny.willems@gmail.com>
Date: Thu, 21 Nov 2024 09:05:27 +0100
Subject: [PATCH] o1vm/MIPS: intro a new column to keep track of values to be
 inverted

---
 o1vm/src/interpreters/mips/column.rs      | 3 +++
 o1vm/src/interpreters/mips/constraints.rs | 8 ++++++++
 o1vm/src/interpreters/mips/interpreter.rs | 2 ++
 o1vm/src/interpreters/mips/witness.rs     | 7 +++++++
 4 files changed, 20 insertions(+)

diff --git a/o1vm/src/interpreters/mips/column.rs b/o1vm/src/interpreters/mips/column.rs
index 1b2acfc830..770dcfa826 100644
--- a/o1vm/src/interpreters/mips/column.rs
+++ b/o1vm/src/interpreters/mips/column.rs
@@ -50,6 +50,9 @@ pub const N_MIPS_COLS: usize = N_MIPS_REL_COLS + N_MIPS_SEL_COLS;
 pub enum ColumnAlias {
     // Can be seen as the abstract indexed variable X_{i}
     ScratchState(usize),
+    // A column whose value needs to be inverted in the final witness.
+    // We're keeping a separate column to perform a batch inversion at the end.
+    ScratchStateInverse(usize),
     InstructionCounter,
     Selector(usize),
 }
diff --git a/o1vm/src/interpreters/mips/constraints.rs b/o1vm/src/interpreters/mips/constraints.rs
index ea02390fd1..a97898dcb5 100644
--- a/o1vm/src/interpreters/mips/constraints.rs
+++ b/o1vm/src/interpreters/mips/constraints.rs
@@ -25,6 +25,7 @@ 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,
     /// A list of constraints, which are multi-variate polynomials over a field,
     /// represented using the expression framework of `kimchi`.
     constraints: Vec<E<Fp>>,
@@ -37,6 +38,7 @@ impl<Fp: Field> Default for Env<Fp> {
     fn default() -> Self {
         Self {
             scratch_state_idx: 0,
+            scratch_state_idx_inverse: 0,
             constraints: Vec::new(),
             lookups: Vec::new(),
             selector: None,
@@ -62,6 +64,12 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         MIPSColumn::ScratchState(scratch_idx)
     }
 
+    fn alloc_scratch_inverse(&mut self) -> Self::Position {
+        let scratch_idx = self.scratch_state_idx_inverse;
+        self.scratch_state_idx_inverse += 1;
+        MIPSColumn::ScratchStateInverse(scratch_idx)
+    }
+
     type Variable = E<Fp>;
 
     fn variable(&self, column: Self::Position) -> Self::Variable {
diff --git a/o1vm/src/interpreters/mips/interpreter.rs b/o1vm/src/interpreters/mips/interpreter.rs
index 5ace8fb260..e20dc51f50 100644
--- a/o1vm/src/interpreters/mips/interpreter.rs
+++ b/o1vm/src/interpreters/mips/interpreter.rs
@@ -172,6 +172,8 @@ pub trait InterpreterEnv {
     /// [crate::interpreters::mips::witness::SCRATCH_SIZE]
     fn alloc_scratch(&mut self) -> Self::Position;
 
+    fn alloc_scratch_inverse(&mut self) -> Self::Position;
+
     type Variable: Clone
         + std::ops::Add<Self::Variable, Output = Self::Variable>
         + std::ops::Sub<Self::Variable, Output = Self::Variable>
diff --git a/o1vm/src/interpreters/mips/witness.rs b/o1vm/src/interpreters/mips/witness.rs
index 2a9650321a..126e259ef5 100644
--- a/o1vm/src/interpreters/mips/witness.rs
+++ b/o1vm/src/interpreters/mips/witness.rs
@@ -113,6 +113,12 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreI
         Column::ScratchState(scratch_idx)
     }
 
+    fn alloc_scratch_inverse(&mut self) -> Self::Position {
+        let scratch_idx = self.scratch_state_idx_inverse;
+        self.scratch_state_idx_inverse += 1;
+        Column::ScratchStateInverse(scratch_idx)
+    }
+
     type Variable = u64;
 
     fn variable(&self, _column: Self::Position) -> Self::Variable {
@@ -927,6 +933,7 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> Env<Fp, PreImageOracle> {
     pub fn write_field_column(&mut self, column: Column, value: Fp) {
         match column {
             Column::ScratchState(idx) => self.scratch_state[idx] = value,
+            Column::ScratchStateInverse(idx) => self.scratch_state_inverse[idx] = value,
             Column::InstructionCounter => panic!("Cannot overwrite the column {:?}", column),
             Column::Selector(s) => self.selector = s,
         }