adding support for batch inverse in parity with mips

o1vm/src/interpreters/riscv32im/constraints.rs

@@ -15,6 +15,7 @@ use kimchi::circuits::{
 
 pub struct Env<F: Field> {
     pub scratch_state_idx: usize,
+    pub scratch_state_idx_inverse: usize,
     pub lookups: Vec<Lookup<E<F>>>,
     pub constraints: Vec<E<F>>,
     pub selector: Option<E<F>>,
@@ -24,6 +25,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,
@@ -49,6 +51,12 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         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 = E<Fp>;
 
     fn variable(&self, column: Self::Position) -> Self::Variable {
@@ -206,8 +214,8 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
             unsafe { self.test_zero(x, pos) }
         };
         let x_inv_or_zero = {
-            let pos = self.alloc_scratch();
-            unsafe { self.inverse_or_zero(x, pos) }
+            let pos = self.alloc_scratch_inverse();
+            self.variable(pos)
         };
         // If x = 0, then res = 1 and x_inv_or_zero = 0
         // If x <> 0, then res = 0 and x_inv_or_zero = x^(-1)
@@ -216,14 +224,6 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         res
     }
 
-    unsafe fn inverse_or_zero(
-        &mut self,
-        _x: &Self::Variable,
-        position: Self::Position,
-    ) -> Self::Variable {
-        self.variable(position)
-    }
-
     fn equal(&mut self, x: &Self::Variable, y: &Self::Variable) -> Self::Variable {
         self.is_zero(&(x.clone() - y.clone()))
     }

o1vm/src/interpreters/riscv32im/interpreter.rs

@@ -561,6 +561,8 @@ pub trait InterpreterEnv {
     /// [crate::interpreters::riscv32im::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>
@@ -1061,21 +1063,6 @@ pub trait InterpreterEnv {
     /// `x`.
     unsafe fn test_zero(&mut self, x: &Self::Variable, position: Self::Position) -> Self::Variable;
 
-    /// Returns `x^(-1)`, or `0` if `x` is `0`, storing the result in `position`.
-    ///
-    /// # Safety
-    ///
-    /// There are no constraints on the returned value; callers must assert the relationship with
-    /// `x`.
-    ///
-    /// The value returned may be a placeholder; callers should be careful not to depend directly
-    /// on the value stored in the variable.
-    unsafe fn inverse_or_zero(
-        &mut self,
-        x: &Self::Variable,
-        position: Self::Position,
-    ) -> Self::Variable;
-
     fn is_zero(&mut self, x: &Self::Variable) -> Self::Variable;
 
     /// Returns 1 if `x` is equal to `y`, or 0 otherwise, storing the result in `position`.

o1vm/src/interpreters/riscv32im/witness.rs

@@ -65,6 +65,12 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         Column::ScratchState(scratch_idx)
     }
 
+    fn alloc_scratch_inverse(&mut self) -> Self::Position {
+        let scratch_inverse_idx = self.scratch_state_inverse_idx;
+        self.scratch_state_inverse_idx += 1;
+        Column::ScratchStateInverse(scratch_inverse_idx)
+    }
+
     type Variable = u64;
 
     fn variable(&self, _column: Self::Position) -> Self::Variable {
@@ -279,44 +285,39 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         res
     }
 
-    unsafe fn inverse_or_zero(
-        &mut self,
-        x: &Self::Variable,
-        position: Self::Position,
-    ) -> Self::Variable {
-        if *x == 0 {
-            self.write_column(position, 0);
-            0
-        } else {
-            self.write_field_column(position, Fp::from(*x).inverse().unwrap());
-            1 // Placeholder value
-        }
-    }
-
     fn is_zero(&mut self, x: &Self::Variable) -> Self::Variable {
         // write the result
         let pos = self.alloc_scratch();
         let res = if *x == 0 { 1 } else { 0 };
         self.write_column(pos, res);
         // write the non deterministic advice inv_or_zero
-        let pos = self.alloc_scratch();
-        let inv_or_zero = if *x == 0 {
-            Fp::zero()
+        let pos = self.alloc_scratch_inverse();
+        if *x == 0 {
+            self.write_field_column(pos, Fp::zero());
         } else {
-            Fp::inverse(&Fp::from(*x)).unwrap()
+            self.write_field_column(pos, Fp::from(*x));
         };
-        self.write_field_column(pos, inv_or_zero);
         // return the result
         res
     }
 
     fn equal(&mut self, x: &Self::Variable, y: &Self::Variable) -> Self::Variable {
-        // To avoid subtraction overflow in the witness interpreter for u32
-        if x > y {
-            self.is_zero(&(*x - *y))
+        // We replicate is_zero(x-y), but working on field elt,
+        // to avoid subtraction overflow in the witness interpreter for u32
+        let to_zero_test = Fp::from(*x) - Fp::from(*y);
+        let res = {
+            let pos = self.alloc_scratch();
+            let is_zero: u64 = if to_zero_test == Fp::zero() { 1 } else { 0 };
+            self.write_column(pos, is_zero);
+            is_zero
+        };
+        let pos = self.alloc_scratch_inverse();
+        if to_zero_test == Fp::zero() {
+            self.write_field_column(pos, Fp::zero());
         } else {
-            self.is_zero(&(*y - *x))
-        }
+            self.write_field_column(pos, to_zero_test);
+        };
+        res
     }
 
     unsafe fn test_less_than(