review: polish and fix cfg normalization

tests/compiler/venom/test_multi_entry_block.py

@@ -1,41 +1,95 @@
+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
 
 
-def test_multi_entry_block():
+def test_multi_entry_block_1():
     ctx = IRFunction()
 
     finish_label = IRLabel("finish")
     target_label = IRLabel("target")
     block_1_label = IRLabel("block_1", ctx)
 
     op = ctx.append_instruction("store", [IRLiteral(10)])
-    sum = ctx.append_instruction("add", [op, op])
-    ctx.append_instruction("jnz", [sum, finish_label, block_1_label], False)
+    acc = ctx.append_instruction("add", [op, op])
+    ctx.append_instruction("jnz", [acc, finish_label, block_1_label], False)
 
-    target_bb = IRBasicBlock(block_1_label, ctx)
+    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)
+
+    target_bb = IRBasicBlock(target_label, ctx)
     ctx.append_basic_block(target_bb)
-    sum = ctx.append_instruction("add", [sum, op])
+    ctx.append_instruction("mul", [acc, acc])
+    ctx.append_instruction("jmp", [finish_label], False)
+
+    finish_bb = IRBasicBlock(finish_label, ctx)
+    ctx.append_basic_block(finish_bb)
+    ctx.append_instruction("stop", [], False)
+
+    calculate_cfg(ctx)
+    assert not ctx.normalized, "CFG should not be normalized"
+
+    NormalizationPass.run_pass(ctx)
+
+    calculate_cfg(ctx)
+    assert ctx.normalized, "CFG should be normalized"
+
+    finish_bb = ctx.get_basic_block(finish_label.value)
+    cfg_in = list(finish_bb.cfg_in.keys())
+    assert cfg_in[0].label.value == "target", "Should contain target"
+    assert cfg_in[1].label.value == "finish_split_global", "Should contain finish_split_global"
+    assert cfg_in[2].label.value == "finish_split_block_1", "Should contain finish_split_block_1"
+
+
+# more complicated one
+def test_multi_entry_block_2():
+    ctx = IRFunction()
+
+    finish_label = IRLabel("finish")
+    target_label = IRLabel("target")
+    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)
+
+    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)
+
+    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", [sum, op], False)
-    ctx.append_instruction("jnz", [sum, finish_label, target_label], False)
+    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)
 
     target_bb = IRBasicBlock(target_label, ctx)
     ctx.append_basic_block(target_bb)
-    _ = ctx.append_instruction("mul", [sum, sum])
+    ctx.append_instruction("mul", [acc, acc])
     ctx.append_instruction("jmp", [finish_label], False)
 
     finish_bb = IRBasicBlock(finish_label, ctx)
     ctx.append_basic_block(finish_bb)
     ctx.append_instruction("stop", [], False)
 
-    assert ctx.cfg_dirty is True, "CFG should be dirty"
+    calculate_cfg(ctx)
+    assert not ctx.normalized, "CFG should not be normalized"
 
     NormalizationPass.run_pass(ctx)
 
-    assert ctx.cfg_dirty is False, "CFG should be clean"
-    assert ctx.normalized is True, "CFG should be normalized"
+    calculate_cfg(ctx)
+    assert ctx.normalized, "CFG should be normalized"
 
     finish_bb = ctx.get_basic_block(finish_label.value)
     cfg_in = list(finish_bb.cfg_in.keys())

vyper/venom/basicblock.py

@@ -281,7 +281,7 @@ def add_cfg_out(self, bb: "IRBasicBlock") -> None:
         # malformed: jnz condition label1 label1
         # (we could handle but it makes a lot of code easier
         # if we have this assumption)
-        assert bb not in in_bb.cfg_out
+        assert bb not in self.cfg_out
         self.cfg_out.add(bb)
 
     @property

vyper/venom/function.py

@@ -43,6 +43,8 @@ def append_basic_block(self, bb: IRBasicBlock) -> IRBasicBlock:
         assert isinstance(bb, IRBasicBlock), f"append_basic_block takes IRBasicBlock, got '{bb}'"
         self.basic_blocks.append(bb)
 
+        # TODO add sanity check somewhere that basic blocks have unique labels
+
         return self.basic_blocks[-1]
 
     def get_basic_block(self, label: Optional[str] = None) -> IRBasicBlock:
@@ -90,8 +92,6 @@ def remove_unreachable_blocks(self) -> int:
         new_basic_blocks = []
         for bb in self.basic_blocks:
             if not bb.is_reachable and bb.label.value != "global":
-                for bb2 in bb.cfg_out:
-                    bb2.remove_cfg_in(bb)
                 removed += 1
             else:
                 new_basic_blocks.append(bb)
@@ -118,24 +118,21 @@ def append_data(self, opcode: str, args: list[IROperand]) -> None:
     @property
     def normalized(self) -> bool:
         """
-        Check if function is normalized.
+        Check if function is normalized. A function is normalized if in the
+        CFG, no basic block simultaneously has multiple inputs and outputs.
+        That is, a basic block can be jumped to *from* multiple blocks, or it
+        can jump *to* multiple blocks, but it cannot simultaneously do both.
+        Having a normalized CFG makes calculation of stack layout easier when
+        emitting assembly.
         """
         for bb in self.basic_blocks:
-            # Ignore if there are no multiple predecessors
             if len(bb.cfg_in) <= 1:
                 continue
 
-            # Check if there is a conditional jump at the end
-            # of one of the predecessors
             for in_bb in bb.cfg_in:
-                jump_inst = in_bb.instructions[-1]
-                if jump_inst.opcode != "jnz":
-                    continue
-
-                # The function is not normalized
-                return False
+                if len(in_bb.cfg_out) > 1:
+                    return False
 
-        # The function is normalized
         return True
 
     def copy(self):

vyper/venom/passes/normalization.py

@@ -1,5 +1,6 @@
-from vyper.venom.analysis import DFG, calculate_cfg
-from vyper.venom.basicblock import IRBasicBlock, IRLabel
+from vyper.exceptions import CompilerPanic
+from vyper.venom.analysis import calculate_cfg
+from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRLabel
 from vyper.venom.function import IRFunction
 from vyper.venom.passes.base_pass import IRPass
 
@@ -15,28 +16,47 @@ class NormalizationPass(IRPass):
 
     def _split_basic_block(self, bb: IRBasicBlock) -> None:
         ctx = self.ctx
-        label_base = bb.label.value
 
         # Iterate over the predecessors of the basic block
         for in_bb in bb.cfg_in:
-            jump_inst = in_bb.instructions[-1]
-            # We are only splitting on contitional jumps
-            if jump_inst.opcode != "jnz":
+            # We are only splitting on conditional jumps
+            if len(in_bb.cfg_out) < 2:
                 continue
 
+            # sanity checks. only jnz can product len(cfg_out) > 1
+            jump_inst = in_bb.instructions[-1]
+            assert jump_inst.opcode == "jnz"
+            # jnz produces cfg_out with length 2
+            assert len(in_bb.cfg_out) == 2
+            assert bb in in_bb.cfg_out
+
+            # find which edge of the jnz targets this block
+            # jnz condition label1 label2
+            jump_inst.operands[1] == jump_inst.operands[2]
+            for i, op in enumerate(jump_inst.operands):
+                if op == bb.label:
+                    edge = i
+                    break
+            else:
+                # none of the edges points to this bb
+                raise CompilerPanic("bad CFG")
+            assert edge in (1, 2)  # the arguments which can be labels
+
             # Create an intermediary basic block and append it
-            split_bb = IRBasicBlock(IRLabel(label_base + "_split_" + in_bb.label.value), ctx)
+            source = in_bb.label.value
+            target = bb.label.value
+            split_bb = IRBasicBlock(IRLabel(f"{target}_split_{source}"), ctx)
+            split_bb.append_instruction(IRInstruction("jmp", [bb.label]))
+
             ctx.append_basic_block(split_bb)
-            ctx.append_instruction("jmp", [bb.label], False)
 
             # Redirect the original conditional jump to the intermediary basic block
-            jump_inst.operands[1] = split_bb.label
+            jump_inst.operands[edge] = split_bb.label
 
             self.changes += 1
 
     def _run_pass(self, ctx: IRFunction) -> int:
         self.ctx = ctx
-        self.dfg = DFG.build_dfg(ctx)
         self.changes = 0
 
         # Ensure that the CFG is up to date
@@ -47,9 +67,10 @@ def _run_pass(self, ctx: IRFunction) -> int:
                 self._split_basic_block(bb)
 
         # Recalculate control flow graph
+        # (perf: could do this only when self.changes > 0, but be paranoid)
         calculate_cfg(ctx)
 
         # Sanity check
-        assert ctx.normalized is True, "Normalization pass failed"
+        assert ctx.normalized, "Normalization pass failed"
 
         return self.changes

vyper/venom/venom_to_assembly.py

@@ -101,6 +101,8 @@ def generate_evm(self, no_optimize: bool = False) -> list[str]:
         calculate_cfg(self.ctx)
         calculate_liveness(self.ctx)
 
+        assert self.ctx.normalized, "Non-normalized CFG!"
+
         self._generate_evm_for_basicblock_r(asm, self.ctx.basic_blocks[0], stack)
 
         # Append postambles