contracts,mipsevm: state format update progress, bugfixes, EVM MIPS-tests passing

mipsevm/evm.go

@@ -80,7 +80,7 @@ type Addresses struct {
 	Challenge  common.Address
 }
 
-func NewEVMEnv(contracts *Contracts, addrs *Addresses) *vm.EVM {
+func NewEVMEnv(contracts *Contracts, addrs *Addresses) (*vm.EVM, *state.StateDB) {
 	chainCfg := params.MainnetChainConfig
 	bc := &testChain{}
 	header := bc.GetHeader(common.Hash{}, 100)
@@ -100,7 +100,7 @@ func NewEVMEnv(contracts *Contracts, addrs *Addresses) *vm.EVM {
 	env.StateDB.SetCode(addrs.MIPSMemory, contracts.MIPSMemory.DeployedBytecode.Object)
 	env.StateDB.SetCode(addrs.Challenge, contracts.Challenge.DeployedBytecode.Object)
 	// TODO: any state to set, or immutables to replace, to link the contracts together?
-	return env
+	return env, state
 }
 
 type testChain struct {

mipsevm/evm_test.go

@@ -3,12 +3,14 @@ package main
 import (
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"math/big"
 	"os"
 	"path"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/require"
@@ -17,7 +19,6 @@ import (
 )
 
 func TestEVM(t *testing.T) {
-	t.Skip("work in progress!")
 
 	testFiles, err := os.ReadDir("test/bin")
 	require.NoError(t, err)
@@ -26,7 +27,7 @@ func TestEVM(t *testing.T) {
 	require.NoError(t, err)
 
 	// the first unlisted source seems to be the ABIDecoderV2 code that the compiler inserts
-	mipsSrcMap, err := contracts.MIPS.SourceMap([]string{"~ABIDecoderV2?", "~compiler?", "../contracts/src/MIPS.sol"})
+	mipsSrcMap, err := contracts.MIPS.SourceMap([]string{"../contracts/src/MIPS.sol", "~compiler?", "../contracts/src/MIPS.sol"})
 	require.NoError(t, err)
 
 	addrs := &Addresses{
@@ -42,14 +43,14 @@ func TestEVM(t *testing.T) {
 				t.Skip("oracle test needs to be updated to use syscall pre-image oracle")
 			}
 
-			env := NewEVMEnv(contracts, addrs)
-			env.Config.Debug = true
+			env, evmState := NewEVMEnv(contracts, addrs)
+			env.Config.Debug = false
 			//env.Config.Tracer = logger.NewMarkdownLogger(&logger.Config{}, os.Stdout)
 			env.Config.Tracer = mipsSrcMap.Tracer(os.Stdout)
 
 			fn := path.Join("test/bin", f.Name())
 			programMem, err := os.ReadFile(fn)
-			state := &State{PC: 0, Memory: make(map[uint32]*Page)}
+			state := &State{PC: 0, NextPC: 4, Memory: make(map[uint32]*Page)}
 			err = state.SetMemoryRange(0, bytes.NewReader(programMem))
 			require.NoError(t, err, "load program into state")
 
@@ -70,52 +71,37 @@ func TestEVM(t *testing.T) {
 			err = HookUnicorn(state, mu, os.Stdout, os.Stderr, al)
 			require.NoError(t, err, "hook unicorn to state")
 
-			// Add hook to stop unicorn once we reached the end of the test (i.e. "ate food")
-			_, err = mu.HookAdd(uc.HOOK_CODE, func(mu uc.Unicorn, addr uint64, size uint32) {
-				if state.PC == endAddr {
-					require.NoError(t, mu.Stop(), "stop test when returned")
-				}
-			}, 0, ^uint64(0))
-			require.NoError(t, err, "")
-
 			so := NewStateCache()
-			for i := 0; i < 1000; i++ {
-				insn := state.GetMemory(state.PC)
-
-				al.Reset() // reset
-				require.NoError(t, RunUnicorn(mu, state.PC, 1))
-				require.LessOrEqual(t, len(al.memReads)+len(al.memWrites), 1, "expecting at most a single mem read or write")
-
-				proofData := make([]byte, 0, 32*2)
-				proofData = append(proofData, uint32ToBytes32(32)...) // length in bytes
-				var tmp [32]byte
-				binary.BigEndian.PutUint32(tmp[0:4], insn) // instruction
-				if len(al.memReads) > 0 {
-					binary.BigEndian.PutUint32(tmp[4:8], state.GetMemory(al.memReads[0]))
+			var stateData []byte
+			var insn uint32
+			var pc uint32
+			var post []byte
+			preCode := func() {
+				insn = state.GetMemory(state.PC)
+				pc = state.PC
+				fmt.Printf("PRE - pc: %08x insn: %08x\n", pc, insn)
+				// remember the pre-state, to repeat it in the EVM during the post processing step
+				stateData = state.EncodeWitness(so)
+				if post != nil {
+					require.Equal(t, hexutil.Bytes(stateData).String(), hexutil.Bytes(post).String(),
+						"unicorn produced different state than EVM")
 				}
-				if len(al.memWrites) > 0 {
-					binary.BigEndian.PutUint32(tmp[4:8], state.GetMemory(al.memWrites[0]))
-				}
-				proofData = append(proofData, tmp[:]...)
 
-				memRoot := state.MerkleizeMemory(so)
+				al.Reset() // reset access list
+			}
+			postCode := func() {
+				fmt.Printf("POST - pc: %08x insn: %08x\n", pc, insn)
 
-				stateData := make([]byte, 0, 44*32)
-				stateData = append(stateData, memRoot[:]...)
-				stateData = append(stateData, make([]byte, 32)...) // TODO preimageKey
-				stateData = append(stateData, make([]byte, 32)...) // TODO preimageOffset
-				for i := 0; i < 32; i++ {
-					stateData = append(stateData, uint32ToBytes32(state.Registers[i])...)
+				var proofData []byte
+				proofData = binary.BigEndian.AppendUint32(proofData, insn)
+				if len(al.memReads) > 0 {
+					proofData = binary.BigEndian.AppendUint32(proofData, al.memReads[0].PreValue)
+				} else if len(al.memWrites) > 0 {
+					proofData = binary.BigEndian.AppendUint32(proofData, al.memWrites[0].PreValue)
+				} else {
+					proofData = append(proofData, make([]byte, 4)...)
 				}
-				stateData = append(stateData, uint32ToBytes32(state.PC)...)
-				stateData = append(stateData, uint32ToBytes32(state.NextPC)...)
-				stateData = append(stateData, uint32ToBytes32(state.LR)...)
-				stateData = append(stateData, uint32ToBytes32(state.LO)...)
-				stateData = append(stateData, uint32ToBytes32(state.HI)...)
-				stateData = append(stateData, uint32ToBytes32(state.Heap)...)
-				stateData = append(stateData, uint8ToBytes32(state.ExitCode)...)
-				stateData = append(stateData, boolToBytes32(state.Exited)...)
-				stateData = append(stateData, uint64ToBytes32(state.Step)...)
+				proofData = append(proofData, make([]byte, 32-4-4)...)
 
 				stateHash := crypto.Keccak256Hash(stateData)
 				var input []byte
@@ -129,14 +115,39 @@ func TestEVM(t *testing.T) {
 				input = append(input, uint32ToBytes32(uint32(len(proofData)))...) // proof data length in bytes
 				input = append(input, proofData[:]...)
 				startingGas := uint64(30_000_000)
+
+				// we take a snapshot so we can clean up the state, and isolate the logs of this instruction run.
+				snap := env.StateDB.Snapshot()
 				ret, leftOverGas, err := env.Call(vm.AccountRef(sender), addrs.MIPS, input, startingGas, big.NewInt(0))
 				require.NoError(t, err, "evm should not fail")
-				t.Logf("step took %d gas", startingGas-leftOverGas)
-				t.Logf("output (state hash): %x", ret)
-				// TODO compare output against unicorn (need to reconstruct state and memory hash)
+				require.Len(t, ret, 32, "expecting 32-byte state hash")
+				// remember state hash, to check it against state
+				postHash := common.Hash(*(*[32]byte)(ret))
+				logs := evmState.Logs()
+				require.Equal(t, 1, len(logs), "expecting a log with post-state")
+				post = logs[0].Data
+				require.Equal(t, crypto.Keccak256Hash(post), postHash, "logged state must be accurate")
+				env.StateDB.RevertToSnapshot(snap)
+
+				t.Logf("EVM step took %d gas, and returned stateHash %s", startingGas-leftOverGas, postHash)
 			}
 
+			firstStep := true
+			_, err = mu.HookAdd(uc.HOOK_CODE, func(mu uc.Unicorn, addr uint64, size uint32) {
+				if state.PC == endAddr {
+					require.NoError(t, mu.Stop(), "stop test when returned")
+				}
+				if !firstStep {
+					postCode()
+				}
+				preCode()
+				firstStep = false
+			}, 0, ^uint64(0))
+			require.NoError(t, err, "hook code")
+
+			err = RunUnicorn(mu, state.PC, 1000)
 			require.NoError(t, err, "must run steps without error")
+
 			// inspect test result
 			done, result := state.GetMemory(baseAddrEnd+4), state.GetMemory(baseAddrEnd+8)
 			require.Equal(t, done, uint32(1), "must be done")
@@ -145,28 +156,8 @@ func TestEVM(t *testing.T) {
 	}
 }
 
-func uint64ToBytes32(v uint64) []byte {
-	var out [32]byte
-	binary.BigEndian.PutUint64(out[32-8:], v)
-	return out[:]
-}
-
 func uint32ToBytes32(v uint32) []byte {
 	var out [32]byte
 	binary.BigEndian.PutUint32(out[32-4:], v)
 	return out[:]
 }
-
-func uint8ToBytes32(v uint8) []byte {
-	var out [32]byte
-	out[31] = v
-	return out[:]
-}
-
-func boolToBytes32(v bool) []byte {
-	var out [32]byte
-	if v {
-		out[31] = 1
-	}
-	return out[:]
-}

mipsevm/patch.go

@@ -11,6 +11,7 @@ import (
 func LoadELF(f *elf.File) (*State, error) {
 	s := &State{
 		PC:        uint32(f.Entry),
+		NextPC:    uint32(f.Entry + 4),
 		HI:        0,
 		LO:        0,
 		Heap:      0x20000000,

mipsevm/solutil.go

@@ -82,7 +82,7 @@ type SourceMap struct {
 func (s *SourceMap) Info(pc uint64) (source string, line uint32, col uint32) {
 	instr := s.Instr[pc]
 	if instr.F < 0 {
-		return
+		return "generated", 0, 0
 	}
 	if instr.F >= int32(len(s.Sources)) {
 		source = "unknown"
@@ -103,7 +103,7 @@ func (s *SourceMap) Info(pc uint64) (source string, line uint32, col uint32) {
 
 func (s *SourceMap) FormattedInfo(pc uint64) string {
 	f, l, c := s.Info(pc)
-	return fmt.Sprintf("%s:%d:%d %v", f, l, c, s.Instr[pc])
+	return fmt.Sprintf("%s:%d:%d", f, l, c)
 }
 
 // ParseSourceMap parses a solidity sourcemap: mapping bytecode indices to source references.
@@ -200,11 +200,21 @@ func (s *SourceMapTracer) CaptureEnter(typ vm.OpCode, from common.Address, to co
 func (s *SourceMapTracer) CaptureExit(output []byte, gasUsed uint64, err error) {}
 
 func (s *SourceMapTracer) CaptureState(pc uint64, op vm.OpCode, gas, cost uint64, scope *vm.ScopeContext, rData []byte, depth int, err error) {
-	fmt.Fprintf(s.out, "%s: pc %x opcode %s  map %v\n", s.srcMap.FormattedInfo(pc), pc, op.String(), s.srcMap.Instr[pc])
+	fmt.Fprintf(s.out, "%-40s : pc %x opcode %s\n", s.srcMap.FormattedInfo(pc), pc, op.String())
 }
 
 func (s *SourceMapTracer) CaptureFault(pc uint64, op vm.OpCode, gas, cost uint64, scope *vm.ScopeContext, depth int, err error) {
-	fmt.Fprintf(s.out, "%s: FAULT %v\n", s.srcMap.FormattedInfo(pc), err)
+	fmt.Fprintf(s.out, "%-40s: pc %x opcode %s FAULT %v\n", s.srcMap.FormattedInfo(pc), pc, op.String(), err)
+	fmt.Println("----")
+	fmt.Fprintf(s.out, "calldata: %x\n", scope.Contract.Input)
+	fmt.Println("----")
+	fmt.Fprintf(s.out, "memory: %x\n", scope.Memory.Data())
+	fmt.Println("----")
+	fmt.Fprintf(s.out, "stack:\n")
+	stack := scope.Stack.Data()
+	for i := range stack {
+		fmt.Fprintf(s.out, "%3d: %x\n", -i, stack[len(stack)-1-i].Bytes32())
+	}
 }
 
 var _ vm.EVMLogger = (*SourceMapTracer)(nil)

mipsevm/state.go

@@ -5,6 +5,8 @@ import (
 	"encoding/hex"
 	"fmt"
 	"io"
+
+	"github.com/ethereum/go-ethereum/common"
 )
 
 const (
@@ -36,27 +38,47 @@ func (p *Page) UnmarshalText(dat []byte) error {
 type State struct {
 	Memory map[uint32]*Page `json:"memory"`
 
-	Registers [32]uint32 `json:"registers"`
+	PreimageKey    common.Hash `json:"preimageKey"`
+	PreimageOffset uint32      `json:"preimageOffset"`
 
 	PC     uint32 `json:"pc"`
 	NextPC uint32 `json:"nextPC"`
-	LR     uint32 `json:"lr"`
-	HI     uint32 `json:"hi"`
 	LO     uint32 `json:"lo"`
+	HI     uint32 `json:"hi"`
 	Heap   uint32 `json:"heap"` // to handle mmap growth
 
 	ExitCode uint8 `json:"exit"`
 	Exited   bool  `json:"exited"`
 
 	Step uint64 `json:"step"`
+
+	Registers [32]uint32 `json:"registers"`
 }
 
-// TODO: VM state pre-image:
-// PC, HI, LO, Heap = 4 * 32/8 = 16 bytes
-// Registers = 32 * 32/8 = 256 bytes
-// Memory tree root = 32 bytes
-// Misc exit/step data = TBD
-// + proof(s) for memory leaf nodes
+func (s *State) EncodeWitness(so StateOracle) []byte {
+	out := make([]byte, 0)
+	memRoot := s.MerkleizeMemory(so)
+	memRoot = common.Hash{31: 42} // TODO need contract to actually write memory
+	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.Heap)
+	out = append(out, s.ExitCode)
+	if s.Exited {
+		out = append(out, 1)
+	} else {
+		out = append(out, 0)
+	}
+	out = binary.BigEndian.AppendUint64(out, s.Step)
+	for _, r := range s.Registers {
+		out = binary.BigEndian.AppendUint32(out, r)
+	}
+	return out
+}
 
 func (s *State) MerkleizeMemory(so StateOracle) [32]byte {
 	// empty parts of the tree are all zero. Precompute the hash of each full-zero range sub-tree level.
@@ -120,20 +142,21 @@ func (s *State) MerkleizeMemory(so StateOracle) [32]byte {
 	return merkleizeMemory(1, 0)
 }
 
-func (s *State) SetMemory(addr uint32, size uint32, v uint32) {
-	for i := size; i > 0; i-- {
-		pageIndex := addr >> pageAddrSize
-		pageAddr := addr & pageAddrMask
-		p, ok := s.Memory[pageIndex]
-		if !ok {
-			// allocate the page if we have not already.
-			// Go may mmap relatively large ranges, but we only allocate the pages just in time.
-			p = &Page{}
-			s.Memory[pageIndex] = p
-		}
-		p[pageAddr] = uint8(v >> (i - 1))
-		addr += 1
+func (s *State) SetMemory(addr uint32, v uint32) {
+	// addr must be aligned to 4 bytes
+	if addr&0x3 != 0 {
+		panic(fmt.Errorf("unaligned memory access: %x", addr))
+	}
+	pageIndex := addr >> pageAddrSize
+	pageAddr := addr & pageAddrMask
+	p, ok := s.Memory[pageIndex]
+	if !ok {
+		// allocate the page if we have not already.
+		// Go may mmap relatively large ranges, but we only allocate the pages just in time.
+		p = &Page{}
+		s.Memory[pageIndex] = p
 	}
+	binary.BigEndian.PutUint32(p[pageAddr:pageAddr+4], v)
 }
 
 func (s *State) GetMemory(addr uint32) uint32 {

mipsevm/state_test.go

@@ -37,7 +37,7 @@ func TestState(t *testing.T) {
 			//state, err := LoadELF(elfProgram)
 			//require.NoError(t, err, "must load ELF into state")
 			programMem, err := os.ReadFile(fn)
-			state := &State{PC: 0, Memory: make(map[uint32]*Page)}
+			state := &State{PC: 0, NextPC: 4, Memory: make(map[uint32]*Page)}
 			err = state.SetMemoryRange(0, bytes.NewReader(programMem))
 			require.NoError(t, err, "load program into state")
 
@@ -69,7 +69,7 @@ func TestState(t *testing.T) {
 					require.NoError(t, mu.Stop(), "stop test when returned")
 				}
 			}, 0, ^uint64(0))
-			require.NoError(t, err, "")
+			require.NoError(t, err, "hook code")
 
 			err = RunUnicorn(mu, state.PC, 1000)
 			require.NoError(t, err, "must run steps without error")

mipsevm/tracer.go

 package main
 
+type MemEntry struct {
+	EffAddr  uint32
+	PreValue uint32
+}
+
 type AccessList struct {
-	memReads  []uint32
-	memWrites []uint32
+	memReads  []MemEntry
+	memWrites []MemEntry
 }
 
 func (al *AccessList) Reset() {
@@ -10,39 +15,39 @@ func (al *AccessList) Reset() {
 	al.memWrites = al.memWrites[:0]
 }
 
-func (al *AccessList) OnRead(addr uint32) {
+func (al *AccessList) OnRead(effAddr uint32, preValue uint32) {
 	// if it matches the last, it's a duplicate; this happens because of multiple callbacks for the same effective addr.
-	if len(al.memReads) > 0 && al.memReads[len(al.memReads)-1] == addr {
+	if len(al.memReads) > 0 && al.memReads[len(al.memReads)-1].EffAddr == effAddr {
 		return
 	}
-	al.memReads = append(al.memReads, addr)
+	al.memReads = append(al.memReads, MemEntry{EffAddr: effAddr, PreValue: preValue})
 }
 
-func (al *AccessList) OnWrite(addr uint32) {
+func (al *AccessList) OnWrite(effAddr uint32, preValue uint32) {
 	// if it matches the last, it's a duplicate; this happens because of multiple callbacks for the same effective addr.
-	if len(al.memWrites) > 0 && al.memWrites[len(al.memWrites)-1] == addr {
+	if len(al.memWrites) > 0 && al.memWrites[len(al.memWrites)-1].EffAddr == effAddr {
 		return
 	}
-	al.memWrites = append(al.memWrites, addr)
+	al.memWrites = append(al.memWrites, MemEntry{EffAddr: effAddr, PreValue: preValue})
 }
 
 var _ Tracer = (*AccessList)(nil)
 
 type Tracer interface {
-	// OnRead remembers reads from the given addr.
+	// OnRead remembers reads from the given effAddr.
 	// Warning: the addr is an effective-addr, i.e. always aligned.
 	// But unicorn will fire it multiple times, for each byte that was changed within the effective addr boundaries.
-	OnRead(addr uint32)
-	// OnWrite remembers writes to the given addr.
+	OnRead(effAddr uint32, value uint32)
+	// OnWrite remembers writes to the given effAddr.
 	// Warning: the addr is an effective-addr, i.e. always aligned.
 	// But unicorn will fire it multiple times, for each byte that was changed within the effective addr boundaries.
-	OnWrite(addr uint32)
+	OnWrite(effAddr uint32, value uint32)
 }
 
 type NoOpTracer struct{}
 
-func (n NoOpTracer) OnRead(addr uint32) {}
+func (n NoOpTracer) OnRead(effAddr uint32, value uint32) {}
 
-func (n NoOpTracer) OnWrite(addr uint32) {}
+func (n NoOpTracer) OnWrite(effAddr uint32, value uint32) {}
 
 var _ Tracer = NoOpTracer{}

mipsevm/unicorn.go

@@ -109,8 +109,8 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer, tr Tracer)
 	}
 
 	_, err = mu.HookAdd(uc.HOOK_MEM_READ, func(mu uc.Unicorn, access int, addr64 uint64, size int, value int64) {
-		addr := uint32(addr64 & 0xFFFFFFFC) // pass effective addr to tracer
-		tr.OnRead(addr)
+		effAddr := uint32(addr64 & 0xFFFFFFFC) // pass effective addr to tracer
+		tr.OnRead(effAddr, st.GetMemory(effAddr))
 	}, 0, ^uint64(0))
 	if err != nil {
 		return fmt.Errorf("failed to set up mem-write hook: %w", err)
@@ -123,9 +123,26 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer, tr Tracer)
 		if size < 0 || size > 4 {
 			panic("invalid mem size")
 		}
-		st.SetMemory(uint32(addr64), uint32(size), uint32(value))
-		addr := uint32(addr64 & 0xFFFFFFFC) // pass effective addr to tracer
-		tr.OnWrite(addr)
+		effAddr := uint32(addr64 & 0xFFFFFFFC)
+		tr.OnWrite(effAddr, st.GetMemory(effAddr))
+
+		rt := value
+		rs := addr64 & 3
+		if size == 1 {
+			mem := st.GetMemory(effAddr)
+			val := uint32((rt & 0xFF) << (24 - (rs&3)*8))
+			mask := 0xFFFFFFFF ^ uint32(0xFF<<(24-(rs&3)*8))
+			st.SetMemory(effAddr, (mem&mask)|val)
+		} else if size == 2 {
+			mem := st.GetMemory(effAddr)
+			val := uint32((rt & 0xFFFF) << (16 - (rs&2)*8))
+			mask := 0xFFFFFFFF ^ uint32(0xFFFF<<(16-(rs&2)*8))
+			st.SetMemory(effAddr, (mem&mask)|val)
+		} else if size == 4 {
+			st.SetMemory(effAddr, uint32(rt))
+		} else {
+			log.Fatal("bad size write to ram")
+		}
 	}, 0, ^uint64(0))
 	if err != nil {
 		return fmt.Errorf("failed to set up mem-write hook: %w", err)
@@ -141,9 +158,57 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer, tr Tracer)
 		for i := 0; i < 32; i++ {
 			st.Registers[i] = uint32(batch[i])
 		}
+		prevPC := st.PC
 		st.PC = uint32(batch[32])
+
+		// We detect if we are potentially in a delay-slot.
+		// If we may be (i.e. last PC is 1 instruction before current),
+		// then parse the last instruction to determine what the next PC would be.
+		// This reflects the handleBranch / handleJump behavior that schedules next-PC.
+		if st.PC == prevPC+4 {
+			st.NextPC = prevPC + 8
+
+			prevInsn := st.GetMemory(prevPC)
+			opcode := prevInsn >> 26
+			switch opcode {
+			case 2, 3: // J/JAL
+				st.NextPC = signExtend(prevInsn&0x03FFFFFF, 25) << 2
+			case 1, 4, 5, 6, 7: // branching
+				rs := st.Registers[(prevInsn>>21)&0x1F]
+				shouldBranch := false
+				switch opcode {
+				case 4, 5:
+					rt := st.Registers[(prevInsn>>16)&0x1F]
+					shouldBranch = (rs == rt && opcode == 4) || (rs != rt && opcode == 5)
+				case 6:
+					shouldBranch = int32(rs) <= 0 // blez
+				case 7:
+					shouldBranch = int32(rs) > 0 // bgtz
+				case 1:
+					rtv := (prevInsn >> 16) & 0x1F
+					if rtv == 0 {
+						shouldBranch = int32(rs) < 0
+					} // bltz
+					if rtv == 1 {
+						shouldBranch = int32(rs) >= 0
+					} // bgez
+				}
+				if shouldBranch {
+					st.NextPC = prevPC + 4 + (signExtend(prevInsn&0xFFFF, 15) << 2)
+				}
+			case 0:
+				if funcv := prevInsn & 0x3f; funcv == 8 || funcv == 9 { // JR/JALR
+					rs := st.Registers[(prevInsn>>21)&0x1F]
+					st.NextPC = rs
+				}
+			}
+		} else {
+			st.NextPC = st.PC + 4
+		}
+
 		st.LO = uint32(batch[33])
 		st.HI = uint32(batch[34])
+		fmt.Printf("pc: 0x%08x\n", st.PC)
 	}, 0, ^uint64(0))
 	if err != nil {
 		return fmt.Errorf("failed to set up instruction hook: %w", err)
@@ -152,6 +217,15 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer, tr Tracer)
 	return nil
 }
 
+func signExtend(v uint32, i uint32) uint32 {
+	mask := ^((uint32(1) << i) - 1)
+	if v&(1<<i) != 0 {
+		return v | mask
+	} else {
+		return v &^ mask
+	}
+}
+
 func RunUnicorn(mu uc.Unicorn, entrypoint uint32, steps uint64) error {
 	return mu.StartWithOptions(uint64(entrypoint), ^uint64(0), &uc.UcOptions{
 		Timeout: 0, // 0 to disable, value is in ms.

mipsevm/unicorn_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
 	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
 )