feat(spec-tools): teach Hardfork to clone forks

Author: SamWilsn

src/ethereum_spec_tools/evm_tools/loaders/fixture_loader.py

@@ -21,6 +21,7 @@
     hex_to_u256,
     hex_to_uint,
 )
+from ethereum_spec_tools.forks import Hardfork
 
 from .fork_loader import ForkLoad
 from .transaction_loader import TransactionLoad
@@ -52,10 +53,18 @@ class Load(BaseLoad):
     _fork_module: str
     fork: ForkLoad
 
-    def __init__(self, network: str, fork_module: str):
+    def __init__(self, network: str, fork_module: str | Hardfork):
         self._network = network
-        self._fork_module = fork_module
-        self.fork = ForkLoad(fork_module)
+        if isinstance(fork_module, Hardfork):
+            self.fork = ForkLoad(fork_module)
+            self._fork_module = fork_module.short_name
+        else:
+            self._fork_module = fork_module
+            for fork in Hardfork.discover():
+                if fork.short_name == fork_module:
+                    self.fork = ForkLoad(fork)
+                    return
+            raise Exception(f"fork `{fork_module}` not found")
 
     def json_to_state(self, raw: Any) -> Any:
         """Converts json state data to a state object."""

src/ethereum_spec_tools/evm_tools/loaders/fork_loader.py

@@ -2,8 +2,7 @@
 Loader for code from the relevant fork.
 """
 
-import importlib
-from typing import Any
+from typing import Any, Final
 
 from ethereum_spec_tools.forks import Hardfork
 
@@ -13,46 +12,38 @@ class ForkLoad:
     Load the functions and classes from the relevant fork.
     """
 
-    _fork_module: str
+    hardfork: Final[Hardfork]
     _forks: list[Hardfork]
 
-    def __init__(self, fork_module: str):
-        self._fork_module = fork_module
+    def __init__(self, hardfork: Hardfork):
+        self.hardfork = hardfork
         self._forks = Hardfork.discover()
 
-    @property
-    def fork_module(self) -> str:
-        """Module that contains the fork code."""
-        return self._fork_module
-
     def _module(self, name: str) -> Any:
         """Imports a module from the fork."""
-        return importlib.import_module(
-            f"ethereum.forks.{self._fork_module}.{name}"
-        )
+        return self.hardfork.module(name)
 
     @property
     def proof_of_stake(self) -> bool:
         """Whether the fork is proof of stake."""
-        for fork in self._forks:
-            if fork.name == "ethereum.forks." + self._fork_module:
-                return fork.consensus.is_pos()
-        raise Exception(f"fork {self._fork_module} not discovered")
+        return self.hardfork.consensus.is_pos()
 
     def is_after_fork(self, target_fork_name: str) -> bool:
         """
         Check if the fork is after the target fork.
 
         Accepts short fork names (e.g. "cancun") instead of full module paths.
         """
-        return_value = False
+        target_fork: Hardfork | None = None
         for fork in self._forks:
-            short_name = fork.short_name
-            if short_name == target_fork_name:
-                return_value = True
-            if short_name == self._fork_module:
+            if fork.short_name == target_fork_name:
+                target_fork = fork
                 break
-        return return_value
+
+        if target_fork is None:
+            raise Exception(f"fork `{target_fork_name}` not found")
+
+        return self.hardfork.mod.FORK_CRITERIA >= target_fork.mod.FORK_CRITERIA
 
     @property
     def BEACON_ROOTS_ADDRESS(self) -> Any:

src/ethereum_spec_tools/evm_tools/t8n/__init__.py

@@ -19,7 +19,7 @@
 from ..loaders.fork_loader import ForkLoad
 from ..utils import (
     FatalError,
-    get_module_name,
+    find_fork,
     get_stream_logger,
     parse_hex_or_int,
 )
@@ -92,7 +92,7 @@ def __init__(
         self.out_file = out_file
         self.in_file = in_file
         self.options = options
-        self.forks = Hardfork.discover()
+        forks = Hardfork.discover()
 
         if "stdin" in (
             options.input_env,
@@ -103,9 +103,7 @@ def __init__(
         else:
             stdin = None
 
-        fork_module, self.fork_block = get_module_name(
-            self.forks, self.options, stdin
-        )
+        fork_module, self.fork_block = find_fork(forks, self.options, stdin)
         self.fork = ForkLoad(fork_module)
 
         tracers = GroupTracer()

src/ethereum_spec_tools/evm_tools/t8n/t8n_types.py

@@ -47,7 +47,7 @@ def __init__(self, t8n: "T8N", stdin: Optional[Dict] = None):
                     data[address][key] = "0x" + hex(int(value))
 
         state = t8n.json_to_state(data)
-        if t8n.fork.fork_module == "dao_fork":
+        if t8n.fork.hardfork.short_name == "dao_fork":
             t8n.fork.apply_dao(state)
 
         self.state = state

src/ethereum_spec_tools/evm_tools/utils.py

@@ -4,6 +4,7 @@
 
 import json
 import logging
+import re
 import sys
 from typing import (
     Any,
@@ -81,9 +82,9 @@ class FatalError(Exception):
     pass
 
 
-def get_module_name(
+def find_fork(
     forks: Sequence[Hardfork], options: Any, stdin: Any
-) -> Tuple[str, int]:
+) -> Tuple[Hardfork, int]:
     """
     Get the module name and the fork block for the given state fork.
     """
@@ -96,6 +97,13 @@ def get_module_name(
     except KeyError:
         pass
 
+    current_fork_block = 0
+    current_fork_module = re.sub(
+        r"(?<!^)(?=[A-Z])",
+        "_",
+        options.state_fork,
+    ).lower()  # CamelCase to snake_case
+
     if exception_config:
         if options.input_env == "stdin":
             assert stdin is not None
@@ -111,15 +119,9 @@ def get_module_name(
                 current_fork_module = fork
                 current_fork_block = fork_block
 
-        return current_fork_module, current_fork_block
-
-    # If the state fork is not an exception, use the fork name.
     for fork in forks:
-        fork_module = fork.name.split(".")[-1]
-        key = "".join(x.title() for x in fork_module.split("_"))
-
-        if key == options.state_fork:
-            return fork_module, 0
+        if current_fork_module == fork.short_name:
+            return fork, current_fork_block
 
     # Neither in exception nor a standard fork name.
     sys.exit(f"Unsupported state fork: {options.state_fork}")

src/ethereum_spec_tools/forks.py

@@ -8,9 +8,13 @@
 import importlib.abc
 import importlib.util
 import pkgutil
+import random
+from contextlib import AbstractContextManager
 from enum import Enum, auto
 from importlib.machinery import ModuleSpec
+from pathlib import Path
 from pkgutil import ModuleInfo
+from tempfile import TemporaryDirectory
 from types import ModuleType
 from typing import (
     TYPE_CHECKING,
@@ -21,12 +25,20 @@
     Optional,
     Type,
     TypeVar,
+    Union,
+    cast,
 )
 
-from ethereum_types.numeric import U256, Uint
+from ethereum_types.numeric import U64, U256, Uint
+from typing_extensions import override
 
 if TYPE_CHECKING:
-    from ethereum.fork_criteria import ForkCriteria
+    from ethereum.fork_criteria import (
+        ByBlockNumber,
+        ByTimestamp,
+        ForkCriteria,
+        Unscheduled,
+    )
 
 
 class ConsensusType(Enum):
@@ -167,6 +179,84 @@ def load_from_json(cls: Type[H], json: Any) -> List[H]:
 
         return cls.load(config)
 
+    @staticmethod
+    def clone(
+        template: H | str,
+        fork_criteria: Union[
+            "ByBlockNumber", "ByTimestamp", "Unscheduled", None
+        ] = None,
+        target_blob_gas_per_block: U64 | None = None,
+        gas_per_blob: U64 | None = None,
+        min_blob_gasprice: Uint | None = None,
+        blob_base_fee_update_fraction: Uint | None = None,
+        max_blob_gas_per_block: U64 | None = None,
+        blob_schedule_target: U64 | None = None,
+    ) -> "TemporaryHardfork":
+        """
+        Create a temporary clone of an existing fork, optionally tweaking its
+        parameters.
+        """
+        from .new_fork.builder import ForkBuilder
+
+        maybe_directory: TemporaryDirectory | None = TemporaryDirectory()
+
+        try:
+            assert maybe_directory is not None
+            directory: TemporaryDirectory = maybe_directory
+
+            if isinstance(template, str):
+                template_name = template
+            else:
+                template_name = template.short_name
+
+            clone_name = (
+                f"{template_name}_clone{random.randrange(1_000_000_000)}"
+            )
+
+            builder = ForkBuilder(template_name, clone_name)
+
+            builder.output = Path(directory.name)
+
+            if fork_criteria is not None:
+                builder.fork_criteria = fork_criteria
+
+            if target_blob_gas_per_block is not None:
+                builder.modify_target_blob_gas_per_block(
+                    target_blob_gas_per_block
+                )
+
+            if gas_per_blob is not None:
+                builder.modify_gas_per_blob(gas_per_blob)
+
+            if min_blob_gasprice is not None:
+                builder.modify_min_blob_gasprice(min_blob_gasprice)
+
+            if blob_base_fee_update_fraction is not None:
+                builder.modify_blob_base_fee_update_fraction(
+                    blob_base_fee_update_fraction
+                )
+
+            if max_blob_gas_per_block is not None:
+                builder.modify_max_blob_gas_per_block(max_blob_gas_per_block)
+
+            if blob_schedule_target is not None:
+                builder.modify_blob_schedule_target(blob_schedule_target)
+
+            builder.build()
+
+            clone_forks = Hardfork.discover([directory.name])
+            if len(clone_forks) != 1:
+                raise Exception("len(clone_forks) != 1")
+            if clone_forks[0].short_name != clone_name:
+                raise Exception("found incorrect fork")
+
+            value = TemporaryHardfork(clone_forks[0].mod, directory)
+            maybe_directory = None
+            return value
+        finally:
+            if maybe_directory is not None:
+                maybe_directory.cleanup()
+
     def __init__(self, mod: ModuleType) -> None:
         self.mod = mod
 
@@ -300,3 +390,28 @@ def walk_packages(self) -> Iterator[ModuleInfo]:
             raise ValueError(f"cannot walk {self.name}, path is None")
 
         return pkgutil.walk_packages([self.path], self.name + ".")
+
+
+class TemporaryHardfork(Hardfork, AbstractContextManager):
+    """
+    Short-lived `Hardfork` located in a temporary directory.
+    """
+
+    directory: TemporaryDirectory | None
+
+    def __init__(self, mod: ModuleType, directory: TemporaryDirectory) -> None:
+        super().__init__(mod)
+        self.directory = directory
+
+    @override
+    def __exit__(self, *args: object, **kwargs: object) -> None:
+        del args
+        del kwargs
+
+        assert self.directory is not None
+        self.directory.cleanup()
+        self.directory = None
+
+        # Intentionally break ourselves. Once the directory is gone, imports
+        # won't work.
+        self.mod = cast(ModuleType, None)