Repository: numba/numba

Issue:
Unexpected behaviour of `and` in guvectorize returned from factory function (bad specialisation of bool as int64?)
- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

Given the following code

```
import numba
import numpy as np


def mk_count_vote(min_samples):
    @numba.guvectorize(["int32[:], int32[:]", "int64[:], int64[:]"], '(n)->()', nopython=True)
    def count_vote(nbr_labels, res):
        max_elem = -1
        max_count = 0
        num_maxes = 0
        cur_elem = -1
        cur_count = 0

        def flush():
            nonlocal max_count, num_maxes, max_elem
            if cur_count > max_count:
                max_count = cur_count
                num_maxes = 1
                max_elem = cur_elem
            elif cur_count == max_count:
                num_maxes += 1

        for nbr_label in nbr_labels:
            if nbr_label == -1:
                break
            elif nbr_label != cur_elem:
                flush()
                cur_elem = nbr_label
                cur_count = 1
            else:
                cur_count += 1
        flush()

        # XXX: Changing the first part to bool(num_maxes == 1) fixes the problem
        if num_maxes == 1 and ((max_count - 1) >= min_samples):
            print("Setting to", max_elem)
            res[0] = max_elem
        else:
            print("Setting to -1")
            res[0] = -1

    return count_vote


ARR = np.array([[-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1,],
 [ 2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
  -1, -1, -1, -1, -1, -1, -1, -1,]], dtype=np.int64)

count_vote = mk_count_vote(3)
print(count_vote(ARR, axis=1))
```

I get:

```
Setting to -1
Setting to -1
[-1 -1]
```

If I go ahead and make the change to the XXX line then I get

```
Setting to -1
Setting to 2
[-1  2]
```

Which is what I would expect the original program to produce. Some tracing reveals that `(num_maxes == 1)` is `int64` while `(max_count - 1) >= min_samples` is `bool`, which is what prompted me to add the `bool(...)` casting. I suppose therefore that this is some kind of bad specialization.

I am running Debian GNU/Linux bullseye/sid and have checked with numba==0.52.0.

Discussion:

Patch:
diff --git a/numba/core/compiler.py b/numba/core/compiler.py
index 163292f9e..a0c03fdb3 100644
--- a/numba/core/compiler.py
+++ b/numba/core/compiler.py
@@ -514,14 +514,17 @@ class DefaultPassBuilder(object):
         pm.add_pass(IRProcessing, "processing IR")
         pm.add_pass(WithLifting, "Handle with contexts")
 
+        # inline closures early in case they are using nonlocal's
+        # see issue #6585.
+        pm.add_pass(InlineClosureLikes,
+                    "inline calls to locally defined closures")
+
         # pre typing
         if not state.flags.no_rewrites:
             pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants")
             pm.add_pass(DeadBranchPrune, "dead branch pruning")
             pm.add_pass(GenericRewrites, "nopython rewrites")
 
-        pm.add_pass(InlineClosureLikes,
-                    "inline calls to locally defined closures")
         # convert any remaining closures into functions
         pm.add_pass(MakeFunctionToJitFunction,
                     "convert make_function into JIT functions")


Test patch:
diff --git a/numba/tests/test_analysis.py b/numba/tests/test_analysis.py
index 35b73bfcc..b09a7c9e2 100644
--- a/numba/tests/test_analysis.py
+++ b/numba/tests/test_analysis.py
@@ -5,7 +5,9 @@ import types as pytypes
 import numpy as np
 from numba.core.compiler import compile_isolated, run_frontend, Flags, StateDict
 from numba import jit, njit
-from numba.core import types, errors, ir, rewrites, ir_utils, utils
+from numba.core import types, errors, ir, rewrites, ir_utils, utils, cpu
+from numba.core import postproc
+from numba.core.inline_closurecall import InlineClosureCallPass
 from numba.tests.support import TestCase, MemoryLeakMixin, SerialMixin
 
 from numba.core.analysis import dead_branch_prune, rewrite_semantic_constants
@@ -61,10 +63,24 @@ class TestBranchPruneBase(MemoryLeakMixin, TestCase):
         before = func_ir.copy()
         if self._DEBUG:
             print("=" * 80)
-            print("before prune")
+            print("before inline")
             func_ir.dump()
 
+        # run closure inlining to ensure that nonlocals in closures are visible
+        inline_pass = InlineClosureCallPass(func_ir,
+                                            cpu.ParallelOptions(False),)
+        inline_pass.run()
+
+        # Remove all Dels, and re-run postproc
+        post_proc = postproc.PostProcessor(func_ir)
+        post_proc.run()
+
         rewrite_semantic_constants(func_ir, args_tys)
+        if self._DEBUG:
+            print("=" * 80)
+            print("before prune")
+            func_ir.dump()
+
         dead_branch_prune(func_ir, args_tys)
 
         after = func_ir
@@ -566,6 +582,44 @@ class TestBranchPrune(TestBranchPruneBase, SerialMixin):
                           [None, None],
                           np.zeros((2, 3)), 1.2, None)
 
+    def test_closure_and_nonlocal_can_prune(self):
+        # Closures must be inlined ahead of branch pruning in case nonlocal
+        # is used. See issue #6585.
+        def impl():
+            x = 1000
+
+            def closure():
+                nonlocal x
+                x = 0
+
+            closure()
+
+            if x == 0:
+                return True
+            else:
+                return False
+
+        self.assert_prune(impl, (), [False,],)
+
+    def test_closure_and_nonlocal_cannot_prune(self):
+        # Closures must be inlined ahead of branch pruning in case nonlocal
+        # is used. See issue #6585.
+        def impl(n):
+            x = 1000
+
+            def closure(t):
+                nonlocal x
+                x = t
+
+            closure(n)
+
+            if x == 0:
+                return True
+            else:
+                return False
+
+        self.assert_prune(impl, (types.int64,), [None,], 1)
+
 
 class TestBranchPrunePredicates(TestBranchPruneBase, SerialMixin):
     # Really important thing to remember... the branch on predicates end up as