refactor: apply suggested changes

Author: LouisTsai-Csie

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
@@ -105,7 +106,8 @@ def generate_repeated_code(
         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.
-        code += Op.STOP * (max_code_size - len(code))
+        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,7 +4,7 @@
 """
 
 import math
-from typing import Dict
+from typing import Any
 
 import pytest
 from execution_testing import (
@@ -25,7 +25,6 @@
     While,
     compute_create2_address,
 )
-from execution_testing.base_types.conversions import NumberConvertible
 
 from tests.benchmark.compute.helpers import (
     XOR_TABLE,
@@ -89,7 +88,9 @@ def test_codecopy(
 
     benchmark_test(
         code_generator=JumpLoopGenerator(
-            setup=setup, attack_block=attack_block
+            setup=setup,
+            attack_block=attack_block,
+            code_padding_opcode=Op.STOP,
         )
     )
 
@@ -348,7 +349,7 @@ def test_extcodecopy_warm(
     ],
 )
 @pytest.mark.parametrize(
-    "empty_account",
+    "empty_code",
     [
         True,
         False,
@@ -372,7 +373,7 @@ def test_ext_account_query_warm(
     benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
     opcode: Op,
-    empty_account: bool,
+    empty_code: bool,
     initial_balance: bool,
     initial_storage: bool,
 ) -> None:
@@ -381,32 +382,24 @@ def test_ext_account_query_warm(
     for an account.
     """
     # Setup
-    target_addr = Address()
-    if not initial_balance and not initial_storage:
-        target_addr = pre.empty_account()
-    elif initial_balance or initial_storage:
-        target_addr = pre.fund_eoa(
-            storage={0: 0x1337} if initial_storage else {0: 0}
-        )
-
     post = {}
-    if not empty_account:
-        code = Op.STOP + Op.JUMPDEST * 100
 
-        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,
-        )
+    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,

tests/benchmark/compute/instruction/test_block_context.py

@@ -43,7 +43,7 @@ def test_block_context_ops(
         1,
         256,
         257,
-        None,
+        pytest.param(None, id="random"),
     ],
 )
 def test_blockhash(

tests/benchmark/compute/instruction/test_keccak.py

@@ -61,7 +61,7 @@ def test_keccak_max_permutations(
 
     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)),
         ),
     )

tests/benchmark/compute/instruction/test_storage.py

@@ -23,59 +23,55 @@
 # SLOAD, SSTORE, TLOAD, TSTORE
 
 
-# `key_mut` indicates the key isn't fixed.
-@pytest.mark.parametrize("key_mut", [True, False])
-# `val_mut` indicates that at the end of each big-loop, the value of the target
-# key changes.
-@pytest.mark.parametrize("val_mut", [True, False])
+@pytest.mark.parametrize("fixed_key", [True, False])
+@pytest.mark.parametrize("fixed_value", [True, False])
 def test_tload(
     benchmark_test: BenchmarkTestFiller,
-    key_mut: bool,
-    val_mut: bool,
+    fixed_key: bool,
+    fixed_value: bool,
 ) -> None:
     """Benchmark TLOAD instruction."""
     setup = Bytecode()
-    if key_mut and val_mut:
+    if not fixed_key and not fixed_value:
         setup = Op.GAS + Op.TSTORE(Op.DUP2, Op.GAS)
         attack_block = Op.TLOAD(Op.DUP1)
-    if key_mut and not val_mut:
+    if not fixed_key and fixed_value:
         attack_block = Op.TLOAD(Op.GAS)
-    if not key_mut and val_mut:
-        setup = Op.TSTORE(Op.CALLVALUE, Op.GAS)
-        attack_block = Op.TLOAD(Op.CALLVALUE)
-    if not key_mut and not val_mut:
-        attack_block = Op.TLOAD(Op.CALLVALUE)
+    if fixed_key and not fixed_value:
+        setup = Op.TSTORE(Op.CALLDATASIZE, Op.GAS)
+        attack_block = Op.TLOAD(Op.CALLDATASIZE)
+    if fixed_key and fixed_value:
+        attack_block = Op.TLOAD(Op.CALLDATASIZE)
 
-    tx_value = 42 if not key_mut and val_mut else 0
+    tx_data = b"42" if fixed_key and not fixed_value else 0
 
     benchmark_test(
         code_generator=ExtCallGenerator(
             setup=setup,
             attack_block=attack_block,
-            contract_balance=tx_value,
+            tx_kwargs={"data": tx_data},
         ),
     )
 
 
-@pytest.mark.parametrize("key_mut", [True, False])
-@pytest.mark.parametrize("dense_val_mut", [True, False])
+@pytest.mark.parametrize("fixed_key", [True, False])
+@pytest.mark.parametrize("fixed_value", [True, False])
 def test_tstore(
     benchmark_test: BenchmarkTestFiller,
-    key_mut: bool,
-    dense_val_mut: bool,
+    fixed_key: bool,
+    fixed_value: bool,
 ) -> None:
     """Benchmark TSTORE instruction."""
     init_key = 42
     setup = Op.PUSH1(init_key)
 
-    # If `dense_val_mut` is set, we use GAS as a cheap way of always
-    # storing a different value than
-    # the previous one.
-    attack_block = Op.TSTORE(Op.DUP2, Op.GAS if dense_val_mut else Op.DUP1)
+    # If fixed_value is False, we use GAS as a cheap way of always
+    # storing a different value than the previous one.
+    attack_block = Op.TSTORE(Op.DUP2, Op.GAS if not fixed_value else Op.DUP1)
 
-    # If `key_mut` is True, we mutate the key on every iteration of the
+    # If fixed_key is False, we mutate the key on every iteration of the
     # big loop.
-    cleanup = Op.POP + Op.GAS if key_mut else Bytecode()
+    cleanup = Op.POP + Op.GAS if not fixed_key else Bytecode()
 
     benchmark_test(
         code_generator=JumpLoopGenerator(

tests/benchmark/compute/instruction/test_system.py

@@ -453,7 +453,9 @@ def test_return_revert(
     )
     benchmark_test(
         code_generator=ExtCallGenerator(
-            setup=mem_preparation, attack_block=opcode(size=return_size)
+            setup=mem_preparation,
+            attack_block=opcode(size=return_size),
+            code_padding_opcode=Op.INVALID,
         ),
     )
 

tests/benchmark/compute/scenario/test_transaction_types.py

@@ -83,10 +83,11 @@ def get_single_receiver_list(
         yield receiver
 
 
+@pytest.fixture
 def ether_transfer_case(
     case_id: str, pre: Alloc, balance: int
 ) -> Tuple[Generator[Address, None, None], Generator[Address, None, None]]:
-    """Generate the test parameters based on the case ID."""
+    """Generate sender and receiver generators based on the test case."""
     if case_id == "a_to_a":
         """Sending to self."""
         senders = get_single_sender_list(pre)
@@ -137,6 +138,9 @@ def test_block_full_of_ether_transfers(
     iteration_count: int,
     transfer_amount: int,
     intrinsic_cost: int,
+    ether_transfer_case: Tuple[
+        Generator[Address, None, None], Generator[Address, None, None]
+    ],
 ) -> None:
     """
     Single test for ether transfer scenarios.
@@ -148,7 +152,7 @@ def test_block_full_of_ether_transfers(
     - a_to_diff_acc: one sender → multiple receivers
     - diff_acc_to_diff_acc: multiple senders → multiple receivers
     """
-    senders, receivers = ether_transfer_case(case_id, pre, balance)
+    senders, receivers = ether_transfer_case
 
     # Create a single block with all transactions
     txs = []