cannon: Extract handleBranch (#10822)

* cannon: Extract handleBranch, outputState

* cannon: Define MIPSInstructions.sol as a library so we can set visibility modifiers

* cannon: Flatten cpu fields in sol files

* cannon: Rework sol cpu field handling

* cannon: Update MIPS contract version

* cannon: Run snapshots, semver lock

* cannon: Update go logic to match sol

* cannon: Move outputState() back into MIPS.sol

* cannon: Mark pure functions

* cannon: Run semver-lock

* cannon: Update state.json format in test_data

* cannon: Fix variable names to match style guide

* cannon: Undo comment formatting change

* cannon: Style fix - use named parameters

* cannon: Run semver-lock

cannon/cmd/run.go

@@ -380,16 +380,16 @@ func Run(ctx *cli.Context) error {
 			delta := time.Since(start)
 			l.Info("processing",
 				"step", step,
-				"pc", mipsevm.HexU32(state.PC),
-				"insn", mipsevm.HexU32(state.Memory.GetMemory(state.PC)),
+				"pc", mipsevm.HexU32(state.Cpu.PC),
+				"insn", mipsevm.HexU32(state.Memory.GetMemory(state.Cpu.PC)),
 				"ips", float64(step-startStep)/(float64(delta)/float64(time.Second)),
 				"pages", state.Memory.PageCount(),
 				"mem", state.Memory.Usage(),
-				"name", meta.LookupSymbol(state.PC),
+				"name", meta.LookupSymbol(state.Cpu.PC),
 			)
 		}
 
-		if sleepCheck(state.PC) { // don't loop forever when we get stuck because of an unexpected bad program
+		if sleepCheck(state.Cpu.PC) { // don't loop forever when we get stuck because of an unexpected bad program
 			return fmt.Errorf("got stuck in Go sleep at step %d", step)
 		}
 
@@ -411,7 +411,7 @@ func Run(ctx *cli.Context) error {
 			}
 			witness, err := stepFn(true)
 			if err != nil {
-				return fmt.Errorf("failed at proof-gen step %d (PC: %08x): %w", step, state.PC, err)
+				return fmt.Errorf("failed at proof-gen step %d (PC: %08x): %w", step, state.Cpu.PC, err)
 			}
 			postStateHash, err := state.EncodeWitness().StateHash()
 			if err != nil {
@@ -435,7 +435,7 @@ func Run(ctx *cli.Context) error {
 		} else {
 			_, err = stepFn(false)
 			if err != nil {
-				return fmt.Errorf("failed at step %d (PC: %08x): %w", step, state.PC, err)
+				return fmt.Errorf("failed at step %d (PC: %08x): %w", step, state.Cpu.PC, err)
 			}
 		}
 

cannon/mipsevm/evm_test.go

@@ -172,7 +172,7 @@ func TestEVM(t *testing.T) {
 			fn := path.Join("open_mips_tests/test/bin", f.Name())
 			programMem, err := os.ReadFile(fn)
 			require.NoError(t, err)
-			state := &State{PC: 0, NextPC: 4, Memory: NewMemory()}
+			state := &State{Cpu: CpuScalars{PC: 0, NextPC: 4}, Memory: NewMemory()}
 			err = state.Memory.SetMemoryRange(0, bytes.NewReader(programMem))
 			require.NoError(t, err, "load program into state")
 
@@ -182,14 +182,14 @@ func TestEVM(t *testing.T) {
 			goState := NewInstrumentedState(state, oracle, os.Stdout, os.Stderr)
 
 			for i := 0; i < 1000; i++ {
-				if goState.state.PC == endAddr {
+				if goState.state.Cpu.PC == endAddr {
 					break
 				}
 				if exitGroup && goState.state.Exited {
 					break
 				}
-				insn := state.Memory.GetMemory(state.PC)
-				t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.PC, insn)
+				insn := state.Memory.GetMemory(state.Cpu.PC)
+				t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.Cpu.PC, insn)
 
 				stepWitness, err := goState.Step(true)
 				require.NoError(t, err)
@@ -201,11 +201,11 @@ func TestEVM(t *testing.T) {
 					"mipsevm produced different state than EVM at step %d", state.Step)
 			}
 			if exitGroup {
-				require.NotEqual(t, uint32(endAddr), goState.state.PC, "must not reach end")
+				require.NotEqual(t, uint32(endAddr), goState.state.Cpu.PC, "must not reach end")
 				require.True(t, goState.state.Exited, "must set exited state")
 				require.Equal(t, uint8(1), goState.state.ExitCode, "must exit with 1")
 			} else {
-				require.Equal(t, uint32(endAddr), state.PC, "must reach end")
+				require.Equal(t, uint32(endAddr), state.Cpu.PC, "must reach end")
 				// inspect test result
 				done, result := state.Memory.GetMemory(baseAddrEnd+4), state.Memory.GetMemory(baseAddrEnd+8)
 				require.Equal(t, done, uint32(1), "must be done")
@@ -233,7 +233,7 @@ func TestEVMSingleStep(t *testing.T) {
 
 	for _, tt := range cases {
 		t.Run(tt.name, func(t *testing.T) {
-			state := &State{PC: tt.pc, NextPC: tt.nextPC, Memory: NewMemory()}
+			state := &State{Cpu: CpuScalars{PC: tt.pc, NextPC: tt.nextPC}, Memory: NewMemory()}
 			state.Memory.SetMemory(tt.pc, tt.insn)
 
 			us := NewInstrumentedState(state, nil, os.Stdout, os.Stderr)
@@ -401,7 +401,7 @@ func TestEVMSysWriteHint(t *testing.T) {
 	for _, tt := range cases {
 		t.Run(tt.name, func(t *testing.T) {
 			oracle := hintTrackingOracle{}
-			state := &State{PC: 0, NextPC: 4, Memory: NewMemory()}
+			state := &State{Cpu: CpuScalars{PC: 0, NextPC: 4}, Memory: NewMemory()}
 
 			state.LastHint = tt.lastHint
 			state.Registers[2] = sysWrite
@@ -448,8 +448,8 @@ func TestEVMFault(t *testing.T) {
 
 	for _, tt := range cases {
 		t.Run(tt.name, func(t *testing.T) {
-			state := &State{PC: 0, NextPC: tt.nextPC, Memory: NewMemory()}
-			initialState := &State{PC: 0, NextPC: tt.nextPC, Memory: state.Memory}
+			state := &State{Cpu: CpuScalars{PC: 0, NextPC: tt.nextPC}, Memory: NewMemory()}
+			initialState := &State{Cpu: CpuScalars{PC: 0, NextPC: tt.nextPC}, Memory: state.Memory}
 			state.Memory.SetMemory(0, tt.insn)
 
 			// set the return address ($ra) to jump into when test completes
@@ -496,9 +496,9 @@ func TestHelloEVM(t *testing.T) {
 		if goState.state.Exited {
 			break
 		}
-		insn := state.Memory.GetMemory(state.PC)
+		insn := state.Memory.GetMemory(state.Cpu.PC)
 		if i%1000 == 0 { // avoid spamming test logs, we are executing many steps
-			t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.PC, insn)
+			t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.Cpu.PC, insn)
 		}
 
 		evm := NewMIPSEVM(contracts, addrs)
@@ -548,9 +548,9 @@ func TestClaimEVM(t *testing.T) {
 			break
 		}
 
-		insn := state.Memory.GetMemory(state.PC)
+		insn := state.Memory.GetMemory(state.Cpu.PC)
 		if i%1000 == 0 { // avoid spamming test logs, we are executing many steps
-			t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.PC, insn)
+			t.Logf("step: %4d pc: 0x%08x insn: 0x%08x", state.Step, state.Cpu.PC, insn)
 		}
 
 		stepWitness, err := goState.Step(true)

cannon/mipsevm/fuzz_evm_test.go

@@ -21,10 +21,12 @@ func FuzzStateSyscallBrk(f *testing.F) {
 		pc = pc & 0xFF_FF_FF_FC // align PC
 		nextPC := pc + 4
 		state := &State{
-			PC:             pc,
-			NextPC:         nextPC,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     pc,
+				NextPC: nextPC,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -44,10 +46,10 @@ func FuzzStateSyscallBrk(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, pc+4, state.PC)
-		require.Equal(t, nextPC+4, state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, pc+4, state.Cpu.PC)
+		require.Equal(t, nextPC+4, state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -71,10 +73,12 @@ func FuzzStateSyscallClone(f *testing.F) {
 		pc = pc & 0xFF_FF_FF_FC // align PC
 		nextPC := pc + 4
 		state := &State{
-			PC:             pc,
-			NextPC:         nextPC,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     pc,
+				NextPC: nextPC,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -93,10 +97,10 @@ func FuzzStateSyscallClone(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, pc+4, state.PC)
-		require.Equal(t, nextPC+4, state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, pc+4, state.Cpu.PC)
+		require.Equal(t, nextPC+4, state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -118,10 +122,12 @@ func FuzzStateSyscallMmap(f *testing.F) {
 	contracts, addrs := testContractsSetup(f)
 	f.Fuzz(func(t *testing.T, addr uint32, siz uint32, heap uint32) {
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           heap,
 			ExitCode:       0,
 			Exited:         false,
@@ -139,10 +145,10 @@ func FuzzStateSyscallMmap(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
 		require.Equal(t, preStateRoot, state.Memory.MerkleRoot())
@@ -179,10 +185,12 @@ func FuzzStateSyscallExitGroup(f *testing.F) {
 		pc = pc & 0xFF_FF_FF_FC // align PC
 		nextPC := pc + 4
 		state := &State{
-			PC:             pc,
-			NextPC:         nextPC,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     pc,
+				NextPC: nextPC,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -200,10 +208,10 @@ func FuzzStateSyscallExitGroup(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, pc, state.PC)
-		require.Equal(t, nextPC, state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, pc, state.Cpu.PC)
+		require.Equal(t, nextPC, state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(exitCode), state.ExitCode)
 		require.Equal(t, true, state.Exited)
@@ -225,10 +233,12 @@ func FuzzStateSyscallFnctl(f *testing.F) {
 	contracts, addrs := testContractsSetup(f)
 	f.Fuzz(func(t *testing.T, fd uint32, cmd uint32) {
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -246,10 +256,10 @@ func FuzzStateSyscallFnctl(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -289,10 +299,12 @@ func FuzzStateHintRead(f *testing.F) {
 	f.Fuzz(func(t *testing.T, addr uint32, count uint32) {
 		preimageData := []byte("hello world")
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -314,10 +326,10 @@ func FuzzStateHintRead(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -342,10 +354,12 @@ func FuzzStatePreimageRead(f *testing.F) {
 			t.SkipNow()
 		}
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -372,10 +386,10 @@ func FuzzStatePreimageRead(f *testing.F) {
 		require.NoError(t, err)
 		require.True(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -403,10 +417,12 @@ func FuzzStateHintWrite(f *testing.F) {
 	f.Fuzz(func(t *testing.T, addr uint32, count uint32, randSeed int64) {
 		preimageData := []byte("hello world")
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -436,10 +452,10 @@ func FuzzStateHintWrite(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)
@@ -461,10 +477,12 @@ func FuzzStatePreimageWrite(f *testing.F) {
 	f.Fuzz(func(t *testing.T, addr uint32, count uint32) {
 		preimageData := []byte("hello world")
 		state := &State{
-			PC:             0,
-			NextPC:         4,
-			LO:             0,
-			HI:             0,
+			Cpu: CpuScalars{
+				PC:     0,
+				NextPC: 4,
+				LO:     0,
+				HI:     0,
+			},
 			Heap:           0,
 			ExitCode:       0,
 			Exited:         false,
@@ -489,10 +507,10 @@ func FuzzStatePreimageWrite(f *testing.F) {
 		require.NoError(t, err)
 		require.False(t, stepWitness.HasPreimage())
 
-		require.Equal(t, uint32(4), state.PC)
-		require.Equal(t, uint32(8), state.NextPC)
-		require.Equal(t, uint32(0), state.LO)
-		require.Equal(t, uint32(0), state.HI)
+		require.Equal(t, uint32(4), state.Cpu.PC)
+		require.Equal(t, uint32(8), state.Cpu.NextPC)
+		require.Equal(t, uint32(0), state.Cpu.LO)
+		require.Equal(t, uint32(0), state.Cpu.HI)
 		require.Equal(t, uint32(0), state.Heap)
 		require.Equal(t, uint8(0), state.ExitCode)
 		require.Equal(t, false, state.Exited)

cannon/mipsevm/instrumented.go

@@ -78,7 +78,7 @@ func (m *InstrumentedState) Step(proof bool) (wit *StepWitness, err error) {
 	m.lastPreimageOffset = ^uint32(0)
 
 	if proof {
-		insnProof := m.state.Memory.MerkleProof(m.state.PC)
+		insnProof := m.state.Memory.MerkleProof(m.state.Cpu.PC)
 		wit = &StepWitness{
 			State:    m.state.EncodeWitness(),
 			MemProof: insnProof[:],

cannon/mipsevm/mips.go

@@ -174,8 +174,8 @@ func (m *InstrumentedState) handleSyscall() error {
 	m.state.Registers[2] = v0
 	m.state.Registers[7] = v1
 
-	m.state.PC = m.state.NextPC
-	m.state.NextPC = m.state.NextPC + 4
+	m.state.Cpu.PC = m.state.Cpu.NextPC
+	m.state.Cpu.NextPC = m.state.Cpu.NextPC + 4
 	return nil
 }
 
@@ -184,7 +184,7 @@ func (m *InstrumentedState) pushStack(target uint32) {
 		return
 	}
 	m.debug.stack = append(m.debug.stack, target)
-	m.debug.caller = append(m.debug.caller, m.state.PC)
+	m.debug.caller = append(m.debug.caller, m.state.Cpu.PC)
 }
 
 func (m *InstrumentedState) popStack() {
@@ -192,7 +192,7 @@ func (m *InstrumentedState) popStack() {
 		return
 	}
 	if len(m.debug.stack) != 0 {
-		fn := m.debug.meta.LookupSymbol(m.state.PC)
+		fn := m.debug.meta.LookupSymbol(m.state.Cpu.PC)
 		topFn := m.debug.meta.LookupSymbol(m.debug.stack[len(m.debug.stack)-1])
 		if fn != topFn {
 			// most likely the function was inlined. Snap back to the last return.
@@ -209,12 +209,12 @@ func (m *InstrumentedState) popStack() {
 			m.debug.caller = m.debug.caller[:len(m.debug.caller)-1]
 		}
 	} else {
-		fmt.Printf("ERROR: stack underflow at pc=%x. step=%d\n", m.state.PC, m.state.Step)
+		fmt.Printf("ERROR: stack underflow at pc=%x. step=%d\n", m.state.Cpu.PC, m.state.Step)
 	}
 }
 
 func (m *InstrumentedState) Traceback() {
-	fmt.Printf("traceback at pc=%x. step=%d\n", m.state.PC, m.state.Step)
+	fmt.Printf("traceback at pc=%x. step=%d\n", m.state.Cpu.PC, m.state.Step)
 	for i := len(m.debug.stack) - 1; i >= 0; i-- {
 		s := m.debug.stack[i]
 		idx := len(m.debug.stack) - i - 1
@@ -222,83 +222,49 @@ func (m *InstrumentedState) Traceback() {
 	}
 }
 
-func (m *InstrumentedState) handleBranch(opcode uint32, insn uint32, rtReg uint32, rs uint32) error {
-	if m.state.NextPC != m.state.PC+4 {
-		panic("branch in delay slot")
-	}
-
-	shouldBranch := false
-	if opcode == 4 || opcode == 5 { // beq/bne
-		rt := m.state.Registers[rtReg]
-		shouldBranch = (rs == rt && opcode == 4) || (rs != rt && opcode == 5)
-	} else if opcode == 6 {
-		shouldBranch = int32(rs) <= 0 // blez
-	} else if opcode == 7 {
-		shouldBranch = int32(rs) > 0 // bgtz
-	} else if opcode == 1 {
-		// regimm
-		rtv := (insn >> 16) & 0x1F
-		if rtv == 0 { // bltz
-			shouldBranch = int32(rs) < 0
-		}
-		if rtv == 1 { // bgez
-			shouldBranch = int32(rs) >= 0
-		}
-	}
-
-	prevPC := m.state.PC
-	m.state.PC = m.state.NextPC // execute the delay slot first
-	if shouldBranch {
-		m.state.NextPC = prevPC + 4 + (signExtend(insn&0xFFFF, 16) << 2) // then continue with the instruction the branch jumps to.
-	} else {
-		m.state.NextPC = m.state.NextPC + 4 // branch not taken
-	}
-	return nil
-}
-
 func (m *InstrumentedState) handleHiLo(fun uint32, rs uint32, rt uint32, storeReg uint32) error {
 	val := uint32(0)
 	switch fun {
 	case 0x10: // mfhi
-		val = m.state.HI
+		val = m.state.Cpu.HI
 	case 0x11: // mthi
-		m.state.HI = rs
+		m.state.Cpu.HI = rs
 	case 0x12: // mflo
-		val = m.state.LO
+		val = m.state.Cpu.LO
 	case 0x13: // mtlo
-		m.state.LO = rs
+		m.state.Cpu.LO = rs
 	case 0x18: // mult
 		acc := uint64(int64(int32(rs)) * int64(int32(rt)))
-		m.state.HI = uint32(acc >> 32)
-		m.state.LO = uint32(acc)
+		m.state.Cpu.HI = uint32(acc >> 32)
+		m.state.Cpu.LO = uint32(acc)
 	case 0x19: // multu
 		acc := uint64(uint64(rs) * uint64(rt))
-		m.state.HI = uint32(acc >> 32)
-		m.state.LO = uint32(acc)
+		m.state.Cpu.HI = uint32(acc >> 32)
+		m.state.Cpu.LO = uint32(acc)
 	case 0x1a: // div
-		m.state.HI = uint32(int32(rs) % int32(rt))
-		m.state.LO = uint32(int32(rs) / int32(rt))
+		m.state.Cpu.HI = uint32(int32(rs) % int32(rt))
+		m.state.Cpu.LO = uint32(int32(rs) / int32(rt))
 	case 0x1b: // divu
-		m.state.HI = rs % rt
-		m.state.LO = rs / rt
+		m.state.Cpu.HI = rs % rt
+		m.state.Cpu.LO = rs / rt
 	}
 
 	if storeReg != 0 {
 		m.state.Registers[storeReg] = val
 	}
 
-	m.state.PC = m.state.NextPC
-	m.state.NextPC = m.state.NextPC + 4
+	m.state.Cpu.PC = m.state.Cpu.NextPC
+	m.state.Cpu.NextPC = m.state.Cpu.NextPC + 4
 	return nil
 }
 
 func (m *InstrumentedState) handleJump(linkReg uint32, dest uint32) error {
-	if m.state.NextPC != m.state.PC+4 {
+	if m.state.Cpu.NextPC != m.state.Cpu.PC+4 {
 		panic("jump in delay slot")
 	}
-	prevPC := m.state.PC
-	m.state.PC = m.state.NextPC
-	m.state.NextPC = dest
+	prevPC := m.state.Cpu.PC
+	m.state.Cpu.PC = m.state.Cpu.NextPC
+	m.state.Cpu.NextPC = dest
 	if linkReg != 0 {
 		m.state.Registers[linkReg] = prevPC + 8 // set the link-register to the instr after the delay slot instruction.
 	}
@@ -312,8 +278,8 @@ func (m *InstrumentedState) handleRd(storeReg uint32, val uint32, conditional bo
 	if storeReg != 0 && conditional {
 		m.state.Registers[storeReg] = val
 	}
-	m.state.PC = m.state.NextPC
-	m.state.NextPC = m.state.NextPC + 4
+	m.state.Cpu.PC = m.state.Cpu.NextPC
+	m.state.Cpu.NextPC = m.state.Cpu.NextPC + 4
 	return nil
 }
 
@@ -323,7 +289,7 @@ func (m *InstrumentedState) mipsStep() error {
 	}
 	m.state.Step += 1
 	// instruction fetch
-	insn := m.state.Memory.GetMemory(m.state.PC)
+	insn := m.state.Memory.GetMemory(m.state.Cpu.PC)
 	opcode := insn >> 26 // 6-bits
 
 	// j-type j/jal
@@ -333,7 +299,7 @@ func (m *InstrumentedState) mipsStep() error {
 			linkReg = 31
 		}
 		// Take top 4 bits of the next PC (its 256 MB region), and concatenate with the 26-bit offset
-		target := (m.state.NextPC & 0xF0000000) | ((insn & 0x03FFFFFF) << 2)
+		target := (m.state.Cpu.NextPC & 0xF0000000) | ((insn & 0x03FFFFFF) << 2)
 		m.pushStack(target)
 		return m.handleJump(linkReg, target)
 	}
@@ -369,7 +335,7 @@ func (m *InstrumentedState) mipsStep() error {
 	}
 
 	if (opcode >= 4 && opcode < 8) || opcode == 1 {
-		return m.handleBranch(opcode, insn, rtReg, rs)
+		return handleBranch(&m.state.Cpu, &m.state.Registers, opcode, insn, rtReg, rs)
 	}
 
 	storeAddr := uint32(0xFF_FF_FF_FF)

cannon/mipsevm/mips_instructions.go

@@ -174,3 +174,37 @@ func signExtend(dat uint32, idx uint32) uint32 {
 		return dat & mask
 	}
 }
+
+func handleBranch(cpu *CpuScalars, registers *[32]uint32, opcode uint32, insn uint32, rtReg uint32, rs uint32) error {
+	if cpu.NextPC != cpu.PC+4 {
+		panic("branch in delay slot")
+	}
+
+	shouldBranch := false
+	if opcode == 4 || opcode == 5 { // beq/bne
+		rt := registers[rtReg]
+		shouldBranch = (rs == rt && opcode == 4) || (rs != rt && opcode == 5)
+	} else if opcode == 6 {
+		shouldBranch = int32(rs) <= 0 // blez
+	} else if opcode == 7 {
+		shouldBranch = int32(rs) > 0 // bgtz
+	} else if opcode == 1 {
+		// regimm
+		rtv := (insn >> 16) & 0x1F
+		if rtv == 0 { // bltz
+			shouldBranch = int32(rs) < 0
+		}
+		if rtv == 1 { // bgez
+			shouldBranch = int32(rs) >= 0
+		}
+	}
+
+	prevPC := cpu.PC
+	cpu.PC = cpu.NextPC // execute the delay slot first
+	if shouldBranch {
+		cpu.NextPC = prevPC + 4 + (signExtend(insn&0xFFFF, 16) << 2) // then continue with the instruction the branch jumps to.
+	} else {
+		cpu.NextPC = cpu.NextPC + 4 // branch not taken
+	}
+	return nil
+}

cannon/mipsevm/patch.go

@@ -12,10 +12,12 @@ const HEAP_START = 0x05000000
 
 func LoadELF(f *elf.File) (*State, error) {
 	s := &State{
-		PC:        uint32(f.Entry),
-		NextPC:    uint32(f.Entry + 4),
-		HI:        0,
-		LO:        0,
+		Cpu: CpuScalars{
+			PC:     uint32(f.Entry),
+			NextPC: uint32(f.Entry + 4),
+			LO:     0,
+			HI:     0,
+		},
 		Heap:      HEAP_START,
 		Registers: [32]uint32{},
 		Memory:    NewMemory(),

cannon/mipsevm/state.go

@@ -12,17 +12,22 @@ import (
 // StateWitnessSize is the size of the state witness encoding in bytes.
 var StateWitnessSize = 226
 
+type CpuScalars struct {
+	PC     uint32 `json:"pc"`
+	NextPC uint32 `json:"nextPC"`
+	LO     uint32 `json:"lo"`
+	HI     uint32 `json:"hi"`
+}
+
 type State struct {
 	Memory *Memory `json:"memory"`
 
 	PreimageKey    common.Hash `json:"preimageKey"`
 	PreimageOffset uint32      `json:"preimageOffset"` // note that the offset includes the 8-byte length prefix
 
-	PC     uint32 `json:"pc"`
-	NextPC uint32 `json:"nextPC"`
-	LO     uint32 `json:"lo"`
-	HI     uint32 `json:"hi"`
-	Heap   uint32 `json:"heap"` // to handle mmap growth
+	Cpu CpuScalars `json:"cpu"`
+
+	Heap uint32 `json:"heap"` // to handle mmap growth
 
 	ExitCode uint8 `json:"exit"`
 	Exited   bool  `json:"exited"`
@@ -54,10 +59,10 @@ func (s *State) EncodeWitness() StateWitness {
 	out = append(out, memRoot[:]...)
 	out = append(out, s.PreimageKey[:]...)
 	out = binary.BigEndian.AppendUint32(out, s.PreimageOffset)
-	out = binary.BigEndian.AppendUint32(out, s.PC)
-	out = binary.BigEndian.AppendUint32(out, s.NextPC)
-	out = binary.BigEndian.AppendUint32(out, s.LO)
-	out = binary.BigEndian.AppendUint32(out, s.HI)
+	out = binary.BigEndian.AppendUint32(out, s.Cpu.PC)
+	out = binary.BigEndian.AppendUint32(out, s.Cpu.NextPC)
+	out = binary.BigEndian.AppendUint32(out, s.Cpu.LO)
+	out = binary.BigEndian.AppendUint32(out, s.Cpu.HI)
 	out = binary.BigEndian.AppendUint32(out, s.Heap)
 	out = append(out, s.ExitCode)
 	if s.Exited {

cannon/mipsevm/state_test.go

@@ -44,7 +44,7 @@ func TestState(t *testing.T) {
 			//require.NoError(t, err, "must load ELF into state")
 			programMem, err := os.ReadFile(fn)
 			require.NoError(t, err)
-			state := &State{PC: 0, NextPC: 4, Memory: NewMemory()}
+			state := &State{Cpu: CpuScalars{PC: 0, NextPC: 4}, Memory: NewMemory()}
 			err = state.Memory.SetMemoryRange(0, bytes.NewReader(programMem))
 			require.NoError(t, err, "load program into state")
 
@@ -54,7 +54,7 @@ func TestState(t *testing.T) {
 			us := NewInstrumentedState(state, oracle, os.Stdout, os.Stderr)
 
 			for i := 0; i < 1000; i++ {
-				if us.state.PC == endAddr {
+				if us.state.Cpu.PC == endAddr {
 					break
 				}
 				if exitGroup && us.state.Exited {
@@ -65,11 +65,11 @@ func TestState(t *testing.T) {
 			}
 
 			if exitGroup {
-				require.NotEqual(t, uint32(endAddr), us.state.PC, "must not reach end")
+				require.NotEqual(t, uint32(endAddr), us.state.Cpu.PC, "must not reach end")
 				require.True(t, us.state.Exited, "must set exited state")
 				require.Equal(t, uint8(1), us.state.ExitCode, "must exit with 1")
 			} else {
-				require.Equal(t, uint32(endAddr), us.state.PC, "must reach end")
+				require.Equal(t, uint32(endAddr), us.state.Cpu.PC, "must reach end")
 				done, result := state.Memory.GetMemory(baseAddrEnd+4), state.Memory.GetMemory(baseAddrEnd+8)
 				// inspect test result
 				require.Equal(t, done, uint32(1), "must be done")

op-challenger/game/fault/trace/cannon/prestate_test.go

@@ -44,15 +44,17 @@ func TestAbsolutePreStateCommitment(t *testing.T) {
 			Memory:         mipsevm.NewMemory(),
 			PreimageKey:    common.HexToHash("cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"),
 			PreimageOffset: 0,
-			PC:             0,
-			NextPC:         1,
-			LO:             0,
-			HI:             0,
-			Heap:           0,
-			ExitCode:       0,
-			Exited:         false,
-			Step:           0,
-			Registers:      [32]uint32{},
+			Cpu: mipsevm.CpuScalars{
+				PC:     0,
+				NextPC: 1,
+				LO:     0,
+				HI:     0,
+			},
+			Heap:      0,
+			ExitCode:  0,
+			Exited:    false,
+			Step:      0,
+			Registers: [32]uint32{},
 		}
 		expected, err := state.EncodeWitness().StateHash()
 		require.NoError(t, err)

op-challenger/game/fault/trace/cannon/test_data/state.json

@@ -2,10 +2,12 @@
   "memory": [],
   "preimageKey": "0xcccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc",
   "preimageOffset": 0,
-  "pc": 0,
-  "nextPC": 1,
-  "lo": 0,
-  "hi": 0,
+  "cpu": {
+    "pc": 0,
+    "nextPC": 1,
+    "lo": 0,
+    "hi": 0
+  },
   "heap": 0,
   "exit": 0,
   "exited": false,

packages/contracts-bedrock/semver-lock.json

@@ -124,8 +124,8 @@
     "sourceCodeHash": "0x3ff4a3f21202478935412d47fd5ef7f94a170402ddc50e5c062013ce5544c83f"
   },
   "src/cannon/MIPS.sol": {
-    "initCodeHash": "0x1c5dbe83af31e70feb906e2bda2bb1d78d3d15012ec6b11ba5643785657af2a6",
-    "sourceCodeHash": "0x9bdc97ff4e51fdec7c3e2113d5b60cd64eeb121a51122bea972789d4a5ac3dfa"
+    "initCodeHash": "0xf0fb656774ff761e2fd9adc523d4ee3dfad6fab63a37d4177543cfc98708b1be",
+    "sourceCodeHash": "0xe6a817aac69c149c24a9ab820853aae3c189eba337ab2e5c75a7cf458b45e308"
   },
   "src/cannon/PreimageOracle.sol": {
     "initCodeHash": "0xe5db668fe41436f53995e910488c7c140766ba8745e19743773ebab508efd090",

packages/contracts-bedrock/src/cannon/MIPS.sol

@@ -4,7 +4,8 @@ pragma solidity 0.8.15;
 import { ISemver } from "src/universal/ISemver.sol";
 import { IPreimageOracle } from "./interfaces/IPreimageOracle.sol";
 import { PreimageKeyLib } from "./PreimageKeyLib.sol";
-import "src/cannon/libraries/MIPSInstructions.sol" as ins;
+import { MIPSInstructions as ins } from "src/cannon/libraries/MIPSInstructions.sol";
+import { MIPSState as st } from "src/cannon/libraries/MIPSState.sol";
 
 /// @title MIPS
 /// @notice The MIPS contract emulates a single MIPS instruction.
@@ -45,7 +46,7 @@ contract MIPS is ISemver {
 
     /// @notice The semantic version of the MIPS contract.
     /// @custom:semver 1.0.1
-    string public constant version = "1.1.0-beta.1";
+    string public constant version = "1.1.0-beta.2";
 
     uint32 internal constant FD_STDIN = 0;
     uint32 internal constant FD_STDOUT = 1;
@@ -301,70 +302,6 @@ contract MIPS is ISemver {
         }
     }
 
-    /// @notice Handles a branch instruction, updating the MIPS state PC where needed.
-    /// @param _opcode The opcode of the branch instruction.
-    /// @param _insn The instruction to be executed.
-    /// @param _rtReg The register to be used for the branch.
-    /// @param _rs The register to be compared with the branch register.
-    /// @return out_ The hashed MIPS state.
-    function handleBranch(uint32 _opcode, uint32 _insn, uint32 _rtReg, uint32 _rs) internal returns (bytes32 out_) {
-        unchecked {
-            // Load state from memory
-            State memory state;
-            assembly {
-                state := 0x80
-            }
-
-            bool shouldBranch = false;
-
-            if (state.nextPC != state.pc + 4) {
-                revert("branch in delay slot");
-            }
-
-            // beq/bne: Branch on equal / not equal
-            if (_opcode == 4 || _opcode == 5) {
-                uint32 rt = state.registers[_rtReg];
-                shouldBranch = (_rs == rt && _opcode == 4) || (_rs != rt && _opcode == 5);
-            }
-            // blez: Branches if instruction is less than or equal to zero
-            else if (_opcode == 6) {
-                shouldBranch = int32(_rs) <= 0;
-            }
-            // bgtz: Branches if instruction is greater than zero
-            else if (_opcode == 7) {
-                shouldBranch = int32(_rs) > 0;
-            }
-            // bltz/bgez: Branch on less than zero / greater than or equal to zero
-            else if (_opcode == 1) {
-                // regimm
-                uint32 rtv = ((_insn >> 16) & 0x1F);
-                if (rtv == 0) {
-                    shouldBranch = int32(_rs) < 0;
-                }
-                if (rtv == 1) {
-                    shouldBranch = int32(_rs) >= 0;
-                }
-            }
-
-            // Update the state's previous PC
-            uint32 prevPC = state.pc;
-
-            // Execute the delay slot first
-            state.pc = state.nextPC;
-
-            // If we should branch, update the PC to the branch target
-            // Otherwise, proceed to the next instruction
-            if (shouldBranch) {
-                state.nextPC = prevPC + 4 + (ins.signExtend(_insn & 0xFFFF, 16) << 2);
-            } else {
-                state.nextPC = state.nextPC + 4;
-            }
-
-            // Return the hash of the resulting state
-            out_ = outputState();
-        }
-    }
-
     /// @notice Handles HI and LO register instructions.
     /// @param _func The function code of the instruction.
     /// @param _rs The value of the RS register.
@@ -730,7 +667,19 @@ contract MIPS is ISemver {
             }
 
             if ((opcode >= 4 && opcode < 8) || opcode == 1) {
-                return handleBranch(opcode, insn, rtReg, rs);
+                st.CpuScalars memory cpu = getCpuScalars(state);
+
+                ins.handleBranch({
+                    _cpu: cpu,
+                    _registers: state.registers,
+                    _opcode: opcode,
+                    _insn: insn,
+                    _rtReg: rtReg,
+                    _rs: rs
+                });
+                setStateCpuScalars(state, cpu);
+
+                return outputState();
             }
 
             uint32 storeAddr = 0xFF_FF_FF_FF;
@@ -796,4 +745,15 @@ contract MIPS is ISemver {
             return handleRd(rdReg, val, true);
         }
     }
+
+    function getCpuScalars(State memory _state) internal pure returns (st.CpuScalars memory) {
+        return st.CpuScalars({ pc: _state.pc, nextPC: _state.nextPC, lo: _state.lo, hi: _state.hi });
+    }
+
+    function setStateCpuScalars(State memory _state, st.CpuScalars memory _cpu) internal pure {
+        _state.pc = _cpu.pc;
+        _state.nextPC = _cpu.nextPC;
+        _state.lo = _cpu.lo;
+        _state.hi = _cpu.hi;
+    }
 }

packages/contracts-bedrock/src/cannon/libraries/MIPSState.sol

+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.15;
+
+library MIPSState {
+    struct CpuScalars {
+        uint32 pc;
+        uint32 nextPC;
+        uint32 lo;
+        uint32 hi;
+    }
+}

packages/contracts-bedrock/test/kontrol/proofs/utils/DeploymentSummaryFaultProofs.sol

@@ -25,7 +25,7 @@ contract DeploymentSummaryFaultProofs is DeploymentSummaryFaultProofsCode {
     address internal constant l1ERC721BridgeProxyAddress = 0xD31598c909d9C935a9e35bA70d9a3DD47d4D5865;
     address internal constant l1StandardBridgeAddress = 0xb7900B27Be8f0E0fF65d1C3A4671e1220437dd2b;
     address internal constant l1StandardBridgeProxyAddress = 0xDeF3bca8c80064589E6787477FFa7Dd616B5574F;
-    address internal constant mipsAddress = 0x1C0e3B8e58dd91536Caf37a6009536255A7816a6;
+    address internal constant mipsAddress = 0x477508A9612B67709Caf812D4356d531ba6a471d;
     address internal constant optimismPortal2Address = 0xfcbb237388CaF5b08175C9927a37aB6450acd535;
     address internal constant optimismPortalProxyAddress = 0x978e3286EB805934215a88694d80b09aDed68D90;
     address internal constant preimageOracleAddress = 0x3bd7E801E51d48c5d94Ea68e8B801DFFC275De75;