Merge pull request #2857 from o1-labs/sai/branch-eq

branch equal instruction implementation
o1-labs · Dec 24, 2024 · 0009117 · 0009117
2 parents f6beae9 + 416e65a
commit 0009117
Showing 1 changed file with 23 additions and 1 deletion.
diff --git a/o1vm/src/interpreters/riscv32im/interpreter.rs b/o1vm/src/interpreters/riscv32im/interpreter.rs
@@ -2218,7 +2218,29 @@ pub fn interpret_sbtype<Env: InterpreterEnv>(env: &mut Env, instr: SBInstruction
 
     match instr {
         SBInstruction::BranchEq => {
-            unimplemented!("BranchEq")
+            // beq: if (x[rs1] == x[rs2]) pc += sext(offset)
+            let local_rs1 = env.read_register(&rs1);
+            let local_rs2 = env.read_register(&rs2);
+
+            let equals = env.equal(&local_rs1, &local_rs2);
+            let offset = (Env::constant(1) - equals.clone()) * Env::constant(4) + equals * imm0_12;
+            let offset = env.sign_extend(&offset, 12);
+            let addr = {
+                let res_scratch = env.alloc_scratch();
+                let overflow_scratch = env.alloc_scratch();
+                let (res, _overflow) = unsafe {
+                    env.add_witness(
+                        &next_instruction_pointer,
+                        &offset,
+                        res_scratch,
+                        overflow_scratch,
+                    )
+                };
+                // FIXME: Requires a range check
+                res
+            };
+            env.set_instruction_pointer(next_instruction_pointer);
+            env.set_next_instruction_pointer(addr);
         }
         SBInstruction::BranchNeq => {
             // bne: if (x[rs1] != x[rs2]) pc += sext(offset)