cannon: Add basic types for MTCannon (ethereum-optimism#11109)

* cannon: Rename StepWitness.MemProofs to ProofData

* cannon: Add MTState type (in progress)

* cannon: Tweak MtState tests to cover more ground

* cannon: Add test for MTState.UpdateCurrentThread()

* cannon: Use constants for byte size vars, set byte slice capacities

* cannon: Add StepsSinceLastContextSwitch field

* cannon: Rename witness offset constants

* cannon: Rename ThreadContext to ThreadState

* cannon: Panic on unimplemented method calls

* cannon: Compute thread stack roots lazily

* cannon: Push initial thread to left stack

Author: mbaxter

cannon/cmd/load_elf.go

@@ -46,7 +46,7 @@ func LoadELF(ctx *cli.Context) error {
 	if elfProgram.Machine != elf.EM_MIPS {
 		return fmt.Errorf("ELF is not big-endian MIPS R3000, but got %q", elfProgram.Machine.String())
 	}
-	state, err := mipsevm.LoadELF(elfProgram)
+	state, err := mipsevm.LoadELF(elfProgram, mipsevm.CreateInitialState)
 	if err != nil {
 		return fmt.Errorf("failed to load ELF data into VM state: %w", err)
 	}

cannon/cmd/run.go

@@ -446,7 +446,7 @@ func Run(ctx *cli.Context) error {
 				Pre:       witness.StateHash,
 				Post:      postStateHash,
 				StateData: witness.State,
-				ProofData: witness.MemProof,
+				ProofData: witness.ProofData,
 			}
 			if witness.HasPreimage() {
 				proof.OracleKey = witness.PreimageKey[:]

cannon/mipsevm/evm_test.go

@@ -113,7 +113,7 @@ func (m *MIPSEVM) Step(t *testing.T, stepWitness *StepWitness, step uint64) []by
 }
 
 func encodeStepInput(t *testing.T, wit *StepWitness, localContext LocalContext, mips *foundry.Artifact) []byte {
-	input, err := mips.ABI.Pack("step", wit.State, wit.MemProof, localContext)
+	input, err := mips.ABI.Pack("step", wit.State, wit.ProofData, localContext)
 	require.NoError(t, err)
 	return input
 }
@@ -485,8 +485,8 @@ func TestEVMFault(t *testing.T) {
 			insnProof := initialState.Memory.MerkleProof(0)
 			encodedWitness, _ := initialState.EncodeWitness()
 			stepWitness := &StepWitness{
-				State:    encodedWitness,
-				MemProof: insnProof[:],
+				State:     encodedWitness,
+				ProofData: insnProof[:],
 			}
 			input := encodeStepInput(t, stepWitness, LocalContext{}, contracts.MIPS)
 			startingGas := uint64(30_000_000)
@@ -509,7 +509,7 @@ func TestHelloEVM(t *testing.T) {
 	elfProgram, err := elf.Open("../example/bin/hello.elf")
 	require.NoError(t, err, "open ELF file")
 
-	state, err := LoadELF(elfProgram)
+	state, err := LoadELF(elfProgram, CreateInitialState)
 	require.NoError(t, err, "load ELF into state")
 
 	err = PatchGo(elfProgram, state)
@@ -560,7 +560,7 @@ func TestClaimEVM(t *testing.T) {
 	elfProgram, err := elf.Open("../example/bin/claim.elf")
 	require.NoError(t, err, "open ELF file")
 
-	state, err := LoadELF(elfProgram)
+	state, err := LoadELF(elfProgram, CreateInitialState)
 	require.NoError(t, err, "load ELF into state")
 
 	err = PatchGo(elfProgram, state)

cannon/mipsevm/instrumented.go

@@ -83,7 +83,7 @@ func (m *InstrumentedState) Step(proof bool) (wit *StepWitness, err error) {
 		wit = &StepWitness{
 			State:     encodedWitness,
 			StateHash: stateHash,
-			MemProof:  insnProof[:],
+			ProofData: insnProof[:],
 		}
 	}
 	err = m.mipsStep()
@@ -92,7 +92,7 @@ func (m *InstrumentedState) Step(proof bool) (wit *StepWitness, err error) {
 	}
 
 	if proof {
-		wit.MemProof = append(wit.MemProof, m.memProof[:]...)
+		wit.ProofData = append(wit.ProofData, m.memProof[:]...)
 		if m.lastPreimageOffset != ^uint32(0) {
 			wit.PreimageOffset = m.lastPreimageOffset
 			wit.PreimageKey = m.lastPreimageKey

cannon/mipsevm/patch.go

@@ -10,21 +10,11 @@ import (
 
 const HEAP_START = 0x05000000
 
-func LoadELF(f *elf.File) (*State, error) {
-	s := &State{
-		Cpu: CpuScalars{
-			PC:     uint32(f.Entry),
-			NextPC: uint32(f.Entry + 4),
-			LO:     0,
-			HI:     0,
-		},
-		Heap:      HEAP_START,
-		Registers: [32]uint32{},
-		Memory:    NewMemory(),
-		ExitCode:  0,
-		Exited:    false,
-		Step:      0,
-	}
+type CreateFPVMState[T FPVMState] func(pc, heapStart uint32) T
+
+func LoadELF[T FPVMState](f *elf.File, initState CreateFPVMState[T]) (T, error) {
+	var empty T
+	s := initState(uint32(f.Entry), HEAP_START)
 
 	for i, prog := range f.Progs {
 		if prog.Type == 0x70000003 { // MIPS_ABIFLAGS
@@ -37,21 +27,21 @@ func LoadELF(f *elf.File) (*State, error) {
 				if prog.Filesz < prog.Memsz {
 					r = io.MultiReader(r, bytes.NewReader(make([]byte, prog.Memsz-prog.Filesz)))
 				} else {
-					return nil, fmt.Errorf("invalid PT_LOAD program segment %d, file size (%d) > mem size (%d)", i, prog.Filesz, prog.Memsz)
+					return empty, fmt.Errorf("invalid PT_LOAD program segment %d, file size (%d) > mem size (%d)", i, prog.Filesz, prog.Memsz)
 				}
 			} else {
-				return nil, fmt.Errorf("program segment %d has different file size (%d) than mem size (%d): filling for non PT_LOAD segments is not supported", i, prog.Filesz, prog.Memsz)
+				return empty, fmt.Errorf("program segment %d has different file size (%d) than mem size (%d): filling for non PT_LOAD segments is not supported", i, prog.Filesz, prog.Memsz)
 			}
 		}
 
 		if prog.Vaddr+prog.Memsz >= uint64(1<<32) {
-			return nil, fmt.Errorf("program %d out of 32-bit mem range: %x - %x (size: %x)", i, prog.Vaddr, prog.Vaddr+prog.Memsz, prog.Memsz)
+			return empty, fmt.Errorf("program %d out of 32-bit mem range: %x - %x (size: %x)", i, prog.Vaddr, prog.Vaddr+prog.Memsz, prog.Memsz)
 		}
 		if prog.Vaddr+prog.Memsz >= HEAP_START {
-			return nil, fmt.Errorf("program %d overlaps with heap: %x - %x (size: %x). The heap start offset must be reconfigured", i, prog.Vaddr, prog.Vaddr+prog.Memsz, prog.Memsz)
+			return empty, fmt.Errorf("program %d overlaps with heap: %x - %x (size: %x). The heap start offset must be reconfigured", i, prog.Vaddr, prog.Vaddr+prog.Memsz, prog.Memsz)
 		}
-		if err := s.Memory.SetMemoryRange(uint32(prog.Vaddr), r); err != nil {
-			return nil, fmt.Errorf("failed to read program segment %d: %w", i, err)
+		if err := s.GetMemory().SetMemoryRange(uint32(prog.Vaddr), r); err != nil {
+			return empty, fmt.Errorf("failed to read program segment %d: %w", i, err)
 		}
 	}
 

cannon/mipsevm/state.go

@@ -10,8 +10,8 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 )
 
-// StateWitnessSize is the size of the state witness encoding in bytes.
-var StateWitnessSize = 226
+// STATE_WITNESS_SIZE is the size of the state witness encoding in bytes.
+const STATE_WITNESS_SIZE = 226
 
 type CpuScalars struct {
 	PC     uint32 `json:"pc"`
@@ -48,6 +48,32 @@ type State struct {
 	LastHint hexutil.Bytes `json:"lastHint,omitempty"`
 }
 
+func CreateEmptyState() *State {
+	return &State{
+		Cpu: CpuScalars{
+			PC:     0,
+			NextPC: 0,
+			LO:     0,
+			HI:     0,
+		},
+		Heap:      0,
+		Registers: [32]uint32{},
+		Memory:    NewMemory(),
+		ExitCode:  0,
+		Exited:    false,
+		Step:      0,
+	}
+}
+
+func CreateInitialState(pc, heapStart uint32) *State {
+	state := CreateEmptyState()
+	state.Cpu.PC = pc
+	state.Cpu.NextPC = pc + 4
+	state.Heap = heapStart
+
+	return state
+}
+
 type stateMarshaling struct {
 	Memory         *Memory       `json:"memory"`
 	PreimageKey    common.Hash   `json:"preimageKey"`
@@ -121,7 +147,7 @@ func (s *State) GetMemory() *Memory {
 }
 
 func (s *State) EncodeWitness() ([]byte, common.Hash) {
-	out := make([]byte, 0)
+	out := make([]byte, 0, STATE_WITNESS_SIZE)
 	memRoot := s.Memory.MerkleRoot()
 	out = append(out, memRoot[:]...)
 	out = append(out, s.PreimageKey[:]...)
@@ -132,11 +158,7 @@ func (s *State) EncodeWitness() ([]byte, common.Hash) {
 	out = binary.BigEndian.AppendUint32(out, s.Cpu.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 = AppendBoolToWitness(out, s.Exited)
 	out = binary.BigEndian.AppendUint64(out, s.Step)
 	for _, r := range s.Registers {
 		out = binary.BigEndian.AppendUint32(out, r)
@@ -154,14 +176,14 @@ const (
 )
 
 func (sw StateWitness) StateHash() (common.Hash, error) {
-	if len(sw) != 226 {
-		return common.Hash{}, fmt.Errorf("Invalid witness length. Got %d, expected 226", len(sw))
+	if len(sw) != STATE_WITNESS_SIZE {
+		return common.Hash{}, fmt.Errorf("Invalid witness length. Got %d, expected %d", len(sw), STATE_WITNESS_SIZE)
 	}
 	return stateHashFromWitness(sw), nil
 }
 
 func stateHashFromWitness(sw []byte) common.Hash {
-	if len(sw) != 226 {
+	if len(sw) != STATE_WITNESS_SIZE {
 		panic("Invalid witness length")
 	}
 	hash := crypto.Keccak256Hash(sw)