o1vm/MIPS: intro a new column to keep track of values to be inverted

o1vm/src/interpreters/mips/column.rs

@@ -9,7 +9,7 @@ use kimchi_msm::{
 use std::ops::{Index, IndexMut};
 use strum::EnumCount;
 
-use super::{ITypeInstruction, JTypeInstruction, RTypeInstruction};
+use super::{witness::SCRATCH_SIZE_INVERSE, ITypeInstruction, JTypeInstruction, RTypeInstruction};
 
 /// The number of hashes performed so far in the block
 pub(crate) const MIPS_HASH_COUNTER_OFF: usize = 80;
@@ -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),
 }
@@ -66,8 +69,12 @@ impl From<ColumnAlias> for usize {
                 assert!(i < SCRATCH_SIZE);
                 i
             }
-            ColumnAlias::InstructionCounter => SCRATCH_SIZE,
-            ColumnAlias::Selector(s) => SCRATCH_SIZE + 1 + s,
+            ColumnAlias::ScratchStateInverse(i) => {
+                assert!(i < SCRATCH_SIZE_INVERSE);
+                SCRATCH_SIZE + i
+            }
+            ColumnAlias::InstructionCounter => SCRATCH_SIZE + SCRATCH_SIZE_INVERSE,
+            ColumnAlias::Selector(s) => SCRATCH_SIZE + SCRATCH_SIZE_INVERSE + 1 + s,
         }
     }
 }
@@ -138,7 +145,11 @@ impl ColumnIndexer for ColumnAlias {
                 assert!(ss < SCRATCH_SIZE);
                 Column::Relation(ss)
             }
-            Self::InstructionCounter => Column::Relation(SCRATCH_SIZE),
+            Self::ScratchStateInverse(ss) => {
+                assert!(ss < SCRATCH_SIZE_INVERSE);
+                Column::Relation(SCRATCH_SIZE + ss)
+            }
+            Self::InstructionCounter => Column::Relation(SCRATCH_SIZE + SCRATCH_SIZE_INVERSE),
             // TODO: what happens with error? It does not have a corresponding alias
             Self::Selector(s) => {
                 assert!(s < N_MIPS_SEL_COLS);

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 {

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>

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,
         }