Fix the equality function in the o1vm interpreter

o1vm/src/interpreters/mips/constraints.rs

@@ -213,6 +213,22 @@ impl<Fp: Field> InterpreterEnv for Env<Fp> {
         self.variable(position)
     }
 
+    fn is_zero(&mut self, x: &Self::Variable) -> Self::Variable {
+        let res = {
+            let pos = self.alloc_scratch();
+            unsafe { self.test_zero(x, pos) }
+        };
+        let x_inv_or_zero = {
+            let pos = self.alloc_scratch();
+            unsafe { self.inverse_or_zero(x, 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)
+        self.add_constraint(x.clone() * x_inv_or_zero.clone() + res.clone() - Self::constant(1));
+        self.add_constraint(x.clone() * res.clone());
+        res
+    }
+
     unsafe fn inverse_or_zero(
         &mut self,
         _x: &Self::Variable,

o1vm/src/interpreters/mips/interpreter.rs

@@ -698,21 +698,7 @@ pub trait InterpreterEnv {
         position: Self::Position,
     ) -> Self::Variable;
 
-    fn is_zero(&mut self, x: &Self::Variable) -> Self::Variable {
-        let res = {
-            let pos = self.alloc_scratch();
-            unsafe { self.test_zero(x, pos) }
-        };
-        let x_inv_or_zero = {
-            let pos = self.alloc_scratch();
-            unsafe { self.inverse_or_zero(x, 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)
-        self.add_constraint(x.clone() * x_inv_or_zero.clone() + res.clone() - Self::constant(1));
-        self.add_constraint(x.clone() * res.clone());
-        res
-    }
+    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`.
     fn equal(&mut self, x: &Self::Variable, y: &Self::Variable) -> Self::Variable;

o1vm/src/interpreters/mips/witness.rs

@@ -361,13 +361,44 @@ impl<Fp: Field, PreImageOracle: PreImageOracleT> InterpreterEnv for Env<Fp, PreI
         }
     }
 
-    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))
+    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()
         } else {
-            self.is_zero(&(*y - *x))
-        }
+            Fp::inverse(&Fp::from(*x)).unwrap()
+        };
+        self.write_field_column(pos, inv_or_zero);
+        // return the result
+        res
+    }
+
+    fn equal(&mut self, x: &Self::Variable, y: &Self::Variable) -> Self::Variable {
+        // 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 _to_zero_test_inv_or_zero = {
+            let pos = self.alloc_scratch();
+            let inv_or_zero = if to_zero_test == Fp::zero() {
+                Fp::zero()
+            } else {
+                Fp::inverse(&to_zero_test).unwrap()
+            };
+            self.write_field_column(pos, inv_or_zero);
+            1 // Placeholder value
+        };
+        res
     }
 
     unsafe fn test_less_than(

o1vm/src/pickles/main.rs

@@ -15,7 +15,7 @@ use o1vm::{
         constraints as mips_constraints,
         interpreter::{self, InterpreterEnv},
         witness::{self as mips_witness},
-        ITypeInstruction, Instruction, RTypeInstruction,
+        Instruction,
     },
     pickles::{
         proof::{Proof, ProofInputs},
@@ -78,27 +78,11 @@ pub fn main() -> ExitCode {
     let mut mips_wit_env =
         mips_witness::Env::<Fp, PreImageOracle>::create(cannon::PAGE_SIZE as usize, state, po);
 
-    // FIXME: This is a hack to skip some instructions that have failing constraints.
-    // It might be related to env.equal.
-    let failing_instructions = [
-        Instruction::RType(RTypeInstruction::SyscallOther),
-        Instruction::RType(RTypeInstruction::SyscallFcntl),
-        Instruction::IType(ITypeInstruction::BranchEq),
-        Instruction::IType(ITypeInstruction::BranchNeq),
-        Instruction::IType(ITypeInstruction::LoadWordLeft),
-        Instruction::IType(ITypeInstruction::LoadWordRight),
-        Instruction::IType(ITypeInstruction::StoreWordLeft),
-        Instruction::IType(ITypeInstruction::StoreWordRight),
-    ];
     let constraints = {
         let mut mips_con_env = mips_constraints::Env::<Fp>::default();
         let mut constraints = Instruction::iter()
             .flat_map(|instr_typ| instr_typ.into_iter())
             .fold(vec![], |mut acc, instr| {
-                if failing_instructions.contains(&instr) {
-                    debug!("Skipping instruction {:?}", instr);
-                    return acc;
-                }
                 interpreter::interpret_instruction(&mut mips_con_env, instr);
                 let selector = mips_con_env.get_selector();
                 let constraints_with_selector: Vec<E<Fp>> = mips_con_env