tests/cancun/eip4788: Update to use pre-deployed contract

src/ethereum_test_forks/base_decorators.py

@@ -0,0 +1,12 @@
+"""
+Decorators for the fork methods.
+"""
+
+
+def prefer_transition_to_method(method):
+    """
+    Decorator to mark a base method that must always call the `fork_to` implementation when
+    transitioning.
+    """
+    method.__prefer_transition_to_method__ = True
+    return method

src/ethereum_test_forks/base_fork.py

@@ -2,7 +2,9 @@
 Abstract base class for Ethereum forks
 """
 from abc import ABC, ABCMeta, abstractmethod
-from typing import Optional, Type
+from typing import Mapping, Optional, Type
+
+from .base_decorators import prefer_transition_to_method
 
 
 class BaseForkMeta(ABCMeta):
@@ -103,6 +105,18 @@ def get_reward(cls, block_number: int = 0, timestamp: int = 0) -> int:
         """
         pass
 
+    @classmethod
+    @prefer_transition_to_method
+    @abstractmethod
+    def pre_allocation(cls, block_number: int = 0, timestamp: int = 0) -> Mapping:
+        """
+        Returns required pre-allocation of accounts.
+
+        This method must always call the `fork_to` method when transitioning, because the
+        allocation can only be set at genesis, and thus cannot be changed at transition time.
+        """
+        pass
+
     # Engine API information abstract methods
     @classmethod
     @abstractmethod

src/ethereum_test_forks/forks/forks.py

@@ -1,7 +1,7 @@
 """
 All Ethereum fork class definitions.
 """
-from typing import Optional
+from typing import Mapping, Optional
 
 from ..base_fork import BaseFork
 
@@ -99,6 +99,15 @@ def get_reward(cls, block_number: int = 0, timestamp: int = 0) -> int:
         """
         return 5_000_000_000_000_000_000
 
+    @classmethod
+    def pre_allocation(cls, block_number: int = 0, timestamp: int = 0) -> Mapping:
+        """
+        Returns whether the fork expects pre-allocation of accounts
+
+        Frontier does not require pre-allocated accounts
+        """
+        return {}
+
 
 class Homestead(Frontier):
     """
@@ -291,6 +300,21 @@ def header_beacon_root_required(cls, block_number: int = 0, timestamp: int = 0)
         """
         return True
 
+    @classmethod
+    def pre_allocation(cls, block_number: int = 0, timestamp: int = 0) -> Mapping:
+        """
+        Cancun requires pre-allocation of the beacon root contract for EIP-4788
+        """
+        new_allocation = {
+            0x0B: {
+                "nonce": 1,
+                "code": "0x3373fffffffffffffffffffffffffffffffffffffffe14604457602036146024575f5f"
+                "fd5b620180005f350680545f35146037575f5ffd5b6201800001545f5260205ff35b426201800042"
+                "06555f3562018000420662018000015500",
+            }
+        }
+        return new_allocation | super(Cancun, cls).pre_allocation(block_number, timestamp)
+
     @classmethod
     def engine_new_payload_version(
         cls, block_number: int = 0, timestamp: int = 0

src/ethereum_test_forks/tests/test_forks.py

@@ -2,7 +2,7 @@
 Test fork utilities.
 """
 
-from typing import cast
+from typing import Mapping, cast
 
 from ..base_fork import Fork
 from ..forks.forks import Berlin, Cancun, Frontier, London, Merge, Shanghai
@@ -17,6 +17,7 @@
     transition_fork_from_to,
     transition_fork_to,
 )
+from ..transition_base_fork import transition_fork
 
 FIRST_DEPLOYED = Frontier
 LAST_DEPLOYED = Shanghai
@@ -112,3 +113,53 @@ def test_deployed_forks():  # noqa: D103
     deployed_forks = get_deployed_forks()
     assert deployed_forks[0] == FIRST_DEPLOYED
     assert deployed_forks[-1] == LAST_DEPLOYED
+
+
+class PrePreAllocFork(Shanghai):
+    """
+    Dummy fork used for testing.
+    """
+
+    @classmethod
+    def pre_allocation(cls, block_number: int = 0, timestamp: int = 0) -> Mapping:
+        """
+        Return some starting point for allocation.
+        """
+        return {"test": "test"}
+
+
+class PreAllocFork(PrePreAllocFork):
+    """
+    Dummy fork used for testing.
+    """
+
+    @classmethod
+    def pre_allocation(cls, block_number: int = 0, timestamp: int = 0) -> Mapping:
+        """
+        Add allocation to the pre-existing one from previous fork.
+        """
+        return {"test2": "test2"} | super(PreAllocFork, cls).pre_allocation(
+            block_number, timestamp
+        )
+
+
+@transition_fork(to_fork=PreAllocFork, at_timestamp=15_000)
+class PreAllocTransitionFork(PrePreAllocFork):
+    """
+    PrePreAllocFork to PreAllocFork transition at Timestamp 15k
+    """
+
+    pass
+
+
+def test_pre_alloc():
+    assert PrePreAllocFork.pre_allocation() == {"test": "test"}
+    assert PreAllocFork.pre_allocation() == {"test": "test", "test2": "test2"}
+    assert PreAllocTransitionFork.pre_allocation() == {
+        "test": "test",
+        "test2": "test2",
+    }
+    assert PreAllocTransitionFork.pre_allocation(block_number=0, timestamp=0) == {
+        "test": "test",
+        "test2": "test2",
+    }

src/ethereum_test_forks/transition_base_fork.py

@@ -51,12 +51,14 @@ class NewTransitionClass(cls, TransitionBaseClass, BaseFork):  # type: ignore
 
         NewTransitionClass.name = lambda: transition_name  # type: ignore
 
-        def make_transition_method(from_fork_method, to_fork_method):
+        def make_transition_method(base_method, from_fork_method, to_fork_method):
             def transition_method(
                 cls,
                 block_number: int = ALWAYS_TRANSITIONED_BLOCK_NUMBER,
                 timestamp: int = ALWAYS_TRANSITIONED_BLOCK_TIMESTAMP,
             ):
+                if getattr(base_method, "__prefer_transition_to_method__", False):
+                    return to_fork_method(block_number, timestamp)
                 return (
                     to_fork_method(block_number, timestamp)
                     if block_number >= at_block and timestamp >= at_timestamp
@@ -70,6 +72,7 @@ def transition_method(
                 NewTransitionClass,
                 method_name,
                 make_transition_method(
+                    getattr(BaseFork, method_name),
                     getattr(from_fork, method_name),
                     getattr(to_fork, method_name),
                 ),

src/ethereum_test_tools/common/types.py

@@ -699,12 +699,54 @@ def with_code(cls: Type, code: BytesConvertible) -> "Account":
         """
         return Account(nonce=1, code=code)
 
+    @classmethod
+    def merge(
+        cls: Type, account_1: "Dict | Account | None", account_2: "Dict | Account | None"
+    ) -> "Account":
+        """
+        Create a merged account from two sources.
+        """
+
+        def to_kwargs_dict(account: "Dict | Account | None") -> Dict:
+            if account is None:
+                return {}
+            if isinstance(account, dict):
+                return account
+            elif isinstance(account, cls):
+                return {
+                    f.name: v for f in fields(cls) if (v := getattr(account, f.name)) is not None
+                }
+            raise TypeError(f"Unexpected type for account merge: {type(account)}")
+
+        kwargs = to_kwargs_dict(account_1)
+        kwargs.update(to_kwargs_dict(account_2))
+
+        return cls(**kwargs)
+
 
-class Alloc(dict, Mapping[FixedSizeBytesConvertible, Account | Dict], SupportsJSON):
+class Alloc(dict, Mapping[Address, Account], SupportsJSON):
     """
     Allocation of accounts in the state, pre and post test execution.
     """
 
+    def __init__(self, d: Mapping[FixedSizeBytesConvertible, Account | Dict] = {}):
+        for address, account in d.items():
+            address = Address(address)
+            assert address not in self, f"Duplicate address in alloc: {address}"
+            self[address] = Account.from_dict(account)
+
+    @classmethod
+    def merge(cls, alloc_1: "Alloc", alloc_2: "Alloc") -> "Alloc":
+        """
+        Returns the merged allocation of two sources.
+        """
+        merged = alloc_1.copy()
+
+        for address, other_account in alloc_2.items():
+            merged[address] = Account.merge(merged.get(address, None), other_account)
+
+        return Alloc(merged)
+
     def __json__(self, encoder: JSONEncoder) -> Mapping[str, Any]:
         """
         Returns the JSON representation of the allocation.

src/ethereum_test_tools/spec/blockchain_test.py

@@ -60,8 +60,9 @@ def make_genesis(
         """
         env = self.genesis_environment.set_fork_requirements(fork)
 
+        pre_alloc = Alloc(fork.pre_allocation(block_number=0, timestamp=Number(env.timestamp)))
         new_alloc, state_root = t8n.calc_state_root(
-            alloc=to_json(Alloc(self.pre)),
+            alloc=to_json(Alloc.merge(pre_alloc, Alloc(self.pre))),
             fork=fork,
             debug_output_path=self.get_next_transition_tool_output_path(),
         )

src/ethereum_test_tools/spec/state_test.py

@@ -66,8 +66,10 @@ def make_genesis(
 
         env = env.set_fork_requirements(fork)
 
+        pre_alloc = Alloc(fork.pre_allocation(block_number=0, timestamp=Number(env.timestamp)))
+
         new_alloc, state_root = t8n.calc_state_root(
-            alloc=to_json(Alloc(self.pre)),
+            alloc=to_json(Alloc.merge(pre_alloc, Alloc(self.pre))),
             fork=fork,
             debug_output_path=self.get_next_transition_tool_output_path(),
         )

src/ethereum_test_tools/tests/test_types.py

@@ -19,6 +19,7 @@
 from ..common.constants import TestPrivateKey
 from ..common.types import (
     Address,
+    Alloc,
     Bloom,
     Bytes,
     FixtureEngineNewPayload,
@@ -282,6 +283,74 @@ def test_account_check_alloc(account: Account, alloc: Dict[Any, Any], should_pas
             account.check_alloc("test", alloc)
 
 
+@pytest.mark.parametrize(
+    ["alloc1", "alloc2", "expected_alloc"],
+    [
+        pytest.param(
+            Alloc(),
+            Alloc(),
+            Alloc(),
+            id="empty_alloc",
+        ),
+        pytest.param(
+            Alloc({0x1: {"nonce": 1}}),
+            Alloc({0x2: {"nonce": 2}}),
+            Alloc({0x1: Account(nonce=1), 0x2: Account(nonce=2)}),
+            id="alloc_different_accounts",
+        ),
+        pytest.param(
+            Alloc({0x2: {"nonce": 1}}),
+            Alloc({"0x02": {"nonce": 2}}),
+            Alloc({0x2: Account(nonce=2)}),
+            id="overwrite_account",
+        ),
+        pytest.param(
+            Alloc({0x2: {"balance": 1}}),
+            Alloc({"0x02": {"nonce": 1}}),
+            Alloc({0x2: Account(balance=1, nonce=1)}),
+            id="mix_account",
+        ),
+    ],
+)
+def test_alloc_append(alloc1: Alloc, alloc2: Alloc, expected_alloc: Alloc):
+    assert Alloc.merge(alloc1, alloc2) == expected_alloc
+
+
+@pytest.mark.parametrize(
+    ["account1", "account2", "expected_account"],
+    [
+        pytest.param(
+            Account(),
+            Account(),
+            Account(),
+            id="empty_accounts",
+        ),
+        pytest.param(
+            None,
+            None,
+            Account(),
+            id="none_accounts",
+        ),
+        pytest.param(
+            Account(nonce=1),
+            Account(code="0x6000"),
+            Account(nonce=1, code="0x6000"),
+            id="accounts_with_different_fields",
+        ),
+        pytest.param(
+            Account(nonce=1),
+            Account(nonce=2),
+            Account(nonce=2),
+            id="accounts_with_different_nonce",
+        ),
+    ],
+)
+def test_account_merge(
+    account1: Account | None, account2: Account | None, expected_account: Account
+):
+    assert Account.merge(account1, account2) == expected_account
+
+
 CHECKSUM_ADDRESS = "0x8a0A19589531694250d570040a0c4B74576919B8"
 
 

src/evm_transition_tool/transition_tool.py

@@ -207,7 +207,7 @@ def get_traces(self) -> List[List[List[Dict]]] | None:
 
     def calc_state_root(
         self, *, alloc: Any, fork: Fork, debug_output_path: str = ""
-    ) -> Tuple[Dict[str, Any], bytes]:
+    ) -> Tuple[Dict, bytes]:
         """
         Calculate the state root for the given `alloc`.
         """