refactor: improve IRBasicBlock builder API

clean up `append_instruction` api so it does a bit of magic on its
arguments and figures out whether or not to allocate a stack variable.

remove `append_instruction()` from IRFunction - automatically appending
to the last basic block could be a bit error prone depending on which
order basic blocks are added to the CFG.

---------

Co-authored-by: Charles Cooper <[email protected]>

tests/unit/compiler/venom/test_duplicate_operands.py

@@ -1,6 +1,5 @@
 from vyper.compiler.settings import OptimizationLevel
 from vyper.venom import generate_assembly_experimental
-from vyper.venom.basicblock import IRLiteral
 from vyper.venom.function import IRFunction
 
 
@@ -17,11 +16,11 @@ def test_duplicate_operands():
     Should compile to: [PUSH1, 10, DUP1, DUP1, DUP1, ADD, MUL, STOP]
     """
     ctx = IRFunction()
-
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    sum = ctx.append_instruction("add", [op, op])
-    ctx.append_instruction("mul", [sum, op])
-    ctx.append_instruction("stop", [], False)
+    bb = ctx.get_basic_block()
+    op = bb.append_instruction("store", 10)
+    sum = bb.append_instruction("add", op, op)
+    bb.append_instruction("mul", sum, op)
+    bb.append_instruction("stop")
 
     asm = generate_assembly_experimental(ctx, OptimizationLevel.CODESIZE)
 

tests/unit/compiler/venom/test_multi_entry_block.py

@@ -1,5 +1,4 @@
 from vyper.venom.analysis import calculate_cfg
-from vyper.venom.basicblock import IRLiteral
 from vyper.venom.function import IRBasicBlock, IRFunction, IRLabel
 from vyper.venom.passes.normalization import NormalizationPass
 
@@ -11,25 +10,26 @@ def test_multi_entry_block_1():
     target_label = IRLabel("target")
     block_1_label = IRLabel("block_1", ctx)
 
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    acc = ctx.append_instruction("add", [op, op])
-    ctx.append_instruction("jnz", [acc, finish_label, block_1_label], False)
+    bb = ctx.get_basic_block()
+    op = bb.append_instruction("store", 10)
+    acc = bb.append_instruction("add", op, op)
+    bb.append_instruction("jnz", acc, finish_label, block_1_label)
 
     block_1 = IRBasicBlock(block_1_label, ctx)
     ctx.append_basic_block(block_1)
-    acc = ctx.append_instruction("add", [acc, op])
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    ctx.append_instruction("mstore", [acc, op], False)
-    ctx.append_instruction("jnz", [acc, finish_label, target_label], False)
+    acc = block_1.append_instruction("add", acc, op)
+    op = block_1.append_instruction("store", 10)
+    block_1.append_instruction("mstore", acc, op)
+    block_1.append_instruction("jnz", acc, finish_label, target_label)
 
     target_bb = IRBasicBlock(target_label, ctx)
     ctx.append_basic_block(target_bb)
-    ctx.append_instruction("mul", [acc, acc])
-    ctx.append_instruction("jmp", [finish_label], False)
+    target_bb.append_instruction("mul", acc, acc)
+    target_bb.append_instruction("jmp", finish_label)
 
     finish_bb = IRBasicBlock(finish_label, ctx)
     ctx.append_basic_block(finish_bb)
-    ctx.append_instruction("stop", [], False)
+    finish_bb.append_instruction("stop")
 
     calculate_cfg(ctx)
     assert not ctx.normalized, "CFG should not be normalized"
@@ -54,33 +54,34 @@ def test_multi_entry_block_2():
     block_1_label = IRLabel("block_1", ctx)
     block_2_label = IRLabel("block_2", ctx)
 
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    acc = ctx.append_instruction("add", [op, op])
-    ctx.append_instruction("jnz", [acc, finish_label, block_1_label], False)
+    bb = ctx.get_basic_block()
+    op = bb.append_instruction("store", 10)
+    acc = bb.append_instruction("add", op, op)
+    bb.append_instruction("jnz", acc, finish_label, block_1_label)
 
     block_1 = IRBasicBlock(block_1_label, ctx)
     ctx.append_basic_block(block_1)
-    acc = ctx.append_instruction("add", [acc, op])
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    ctx.append_instruction("mstore", [acc, op], False)
-    ctx.append_instruction("jnz", [acc, target_label, finish_label], False)
+    acc = block_1.append_instruction("add", acc, op)
+    op = block_1.append_instruction("store", 10)
+    block_1.append_instruction("mstore", acc, op)
+    block_1.append_instruction("jnz", acc, target_label, finish_label)
 
     block_2 = IRBasicBlock(block_2_label, ctx)
     ctx.append_basic_block(block_2)
-    acc = ctx.append_instruction("add", [acc, op])
-    op = ctx.append_instruction("store", [IRLiteral(10)])
-    ctx.append_instruction("mstore", [acc, op], False)
-    # switch the order of the labels, for fun
-    ctx.append_instruction("jnz", [acc, finish_label, target_label], False)
+    acc = block_2.append_instruction("add", acc, op)
+    op = block_2.append_instruction("store", 10)
+    block_2.append_instruction("mstore", acc, op)
+    # switch the order of the labels, for fun and profit
+    block_2.append_instruction("jnz", acc, finish_label, target_label)
 
     target_bb = IRBasicBlock(target_label, ctx)
     ctx.append_basic_block(target_bb)
-    ctx.append_instruction("mul", [acc, acc])
-    ctx.append_instruction("jmp", [finish_label], False)
+    target_bb.append_instruction("mul", acc, acc)
+    target_bb.append_instruction("jmp", finish_label)
 
     finish_bb = IRBasicBlock(finish_label, ctx)
     ctx.append_basic_block(finish_bb)
-    ctx.append_instruction("stop", [], False)
+    finish_bb.append_instruction("stop")
 
     calculate_cfg(ctx)
     assert not ctx.normalized, "CFG should not be normalized"

vyper/venom/analysis.py

@@ -2,7 +2,7 @@
 from vyper.utils import OrderedSet
 from vyper.venom.basicblock import (
     BB_TERMINATORS,
-    CFG_ALTERING_OPS,
+    CFG_ALTERING_INSTRUCTIONS,
     IRBasicBlock,
     IRInstruction,
     IRVariable,
@@ -55,7 +55,7 @@ def calculate_cfg(ctx: IRFunction) -> None:
         assert last_inst.opcode in BB_TERMINATORS, f"Last instruction should be a terminator {bb}"
 
         for inst in bb.instructions:
-            if inst.opcode in CFG_ALTERING_OPS:
+            if inst.opcode in CFG_ALTERING_INSTRUCTIONS:
                 ops = inst.get_label_operands()
                 for op in ops:
                     ctx.get_basic_block(op.value).add_cfg_in(bb)

vyper/venom/basicblock.py

@@ -1,5 +1,5 @@
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, Iterator, Optional
+from typing import TYPE_CHECKING, Any, Iterator, Optional, Union
 
 from vyper.utils import OrderedSet
 
@@ -31,17 +31,40 @@
     ]
 )
 
-CFG_ALTERING_OPS = frozenset(["jmp", "jnz", "call", "staticcall", "invoke", "deploy"])
+NO_OUTPUT_INSTRUCTIONS = frozenset(
+    [
+        "deploy",
+        "mstore",
+        "sstore",
+        "dstore",
+        "istore",
+        "dloadbytes",
+        "calldatacopy",
+        "codecopy",
+        "return",
+        "ret",
+        "revert",
+        "assert",
+        "selfdestruct",
+        "stop",
+        "invalid",
+        "invoke",
+        "jmp",
+        "jnz",
+        "log",
+    ]
+)
 
+CFG_ALTERING_INSTRUCTIONS = frozenset(["jmp", "jnz", "call", "staticcall", "invoke", "deploy"])
 
 if TYPE_CHECKING:
     from vyper.venom.function import IRFunction
 
 
 class IRDebugInfo:
     """
-    IRDebugInfo represents debug information in IR, used to annotate IR instructions
-    with source code information when printing IR.
+    IRDebugInfo represents debug information in IR, used to annotate IR
+    instructions with source code information when printing IR.
     """
 
     line_no: int
@@ -83,7 +106,7 @@ class IRLiteral(IRValue):
     value: int
 
     def __init__(self, value: int) -> None:
-        assert isinstance(value, str) or isinstance(value, int), "value must be an int"
+        assert isinstance(value, int), "value must be an int"
         self.value = value
 
     def __repr__(self) -> str:
@@ -170,7 +193,7 @@ def __init__(
         assert isinstance(operands, list | Iterator), "operands must be a list"
         self.opcode = opcode
         self.volatile = opcode in VOLATILE_INSTRUCTIONS
-        self.operands = [op for op in operands]  # in case we get an iterator
+        self.operands = list(operands)  # in case we get an iterator
         self.output = output
         self.liveness = OrderedSet()
         self.dup_requirements = OrderedSet()
@@ -233,6 +256,14 @@ def __repr__(self) -> str:
         return s
 
 
+def _ir_operand_from_value(val: Any) -> IROperand:
+    if isinstance(val, IROperand):
+        return val
+
+    assert isinstance(val, int)
+    return IRLiteral(val)
+
+
 class IRBasicBlock:
     """
     IRBasicBlock represents a basic block in IR. Each basic block has a label and
@@ -243,8 +274,8 @@ class IRBasicBlock:
         %2 = mul %1, 2
     is represented as:
         bb = IRBasicBlock("bb", function)
-        bb.append_instruction(IRInstruction("add", ["%0", "1"], "%1"))
-        bb.append_instruction(IRInstruction("mul", ["%1", "2"], "%2"))
+        r1 = bb.append_instruction("add", "%0", "1")
+        r2 = bb.append_instruction("mul", r1, "2")
 
     The label of a basic block is used to refer to it from other basic blocks
     in order to branch to it.
@@ -296,10 +327,41 @@ def remove_cfg_out(self, bb: "IRBasicBlock") -> None:
     def is_reachable(self) -> bool:
         return len(self.cfg_in) > 0
 
-    def append_instruction(self, instruction: IRInstruction) -> None:
-        assert isinstance(instruction, IRInstruction), "instruction must be an IRInstruction"
-        instruction.parent = self
-        self.instructions.append(instruction)
+    def append_instruction(self, opcode: str, *args: Union[IROperand, int]) -> Optional[IRVariable]:
+        """
+        Append an instruction to the basic block
+
+        Returns the output variable if the instruction supports one
+        """
+        ret = self.parent.get_next_variable() if opcode not in NO_OUTPUT_INSTRUCTIONS else None
+
+        # Wrap raw integers in IRLiterals
+        inst_args = [_ir_operand_from_value(arg) for arg in args]
+
+        inst = IRInstruction(opcode, inst_args, ret)
+        inst.parent = self
+        self.instructions.append(inst)
+        return ret
+
+    def append_invoke_instruction(
+        self, args: list[IROperand | int], returns: bool
+    ) -> Optional[IRVariable]:
+        """
+        Append an instruction to the basic block
+
+        Returns the output variable if the instruction supports one
+        """
+        ret = None
+        if returns:
+            ret = self.parent.get_next_variable()
+
+        # Wrap raw integers in IRLiterals
+        inst_args = [_ir_operand_from_value(arg) for arg in args]
+
+        inst = IRInstruction("invoke", inst_args, ret)
+        inst.parent = self
+        self.instructions.append(inst)
+        return ret
 
     def insert_instruction(self, instruction: IRInstruction, index: int) -> None:
         assert isinstance(instruction, IRInstruction), "instruction must be an IRInstruction"

vyper/venom/function.py

@@ -98,17 +98,6 @@ def remove_unreachable_blocks(self) -> int:
         self.basic_blocks = new_basic_blocks
         return removed
 
-    def append_instruction(
-        self, opcode: str, args: list[IROperand], do_ret: bool = True
-    ) -> Optional[IRVariable]:
-        """
-        Append instruction to last basic block.
-        """
-        ret = self.get_next_variable() if do_ret else None
-        inst = IRInstruction(opcode, args, ret)  # type: ignore
-        self.get_basic_block().append_instruction(inst)
-        return ret
-
     def append_data(self, opcode: str, args: list[IROperand]) -> None:
         """
         Append data