refactor: update benchmark code generator and test wrapper

Author: LouisTsai-Csie

src/ethereum_test_benchmark/__init__.py

@@ -0,0 +1,13 @@
+"""Benchmark code generator classes for creating optimized bytecode patterns."""
+
+from .benchmark_code_generator import (
+    BenchmarkCodeGenerator,
+    ExtCallGenerator,
+    JumpLoopGenerator,
+)
+
+__all__ = (
+    "BenchmarkCodeGenerator",
+    "ExtCallGenerator",
+    "JumpLoopGenerator",
+)

src/ethereum_test_tools/benchmark_code_generator.py renamed to src/ethereum_test_benchmark/benchmark_code_generator.py

@@ -1,26 +1,26 @@
 """Benchmark code generator classes for creating optimized bytecode patterns."""
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from dataclasses import dataclass, field
 
 from ethereum_test_forks import Fork
-from ethereum_test_tools import Alloc, Bytecode, Transaction
-from ethereum_test_tools.vm.opcode import Opcodes as Op
+from ethereum_test_types import Alloc, Transaction
+from ethereum_test_vm import Bytecode
+from ethereum_test_vm.opcode import Opcodes as Op
 
 
+@dataclass
 class BenchmarkCodeGenerator(ABC):
     """Abstract base class for generating benchmark bytecode."""
 
-    def __init__(
-        self,
-        fork: Fork,
-        attack_block: Bytecode,
-        setup: Optional[Bytecode] = None,
-    ):
-        """Initialize with fork, attack block, and optional setup bytecode."""
-        self.fork = fork
-        self.setup = setup or Bytecode()
-        self.attack_block = attack_block
+    fork: Fork
+    attack_block: Bytecode
+    setup: Bytecode = field(default_factory=Bytecode)
+
+    @abstractmethod
+    def deploy_contracts(self, pre: Alloc) -> None:
+        """Deploy any contracts needed for the benchmark."""
+        pass
 
     @abstractmethod
     def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
@@ -29,14 +29,14 @@ def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
 
     def generate_repeated_code(self, repeated_code: Bytecode, setup: Bytecode) -> Bytecode:
         """Calculate the maximum number of iterations that can fit in the code size limit."""
+        assert len(repeated_code) > 0, "repeated_code cannot be empty"
         max_code_size = self.fork.max_code_size()
 
-        overhead = len(Op.JUMPDEST) + len(Op.JUMP(len(setup)))
+        overhead = len(setup) + len(Op.JUMPDEST) + len(Op.JUMP(len(setup)))
         available_space = max_code_size - overhead
-        max_iterations = available_space // len(repeated_code) if len(repeated_code) > 0 else 0
+        max_iterations = available_space // len(repeated_code)
 
         code = setup + Op.JUMPDEST + repeated_code * max_iterations + Op.JUMP(len(setup))
-
         self._validate_code_size(code)
 
         return code
@@ -50,47 +50,62 @@ def _validate_code_size(self, code: Bytecode) -> None:
             )
 
 
+@dataclass
 class JumpLoopGenerator(BenchmarkCodeGenerator):
     """Generates bytecode that loops execution using JUMP operations."""
 
-    def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
-        """Generate transaction with looping bytecode pattern."""
+    def deploy_contracts(self, pre: Alloc) -> None:
+        """Deploy the looping contract."""
         # Benchmark Test Structure:
         # setup + JUMPDEST + attack + attack + ... + attack + JUMP(setup_length)
-
         code = self.generate_repeated_code(self.attack_block, self.setup)
+        self._contract_address = pre.deploy_contract(code=code)
+
+    def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
+        """Generate transaction that executes the looping contract."""
+        if not hasattr(self, "_contract_address"):
+            raise ValueError("deploy_contracts must be called before generate_transaction")
 
         return Transaction(
-            to=pre.deploy_contract(code=code),
-            gas_limit=self.fork.transaction_gas_limit_cap() or 30_000_000,
+            to=self._contract_address,
+            gas_limit=gas_limit,
             sender=pre.fund_eoa(),
         )
 
 
+@dataclass
 class ExtCallGenerator(BenchmarkCodeGenerator):
     """Generates bytecode that fills the contract to maximum allowed code size."""
 
-    def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
-        """Generate transaction with maximal code size coverage."""
+    def deploy_contracts(self, pre: Alloc) -> None:
+        """Deploy both target and caller contracts."""
         # Benchmark Test Structure:
         # There are two contracts:
         # 1. The target contract that executes certain operation but not loop (e.g. PUSH)
         # 2. The loop contract that calls the target contract in a loop
-        #
-        # attack = POP(STATICCALL(GAS, target_contract_address, 0, 0, 0, 0))
-        # setup + JUMPDEST + attack + attack + ... + attack + JUMP(setup_lengt)
-        # This could optimize the gas consumption and increase the cycle count.
 
         max_stack_height = self.fork.max_stack_height()
 
-        target_contract_address = pre.deploy_contract(code=self.attack_block * max_stack_height)
+        # Deploy target contract that contains the actual attack block
+        self._target_contract_address = pre.deploy_contract(
+            code=self.attack_block * max_stack_height
+        )
 
-        code_sequence = Op.POP(Op.STATICCALL(Op.GAS, target_contract_address, 0, 0, 0, 0))
+        # Create caller contract that repeatedly calls the target contract
+        # attack = POP(STATICCALL(GAS, target_contract_address, 0, 0, 0, 0))
+        # setup + JUMPDEST + attack + attack + ... + attack + JUMP(setup_length)
+        code_sequence = Op.POP(Op.STATICCALL(Op.GAS, self._target_contract_address, 0, 0, 0, 0))
+
+        caller_code = self.generate_repeated_code(code_sequence, Bytecode())
+        self._contract_address = pre.deploy_contract(code=caller_code)
 
-        code = self.generate_repeated_code(code_sequence, Bytecode())
+    def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
+        """Generate transaction that executes the caller contract."""
+        if not hasattr(self, "_contract_address"):
+            raise ValueError("deploy_contracts must be called before generate_transaction")
 
         return Transaction(
-            to=pre.deploy_contract(code=code),
-            gas_limit=self.fork.transaction_gas_limit_cap() or 30_000_000,
+            to=self._contract_address,
+            gas_limit=gas_limit,
             sender=pre.fund_eoa(),
         )

src/ethereum_test_specs/benchmark.py

@@ -1,9 +1,12 @@
 """Ethereum benchmark test spec definition and filler."""
 
-from typing import Callable, ClassVar, Dict, Generator, List, Optional, Sequence, Type
+from contextlib import contextmanager
+from contextvars import ContextVar
+from enum import Enum
+from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Sequence, Type
 
 import pytest
-from pydantic import Field
+from pydantic import ConfigDict, Field
 
 from ethereum_clis import TransitionTool
 from ethereum_test_base_types import HexNumber
@@ -29,9 +32,74 @@
 from .blockchain import Block, BlockchainTest
 
 
+class BenchmarkPhase(Enum):
+    """Phases of a benchmark test."""
+
+    SETUP = "setup"
+    EXECUTION = "execution"
+
+
+_current_phase: ContextVar[Optional[BenchmarkPhase]] = ContextVar("benchmark_phase", default=None)
+
+
+class BenchmarkManager:
+    """Context manager for managing benchmark test phases."""
+
+    def __init__(self):
+        """Initialize the BenchmarkManager with empty transaction and block lists."""
+        self.setup_transactions: List[Transaction] = []
+        self.setup_blocks: List[Block] = []
+        self.execution_transactions: List[Transaction] = []
+        self.execution_blocks: List[Block] = []
+
+    @contextmanager
+    def setup(self):
+        """Context manager for the setup phase of a benchmark test."""
+        token = _current_phase.set(BenchmarkPhase.SETUP)
+        try:
+            yield self
+        finally:
+            _current_phase.reset(token)
+
+    @contextmanager
+    def execution(self):
+        """Context manager for the execution phase of a benchmark test."""
+        token = _current_phase.set(BenchmarkPhase.EXECUTION)
+        try:
+            yield self
+        finally:
+            _current_phase.reset(token)
+
+    def add_transaction(self, tx: Transaction):
+        """Add a transaction to the current phase."""
+        current_phase = _current_phase.get()
+        if current_phase == BenchmarkPhase.SETUP:
+            self.setup_transactions.append(tx)
+        elif current_phase == BenchmarkPhase.EXECUTION:
+            self.execution_transactions.append(tx)
+        else:
+            self.setup_transactions.append(tx)
+
+    def add_block(self, block: Block):
+        """Add a block to the current phase."""
+        current_phase = _current_phase.get()
+        if current_phase == BenchmarkPhase.SETUP:
+            self.setup_blocks.append(block)
+        elif current_phase == BenchmarkPhase.EXECUTION:
+            self.execution_blocks.append(block)
+        else:
+            self.setup_blocks.append(block)
+
+    def get_current_phase(self) -> Optional[BenchmarkPhase]:
+        """Get the current benchmark phase."""
+        return _current_phase.get()
+
+
 class BenchmarkTest(BaseTest):
     """Test type designed specifically for benchmark test cases."""
 
+    model_config = ConfigDict(extra="forbid")
+
     pre: Alloc
     post: Alloc
     tx: Optional[Transaction] = None
@@ -41,6 +109,9 @@ class BenchmarkTest(BaseTest):
     ) = None
     env: Environment = Field(default_factory=Environment)
     expected_benchmark_gas_used: int | None = None
+    gas_benchmark_value: int
+    benchmark_manager: Optional[Any] = Field(default=None, exclude=True)
+    code_generator: Optional[Any] = Field(default=None, exclude=True)
 
     supported_fixture_formats: ClassVar[Sequence[FixtureFormat | LabeledFixtureFormat]] = [
         BlockchainFixture,
@@ -86,26 +157,81 @@ def get_genesis_environment(self, fork: Fork) -> Environment:
 
     def split_transaction(self, tx: Transaction, gas_limit_cap: int | None) -> List[Transaction]:
         """Split a transaction that exceeds the gas limit cap into multiple transactions."""
-        if (gas_limit_cap is None) or (tx.gas_limit <= gas_limit_cap):
+        if gas_limit_cap is None:
+            tx.gas_limit = HexNumber(self.gas_benchmark_value)
+            return [tx]
+
+        if gas_limit_cap >= self.gas_benchmark_value:
+            tx.gas_limit = HexNumber(min(tx.gas_limit, self.gas_benchmark_value))
             return [tx]
 
-        total_gas = int(self.expected_benchmark_gas_used or self.env.gas_limit)
-        print(f"total_gas: {total_gas}")
-        num_splits = total_gas // gas_limit_cap
+        remaining_gas = self.gas_benchmark_value
+        num_splits = remaining_gas // gas_limit_cap + int(remaining_gas % gas_limit_cap)
 
         split_transactions = []
         for i in range(num_splits):
             split_tx = tx.model_copy()
-            total_gas -= gas_limit_cap
-            split_tx.gas_limit = HexNumber(total_gas if i == num_splits - 1 else gas_limit_cap)
+            split_tx.gas_limit = HexNumber(remaining_gas if i == num_splits - 1 else gas_limit_cap)
+            remaining_gas -= gas_limit_cap
             split_tx.nonce = HexNumber(tx.nonce + i)
             split_transactions.append(split_tx)
 
         return split_transactions
 
+    def generate_blocks_from_code_generator(self, fork: Fork) -> List[Block]:
+        """Generate blocks using the code generator."""
+        if self.code_generator is None:
+            return []
+
+        self.code_generator.deploy_contracts(self.pre)
+        gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
+        benchmark_tx = self.code_generator.generate_transaction(self.pre, gas_limit)
+
+        execution_txs = self.split_transaction(benchmark_tx, gas_limit)
+        execution_block = Block(txs=execution_txs)
+
+        return [execution_block]
+
     def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
         """Create a BlockchainTest from this BenchmarkTest."""
-        if self.blocks is not None:
+        if self.code_generator is not None:
+            generated_blocks = self.generate_blocks_from_code_generator(fork)
+            return BlockchainTest.from_test(
+                base_test=self,
+                genesis_environment=self.env,
+                pre=self.pre,
+                post=self.post,
+                blocks=generated_blocks,
+            )
+
+        elif self.benchmark_manager is not None:
+            all_blocks = []
+            gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
+
+            if self.benchmark_manager.setup_blocks:
+                all_blocks.extend(self.benchmark_manager.setup_blocks)
+            elif self.benchmark_manager.setup_transactions:
+                setup_txs = []
+                for tx in self.benchmark_manager.setup_transactions:
+                    setup_txs.extend(self.split_transaction(tx, gas_limit))
+                all_blocks.append(Block(txs=setup_txs))
+
+            if self.benchmark_manager.execution_blocks:
+                all_blocks.extend(self.benchmark_manager.execution_blocks)
+            elif self.benchmark_manager.execution_transactions:
+                execution_txs = []
+                for tx in self.benchmark_manager.execution_transactions:
+                    execution_txs.extend(self.split_transaction(tx, gas_limit))
+                all_blocks.append(Block(txs=execution_txs))
+
+            return BlockchainTest.from_test(
+                base_test=self,
+                genesis_environment=self.env,
+                pre=self.pre,
+                post=self.post,
+                blocks=all_blocks,
+            )
+        elif self.blocks is not None:
             return BlockchainTest.from_test(
                 base_test=self,
                 genesis_environment=self.env,
@@ -114,9 +240,9 @@ def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
                 blocks=self.blocks,
             )
         elif self.tx is not None:
-            gas_limit_cap = fork.transaction_gas_limit_cap()
+            gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
 
-            transactions = self.split_transaction(self.tx, gas_limit_cap)
+            transactions = self.split_transaction(self.tx, gas_limit)
 
             blocks = [Block(txs=transactions)]
 
@@ -129,7 +255,7 @@ def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
             )
         else:
             raise ValueError(
-                "Cannot create BlockchainTest without transactions, blocks, or code_generator"
+                "Cannot create BlockchainTest without transactions, blocks, or benchmark_manager"
             )
 
     def generate(
@@ -162,5 +288,10 @@ def execute(
         raise Exception(f"Unsupported execute format: {execute_format}")
 
 
+def create_benchmark_manager() -> BenchmarkManager:
+    """Create a new BenchmarkManager instance for phase-aware benchmark testing."""
+    return BenchmarkManager()
+
+
 BenchmarkTestSpec = Callable[[str], Generator[BenchmarkTest, None, None]]
 BenchmarkTestFiller = Type[BenchmarkTest]

src/ethereum_test_tools/__init__.py

@@ -16,6 +16,11 @@
     TestPrivateKey2,
 )
 from ethereum_test_base_types.reference_spec import ReferenceSpec, ReferenceSpecTypes
+from ethereum_test_benchmark import (
+    BenchmarkCodeGenerator,
+    ExtCallGenerator,
+    JumpLoopGenerator,
+)
 from ethereum_test_exceptions import (
     BlockException,
     EngineAPIError,
@@ -86,11 +91,6 @@
     call_return_code,
 )
 
-from .benchmark_code_generator import (
-    BenchmarkCodeGenerator,
-    ExtCallGenerator,
-    JumpLoopGenerator,
-)
 from .code import (
     CalldataCase,
     Case,