tests(benchmark): add missing scenario and optimization

packages/testing/src/execution_testing/benchmark/benchmark_code_generator.py

@@ -16,6 +16,8 @@
 class JumpLoopGenerator(BenchmarkCodeGenerator):
     """Generates bytecode that loops execution using JUMP operations."""
 
+    contract_balance: int = 0
+
     def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         """Deploy the looping contract."""
         # Benchmark Test Structure:
@@ -28,7 +30,9 @@ def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
             cleanup=self.cleanup,
             fork=fork,
         )
-        self._contract_address = pre.deploy_contract(code=code)
+        self._contract_address = pre.deploy_contract(
+            code=code, balance=self.contract_balance
+        )
         return self._contract_address
 
 
@@ -49,20 +53,39 @@ def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         #    but not loop (e.g. PUSH)
         # 2. The loop contract that calls the target contract in a loop
 
-        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
+        attack_block_stack_delta = (
+            self.attack_block.pushed_stack_items
+            - self.attack_block.popped_stack_items
+        )
+        assert attack_block_stack_delta >= 0, (
+            "attack block stack delta must be non-negative"
+        )
+
+        setup_stack_delta = (
+            self.setup.pushed_stack_items - self.setup.popped_stack_items
+        )
+        assert setup_stack_delta >= 0, "setup stack delta must be non-negative"
 
         max_iterations = fork.max_code_size() // len(self.attack_block)
+        max_stack_height = fork.max_stack_height() - setup_stack_delta
 
-        if stack_delta > 0:
+        if attack_block_stack_delta > 0:
             max_iterations = min(
-                fork.max_stack_height() // stack_delta, max_iterations
+                max_stack_height // attack_block_stack_delta, max_iterations
             )
 
+        code = self.setup + self.attack_block * max_iterations
+        # Pad the code to the maximum code size.
+        if self.code_padding_opcode is not None:
+            code += self.code_padding_opcode * (
+                fork.max_code_size() - len(code)
+            )
+
+        self._validate_code_size(code, fork)
+
         # 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=code,
             balance=self.contract_balance,
         )
 
@@ -74,11 +97,22 @@ def deploy_contracts(self, *, pre: Alloc, fork: Fork) -> Address:
         # setup + JUMPDEST + attack + attack + ... + attack +
         # JUMP(setup_length)
         code_sequence = Op.POP(
-            Op.STATICCALL(Op.GAS, self._target_contract_address, 0, 0, 0, 0)
+            Op.STATICCALL(
+                Op.GAS,
+                self._target_contract_address,
+                Op.PUSH0,
+                Op.CALLDATASIZE,
+                Op.PUSH0,
+                Op.PUSH0,
+            )
         )
 
         caller_code = self.generate_repeated_code(
-            repeated_code=code_sequence, cleanup=self.cleanup, fork=fork
+            setup=Op.CALLDATACOPY(Op.PUSH0, Op.PUSH0, Op.CALLDATASIZE),
+            repeated_code=code_sequence,
+            cleanup=self.cleanup,
+            fork=fork,
         )
+
         self._contract_address = pre.deploy_contract(code=caller_code)
         return self._contract_address

packages/testing/src/execution_testing/specs/benchmark.py

@@ -53,6 +53,7 @@ class BenchmarkCodeGenerator(ABC):
     setup: Bytecode = field(default_factory=Bytecode)
     cleanup: Bytecode = field(default_factory=Bytecode)
     tx_kwargs: Dict[str, Any] = field(default_factory=dict)
+    code_padding_opcode: Op | None = None
     _contract_address: Address | None = None
 
     @abstractmethod
@@ -104,6 +105,9 @@ def generate_repeated_code(
         # 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
+        # Pad the code to the maximum code size.
+        if self.code_padding_opcode is not None:
+            code += self.code_padding_opcode * (max_code_size - len(code))
         self._validate_code_size(code, fork)
 
         return code

tests/benchmark/compute/instruction/test_account_query.py

@@ -4,6 +4,7 @@
 """
 
 import math
+from typing import Any
 
 import pytest
 from execution_testing import (
@@ -72,7 +73,6 @@ def test_codesize(
 )
 def test_codecopy(
     benchmark_test: BenchmarkTestFiller,
-    pre: Alloc,
     fork: Fork,
     max_code_size_ratio: float,
     fixed_src_dst: bool,
@@ -86,26 +86,14 @@ def test_codecopy(
     src_dst = 0 if fixed_src_dst else Op.MOD(Op.GAS, 7)
     attack_block = Op.CODECOPY(src_dst, src_dst, Op.DUP1)  # DUP1 copies size.
 
-    code = JumpLoopGenerator(
-        setup=setup, attack_block=attack_block
-    ).generate_repeated_code(
-        repeated_code=attack_block, setup=setup, fork=fork
-    )
-
-    # Pad the generated code to ensure the contract size matches the maximum
-    # The content of the padding bytes is arbitrary.
-    code += Op.INVALID * (max_code_size - len(code))
-    assert len(code) == max_code_size, (
-        f"Code size {len(code)} is not equal to max code size {max_code_size}."
-    )
-
-    tx = Transaction(
-        to=pre.deploy_contract(code=code),
-        sender=pre.fund_eoa(),
+    benchmark_test(
+        code_generator=JumpLoopGenerator(
+            setup=setup,
+            attack_block=attack_block,
+            code_padding_opcode=Op.STOP,
+        )
     )
 
-    benchmark_test(tx=tx)
-
 
 @pytest.mark.parametrize(
     "opcode",
@@ -361,7 +349,21 @@ def test_extcodecopy_warm(
     ],
 )
 @pytest.mark.parametrize(
-    "absent_target",
+    "empty_code",
+    [
+        True,
+        False,
+    ],
+)
+@pytest.mark.parametrize(
+    "initial_balance",
+    [
+        True,
+        False,
+    ],
+)
+@pytest.mark.parametrize(
+    "initial_storage",
     [
         True,
         False,
@@ -371,27 +373,39 @@ def test_ext_account_query_warm(
     benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
     opcode: Op,
-    absent_target: bool,
+    empty_code: bool,
+    initial_balance: bool,
+    initial_storage: bool,
 ) -> None:
     """
     Test running a block with as many stateful opcodes doing warm access
     for an account.
     """
     # Setup
-    target_addr = pre.empty_account()
     post = {}
-    if not absent_target:
-        code = Op.STOP + Op.JUMPDEST * 100
-        target_addr = pre.deploy_contract(balance=100, code=code)
-        post[target_addr] = Account(balance=100, code=code)
 
-    # Execution
-    setup = Op.MSTORE(0, target_addr)
-    attack_block = Op.POP(opcode(address=Op.MLOAD(0)))
+    if not initial_balance and not initial_storage and empty_code:
+        target_addr = pre.empty_account()
+    else:
+        kwargs: dict[str, Any] = {}
+        if initial_balance:
+            kwargs["balance"] = 100
+        if initial_storage:
+            kwargs["storage"] = {0: 0x1337}
+
+        if empty_code:
+            target_addr = pre.fund_eoa(**kwargs)
+        else:
+            code = Op.STOP + Op.JUMPDEST * 100
+            kwargs["code"] = code
+            target_addr = pre.deploy_contract(**kwargs)
+            post[target_addr] = Account(**kwargs)
+
     benchmark_test(
         post=post,
         code_generator=JumpLoopGenerator(
-            setup=setup, attack_block=attack_block
+            setup=Op.MSTORE(0, target_addr),
+            attack_block=Op.POP(opcode(address=Op.MLOAD(0))),
         ),
     )
 

tests/benchmark/compute/instruction/test_block_context.py

@@ -36,14 +36,29 @@ def test_block_context_ops(
     )
 
 
+@pytest.mark.parametrize(
+    "index",
+    [
+        0,
+        1,
+        256,
+        257,
+        pytest.param(None, id="random"),
+    ],
+)
 def test_blockhash(
     benchmark_test: BenchmarkTestFiller,
+    index: int | None,
 ) -> None:
     """Benchmark BLOCKHASH instruction accessing oldest allowed block."""
     # Create 256 dummy blocks to fill the blockhash window.
     blocks = [Block()] * 256
 
+    block_number = Op.AND(Op.GAS, 0xFF) if index is None else index
+
     benchmark_test(
         setup_blocks=blocks,
-        code_generator=ExtCallGenerator(attack_block=Op.BLOCKHASH(1)),
+        code_generator=ExtCallGenerator(
+            attack_block=Op.BLOCKHASH(block_number)
+        ),
     )

tests/benchmark/compute/instruction/test_call_context.py

@@ -43,8 +43,8 @@ def test_calldatasize(
 ) -> None:
     """Benchmark CALLDATASIZE instruction."""
     benchmark_test(
-        code_generator=JumpLoopGenerator(
-            attack_block=Op.POP(Op.CALLDATASIZE),
+        code_generator=ExtCallGenerator(
+            attack_block=Op.CALLDATASIZE,
             tx_kwargs={"data": b"\x00" * calldata_length},
         ),
     )
@@ -149,7 +149,7 @@ def test_calldatacopy(
     size: int,
     fixed_src_dst: bool,
     non_zero_data: bool,
-    gas_benchmark_value: int,
+    tx_gas_limit: int,
 ) -> None:
     """Benchmark CALLDATACOPY instruction."""
     if size == 0 and non_zero_data:
@@ -162,7 +162,7 @@ def test_calldatacopy(
 
     intrinsic_gas_calculator = fork.transaction_intrinsic_cost_calculator()
     min_gas = intrinsic_gas_calculator(calldata=data)
-    if min_gas > gas_benchmark_value:
+    if min_gas > tx_gas_limit:
         pytest.skip(
             "Minimum gas required for calldata ({min_gas}) is greater "
             "than the gas limit"
@@ -210,7 +210,6 @@ def test_calldatacopy(
 
     tx = Transaction(
         to=tx_target,
-        gas_limit=gas_benchmark_value,
         data=data,
         sender=pre.fund_eoa(),
     )
@@ -311,11 +310,6 @@ def test_returndatacopy(
     )
     dst = 0 if fixed_dst else Op.MOD(Op.GAS, 7)
 
-    # We create the contract that will be doing the RETURNDATACOPY multiple
-    # times.
-    returndata_gen = (
-        Op.STATICCALL(address=helper_contract) if size > 0 else Bytecode()
-    )
     attack_block = Op.RETURNDATACOPY(dst, Op.PUSH0, Op.RETURNDATASIZE)
 
     benchmark_test(

tests/benchmark/compute/instruction/test_control_flow.py

@@ -22,6 +22,15 @@ def test_gas_op(
     )
 
 
+def test_pc_op(
+    benchmark_test: BenchmarkTestFiller,
+) -> None:
+    """Benchmark PC instruction."""
+    benchmark_test(
+        code_generator=ExtCallGenerator(attack_block=Op.PC),
+    )
+
+
 def test_jumps(
     benchmark_test: BenchmarkTestFiller,
     pre: Alloc,

tests/benchmark/compute/instruction/test_keccak.py

@@ -2,6 +2,7 @@
 
 import math
 
+import pytest
 from execution_testing import (
     BenchmarkTestFiller,
     Fork,
@@ -15,15 +16,15 @@
 KECCAK_RATE = 136
 
 
-def test_keccak(
+def test_keccak_max_permutations(
     benchmark_test: BenchmarkTestFiller,
     fork: Fork,
-    gas_benchmark_value: int,
+    tx_gas_limit: int,
 ) -> None:
-    """Benchmark KECCAK256 instruction."""
+    """Benchmark KECCAK256 instruction to maximize permutations per block."""
     # Intrinsic gas cost is paid once.
     intrinsic_gas_calculator = fork.transaction_intrinsic_cost_calculator()
-    available_gas = gas_benchmark_value - intrinsic_gas_calculator()
+    available_gas = tx_gas_limit - intrinsic_gas_calculator()
 
     gsc = fork.gas_costs()
     mem_exp_gas_calculator = fork.memory_expansion_gas_calculator()
@@ -64,3 +65,20 @@ def test_keccak(
             attack_block=Op.POP(Op.SHA3(Op.PUSH0, Op.DUP1)),
         ),
     )
+
+
+@pytest.mark.parametrize("mem_alloc", [b"", b"ff", b"ff" * 32])
+@pytest.mark.parametrize("offset", [0, 31, 1024])
+def test_keccak(
+    benchmark_test: BenchmarkTestFiller,
+    offset: int,
+    mem_alloc: bytes,
+) -> None:
+    """Benchmark KECCAK256 instruction with diff input data and offsets."""
+    benchmark_test(
+        code_generator=JumpLoopGenerator(
+            setup=Op.CALLDATACOPY(offset, Op.PUSH0, Op.CALLDATASIZE),
+            attack_block=Op.POP(Op.SHA3(offset, Op.CALLDATASIZE)),
+            tx_kwargs={"data": mem_alloc},
+        ),
+    )