refactor(tests/benchmark): Optimize generators usages

Author: marioevz

src/ethereum_test_benchmark/benchmark_code_generator.py

@@ -3,46 +3,42 @@
 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) -> Address:
+    def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         """Deploy the looping contract."""
         # Benchmark Test Structure:
         # setup + JUMPDEST +
         # attack + attack + ... + attack +
         # cleanup + JUMP(setup_length)
-        code = self.generate_repeated_code(self.attack_block, self.setup, self.cleanup, fork)
+        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
 
-    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"):
-            self.deploy_contracts(pre, fork)
-
-        return Transaction(
-            to=self._contract_address,
-            gas_limit=gas_limit,
-            sender=pre.fund_eoa(),
-        )
-
 
+@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) -> Address:
+    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:
@@ -56,7 +52,8 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> Address:
 
         # Deploy target contract that contains the actual attack block
         self._target_contract_address = pre.deploy_contract(
-            code=self.setup + 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
@@ -68,17 +65,8 @@ def deploy_contracts(self, pre: Alloc, fork: Fork) -> Address:
         # 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.cleanup, fork)
+        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
-
-    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"):
-            self.deploy_contracts(pre, fork)
-
-        return Transaction(
-            to=self._contract_address,
-            gas_limit=gas_limit,
-            sender=pre.fund_eoa(),
-        )

src/ethereum_test_specs/benchmark.py

@@ -3,7 +3,7 @@
 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
@@ -41,27 +41,44 @@ 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) -> Address:
+    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, fork: Fork) -> 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, cleanup: 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()
-
+        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)
@@ -87,7 +104,7 @@ 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
     blocks: List[Block] | None = None
@@ -118,6 +135,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 +203,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, fork)
+        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)

tests/benchmark/test_worst_blocks.py

@@ -252,8 +252,6 @@ def test_block_full_data(
         )
 
     benchmark_test(
-        pre=pre,
-        post={},
         blocks=[Block(txs=txs)],
         expected_benchmark_gas_used=total_gas_used,
     )
@@ -361,8 +359,6 @@ def test_block_full_access_list_and_data(
         )
 
     benchmark_test(
-        pre=pre,
-        post={},
         blocks=[Block(txs=txs)],
         expected_benchmark_gas_used=total_gas_used,
     )

tests/benchmark/test_worst_bytecode.py

@@ -15,7 +15,6 @@
     Block,
     BlockchainTestFiller,
     Bytecode,
-    Bytes,
     Environment,
     Hash,
     Transaction,
@@ -246,10 +245,8 @@ def test_worst_bytecode_single_opcode(
 )
 def test_worst_initcode_jumpdest_analysis(
     benchmark_test: BenchmarkTestFiller,
-    pre: Alloc,
     fork: Fork,
     pattern: Bytecode,
-    gas_benchmark_value: int,
 ) -> None:
     """
     Test the jumpdest analysis performance of the initcode.
@@ -298,15 +295,12 @@ def test_worst_initcode_jumpdest_analysis(
 
     setup = code_prepare_initcode + Op.PUSH0
 
-    tx = JumpLoopGenerator(
-        setup=setup, attack_block=attack_block, cleanup=Bytecode()
-    ).generate_transaction(pre, gas_benchmark_value, fork)
-    tx.data = Bytes(tx_data)
-
     benchmark_test(
-        pre=pre,
-        post={},
-        tx=tx,
+        code_generator=JumpLoopGenerator(
+            setup=setup,
+            attack_block=attack_block,
+            tx_kwargs={"data": tx_data},
+        ),
     )
 
 
@@ -408,7 +402,7 @@ def test_worst_create(
     )
 
     code = JumpLoopGenerator(setup=setup, attack_block=attack_block).generate_repeated_code(
-        repeated_code=attack_block, setup=setup, cleanup=Bytecode(), fork=fork
+        repeated_code=attack_block, setup=setup, fork=fork
     )
 
     tx = Transaction(
@@ -417,11 +411,7 @@ def test_worst_create(
         sender=pre.fund_eoa(),
     )
 
-    benchmark_test(
-        pre=pre,
-        post={},
-        tx=tx,
-    )
+    benchmark_test(tx=tx)
 
 
 @pytest.mark.parametrize(
@@ -483,7 +473,5 @@ def test_worst_creates_collisions(
             pre.deploy_contract(address=addr, code=Op.INVALID)
 
     benchmark_test(
-        pre=pre,
-        post={},
         code_generator=JumpLoopGenerator(setup=setup, attack_block=attack_block),
     )