Repository: numba/numba

Issue:
AssertionError when processing bytecode with embedded if-expression
This is on Python 2.6:

```
>>> @jit(nopython=True)
... def f1(a, b):
...     c = 5 + (a if a>b else b)/2.0
...     return c
...
>>> f1(2, 3)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "numba\dispatcher.py", line 125, in _compile_and_call
    self.jit(sig)
  File "numba\dispatcher.py", line 120, in jit
    return self.compile(sig, **kws)
  File "numba\dispatcher.py", line 108, in compile
    flags=flags, locals=locs)
  File "numba\compiler.py", line 94, in compile_extra
    return_type, flags, locals)
  File "numba\compiler.py", line 99, in compile_bytecode
    interp = translate_stage(bc)
  File "numba\compiler.py", line 245, in translate_stage
    interp = interpreter.Interpreter(bytecode=bytecode)
  File "numba\interpreter.py", line 26, in __init__
    self.dfa.run()
  File "numba\dataflow.py", line 23, in run
    self.infos[blk.offset] = self.run_on_block(blk)
  File "numba\dataflow.py", line 29, in run_on_block
    self.dispatch(info, inst)
  File "numba\dataflow.py", line 39, in dispatch
    fn(info, inst)
  File "numba\dataflow.py", line 222, in _binaryop
    lhs = info.pop()
  File "numba\dataflow.py", line 487, in pop
    assert not self.insts
AssertionError
```



Discussion:
Reproduced on 2.7 and 3.4.

This looks unfortunately like a more general issue in the control flow / dataflow analysis passes. Each BlockInfo pretends to have a separate stack and is oblivious of other blocks' respective stack effects. The _insert_phi() method is then unable to propagate incoming values accross more than one block, but here one of the operands of the BINARY_ADD is coming from two blocks before.


Resolving patch:
diff --git a/numba/controlflow.py b/numba/controlflow.py
--- a/numba/controlflow.py
+++ b/numba/controlflow.py
@@ -1,5 +1,7 @@
 from __future__ import print_function, division, absolute_import
+
 import functools
+
 from numba import utils
 
 
@@ -10,12 +12,16 @@ class CFBlock(object):
     def __init__(self, offset):
         self.offset = offset
         self.body = []
-        self.outgoing = set()
-        self.incoming = set()
+        # A map of jumps to outgoing blocks (successors):
+        #   { offset of outgoing block -> number of stack pops }
+        self.outgoing_jumps = {}
+        # A map of jumps to incoming blocks (predecessors):
+        #   { offset of incoming block -> number of stack pops }
+        self.incoming_jumps = {}
         self.terminating = False
 
     def __repr__(self):
-        args = self.body, self.outgoing, self.incoming
+        args = self.body, sorted(self.outgoing_jumps), sorted(self.incoming_jumps)
         return "block(body: %s, outgoing: %s, incoming: %s)" % args
 
     def __iter__(self):
@@ -67,6 +73,14 @@ def iterliveblocks(self):
             if i in self.liveblocks:
                 yield self.blocks[i]
 
+    def incoming_blocks(self, block):
+        """
+        Yield (incoming block, number of stack pops) pairs for *block*.
+        """
+        for i, pops in block.incoming_jumps.items():
+            if i in self.liveblocks:
+                yield self.blocks[i], pops
+
     def run(self):
         for inst in self._iter_inst():
             fname = "op_%s" % inst.opname
@@ -79,13 +93,13 @@ def run(self):
         # Close all blocks
         for cur, nxt in zip(self.blockseq, self.blockseq[1:]):
             blk = self.blocks[cur]
-            if not blk.outgoing and not blk.terminating:
-                blk.outgoing.add(nxt)
+            if not blk.outgoing_jumps and not blk.terminating:
+                blk.outgoing_jumps[nxt] = 0
 
         # Fill incoming
         for b in utils.dict_itervalues(self.blocks):
-            for out in b.outgoing:
-                self.blocks[out].incoming.add(b.offset)
+            for out, pops in b.outgoing_jumps.items():
+                self.blocks[out].incoming_jumps[b.offset] = pops
 
         # Find liveblocks
         self.dead_block_elimin()
@@ -121,7 +135,7 @@ def dead_block_elimin(self):
         while pending:
             cur = pending.pop()
             blk = self.blocks[cur]
-            outgoing = set(blk.outgoing)
+            outgoing = set(blk.outgoing_jumps)
             liveset |= outgoing
             pending |= outgoing - finished
             finished.add(cur)
@@ -129,8 +143,12 @@ def dead_block_elimin(self):
         for offset in liveset:
             self.liveblocks[offset] = self.blocks[offset]
 
-    def jump(self, target):
-        self._curblock.outgoing.add(target)
+    def jump(self, target, pops=0):
+        """
+        Register a jump (conditional or not) to *target* offset.
+        *pops* is the number of stack pops implied by the jump (default 0).
+        """
+        self._curblock.outgoing_jumps[target] = pops
 
     def _iter_inst(self):
         for inst in self.bytecode:
@@ -177,8 +195,14 @@ def _op_ABSOLUTE_JUMP_IF(self, inst):
     op_POP_JUMP_IF_TRUE = _op_ABSOLUTE_JUMP_IF
     op_JUMP_IF_FALSE = _op_ABSOLUTE_JUMP_IF
     op_JUMP_IF_TRUE = _op_ABSOLUTE_JUMP_IF
-    op_JUMP_IF_FALSE_OR_POP = _op_ABSOLUTE_JUMP_IF
-    op_JUMP_IF_TRUE_OR_POP = _op_ABSOLUTE_JUMP_IF
+
+    def _op_ABSOLUTE_JUMP_OR_POP(self, inst):
+        self.jump(inst.get_jump_target())
+        self.jump(inst.next, pops=1)
+        self._force_new_block = True
+
+    op_JUMP_IF_FALSE_OR_POP = _op_ABSOLUTE_JUMP_OR_POP
+    op_JUMP_IF_TRUE_OR_POP = _op_ABSOLUTE_JUMP_OR_POP
 
     def op_JUMP_ABSOLUTE(self, inst):
         self.jump(inst.get_jump_target())
@@ -221,7 +245,7 @@ def find_dominators(blocks):
         changed = False
         for blk in remainblks:
             d = doms[blk]
-            ps = [doms[p] for p in blocks[blk].incoming if p in doms]
+            ps = [doms[p] for p in blocks[blk].incoming_jumps if p in doms]
             if not ps:
                 p = set()
             else:
diff --git a/numba/dataflow.py b/numba/dataflow.py
--- a/numba/dataflow.py
+++ b/numba/dataflow.py
@@ -1,8 +1,11 @@
 from __future__ import print_function, division, absolute_import
+
+import collections
 from pprint import pprint
-from numba import utils
 import warnings
 
+from numba import utils
+
 
 class DataFlowAnalysis(object):
     """
@@ -15,15 +18,45 @@ class DataFlowAnalysis(object):
     def __init__(self, cfa):
         self.cfa = cfa
         self.bytecode = cfa.bytecode
+        # { block offset -> BlockInfo }
         self.infos = {}
         self.syntax_blocks = []
 
     def run(self):
-        for blk in self.cfa.iterblocks():
+        for blk in self.cfa.iterliveblocks():
             self.infos[blk.offset] = self.run_on_block(blk)
 
     def run_on_block(self, blk):
-        info = BlockInfo(blk.offset)
+        incoming_blocks = []
+        info = BlockInfo(blk.offset, incoming_blocks)
+
+        for ib, pops in self.cfa.incoming_blocks(blk):
+            # By nature of Python bytecode, there will be no incoming
+            # variables from subsequent blocks.  This is an easy way
+            # of breaking the potential circularity of the problem.
+            if ib.offset >= blk.offset:
+                continue
+            ib = self.infos[ib.offset]
+            incoming_blocks.append(ib)
+
+            # Compute stack offset at block entry
+            # The stack effect of our predecessors should be known
+            assert ib.stack_offset is not None, ib
+            new_offset = ib.stack_offset + ib.stack_effect - pops
+            if new_offset < 0:
+                raise RuntimeError("computed negative stack offset for %s"
+                                   % blk)
+            if info.stack_offset is None:
+                info.stack_offset = new_offset
+            elif info.stack_offset != new_offset:
+                warnings.warn("inconsistent stack offset for %s" % blk,
+                              RuntimeWarning)
+
+        if info.stack_offset is None:
+            # No incoming blocks => assume it's the entry block
+            info.stack_offset = 0
+        info.stack_effect = 0
+
         for offset in blk:
             inst = self.bytecode[offset]
             self.dispatch(info, inst)
@@ -121,7 +154,7 @@ def op_BUILD_SET(self, info, inst):
         info.push(res)
 
     def op_POP_TOP(self, info, inst):
-        info.pop()
+        info.pop(discard=True)
 
     def op_STORE_ATTR(self, info, inst):
         target = info.pop()
@@ -143,7 +176,7 @@ def op_LOAD_FAST(self, info, inst):
         info.push(name)
 
     def op_LOAD_CONST(self, info, inst):
-        res = info.make_temp()
+        res = info.make_temp('const')
         info.append(inst, res=res)
         info.push(res)
 
@@ -475,41 +508,85 @@ def __init__(self):
 
 
 class BlockInfo(object):
-    def __init__(self, offset):
+    def __init__(self, offset, incoming_blocks):
         self.offset = offset
+        # The list of incoming BlockInfo objects (obtained by control
+        # flow analysis).
+        self.incoming_blocks = incoming_blocks
         self.stack = []
-        self.incomings = []
+        # Outgoing variables from this block:
+        #   { outgoing phi name -> var name }
+        self.outgoing_phis = {}
         self.insts = []
         self.tempct = 0
         self._term = None
+        self.stack_offset = None
+        self.stack_effect = 0
+
+    def __repr__(self):
+        return "<%s at offset %d>" % (self.__class__.__name__, self.offset)
 
     def dump(self):
         print("offset", self.offset, "{")
         print("  stack: ", end='')
         pprint(self.stack)
-        print("  incomings: ", end='')
-        pprint(self.incomings)
         pprint(self.insts)
         print("}")
 
-    def make_temp(self):
+    def make_temp(self, prefix=''):
         self.tempct += 1
-        name = '$%d.%d' % (self.offset, self.tempct)
+        name = '$%s%d.%d' % (prefix, self.offset, self.tempct)
         return name
 
     def push(self, val):
+        self.stack_effect += 1
         self.stack.append(val)
 
-    def pop(self):
-        # TODO: lingering incoming values
+    def pop(self, discard=False):
+        """
+        Pop a variable from the stack, or request it from incoming blocks if
+        the stack is empty.
+        If *discard* is true, the variable isn't meant to be used anymore,
+        which allows reducing the number of temporaries created.
+        """
         if not self.stack:
-            assert not self.insts
-            ret = self.make_temp()
-            self.incomings.append(ret)
+            self.stack_offset -= 1
+            if not discard:
+                return self.make_incoming()
         else:
-            ret = self.stack.pop()
+            self.stack_effect -= 1
+            return self.stack.pop()
+
+    def make_incoming(self):
+        """
+        Create an incoming variable (due to not enough values being
+        available on our stack) and request its assignment from our
+        incoming blocks' own stacks.
+        """
+        assert self.incoming_blocks
+        ret = self.make_temp('phi')
+        for ib in self.incoming_blocks:
+            stack_index = self.stack_offset + self.stack_effect
+            ib.request_outgoing(self, ret, stack_index)
         return ret
 
+    def request_outgoing(self, outgoing_block, phiname, stack_index):
+        """
+        Request the assignment of the next available stack variable
+        for block *outgoing_block* with target name *phiname*.
+        """
+        if phiname in self.outgoing_phis:
+            # If phiname was already requested, ignore this new request
+            # (can happen with a diamond-shaped block flow structure).
+            return
+        if stack_index < self.stack_offset:
+            assert self.incoming_blocks
+            for ib in self.incoming_blocks:
+                ib.request_outgoing(self, phiname, stack_index)
+        else:
+            varname = self.stack[stack_index - self.stack_offset]
+            self.outgoing_phis[phiname] = varname
+
     @property
     def tos(self):
         r = self.pop()
diff --git a/numba/interpreter.py b/numba/interpreter.py
--- a/numba/interpreter.py
+++ b/numba/interpreter.py
@@ -89,6 +89,7 @@ def _iter_inst(self):
                 self.loc = ir.Loc(filename=self.bytecode.filename,
                                   line=inst.lineno)
                 yield inst, kws
+            self._end_current_block()
 
     def _start_new_block(self, inst):
         self.loc = ir.Loc(filename=self.bytecode.filename, line=inst.lineno)
@@ -100,38 +101,27 @@ def _start_new_block(self, inst):
             oldblock.append(jmp)
             # Get DFA block info
         self.dfainfo = self.dfa.infos[self.current_block_offset]
-        # Insert PHI
-        self._insert_phi()
         # Notify listeners for the new block
         for fn in utils.dict_itervalues(self._block_actions):
             fn(self.current_block_offset, self.current_block)
 
-    def _insert_phi(self):
-        if self.dfainfo.incomings:
-            assert len(self.dfainfo.incomings) == 1
-            incomings = self.cfa.blocks[self.current_block_offset].incoming
-            phivar = self.dfainfo.incomings[0]
-            if len(incomings) == 1:
-                ib = utils.iter_next(iter(incomings))
-                lingering = self.dfa.infos[ib].stack
-                assert len(lingering) == 1
-                iv = lingering[0]
-                self.store(self.get(iv), phivar)
-            else:
-                # Invert the PHI node
-                for ib in incomings:
-                    lingering = self.dfa.infos[ib].stack
-                    assert len(lingering) == 1
-                    iv = lingering[0]
-
-                    # Add assignment in incoming block to forward the value
-                    target = self.current_scope.get_or_define('$phi' + phivar,
-                                                              loc=self.loc)
-                    stmt = ir.Assign(value=self.get(iv), target=target,
-                                     loc=self.loc)
-                    self.blocks[ib].insert_before_terminator(stmt)
+    def _end_current_block(self):
+        self._insert_outgoing_phis()
 
-                self.store(target, phivar)
+    def _insert_outgoing_phis(self):
+        """
+        Add assignments to forward requested outgoing values
+        to subsequent blocks.
+        """
+        for phiname, varname in self.dfainfo.outgoing_phis.items():
+            target = self.current_scope.get_or_define(phiname,
+                                                      loc=self.loc)
+            stmt = ir.Assign(value=self.get(varname), target=target,
+                             loc=self.loc)
+            if not self.current_block.is_terminated:
+                self.current_block.append(stmt)
+            else:
+                self.current_block.insert_before_terminator(stmt)
 
     def get_global_value(self, name):
         """
@@ -741,7 +731,7 @@ def _determine_while_condition(self, branches):
         else:
             return
             # Which is the exit of the loop
-        if br not in self.cfa.blocks[loop.exit].incoming:
+        if br not in self.cfa.blocks[loop.exit].incoming_jumps:
             return
 
         # Therefore, current block is a while loop condition


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -4,9 +4,12 @@
 
 import contextlib
 
+from numba import types, typing, utils
+from numba.compiler import compile_extra, compile_isolated, Flags, DEFAULT_FLAGS
 from numba import types, utils
 from numba.compiler import compile_isolated, Flags
 from numba.lowering import LoweringError
+from numba.targets import cpu
 from numba.typeinfer import TypingError
 import numba.unittest_support as unittest
 
@@ -20,6 +23,32 @@
 no_pyobj_flags = Flags()
 
 
+class CompilationCache(object):
+    """
+    A cache of compilation results for various signatures and flags.
+    This can make tests significantly faster (or less slow).
+    """
+
+    def __init__(self):
+        self.typingctx = typing.Context()
+        self.targetctx = cpu.CPUContext(self.typingctx)
+        self.cr_cache = {}
+
+    def compile(self, func, args, return_type=None, flags=DEFAULT_FLAGS):
+        """
+        Compile the function or retrieve an already compiled result
+        from the cache.
+        """
+        cache_key = (func, args, return_type, flags)
+        try:
+            cr = self.cr_cache[cache_key]
+        except KeyError:
+            cr = compile_extra(self.typingctx, self.targetctx, func,
+                               args, return_type, flags, locals={})
+            self.cr_cache[cache_key] = cr
+        return cr
+
+
 class TestCase(unittest.TestCase):
 
     @contextlib.contextmanager
diff --git a/numba/tests/test_dataflow.py b/numba/tests/test_dataflow.py
--- a/numba/tests/test_dataflow.py
+++ b/numba/tests/test_dataflow.py
@@ -1,12 +1,20 @@
 from __future__ import print_function
+
 import numba.unittest_support as unittest
 from numba.compiler import compile_isolated, Flags
+from numba.utils import PYVERSION
 from numba import types
+from .support import TestCase, CompilationCache
+
 
+enable_pyobj_flags = Flags()
+enable_pyobj_flags.set("enable_pyobject")
 
 force_pyobj_flags = Flags()
 force_pyobj_flags.set("force_pyobject")
 
+no_pyobj_flags = Flags()
+
 
 def assignments(a):
     b = c = str(a)
@@ -18,24 +26,87 @@ def assignments2(a):
     return b + c + d
 
 
-class TestDataFlow(unittest.TestCase):
+# Use cases for issue #503
+
+def var_propagate1(a, b):
+    c = (a if a > b else b) + 5
+    return c
+
+
+def var_propagate2(a, b):
+    c = 5 + (a if a > b else b + 12) / 2.0
+    return c
+
+
+def var_propagate3(a, b):
+    c = 5 + (a > b and a or b)
+    return c
+
+
+def var_propagate4(a, b):
+    c = 5 + (a - 1 and b + 1) or (a + 1 and b - 1)
+    return c
+
+
+
+class TestDataFlow(TestCase):
+
+    def setUp(self):
+        self.cache = CompilationCache()
+
     def test_assignments(self, flags=force_pyobj_flags):
         pyfunc = assignments
         cr = compile_isolated(pyfunc, (types.int32,), flags=flags)
         cfunc = cr.entry_point
         for x in [-1, 0, 1]:
-            self.assertEqual(pyfunc(x), cfunc(x))
+            self.assertPreciseEqual(pyfunc(x), cfunc(x))
 
     def test_assignments2(self, flags=force_pyobj_flags):
         pyfunc = assignments2
         cr = compile_isolated(pyfunc, (types.int32,), flags=flags)
         cfunc = cr.entry_point
         for x in [-1, 0, 1]:
-            self.assertEqual(pyfunc(x), cfunc(x))
+            self.assertPreciseEqual(pyfunc(x), cfunc(x))
 
         if flags is force_pyobj_flags:
             cfunc("a")
 
+    # The dataflow analysis must be good enough for native mode
+    # compilation to succeed, hence the no_pyobj_flags in the following tests.
+
+    def run_propagate_func(self, pyfunc, args):
+        cr = self.cache.compile(pyfunc, (types.int32, types.int32),
+                                flags=no_pyobj_flags)
+        cfunc = cr.entry_point
+        self.assertPreciseEqual(cfunc(*args), pyfunc(*args))
+
+    def test_var_propagate1(self):
+        self.run_propagate_func(var_propagate1, (2, 3))
+        self.run_propagate_func(var_propagate1, (3, 2))
+
+    def test_var_propagate2(self):
+        self.run_propagate_func(var_propagate2, (2, 3))
+        self.run_propagate_func(var_propagate2, (3, 2))
+
+    def test_var_propagate3(self):
+        self.run_propagate_func(var_propagate3, (2, 3))
+        self.run_propagate_func(var_propagate3, (3, 2))
+        self.run_propagate_func(var_propagate3, (2, 0))
+        self.run_propagate_func(var_propagate3, (-1, 0))
+        self.run_propagate_func(var_propagate3, (0, 2))
+        self.run_propagate_func(var_propagate3, (0, -1))
+
+    def test_var_propagate4(self):
+        self.run_propagate_func(var_propagate4, (1, 1))
+        self.run_propagate_func(var_propagate4, (1, 0))
+        self.run_propagate_func(var_propagate4, (1, -1))
+        self.run_propagate_func(var_propagate4, (0, 1))
+        self.run_propagate_func(var_propagate4, (0, 0))
+        self.run_propagate_func(var_propagate4, (0, -1))
+        self.run_propagate_func(var_propagate4, (-1, 1))
+        self.run_propagate_func(var_propagate4, (-1, 0))
+        self.run_propagate_func(var_propagate4, (-1, -1))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/numba/tests/test_ufuncs.py b/numba/tests/test_ufuncs.py
--- a/numba/tests/test_ufuncs.py
+++ b/numba/tests/test_ufuncs.py
@@ -8,12 +8,12 @@
 import functools
 
 import numba.unittest_support as unittest
-from numba.compiler import compile_extra, compile_isolated, Flags, DEFAULT_FLAGS
-from numba import types, typing, utils
+from numba.compiler import compile_isolated, Flags, DEFAULT_FLAGS
+from numba import types, utils
 from numba.config import PYVERSION
 from numba.targets import cpu
 from numba.typeinfer import TypingError
-from numba.tests.support import TestCase
+from numba.tests.support import TestCase, CompilationCache
 
 
 is32bits = tuple.__itemsize__ == 4
@@ -51,32 +51,6 @@ def _make_binary_ufunc_usecase(ufunc_name):
     return fn
 
 
-class CompilationCache(object):
-    """
-    A cache of compilation results for various signatures and flags.
-    This makes the tests significantly faster (or less slow).
-    """
-
-    def __init__(self):
-        self.typingctx = typing.Context()
-        self.targetctx = cpu.CPUContext(self.typingctx)
-        self.cr_cache = {}
-
-    def compile(self, func, args, return_type=None, flags=DEFAULT_FLAGS):
-        """
-        Compile the function or retrieve an already compiled result
-        from the cache.
-        """
-        cache_key = (func, args, return_type, flags)
-        try:
-            cr = self.cr_cache[cache_key]
-        except KeyError:
-            cr = compile_extra(self.typingctx, self.targetctx, func,
-                               args, return_type, flags, locals={})
-            self.cr_cache[cache_key] = cr
-        return cr
-
-
 class TestUFuncs(TestCase):
 
     def setUp(self):