Repository: numba/numba

Issue:
Miscompile on `main` as a result of #9711
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using ~the latest released version~ main branch of Numba (most recent is
  visible in the release notes
  (https://numba.readthedocs.io/en/stable/release-notes-overview.html).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

PR #9711 introduced the use of the branch pruning after the Numba-IR is converted to SSA form. Whilst the branch pruning pass "understands" SSA, as in, once it decides to do a transform it will ensure the transformed IR is SSA compliant, it does not correctly handle analysis of this use case:

```python
from numba import jit

CONST = 32

@jit
def foo():
    conditional = CONST // 2
    while conditional:
        print(conditional)
        conditional //= 2

foo()
```
it should print
```
16
8
4
2
1
```
but it currently (as of 717308e) prints:
```
1
```
The branch pruner analysis uses `get_definition` to find if there's a single definition of the predicate name, which in SSA form is always the case, hence it erroneously prunes the branch structure associated with the `while` and then strange things happen!

Discussion:

Patch:
diff --git a/docs/upcoming_changes/9758.bug_fix.rst b/docs/upcoming_changes/9758.bug_fix.rst
new file mode 100644
index 000000000..2fe87b946
--- /dev/null
+++ b/docs/upcoming_changes/9758.bug_fix.rst
@@ -0,0 +1,5 @@
+Fix miscompile in branch pruning of SSA form IR.
+------------------------------------------------
+
+A miscompile occurring when a binop expression with constant arguments is used
+as a predicate in SSA form IR is now fixed.
diff --git a/numba/core/analysis.py b/numba/core/analysis.py
index 1ea3322e6..4e6691b5b 100644
--- a/numba/core/analysis.py
+++ b/numba/core/analysis.py
@@ -354,8 +354,12 @@ 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)
-        taken = do_prune(take_truebr, blk)
-        return True, taken
+        do_prune(take_truebr, blk)
+        # It is not safe to rewrite the predicate to a nominal value based on
+        # which branch is taken, the rewritten const predicate needs to
+        # hold the actual computed const value as something else may refer to
+        # it!
+        return True, take_truebr
 
     def prune_by_predicate(branch, pred, blk):
         try:


Test patch:
diff --git a/numba/tests/test_remove_dead.py b/numba/tests/test_remove_dead.py
index 3eaab0f23..b1ab55ce2 100644
--- a/numba/tests/test_remove_dead.py
+++ b/numba/tests/test_remove_dead.py
@@ -5,7 +5,7 @@
 
 import numba
 import numba.parfors.parfor
-from numba import njit
+from numba import njit, jit
 from numba.core import ir_utils
 from numba.core import types, ir,  compiler
 from numba.core.registry import cpu_target
@@ -399,6 +399,37 @@ class TestSSADeadBranchPrune(TestCase):
 
         self.assertEqual(make(), make(lambda x: x))
 
+    def test_issue_9742(self):
+        CONST = 32
+
+        @jit
+        def foo():
+            # This is a prune by value case, conditional is a compile time
+            # evaluatable constant.
+            conditional = CONST // 2
+            collect = []
+            while conditional:
+                collect.append(conditional)
+                conditional //= 2
+
+            return collect
+
+        self.assertEqual(foo(), foo.py_func())
+
+    def test_issue_9742_variant(self):
+        CONST = 32
+
+        @jit
+        def foo():
+            collect = []
+            # This is a prune by value case, conditional is a compile time
+            # evaluatable constant.
+            x = CONST + 1
+            if x:
+                collect.append(x)
+            return collect
+
+        self.assertEqual(foo(), foo.py_func())
 
 
 if __name__ == "__main__":