feat(vm): evm memory variable abstraction (#1609)

* feat(vm): Implement `MemoryVariable`

* refactor(tests): Use MemoryVariable in tests/frontier/scenarios/common.py

* Rename "store" to "set"

* Fixup

* refactor

* chore: ruff 79 line lenght.

---------

Co-authored-by: spencer-tb <[email protected]>

Author: marioevz

docs/CHANGELOG.md

@@ -31,6 +31,7 @@ Test fixtures for use by clients are available for each release on the [Github r
 - 🐞 Use `engine_newPayloadV5` for `>=Amsterdam` forks in `consume engine` ([#2170](https://github.com/ethereum/execution-spec-tests/pull/2170)).
 - 🔀 Refactor EIP-7928 (BAL) absence checks into a friendlier class-based DevEx ([#2175](https://github.com/ethereum/execution-spec-tests/pull/2175)).
 - 🐞 Tighten up validation for empty lists on Block-Level Access List tests ([#2118](https://github.com/ethereum/execution-spec-tests/pull/2118)).
+- ✨ Added the `MemoryVariable` EVM abstraction to generate more readable bytecode when there's heavy use of variables that are stored in memory ([#1609](https://github.com/ethereum/execution-spec-tests/pull/1609)).
 
 ### 🧪 Test Cases
 

src/ethereum_test_tools/__init__.py

@@ -82,6 +82,7 @@
     EVMCodeType,
     Macro,
     Macros,
+    MemoryVariable,
     Opcode,
     OpcodeCallArg,
     Opcodes,
@@ -157,6 +158,7 @@
     "JumpLoopGenerator",
     "Macro",
     "Macros",
+    "MemoryVariable",
     "NetworkWrappedTransaction",
     "Opcode",
     "OpcodeCallArg",

src/ethereum_test_vm/__init__.py

@@ -2,14 +2,15 @@
 
 from .bytecode import Bytecode
 from .evm_types import EVMCodeType
-from .helpers import call_return_code
+from .helpers import MemoryVariable, call_return_code
 from .opcodes import Macro, Macros, Opcode, OpcodeCallArg, Opcodes, UndefinedOpcodes
 
 __all__ = (
     "Bytecode",
     "EVMCodeType",
     "Macro",
     "Macros",
+    "MemoryVariable",
     "Opcode",
     "OpcodeCallArg",
     "Opcodes",

src/ethereum_test_vm/helpers.py

@@ -1,8 +1,92 @@
 """Helper functions for the EVM."""
 
+from .bytecode import Bytecode
 from .opcodes import Opcodes as Op
 
 
+class MemoryVariable(Bytecode):
+    """
+    Variable abstraction to help keep track values that are stored in memory.
+
+    To use, simply declare a variable with an unique offset that is not used
+    by any other variable.
+
+    The variable then can be used in-place to read the value from memory:
+
+    ```python
+    v = MemoryVariable(128)
+
+    bytecode = Op.ADD(v, Op.CALLDATASIZE())
+    ```
+
+    The previous example is equivalent to:
+
+    ```python
+    bytecode = Op.ADD(Op.MLOAD(offset=128), Op.CALLDATASIZE())
+    ```
+
+    The variable also contains methods to add and subtract values from the
+    memory offset.
+
+    ```python
+    v = MemoryVariable(128)
+
+    bytecode = (
+        v.set(0xff)
+        + v.add(1)
+        + v.return_value()
+    )
+    ```
+
+    The previous example is equivalent to:
+
+    ```python
+    bytecode = (
+        Op.MSTORE(offset=128, value=0xff)
+        + Op.MSTORE(offset=128, value=Op.ADD(Op.MLOAD(offset=128), 1))
+        + Op.RETURN(offset=128, size=32)
+    )
+    ```
+
+    """
+
+    offset: int
+
+    def __new__(cls, offset: int):
+        """
+        Instantiate a new EVM memory variable.
+
+        When used with normal bytecode, this class simply returns the MLOAD
+        with the provided offset.
+        """
+        instance = super().__new__(cls, Op.MLOAD(offset=offset))
+        instance.offset = offset
+        return instance
+
+    def set(self, value: int | Bytecode) -> Bytecode:
+        """Set the given value at the memory location of this variable."""
+        return Op.MSTORE(offset=self.offset, value=value)
+
+    def add(self, value: int | Bytecode) -> Bytecode:
+        """In-place add the given value to the one currently in memory."""
+        return Op.MSTORE(offset=self.offset, value=Op.ADD(Op.MLOAD(offset=self.offset), value))
+
+    def sub(self, value: int | Bytecode) -> Bytecode:
+        """
+        In-place subtract the given value from the one currently
+        in memory.
+        """
+        return Op.MSTORE(offset=self.offset, value=Op.SUB(Op.MLOAD(offset=self.offset), value))
+
+    def store_value(self, key: int | Bytecode) -> Bytecode:
+        """Op.SSTORE the value that is currently in memory."""
+        return Op.SSTORE(key, Op.MLOAD(offset=self.offset))
+
+    def return_value(self) -> Bytecode:
+        """Op.RETURN the value that is currently in memory."""
+        return Op.RETURN(offset=self.offset, size=32)
+
+
 def call_return_code(opcode: Op, success: bool, *, revert: bool = False) -> int:
     """Return return code for a CALL operation."""
     if opcode in [Op.CALL, Op.CALLCODE, Op.DELEGATECALL, Op.STATICCALL]:

tests/frontier/scenarios/common.py

@@ -5,7 +5,7 @@
 from enum import Enum
 
 from ethereum_test_forks import Fork, Frontier
-from ethereum_test_tools import Address, Alloc, Bytecode, Conditional
+from ethereum_test_tools import Address, Alloc, Bytecode, Conditional, MemoryVariable
 from ethereum_test_vm import Opcodes as Op
 
 
@@ -196,20 +196,41 @@ def make_gas_hash_contract(pre: Alloc) -> Address:
     So that if we can't check exact value in expect section, we at least
     could spend unique gas amount.
     """
+    # EVM memory variables
+    byte_offset = MemoryVariable(0)
+    current_byte = MemoryVariable(32)
+
+    # Code for memory initialization
+    initialize_code = byte_offset.set(0)
+    calldata_copy = Op.JUMPDEST + Op.CALLDATACOPY(
+        dest_offset=current_byte.offset + 32 - 1,
+        offset=byte_offset,
+        size=1,
+    )
+
+    # Code offsets
+    offset_calldata_copy = len(initialize_code)
+    offset_conditional = offset_calldata_copy + len(calldata_copy)
+
+    # Deploy contract
     gas_hash_address = pre.deploy_contract(
-        code=Op.MSTORE(0, 0)
-        + Op.JUMPDEST
-        + Op.CALLDATACOPY(63, Op.MLOAD(0), 1)
-        + Op.JUMPDEST
+        code=initialize_code
+        + calldata_copy  # offset_calldata_copy
+        + Op.JUMPDEST  # offset_conditional
         + Conditional(
-            condition=Op.ISZERO(Op.MLOAD(32)),
-            if_true=Op.MSTORE(0, Op.ADD(1, Op.MLOAD(0)))
-            + Conditional(
-                condition=Op.GT(Op.MLOAD(0), 32),
-                if_true=Op.RETURN(0, 0),
-                if_false=Op.JUMP(5),
+            condition=Op.ISZERO(current_byte),
+            if_true=(
+                # Increase the calldata byte offset, and if it's greater than
+                # the calldata size, return, otherwise jump to the calldata
+                # copy code and read the next byte.
+                byte_offset.add(1)
+                + Conditional(
+                    condition=Op.GT(byte_offset, Op.CALLDATASIZE()),
+                    if_true=Op.RETURN(offset=0, size=0),
+                    if_false=Op.JUMP(offset_calldata_copy),
+                )
             ),
-            if_false=Op.MSTORE(32, Op.SUB(Op.MLOAD(32), 1)) + Op.JUMP(14),
+            if_false=(current_byte.sub(1) + Op.JUMP(offset_conditional)),
         )
     )
     return gas_hash_address
@@ -241,19 +262,25 @@ def make_invalid_opcode_contract(pre: Alloc, fork: Fork) -> Address:
         if op not in valid_opcode_values:
             invalid_opcodes.append(op)
 
+    results_sum = MemoryVariable(0)
+    current_opcode = MemoryVariable(32)
+
     code = Bytecode(
         sum(
-            Op.MSTORE(64, opcode)
-            + Op.MSTORE(
-                32,
-                Op.CALL(gas=50000, address=invalid_opcode_caller, args_offset=64, args_size=32),
+            current_opcode.set(opcode)
+            + results_sum.add(
+                Op.CALL(
+                    gas=50000,
+                    address=invalid_opcode_caller,
+                    args_offset=current_opcode.offset,
+                    args_size=32,
+                ),
             )
-            + Op.MSTORE(0, Op.ADD(Op.MLOAD(0), Op.MLOAD(32)))
             for opcode in invalid_opcodes
         )
         # If any of invalid instructions works, mstore[0] will be > 1
-        + Op.MSTORE(0, Op.ADD(Op.MLOAD(0), 1))
-        + Op.RETURN(0, 32)
+        + results_sum.add(1)
+        + results_sum.return_value()
     )
 
     return pre.deploy_contract(code=code)