Repository: numba/numba

Issue:
is None branch incorrectly pruned in numba 0.43!
The following code does not execute the is None branch in numba 0.43 on Mac OSX:

~~~python
import numpy as np
import numba

@numba.jit
def bug(a,b):
    if a.ndim == 1:
        if b is None:
            # FIXME Numba 0.43 - this branch is not executed when run using numba 0.43!
            print('should get here')
            return None
        print('************** BUG')

    return []

bug(np.arange(10),None)
~~~

Obviously, this is **VERY BAD**.

Here is my conda environment:

~~~sh
$ conda list
# packages in environment at /Users/cbarber/miniconda3/envs/numba-bug:
#
# Name                    Version                   Build  Channel
blas                      1.1                    openblas    
ca-certificates           2019.3.9             hecc5488_0    
certifi                   2019.3.9                 py36_0    
intel-openmp              2019.1                      144  
libblas                   3.8.0                     4_mkl    
libcblas                  3.8.0                     4_mkl    
libcxx                    7.0.0                h2d50403_1    
libffi                    3.2.1             h0a44026_1005    
libgfortran               3.0.1                         0    
liblapack                 3.8.0                     4_mkl    
llvm-meta                 7.0.0                         0    
llvmlite                  0.28.0           py36hf292337_0    
mkl                       2019.1                      144  
ncurses                   6.1               h0a44026_1002    
numba                     0.43.0           py36h0a44026_0    h
numpy                     1.15.4          py36hbb3c62a_1002    
openblas                  0.3.3             hdc02c5d_1001    
openssl                   1.1.1b               h1de35cc_1    
pip                       19.0.3                   py36_0    
python                    3.6.7             h8dc6b48_1004    
readline                  7.0               hcfe32e1_1001    
scipy                     1.2.1           py36_blas_openblash486cb9f_0    
setuptools                40.8.0                   py36_0    
sqlite                    3.26.0            h1765d9f_1001    
tk                        8.6.9             ha441bb4_1000    
wheel                     0.33.1                   py36_0    
xz                        5.2.4             h1de35cc_1001    
zlib                      1.2.11            h1de35cc_1004    
~~~



Discussion:

Patch:
diff --git a/numba/analysis.py b/numba/analysis.py
index 08449494c..add81bacf 100644
--- a/numba/analysis.py
+++ b/numba/analysis.py
@@ -7,7 +7,7 @@ from collections import namedtuple, defaultdict
 
 from numba import ir
 from numba.controlflow import CFGraph
-from numba import types
+from numba import types, consts
 
 #
 # Analysis related to variable lifetime
@@ -293,6 +293,7 @@ def dead_branch_prune(func_ir, called_args):
         # replace the branch with a direct jump
         jmp = ir.Jump(keep, loc=branch.loc)
         blk.body[-1] = jmp
+        return 1 if keep == branch.truebr else 0
 
     def prune_by_type(branch, condition, blk, *conds):
         # this prunes a given branch and fixes up the IR
@@ -306,9 +307,9 @@ def dead_branch_prune(func_ir, called_args):
                 kill = branch.falsebr if take_truebr else branch.truebr
                 print("Pruning %s" % kill, branch, lhs_cond, rhs_cond,
                       condition.fn)
-            do_prune(take_truebr, blk)
-            return True
-        return False
+            taken = do_prune(take_truebr, blk)
+            return True, taken
+        return False, None
 
     def prune_by_value(branch, condition, blk, *conds):
         lhs_cond, rhs_cond = conds
@@ -316,8 +317,8 @@ def dead_branch_prune(func_ir, called_args):
         if DEBUG > 0:
             kill = branch.falsebr if take_truebr else branch.truebr
             print("Pruning %s" % kill, branch, lhs_cond, rhs_cond, condition.fn)
-        do_prune(take_truebr, blk)
-        return True
+        taken = do_prune(take_truebr, blk)
+        return True, taken
 
     class Unknown(object):
         pass
@@ -382,9 +383,10 @@ def dead_branch_prune(func_ir, called_args):
             # lhs/rhs are consts
             if len(const_conds) == 2:
                 # prune the branch, switch the branch for an unconditional jump
-                if(prune(branch, condition, blk, *const_conds)):
+                prune_stat, taken = prune(branch, condition, blk, *const_conds)
+                if(prune_stat):
                     # add the condition to the list of nullified conditions
-                    nullified_conditions.append(condition)
+                    nullified_conditions.append((condition, taken))
 
     # 'ERE BE DRAGONS...
     # It is the evaluation of the condition expression that often trips up type
@@ -395,18 +397,34 @@ def dead_branch_prune(func_ir, called_args):
     # information to run DCE safely. Upshot of all this is the condition gets
     # rewritten below into a benign const that typing will be happy with and DCE
     # can remove it and its reference post typing when it is safe to do so
-    # (if desired).
-    for _, condition, blk in branch_info:
-        if condition in nullified_conditions:
+    # (if desired). It is required that the const is assigned a value that
+    # indicates the branch taken as its mutated value would be read in the case
+    # of object mode fall back in place of the condition itself. For
+    # completeness the func_ir._definitions and ._consts are also updated to
+    # make the IR state self consistent.
+
+    deadcond = [x[0] for x in nullified_conditions]
+    for _, cond, blk in branch_info:
+        if cond in deadcond:
             for x in blk.body:
-                if isinstance(x, ir.Assign) and x.value is condition:
-                    x.value = ir.Const(0, loc=x.loc)
+                if isinstance(x, ir.Assign) and x.value is cond:
+                    # rewrite the condition as a true/false bit
+                    branch_bit = nullified_conditions[deadcond.index(cond)][1]
+                    x.value = ir.Const(branch_bit, loc=x.loc)
+                    # update the specific definition to the new const
+                    defns = func_ir._definitions[x.target.name]
+                    repl_idx = defns.index(cond)
+                    defns[repl_idx] = x.value
 
     # Remove dead blocks, this is safe as it relies on the CFG only.
     cfg = compute_cfg_from_blocks(func_ir.blocks)
     for dead in cfg.dead_nodes():
         del func_ir.blocks[dead]
 
+    # if conditions were nullified then consts were rewritten, update
+    if nullified_conditions:
+        func_ir._consts = consts.ConstantInference(func_ir)
+
     if DEBUG > 1:
         print("after".center(80, '-'))
         print(func_ir.dump())


Test patch:
diff --git a/numba/tests/test_analysis.py b/numba/tests/test_analysis.py
index a19bf30df..062def4f9 100644
--- a/numba/tests/test_analysis.py
+++ b/numba/tests/test_analysis.py
@@ -1,8 +1,9 @@
 # Tests numba.analysis functions
 from __future__ import print_function, absolute_import, division
 
+import numpy as np
 from numba.compiler import compile_isolated, run_frontend
-from numba import types, rewrites, ir
+from numba import types, rewrites, ir, jit, ir_utils
 from .support import TestCase, MemoryLeakMixin
 
 
@@ -33,6 +34,14 @@ class TestBranchPrune(MemoryLeakMixin, TestCase):
     """
     _DEBUG = False
 
+    # find *all* branches
+    def find_branches(self, the_ir):
+        branches = []
+        for blk in the_ir.blocks.values():
+            tmp = [_ for _ in blk.find_insts(cls=ir.Branch)]
+            branches.extend(tmp)
+        return branches
+
     def assert_prune(self, func, args_tys, prune, *args):
         # This checks that the expected pruned branches have indeed been pruned.
         # func is a python function to assess
@@ -59,15 +68,7 @@ class TestBranchPrune(MemoryLeakMixin, TestCase):
             print("after prune")
             func_ir.dump()
 
-        # find *all* branches
-        def find_branches(the_ir):
-            branches = []
-            for blk in the_ir.blocks.values():
-                tmp = [_ for _ in blk.find_insts(cls=ir.Branch)]
-                branches.extend(tmp)
-            return branches
-
-        before_branches = find_branches(before)
+        before_branches = self.find_branches(before)
         self.assertEqual(len(before_branches), len(prune))
 
         # what is expected to be pruned
@@ -317,3 +318,61 @@ class TestBranchPrune(MemoryLeakMixin, TestCase):
                           1000)
         self.assert_prune(impl, (types.IntegerLiteral(0),), [None, None], 0)
         self.assert_prune(impl, (types.NoneType('none'),), [True, False], None)
+
+    def test_cond_rewrite_is_correct(self):
+        # this checks that when a condition is replaced, it is replace by a
+        # true/false bit that correctly represents the evaluated condition
+        def fn(x):
+            if x is None:
+                return 10
+            return 12
+
+        def check(func, arg_tys, bit_val):
+            func_ir = compile_to_ir(func)
+
+            # check there is 1 branch
+            before_branches = self.find_branches(func_ir)
+            self.assertEqual(len(before_branches), 1)
+
+            # check the condition in the branch is a binop
+            condition_var = before_branches[0].cond
+            condition_defn = ir_utils.get_definition(func_ir, condition_var)
+            self.assertEqual(condition_defn.op, 'binop')
+
+            # do the prune, this should kill the dead branch and rewrite the
+            #'condition to a true/false const bit
+            if self._DEBUG:
+                print("=" * 80)
+                print("before prune")
+                func_ir.dump()
+            dead_branch_prune(func_ir, arg_tys)
+            if self._DEBUG:
+                print("=" * 80)
+                print("after prune")
+                func_ir.dump()
+
+            # after mutation, the condition should be a const value `bit_val`
+            new_condition_defn = ir_utils.get_definition(func_ir, condition_var)
+            self.assertTrue(isinstance(new_condition_defn, ir.Const))
+            self.assertEqual(new_condition_defn.value, bit_val)
+
+        check(fn, (types.NoneType('none'),), 1)
+        check(fn, (types.IntegerLiteral(10),), 0)
+
+    def test_obj_mode_fallback(self):
+        # see issue #3879, this checks that object mode fall back doesn't suffer
+        # from the IR mutation
+
+        @jit
+        def bug(a,b):
+            if a.ndim == 1:
+                if b is None:
+                    return 10
+                return 12
+            return []
+
+        self.assertEqual(bug(np.arange(10), 4), 12)
+        self.assertEqual(bug(np.arange(10), None), 10)
+        self.assertEqual(bug(np.arange(10).reshape((2, 5)), 10), [])
+        self.assertEqual(bug(np.arange(10).reshape((2, 5)), None), [])
+        self.assertFalse(bug.nopython_signatures)
diff --git a/numba/tests/test_jitclasses.py b/numba/tests/test_jitclasses.py
index 77af33bef..a08656c89 100644
--- a/numba/tests/test_jitclasses.py
+++ b/numba/tests/test_jitclasses.py
@@ -20,7 +20,6 @@ from numba.errors import LoweringError
 if sys.version_info >= (3,):
     from .jitclass_usecases import TestClass1, TestClass2
 
-
 def _get_meminfo(box):
     ptr = _box.box_get_meminfoptr(box)
     mi = MemInfo(ptr)
@@ -868,5 +867,6 @@ class TestJitClass(TestCase, MemoryLeakMixin):
         self.assertEqual(t[6:6:1], 3)
 
 
+
 if __name__ == '__main__':
     unittest.main()