refactor[venom]: add new venom test machinery

tests/functional/venom/parser/test_parsing.py

@@ -1,49 +1,9 @@
+from tests.venom_utils import assert_ctx_eq
 from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRLabel, IRLiteral, IRVariable
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
 from vyper.venom.parser import parse_venom
 
-# TODO: Refactor tests with these helpers
-
-
-def instructions_eq(i1: IRInstruction, i2: IRInstruction) -> bool:
-    return i1.output == i2.output and i1.opcode == i2.opcode and i1.operands == i2.operands
-
-
-def assert_bb_eq(bb1: IRBasicBlock, bb2: IRBasicBlock):
-    assert bb1.label.value == bb2.label.value
-    assert len(bb1.instructions) == len(bb2.instructions)
-    for i1, i2 in zip(bb1.instructions, bb2.instructions):
-        assert instructions_eq(i1, i2), f"[{i1}] != [{i2}]"
-
-
-def assert_fn_eq(fn1: IRFunction, fn2: IRFunction):
-    assert fn1.name.value == fn2.name.value
-    assert fn1.last_variable == fn2.last_variable
-    assert len(fn1._basic_block_dict) == len(fn2._basic_block_dict)
-
-    for name1, bb1 in fn1._basic_block_dict.items():
-        assert name1 in fn2._basic_block_dict
-        assert_bb_eq(bb1, fn2._basic_block_dict[name1])
-
-    # check function entry is the same
-    assert fn1.entry.label == fn2.entry.label
-
-
-def assert_ctx_eq(ctx1: IRContext, ctx2: IRContext):
-    assert ctx1.last_label == ctx2.last_label
-    assert len(ctx1.functions) == len(ctx2.functions)
-    for label1, fn1 in ctx1.functions.items():
-        assert label1 in ctx2.functions
-        assert_fn_eq(fn1, ctx2.functions[label1])
-
-    # check entry function is the same
-    assert next(iter(ctx1.functions.keys())) == next(iter(ctx2.functions.keys()))
-
-    assert len(ctx1.data_segment) == len(ctx2.data_segment)
-    for d1, d2 in zip(ctx1.data_segment, ctx2.data_segment):
-        assert instructions_eq(d1, d2), f"data: [{d1}] != [{d2}]"
-
 
 def test_single_bb():
     source = """
@@ -70,7 +30,7 @@ def test_multi_bb_single_fn():
     function start {
         start:
             %1 = callvalue
-            jnz @fine, @has_callvalue, %1
+            jnz %1, @fine, @has_callvalue
         fine:
             %2 = calldataload 4
             %4 = add %2, 279387
@@ -89,7 +49,7 @@ def test_multi_bb_single_fn():
 
     start_bb = start_fn.get_basic_block("start")
     start_bb.append_instruction("callvalue", ret=IRVariable("1"))
-    start_bb.append_instruction("jnz", IRVariable("1"), IRLabel("has_callvalue"), IRLabel("fine"))
+    start_bb.append_instruction("jnz", IRVariable("1"), IRLabel("fine"), IRLabel("has_callvalue"))
 
     start_fn.append_basic_block(fine_bb := IRBasicBlock(IRLabel("fine"), start_fn))
     fine_bb.append_instruction("calldataload", IRLiteral(4), ret=IRVariable("2"))
@@ -160,7 +120,7 @@ def test_multi_function():
         check_cv:
             %1 = callvalue
             %2 = param
-            jnz @no_value, @has_value, %1
+            jnz %1, @has_value, @no_value
         no_value:
             ret %2
         has_value:
@@ -218,7 +178,7 @@ def test_multi_function_and_data():
         check_cv:
             %1 = callvalue
             %2 = param
-            jnz @no_value, @has_value, %1
+            jnz %1, @has_value, @no_value
         no_value:
             ret %2
         has_value:

tests/unit/compiler/venom/test_simplify_cfg.py

@@ -1,48 +1,37 @@
+from tests.venom_utils import assert_ctx_eq, parse_venom
 from vyper.venom.analysis import IRAnalysesCache
-from vyper.venom.basicblock import IRBasicBlock, IRLabel, IRLiteral
-from vyper.venom.context import IRContext
 from vyper.venom.passes import SCCP, SimplifyCFGPass
 
 
 def test_phi_reduction_after_block_pruning():
-    ctx = IRContext()
-    fn = ctx.create_function("_global")
-
-    bb = fn.get_basic_block()
-
-    br1 = IRBasicBlock(IRLabel("then"), fn)
-    fn.append_basic_block(br1)
-    br2 = IRBasicBlock(IRLabel("else"), fn)
-    fn.append_basic_block(br2)
-
-    join = IRBasicBlock(IRLabel("join"), fn)
-    fn.append_basic_block(join)
-
-    true = IRLiteral(1)
-    bb.append_instruction("jnz", true, br1.label, br2.label)
-
-    op1 = br1.append_instruction("store", 1)
-    op2 = br2.append_instruction("store", 2)
-
-    br1.append_instruction("jmp", join.label)
-    br2.append_instruction("jmp", join.label)
-
-    join.append_instruction("phi", br1.label, op1, br2.label, op2)
-    join.append_instruction("stop")
-
-    ac = IRAnalysesCache(fn)
-    SCCP(ac, fn).run_pass()
-    SimplifyCFGPass(ac, fn).run_pass()
-
-    bbs = list(fn.get_basic_blocks())
-
-    assert len(bbs) == 1
-    final_bb = bbs[0]
-
-    inst0, inst1, inst2 = final_bb.instructions
-
-    assert inst0.opcode == "store"
-    assert inst0.operands == [IRLiteral(1)]
-    assert inst1.opcode == "store"
-    assert inst1.operands == [inst0.output]
-    assert inst2.opcode == "stop"
+    pre = """
+    function _global {
+        _global:
+            jnz 1, @then, @else
+        then:
+            %1 = 1
+            jmp @join
+        else:
+            %2 = 2
+            jmp @join
+        join:
+            %3 = phi @then, %1, @else, %2
+            stop
+    }
+    """
+    post = """
+    function _global {
+        _global:
+            %1 = 1
+            %3 = %1
+            stop
+    }
+    """
+    ctx1 = parse_venom(pre)
+    for fn in ctx1.functions.values():
+        ac = IRAnalysesCache(fn)
+        SCCP(ac, fn).run_pass()
+        SimplifyCFGPass(ac, fn).run_pass()
+
+    ctx2 = parse_venom(post)
+    assert_ctx_eq(ctx1, ctx2)

tests/venom_utils.py

@@ -0,0 +1,50 @@
+from vyper.venom.basicblock import IRBasicBlock, IRInstruction
+from vyper.venom.context import IRContext
+from vyper.venom.function import IRFunction
+from vyper.venom.parser import parse_venom
+
+
+def parse_from_basic_block(source: str, funcname="_global"):
+    """
+    Parse an IRContext from a basic block
+    """
+    source = f"function {funcname} {{\n{source}\n}}"
+    return parse_venom(source)
+
+
+def instructions_eq(i1: IRInstruction, i2: IRInstruction) -> bool:
+    return i1.output == i2.output and i1.opcode == i2.opcode and i1.operands == i2.operands
+
+
+def assert_bb_eq(bb1: IRBasicBlock, bb2: IRBasicBlock):
+    assert bb1.label.value == bb2.label.value
+    assert len(bb1.instructions) == len(bb2.instructions)
+    for i1, i2 in zip(bb1.instructions, bb2.instructions):
+        assert instructions_eq(i1, i2), f"[{i1}] != [{i2}]"
+
+
+def assert_fn_eq(fn1: IRFunction, fn2: IRFunction):
+    assert fn1.name.value == fn2.name.value
+    assert len(fn1._basic_block_dict) == len(fn2._basic_block_dict)
+
+    for name1, bb1 in fn1._basic_block_dict.items():
+        assert name1 in fn2._basic_block_dict
+        assert_bb_eq(bb1, fn2._basic_block_dict[name1])
+
+    # check function entry is the same
+    assert fn1.entry.label == fn2.entry.label
+
+
+def assert_ctx_eq(ctx1: IRContext, ctx2: IRContext):
+    assert ctx1.last_label == ctx2.last_label
+    assert len(ctx1.functions) == len(ctx2.functions)
+    for label1, fn1 in ctx1.functions.items():
+        assert label1 in ctx2.functions
+        assert_fn_eq(fn1, ctx2.functions[label1])
+
+    # check entry function is the same
+    assert next(iter(ctx1.functions.keys())) == next(iter(ctx2.functions.keys()))
+
+    assert len(ctx1.data_segment) == len(ctx2.data_segment)
+    for d1, d2 in zip(ctx1.data_segment, ctx2.data_segment):
+        assert instructions_eq(d1, d2), f"data: [{d1}] != [{d2}]"

vyper/venom/basicblock.py

@@ -381,7 +381,7 @@ def __repr__(self) -> str:
         opcode = f"{self.opcode} " if self.opcode != "store" else ""
         s += opcode
         operands = self.operands
-        if opcode not in ["jmp", "jnz", "invoke"]:
+        if opcode not in ("jmp", "jnz", "invoke"):
             operands = reversed(operands)  # type: ignore
         s += ", ".join(
             [(f"label %{op}" if isinstance(op, IRLabel) else str(op)) for op in operands]

vyper/venom/parser.py

@@ -19,8 +19,7 @@
     %import common.WS
     %import common.INT
 
-    # TODO: make data_section optional -- `function* data_section?`
-    start: function* data_section
+    start: function* data_section?
 
     # TODO: consider making entry block implicit, e.g.
     # `"{" instruction+ block* "}"`
@@ -32,7 +31,7 @@
 
     statement: instruction | assignment
     assignment: VAR_IDENT "=" expr
-    expr: instruction | CONST
+    expr: instruction | operand
     instruction: OPCODE operands_list?
 
     operands_list: operand ("," operand)*
@@ -45,7 +44,10 @@
     LABEL: "@" NAME
     NAME: (DIGIT|LETTER|"_")+
 
+    COMMENT: ";"  /[^\\n]/*
+
     %ignore WS
+    %ignore COMMENT
     """
 )
 
@@ -72,15 +74,24 @@ def _set_last_label(ctx: IRContext):
 def _ensure_terminated(bb):
     # Since "revert" is not considered terminal explicitly check for it to ensure basic
     # blocks are terminating
-    if not bb.is_terminated and any(inst.opcode == "revert" for inst in bb.instructions):
-        bb.append_instruction("stop")
+    if not bb.is_terminated:
+        if any(inst.opcode == "revert" for inst in bb.instructions):
+            bb.append_instruction("stop")
+        # TODO: raise error if still not terminated.
+
+
+class _DataSegment:
+    def __init__(self, instructions):
+        self.instructions = instructions
 
 
 class VenomTransformer(Transformer):
     def start(self, children) -> IRContext:
         ctx = IRContext()
-        funcs = children[:-1]
-        data_section = children[-1]
+        data_section = []
+        if isinstance(children[-1], _DataSegment):
+            data_section = children.pop().instructions
+        funcs = children
         for fn_name, blocks in funcs:
             fn = ctx.create_function(fn_name)
             fn._basic_block_dict.clear()
@@ -110,7 +121,7 @@ def statement(self, children):
         return children[0]
 
     def data_section(self, children):
-        return children
+        return _DataSegment(children)
 
     def block(self, children) -> tuple[str, list[IRInstruction]]:
         label, *instructions = children
@@ -121,7 +132,7 @@ def assignment(self, children) -> IRInstruction:
         if isinstance(value, IRInstruction):
             value.output = to
             return value
-        if isinstance(value, IRLiteral):
+        if isinstance(value, (IRLiteral, IRVariable)):
             return IRInstruction("store", [value], output=to)
         raise TypeError(f"Unexpected value {value} of type {type(value)}")
 
@@ -130,15 +141,18 @@ def expr(self, children):
 
     def instruction(self, children) -> IRInstruction:
         if len(children) == 1:
-            name = children[0]
+            opcode = children[0]
             operands = []
         else:
             assert len(children) == 2
-            name, operands = children
+            opcode, operands = children
 
         # reverse operands, venom internally represents top of stack
         # as rightmost operand
-        return IRInstruction(name, reversed(operands))
+        if opcode not in ("jmp", "jnz", "invoke", "phi"):
+            # special cases: operands with labels look better un-reversed
+            operands.reverse()
+        return IRInstruction(opcode, operands)
 
     def operands_list(self, children) -> list[IROperand]:
         return children