Make the handling of ControlOp targets more generic (#10920)

Add a `targets` property and a `set_target` method to `ControlOp` to
enable generic processing. Update CFG extraction, CFG cleanup, and
(most importantly) refcount insertion to operate generically.

This will make it easier to add new control ops, such as switch
statements.

Author: msullivan

mypyc/analysis/dataflow.py

@@ -53,13 +53,8 @@ def get_cfg(blocks: List[BasicBlock]) -> CFG:
         assert not any(isinstance(op, ControlOp) for op in block.ops[:-1]), (
             "Control-flow ops must be at the end of blocks")
 
-        last = block.ops[-1]
-        if isinstance(last, Branch):
-            succ = [last.true, last.false]
-        elif isinstance(last, Goto):
-            succ = [last.label]
-        else:
-            succ = []
+        succ = list(block.terminator.targets())
+        if not succ:
             exits.add(block)
 
         # Errors can occur anywhere inside a block, which means that
@@ -104,12 +99,8 @@ def cleanup_cfg(blocks: List[BasicBlock]) -> None:
     while changed:
         # First collapse any jumps to basic block that only contain a goto
         for block in blocks:
-            term = block.ops[-1]
-            if isinstance(term, Goto):
-                term.label = get_real_target(term.label)
-            elif isinstance(term, Branch):
-                term.true = get_real_target(term.true)
-                term.false = get_real_target(term.false)
+            for i, tgt in enumerate(block.terminator.targets()):
+                block.terminator.set_target(i, get_real_target(tgt))
 
         # Then delete any blocks that have no predecessors
         changed = False

mypyc/ir/ops.py

@@ -75,6 +75,12 @@ def terminated(self) -> bool:
         """
         return bool(self.ops) and isinstance(self.ops[-1], ControlOp)
 
+    @property
+    def terminator(self) -> 'ControlOp':
+        """The terminator operation of the block."""
+        assert bool(self.ops) and isinstance(self.ops[-1], ControlOp)
+        return self.ops[-1]
+
 
 # Never generates an exception
 ERR_NEVER: Final = 0
@@ -260,12 +266,15 @@ def accept(self, visitor: 'OpVisitor[T]') -> T:
 
 
 class ControlOp(Op):
-    """Control flow operation.
+    """Control flow operation."""
 
-    This is Basically just for class hierarchy organization.
+    def targets(self) -> Sequence[BasicBlock]:
+        """Get all basic block targets of the control operation."""
+        return ()
 
-    We could plausibly have a targets() method if we wanted.
-    """
+    def set_target(self, i: int, new: BasicBlock) -> None:
+        """Update a basic block target."""
+        raise AssertionError("Invalid set_target({}, {})".format(self, i))
 
 
 class Goto(ControlOp):
@@ -277,6 +286,13 @@ def __init__(self, label: BasicBlock, line: int = -1) -> None:
         super().__init__(line)
         self.label = label
 
+    def targets(self) -> Sequence[BasicBlock]:
+        return (self.label,)
+
+    def set_target(self, i: int, new: BasicBlock) -> None:
+        assert i == 0
+        self.label = new
+
     def __repr__(self) -> str:
         return '<Goto %s>' % self.label.label
 
@@ -327,6 +343,16 @@ def __init__(self,
         # If True, the condition is expected to be usually False (for optimization purposes)
         self.rare = rare
 
+    def targets(self) -> Sequence[BasicBlock]:
+        return (self.true, self.false)
+
+    def set_target(self, i: int, new: BasicBlock) -> None:
+        assert i == 0 or i == 1
+        if i == 0:
+            self.true = new
+        elif i == 1:
+            self.false = new
+
     def sources(self) -> List[Value]:
         return [self.value]
 

mypyc/transform/refcount.py

@@ -33,12 +33,13 @@
 from mypyc.ir.func_ir import FuncIR, all_values
 
 
-DecIncs = Tuple[Tuple[Tuple[Value, bool], ...], Tuple[Value, ...]]
+Decs = Tuple[Tuple[Value, bool], ...]
+Incs = Tuple[Value, ...]
 
-# A of basic blocks that decrement and increment specific values and
-# then jump to some target block. This lets us cut down on how much
-# code we generate in some circumstances.
-BlockCache = Dict[Tuple[BasicBlock, DecIncs], BasicBlock]
+# A cache of basic blocks that decrement and increment specific values
+# and then jump to some target block. This lets us cut down on how
+# much code we generate in some circumstances.
+BlockCache = Dict[Tuple[BasicBlock, Decs, Incs], BasicBlock]
 
 
 def insert_ref_count_opcodes(ir: FuncIR) -> None:
@@ -161,36 +162,25 @@ def f(a: int) -> None
     source_live_regs = pre_live[prev_key]
     source_borrowed = post_borrow[prev_key]
     source_defined = post_must_defined[prev_key]
-    if isinstance(block.ops[-1], Branch):
-        branch = block.ops[-1]
+
+    term = block.terminator
+    for i, target in enumerate(term.targets()):
         # HAX: After we've checked against an error value the value we must not touch the
         #      refcount since it will be a null pointer. The correct way to do this would be
         #      to perform data flow analysis on whether a value can be null (or is always
         #      null).
-        if branch.op == Branch.IS_ERROR:
-            omitted = {branch.value}
+        omitted: Iterable[Value]
+        if isinstance(term, Branch) and term.op == Branch.IS_ERROR and i == 0:
+            omitted = (term.value,)
         else:
-            omitted = set()
-        true_decincs = (
-            after_branch_decrefs(
-                branch.true, pre_live, source_defined,
-                source_borrowed, source_live_regs, ordering, omitted),
-            after_branch_increfs(
-                branch.true, pre_live, pre_borrow, source_borrowed, ordering))
-        branch.true = add_block(true_decincs, cache, blocks, branch.true)
-
-        false_decincs = (
-            after_branch_decrefs(
-                branch.false, pre_live, source_defined, source_borrowed, source_live_regs,
-                ordering),
-            after_branch_increfs(
-                branch.false, pre_live, pre_borrow, source_borrowed, ordering))
-        branch.false = add_block(false_decincs, cache, blocks, branch.false)
-    elif isinstance(block.ops[-1], Goto):
-        goto = block.ops[-1]
-        new_decincs = ((), after_branch_increfs(
-            goto.label, pre_live, pre_borrow, source_borrowed, ordering))
-        goto.label = add_block(new_decincs, cache, blocks, goto.label)
+            omitted = ()
+
+        decs = after_branch_decrefs(
+            target, pre_live, source_defined,
+            source_borrowed, source_live_regs, ordering, omitted)
+        incs = after_branch_increfs(
+            target, pre_live, pre_borrow, source_borrowed, ordering)
+        term.set_target(i, add_block(decs, incs, cache, blocks, target))
 
 
 def after_branch_decrefs(label: BasicBlock,
@@ -199,7 +189,7 @@ def after_branch_decrefs(label: BasicBlock,
                          source_borrowed: Set[Value],
                          source_live_regs: Set[Value],
                          ordering: Dict[Value, int],
-                         omitted: Iterable[Value] = ()) -> Tuple[Tuple[Value, bool], ...]:
+                         omitted: Iterable[Value]) -> Tuple[Tuple[Value, bool], ...]:
     target_pre_live = pre_live[label, 0]
     decref = source_live_regs - target_pre_live - source_borrowed
     if decref:
@@ -224,22 +214,21 @@ def after_branch_increfs(label: BasicBlock,
     return ()
 
 
-def add_block(decincs: DecIncs, cache: BlockCache,
+def add_block(decs: Decs, incs: Incs, cache: BlockCache,
               blocks: List[BasicBlock], label: BasicBlock) -> BasicBlock:
-    decs, incs = decincs
     if not decs and not incs:
         return label
 
     # TODO: be able to share *partial* results
-    if (label, decincs) in cache:
-        return cache[label, decincs]
+    if (label, decs, incs) in cache:
+        return cache[label, decs, incs]
 
     block = BasicBlock()
     blocks.append(block)
     block.ops.extend(DecRef(reg, is_xdec=xdec) for reg, xdec in decs)
     block.ops.extend(IncRef(reg) for reg in incs)
     block.ops.append(Goto(label))
-    cache[label, decincs] = block
+    cache[label, decs, incs] = block
     return block