refactor(tests): update MSIZE and SELFBALANCE benchmark test (#1771)

* refactor: create helper function for the repeated code sequence

* refactor: update memsize opcode

* refactor: update selfbalance opcode

Author: LouisTsai-Csie

tests/zkevm/helpers.py

@@ -0,0 +1,23 @@
+"""Helper functions for the EVM benchmark worst-case tests."""
+
+from ethereum_test_forks import Fork
+from ethereum_test_tools import Bytecode
+from ethereum_test_tools.vm.opcode import Opcodes as Op
+
+
+def code_loop_precompile_call(calldata: Bytecode, attack_block: Bytecode, fork: Fork):
+    """Create a code loop that calls a precompile with the given calldata."""
+    max_code_size = fork.max_code_size()
+
+    # The attack contract is: CALLDATA_PREP + #JUMPDEST + [attack_block]* + JUMP(#)
+    jumpdest = Op.JUMPDEST
+    jump_back = Op.JUMP(len(calldata))
+    max_iters_loop = (max_code_size - len(calldata) - len(jumpdest) - len(jump_back)) // len(
+        attack_block
+    )
+    code = calldata + jumpdest + sum([attack_block] * max_iters_loop) + jump_back
+    if len(code) > max_code_size:
+        # Must never happen, but keep it as a sanity check.
+        raise ValueError(f"Code size {len(code)} exceeds maximum code size {max_code_size}")
+
+    return code

tests/zkevm/test_worst_compute.py

@@ -39,6 +39,8 @@
 from tests.prague.eip2537_bls_12_381_precompiles import spec as bls12381_spec
 from tests.prague.eip2537_bls_12_381_precompiles.spec import BytesConcatenation
 
+from .helpers import code_loop_precompile_call
+
 REFERENCE_SPEC_GIT_PATH = "TODO"
 REFERENCE_SPEC_VERSION = "TODO"
 
@@ -311,18 +313,20 @@ def test_worst_msize(
     The `mem_size` parameter indicates by how much the memory is expanded.
     """
     env = Environment()
-    max_code_size = fork.max_code_size()
+    max_stack_height = fork.max_stack_height()
 
-    # We use CALLVALUE for the parameter since is 1 gas cheaper than PUSHX.
-    code_prefix = Op.MLOAD(Op.CALLVALUE) + Op.JUMPDEST
-    iter_loop = Op.POP(Op.MSIZE)
-    code_suffix = Op.JUMP(len(code_prefix) - 1)
-    code_iter_len = (max_code_size - len(code_prefix) - len(code_suffix)) // len(iter_loop)
-    code = code_prefix + iter_loop * code_iter_len + code_suffix
-    assert len(code) <= max_code_size
+    code_sequence = Op.MLOAD(Op.CALLVALUE) + Op.POP + Op.MSIZE * max_stack_height
+    target_address = pre.deploy_contract(code=code_sequence)
+
+    calldata = Bytecode()
+    attack_block = Op.POP(Op.STATICCALL(Op.GAS, target_address, 0, 0, 0, 0))
+    code = code_loop_precompile_call(calldata, attack_block, fork)
+    assert len(code) <= fork.max_code_size()
+
+    code_address = pre.deploy_contract(code=code)
 
     tx = Transaction(
-        to=pre.deploy_contract(code=bytes(code)),
+        to=code_address,
         gas_limit=env.gas_limit,
         sender=pre.fund_eoa(),
         value=mem_size,
@@ -823,24 +827,6 @@ def test_worst_precompile_fixed_cost(
     )
 
 
-def code_loop_precompile_call(calldata: Bytecode, attack_block: Bytecode, fork: Fork):
-    """Create a code loop that calls a precompile with the given calldata."""
-    max_code_size = fork.max_code_size()
-
-    # The attack contract is: CALLDATA_PREP + #JUMPDEST + [attack_block]* + JUMP(#)
-    jumpdest = Op.JUMPDEST
-    jump_back = Op.JUMP(len(calldata))
-    max_iters_loop = (max_code_size - len(calldata) - len(jumpdest) - len(jump_back)) // len(
-        attack_block
-    )
-    code = calldata + jumpdest + sum([attack_block] * max_iters_loop) + jump_back
-    if len(code) > max_code_size:
-        # Must never happen, but keep it as a sanity check.
-        raise ValueError(f"Code size {len(code)} exceeds maximum code size {max_code_size}")
-
-    return code
-
-
 @pytest.mark.zkevm
 @pytest.mark.valid_from("Cancun")
 @pytest.mark.slow

tests/zkevm/test_worst_stateful_opcodes.py

@@ -27,6 +27,8 @@
 )
 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"
 
@@ -415,16 +417,25 @@ def test_worst_blockhash(
 def test_worst_selfbalance(
     state_test: StateTestFiller,
     pre: Alloc,
+    fork: Fork,
 ):
     """Test running a block with as many SELFBALANCE opcodes as possible."""
     env = Environment()
+    max_stack_height = fork.max_stack_height()
+
+    code_sequence = Op.SELFBALANCE * max_stack_height
+    target_address = pre.deploy_contract(code=code_sequence)
+
+    calldata = Bytecode()
+    attack_block = Op.POP(Op.STATICCALL(Op.GAS, target_address, 0, 0, 0, 0))
+
+    code = code_loop_precompile_call(calldata, attack_block, fork)
+    assert len(code) <= fork.max_code_size()
+
+    code_address = pre.deploy_contract(code=code)
 
-    execution_code = While(
-        body=Op.POP(Op.SELFBALANCE),
-    )
-    execution_code_address = pre.deploy_contract(code=execution_code)
     tx = Transaction(
-        to=execution_code_address,
+        to=code_address,
         gas_limit=env.gas_limit,
         sender=pre.fund_eoa(),
     )