A bit of EVM in Lean4

evm1.lean

namespace EthereumVm

inductive OpCode where
  | STOP
  | ADD
  | MUL
  | PUSH1 (value : UInt8)
  | MLOAD
  | MSTORE
  | MSTORE8
  | SLOAD
  | SSTORE

structure Stack where
  items : List UInt8

structure Memory where
  data : Array UInt8 := Array.empty

structure Storage where
  data : Array UInt8 := Array.empty

structure VMState where
  stack : Stack
  memory : Memory
  storage : Storage
  pc : Nat
  code : List UInt8
  halted : Bool

def Stack.push (s : Stack) (v : UInt8) : Stack :=
  { items := v :: s.items }

def Stack.pop (s : Stack) : Option (UInt8 × Stack) :=
  match s.items with
  | [] => none
  | x :: xs => some (x, { items := xs })

def Memory.read (mem : Memory) (offset : Nat) : UInt8 :=
  if h : offset < mem.data.size then
    mem.data.get ⟨offset, h⟩
  else
    0

def Memory.write (mem : Memory) (offset : Nat) (value : UInt8) : Memory :=
  let newSize := Nat.max (offset + 1) mem.data.size
  let newData := Array.mkArray newSize 0
  let newData := newData.set! offset value
  let newData := newData.mapIdx fun i v => 
    if i < offset then mem.data[i]!
    else if i > offset then mem.data[i]!
    else v
  { data := newData }

def Memory.write8 (mem : Memory) (offset : Nat) (value : UInt8) : Memory :=
  mem.write offset value

def Storage.read (storage : Storage) (offset : Nat) : UInt8 :=
  if h : offset < storage.data.size then
    storage.data.get ⟨offset, h⟩
  else
    0

def Storage.write (storage : Storage) (offset : Nat) (value : UInt8) : Storage :=
  let newSize := Nat.max (offset + 1) storage.data.size
  let newData := Array.mkArray newSize 0
  let newData := newData.set! offset value
  let newData := newData.mapIdx fun i v =>
    if i < offset then storage.data[i]!
    else if i > offset then storage.data[i]!
    else v
  { data := newData }

def decodeOpCode (byte : UInt8) : OpCode :=
  match byte with
  | 0x00 => OpCode.STOP
  | 0x01 => OpCode.ADD
  | 0x02 => OpCode.MUL
  | 0x51 => OpCode.MLOAD
  | 0x52 => OpCode.MSTORE
  | 0x53 => OpCode.MSTORE8
  | 0x54 => OpCode.SLOAD
  | 0x55 => OpCode.SSTORE
  | 0x60 => OpCode.PUSH1 0
  | _ => OpCode.STOP

def executeOpCode (state : VMState) : VMState :=
  if state.halted then state
  else
    match state.code[state.pc]? with
    | none => { state with halted := true }
    | some byte =>
      let op := decodeOpCode byte
      match op with
      | OpCode.STOP => { state with halted := true }
      | OpCode.ADD =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (a, s1) =>
          match s1.pop with
          | none => { state with halted := true }
          | some (b, s2) =>
            let result := a + b
            { state with
              stack := s2.push result,
              pc := state.pc + 1
            }
      | OpCode.MUL =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (a, s1) =>
          match s1.pop with
          | none => { state with halted := true }
          | some (b, s2) =>
            let result := a * b
            { state with
              stack := s2.push result,
              pc := state.pc + 1
            }
      | OpCode.PUSH1 _ =>
        match state.code[state.pc + 1]? with
        | none => { state with halted := true }
        | some value =>
          { state with
            stack := state.stack.push value,
            pc := state.pc + 2
          }
      | OpCode.MLOAD =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (offset, s) =>
          let value := state.memory.read offset.toNat
          { state with
            stack := s.push value,
            pc := state.pc + 1
          }
      | OpCode.MSTORE =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (offset, s1) =>
          match s1.pop with
          | none => { state with halted := true }
          | some (value, s2) =>
            { state with
              memory := state.memory.write offset.toNat value,
              stack := s2,
              pc := state.pc + 1
            }
      | OpCode.MSTORE8 =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (offset, s1) =>
          match s1.pop with
          | none => { state with halted := true }
          | some (value, s2) =>
            { state with
              memory := state.memory.write8 offset.toNat value,
              stack := s2,
              pc := state.pc + 1
            }
      | OpCode.SLOAD =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (offset, s) =>
          let value := state.storage.read offset.toNat
          { state with
            stack := s.push value,
            pc := state.pc + 1
          }
      | OpCode.SSTORE =>
        match state.stack.pop with
        | none => { state with halted := true }
        | some (offset, s1) =>
          match s1.pop with
          | none => { state with halted := true }
          | some (value, s2) =>
            { state with
              storage := state.storage.write offset.toNat value,
              stack := s2,
              pc := state.pc + 1
            }

def executeVM (fuel : Nat) (state : VMState) : VMState :=
  match fuel with
  | 0 => state
  | fuel + 1 =>
    if state.halted then state
    else executeVM fuel (executeOpCode state)
termination_by fuel

def opCodeToUInt8 (op : OpCode) : UInt8 :=
  match op with
  | OpCode.STOP => 0x00
  | OpCode.ADD => 0x01
  | OpCode.MUL => 0x02
  | OpCode.PUSH1 _ => 0x60
  | OpCode.MLOAD => 0x51
  | OpCode.MSTORE => 0x52
  | OpCode.MSTORE8 => 0x53
  | OpCode.SLOAD => 0x54
  | OpCode.SSTORE => 0x55

def opCodeListToUInt8List (code : List OpCode) : List UInt8 :=
  code.foldl (fun acc op => 
    match op with
    | OpCode.PUSH1 v => acc ++ [opCodeToUInt8 op, v]
    | _ => acc ++ [opCodeToUInt8 op]
  ) []

def uint8ToHexString (n : UInt8) : String :=
  let hexChars := "0123456789ABCDEF"
  let highNibble := n.toNat / 16
  let lowNibble := n.toNat % 16
  s!"{hexChars.get ⟨highNibble⟩}{hexChars.get ⟨lowNibble⟩}"

def opCodeToHex (op : OpCode) : String :=
  match op with
  | OpCode.STOP => "00"
  | OpCode.ADD => "01"
  | OpCode.MUL => "02"
  | OpCode.PUSH1 v => s!"60{uint8ToHexString v}"  -- Remove space after 60
  | OpCode.MLOAD => "51"
  | OpCode.MSTORE => "52"
  | OpCode.MSTORE8 => "53"
  | OpCode.SLOAD => "54"
  | OpCode.SSTORE => "55"

def generateBytecode (code : List OpCode) : String :=
  let bytecode := String.join (code.map opCodeToHex)
  if bytecode.length % 2 == 1 then
    "0" ++ bytecode  -- Prepend a '0' if the length is odd
  else
    bytecode

def writeBytecodeToFile (filename : String) (code : List OpCode) : IO Unit := do
  let bytecode := generateBytecode code
  IO.FS.writeFile filename bytecode

def verifyExecution (state : VMState) : String :=
  let expectedStack := [100,42,255, 10, 10]
  let expectedMemory := [5, 10, 10, 255]
  let expectedStorage := [42, 100]
  
  let stackCorrect := state.stack.items == expectedStack
  let memoryCorrect := state.memory.data.toList == expectedMemory
    let storageCorrect := state.storage.data.toList == expectedStorage

  if stackCorrect && memoryCorrect && storageCorrect then
    "Execution verified: The result is correct."
  else
    let stackMsg := 
      if stackCorrect then 
        "Stack is correct"
      else 
        s!"Stack mismatch. Expected: {expectedStack}, Got: {state.stack.items}"
    
    let memoryMsg := 
      if memoryCorrect then 
        "Memory is correct"
      else 
        s!"Memory mismatch. Expected: {expectedMemory}, Got: {state.memory.data.toList}"
    
    let storageMsg := 
      if storageCorrect then 
        "Storage is correct"
      else 
        s!"Storage mismatch. Expected: {expectedStorage}, Got: {state.storage.data.toList}"
    
    s!"Execution verification failed:\n{stackMsg}\n{memoryMsg}\n{storageMsg}\n"

def runExample : IO Unit := do
  let symbolicCode := [
    -- Store 5 at memory address 0
    OpCode.PUSH1 5,
    OpCode.PUSH1 0,
    OpCode.MSTORE,
    
    -- Store 10 at memory address 1
    OpCode.PUSH1 10,
    OpCode.PUSH1 1,
    OpCode.MSTORE,
    
    -- Load value from memory address 0, multiply by 2, and store at address 2
    OpCode.PUSH1 0,
    OpCode.MLOAD,
    OpCode.PUSH1 2,
    OpCode.MUL,
    OpCode.PUSH1 2,
    OpCode.MSTORE,
    
    -- Store 255 (0xFF) at memory address 3 using MSTORE8
    OpCode.PUSH1 255,
    OpCode.PUSH1 3,
    OpCode.MSTORE8,
    
    -- Store 42 at storage address 0
    OpCode.PUSH1 42,
    OpCode.PUSH1 0,
    OpCode.SSTORE,
    
    -- Store 100 at storage address 1
    OpCode.PUSH1 100,
    OpCode.PUSH1 1,
    OpCode.SSTORE,

    -- Load values from addresses 1, 2, and 3, and push them to the stack
    OpCode.PUSH1 1,
    OpCode.MLOAD,
    OpCode.PUSH1 2,
    OpCode.MLOAD,
    OpCode.PUSH1 3,
    OpCode.MLOAD,

    -- Load values from storage addresses 0 and 1, and push them to the stack
    OpCode.PUSH1 0,
    OpCode.SLOAD,
    OpCode.PUSH1 1,
    OpCode.SLOAD,
    OpCode.STOP
  ]

  let initialState : VMState := {
    stack := { items := [] },
    memory := { data := Array.empty },
    storage := { data := Array.empty },
    code := opCodeListToUInt8List symbolicCode,
    pc := 0,
    halted := false
  }
  
  let finalState := executeVM 1000 initialState
  
  IO.println s!"Final stack: {finalState.stack.items}"
  IO.println s!"Final memory: {finalState.memory.data.toList}"
  IO.println s!"Final storage: {finalState.storage.data.toList}"
  IO.println s!"Halted: {finalState.halted}"

  let executionResult := verifyExecution finalState
  IO.println executionResult

  IO.println s!"Program counter: {finalState.pc}"
  IO.println s!"Code length: {finalState.code.length}"

  -- Generate and write bytecode to file
  let bytecodeFilename := "ethereum_bytecode.txt"
  writeBytecodeToFile bytecodeFilename symbolicCode

  IO.println s!"Bytecode written to {bytecodeFilename}"

end EthereumVm