From f5797b53b20889f3392c7f53c001ce26ca79d6a6 Mon Sep 17 00:00:00 2001
From: Mario Vega <marioevz@gmail.com>
Date: Wed, 10 Jul 2024 18:23:42 +0000
Subject: [PATCH] feat(fw): create `MSTORE` macro

---
 src/ethereum_test_tools/tests/test_vm.py | 43 ++++++++++++++++++++++
 src/ethereum_test_tools/vm/opcode.py     | 46 ++++++++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/src/ethereum_test_tools/tests/test_vm.py b/src/ethereum_test_tools/tests/test_vm.py
index 7555797979..6b0a60ec05 100644
--- a/src/ethereum_test_tools/tests/test_vm.py
+++ b/src/ethereum_test_tools/tests/test_vm.py
@@ -284,6 +284,49 @@
             b"\x60\x00\x60\x00\x60\x01\xF9",
             id="Op.EXTDELEGATECALL(address=1)",
         ),
+        pytest.param(
+            Om.MSTORE(b""),
+            b"",
+            id='Om.MSTORE(b"")',
+        ),
+        pytest.param(
+            Om.MSTORE(bytes(range(32))),
+            bytes(Op.MSTORE(0, bytes(range(32)))),
+            id="Om.MSTORE(bytes(range(32)))",
+        ),
+        pytest.param(
+            Om.MSTORE(bytes(range(64))),
+            bytes(Op.MSTORE(0, bytes(range(32))) + Op.MSTORE(32, bytes(range(32, 64)))),
+            id="Om.MSTORE(bytes(range(64)))",
+        ),
+        pytest.param(
+            Om.MSTORE(bytes(range(33))),
+            bytes(
+                Op.MSTORE(0, bytes(range(32)))
+                + Op.MLOAD(32)
+                + Op.PUSH31[-1]
+                + Op.AND
+                + Op.PUSH32[b"\x20".ljust(32, b"\x00")]
+                + Op.OR
+                + Op.PUSH1(32)
+                + Op.MSTORE
+            ),
+            id="Om.MSTORE(bytes(range(33)))",
+        ),
+        pytest.param(
+            Om.MSTORE(bytes(range(63))),
+            bytes(
+                Op.MSTORE(0, bytes(range(32)))
+                + Op.MLOAD(32)
+                + Op.PUSH1[-1]
+                + Op.AND
+                + Op.PUSH32[bytes(range(32, 63)).ljust(32, b"\x00")]
+                + Op.OR
+                + Op.PUSH1(32)
+                + Op.MSTORE
+            ),
+            id="Om.MSTORE(bytes(range(63)))",
+        ),
     ],
 )
 def test_opcodes(opcodes: bytes, expected: bytes):
diff --git a/src/ethereum_test_tools/vm/opcode.py b/src/ethereum_test_tools/vm/opcode.py
index e9c5dee7d8..c26427dcf0 100644
--- a/src/ethereum_test_tools/vm/opcode.py
+++ b/src/ethereum_test_tools/vm/opcode.py
@@ -5706,6 +5706,35 @@ class Opcodes(Opcode, Enum):
 ]
 
 
+def _mstore_operation(data: OpcodeCallArg = b"", offset: OpcodeCallArg = 0) -> Bytecode:
+    """
+    Helper function to generate the bytecode that stores an arbitrary amount of data in memory.
+    """
+    assert isinstance(offset, int)
+    if isinstance(data, int):
+        data = data.to_bytes(32, "big")
+    data = to_bytes(data)  # type: ignore
+    bytecode = Bytecode()
+    for i in range(0, len(data), 32):
+        chunk = data[i : i + 32]
+        if len(chunk) == 32:
+            bytecode += Opcodes.MSTORE(offset, chunk)
+        else:
+            # We need to MLOAD the existing data at the offset and then do a bitwise OR with the
+            # new data to store it in memory.
+            bytecode += Opcodes.MLOAD(offset)
+            # Create a mask to zero out the leftmost bytes of the existing data.
+            mask_size = 32 - len(chunk)
+            bytecode += _push_opcodes_byte_list[mask_size - 1][-1]
+            bytecode += Opcodes.AND
+            bytecode += Opcodes.PUSH32[chunk.ljust(32, b"\x00")]
+            bytecode += Opcodes.OR
+            bytecode += _stack_argument_to_bytecode(offset)
+            bytecode += Opcodes.MSTORE
+        offset += len(chunk)
+    return bytecode
+
+
 class Macros(Macro, Enum):
     """
     Enum containing all macros.
@@ -5745,6 +5774,23 @@ class Macros(Macro, Enum):
     SHA3(0, 100000000000)
     """
 
+    MSTORE = Macro(lambda_operation=_mstore_operation)
+    """
+    MSTORE(data, offset)
+    ----
+
+    Place data of arbitrary length into memory at a given offset.
+
+    Inputs
+    ----
+    - data: The data to store in memory. Can be an integer or bytes.
+    - offset: The offset in memory to store the data.
+
+    Outputs
+    ----
+    - None
+    """
+
 
 class UndefinedOpcodes(Opcode, Enum):
     """