Repository: numba/numba

Issue:
A return statement in a nested loop may cause error in loop lifting
The following code
```
import numba
@numba.jit
def numbafunc(arg):
	for i in range(10):
		for j in range(10):
			return

numbafunc([])
```
causes the following error
```
Traceback (most recent call last):
  File "numbabug.py", line 11, in <module>
    numbafunc([])
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 307, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 579, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 80, in compile
    flags=flags, locals=self.locals)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 763, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 360, in compile_extra
    return self._compile_bytecode()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 722, in _compile_bytecode
    return self._compile_core()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 246, in run
    raise patched_exception
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 238, in run
    stage()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 431, in stage_objectmode_frontend
    cres = self.frontend_looplift()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 411, in frontend_looplift
    flags=loop_flags)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/transforms.py", line 191, in loop_lifting
    func_ir.variable_lifetime.livemap)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/transforms.py", line 59, in _loop_lift_get_candidate_infos
    [(returnto, _)] = cfg.successors(an_exit)  # requirement checked earlier
ValueError: Failed at object (object mode frontend)
not enough values to unpack (expected 1, got 0)
```

Looking at the output of cfg.dump(), it seems that the loop has two exit nodes, and one of them (the return statement) does not have a successor. If that exit comes first out of the iterator, the error will occur. Random changes to the code may cause the other exit to come first, which avoids the error.
A return statement in a nested loop may cause error in loop lifting
The following code
```
import numba
@numba.jit
def numbafunc(arg):
	for i in range(10):
		for j in range(10):
			return

numbafunc([])
```
causes the following error
```
Traceback (most recent call last):
  File "numbabug.py", line 11, in <module>
    numbafunc([])
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 307, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 579, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/dispatcher.py", line 80, in compile
    flags=flags, locals=self.locals)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 763, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 360, in compile_extra
    return self._compile_bytecode()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 722, in _compile_bytecode
    return self._compile_core()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 246, in run
    raise patched_exception
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 238, in run
    stage()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 431, in stage_objectmode_frontend
    cres = self.frontend_looplift()
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/compiler.py", line 411, in frontend_looplift
    flags=loop_flags)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/transforms.py", line 191, in loop_lifting
    func_ir.variable_lifetime.livemap)
  File "/home/tuukka/.local/lib/python3.5/site-packages/numba/transforms.py", line 59, in _loop_lift_get_candidate_infos
    [(returnto, _)] = cfg.successors(an_exit)  # requirement checked earlier
ValueError: Failed at object (object mode frontend)
not enough values to unpack (expected 1, got 0)
```

Looking at the output of cfg.dump(), it seems that the loop has two exit nodes, and one of them (the return statement) does not have a successor. If that exit comes first out of the iterator, the error will occur. Random changes to the code may cause the other exit to come first, which avoids the error.


Discussion:
Thanks for the report. I can reproduce and think your analysis is correct.
Thanks for the report. I can reproduce and think your analysis is correct.

Resolving patch:
diff --git a/numba/transforms.py b/numba/transforms.py
--- a/numba/transforms.py
+++ b/numba/transforms.py
@@ -20,7 +20,13 @@ def same_exit_point(loop):
         "all exits must point to the same location"
         outedges = set()
         for k in loop.exits:
-            outedges |= set(x for x, _ in cfg.successors(k))
+            succs = set(x for x, _ in cfg.successors(k))
+            if not succs:
+                # If the exit point has no successor, it contains an return
+                # statement, which is not handled by the looplifting code.
+                # Thus, this loop is not a candidate.
+                return False
+            outedges |= succs
         return len(outedges) == 1
 
     def one_entry(loop):


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
@@ -70,6 +70,12 @@ def lift_gen1(x):
         a[i] = x
     yield np.sum(a)
 
+def lift_issue2561():
+    np.empty(1)   # This forces objectmode because no nrt
+    for i in range(10):
+        for j in range(10):
+            return 1
+    return 2
 
 def reject1(x):
     a = np.arange(4)
@@ -201,6 +207,9 @@ def test_lift4(self):
     def test_lift5(self):
         self.check_lift_ok(lift5, (types.intp,), (123,))
 
+    def test_lift_issue2561(self):
+        self.check_no_lift(lift_issue2561, (), ())
+
     @tag('important')
     def test_lift_gen1(self):
         self.check_lift_generator_ok(lift_gen1, (types.intp,), (123,))