WiP: Add flag to count opcodes during execution

Author: SamWilsn

src/ethereum_spec_tools/evm_tools/t8n/__init__.py

@@ -7,7 +7,7 @@
 import json
 import os
 from functools import partial
-from typing import Any, TextIO
+from typing import Any, Final, TextIO
 
 from ethereum_rlp import rlp
 from ethereum_types.numeric import U64, U256, Uint
@@ -25,7 +25,10 @@
     parse_hex_or_int,
 )
 from .env import Env
-from .evm_trace.eip3155 import evm_trace
+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 evm_trace as evm_trace_eip3155
+from .evm_trace.group import GroupTracer
 from .t8n_types import Alloc, Result, Txs
 
 
@@ -73,12 +76,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:
@@ -101,19 +108,34 @@ 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(
                 partial(
-                    evm_trace,
+                    evm_trace_eip3155,
                     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__(
@@ -349,6 +371,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["opcodes"] = 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,36 @@
+"""
+EVM trace implementation that fans out to many concrete trace implementations.
+"""
+from typing import Final, 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)