feat(state): wire MPT state root computation

lib/eevm/database.ex

@@ -120,6 +120,16 @@ defmodule EEVM.Database do
               key :: non_neg_integer()
             ) :: non_neg_integer()
 
+  @doc "Return all known account addresses in the backend state."
+  @callback account_addresses(state :: term()) :: [non_neg_integer()]
+
+  @doc "Return all storage slots currently tracked for an address."
+  @callback storage_slots(state :: term(), address :: non_neg_integer()) ::
+              [{non_neg_integer(), non_neg_integer()}]
+
+  @doc "Return addresses that currently have tracked storage entries."
+  @callback storage_addresses(state :: term()) :: [non_neg_integer()]
+
   @doc "Store a 256-bit value into a storage slot for the given contract address."
   @callback storage_store(
               state :: term(),
@@ -217,6 +227,21 @@ defmodule EEVM.Database do
     impl.storage_load(state, address, key)
   end
 
+  @spec account_addresses(t()) :: [non_neg_integer()]
+  def account_addresses(%__MODULE__{impl: impl, state: state}) do
+    impl.account_addresses(state)
+  end
+
+  @spec storage_slots(t(), non_neg_integer()) :: [{non_neg_integer(), non_neg_integer()}]
+  def storage_slots(%__MODULE__{impl: impl, state: state}, address) do
+    impl.storage_slots(state, address)
+  end
+
+  @spec storage_addresses(t()) :: [non_neg_integer()]
+  def storage_addresses(%__MODULE__{impl: impl, state: state}) do
+    impl.storage_addresses(state)
+  end
+
   @spec storage_store(t(), non_neg_integer(), non_neg_integer(), non_neg_integer()) :: t()
   def storage_store(%__MODULE__{impl: impl, state: state} = db, address, key, value) do
     %{db | state: impl.storage_store(state, address, key, value)}

lib/eevm/database/in_memory.ex

@@ -180,6 +180,23 @@ defmodule EEVM.Database.InMemory do
     |> Map.get(key, 0)
   end
 
+  @impl true
+  def account_addresses(%{accounts: accounts}) do
+    Map.keys(accounts)
+  end
+
+  @impl true
+  def storage_slots(%{storage: storage}, address) do
+    storage
+    |> Map.get(address, %{})
+    |> Map.to_list()
+  end
+
+  @impl true
+  def storage_addresses(%{storage: storage}) do
+    Map.keys(storage)
+  end
+
   @impl true
   def storage_store(%{storage: storage} = state, address, key, value) do
     address_storage = Map.get(storage, address, %{})

lib/eevm/state_root.ex

@@ -0,0 +1,73 @@
+defmodule EEVM.StateRoot do
+  @moduledoc """
+  Computes Ethereum storage roots and state roots from the current `EEVM.Database`.
+  """
+
+  alias EEVM.Database
+  alias EEVM.MPT.Trie
+
+  @storage_key_bytes 32
+  @address_bytes 20
+
+  @spec compute_storage_root(Database.t(), non_neg_integer()) :: binary()
+  def compute_storage_root(%Database{} = db, address) when is_integer(address) and address >= 0 do
+    entries =
+      db
+      |> Database.storage_slots(address)
+      |> Enum.reject(fn {_slot, value} -> value == 0 end)
+      |> Enum.sort_by(fn {slot, _value} -> slot end)
+      |> Enum.map(fn {slot, value} ->
+        {encode_uint(slot, @storage_key_bytes), ExRLP.encode(value)}
+      end)
+
+    Trie.secure_root_hash(entries)
+  end
+
+  @spec compute_state_root(Database.t()) :: binary()
+  def compute_state_root(%Database{} = db) do
+    account_addresses =
+      db
+      |> Database.account_addresses()
+      |> Enum.concat(storage_backed_addresses(db))
+      |> Enum.uniq()
+      |> Enum.sort()
+
+    entries =
+      Enum.map(account_addresses, fn address ->
+        {encode_uint(address, @address_bytes), encode_account(db, address)}
+      end)
+
+    Trie.secure_root_hash(entries)
+  end
+
+  defp encode_account(db, address) do
+    nonce = Database.get_nonce(db, address)
+    balance = Database.get_balance(db, address)
+    storage_root = compute_storage_root(db, address)
+    code_hash = ExKeccak.hash_256(Database.get_code(db, address))
+
+    ExRLP.encode([nonce, balance, storage_root, code_hash])
+  end
+
+  defp storage_backed_addresses(db) do
+    db
+    |> Database.storage_addresses()
+    |> Enum.flat_map(fn address ->
+      case Database.storage_slots(db, address)
+           |> Enum.reject(fn {_slot, value} -> value == 0 end) do
+        [] -> []
+        _slots -> [address]
+      end
+    end)
+  end
+
+  defp encode_uint(value, bytes) when is_integer(value) and value >= 0 do
+    encoded = :binary.encode_unsigned(value)
+
+    cond do
+      byte_size(encoded) == bytes -> encoded
+      byte_size(encoded) < bytes -> <<0::size((bytes - byte_size(encoded)) * 8), encoded::binary>>
+      true -> binary_part(encoded, byte_size(encoded) - bytes, bytes)
+    end
+  end
+end

test/database_test.exs

@@ -208,6 +208,25 @@ defmodule EEVM.DatabaseTest do
       db = Database.storage_store(db, 0xAA, 0, 99)
       assert Database.storage_load(db, 0xAA, 0) == 99
     end
+
+    test "account_addresses/1 returns all account keys" do
+      db = InMemory.new(accounts: %{0xAA => %{balance: 1}, 0xBB => %{nonce: 2}})
+
+      assert db |> Database.account_addresses() |> Enum.sort() == [0xAA, 0xBB]
+    end
+
+    test "storage_slots/2 returns tracked slot entries" do
+      db = InMemory.new(storage: %{0xAA => %{0 => 10, 9 => 11}})
+
+      assert db |> Database.storage_slots(0xAA) |> Enum.sort() == [{0, 10}, {9, 11}]
+      assert Database.storage_slots(db, 0xBB) == []
+    end
+
+    test "storage_addresses/1 returns addresses with tracked storage" do
+      db = InMemory.new(storage: %{0xAA => %{0 => 1}, 0xBB => %{2 => 2}})
+
+      assert db |> Database.storage_addresses() |> Enum.sort() == [0xAA, 0xBB]
+    end
   end
 
   describe "Database struct" do

test/state_root_test.exs

@@ -0,0 +1,79 @@
+defmodule EEVM.StateRootTest do
+  use ExUnit.Case, async: true
+
+  alias EEVM.Database.InMemory
+  alias EEVM.MPT.Trie
+  alias EEVM.StateRoot
+
+  test "compute_storage_root/2 returns empty trie hash for empty storage" do
+    db = InMemory.new()
+
+    assert StateRoot.compute_storage_root(db, 0xAA) == Trie.secure_root_hash([])
+  end
+
+  test "compute_storage_root/2 includes only non-zero slots" do
+    db =
+      InMemory.new(
+        storage: %{
+          0xAA => %{
+            0 => 0,
+            1 => 42,
+            5 => 999
+          }
+        }
+      )
+
+    expected =
+      Trie.secure_root_hash([
+        {encode_uint(1, 32), ExRLP.encode(42)},
+        {encode_uint(5, 32), ExRLP.encode(999)}
+      ])
+
+    assert StateRoot.compute_storage_root(db, 0xAA) == expected
+  end
+
+  test "compute_state_root/1 returns empty trie hash for empty world state" do
+    db = InMemory.new()
+
+    assert StateRoot.compute_state_root(db) == Trie.secure_root_hash([])
+  end
+
+  test "compute_state_root/1 computes root from account tuple" do
+    address = 0x1234
+
+    db =
+      InMemory.new(accounts: %{address => %{balance: 7, nonce: 3, code: <<0x60, 0x00, 0x00>>}})
+
+    storage_root = Trie.secure_root_hash([])
+    code_hash = ExKeccak.hash_256(<<0x60, 0x00, 0x00>>)
+    account_value = ExRLP.encode([3, 7, storage_root, code_hash])
+
+    expected = Trie.secure_root_hash([{encode_uint(address, 20), account_value}])
+
+    assert StateRoot.compute_state_root(db) == expected
+  end
+
+  test "compute_state_root/1 includes storage-backed accounts" do
+    address = 0xCAFE
+
+    db =
+      InMemory.new(
+        storage: %{
+          address => %{9 => 1, 10 => 0}
+        }
+      )
+
+    storage_root = Trie.secure_root_hash([{encode_uint(9, 32), ExRLP.encode(1)}])
+    code_hash = ExKeccak.hash_256(<<>>)
+    account_value = ExRLP.encode([0, 0, storage_root, code_hash])
+
+    expected = Trie.secure_root_hash([{encode_uint(address, 20), account_value}])
+
+    assert StateRoot.compute_state_root(db) == expected
+  end
+
+  defp encode_uint(value, bytes) do
+    encoded = :binary.encode_unsigned(value)
+    <<0::size((bytes - byte_size(encoded)) * 8), encoded::binary>>
+  end
+end