feat(test-forks): add --state-trie state-commitment override flag

packages/testing/src/execution_testing/cli/pytest_commands/plugins/forks/forks.py

@@ -26,6 +26,7 @@
 from _pytest.mark.structures import ParameterSet
 from pytest import Mark, Metafunc, StashKey
 
+from execution_testing.base_types import StateCommitment
 from execution_testing.client_clis import TransitionTool
 from execution_testing.forks import (
     ALL_FORKS,
@@ -40,6 +41,7 @@
     get_transition_forks,
     transition_fork_to,
 )
+from execution_testing.forks.base_fork import BaseFork
 from execution_testing.logging import (
     get_logger,
 )
@@ -50,6 +52,13 @@
 # (see `get_unsupported_forks`).
 unsupported_forks_key: StashKey[FrozenSet[Fork | TransitionFork]] = StashKey()
 
+STATE_TRIE_CHOICES: Dict[str, StateCommitment] = {
+    "mpt": StateCommitment.MPT,
+    "merkle": StateCommitment.MPT,
+    "binary": StateCommitment.BINARY,
+    "bmt": StateCommitment.BINARY,
+}
+
 
 def pytest_addoption(parser: pytest.Parser) -> None:
     """Add command-line options to pytest."""
@@ -84,6 +93,20 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         default="",
         help="Fill tests until and including the specified fork.",
     )
+    fork_group.addoption(
+        "--state-trie",
+        action="store",
+        dest="state_trie",
+        default=None,
+        type=str.lower,
+        choices=list(STATE_TRIE_CHOICES),
+        help=(
+            "Override the state-commitment scheme used to compute state "
+            "roots for every fork: 'mpt'/'merkle' for the Merkle-Patricia "
+            "trie, 'binary'/'bmt' for the binary tree. By default, each "
+            "fork defines its own scheme."
+        ),
+    )
 
 
 @dataclass(kw_only=True)
@@ -569,6 +592,9 @@ def get_fork_option(
     forks_until = get_fork_option(config, "forks_until", "--until")
     show_fork_help = config.getoption("show_fork_help")
 
+    if state_trie := config.getoption("state_trie"):
+        BaseFork.set_state_commitment_override(STATE_TRIE_CHOICES[state_trie])
+
     dev_forks_help = textwrap.dedent(
         "To run tests for a fork under active development, it must be "
         "specified explicitly via --until=FORK.\n"
@@ -628,6 +654,11 @@ def get_fork_option(
         )
 
 
+def pytest_unconfigure() -> None:
+    """Reset global fork state derived from command-line options."""
+    BaseFork.set_state_commitment_override(None)
+
+
 def get_unsupported_forks(
     config: pytest.Config,
 ) -> FrozenSet[Fork | TransitionFork]:

packages/testing/src/execution_testing/forks/base_fork.py

@@ -346,6 +346,7 @@ class BaseFork(ForkOpcodeInterface, metaclass=BaseForkMeta):
     _deployed: ClassVar[bool] = True
     _enabled_eips: ClassVar[Set[int]] = set()
     _enabling_forks: ClassVar[Set[Type["BaseFork"]]] = set()
+    _state_commitment_override: ClassVar[Optional[StateCommitment]] = None
 
     # Method version bumps
     _engine_new_payload_version_bump: ClassVar[bool] = False
@@ -461,9 +462,27 @@ def __init_subclass__(
 
     @classmethod
     def state_commitment(cls) -> StateCommitment:
-        """Return the state-commitment scheme for the state root."""
+        """
+        Return the state-commitment scheme for the state root.
+
+        The `--state-trie` command-line option overrides the scheme for
+        every fork.
+        """
+        if BaseFork._state_commitment_override is not None:
+            return BaseFork._state_commitment_override
         return StateCommitment.MPT
 
+    @classmethod
+    def set_state_commitment_override(
+        cls, commitment: Optional[StateCommitment]
+    ) -> None:
+        """
+        Force every fork to report `commitment` from `state_commitment`.
+
+        `None` restores the fork-defined scheme.
+        """
+        BaseFork._state_commitment_override = commitment
+
     # Header information abstract methods
     @classmethod
     @abstractmethod

packages/testing/src/execution_testing/forks/tests/test_forks.py

@@ -5,9 +5,10 @@
 import pytest
 from pydantic import BaseModel
 
-from execution_testing.base_types import BlobSchedule
+from execution_testing.base_types import BlobSchedule, StateCommitment
 from execution_testing.vm import Opcodes
 
+from ..base_fork import BaseFork
 from ..forks.eips.paris.eip_3675 import EIP3675
 from ..forks.forks import (
     BPO1,
@@ -826,3 +827,23 @@ def test_oog_budget_lift() -> None:
         )
         == 3 * sstore + 2 * create + code_64
     )
+
+
+def test_state_commitment_override() -> None:
+    """
+    `--state-trie` forces the scheme on every fork; `None` restores the
+    fork-defined default.
+    """
+    assert Frontier.state_commitment() == StateCommitment.MPT
+    assert Amsterdam.state_commitment() == StateCommitment.MPT
+    BaseFork.set_state_commitment_override(StateCommitment.BINARY)
+    try:
+        assert Frontier.state_commitment() == StateCommitment.BINARY
+        assert Amsterdam.state_commitment() == StateCommitment.BINARY
+        assert (
+            BerlinToLondonAt5.transitions_from().state_commitment()
+            == StateCommitment.BINARY
+        )
+    finally:
+        BaseFork.set_state_commitment_override(None)
+    assert Amsterdam.state_commitment() == StateCommitment.MPT