Add flag to count opcodes during execution

Author: SamWilsn

src/ethereum_spec_tools/evm_tools/daemon.py

@@ -80,6 +80,11 @@ def do_POST(self) -> None:
                 ]
             )
 
+        count_opcodes = content.get("count-opcodes", False)
+        if count_opcodes:
+            # send full opcode counts if ``count-opcodes`` is ``True``
+            args.extend(["--opcodes.count", "stdout"])
+
         query_string = urlparse(self.path).query
         if query_string:
             query = parse_qs(

src/ethereum_spec_tools/evm_tools/t8n/__init__.py

@@ -6,7 +6,7 @@
 import fnmatch
 import json
 import os
-from typing import Any, TextIO
+from typing import Any, Final, TextIO
 
 from ethereum_rlp import rlp
 from ethereum_types.numeric import U64, U256, Uint
@@ -24,7 +24,10 @@
     parse_hex_or_int,
 )
 from .env import Env
+from .evm_trace.count import evm_trace as evm_trace_count
+from .evm_trace.count import results as count_results
 from .evm_trace.eip3155 import Eip3155Tracer
+from .evm_trace.group import GroupTracer
 from .t8n_types import Alloc, Result, Txs
 
 
@@ -72,12 +75,16 @@ def t8n_arguments(subparsers: argparse._SubParsersAction) -> None:
     t8n_parser.add_argument("--trace.nostack", action="store_true")
     t8n_parser.add_argument("--trace.returndata", action="store_true")
 
+    t8n_parser.add_argument("--opcode.count", dest="opcode_count", type=str)
+
     t8n_parser.add_argument("--state-test", action="store_true")
 
 
 class T8N(Load):
     """The class that carries out the transition"""
 
+    tracers: Final[GroupTracer | None]
+
     def __init__(
         self, options: Any, out_file: TextIO, in_file: TextIO
     ) -> None:
@@ -100,18 +107,33 @@ def __init__(
         )
         self.fork = ForkLoad(fork_module)
 
+        tracers = GroupTracer()
+
         if self.options.trace:
             trace_memory = getattr(self.options, "trace.memory", False)
             trace_stack = not getattr(self.options, "trace.nostack", False)
             trace_return_data = getattr(self.options, "trace.returndata")
-            trace.set_evm_trace(
+            tracers.add(
                 Eip3155Tracer(
                     trace_memory=trace_memory,
                     trace_stack=trace_stack,
                     trace_return_data=trace_return_data,
                     output_basedir=self.options.output_basedir,
                 )
             )
+
+        if self.options.opcode_count is not None:
+            tracers.add(evm_trace_count)
+
+        maybe_tracers: GroupTracer | None
+        if tracers.tracers:
+            trace.set_evm_trace(tracers)
+            maybe_tracers = tracers
+        else:
+            maybe_tracers = None
+
+        self.tracers = maybe_tracers
+
         self.logger = get_stream_logger("T8N")
 
         super().__init__(
@@ -310,7 +332,7 @@ def run(self) -> int:
         json_state = self.alloc.to_json()
         json_result = self.result.to_json()
 
-        json_output = {}
+        json_output: dict[str, object] = {}
 
         if self.options.output_body == "stdout":
             txs_rlp = "0x" + rlp.encode(self.txs.all_txs).hex()
@@ -347,6 +369,18 @@ def run(self) -> int:
                 json.dump(json_result, f, indent=4)
             self.logger.info(f"Wrote result to {result_output_path}")
 
+        opcode_count_results = count_results()
+        if self.options.opcode_count == "stdout":
+            json_output["opcodeCount"] = opcode_count_results
+        elif self.options.opcode_count is not None:
+            result_output_path = os.path.join(
+                self.options.output_basedir,
+                self.options.opcode_count,
+            )
+            with open(result_output_path, "w") as f:
+                json.dump(opcode_count_results, f, indent=4)
+            self.logger.info(f"Wrote opcode counts to {result_output_path}")
+
         if json_output:
             json.dump(json_output, self.out_file, indent=4)
 

src/ethereum_spec_tools/evm_tools/t8n/evm_trace/count.py

@@ -0,0 +1,46 @@
+"""
+EVM trace implementation that counts how many times each opcode is executed.
+"""
+from collections import defaultdict
+from typing import Any, TypeAlias
+
+from ethereum.trace import OpStart, TraceEvent
+
+from .protocols import Evm
+
+_ActiveTraces: TypeAlias = tuple[object, dict[str, int]]
+_active_traces: _ActiveTraces | None = None
+
+
+def evm_trace(
+    evm: Any,
+    event: TraceEvent,
+) -> None:
+    """
+    Create a trace of the event.
+    """
+    global _active_traces
+
+    if not isinstance(event, OpStart):
+        return
+
+    assert isinstance(evm, Evm)
+
+    if _active_traces and _active_traces[0] is evm.message.tx_env:
+        traces = _active_traces[1]
+    else:
+        traces = defaultdict(lambda: 0)
+        _active_traces = (evm.message.tx_env, traces)
+
+    traces[event.op.name] += 1
+
+
+def results() -> dict[str, int]:
+    """
+    Take and clear the current opcode counts.
+    """
+    global _active_traces
+
+    results = _active_traces
+    _active_traces = None
+    return {} if results is None else results[1]

src/ethereum_spec_tools/evm_tools/t8n/evm_trace/group.py

@@ -0,0 +1,38 @@
+"""
+EVM trace implementation that fans out to many concrete trace implementations.
+"""
+from typing import Final
+
+from typing_extensions import override
+
+from ethereum.trace import EvmTracer, TraceEvent
+
+
+class GroupTracer(EvmTracer):
+    """
+    EVM trace implementation that fans out to many concrete trace
+    implementations.
+    """
+
+    tracers: Final[set[EvmTracer]]
+
+    def __init__(self) -> None:
+        self.tracers = set()
+
+    def add(self, tracer: EvmTracer) -> None:
+        """
+        Insert a new tracer.
+        """
+        self.tracers.add(tracer)
+
+    @override
+    def __call__(
+        self,
+        evm: object,
+        event: TraceEvent,
+    ) -> None:
+        """
+        Record a trace event.
+        """
+        for tracer in self.tracers:
+            tracer(evm, event)

tests/conftest.py

@@ -2,15 +2,15 @@
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Final, Optional, Set
+from typing import Callable, Final, Optional, Set
 
 import git
 import requests_cache
 from _pytest.config import Config
 from _pytest.config.argparsing import Parser
 from filelock import FileLock
 from git.exc import GitCommandError, InvalidGitRepositoryError
-from pytest import Session, StashKey
+from pytest import Session, StashKey, fixture
 from requests_cache import CachedSession
 from requests_cache.backends.sqlite import SQLiteCache
 from typing_extensions import Self
@@ -26,6 +26,14 @@ def get_xdist_worker_id(request_or_session: object) -> str:  # noqa: U100
         return "master"
 
 
+@fixture()
+def root_relative() -> Callable[[str | Path], Path]:
+    def _(path: str | Path) -> Path:
+        return Path(__file__).parent / path
+
+    return _
+
+
 def pytest_addoption(parser: Parser) -> None:
     """
     Accept --evm-trace option in pytest.

tests/evm_tools/test_count_opcodes.py

@@ -0,0 +1,46 @@
+import json
+from io import StringIO
+from pathlib import Path
+from typing import Callable
+
+from ethereum_spec_tools.evm_tools import create_parser
+from ethereum_spec_tools.evm_tools.t8n import T8N
+
+parser = create_parser()
+
+
+def test_count_opcodes(root_relative: Callable[[str | Path], Path]) -> None:
+    base_path = root_relative(
+        "fixtures/evm_tools_testdata/t8n/fixtures/testdata/2"
+    )
+
+    options = parser.parse_args(
+        [
+            "t8n",
+            f"--input.env={base_path / 'env.json'}",
+            f"--input.alloc={base_path / 'alloc.json'}",
+            f"--input.txs={base_path / 'txs.json'}",
+            "--output.result=stdout",
+            "--output.body=stdout",
+            "--output.alloc=stdout",
+            "--opcode.count=stdout",
+            "--state-test",
+        ]
+    )
+
+    in_file = StringIO()
+    out_file = StringIO()
+
+    t8n_tool = T8N(options, out_file=out_file, in_file=in_file)
+    exit_code = t8n_tool.run()
+    assert 0 == exit_code
+
+    results = json.loads(out_file.getvalue())
+
+    assert results["opcodeCount"] == {
+        "PUSH1": 5,
+        "MSTORE8": 1,
+        "CREATE": 1,
+        "ADD": 1,
+        "SELFDESTRUCT": 1,
+    }