refactor: optimize benchmark cases

Author: LouisTsai-Csie

tests/benchmark/compute/instruction/test_account_query.py

@@ -4,6 +4,7 @@
 """
 
 import math
+from typing import Dict
 
 import pytest
 from execution_testing import (
@@ -24,6 +25,7 @@
     While,
     compute_create2_address,
 )
+from execution_testing.base_types.conversions import NumberConvertible
 
 from tests.benchmark.compute.helpers import (
     XOR_TABLE,
@@ -72,7 +74,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 +87,12 @@ 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
+        )
     )
 
-    benchmark_test(tx=tx)
-
 
 @pytest.mark.parametrize(
     "opcode",
@@ -361,7 +348,21 @@ def test_extcodecopy_warm(
     ],
 )
 @pytest.mark.parametrize(
-    "absent_target",
+    "empty_account",
+    [
+        True,
+        False,
+    ],
+)
+@pytest.mark.parametrize(
+    "initial_balance",
+    [
+        True,
+        False,
+    ],
+)
+@pytest.mark.parametrize(
+    "initial_storage",
     [
         True,
         False,
@@ -371,27 +372,40 @@ def test_ext_account_query_warm(
     benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
     opcode: Op,
-    absent_target: bool,
+    empty_account: 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()
+    target_addr = pre.empty_account() if initial_balance else pre.fund_eoa()
     post = {}
-    if not absent_target:
+    if not empty_account:
         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)))
+        storage: Dict[NumberConvertible, NumberConvertible] = (
+            {0: 0x1337} if initial_storage else {0: 0}
+        )
+        target_addr = pre.deploy_contract(
+            balance=initial_balance,
+            code=code,
+            storage=storage,
+        )
+
+        post[target_addr] = Account(
+            balance=initial_balance,
+            code=code,
+            storage=storage,
+        )
+
     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,
+        None,
+    ],
+)
 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

@@ -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()
@@ -60,7 +61,24 @@ def test_keccak(
 
     benchmark_test(
         code_generator=JumpLoopGenerator(
-            setup=Op.PUSH20[optimal_input_length],
+            setup=Op.PUSH20(optimal_input_length),
             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},
+        ),
+    )

tests/benchmark/compute/instruction/test_memory.py

@@ -18,14 +18,10 @@ def test_msize(
     benchmark_test: BenchmarkTestFiller,
     mem_size: int,
 ) -> None:
-    """
-    Benchmark MSIZE instruction.
-
-    - mem_size: by how much the memory is expanded.
-    """
+    """Benchmark MSIZE instruction."""
     benchmark_test(
         code_generator=ExtCallGenerator(
-            setup=Op.MLOAD(Op.SELFBALANCE) + Op.POP,
+            setup=Op.POP(Op.MLOAD(Op.SELFBALANCE)),
             attack_block=Op.MSIZE,
             contract_balance=mem_size,
         ),
@@ -99,6 +95,6 @@ def test_mcopy(
     )
     benchmark_test(
         code_generator=JumpLoopGenerator(
-            setup=mem_touch, attack_block=attack_block, cleanup=mem_touch
+            attack_block=attack_block, cleanup=mem_touch
         ),
     )

tests/benchmark/compute/instruction/test_storage.py

@@ -7,6 +7,7 @@
     Block,
     Bytecode,
     Environment,
+    ExtCallGenerator,
     Fork,
     JumpLoopGenerator,
     Op,
@@ -33,36 +34,25 @@ def test_tload(
     val_mut: bool,
 ) -> None:
     """Benchmark TLOAD instruction."""
-    start_key = 41
-    code_key_mut = Bytecode()
-    code_val_mut = Bytecode()
     setup = Bytecode()
     if key_mut and val_mut:
-        setup = Op.PUSH1(start_key)
-        attack_block = Op.POP(Op.TLOAD(Op.DUP1))
-        code_key_mut = Op.POP + Op.GAS
-        code_val_mut = Op.TSTORE(Op.DUP2, Op.GAS)
+        setup = Op.GAS + Op.TSTORE(Op.DUP2, Op.GAS)
+        attack_block = Op.TLOAD(Op.DUP1)
     if key_mut and not val_mut:
-        attack_block = Op.POP(Op.TLOAD(Op.GAS))
+        attack_block = Op.TLOAD(Op.GAS)
     if not key_mut and val_mut:
-        attack_block = Op.POP(Op.TLOAD(Op.CALLVALUE))
-        code_val_mut = Op.TSTORE(
-            Op.CALLVALUE, Op.GAS
-        )  # CALLVALUE configured in the tx
+        setup = Op.TSTORE(Op.CALLVALUE, Op.GAS)
+        attack_block = Op.TLOAD(Op.CALLVALUE)
     if not key_mut and not val_mut:
-        attack_block = Op.POP(Op.TLOAD(Op.CALLVALUE))
+        attack_block = Op.TLOAD(Op.CALLVALUE)
 
-    cleanup = code_key_mut + code_val_mut
-    tx_value = start_key if not key_mut and val_mut else 0
+    tx_value = 42 if not key_mut and val_mut else 0
 
     benchmark_test(
-        code_generator=JumpLoopGenerator(
+        code_generator=ExtCallGenerator(
             setup=setup,
             attack_block=attack_block,
-            cleanup=cleanup,
-            tx_kwargs={
-                "value": tx_value,
-            },
+            contract_balance=tx_value,
         ),
     )