feat(benchmarking): Optimize IteratingBytecode (ethereum#2184)

* feat(testing/vm): Cache bytecode's keccak

* feat(testing/vm): Cache bytecode gas cost/refund

* feat(tests/benchmark): Cache created contract address

* feat(tests/benchmarking): Cache access list

* refactor: refactor caching with module-level ``@cache``

---------

Co-authored-by: fselmo <fselmo2@gmail.com>

Author: marioevz

packages/testing/src/execution_testing/test_types/helpers.py

@@ -11,7 +11,7 @@
     FixedSizeBytesConvertible,
 )
 from execution_testing.forks import Fork
-from execution_testing.vm import Op
+from execution_testing.vm import Bytecode, Op
 
 from .account_types import EOA
 from .utils import int_to_bytes
@@ -88,14 +88,18 @@ def compute_create_address(
 def compute_create2_address(
     address: FixedSizeBytesConvertible,
     salt: FixedSizeBytesConvertible,
-    initcode: BytesConvertible,
+    initcode: Bytecode | BytesConvertible,
 ) -> Address:
     """
     Compute address of the resulting contract created using the `CREATE2`
     opcode.
     """
+    if isinstance(initcode, Bytecode):
+        initcode_hash = initcode.keccak256()
+    else:
+        initcode_hash = Bytes(initcode).keccak256()
     hash_bytes = Bytes(
-        b"\xff" + Address(address) + Hash(salt) + Bytes(initcode).keccak256()
+        b"\xff" + Address(address) + Hash(salt) + initcode_hash
     ).keccak256()
     return Address(hash_bytes[-20:])
 

packages/testing/src/execution_testing/vm/bytecode.py

@@ -1,5 +1,6 @@
 """Ethereum Virtual Machine bytecode primitives and utilities."""
 
+from functools import cache
 from typing import Any, List, Self, SupportsBytes, Type
 
 from pydantic import GetCoreSchemaHandler
@@ -33,6 +34,7 @@ class Bytecode:
 
     _name_: str = ""
     _bytes_: bytes
+    _keccak_256_: Hash | None = None
 
     popped_stack_items: int
     pushed_stack_items: int
@@ -261,7 +263,9 @@ def hex(self) -> str:
 
     def keccak256(self) -> Hash:
         """Return the keccak256 hash of the opcode byte representation."""
-        return Bytes(self._bytes_).keccak256()
+        if self._keccak_256_ is None:
+            self._keccak_256_ = Bytes(self._bytes_).keccak256()
+        return self._keccak_256_
 
     def gas_cost(
         self,
@@ -271,13 +275,7 @@ def gas_cost(
         timestamp: int = 0,
     ) -> int:
         """Use a fork object to calculate the gas used by this bytecode."""
-        opcode_gas_calculator = fork.opcode_gas_calculator(
-            block_number=block_number, timestamp=timestamp
-        )
-        total_gas = 0
-        for opcode in self.opcode_list:
-            total_gas += opcode_gas_calculator(opcode)
-        return total_gas
+        return _bytecode_gas_cost(self, fork, block_number, timestamp)
 
     def refund(
         self,
@@ -287,13 +285,7 @@ def refund(
         timestamp: int = 0,
     ) -> int:
         """Use a fork object to calculate the gas refund from this bytecode."""
-        opcode_refund_calculator = fork.opcode_refund_calculator(
-            block_number=block_number, timestamp=timestamp
-        )
-        total_refund = 0
-        for opcode in self.opcode_list:
-            total_refund += opcode_refund_calculator(opcode)
-        return total_refund
+        return _bytecode_refund(self, fork, block_number, timestamp)
 
     @classmethod
     def __get_pydantic_core_schema__(
@@ -310,3 +302,37 @@ def __get_pydantic_core_schema__(
                 info_arg=False,
             ),
         )
+
+
+@cache
+def _bytecode_gas_cost(
+    bytecode: Bytecode,
+    fork: Type[ForkOpcodeInterface],
+    block_number: int,
+    timestamp: int,
+) -> int:
+    """Return cached gas cost for the given bytecode and fork parameters."""
+    opcode_gas_calculator = fork.opcode_gas_calculator(
+        block_number=block_number, timestamp=timestamp
+    )
+    total_gas = 0
+    for opcode in bytecode.opcode_list:
+        total_gas += opcode_gas_calculator(opcode)
+    return total_gas
+
+
+@cache
+def _bytecode_refund(
+    bytecode: Bytecode,
+    fork: Type[ForkOpcodeInterface],
+    block_number: int,
+    timestamp: int,
+) -> int:
+    """Return cached gas refund for the given bytecode and fork parameters."""
+    opcode_refund_calculator = fork.opcode_refund_calculator(
+        block_number=block_number, timestamp=timestamp
+    )
+    total_refund = 0
+    for opcode in bytecode.opcode_list:
+        total_refund += opcode_refund_calculator(opcode)
+    return total_refund

tests/benchmark/compute/helpers.py

@@ -2,7 +2,7 @@
 
 import math
 from enum import Enum, auto
-from typing import Generator, Self, Sequence, cast
+from typing import Dict, Generator, Self, Sequence, cast
 
 from execution_testing import (
     EOA,
@@ -349,6 +349,8 @@ class CustomSizedContractFactory(IteratingBytecode):
 
     _cached_address: Address
     """Cached address to avoid expensive recomputation."""
+    _cached_created_contracts: Dict[int, Address]
+    """Cached created contract addresses to avoid expensive recomputation."""
     contract_size: int
     """The size of the contracts to deploy."""
 
@@ -418,6 +420,7 @@ def __new__(
             initcode=Initcode(deploy_code=instance),
             fork=fork,
         )
+        instance._cached_created_contracts = {}
         instance.contract_size = initcode.contract_size
         deployed_address = pre.deterministic_deploy_contract(
             deploy_code=instance
@@ -494,8 +497,10 @@ def address(self) -> Address:
 
     def created_contract_address(self, *, salt: int) -> Address:
         """Get the deterministic address of the created contract."""
-        return compute_create2_address(
-            address=self.address(),
-            salt=salt,
-            initcode=self.initcode,
-        )
+        if salt not in self._cached_created_contracts:
+            self._cached_created_contracts[salt] = compute_create2_address(
+                address=self.address(),
+                salt=salt,
+                initcode=self.initcode,
+            )
+        return self._cached_created_contracts[salt]

tests/benchmark/compute/instruction/test_account_query.py

@@ -13,7 +13,7 @@
 """
 
 import math
-from typing import Any
+from typing import Any, Dict
 
 import pytest
 from execution_testing import (
@@ -539,17 +539,25 @@ def calldata(iteration_count: int, start_iteration: int) -> bytes:
     # Access list generator for warm access tests.
     # When access_warm=True, include all contract addresses that will be
     # accessed in each transaction to warm them up via access list.
+    # Note: This access list generation is very expensive due to the binary
+    # search, which builds different access lists using the same elements
+    # over and over. Caching the elements helps a bit.
+    access_list_cache: Dict[int, AccessList] = {}
+
     def access_list_generator(
         iteration_count: int, start_iteration: int
     ) -> list[AccessList] | None:
         if not access_warm:
             return None
         return [
-            AccessList(
-                address=custom_sized_contract_factory.created_contract_address(
-                    salt=i
+            access_list_cache.setdefault(
+                i,
+                AccessList(
+                    address=custom_sized_contract_factory.created_contract_address(
+                        salt=i
+                    ),
+                    storage_keys=[],
                 ),
-                storage_keys=[],
             )
             for i in range(start_iteration, start_iteration + iteration_count)
         ]