Repository: numba/numba

Issue:
Control flow analysis failure in object mode when loop nested inside if statement
When Numba attempts to lift loops that are nested inside of an if-else construct, the control flow analysis will fail.  Here is a simple example:

``` python
from numba import jit, void, int32, double

@jit(void(int32, double[:]), forceobj=True)
def forloop_with_if(u, a):
    if u == 0:
        for i in range(a.shape[0]):
            a[i] = a[i] * 2.0
    else:
        for i in range(a.shape[0]):
            a[i] = a[i] + 1.0
```

The resulting exception is:

```
Traceback (most recent call last):
  File "mincase.py", line 4, in <module>
    @jit(void(int32, double[:]), forceobj=True)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/decorators.py", line 128, in wrapper
    disp.compile(sig)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/dispatcher.py", line 235, in compile
    flags=flags, locals=locs)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/compiler.py", line 146, in compile_extra
    func_attr=func_attr)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/compiler.py", line 215, in compile_bytecode
    func_attr=func_attr)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/compiler.py", line 154, in compile_bytecode
    interp = translate_stage(bc)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/compiler.py", line 330, in translate_stage
    interp = interpreter.Interpreter(bytecode=bytecode)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/interpreter.py", line 71, in __init__
    self.cfa.run()
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/controlflow.py", line 520, in run
    graph.add_edge(b.offset, out, pops)
  File "/Users/stan/anaconda/envs/numba_dev/lib/python2.7/site-packages/numba/controlflow.py", line 80, in add_edge
    assert dest in self._nodes
AssertionError
```

This was reported on the numba-users mailing list: https://groups.google.com/a/continuum.io/d/msgid/numba-users/52256cd0-6b65-48d4-8193-3b1f51ae40c4%40continuum.io



Discussion:
Just to note that compilation succeeds if this code is refactored as two separate if statements with no else clauses.


Resolving patch:
diff --git a/numba/bytecode.py b/numba/bytecode.py
--- a/numba/bytecode.py
+++ b/numba/bytecode.py
@@ -197,6 +197,18 @@ def get_jump_target(self):
     def __repr__(self):
         return '%s(arg=%s, lineno=%d)' % (self.opname, self.arg, self.lineno)
 
+    @property
+    def block_effect(self):
+        """Effect of the block stack
+        Returns +1 (push), 0 (none) or -1 (pop)
+        """
+        if self.opname.startswith('SETUP_'):
+            return 1
+        elif self.opname == 'POP_BLOCK':
+            return -1
+        else:
+            return 0
+
 
 class ByteCodeIter(object):
     def __init__(self, code):
diff --git a/numba/looplifting.py b/numba/looplifting.py
--- a/numba/looplifting.py
+++ b/numba/looplifting.py
@@ -206,7 +206,10 @@ def separate_loops(bytecode, outer, loops):
         if endloop is None:
             if inst.opname == 'SETUP_LOOP':
                 cur = [inst]
-                endloop = inst.next + inst.arg
+                # Python may set the end of loop to the final jump destination
+                # when nested in a if-else.  We need to scan the bytecode to
+                # find the actual end of loop
+                endloop = _scan_real_end_loop(bytecode, inst)
             else:
                 outer.append(inst)
         else:
@@ -219,4 +222,21 @@ def separate_loops(bytecode, outer, loops):
                         break
                 else:
                     loops.append(cur)
-                endloop = None
\ No newline at end of file
+                endloop = None
+
+
+def _scan_real_end_loop(bytecode, setuploop_inst):
+    """Find the end of loop.
+    Return the instruction offset.
+    """
+    start = setuploop_inst.next
+    end = start + setuploop_inst.arg
+    offset = start
+    depth = 0
+    while offset < end:
+        inst = bytecode[offset]
+        depth += inst.block_effect
+        if depth < 0:
+            return inst.next
+        offset = inst.next
+


Test patch:
diff --git a/numba/tests/test_looplifting.py b/numba/tests/test_looplifting.py
--- a/numba/tests/test_looplifting.py
+++ b/numba/tests/test_looplifting.py
@@ -53,7 +53,6 @@ def reject_npm1(x):
 
 
 class TestLoopLifting(TestCase):
-
     def check_lift_ok(self, pyfunc, argtypes, args):
         """
         Check that pyfunc can loop-lift even in nopython mode.
@@ -107,5 +106,26 @@ def test_reject_npm1(self):
         self.check_no_lift_nopython(reject_npm1, (types.intp,), (123,))
 
 
+class TestLoopLiftingInAction(TestCase):
+    def test_issue_734(self):
+        from numba import jit, void, int32, double
+
+        @jit(void(int32, double[:]), forceobj=True)
+        def forloop_with_if(u, a):
+            if u == 0:
+                for i in range(a.shape[0]):
+                    a[i] = a[i] * 2.0
+            else:
+                for i in range(a.shape[0]):
+                    a[i] = a[i] + 1.0
+
+        for u in (0, 1):
+            nb_a = np.arange(10, dtype='int32')
+            np_a = np.arange(10, dtype='int32')
+            forloop_with_if(u, nb_a)
+            forloop_with_if.py_func(u, np_a)
+            self.assertTrue(np.all(nb_a == np_a))
+
+
 if __name__ == '__main__':
     unittest.main()