Repository: numba/numba
Pull request: Fix issue #2561, handle missing successor on loop exit
This attempts to fix #2561, there are cases where there is no
successor present for a loop exit node (quite possibly due to
bytecode optimisations). This patch guards against this case by
trying to find alternative successors and dropping the looplift
candidate if none are found.

Fixes #2561.

Problem statement:
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.

Hints from issue discussion:
Thanks for the report. I can reproduce and think your analysis is correct.

Commit message: Fix issue #2561, handle missing successor on loop exit

This attempts to fix #2561, there are cases where there is no
successor present for a loop exit node (quite possibly due to
bytecode optimisations). This patch guards against this case by
trying to find alternative successors and dropping the looplift
candidate if none are found.

Fixes #2561.

Patch to review:
diff --git a/numba/tests/test_looplifting.py b/numba/tests/test_looplifting.py
index a754f4869ed..b6c1cc8f9e4 100644
--- a/numba/tests/test_looplifting.py
+++ b/numba/tests/test_looplifting.py
@@ -70,6 +70,10 @@ def lift_gen1(x):
         a[i] = x
     yield np.sum(a)
 
+def lift_issue2561(x):
+    for i in range(10):
+        for j in range(10):
+            return
 
 def reject1(x):
     a = np.arange(4)
@@ -201,6 +205,9 @@ def test_lift4(self):
     def test_lift5(self):
         self.check_lift_ok(lift5, (types.intp,), (123,))
 
+    def test_lift_issue2561(self):
+        self.check_lift_ok(lift_issue2561, (types.List,), ([],))
+
     @tag('important')
     def test_lift_gen1(self):
         self.check_lift_generator_ok(lift_gen1, (types.intp,), (123,))
diff --git a/numba/transforms.py b/numba/transforms.py
index 02039af3e58..b2af602c302 100644
--- a/numba/transforms.py
+++ b/numba/transforms.py
@@ -53,8 +53,16 @@ def _loop_lift_get_candidate_infos(cfg, blocks, livemap):
     loopinfos = []
     for loop in loops:
         [callfrom] = loop.entries   # requirement checked earlier
-        an_exit = next(iter(loop.exits))  # anyone of the exit block
-        [(returnto, _)] = cfg.successors(an_exit)  # requirement checked earlier
+        # This loop exists to handle missing cfg.successors, only *an* exit is
+        # needed. Walk in stable order, higher exits first.
+        for an_exit in iter(sorted(loop.exits)):
+            # requirement checked earlier
+            ret = [x for x in cfg.successors(an_exit)]
+            if ret:
+                break
+        else:
+            continue # drop this loop from being liftable 
+        [(returnto, _)] = ret
         # note: sorted for stable ordering
         inputs = sorted(livemap[callfrom])
         outputs = sorted(livemap[returnto])


Code review comments:
diff_hunk: @@ -53,8 +53,16 @@ def _loop_lift_get_candidate_infos(cfg, blocks, livemap):
     loopinfos = []
     for loop in loops:
         [callfrom] = loop.entries   # requirement checked earlier
-        an_exit = next(iter(loop.exits))  # anyone of the exit block
-        [(returnto, _)] = cfg.successors(an_exit)  # requirement checked earlier
+        # This loop exists to handle missing cfg.successors, only *an* exit is
+        # needed. Walk in stable order, higher exits first.
+        for an_exit in iter(sorted(loop.exits)):
+            # requirement checked earlier
+            ret = [x for x in cfg.successors(an_exit)]
+            if ret:
+                break
+        else:
+            continue # drop this loop from being liftable 
path: numba/transforms.py
text: @user1:
Since the return statement in loop is not being handled by the looplifting logic, we shouldn't allow it for looplifting.  We should make it illegal to looplift if the loop contains a return statement.

diff_hunk: @@ -201,6 +205,9 @@ def test_lift4(self):
     def test_lift5(self):
         self.check_lift_ok(lift5, (types.intp,), (123,))
 
+    def test_lift_issue2561(self):
+        self.check_lift_ok(lift_issue2561, (types.List,), ([],))
path: numba/tests/test_looplifting.py
text: @user1:
`types.List` here will trigger an assertion in the type-inference because it is not an instance of `types.Type`.  That assertion is why the function is dropping into object-mode and into looplifting, but it seems like a bug in the test.

diff_hunk: @@ -70,6 +70,10 @@ def lift_gen1(x):
         a[i] = x
     yield np.sum(a)
 
+def lift_issue2561(x):
+    for i in range(10):
+        for j in range(10):
+            return
path: numba/tests/test_looplifting.py
text: @user1:
If we added and check the return value, we will find that the looplifting is not handling the `return` statement properly.