filler: allow genesis state root calc to modify pre-alloc

Author: marioevz

src/ethereum_test_tools/common/types.py

@@ -683,6 +683,15 @@ def check_alloc(self: "Account", address: str, alloc: dict):
             actual_storage = Storage(alloc["storage"]) if "storage" in alloc else Storage({})
             expected_storage.must_be_equal(address=address, other=actual_storage)
 
+    @classmethod
+    def from_dict(cls: Type, data: "Dict | Account") -> "Account":
+        """
+        Create account from dictionary.
+        """
+        if isinstance(data, cls):
+            return data
+        return cls(**data)
+
     @classmethod
     def with_code(cls: Type, code: BytesConvertible) -> "Account":
         """
@@ -691,7 +700,7 @@ def with_code(cls: Type, code: BytesConvertible) -> "Account":
         return Account(nonce=1, code=code)
 
 
-class Alloc(dict, Mapping[FixedSizeBytesConvertible, Account], SupportsJSON):
+class Alloc(dict, Mapping[FixedSizeBytesConvertible, Account | Dict], SupportsJSON):
     """
     Allocation of accounts in the state, pre and post test execution.
     """
@@ -700,7 +709,9 @@ def __json__(self, encoder: JSONEncoder) -> Mapping[str, Any]:
         """
         Returns the JSON representation of the allocation.
         """
-        return encoder.default({Address(address): account for address, account in self.items()})
+        return encoder.default(
+            {Address(address): Account.from_dict(account) for address, account in self.items()}
+        )
 
 
 def alloc_to_accounts(got_alloc: Dict[str, Any]) -> Mapping[str, Account]:

src/ethereum_test_tools/filling/fill.py

@@ -1,7 +1,6 @@
 """
 Filler object definitions.
 """
-from copy import copy
 from typing import List, Optional
 
 from ethereum_test_forks import Fork
@@ -25,11 +24,12 @@ def fill_test(
     """
     t8n.reset_traces()
 
-    genesis_rlp, genesis = test_spec.make_genesis(t8n, fork)
+    pre, genesis_rlp, genesis = test_spec.make_genesis(t8n, fork)
 
     (blocks, head, alloc) = test_spec.make_blocks(
         t8n,
         genesis,
+        pre,
         fork,
         eips=eips,
     )
@@ -41,7 +41,7 @@ def fill_test(
         genesis_rlp=genesis_rlp,
         head=head,
         fork="+".join([fork_name] + [str(eip) for eip in eips]) if eips is not None else fork_name,
-        pre_state=copy(test_spec.pre),
+        pre_state=pre,
         post_state=alloc_to_accounts(alloc),
         seal_engine=engine,
         name=test_spec.tag,

src/ethereum_test_tools/spec/base_test.py

@@ -10,7 +10,7 @@
 from ethereum_test_forks import Fork
 from evm_transition_tool import TransitionTool
 
-from ..common import Account, Address, Bytes, FixtureBlock, FixtureHeader, Hash, Transaction
+from ..common import Account, Address, Alloc, Bytes, FixtureBlock, FixtureHeader, Hash, Transaction
 
 
 def verify_transactions(txs: List[Transaction] | None, result) -> List[int]:
@@ -91,7 +91,7 @@ def make_genesis(
         self,
         t8n: TransitionTool,
         fork: Fork,
-    ) -> Tuple[Bytes, FixtureHeader]:
+    ) -> Tuple[Alloc, Bytes, FixtureHeader]:
         """
         Create a genesis block from the test definition.
         """
@@ -102,6 +102,7 @@ def make_blocks(
         self,
         t8n: TransitionTool,
         genesis: FixtureHeader,
+        pre: Alloc,
         fork: Fork,
         chain_id: int = 1,
         eips: Optional[List[int]] = None,

src/ethereum_test_tools/spec/blockchain_test.py

@@ -54,23 +54,22 @@ def make_genesis(
         self,
         t8n: TransitionTool,
         fork: Fork,
-    ) -> Tuple[Bytes, FixtureHeader]:
+    ) -> Tuple[Alloc, Bytes, FixtureHeader]:
         """
         Create a genesis block from the state test definition.
         """
         env = self.genesis_environment.set_fork_requirements(fork)
 
+        new_alloc, state_root = t8n.calc_state_root(
+            alloc=to_json(Alloc(self.pre)),
+            fork=fork,
+            debug_output_path=self.get_next_transition_tool_output_path(),
+        )
         genesis = FixtureHeader(
             parent_hash=Hash(0),
             ommers_hash=Hash(EmptyOmmersRoot),
             coinbase=Address(0),
-            state_root=Hash(
-                t8n.calc_state_root(
-                    alloc=to_json(Alloc(self.pre)),
-                    fork=fork,
-                    debug_output_path=self.get_next_transition_tool_output_path(),
-                )
-            ),
+            state_root=Hash(state_root),
             transactions_root=Hash(EmptyTrieRoot),
             receipt_root=Hash(EmptyTrieRoot),
             bloom=Bloom(0),
@@ -103,7 +102,7 @@ def make_genesis(
             withdrawals=env.withdrawals,
         )
 
-        return genesis_rlp, genesis
+        return Alloc(new_alloc), genesis_rlp, genesis
 
     def make_block(
         self,
@@ -256,6 +255,7 @@ def make_blocks(
         self,
         t8n: TransitionTool,
         genesis: FixtureHeader,
+        pre: Alloc,
         fork: Fork,
         chain_id=1,
         eips: Optional[List[int]] = None,
@@ -265,7 +265,7 @@ def make_blocks(
         Performs checks against the expected behavior of the test.
         Raises exception on invalid test behavior.
         """
-        alloc = to_json(Alloc(self.pre))
+        alloc = to_json(pre)
         env = Environment.from_parent_header(genesis)
         blocks: List[FixtureBlock] = []
         head = genesis.hash if genesis.hash is not None else Hash(0)

src/ethereum_test_tools/spec/state_test.py

@@ -54,7 +54,7 @@ def make_genesis(
         self,
         t8n: TransitionTool,
         fork: Fork,
-    ) -> Tuple[Bytes, FixtureHeader]:
+    ) -> Tuple[Alloc, Bytes, FixtureHeader]:
         """
         Create a genesis block from the state test definition.
         """
@@ -66,17 +66,17 @@ def make_genesis(
 
         env = env.set_fork_requirements(fork)
 
+        new_alloc, state_root = t8n.calc_state_root(
+            alloc=to_json(Alloc(self.pre)),
+            fork=fork,
+            debug_output_path=self.get_next_transition_tool_output_path(),
+        )
+
         genesis = FixtureHeader(
             parent_hash=Hash(0),
             ommers_hash=Hash(EmptyOmmersRoot),
             coinbase=Address(0),
-            state_root=Hash(
-                t8n.calc_state_root(
-                    alloc=to_json(Alloc(self.pre)),
-                    fork=fork,
-                    debug_output_path=self.get_next_transition_tool_output_path(),
-                )
-            ),
+            state_root=Hash(state_root),
             transactions_root=Hash(EmptyTrieRoot),
             receipt_root=Hash(EmptyTrieRoot),
             bloom=Bloom(0),
@@ -109,12 +109,13 @@ def make_genesis(
             withdrawals=env.withdrawals,
         )
 
-        return genesis_rlp, genesis
+        return Alloc(new_alloc), genesis_rlp, genesis
 
     def make_blocks(
         self,
         t8n: TransitionTool,
         genesis: FixtureHeader,
+        pre: Alloc,
         fork: Fork,
         chain_id=1,
         eips: Optional[List[int]] = None,
@@ -130,7 +131,7 @@ def make_blocks(
         txs = [tx.with_signature_and_sender() for tx in self.txs] if self.txs is not None else []
 
         alloc, result = t8n.evaluate(
-            alloc=to_json(Alloc(self.pre)),
+            alloc=to_json(pre),
             txs=to_json(txs),
             env=to_json(env),
             fork_name=fork.fork(block_number=Number(env.number), timestamp=Number(env.timestamp)),

src/ethereum_test_tools/tests/test_filler.py

@@ -76,7 +76,9 @@ def test_make_genesis(fork: Fork, hash: bytes):
 
     t8n = GethTransitionTool()
 
-    _, genesis = StateTest(env=env, pre=pre, post={}, txs=[], tag="some_state_test").make_genesis(
+    _, _, genesis = StateTest(
+        env=env, pre=pre, post={}, txs=[], tag="some_state_test"
+    ).make_genesis(
         t8n,
         fork,
     )

src/evm_transition_tool/tests/test_evaluate.py

@@ -76,7 +76,7 @@ class TestEnv:
 
     env = TestEnv()
     env.base_fee = base_fee
-    assert t8n.calc_state_root(alloc=alloc, fork=fork).startswith(hash)
+    assert t8n.calc_state_root(alloc=alloc, fork=fork)[1].startswith(hash)
 
 
 @pytest.mark.parametrize("evm_tool", [GethTransitionTool])

src/evm_transition_tool/transition_tool.py

@@ -205,7 +205,9 @@ def get_traces(self) -> List[List[List[Dict]]] | None:
         """
         return self.traces
 
-    def calc_state_root(self, *, alloc: Any, fork: Fork, debug_output_path: str = "") -> bytes:
+    def calc_state_root(
+        self, *, alloc: Any, fork: Fork, debug_output_path: str = ""
+    ) -> Tuple[Dict[str, Any], bytes]:
         """
         Calculate the state root for the given `alloc`.
         """
@@ -234,7 +236,7 @@ def calc_state_root(self, *, alloc: Any, fork: Fork, debug_output_path: str = ""
                 "beaconRoot"
             ] = "0x0000000000000000000000000000000000000000000000000000000000000000"
 
-        _, result = self.evaluate(
+        new_alloc, result = self.evaluate(
             alloc=alloc,
             txs=[],
             env=env,
@@ -244,7 +246,7 @@ def calc_state_root(self, *, alloc: Any, fork: Fork, debug_output_path: str = ""
         state_root = result.get("stateRoot")
         if state_root is None or not isinstance(state_root, str):
             raise Exception("Unable to calculate state root")
-        return bytes.fromhex(state_root[2:])
+        return new_alloc, bytes.fromhex(state_root[2:])
 
     def calc_withdrawals_root(
         self, *, withdrawals: Any, fork: Fork, debug_output_path: str = ""