refactor: update worst memory scenario

Author: LouisTsai-Csie

tests/benchmark/test_worst_memory.py

@@ -5,29 +5,30 @@
 Tests that benchmark EVMs in the worst-case memory opcodes.
 """
 
+from enum import auto
+
 import pytest
 
 from ethereum_test_base_types.base_types import Bytes
+from ethereum_test_benchmark.benchmark_code_generator import JumpLoopGenerator
 from ethereum_test_forks import Fork
 from ethereum_test_tools import (
     Alloc,
+    BenchmarkTestFiller,
     Bytecode,
-    StateTestFiller,
     Transaction,
 )
 from ethereum_test_tools.vm.opcode import Opcodes as Op
 
-from .helpers import code_loop_precompile_call
-
 REFERENCE_SPEC_GIT_PATH = "TODO"
 REFERENCE_SPEC_VERSION = "TODO"
 
 
 class CallDataOrigin:
     """Enum for calldata origins."""
 
-    TRANSACTION = 1
-    CALL = 2
+    TRANSACTION = auto()
+    CALL = auto()
 
 
 @pytest.mark.parametrize(
@@ -61,7 +62,7 @@ class CallDataOrigin:
     ],
 )
 def test_worst_calldatacopy(
-    state_test: StateTestFiller,
+    benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
     fork: Fork,
     origin: CallDataOrigin,
@@ -87,13 +88,15 @@ def test_worst_calldatacopy(
     #
     # If `non_zero_data` is True, we leverage CALLDATASIZE for the copy length. Otherwise, since we
     # don't send zero data explicitly via calldata, PUSH the target size and use DUP1 to copy it.
-    prefix = Bytecode() if non_zero_data or size == 0 else Op.PUSH3(size)
+    setup = Bytecode() if non_zero_data or size == 0 else Op.PUSH3(size)
     src_dst = 0 if fixed_src_dst else Op.MOD(Op.GAS, 7)
     attack_block = Op.CALLDATACOPY(
         src_dst, src_dst, Op.CALLDATASIZE if non_zero_data or size == 0 else Op.DUP1
     )
-    code = code_loop_precompile_call(prefix, attack_block, fork)
-    code_address = pre.deploy_contract(code=code)
+
+    code_address = JumpLoopGenerator(setup=setup, attack_block=attack_block).deploy_contracts(
+        pre, fork
+    )
 
     tx_target = code_address
 
@@ -102,15 +105,17 @@ def test_worst_calldatacopy(
     if origin == CallDataOrigin.CALL:
         # If `non_zero_data` is False we leverage just using zeroed memory. Otherwise, we
         # copy the calldata received from the transaction.
-        prefix = (
+        setup = (
             Op.CALLDATACOPY(Op.PUSH0, Op.PUSH0, Op.CALLDATASIZE) if non_zero_data else Bytecode()
         ) + Op.JUMPDEST
         arg_size = Op.CALLDATASIZE if non_zero_data else size
-        code = prefix + Op.STATICCALL(
+        attack_block = Op.STATICCALL(
             address=code_address, args_offset=Op.PUSH0, args_size=arg_size
         )
-        code += Op.JUMP(len(prefix) - 1)
-        tx_target = pre.deploy_contract(code=code)
+
+        tx_target = JumpLoopGenerator(setup=setup, attack_block=attack_block).deploy_contracts(
+            pre, fork
+        )
 
     tx = Transaction(
         to=tx_target,
@@ -119,7 +124,7 @@ def test_worst_calldatacopy(
         sender=pre.fund_eoa(),
     )
 
-    state_test(
+    benchmark_test(
         pre=pre,
         post={},
         tx=tx,
@@ -144,38 +149,40 @@ def test_worst_calldatacopy(
     ],
 )
 def test_worst_codecopy(
-    state_test: StateTestFiller,
+    benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
     fork: Fork,
     max_code_size_ratio: float,
     fixed_src_dst: bool,
-    gas_benchmark_value: int,
+    dq,
 ):
     """Test running a block filled with CODECOPY executions."""
     max_code_size = fork.max_code_size()
 
     size = int(max_code_size * max_code_size_ratio)
 
-    code_prefix = Op.PUSH32(size)
+    setup = Op.PUSH32(size)
     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 = code_loop_precompile_call(code_prefix, attack_block, fork)
+
+    code = JumpLoopGenerator(setup=setup, attack_block=attack_block).generate_repeated_code(
+        attack_block, Bytecode(), Bytecode(), fork
+    )
 
     # The code generated above is not guaranteed to be of max_code_size, so we pad it since
     # a test parameter targets CODECOPYing a contract with max code size. Padded bytecode values
     # are not relevant.
-    code = code + Op.INVALID * (max_code_size - len(code))
+    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),
-        gas_limit=gas_benchmark_value,
         sender=pre.fund_eoa(),
     )
 
-    state_test(
+    benchmark_test(
         pre=pre,
         post={},
         tx=tx,
@@ -199,16 +206,12 @@ def test_worst_codecopy(
     ],
 )
 def test_worst_returndatacopy(
-    state_test: StateTestFiller,
+    benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
-    fork: Fork,
     size: int,
     fixed_dst: bool,
-    gas_benchmark_value: int,
 ):
     """Test running a block filled with RETURNDATACOPY executions."""
-    max_code_size = fork.max_code_size()
-
     # Create the contract that will RETURN the data that will be used for RETURNDATACOPY.
     # Random-ish data is injected at different points in memory to avoid making the content
     # predictable. If `size` is 0, this helper contract won't be used.
@@ -220,13 +223,13 @@ def test_worst_returndatacopy(
     )
     helper_contract = pre.deploy_contract(code=code)
 
-    # We create the contract that will be doing the RETURNDATACOPY multiple times.
     returndata_gen = Op.STATICCALL(address=helper_contract) if size > 0 else Bytecode()
     dst = 0 if fixed_dst else Op.MOD(Op.GAS, 7)
-    attack_iter = Op.RETURNDATACOPY(dst, Op.PUSH0, Op.RETURNDATASIZE)
 
-    jumpdest = Op.JUMPDEST
-    jump_back = Op.JUMP(len(returndata_gen))
+    # 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)
+
     # The attack loop is constructed as:
     # ```
     # JUMPDEST(#)
@@ -238,30 +241,13 @@ def test_worst_returndatacopy(
     # ```
     # The goal is that once per (big) loop iteration, the helper contract is called to
     # generate fresh returndata to continue calling RETURNDATACOPY.
-    max_iters_loop = (
-        max_code_size - 2 * len(returndata_gen) - len(jumpdest) - len(jump_back)
-    ) // len(attack_iter)
-    code = (
-        returndata_gen
-        + jumpdest
-        + sum([attack_iter] * max_iters_loop)
-        + returndata_gen
-        + jump_back
-    )
-    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),
-        gas_limit=gas_benchmark_value,
-        sender=pre.fund_eoa(),
-    )
 
-    state_test(
+    benchmark_test(
         pre=pre,
         post={},
-        tx=tx,
+        code_generator=JumpLoopGenerator(
+            setup=returndata_gen, attack_block=attack_block, cleanup=returndata_gen
+        ),
     )
 
 
@@ -282,42 +268,24 @@ def test_worst_returndatacopy(
     ],
 )
 def test_worst_mcopy(
-    state_test: StateTestFiller,
+    benchmark_test: BenchmarkTestFiller,
     pre: Alloc,
-    fork: Fork,
     size: int,
     fixed_src_dst: bool,
-    gas_benchmark_value: int,
 ):
     """Test running a block filled with MCOPY executions."""
-    max_code_size = fork.max_code_size()
+    src_dst = 0 if fixed_src_dst else Op.MOD(Op.GAS, 7)
+    attack_block = Op.MCOPY(src_dst, src_dst, size)
 
     mem_touch = (
         Op.MSTORE8(0, Op.GAS) + Op.MSTORE8(size // 2, Op.GAS) + Op.MSTORE8(size - 1, Op.GAS)
         if size > 0
         else Bytecode()
     )
-    src_dst = 0 if fixed_src_dst else Op.MOD(Op.GAS, 7)
-    attack_block = Op.MCOPY(src_dst, src_dst, size)
-
-    jumpdest = Op.JUMPDEST
-    jump_back = Op.JUMP(len(mem_touch))
-    max_iters_loop = (max_code_size - 2 * len(mem_touch) - len(jumpdest) - len(jump_back)) // len(
-        attack_block
-    )
-    code = mem_touch + jumpdest + sum([attack_block] * max_iters_loop) + mem_touch + jump_back
-    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),
-        gas_limit=gas_benchmark_value,
-        sender=pre.fund_eoa(),
-    )
-
-    state_test(
+    benchmark_test(
         pre=pre,
         post={},
-        tx=tx,
+        code_generator=JumpLoopGenerator(
+            setup=mem_touch, attack_block=attack_block, cleanup=mem_touch
+        ),
     )