tools: add optional prefix code to Initcode (#282)

ethereum · Sep 1, 2023 · 5bf8633 · 5bf8633
1 parent fa1ca88
commit 5bf8633
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 15 deletions.
diff --git a/src/ethereum_test_tools/code/generators.py b/src/ethereum_test_tools/code/generators.py
@@ -47,19 +47,22 @@ 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,
     ):
         """
         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
@@ -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",
         ),
     ],
 )