Repository: numba/numba

Issue:
Generator function with branch inside loop can cause KeyError in lower_yield_suspend
System: x86_64 GNU/Linux, LLVM 3.6.1, Python 3.4.3, llvmlite 0.6.0, numba 0.20.0
System: x86_64 GNU/Linux, LLVM 3.6.1, Python 3.4.3, llvmlite 0.8.0, numba 0.22.0

Using a branch inside a loop inside a generator function results in a `KeyError` caused by `numba.generators.LowerYield.lower_yield_suspend` (calling `self.lower.loadvar`) if the conditional statement references a variable declared after a `yield` statement.
This includes explicitly defined variables (see example below) as well as implicit ones, i.e. results of boolean expressions inside an `if`-statement (e.g. as in `if True is True:`).

The following results in raising said `KeyError`:

``` python
@njit
def f():
    while True:
        yield
        p = True
        if p:
            pass
```

Whereas storing the variable prior to the `yield` statement works fine since `p` is in `lower.varmap`:

``` python
@njit
def f():
    while True:
        p = True
        yield
        if p:
            pass
```

Since implicit variables are likewise affected, the following will not compile either:

``` python
@njit
def f():
    while True:
        yield
        if True is True:
            pass
```

Generator function with branch inside loop can cause KeyError in lower_yield_suspend
System: x86_64 GNU/Linux, LLVM 3.6.1, Python 3.4.3, llvmlite 0.6.0, numba 0.20.0
System: x86_64 GNU/Linux, LLVM 3.6.1, Python 3.4.3, llvmlite 0.8.0, numba 0.22.0

Using a branch inside a loop inside a generator function results in a `KeyError` caused by `numba.generators.LowerYield.lower_yield_suspend` (calling `self.lower.loadvar`) if the conditional statement references a variable declared after a `yield` statement.
This includes explicitly defined variables (see example below) as well as implicit ones, i.e. results of boolean expressions inside an `if`-statement (e.g. as in `if True is True:`).

The following results in raising said `KeyError`:

``` python
@njit
def f():
    while True:
        yield
        p = True
        if p:
            pass
```

Whereas storing the variable prior to the `yield` statement works fine since `p` is in `lower.varmap`:

``` python
@njit
def f():
    while True:
        p = True
        yield
        if p:
            pass
```

Since implicit variables are likewise affected, the following will not compile either:

``` python
@njit
def f():
    while True:
        yield
        if True is True:
            pass
```



Discussion:
Any idea how hard this would be to fix?
With current master, i.e `0.28.0.dev+141.g03a378f`, I get the following traceback:

```
Traceback (most recent call last):
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 237, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 410, in lower_assign
    res = self.lower_yield(ty, value)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 420, in lower_yield
    y.lower_yield_suspend()
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 320, in lower_yield_suspend
    val = self.lower.loadvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 875, in loadvar
    ptr = self.getvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 869, in getvar
    return self.varmap[name]
KeyError: 'p'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "gen_loop_branch_1259.py", line 11, in <module>
    f()
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 281, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 527, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 76, in compile
    flags=flags, locals=self.locals)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 725, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 369, in compile_extra
    return self.compile_bytecode(bc, func_attr=self.func_attr)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 378, in compile_bytecode
    return self._compile_bytecode()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 690, in _compile_bytecode
    return self._compile_core()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 677, in _compile_core
    res = pm.run(self.status)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 257, in run
    raise patched_exception
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 249, in run
    stage()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 617, in stage_nopython_backend
    self._backend(lowerfn, objectmode=False)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 572, in _backend
    lowered = lowerfn()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 559, in backend_nopython_mode
    self.flags)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 845, in native_lowering_stage
    lower.lower()
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 131, in lower
    self.genlower.lower_next_func(self)
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 149, in lower_next_func
    entry_block_tail = lower.lower_function_body()
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 185, in lower_function_body
    self.lower_block(block)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/miniconda3/envs/numba-git/lib/python3.5/contextlib.py", line 77, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 249, in new_error_context
    six.reraise(type(newerr), newerr, sys.exc_info()[2])
  File "/home/maba/dev/numba/git/numba/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 237, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 410, in lower_assign
    res = self.lower_yield(ty, value)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 420, in lower_yield
    y.lower_yield_suspend()
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 320, in lower_yield_suspend
    val = self.lower.loadvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 875, in loadvar
    ptr = self.getvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 869, in getvar
    return self.varmap[name]
numba.errors.LoweringError: Failed at nopython (nopython mode backend)
'p'
File "gen_loop_branch_1259.py", line 6
[1] During: lowering "$3.2 = yield $const3.1" at gen_loop_branch_1259.py (6)
```

for file `gen_loop_branch_1259.py`:

``` python
from numba import njit

@njit
def f():
    while True:
        yield
        p = True
        if p:
            pass

f()
```

@sklam: thoughts?

The compiler is incorrectly thinking that `p` is alive at the `yield` and is trying to suspend it.

The reason for this bug:

When a variable is used in a branch, its lifetime must at least be extended to the start of the branch targets.  This extends the lifetime back to the loop body (at `yield`) due to `branch p, <if body>, <while body>`.

Any idea how hard this would be to fix?
With current master, i.e `0.28.0.dev+141.g03a378f`, I get the following traceback:

```
Traceback (most recent call last):
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 237, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 410, in lower_assign
    res = self.lower_yield(ty, value)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 420, in lower_yield
    y.lower_yield_suspend()
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 320, in lower_yield_suspend
    val = self.lower.loadvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 875, in loadvar
    ptr = self.getvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 869, in getvar
    return self.varmap[name]
KeyError: 'p'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "gen_loop_branch_1259.py", line 11, in <module>
    f()
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 281, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 527, in compile
    cres = self._compiler.compile(args, return_type)
  File "/home/maba/dev/numba/git/numba/numba/dispatcher.py", line 76, in compile
    flags=flags, locals=self.locals)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 725, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 369, in compile_extra
    return self.compile_bytecode(bc, func_attr=self.func_attr)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 378, in compile_bytecode
    return self._compile_bytecode()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 690, in _compile_bytecode
    return self._compile_core()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 677, in _compile_core
    res = pm.run(self.status)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 257, in run
    raise patched_exception
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 249, in run
    stage()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 617, in stage_nopython_backend
    self._backend(lowerfn, objectmode=False)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 572, in _backend
    lowered = lowerfn()
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 559, in backend_nopython_mode
    self.flags)
  File "/home/maba/dev/numba/git/numba/numba/compiler.py", line 845, in native_lowering_stage
    lower.lower()
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 131, in lower
    self.genlower.lower_next_func(self)
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 149, in lower_next_func
    entry_block_tail = lower.lower_function_body()
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 185, in lower_function_body
    self.lower_block(block)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/miniconda3/envs/numba-git/lib/python3.5/contextlib.py", line 77, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 249, in new_error_context
    six.reraise(type(newerr), newerr, sys.exc_info()[2])
  File "/home/maba/dev/numba/git/numba/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
  File "/home/maba/dev/numba/git/numba/numba/errors.py", line 243, in new_error_context
    yield
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 200, in lower_block
    self.lower_inst(inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 237, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 410, in lower_assign
    res = self.lower_yield(ty, value)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 420, in lower_yield
    y.lower_yield_suspend()
  File "/home/maba/dev/numba/git/numba/numba/generators.py", line 320, in lower_yield_suspend
    val = self.lower.loadvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 875, in loadvar
    ptr = self.getvar(name)
  File "/home/maba/dev/numba/git/numba/numba/lowering.py", line 869, in getvar
    return self.varmap[name]
numba.errors.LoweringError: Failed at nopython (nopython mode backend)
'p'
File "gen_loop_branch_1259.py", line 6
[1] During: lowering "$3.2 = yield $const3.1" at gen_loop_branch_1259.py (6)
```

for file `gen_loop_branch_1259.py`:

``` python
from numba import njit

@njit
def f():
    while True:
        yield
        p = True
        if p:
            pass

f()
```

@sklam: thoughts?

The compiler is incorrectly thinking that `p` is alive at the `yield` and is trying to suspend it.

The reason for this bug:

When a variable is used in a branch, its lifetime must at least be extended to the start of the branch targets.  This extends the lifetime back to the loop body (at `yield`) due to `branch p, <if body>, <while body>`.


Resolving patch:
diff --git a/numba/interpreter.py b/numba/interpreter.py
--- a/numba/interpreter.py
+++ b/numba/interpreter.py
@@ -225,6 +225,10 @@ def _post_processing(self):
         """
         Post-processing and analysis on generated IR
         """
+        from . import transforms
+
+        self.blocks = transforms.canonicalize_cfg(self.blocks)
+
         # emit del nodes
         self._insert_var_dels()
 
diff --git a/numba/transforms.py b/numba/transforms.py
--- a/numba/transforms.py
+++ b/numba/transforms.py
@@ -196,3 +196,80 @@ def loop_lifting(interp, typingctx, targetctx, flags, locals):
 
     return main, loops
 
+
+def canonicalize_cfg_single_backedge(blocks):
+    """
+    Rewrite loops that have multiple backedges.
+    """
+    cfg = compute_cfg_from_blocks(blocks)
+    newblocks = blocks.copy()
+
+    def new_block_id():
+        return max(newblocks.keys()) + 1
+
+    def has_multiple_backedges(loop):
+        count = 0
+        for k in loop.body:
+            blk = blocks[k]
+            edges = blk.terminator.get_targets()
+            # is a backedge?
+            if loop.header in edges:
+                count += 1
+                if count > 1:
+                    # early exit
+                    return True
+        return False
+
+    def yield_loops_with_multiple_backedges():
+        for lp in cfg.loops().values():
+            if has_multiple_backedges(lp):
+                yield lp
+
+    def replace_target(term, src, dst):
+        def replace(target):
+            return (dst if target == src else target)
+
+        if isinstance(term, ir.Branch):
+            return ir.Branch(cond=term.cond,
+                             truebr=replace(term.truebr),
+                             falsebr=replace(term.falsebr),
+                             loc=term.loc)
+        elif isinstance(term, ir.Jump):
+            return ir.Jump(target=replace(term.target), loc=term.loc)
+        else:
+            assert not term.get_targets()
+            return term
+
+    def rewrite_single_backedge(loop):
+        """
+        Add new tail block that gathers all the backedges
+        """
+        header = loop.header
+        tailkey = new_block_id()
+        for blkkey in loop.body:
+            blk = newblocks[blkkey]
+            if header in blk.terminator.get_targets():
+                newblk = blk.copy()
+                # rewrite backedge into jumps to new tail block
+                newblk.body[-1] = replace_target(blk.terminator, header,
+                                                 tailkey)
+                newblocks[blkkey] = newblk
+        # create new tail block
+        entryblk = newblocks[header]
+        tailblk = ir.Block(scope=entryblk.scope, loc=entryblk.loc)
+        # add backedge
+        tailblk.append(ir.Jump(target=header, loc=tailblk.loc))
+        newblocks[tailkey] = tailblk
+
+    for loop in yield_loops_with_multiple_backedges():
+        rewrite_single_backedge(loop)
+
+    return newblocks
+
+
+def canonicalize_cfg(blocks):
+    """
+    Rewrite the given blocks to canonicalize the CFG.
+    Returns a new dictionary of blocks.
+    """
+    return canonicalize_cfg_single_backedge(blocks)


Test patch:
diff --git a/numba/tests/test_obj_lifetime.py b/numba/tests/test_obj_lifetime.py
--- a/numba/tests/test_obj_lifetime.py
+++ b/numba/tests/test_obj_lifetime.py
@@ -223,6 +223,20 @@ def del_before_definition(rec):
     return -1
 
 
+def inf_loop_multiple_back_edge(rec):
+    """
+    test to reveal bug of invalid liveness when infinite loop has multiple
+    backedge.
+    """
+    while True:
+        rec.mark("yield")
+        yield
+        p = rec('p')
+        if p:
+            rec.mark('bra')
+            pass
+
+
 class TestObjLifetime(TestCase):
     """
     Test lifetime of Python objects inside jit-compiled functions.
@@ -355,6 +369,19 @@ def test_raising3(self):
             rec = self.compile_and_record(raising_usecase3, raises=MyError)
             self.assertFalse(rec.alive)
 
+    def test_inf_loop_multiple_back_edge(self):
+        cfunc = self.compile(inf_loop_multiple_back_edge)
+        rec = RefRecorder()
+        iterator = iter(cfunc(rec))
+        next(iterator)
+        self.assertEqual(rec.alive, [])
+        next(iterator)
+        self.assertEqual(rec.alive, [])
+        next(iterator)
+        self.assertEqual(rec.alive, [])
+        self.assertEqual(rec.recorded,
+                         ['yield', 'p', 'bra', 'yield', 'p', 'bra', 'yield'])
+
 
 if __name__ == "__main__":
     unittest.main()