Merge pull request #239 from marioevz/conditional-generator

tools/code: Add conditional bytecode generator

Author: spencer-tb

src/ethereum_test_tools/__init__.py

@@ -3,7 +3,7 @@
 tests.
 """
 
-from .code import Code, CodeGasMeasure, Initcode, Yul, YulCompiler
+from .code import Code, CodeGasMeasure, Conditional, Initcode, Yul, YulCompiler
 from .common import (
     AccessList,
     Account,
@@ -59,6 +59,7 @@
     "BlockchainTestFiller",
     "Code",
     "CodeGasMeasure",
+    "Conditional",
     "EngineAPIError",
     "Environment",
     "Fixture",

src/ethereum_test_tools/code/__init__.py

@@ -2,12 +2,13 @@
 Code related utilities and classes.
 """
 from .code import Code
-from .generators import CodeGasMeasure, Initcode
+from .generators import CodeGasMeasure, Conditional, Initcode
 from .yul import Yul, YulCompiler
 
 __all__ = (
     "Code",
     "CodeGasMeasure",
+    "Conditional",
     "Initcode",
     "Yul",
     "YulCompiler",

src/ethereum_test_tools/code/generators.py

@@ -7,6 +7,7 @@
 
 from ..common.conversions import to_bytes
 from ..common.helpers import ceiling_division
+from ..vm.opcode import Opcodes as Op
 from .code import Code
 
 GAS_PER_DEPLOYED_CODE_BYTE = 0xC8
@@ -181,3 +182,51 @@ def __post_init__(self):
             ]
         )
         self.bytecode = res
+
+
+@dataclass(kw_only=True)
+class Conditional(Code):
+    """
+    Helper class used to generate conditional bytecode.
+    """
+
+    condition: str | bytes | SupportsBytes
+    """
+    Condition bytecode which must return the true or false condition of the conditional statement.
+    """
+
+    if_true: str | bytes | SupportsBytes
+    """
+    Bytecode to execute if the condition is true.
+    """
+
+    if_false: str | bytes | SupportsBytes
+    """
+    Bytecode to execute if the condition is false.
+    """
+
+    def __post_init__(self):
+        """
+        Assemble the conditional bytecode by generating the necessary jumps and
+        jumpdests surrounding the condition and the two possible execution
+        paths.
+
+        In the future, PC usage should be replaced by using RJUMP and RJUMPI
+        """
+        condition_bytes = to_bytes(self.condition)
+        if_true_bytes = to_bytes(self.if_true)
+        if_false_bytes = to_bytes(self.if_false)
+
+        # First we append a jumpdest to the start of the true branch
+        if_true_bytes = Op.JUMPDEST + if_true_bytes
+
+        # Then we append the unconditional jump to the end of the false branch, used to skip the
+        # true branch
+        if_false_bytes += Op.JUMP(Op.ADD(Op.PC, len(if_true_bytes) + 3))
+
+        # Then we need to do the conditional jump by skipping the false branch
+        condition_bytes = Op.JUMPI(Op.ADD(Op.PC, len(if_false_bytes) + 3), condition_bytes)
+
+        # Finally we append the true and false branches, and the condition, plus the jumpdest at
+        # the very end
+        self.bytecode = condition_bytes + if_false_bytes + if_true_bytes + Op.JUMPDEST

src/ethereum_test_tools/tests/test_code.py

@@ -10,7 +10,8 @@
 
 from ethereum_test_forks import Fork, Homestead, Shanghai, forks_from_until, get_deployed_forks
 
-from ..code import Code, Initcode, Yul
+from ..code import Code, Conditional, Initcode, Yul
+from ..vm.opcode import Opcodes as Op
 
 
 @pytest.mark.parametrize(
@@ -239,3 +240,23 @@ def test_yul(
 )
 def test_initcode(initcode: Initcode, bytecode: bytes):
     assert bytes(initcode) == bytecode
+
+
+@pytest.mark.parametrize(
+    "conditional_bytecode,expected",
+    [
+        (
+            Conditional(
+                condition=Op.CALLDATALOAD(0),
+                if_true=Op.MSTORE(0, Op.SLOAD(0)) + Op.RETURN(0, 32),
+                if_false=Op.SSTORE(0, 69),
+            ),
+            bytes.fromhex("600035600d5801576045600055600f5801565b60005460005260206000f35b"),
+        ),
+    ],
+)
+def test_opcodes_if(conditional_bytecode: bytes, expected: bytes):
+    """
+    Test that the if opcode macro is transformed into bytecode as expected.
+    """
+    assert bytes(conditional_bytecode) == expected