refactor[venom]: make venom repr parseable

tests/functional/venom/test_venom_repr.py

@@ -0,0 +1,36 @@
+import glob
+
+import pytest
+
+from tests.venom_utils import assert_ctx_eq, parse_venom
+from vyper.compiler import compile_code
+from vyper.compiler.settings import OptimizationLevel
+from vyper.venom.context import IRContext
+
+"""
+Check that venom text format round-trips through parser
+"""
+
+
+def get_example_vy_filenames():
+    return glob.glob("**/*.vy", root_dir="examples/", recursive=True)
+
+
+@pytest.mark.parametrize("vy_filename", get_example_vy_filenames())
+def test_round_trip(vy_filename, optimize, request):
+    if optimize == OptimizationLevel.CODESIZE:
+        # codesize optimization issues things like `db b"\x12\x34"` which we
+        # don't handle.
+        request.node.add_marker(pytest.mark.xfail(strict=False, reason="unimplemented in parser"))
+
+    path = f"examples/{vy_filename}"
+    with open(path) as f:
+        vyper_source = f.read()
+
+    out = compile_code(vyper_source, output_formats=["bb_runtime"])
+    bb_runtime = out["bb_runtime"]
+    venom_code = IRContext.__repr__(bb_runtime)
+
+    ctx = parse_venom(venom_code)
+
+    assert_ctx_eq(bb_runtime, ctx)

tests/venom_utils.py

@@ -36,7 +36,6 @@ def assert_fn_eq(fn1: IRFunction, fn2: IRFunction):
 
 
 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

vyper/venom/basicblock.py

@@ -1,3 +1,5 @@
+import json
+import re
 from typing import TYPE_CHECKING, Any, Iterator, Optional, Union
 
 import vyper.venom.effects as effects
@@ -105,7 +107,7 @@ def __init__(self, line_no: int, src: str) -> None:
 
     def __repr__(self) -> str:
         src = self.src if self.src else ""
-        return f"\t# line {self.line_no}: {src}".expandtabs(20)
+        return f"\t; line {self.line_no}: {src}".expandtabs(20)
 
 
 class IROperand:
@@ -183,10 +185,19 @@ class IRLabel(IROperand):
     value: str
 
     def __init__(self, value: str, is_symbol: bool = False) -> None:
-        assert isinstance(value, str), "value must be an str"
+        assert isinstance(value, str), f"not a str: {value} ({type(value)})"
+        assert len(value) > 0
         self.value = value
         self.is_symbol = is_symbol
 
+    _IS_IDENTIFIER = re.compile("[0-9a-zA-Z_]*")
+
+    def __repr__(self):
+        if self.__class__._IS_IDENTIFIER.fullmatch(self.value):
+            return self.value
+
+        return json.dumps(self.value)  # escape it
+
 
 class IRInstruction:
     """
@@ -360,35 +371,20 @@ def get_ast_source(self) -> Optional[IRnode]:
                 return inst.ast_source
         return self.parent.parent.ast_source
 
-    def str_short(self) -> str:
-        s = ""
-        if self.output:
-            s += f"{self.output} = "
-        opcode = f"{self.opcode} " if self.opcode != "store" else ""
-        s += opcode
-        operands = self.operands
-        if opcode not in ["jmp", "jnz", "invoke"]:
-            operands = list(reversed(operands))
-        s += ", ".join(
-            [(f"label %{op}" if isinstance(op, IRLabel) else str(op)) for op in operands]
-        )
-        return s
-
     def __repr__(self) -> str:
         s = ""
         if self.output:
             s += f"{self.output} = "
         opcode = f"{self.opcode} " if self.opcode != "store" else ""
         s += opcode
         operands = self.operands
-        if opcode not in ("jmp", "jnz", "invoke"):
+        if self.opcode not in ("jmp", "jnz", "invoke", "phi"):
+            # TODO: for invoke, maybe reverse the non-label instructions
             operands = reversed(operands)  # type: ignore
-        s += ", ".join(
-            [(f"label %{op}" if isinstance(op, IRLabel) else str(op)) for op in operands]
-        )
+        s += ", ".join([(f"@{op}" if isinstance(op, IRLabel) else str(op)) for op in operands])
 
         if self.annotation:
-            s += f" <{self.annotation}>"
+            s += f" ; {self.annotation}"
 
         return f"{s: <30}"
 
@@ -644,10 +640,8 @@ def copy(self):
         return bb
 
     def __repr__(self) -> str:
-        s = (
-            f"{repr(self.label)}:  IN={[bb.label for bb in self.cfg_in]}"
-            f" OUT={[bb.label for bb in self.cfg_out]} => {self.out_vars}\n"
-        )
+        s = f"{self.label}:  ; IN={[bb.label for bb in self.cfg_in]}"
+        s += f" OUT={[bb.label for bb in self.cfg_out]} => {self.out_vars}\n"
         for instruction in self.instructions:
-            s += f"    {str(instruction).strip()}\n"
+            s += f"  {str(instruction).strip()}\n"
         return s

vyper/venom/context.py

@@ -62,14 +62,14 @@ def as_graph(self) -> str:
         return "\n".join(s)
 
     def __repr__(self) -> str:
-        s = ["IRContext:"]
+        s = []
         for fn in self.functions.values():
             s.append(fn.__repr__())
             s.append("\n")
 
         if len(self.data_segment) > 0:
-            s.append("\nData segment:")
+            s.append("\n[data]")
             for inst in self.data_segment:
-                s.append(f"{inst}")
+                s.append(f"  {inst}")
 
         return "\n".join(s)

vyper/venom/function.py

@@ -1,3 +1,4 @@
+import textwrap
 from typing import Iterator, Optional
 
 from vyper.codegen.ir_node import IRnode
@@ -41,7 +42,7 @@ def append_basic_block(self, bb: IRBasicBlock):
         Append basic block to function.
         """
         assert isinstance(bb, IRBasicBlock), bb
-        assert bb.label.name not in self._basic_block_dict
+        assert bb.label.name not in self._basic_block_dict, bb.label
         self._basic_block_dict[bb.label.name] = bb
 
     def remove_basic_block(self, bb: IRBasicBlock):
@@ -222,7 +223,9 @@ def _make_label(bb):
         return "\n".join(ret)
 
     def __repr__(self) -> str:
-        str = f"IRFunction: {self.name}\n"
+        ret = f"function {self.name} {{\n"
         for bb in self.get_basic_blocks():
-            str += f"{bb}\n"
-        return str.strip()
+            bb_str = textwrap.indent(str(bb), "  ")
+            ret += f"{bb_str}\n"
+        ret = ret.strip() + "\n}"
+        return ret.strip()

vyper/venom/ir_node_to_venom.py

@@ -369,10 +369,7 @@ def _convert_ir_bb(fn, ir, symbols):
         label = IRLabel(ir.args[0].value)
         ctx.append_data("dbname", [label])
         for c in ir.args[1:]:
-            if isinstance(c, int):
-                assert 0 <= c <= 255, "data with invalid size"
-                ctx.append_data("db", [c])  # type: ignore
-            elif isinstance(c.value, bytes):
+            if isinstance(c.value, bytes):
                 ctx.append_data("db", [c.value])  # type: ignore
             elif isinstance(c, IRnode):
                 data = _convert_ir_bb(fn, c, symbols)

vyper/venom/parser.py

@@ -1,3 +1,5 @@
+import json
+
 from lark import Lark, Transformer
 
 from vyper.venom.basicblock import (
@@ -11,13 +13,14 @@
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
 
-VENOM_PARSER = Lark(
-    """
+VENOM_GRAMMAR = """
     %import common.CNAME
     %import common.DIGIT
     %import common.LETTER
     %import common.WS
     %import common.INT
+    %import common.SIGNED_INT
+    %import common.ESCAPED_STRING
 
     # Allow multiple comment styles
     COMMENT: ";" /[^\\n]*/ | "//" /[^\\n]*/ | "#" /[^\\n]*/
@@ -26,13 +29,13 @@
 
     # TODO: consider making entry block implicit, e.g.
     # `"{" instruction+ block* "}"`
-    function: "function" NAME "{" block* "}"
+    function: "function" LABEL_IDENT "{" block* "}"
 
     data_section: "[data]" instruction*
 
-    block: NAME ":" statement*
+    block: LABEL_IDENT ":" "\\n" statement*
 
-    statement: instruction | assignment
+    statement: (instruction | assignment) "\\n"
     assignment: VAR_IDENT "=" expr
     expr: instruction | operand
     instruction: OPCODE operands_list?
@@ -41,16 +44,22 @@
 
     operand: VAR_IDENT | CONST | LABEL
 
-    CONST: INT
+    CONST: SIGNED_INT
     OPCODE: CNAME
-    VAR_IDENT: "%" NAME
-    LABEL: "@" NAME
+    VAR_IDENT: "%" NAME (":" INT)?
+
+    # handy for identifier to be an escaped string sometimes
+    # (especially for machine-generated labels)
+    LABEL_IDENT: (NAME | ESCAPED_STRING)
+    LABEL: "@" LABEL_IDENT
+
     NAME: (DIGIT|LETTER|"_")+
 
     %ignore WS
     %ignore COMMENT
     """
-)
+
+VENOM_PARSER = Lark(VENOM_GRAMMAR)
 
 
 def _set_last_var(fn: IRFunction):
@@ -100,7 +109,7 @@ def start(self, children) -> IRContext:
             fn._basic_block_dict.clear()
 
             for block_name, instructions in blocks:
-                bb = IRBasicBlock(IRLabel(block_name), fn)
+                bb = IRBasicBlock(IRLabel(block_name, True), fn)
                 fn.append_basic_block(bb)
 
                 for instruction in instructions:
@@ -166,16 +175,26 @@ def operand(self, children) -> IROperand:
     def OPCODE(self, token):
         return token.value
 
+    def LABEL_IDENT(self, label) -> str:
+        if label.startswith('"'):
+            # unescape the escaped string
+            label = json.loads(label)
+        return label
+
     def LABEL(self, label) -> IRLabel:
-        return IRLabel(label[1:])
+        label = label[1:]
+        if label.startswith('"'):
+            # unescape the escaped string
+            label = json.loads(label)
+        return IRLabel(label, True)
 
     def VAR_IDENT(self, var_ident) -> IRVariable:
         parts = var_ident[1:].split(":", maxsplit=1)
         assert 1 <= len(parts) <= 2
         varname = parts[0]
         version = None
         if len(parts) > 1:
-            version = int(parts[1])
+            version = parts[1]
         return IRVariable(varname, version=version)
 
     def CONST(self, val) -> IRLiteral:

vyper/venom/venom_to_assembly.py

@@ -185,8 +185,8 @@ def generate_evm(self, no_optimize: bool = False) -> list[str]:
             data_segments: dict = dict()
             for inst in ctx.data_segment:
                 if inst.opcode == "dbname":
-                    label = inst.operands[0].value
-                    data_segments[label] = [DataHeader(f"_sym_{label}")]
+                    label = inst.operands[0]
+                    data_segments[label] = [DataHeader(f"_sym_{label.value}")]
                 elif inst.opcode == "db":
                     data = inst.operands[0]
                     if isinstance(data, IRLabel):
@@ -293,7 +293,7 @@ def _generate_evm_for_basicblock_r(
         asm = []
 
         # assembly entry point into the block
-        asm.append(f"_sym_{basicblock.label}")
+        asm.append(f"_sym_{basicblock.label.value}")
         asm.append("JUMPDEST")
 
         if len(basicblock.cfg_in) == 1: