refactor(benchmark): add cleanup parameter to code generation logic

Author: LouisTsai-Csie

src/ethereum_test_benchmark/benchmark_code_generator.py

@@ -14,7 +14,7 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> 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, fork)
+        code = self.generate_repeated_code(self.attack_block, self.setup, self.cleanup, fork)
         self._contract_address = pre.deploy_contract(code=code)
 
     def generate_transaction(self, pre: Alloc, gas_limit: int, fork: Fork) -> Transaction:
@@ -53,7 +53,7 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> None:
         # 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(), fork)
+        caller_code = self.generate_repeated_code(code_sequence, Bytecode(), self.cleanup, fork)
         self._contract_address = pre.deploy_contract(code=caller_code)
 
     def generate_transaction(self, pre: Alloc, gas_limit: int, fork: Fork) -> Transaction:

src/ethereum_test_specs/benchmark.py

@@ -40,6 +40,7 @@ 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:
@@ -52,17 +53,17 @@ 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 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)))
+        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

tests/benchmark/test_worst_memory.py

@@ -92,7 +92,7 @@ def test_worst_calldatacopy(
         src_dst, src_dst, Op.CALLDATASIZE if non_zero_data or size == 0 else Op.DUP1
     )
     code = JumpLoopGenerator(setup=prefix, attack_block=attack_block).generate_repeated_code(
-        attack_block, Bytecode(), fork
+        attack_block, Bytecode(), Bytecode(), fork
     )
     code_address = pre.deploy_contract(code=code)
 
@@ -162,7 +162,7 @@ def test_worst_codecopy(
     attack_block = Op.CODECOPY(src_dst, src_dst, Op.DUP1)  # DUP1 copies size.
 
     code = JumpLoopGenerator(setup=code_prefix, attack_block=attack_block).generate_repeated_code(
-        attack_block, Bytecode(), fork
+        attack_block, Bytecode(), Bytecode(), fork
     )
 
     # The code generated above is not guaranteed to be of max_code_size, so we pad it since