More work on CFA for try-finally in the presence of loops.

llpython/byte_control.py

@@ -1,5 +1,6 @@
 #! /usr/bin/env python
 # ______________________________________________________________________
+
 from __future__ import absolute_import
 import opcode
 from . import opcode_util
@@ -13,9 +14,10 @@ from .control_flow import ControlFlowGraph
 
 # The following opcodes branch based on the control (a.k.a. frame)
 # stack:
-RETURN_VALUE, CONTINUE_LOOP, BREAK_LOOP, END_FINALLY = (
+RETURN_VALUE, CONTINUE_LOOP, BREAK_LOOP, END_FINALLY, RAISE_VARARGS = (
     opcode.opmap[opname] for opname in (
-        'RETURN_VALUE', 'CONTINUE_LOOP', 'BREAK_LOOP', 'END_FINALLY'))
+        'RETURN_VALUE', 'CONTINUE_LOOP', 'BREAK_LOOP', 'END_FINALLY',
+        'RAISE_VARARGS'))
 
 # The following opcodes push a new frame on the control stack:
 SETUP_EXCEPT, SETUP_FINALLY, SETUP_LOOP, SETUP_WITH = (
@@ -24,7 +26,11 @@ SETUP_EXCEPT, SETUP_FINALLY, SETUP_LOOP, SETUP_WITH = (
 
 WHY_NOT = 1
 WHY_EXCEPTION = WHY_NOT << 1
-WHY_RERAISE = WHY_EXCEPTION << 1
+WHY_RERAISE = WHY_EXCEPTION << 1 # We don't worry about this code
+                                 # during CFA, since its primary use
+                                 # is to log traceback information;
+                                 # WHY_RERAISE's bytecode control flow
+                                 # is the same as WHY_EXCEPTION.
 WHY_RETURN = WHY_RERAISE << 1
 WHY_BREAK = WHY_RETURN << 1
 WHY_CONTINUE = WHY_BREAK << 1
@@ -41,10 +47,10 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
     which is used by later transformers for dataflow analysis.
     '''
     def visit (self, flow, nargs = 0, *args, **kws):
-        '''Given a bytecode flow, and an optional number of arguments,
+        '''Given a map of bytecode basic blocks, and an optional
-        return a :py:class:`llpython.control_flow.ControlFlowGraph`
+        number of arguments, return a
-        instance describing the full control flow of the bytecode
+        :py:class:`llpython.control_flow.ControlFlowGraph` instance
-        flow.'''
+        describing the full control flow of the bytecode flow.'''
         self.nargs = nargs
         ret_val = super(ControlFlowBuilder, self).visit(flow, *args, **kws)
         del self.nargs
@@ -57,6 +63,13 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
         self.block_list.sort()
         self.cfg = ControlFlowGraph()
         self.control_stack = []
+        self.continue_targets = {} # Map from SETUP_LOOP addresses to
+                                   # start of loop addresses, based on
+                                   # observed CONTINUE_LOOP opcodes.
+        self.break_targets = set() # Set of SETUP_LOOP address that
+                                   # have at least one observed
+                                   # BREAK_LOOP opcode corresponding
+                                   # to them.
         for block in self.block_list:
             self.cfg.add_block(block, blocks[block])
 
@@ -66,6 +79,7 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
         self.cfg.compute_dataflow()
         self.cfg.update_for_ssa()
         ret_val = self.cfg
+        del self.continue_targets
         del self.control_stack
         del self.cfg
         del self.block_list
@@ -91,29 +105,37 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
 
         Returns True if an edge was added to the CFG, False otherwise.
         Based on the opcode the return result may mean different
-        things (for example: if why == WHY_RETURN, then the function
+        things (for example: if why == WHY_RETURN, then a False return
-        returns)."""
+        result means the function returned, and no edge was
+        generated)."""
         ret_val = False
         if len(self.control_stack) > 0:
             handlers = set((SETUP_FINALLY, SETUP_WITH))
             if why == WHY_EXCEPTION:
                 handlers.add(SETUP_EXCEPT)
-            reverse_stack = self.control_stack[::-1]
+            reversed_stack = reversed(self.control_stack)
             target = None
-            for handler_i, handler_op, handler_arg in reverse_stack:
+            for handler_i, handler_op, handler_arg in reversed_stack:
                 if handler_op in handlers:
                     target = handler_i + handler_arg + 3
                 elif handler_op == SETUP_LOOP:
                     if why == WHY_CONTINUE:
-                        if op == CONTINUE_LOOP:
+                        # Only generate a WHY_CONTINUE edge if a continue
-                            target = i + arg + 3
+                        # statement has been observed for this loop.
+                        if handler_i not in self.continue_targets:
+                            break
+                        elif op == CONTINUE_LOOP:
+                            target = arg
+                            assert target == self.continue_targets[handler_i]
                         else:
-                            # XXX This isn't going to be correct for
+                            target = self.continue_targets[handler_i]
-                            # for-loops, or really long while-loops
-                            # (which use EXTENDED_ARG):
-                            target = handler_i + 3
                     elif why == WHY_BREAK:
-                        target = handler_i + handler_arg + 3
+                        # Only generate a WHY_BREAK edge if a break
+                        # statement has been observed for this loop.
+                        if handler_i not in self.break_targets:
+                            break
+                        else:
+                            target = handler_i + handler_arg + 3
                 if target is not None:
                     self.cfg.add_edge(block, target)
                     ret_val = True
@@ -137,11 +159,12 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
                                                    WHY_BREAK)
             assert branched, ("Attempted to break outside of loop %r" %
                               (self.blocks[block][-1],))
+        elif op == RAISE_VARARGS:
+            self._generate_handler_edge(block, i, op, arg, WHY_EXCEPTION)
         elif op == END_FINALLY:
-            # XXX Should we detect cases where return, continue, and
+            # The following does a lot of redundant traversal of the
-            # break appear inside the try-block?  This would create
+            # simulated frame stack, but it works, and keeps a lot of
-            # more accurate control flow graphs by eliding edges we
+            # special case logic out of _generate_handler_edge().
-            # know won't be taken.
             self._generate_handler_edge(block, i, op, arg, WHY_EXCEPTION)
             self._generate_handler_edge(block, i, op, arg, WHY_RETURN)
             self._generate_handler_edge(block, i, op, arg, WHY_BREAK)
@@ -156,6 +179,9 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
         if op in opcode_util.hascbranch or goto_next:
             self.cfg.add_edge(block, self._get_next_block(block))
 
+    # ____________________________________________________________
+    # LOAD/STORE_FAST
+
     def op_LOAD_FAST (self, i, op, arg, *args, **kws):
         self.cfg.blocks_reads[self.block].add(arg)
         return super(ControlFlowBuilder, self).op_LOAD_FAST(i, op, arg, *args,
@@ -166,6 +192,44 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
         return super(ControlFlowBuilder, self).op_STORE_FAST(i, op, arg, *args,
                                                              **kws)
 
+    # ____________________________________________________________
+    # *_LOOP: Special loop control flow.
+
+    def _get_current_loop (self):
+        for handler in reversed(self.control_stack):
+            if handler[1] == SETUP_LOOP:
+                return handler
+        return None, None, None
+
+    def op_BREAK_LOOP (self, i, op, arg, *args, **kws):
+        handler_i, _, _ = self._get_current_loop()
+        assert handler_i is not None
+        self.break_targets.add(handler_i)
+
+    def op_CONTINUE_LOOP (self, i, op, arg, *args, **kws):
+        """
+        CONTINUE_LOOP has to be handled differently than BREAK_LOOP,
+        since its argument specifies where the start of the loop is
+        (in the case of for-loops, FOR_ITER defines the true start of
+        the loop, instead of SETUP_LOOP.)
+        """
+        handler_i, _, _ = self._get_current_loop()
+        assert handler_i is not None
+        if handler_i in self.continue_targets:
+            assert arg == self.continue_targets[handler_i]
+        else:
+            self.continue_targets[handler_i] = arg
+
+    # ____________________________________________________________
+    # POP_BLOCK
+
+    def op_POP_BLOCK (self, i, op, arg, *args, **kws):
+        self.control_stack.pop()
+        return super(ControlFlowBuilder, self).op_POP_BLOCK(i, op, arg, *args,
+                                                            **kws)
+    # ____________________________________________________________
+    # SETUP_*
+
     def op_SETUP_EXCEPT (self, i, op, arg, *args, **kws):
         self.control_stack.append((i, op, arg))
         return super(ControlFlowBuilder, self).op_SETUP_EXCEPT(i, op, arg,
@@ -186,11 +250,6 @@ class ControlFlowBuilder (BenignBytecodeVisitorMixin, BasicBlockVisitor):
         return super(ControlFlowBuilder, self).op_SETUP_WITH(i, op, arg, *args,
                                                              **kws)
 
-    def op_POP_BLOCK (self, i, op, arg, *args, **kws):
-        self.control_stack.pop()
-        return super(ControlFlowBuilder, self).op_POP_BLOCK(i, op, arg, *args,
-                                                            **kws)
-
 # ______________________________________________________________________
 
 def build_cfg (func):

llpython/opcode_util.py

@@ -9,9 +9,11 @@ import opcode
 
 # Note that opcode.hasjrel and opcode.hasjabs applies only to opcodes
 # that calculate a jump point based on the argument.  This ignores
-# jumps that use the frame stack to calculate their targets.
+# jumps that use the frame stack to calculate their targets, and
+# exceptions.
 
-NON_ARG_JUMP_NAMES = 'BREAK_LOOP', 'RETURN_VALUE', 'END_FINALLY'
+NON_ARG_JUMP_NAMES = ('BREAK_LOOP', 'RETURN_VALUE', 'END_FINALLY',
+                      'RAISE_VARARGS')
 NON_ARG_JUMPS = [opcode.opmap[opname]
                  for opname in NON_ARG_JUMP_NAMES
                  if opname in opcode.opmap]
@@ -118,7 +120,7 @@ OPCODE_MAP = {
     'PRINT_ITEM_TO': (2, None, 1),
     'PRINT_NEWLINE': (0, None, 1),
     'PRINT_NEWLINE_TO': (1, None, 1),
-    'RAISE_VARARGS': (None, None, None),
+    'RAISE_VARARGS': (-1, None, 1),
     'RETURN_VALUE': (1, None, 1),
     'ROT_FOUR': (None, None, None),
     'ROT_THREE': (None, None, None),

llpython/tests/test_byte_control.py

@@ -3,6 +3,7 @@
 
 from __future__ import absolute_import
 
+import sys
 import unittest
 
 from llpython import byte_control
@@ -10,6 +11,11 @@ from llpython import byte_control
 # ______________________________________________________________________
 # Global data
 
+# Technically we could also compute if we need special logic for the
+# old bytecode compiler by scanning for JUMP_IF_TRUE and JUMP_IF_FALSE
+# opcodes.  These opcodes require additional POP_TOP's be inserted.
+OLD_BYTECODE_COMPILER = sys.version_info < (2, 7)
+
 got_done = 0
 
 # ______________________________________________________________________
@@ -19,6 +25,7 @@ def do_something():
     global got_done
     got_done += 1
     print("Something good got done.")
+    return got_done
 
 # ____________________________________________________________
 
@@ -47,6 +54,62 @@ def try_finally_1(m, n): # why == WHY_BREAK
             do_something()
     return i
 
+# ____________________________________________________________
+
+def try_finally_2(m, n): # why == WHY_CONTINUE
+    i = m
+    while i < n:
+        try:
+            if i == 101:
+                i += 200
+                continue
+        finally:
+            do_something()
+        i += 1
+    return i
+
+# ____________________________________________________________
+
+def try_finally_3(m, n): # why == WHY_EXCEPTION (or WHY_RETURN)
+    d = {}
+    try:
+        return d[n] - d[m]
+    finally:
+        do_something()
+    return do_something_else()
+
+# ____________________________________________________________
+
+def try_finally_4(m, n): # why == WHY_NOT
+    try:
+        rv = n - m
+    finally:
+        do_something()
+    return rv
+
+# ____________________________________________________________
+
+def try_finally_5(m, n):
+    for i in range(m, n):
+        try:
+            if i == 99:
+                break
+            elif i == 121:
+                continue
+            elif i == 86:
+                return
+            elif i < -102:
+                raise ValueError(i)
+            else:
+                try:
+                    if i < 0:
+                        raise ValueError(i)
+                finally:
+                    do_something()
+        finally:
+            do_something()
+    return do_something()
+
 # ______________________________________________________________________
 # Class (test case) definition(s)
 
@@ -55,22 +118,151 @@ class TestByteControl(unittest.TestCase):
         assert len(test_cfg.blocks) == block_count
         block_keys = list(test_cfg.blocks.keys())
         block_keys.sort()
-        # TODO: Ensure unexpected edges cause error.
+        expected_blocks_in = dict((block_key, set())
+                                  for block_key in block_keys)
+        expected_blocks_out = dict((block_key, set())
+                                   for block_key in block_keys)
         for from_block_ofs, to_block_ofs in edges:
             from_block = block_keys[from_block_ofs]
             to_block = block_keys[to_block_ofs]
-            assert from_block in test_cfg.blocks_in[to_block]
+            expected_blocks_in[to_block].add(from_block)
-            assert to_block in test_cfg.blocks_out[from_block]
+            expected_blocks_out[from_block].add(to_block)
+        for block_key in block_keys:
+            expected_in = expected_blocks_in[block_key]
+            test_in = test_cfg.blocks_in[block_key]
+            self.assertEqual(
+                expected_in, test_in, '%r != %r for blocks_in[%d]' % (
+                    test_in, expected_in, block_key))
+            expected_out = expected_blocks_out[block_key]
+            test_out = test_cfg.blocks_out[block_key]
+            self.assertEqual(
+                expected_out, test_out, '%r != %r for set blocks_out[%d]' % (
+                    test_out, expected_out, block_key))
+
+    def test_raise(self):
+        cfg = byte_control.build_cfg(do_something_else)
+        self.fail_unless_cfg_match(cfg, 2, ())
 
     def test_try_finally_0(self):
+        """
+        Expected CFG (Python 2.7+):
+        digraph CFG_try_finally_0 {
+            BLOCK_0 -> BLOCK_3; // 0 -> 1
+            BLOCK_0 -> BLOCK_15; // 0 -> 3
+            BLOCK_3 -> BLOCK_15; // 1 -> 3
+            BLOCK_11 -> BLOCK_15; // 2 -> 3
+            BLOCK_15 -> BLOCK_23; // 3 -> 4, why == WHY_NOT
+            BLOCK_23; // 4
+        }
+        (Possibly terminal blocks: 15, 23.)
+        """
         cfg = byte_control.build_cfg(try_finally_0)
         self.fail_unless_cfg_match(cfg, 5, ((0, 1), (0, 3), (1, 3), (2, 3),
                                             (3, 4)))
 
     def test_try_finally_1(self):
+        """
+        Expected CFG (Python 2.7+):
+        digraph CFG_try_finally_1 {
+            BLOCK_0 -> BLOCK_9; // 0 -> 1
+            BLOCK_0 -> BLOCK_63; // 0 -> 11
+            BLOCK_9 -> BLOCK_22; // 1 -> 2
+            BLOCK_22 -> BLOCK_25; // 2 -> 3
+            BLOCK_22 -> BLOCK_62; // 2 -> 10
+            BLOCK_25 -> BLOCK_31; // 3 -> 4
+            BLOCK_25 -> BLOCK_51; // 3 -> 8
+            BLOCK_31 -> BLOCK_43; // 4 -> 5
+            BLOCK_31 -> BLOCK_47; // 4 -> 7
+            BLOCK_43 -> BLOCK_51; // 5 -> 8
+            BLOCK_44 -> BLOCK_47; // 6 -> 7
+            BLOCK_47 -> BLOCK_51; // 7 -> 8
+            BLOCK_51 -> BLOCK_59; // 8 -> 9, why == WHY_NOT
+            BLOCK_51 -> BLOCK_63; // 8 -> 11, why == WHY_BREAK, WHY_RETURN, ...
+            BLOCK_59 -> BLOCK_22; // 9 -> 2
+            BLOCK_62 -> BLOCK_63; // 10 -> 11
+            BLOCK_63; // 11
+        }
+        (Possibly terminal blocks: 51, 63.)
+        """
         cfg = byte_control.build_cfg(try_finally_1)
-        # TODO: Translate known graph to offsets...
+        if not OLD_BYTECODE_COMPILER:
-        self.fail_unless_cfg_match(cfg, 12, ())
+            self.fail_unless_cfg_match(
+                cfg, 12, ((0, 1), (0, 11), (1, 2), (2, 3), (2, 10), (3, 4),
+                          (3, 8), (4, 5), (4, 7), (5, 8), (6, 7), (7, 8),
+                          (8, 9), (8, 11), (9, 2), (10, 11)))
+        else:
+            self.fail_unless_cfg_match(
+                cfg, 13, ((0, 1), (0, 12), (1, 2), (2, 3), (2, 11), (3, 4),
+                          (3, 9), (4, 5), (4, 7), (5, 9), (6, 8), (7, 8),
+                          (8, 9), (9, 10), (9, 12), (10, 2), (11, 12)))
+
+    def test_try_finally_2(self):
+        """
+        Expected CFG (Python 2.7+):
+        digraph CFG_try_finally_2 {
+            BLOCK_0 -> BLOCK_9; // 0 -> 1
+            BLOCK_0 -> BLOCK_78; // 0 -> 10
+            BLOCK_9 -> BLOCK_21; // 1 -> 2
+            BLOCK_9 -> BLOCK_77; // 1 -> 9
+            BLOCK_21 -> BLOCK_24; // 2 -> 3
+            BLOCK_21 -> BLOCK_56; // 2 -> 7
+            BLOCK_24 -> BLOCK_36; // 3 -> 4
+            BLOCK_24 -> BLOCK_52; // 3 -> 6
+            BLOCK_36 -> BLOCK_56; // 4 -> 7
+            BLOCK_49 -> BLOCK_52; // 5 -> 6
+            BLOCK_52 -> BLOCK_56; // 6 -> 7
+            BLOCK_56 -> BLOCK_9; // 7 -> 1, why == WHY_CONTINUE
+            BLOCK_56 -> BLOCK_64; // 7 -> 8, why == WHY_NOT
+            BLOCK_64 -> BLOCK_9; // 8 -> 1
+            BLOCK_77 -> BLOCK_78; // 9 -> 10
+            BLOCK_78; // 10
+        }
+        (Possibly terminal blocks: 56, 78.)
+        """
+        cfg = byte_control.build_cfg(try_finally_2)
+        if not OLD_BYTECODE_COMPILER:
+            self.fail_unless_cfg_match(
+                cfg, 11, ((0, 1), (0, 10), (1, 2), (1, 9), (2, 3), (2, 7),
+                          (3, 4), (3, 6), (4, 7), (5, 6), (6, 7), (7, 1),
+                          (7, 8), (8, 1), (9, 10)))
+
+    def test_try_finally_3(self):
+        """
+        Expected (Python 2.7+):
+        digraph CFG_foo3 {
+            BLOCK_0 -> BLOCK_9; // 0 -> 1
+            BLOCK_0 -> BLOCK_29; // 0 -> 3
+            BLOCK_9 -> BLOCK_29; // 1 -> 3
+            BLOCK_25 -> BLOCK_29; // 2 -> 3
+            BLOCK_29 -> BLOCK_37; // 3 -> 4, why == WHY_NOT
+            BLOCK_37; // 4
+        }
+        (Possibly terminal blocks: 29, 37.)
+        """
+        cfg = byte_control.build_cfg(try_finally_3)
+        self.fail_unless_cfg_match(cfg, 5, ((0, 1), (0, 3), (1, 3), (2, 3),
+                                            (3, 4)))
+
+    def test_try_finally_4(self):
+        """
+        Expected:
+        digraph CFG_foo4 {
+           BLOCK_0 -> BLOCK_3; // 0 -> 1
+           BLOCK_0 -> BLOCK_17; // 0 -> 2
+           BLOCK_3 -> BLOCK_17; // 1 -> 2
+           BLOCK_17 -> BLOCK_25; // 2 -> 3, why == WHY_NOT
+           BLOCK_25; // 3
+        }
+        (Possibly terminal blocks: 17, 25.)
+        """
+        cfg = byte_control.build_cfg(try_finally_4)
+        self.fail_unless_cfg_match(cfg, 4, ((0, 1), (0, 2), (1, 2), (2, 3)))
+
+    def test_try_finally_5(self):
+        pass # TODO: fix ControlFlowGraph.update_for_ssa which
+             # diverges on the CFG for the following:
+        #cfg = byte_control.build_cfg(try_finally_5)
+        #self.fail_unless_cfg_match(cfg, XXX, ())
 
 # ______________________________________________________________________