cannon: Fix j/jal VM emulation

Fixes two bugs:
- Jump offsets should be zero-extended
- Jumps are based on the PC-region. The effective address contains
 the current 256 MB-aligned region.

cannon/mipsevm/evm_test.go

@@ -165,6 +165,43 @@ func TestEVM(t *testing.T) {
 	}
 }
 
+func TestEVMSingleStep(t *testing.T) {
+	contracts, addrs := testContractsSetup(t)
+	var tracer vm.EVMLogger
+	//tracer = SourceMapTracer(t, contracts, addrs)
+
+	type testInput struct {
+		name   string
+		pc     uint32
+		nextPC uint32
+		insn   uint32
+	}
+	cases := []testInput{
+		{"j MSB set target", 0, 4, 0x0A_00_00_02},                         // j 0x02_00_00_02
+		{"j non-zero PC region", 0x10000000, 0x10000004, 0x08_00_00_02},   // j 0x2
+		{"jal MSB set target", 0, 4, 0x0E_00_00_02},                       // jal 0x02_00_00_02
+		{"jal non-zero PC region", 0x10000000, 0x10000004, 0x0C_00_00_02}, // jal 0x2
+	}
+
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			state := &State{PC: tt.pc, NextPC: tt.nextPC, Memory: NewMemory()}
+			state.Memory.SetMemory(tt.pc, tt.insn)
+
+			us := NewInstrumentedState(state, nil, os.Stdout, os.Stderr)
+			stepWitness, err := us.Step(true)
+			require.NoError(t, err)
+
+			evm := NewMIPSEVM(contracts, addrs)
+			evm.SetTracer(tracer)
+			evmPost := evm.Step(t, stepWitness)
+			goPost := us.state.EncodeWitness()
+			require.Equal(t, hexutil.Bytes(goPost).String(), hexutil.Bytes(evmPost).String(),
+				"mipsevm produced different state than EVM")
+		})
+	}
+}
+
 func TestEVMFault(t *testing.T) {
 	contracts, addrs := testContractsSetup(t)
 	var tracer vm.EVMLogger // no-tracer by default, but see SourceMapTracer and MarkdownTracer

cannon/mipsevm/mips.go

@@ -285,13 +285,13 @@ func (m *InstrumentedState) mipsStep() error {
 
 	// j-type j/jal
 	if opcode == 2 || opcode == 3 {
-		// TODO likely bug in original code: MIPS spec says this should be in the "current" region;
-		// a 256 MB aligned region (i.e. use top 4 bits of branch delay slot (pc+4))
 		linkReg := uint32(0)
 		if opcode == 3 {
 			linkReg = 31
 		}
-		return m.handleJump(linkReg, SE(insn&0x03FFFFFF, 26)<<2)
+		// Take top 4 bits of current PC (its 256 MB region), and concatenate with 26-bit offset
+		target := (m.state.NextPC & 0xF0000000) | ((insn & 0x03FFFFFF) << 2)
+		return m.handleJump(linkReg, target)
 	}
 
 	// register fetch
@@ -396,7 +396,6 @@ func (m *InstrumentedState) mipsStep() error {
 func execute(insn uint32, rs uint32, rt uint32, mem uint32) uint32 {
 	opcode := insn >> 26 // 6-bits
 	fun := insn & 0x3f   // 6-bits
-	// TODO(CLI-4136): deref the immed into a register
 
 	if opcode < 0x20 {
 		// transform ArithLogI

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

@@ -667,9 +667,9 @@ contract MIPS {
 
             // j-type j/jal
             if (opcode == 2 || opcode == 3) {
-                // TODO(CLI-4136): likely bug in original code: MIPS spec says this should be in the "current" region;
-                // a 256 MB aligned region (i.e. use top 4 bits of branch delay slot (pc+4))
-                return handleJump(opcode == 2 ? 0 : 31, SE(insn & 0x03FFFFFF, 26) << 2);
+                // Take top 4 bits of current PC (its 256 MB region), and concatenate with 26-bit offset
+                uint32 target = (state.nextPC & 0xF0000000) | (insn & 0x03FFFFFF) << 2;
+                return handleJump(opcode == 2 ? 0 : 31, target);
             }
 
             // register fetch
@@ -775,7 +775,6 @@ contract MIPS {
         unchecked {
             uint32 opcode = insn >> 26; // 6-bits
             uint32 func = insn & 0x3f; // 6-bits
-            // TODO(CLI-4136): deref the immed into a register
 
             if (opcode < 0x20) {
                 // transform ArithLogI

packages/contracts-bedrock/test/MIPS.t.sol

@@ -986,7 +986,7 @@ contract MIPS_Test is CommonTest {
     }
 
     function test_jump_succeeds() external {
-        uint16 label = 0x2;
+        uint32 label = 0x02_00_00_02; // set the 26th bit to assert no sign extension
         uint32 insn = uint32(0x08_00_00_00) | label; // j label
         (MIPS.State memory state, bytes memory proof) = constructMIPSState(0, insn, 0x4, 0);
 
@@ -1000,23 +1000,51 @@ contract MIPS_Test is CommonTest {
         assertEq(postState, outputState(expect), "unexpected post state");
     }
 
+    function test_jump_nonzeroRegion_succeeds() external {
+        uint32 pcRegion1 = 0x10000000;
+        uint32 label = 0x2;
+        uint32 insn = uint32(0x08_00_00_00) | label; // j label
+        (MIPS.State memory state, bytes memory proof) = constructMIPSState(pcRegion1, insn, 0x4, 0);
+
+        MIPS.State memory expect;
+        expect.memRoot = state.memRoot;
+        expect.pc = state.nextPC;
+        expect.nextPC = (state.nextPC & 0xF0_00_00_00) | (uint32(label) << 2);
+        expect.step = state.step + 1;
+
+        bytes memory witness = encodeState(state);
+        bytes32 postState = mips.step(witness, proof);
+        assertEq(postState, outputState(expect), "unexpected post state");
+    }
+
     function test_jal_succeeds() external {
-        uint32 pc = 0x0;
-        uint16 label = 0x2;
+        uint32 label = 0x02_00_00_02; // set the 26th bit to assert no sign extension
         uint32 insn = uint32(0x0c_00_00_00) | label; // jal label
-        (bytes32 memRoot, bytes memory proof) = ffi.getCannonMemoryProof(pc, insn);
-
-        MIPS.State memory state;
-        state.pc = 0;
-        state.nextPC = 4;
-        state.memRoot = memRoot;
+        (MIPS.State memory state, bytes memory proof) = constructMIPSState(0, insn, 0x4, 0);
 
         MIPS.State memory expect;
         expect.memRoot = state.memRoot;
         expect.pc = state.nextPC;
         expect.nextPC = label << 2;
         expect.step = state.step + 1;
-        expect.registers[31] = state.pc + 8;
+        expect.registers[31] = state.pc + 8; // ra
+
+        bytes32 postState = mips.step(encodeState(state), proof);
+        assertEq(postState, outputState(expect), "unexpected post state");
+    }
+
+    function test_jal_nonzeroRegion_succeeds() external {
+        uint32 pcRegion1 = 0x10000000;
+        uint32 label = 0x2;
+        uint32 insn = uint32(0x0c_00_00_00) | label; // jal label
+        (MIPS.State memory state, bytes memory proof) = constructMIPSState(pcRegion1, insn, 0x4, 0);
+
+        MIPS.State memory expect;
+        expect.memRoot = state.memRoot;
+        expect.pc = state.nextPC;
+        expect.nextPC = (state.nextPC & 0xF0_00_00_00) | (uint32(label) << 2);
+        expect.step = state.step + 1;
+        expect.registers[31] = state.pc + 8; // ra
 
         bytes32 postState = mips.step(encodeState(state), proof);
         assertEq(postState, outputState(expect), "unexpected post state");