refactor(benchmark): update to benchmark test wrapper

src/ethereum_test_benchmark/benchmark_code_generator.py

@@ -3,57 +3,62 @@
 optimized bytecode patterns.
 """
 
+from dataclasses import dataclass
+
+from ethereum_test_base_types import Address
 from ethereum_test_forks import Fork
 from ethereum_test_specs.benchmark import BenchmarkCodeGenerator
-from ethereum_test_types import Alloc, Transaction
-from ethereum_test_vm import Bytecode
+from ethereum_test_types import Alloc
 from ethereum_test_vm.opcodes import Opcodes as Op
 
 
+@dataclass(kw_only=True)
 class JumpLoopGenerator(BenchmarkCodeGenerator):
     """Generates bytecode that loops execution using JUMP operations."""
 
-    def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
+    def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         """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, fork)
-        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=self._contract_address,
-            gas_limit=gas_limit,
-            sender=pre.fund_eoa(),
+        # setup + JUMPDEST +
+        # attack + attack + ... + attack +
+        # cleanup + JUMP(setup_length)
+        code = self.generate_repeated_code(
+            repeated_code=self.attack_block, setup=self.setup, cleanup=self.cleanup, fork=fork
         )
+        self._contract_address = pre.deploy_contract(code=code)
+        return self._contract_address
 
 
+@dataclass(kw_only=True)
 class ExtCallGenerator(BenchmarkCodeGenerator):
     """
     Generates bytecode that fills the contract to
     maximum allowed code size.
     """
 
-    def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
+    contract_balance: int = 0
+
+    def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         """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
 
-        max_iterations = min(
-            fork.max_stack_height(), fork.max_code_size() // len(self.attack_block)
-        )
+        pushed_stack_items = self.attack_block.pushed_stack_items
+        popped_stack_items = self.attack_block.popped_stack_items
+        stack_delta = pushed_stack_items - popped_stack_items
+
+        max_iterations = fork.max_code_size() // len(self.attack_block)
+
+        if stack_delta > 0:
+            max_iterations = min(fork.max_stack_height() // stack_delta, max_iterations)
 
         # Deploy target contract that contains the actual attack block
         self._target_contract_address = pre.deploy_contract(
-            code=self.attack_block * max_iterations
+            code=self.setup + self.attack_block * max_iterations,
+            balance=self.contract_balance,
         )
 
         # Create caller contract that repeatedly calls the target contract
@@ -65,16 +70,8 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
         # 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(), fork)
-        self._contract_address = pre.deploy_contract(code=caller_code)
-
-    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=self._contract_address,
-            gas_limit=gas_limit,
-            sender=pre.fund_eoa(),
+        caller_code = self.generate_repeated_code(
+            repeated_code=code_sequence, cleanup=self.cleanup, fork=fork
         )
+        self._contract_address = pre.deploy_contract(code=caller_code)
+        return self._contract_address

src/ethereum_test_specs/benchmark.py

@@ -3,13 +3,13 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Callable, ClassVar, Dict, Generator, List, Sequence, Type
+from typing import Any, Callable, ClassVar, Dict, Generator, List, Sequence, Type
 
 import pytest
 from pydantic import ConfigDict, Field
 
 from ethereum_clis import TransitionTool
-from ethereum_test_base_types import HexNumber
+from ethereum_test_base_types import Address, HexNumber
 from ethereum_test_exceptions import BlockException, TransactionException
 from ethereum_test_execution import (
     BaseExecute,
@@ -40,32 +40,52 @@ class BenchmarkCodeGenerator(ABC):
 
     attack_block: Bytecode
     setup: Bytecode = field(default_factory=Bytecode)
+    cleanup: Bytecode = field(default_factory=Bytecode)
+    tx_kwargs: Dict[str, Any] = field(default_factory=dict)
+    _contract_address: Address | None = None
 
     @abstractmethod
-    def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
+    def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         """Deploy any contracts needed for the benchmark."""
         ...
 
-    @abstractmethod
-    def generate_transaction(self, pre: Alloc, gas_limit: int) -> Transaction:
-        """Generate a transaction with the specified gas limit."""
-        ...
+    def generate_transaction(self, *, pre: Alloc, gas_benchmark_value: int) -> Transaction:
+        """Generate transaction that executes the looping contract."""
+        assert self._contract_address is not None
+        if "gas_limit" not in self.tx_kwargs:
+            self.tx_kwargs["gas_limit"] = gas_benchmark_value
+
+        return Transaction(
+            to=self._contract_address,
+            sender=pre.fund_eoa(),
+            **self.tx_kwargs,
+        )
 
     def generate_repeated_code(
-        self, repeated_code: Bytecode, setup: Bytecode, fork: Fork
+        self,
+        *,
+        repeated_code: Bytecode,
+        setup: Bytecode | None = None,
+        cleanup: Bytecode | None = None,
+        fork: Fork,
     ) -> 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 = fork.max_code_size()
-
-        overhead = len(setup) + len(Op.JUMPDEST) + len(Op.JUMP(len(setup)))
+        if setup is None:
+            setup = Bytecode()
+        if cleanup is None:
+            cleanup = Bytecode()
+        overhead = len(setup) + len(Op.JUMPDEST) + len(cleanup) + len(Op.JUMP(len(setup)))
         available_space = max_code_size - overhead
         max_iterations = available_space // len(repeated_code)
 
-        code = setup + Op.JUMPDEST + repeated_code * max_iterations + Op.JUMP(len(setup))
+        # TODO: Unify the PUSH0 and PUSH1 usage.
+        code = setup + Op.JUMPDEST + repeated_code * max_iterations + cleanup
+        code += Op.JUMP(len(setup)) if len(setup) > 0 else Op.PUSH0 + Op.JUMP
         self._validate_code_size(code, fork)
 
         return code
@@ -84,9 +104,10 @@ class BenchmarkTest(BaseTest):
 
     model_config = ConfigDict(extra="forbid")
 
-    pre: Alloc
+    pre: Alloc = Field(default_factory=Alloc)
     post: Alloc = Field(default_factory=Alloc)
     tx: Transaction | None = None
+    setup_blocks: List[Block] = Field(default_factory=list)
     blocks: List[Block] | None = None
     block_exception: (
         List[TransactionException | BlockException] | TransactionException | BlockException | None
@@ -115,6 +136,14 @@ class BenchmarkTest(BaseTest):
         "blockchain_test_only": "Only generate a blockchain test fixture",
     }
 
+    def model_post_init(self, __context: Any, /) -> None:
+        """
+        Model post-init to assert that the custom pre-allocation was
+        provided and the default was not used.
+        """
+        super().model_post_init(__context)
+        assert "pre" in self.model_fields_set, "pre allocation was not provided"
+
     @classmethod
     def pytest_parameter_name(cls) -> str:
         """
@@ -178,9 +207,11 @@ def generate_blocks_from_code_generator(self, fork: Fork) -> List[Block]:
         if self.code_generator is None:
             raise Exception("Code generator is not set")
 
-        self.code_generator.deploy_contracts(self.pre, fork)
+        self.code_generator.deploy_contracts(pre=self.pre, fork=fork)
         gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
-        benchmark_tx = self.code_generator.generate_transaction(self.pre, gas_limit)
+        benchmark_tx = self.code_generator.generate_transaction(
+            pre=self.pre, gas_benchmark_value=gas_limit
+        )
 
         execution_txs = self.split_transaction(benchmark_tx, gas_limit)
         execution_block = Block(txs=execution_txs)
@@ -204,39 +235,34 @@ def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
                 f"Exactly one must be set, but got {len(set_props)}: {', '.join(set_props)}"
             )
 
+        blocks: List[Block] = self.setup_blocks
+
         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,
-            )
+            blocks += generated_blocks
+
         elif self.blocks is not None:
-            return BlockchainTest.from_test(
-                base_test=self,
-                genesis_environment=self.env,
-                pre=self.pre,
-                post=self.post,
-                blocks=self.blocks,
-            )
+            blocks += self.blocks
+
         elif self.tx is not None:
             gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
 
             transactions = self.split_transaction(self.tx, gas_limit)
 
-            blocks = [Block(txs=transactions)]
+            blocks.append(Block(txs=transactions))
 
-            return BlockchainTest.from_test(
-                base_test=self,
-                pre=self.pre,
-                post=self.post,
-                blocks=blocks,
-                genesis_environment=self.env,
-            )
         else:
-            raise ValueError("Cannot create BlockchainTest without transactions or blocks")
+            raise ValueError(
+                "Cannot create BlockchainTest without a code generator, transactions, or blocks"
+            )
+
+        return BlockchainTest.from_test(
+            base_test=self,
+            genesis_environment=self.env,
+            pre=self.pre,
+            post=self.post,
+            blocks=blocks,
+        )
 
     def generate(
         self,