Merge pull request #1380 from o1-labs/feature/mips-constraints-prep

Fixup environment trait to avoid abstraction leaks
o1-labs · Dec 5, 2023 · 21c2834 · 21c2834
2 parents 6233800 + d7d4f44
commit 21c2834
Showing 1 changed file with 22 additions and 16 deletions.
diff --git a/optimism/src/mips/interpreter.rs b/optimism/src/mips/interpreter.rs
@@ -196,10 +196,8 @@ pub enum ITypeInstruction {
 pub trait InterpreterEnv {
     type Variable: Clone
         + std::ops::Add<Self::Variable, Output = Self::Variable>
-        + std::ops::Mul<u32, Output = Self::Variable>
-        + std::ops::Shl<u32, Output = Self::Variable>
-        + std::ops::BitAnd<u32, Output = Self::Variable>
-        + std::fmt::Display;
+        + std::ops::Mul<Self::Variable, Output = Self::Variable>
+        + std::fmt::Debug;
 
     fn overwrite_register_checked(&mut self, register_idx: &Self::Variable, value: &Self::Variable);
 
@@ -211,8 +209,8 @@ pub trait InterpreterEnv {
 
     fn get_immediate(&self) -> Self::Variable {
         // The immediate value is the last 16bits
-        (self.get_instruction_part(InstructionPart::RD) << 11)
-            + (self.get_instruction_part(InstructionPart::Shamt) << 6)
+        (self.get_instruction_part(InstructionPart::RD) * Self::constant(1 << 11))
+            + (self.get_instruction_part(InstructionPart::Shamt) * Self::constant(1 << 6))
             + (self.get_instruction_part(InstructionPart::Funct))
     }
 
@@ -290,12 +288,12 @@ pub fn interpret_jtype<Env: InterpreterEnv>(env: &mut Env, instr: JTypeInstructi
             // > bits (which would always be 00, since addresses are always
             // > divisible by 4).
             // Source: https://max.cs.kzoo.edu/cs230/Resources/MIPS/MachineXL/InstructionFormats.html
-            let addr = (env.get_instruction_part(InstructionPart::RS) << 21)
-                + (env.get_instruction_part(InstructionPart::RT) << 16)
-                + (env.get_instruction_part(InstructionPart::RD) << 11)
-                + (env.get_instruction_part(InstructionPart::Shamt) << 6)
+            let addr = (env.get_instruction_part(InstructionPart::RS) * Env::constant(1 << 21))
+                + (env.get_instruction_part(InstructionPart::RT) * Env::constant(1 << 16))
+                + (env.get_instruction_part(InstructionPart::RD) * Env::constant(1 << 11))
+                + (env.get_instruction_part(InstructionPart::Shamt) * Env::constant(1 << 6))
                 + (env.get_instruction_part(InstructionPart::Funct));
-            env.set_instruction_pointer(addr * 4);
+            env.set_instruction_pointer(addr * Env::constant(4));
             // REMOVEME: when all jtype instructions are implemented.
             return;
         }
@@ -326,7 +324,7 @@ pub fn interpret_itype<Env: InterpreterEnv>(env: &mut Env, instr: ITypeInstructi
         ITypeInstruction::AddImmediateUnsigned => {
             let rs = env.get_instruction_part(InstructionPart::RS);
             let rt = env.get_instruction_part(InstructionPart::RT);
-            debug!("Fetching register: {}", rs);
+            debug!("Fetching register: {:?}", rs);
             let register_rs = env.fetch_register_checked(&rs);
             let immediate = env.get_immediate();
             let res = register_rs + immediate;
@@ -344,7 +342,7 @@ pub fn interpret_itype<Env: InterpreterEnv>(env: &mut Env, instr: ITypeInstructi
         ITypeInstruction::LoadUpperImmediate => {
             // lui $reg, [most significant 16 bits of immediate]
             let rt = env.get_instruction_part(InstructionPart::RT);
-            let immediate_value = env.get_immediate() << 16;
+            let immediate_value = env.get_immediate() * Env::constant(1 << 16);
             env.overwrite_register_checked(&rt, &immediate_value);
             env.set_instruction_pointer(env.get_instruction_pointer() + Env::constant(4u32));
             // TODO: update next_instruction_pointer
@@ -358,15 +356,23 @@ pub fn interpret_itype<Env: InterpreterEnv>(env: &mut Env, instr: ITypeInstructi
             let addr = env.get_instruction_part(InstructionPart::RS);
             let offset = env.get_immediate();
             let addr_with_offset = addr.clone() + offset.clone();
-            debug!("lw {}, {}({})", dest.clone(), offset.clone(), addr.clone());
+            debug!(
+                "lw {:?}, {:?}({:?})",
+                dest.clone(),
+                offset.clone(),
+                addr.clone()
+            );
             // We load 4 bytes, i.e. one word.
             let v0 = env.fetch_memory(&addr_with_offset);
             let v1 = env.fetch_memory(&(addr_with_offset.clone() + Env::constant(1)));
             let v2 = env.fetch_memory(&(addr_with_offset.clone() + Env::constant(2)));
             let v3 = env.fetch_memory(&(addr_with_offset.clone() + Env::constant(3)));
-            let value = (v0 << 24) + (v1 << 16) + (v2 << 8) + v3;
+            let value = (v0 * Env::constant(1 << 24))
+                + (v1 * Env::constant(1 << 16))
+                + (v2 * Env::constant(1 << 8))
+                + v3;
             debug!(
-                "Loaded 32 bits value from {}: {}",
+                "Loaded 32 bits value from {:?}: {:?}",
                 addr_with_offset.clone(),
                 value
             );