refactor(benchmark): update to benchmark test wrapper

src/ethereum_test_benchmark/benchmark_code_generator.py

@@ -3,6 +3,7 @@
 optimized bytecode patterns.
 """
 
+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
@@ -13,18 +14,20 @@
 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)
+        # setup + JUMPDEST +
+        # attack + attack + ... + attack +
+        # cleanup + JUMP(setup_length)
+        code = self.generate_repeated_code(self.attack_block, self.setup, self.cleanup, fork)
         self._contract_address = pre.deploy_contract(code=code)
+        return self._contract_address
 
     def generate_transaction(self, pre: Alloc, gas_limit: int, fork: Fork) -> Transaction:
         """Generate transaction that executes the looping contract."""
         if not hasattr(self, "_contract_address"):
-            raise ValueError("deploy_contracts must be called before generate_transaction")
+            self.deploy_contracts(pre, fork)
 
         return Transaction(
             to=self._contract_address,
@@ -39,7 +42,7 @@ class ExtCallGenerator(BenchmarkCodeGenerator):
     maximum allowed code size.
     """
 
-    def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
+    def deploy_contracts(self, pre: Alloc, fork: Fork) -> Address:
         """Deploy both target and caller contracts."""
         # Benchmark Test Structure:
         # There are two contracts:
@@ -53,7 +56,7 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
 
         # 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
         )
 
         # Create caller contract that repeatedly calls the target contract
@@ -65,13 +68,14 @@ 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)
+        caller_code = self.generate_repeated_code(code_sequence, Bytecode(), self.cleanup, fork)
         self._contract_address = pre.deploy_contract(code=caller_code)
+        return self._contract_address
 
     def generate_transaction(self, pre: Alloc, gas_limit: int, fork: Fork) -> Transaction:
         """Generate transaction that executes the caller contract."""
         if not hasattr(self, "_contract_address"):
-            raise ValueError("deploy_contracts must be called before generate_transaction")
+            self.deploy_contracts(pre, fork)
 
         return Transaction(
             to=self._contract_address,

src/ethereum_test_specs/benchmark.py

@@ -9,7 +9,7 @@
 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,9 +40,10 @@ class BenchmarkCodeGenerator(ABC):
 
     attack_block: Bytecode
     setup: Bytecode = field(default_factory=Bytecode)
+    cleanup: Bytecode = field(default_factory=Bytecode)
 
     @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."""
         ...
 
@@ -52,7 +53,7 @@ def generate_transaction(self, pre: Alloc, gas_limit: int, fork: Fork) -> Transa
         ...
 
     def generate_repeated_code(
-        self, repeated_code: Bytecode, setup: Bytecode, fork: Fork
+        self, repeated_code: Bytecode, setup: Bytecode, cleanup: Bytecode, fork: Fork
     ) -> Bytecode:
         """
         Calculate the maximum number of iterations that
@@ -61,11 +62,11 @@ def generate_repeated_code(
         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)))
+        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))
+        code = setup + Op.JUMPDEST + repeated_code * max_iterations + cleanup + Op.JUMP(len(setup))
         self._validate_code_size(code, fork)
 
         return code