From 5bf863348bc426fcea1a8989e6e61157fc60d385 Mon Sep 17 00:00:00 2001
From: Mario Vega <marioevz@gmail.com>
Date: Fri, 1 Sep 2023 08:22:05 -0600
Subject: [PATCH] tools: add optional prefix code to `Initcode` (#282)

---
 src/ethereum_test_tools/code/generators.py | 29 ++++++----
 src/ethereum_test_tools/tests/test_code.py | 61 ++++++++++++++++++++--
 2 files changed, 75 insertions(+), 15 deletions(-)

diff --git a/src/ethereum_test_tools/code/generators.py b/src/ethereum_test_tools/code/generators.py
index bccf5d52dd..35fc2348c6 100644
--- a/src/ethereum_test_tools/code/generators.py
+++ b/src/ethereum_test_tools/code/generators.py
@@ -47,6 +47,8 @@ def __init__(
         *,
         deploy_code: str | bytes | SupportsBytes,
         initcode_length: Optional[int] = None,
+        initcode_prefix: str | bytes | SupportsBytes = b"",
+        initcode_prefix_execution_gas: int = 0,
         padding_byte: int = 0x00,
         name: Optional[str] = None,
     ):
@@ -54,12 +56,13 @@ def __init__(
         Generate legacy initcode that inits a contract with the specified code.
         The initcode can be padded to a specified length for testing purposes.
         """
-        self.execution_gas = 0
+        self.execution_gas = initcode_prefix_execution_gas
         self.deploy_code = deploy_code
         deploy_code_bytes = to_bytes(self.deploy_code)
         code_length = len(deploy_code_bytes)
 
-        initcode = bytearray()
+        initcode_prefix_bytes = to_bytes(initcode_prefix)
+        initcode = bytearray(initcode_prefix_bytes)
 
         # PUSH2: length=<bytecode length>
         initcode.append(0x61)
@@ -75,9 +78,11 @@ def __init__(
         initcode.append(0x81)
         self.execution_gas += 3
 
-        # PUSH1: initcode_length=11 (constant)
+        # PUSH1: initcode_length=11 + len(initcode_prefix_bytes) (constant)
+        no_prefix_length = 0x0B
+        assert no_prefix_length + len(initcode_prefix_bytes) <= 0xFF, "initcode prefix too long"
         initcode.append(0x60)
-        initcode.append(0x0B)
+        initcode.append(no_prefix_length + len(initcode_prefix_bytes))
         self.execution_gas += 3
 
         # DUP3
@@ -97,19 +102,21 @@ def __init__(
         initcode.append(0xF3)
         self.execution_gas += 0
 
-        pre_padding_bytes = bytes(initcode) + deploy_code_bytes
+        initcode_plus_deploy_code = bytes(initcode) + deploy_code_bytes
+        padding_bytes = bytes()
 
         if initcode_length is not None:
-            if len(pre_padding_bytes) > initcode_length:
-                raise Exception("Invalid specified length for initcode")
+            assert initcode_length >= len(
+                initcode_plus_deploy_code
+            ), "specified invalid lenght for initcode"
 
-            padding_bytes = bytes([padding_byte] * (initcode_length - len(pre_padding_bytes)))
-        else:
-            padding_bytes = bytes()
+            padding_bytes = bytes(
+                [padding_byte] * (initcode_length - len(initcode_plus_deploy_code))
+            )
 
         self.deployment_gas = GAS_PER_DEPLOYED_CODE_BYTE * len(deploy_code_bytes)
 
-        super().__init__(pre_padding_bytes + padding_bytes, name=name)
+        super().__init__(initcode_plus_deploy_code + padding_bytes, name=name)
 
 
 @dataclass(kw_only=True)
diff --git a/src/ethereum_test_tools/tests/test_code.py b/src/ethereum_test_tools/tests/test_code.py
index feac14a5d2..c612ab3a2b 100644
--- a/src/ethereum_test_tools/tests/test_code.py
+++ b/src/ethereum_test_tools/tests/test_code.py
@@ -174,7 +174,7 @@ def test_yul(
 @pytest.mark.parametrize(
     "initcode,bytecode",
     [
-        (
+        pytest.param(
             Initcode(deploy_code=bytes()),
             bytes(
                 [
@@ -191,8 +191,29 @@ def test_yul(
                     0xF3,
                 ]
             ),
+            id="empty-deployed-code",
         ),
-        (
+        pytest.param(
+            Initcode(deploy_code=bytes(), initcode_prefix=bytes([0x00])),
+            bytes(
+                [
+                    0x00,
+                    0x61,
+                    0x00,
+                    0x00,
+                    0x60,
+                    0x00,
+                    0x81,
+                    0x60,
+                    0x0C,
+                    0x82,
+                    0x39,
+                    0xF3,
+                ]
+            ),
+            id="empty-deployed-code-with-prefix",
+        ),
+        pytest.param(
             Initcode(deploy_code=bytes(), initcode_length=20),
             bytes(
                 [
@@ -210,8 +231,9 @@ def test_yul(
                 ]
                 + [0x00] * 9  # padding
             ),
+            id="empty-deployed-code-with-padding",
         ),
-        (
+        pytest.param(
             Initcode(deploy_code=bytes([0x00]), initcode_length=20),
             bytes(
                 [
@@ -227,9 +249,40 @@ def test_yul(
                     0x39,
                     0xF3,
                 ]
-                + [0x00]
+                + [0x00]  # deployed code
                 + [0x00] * 8  # padding
             ),
+            id="single-byte-deployed-code-with-padding",
+        ),
+        pytest.param(
+            Initcode(
+                deploy_code=bytes([0x00]),
+                initcode_prefix=Op.SSTORE(0, 1),
+                initcode_length=20,
+            ),
+            bytes(
+                [
+                    0x60,
+                    0x01,
+                    0x60,
+                    0x00,
+                    0x55,
+                    0x61,
+                    0x00,
+                    0x01,
+                    0x60,
+                    0x00,
+                    0x81,
+                    0x60,
+                    0x10,
+                    0x82,
+                    0x39,
+                    0xF3,
+                ]
+                + [0x00]  # deployed code
+                + [0x00] * 3  # padding
+            ),
+            id="single-byte-deployed-code-with-padding-and-prefix",
         ),
     ],
 )