feat: state root

cairo/ethereum/cancun/fork_types.cairo

@@ -143,6 +143,10 @@ struct MappingAddressBytes32 {
     value: MappingAddressBytes32Struct*,
 }
 
+struct OptionalMappingAddressBytes32 {
+    value: MappingAddressBytes32Struct*,
+}
+
 struct ListTupleAddressBytes32 {
     value: ListTupleAddressBytes32Struct*,
 }

cairo/ethereum/cancun/state.cairo

@@ -17,6 +17,7 @@ from ethereum.cancun.fork_types import (
     MappingAddressAccountStruct,
     AddressAccountDictAccess,
     MappingAddressBytes32,
+    OptionalMappingAddressBytes32,
     MappingAddressBytes32Struct,
     AddressBytes32DictAccess,
     SetAddress,
@@ -63,7 +64,7 @@ from ethereum.cancun.trie import (
     copy_TrieTupleAddressBytes32U256,
 )
 from ethereum.cancun.blocks import Withdrawal
-from ethereum_types.bytes import Bytes, Bytes32
+from ethereum_types.bytes import Bytes, Bytes32, Bytes32Struct
 from ethereum_types.numeric import U256, U256Struct, Bool, bool, Uint
 from ethereum.utils.numeric import U256_le, U256_sub, U256_add, U256_mul
 from cairo_core.comparison import is_zero
@@ -81,6 +82,9 @@ from legacy.utils.dict import (
     dict_squash,
 )
 
+const EMPTY_ROOT_LOW = 0x6ef8c092e64583ffa655cc1b171fe856;
+const EMPTY_ROOT_HIGH = 0x21b463e3b52f6201c0ad6c991be0485b;
+
 struct AddressTrieBytes32U256DictAccess {
     key: Address,
     prev_value: TrieBytes32U256,
@@ -978,7 +982,8 @@ func storage_roots{
 
     // Create a Mapping[Address, Bytes32] that will contain the storage root of each address's
     // storage trie
-    let (map_addr_storage_root_start) = default_dict_new(0);
+    tempvar EMPTY_ROOT_PTR = new Bytes32Struct(EMPTY_ROOT_LOW, EMPTY_ROOT_HIGH);
+    let (map_addr_storage_root_start) = default_dict_new(cast(EMPTY_ROOT_PTR, felt));
     tempvar map_addr_storage_root = MappingAddressBytes32(
         new MappingAddressBytes32Struct(
             dict_ptr_start=cast(map_addr_storage_root_start, AddressBytes32DictAccess*),
@@ -1040,7 +1045,9 @@ func build_map_addr_storage_root{
         ),
     );
 
-    let storage_root = root(union_trie);
+    let storage_root = root(
+        union_trie, OptionalMappingAddressBytes32(cast(0, MappingAddressBytes32Struct*))
+    );
 
     let dict_ptr = cast(map_addr_storage_root.value.dict_ptr, DictAccess*);
     dict_write{dict_ptr=dict_ptr}(address.value, cast(storage_root.value, felt));
@@ -1138,3 +1145,35 @@ func build_storage_trie_for_address{
 
     return ();
 }
+
+func state_root{
+    range_check_ptr,
+    bitwise_ptr: BitwiseBuiltin*,
+    keccak_ptr: KeccakBuiltin*,
+    poseidon_ptr: PoseidonBuiltin*,
+}(state: State) -> Bytes32 {
+    alloc_locals;
+
+    // TODO: assert the state is empty
+
+    let storage_roots_ = storage_roots(state);
+
+    tempvar trie_union = EthereumTries(
+        new EthereumTriesEnum(
+            account=state.value._main_trie,
+            storage=TrieBytes32U256(cast(0, TrieBytes32U256Struct*)),
+            transaction=TrieBytesOptionalUnionBytesLegacyTransaction(
+                cast(0, TrieBytesOptionalUnionBytesLegacyTransactionStruct*)
+            ),
+            receipt=TrieBytesOptionalUnionBytesReceipt(
+                cast(0, TrieBytesOptionalUnionBytesReceiptStruct*)
+            ),
+            withdrawal=TrieBytesOptionalUnionBytesWithdrawal(
+                cast(0, TrieBytesOptionalUnionBytesWithdrawalStruct*)
+            ),
+        ),
+    );
+
+    let state_root = root(trie_union, OptionalMappingAddressBytes32(storage_roots_.value));
+    return state_root;
+}

cairo/ethereum/cancun/trie.cairo

@@ -11,7 +11,7 @@ from starkware.cairo.common.cairo_builtins import KeccakBuiltin
 from starkware.cairo.common.memcpy import memcpy
 
 from legacy.utils.bytes import uint256_to_bytes32_little
-from legacy.utils.dict import hashdict_read, hashdict_write, dict_new_empty, dict_squash
+from legacy.utils.dict import hashdict_read, hashdict_write, dict_new_empty, dict_squash, dict_read
 from ethereum.crypto.hash import keccak256
 from ethereum.utils.numeric import min
 from ethereum_rlp.rlp import encode, _encode_bytes, _encode
@@ -62,6 +62,10 @@ from ethereum.cancun.fork_types import (
     AddressAccountDictAccess,
     MappingTupleAddressBytes32U256,
     MappingTupleAddressBytes32U256Struct,
+    OptionalMappingAddressBytes32,
+    MappingAddressBytes32,
+    MappingAddressBytes32Struct,
+    AddressBytes32DictAccess,
     Root,
 )
 from ethereum.cancun.transactions_types import LegacyTransaction, LegacyTransactionStruct
@@ -959,7 +963,7 @@ func _prepare_trie{
     bitwise_ptr: BitwiseBuiltin*,
     keccak_ptr: KeccakBuiltin*,
     poseidon_ptr: PoseidonBuiltin*,
-}(trie_union: EthereumTries) -> MappingBytesBytes {
+}(trie_union: EthereumTries, storage_roots_: OptionalMappingAddressBytes32) -> MappingBytesBytes {
     alloc_locals;
 
     let (local mapping_ptr_start: BytesBytesDictAccess*) = default_dict_new(0);
@@ -982,9 +986,15 @@ func _prepare_trie{
     raise('Invalid trie union');
 
     account:
+    if (cast(storage_roots_.value, felt) == 0) {
+        raise('Missing Storage Roots');
+    }
     let account_trie = trie_union.value.account;
     _prepare_trie_inner_account(
-        account_trie, account_trie.value._data.value.dict_ptr_start, mapping_ptr_start
+        account_trie,
+        account_trie.value._data.value.dict_ptr_start,
+        mapping_ptr_start,
+        MappingAddressBytes32(storage_roots_.value),
     );
     jmp end;
 
@@ -1042,6 +1052,7 @@ func _prepare_trie_inner_account{
     trie: TrieAddressOptionalAccount,
     dict_ptr: AddressAccountDictAccess*,
     mapping_ptr_end: BytesBytesDictAccess*,
+    storage_roots_: MappingAddressBytes32,
 ) -> BytesBytesDictAccess* {
     alloc_locals;
 
@@ -1052,15 +1063,26 @@ func _prepare_trie_inner_account{
     // Skip all None values, which are deleted trie entries
     if (cast(dict_ptr.new_value.value, felt) == 0) {
         return _prepare_trie_inner_account(
-            trie, dict_ptr + AddressAccountDictAccess.SIZE, mapping_ptr_end
+            trie, dict_ptr + AddressAccountDictAccess.SIZE, mapping_ptr_end, storage_roots_
         );
     }
 
+    let storage_roots_ptr = cast(storage_roots_.value.dict_ptr, DictAccess*);
+    let (storage_root_ptr) = dict_read{dict_ptr=storage_roots_ptr}(dict_ptr.key.value);
+    let storage_root_b32 = Bytes32(cast(storage_root_ptr, Bytes32Struct*));
+    tempvar storage_roots_ = MappingAddressBytes32(
+        new MappingAddressBytes32Struct(
+            storage_roots_.value.dict_ptr_start,
+            cast(storage_roots_ptr, AddressBytes32DictAccess*),
+            storage_roots_.value.parent_dict,
+        ),
+    );
+    let storage_root = Bytes32_to_Bytes(storage_root_b32);
+
     let preimage = Bytes20_to_Bytes(dict_ptr.key);
     let value = dict_ptr.new_value;
-    // TODO: get storage root
+
     let (buffer: felt*) = alloc();
-    tempvar storage_root = Bytes(new BytesStruct(buffer, 0));
     tempvar node = Node(
         new NodeEnum(
             account=value,
@@ -1107,6 +1129,7 @@ func _prepare_trie_inner_account{
         trie,
         dict_ptr + AddressAccountDictAccess.SIZE,
         cast(mapping_dict_ptr, BytesBytesDictAccess*),
+        storage_roots_,
     );
 }
 
@@ -1458,11 +1481,10 @@ func root{
     bitwise_ptr: BitwiseBuiltin*,
     keccak_ptr: KeccakBuiltin*,
     poseidon_ptr: PoseidonBuiltin*,
-}(trie_union: EthereumTries) -> Root {
+}(trie_union: EthereumTries, storage_roots_: OptionalMappingAddressBytes32) -> Root {
     alloc_locals;
 
-    // TODO: get_storage_root
-    let obj = _prepare_trie(trie_union);
+    let obj = _prepare_trie(trie_union, storage_roots_);
     let patricialized = patricialize(obj, Uint(0));
     let root_node = encode_internal_node(patricialized);
     let rlp_encoded_root_node = encode(root_node);

cairo/tests/ethereum/cancun/test_state.py

@@ -30,13 +30,14 @@
     set_code,
     set_storage,
     set_transient_storage,
+    state_root,
     storage_root,
     touch_account,
 )
 from ethereum.cancun.trie import Trie, copy_trie
 from ethereum_types.bytes import Bytes32
 from ethereum_types.numeric import U256
-from hypothesis import given, settings
+from hypothesis import HealthCheck, given, settings
 from hypothesis import strategies as st
 from hypothesis.strategies import composite
 
@@ -564,6 +565,17 @@ def test_commit_transaction(
         assert transient_storage_cairo == transient_storage
 
 
+class TestRoot:
+    @given(state=...)
+    @settings(max_examples=1, suppress_health_check=[HealthCheck.filter_too_much])
+    def test_state_root(self, cairo_run, state: State):
+        # The state from the strategy contains a snapshot. Remove it
+        state._snapshots = []
+        state_root_cairo = cairo_run("state_root", state)
+        state_root_py = state_root(state)
+        assert state_root_cairo == state_root_py
+
+
 class TestStorageRoots:
     @given(state=...)
     def test_storage_roots(self, cairo_run, state: State):

cairo/tests/ethereum/cancun/test_trie.py

@@ -1,8 +1,10 @@
+from collections import defaultdict
 from typing import Mapping, Optional, Tuple, Union
 
 import pytest
 from ethereum.cancun.blocks import Receipt, Withdrawal
-from ethereum.cancun.fork_types import Account, Address
+from ethereum.cancun.fork_types import Account, Address, Root
+from ethereum.cancun.state import EMPTY_TRIE_ROOT
 from ethereum.cancun.transactions import LegacyTransaction
 from ethereum.cancun.trie import (
     BranchNode,
@@ -173,8 +175,13 @@ def test_internal_node_branch_node(self, cairo_run, branch_node: BranchNode):
         result = cairo_run("InternalNodeImpl.branch_node", branch_node)
         assert result == branch_node
 
-    @given(trie_with_none=prepare_trie_strategy())
-    def test_prepare_trie(self, cairo_run, trie_with_none: EthereumTries):
+    @given(trie_with_none=prepare_trie_strategy(), storage_tries=...)
+    def test_prepare_trie(
+        self,
+        cairo_run,
+        trie_with_none: EthereumTries,
+        storage_tries: Mapping[Address, Trie[Bytes32, U256]],
+    ):
         key_type, _ = trie_with_none.__orig_class__.__args__
 
         # Python expects tries not to have None values (as writing the default None suppresses the key)
@@ -186,26 +193,36 @@ def test_prepare_trie(self, cairo_run, trie_with_none: EthereumTries):
             _data={k: v for k, v in trie_with_none._data.items() if v is not None},
         )
 
-        # TODO: compute storage root
         if key_type is Address:
+            # Cairo expects a Dict[Address, Root] as storage_roots - not a callable function
+            storage_roots = defaultdict(
+                lambda: U256.from_le_bytes(EMPTY_TRIE_ROOT),
+                {address: root(storage_tries[address]) for address in storage_tries},
+            )
 
-            def get_storage_root(_address):
-                return b""
+            def get_storage_root(address: Address) -> Root:
+                return storage_roots[address].to_le_bytes32()
 
         else:
+            storage_roots = None
             get_storage_root = None
 
         try:
-            result_cairo = cairo_run("_prepare_trie", trie_with_none, get_storage_root)
+            result_cairo = cairo_run("_prepare_trie", trie_with_none, storage_roots)
         except Exception as e:
             with strict_raises(type(e)):
                 _prepare_trie(trie, get_storage_root)
             return
 
         assert result_cairo == _prepare_trie(trie, get_storage_root)
 
-    @given(trie_with_none=prepare_trie_strategy())
-    def test_root(self, cairo_run, trie_with_none: EthereumTries):
+    @given(trie_with_none=prepare_trie_strategy(), storage_tries=...)
+    def test_root(
+        self,
+        cairo_run,
+        trie_with_none: EthereumTries,
+        storage_tries: Mapping[Address, Trie[Bytes32, U256]],
+    ):
         key_type, _ = trie_with_none.__orig_class__.__args__
 
         # Python expects tries not to have None values (as writing the default None suppresses the key)
@@ -217,17 +234,22 @@ def test_root(self, cairo_run, trie_with_none: EthereumTries):
             _data={k: v for k, v in trie_with_none._data.items() if v is not None},
         )
 
-        # TODO: compute storage root
         if key_type is Address:
+            # Cairo expects a Dict[Address, Root] as storage_roots - not a callable function
+            storage_roots = defaultdict(
+                lambda: U256.from_le_bytes(EMPTY_TRIE_ROOT),
+                {address: root(storage_tries[address]) for address in storage_tries},
+            )
 
-            def get_storage_root(_address):
-                return b""
+            def get_storage_root(address: Address) -> Root:
+                return storage_roots[address].to_le_bytes32()
 
         else:
+            storage_roots = None
             get_storage_root = None
 
         try:
-            result_cairo = cairo_run("root", trie_with_none, get_storage_root)
+            result_cairo = cairo_run("root", trie_with_none, storage_roots)
         except Exception as e:
             with strict_raises(type(e)):
                 root(trie, get_storage_root)

cairo/tests/utils/args_gen.py

@@ -499,6 +499,9 @@ def __eq__(self, other):
     ("ethereum", "cancun", "fork_types", "MappingAddressBytes32"): Mapping[
         Address, Bytes32
     ],
+    ("ethereum", "cancun", "fork_types", "OptionalMappingAddressBytes32"): Optional[
+        Mapping[Address, Bytes32]
+    ],
     ("ethereum", "cancun", "fork_types", "Address"): Address,
     ("ethereum", "cancun", "fork_types", "SetAddress"): Set[Address],
     ("ethereum", "cancun", "fork_types", "Root"): Root,
@@ -1131,7 +1134,17 @@ def generate_dict_arg(
     }
 
     if isinstance_with_generic(arg, defaultdict):
-        data = defaultdict(arg.default_factory, data)
+        default_value = _gen_arg(
+            dict_manager,
+            segments,
+            type(arg.default_factory()),
+            arg.default_factory(),
+        )
+
+        def default_factory():
+            return default_value
+
+        data = defaultdict(default_factory, data)
 
     # This is required for tests where we read data from DictAccess segments while no dict method has been used.
     # Equivalent to doing an initial dict_read of all keys.