GH-130415: Optimize JIT path for _TO_BOOL_INT branching

Lib/test/test_capi/test_opt.py

@@ -1437,6 +1437,67 @@ def crash_addition():
 
         crash_addition()
 
+    def test_narrow_type_to_constant_int_zero(self):
+        def f(n):
+            trace = []
+            for i in range(n):
+                # f is always (int) 0, but we can only prove that it's a integer:
+                f = i - i # at this point python knows f is an int, but doesnt know that it is 0 (we know it is 0 though)
+                trace.append("A")
+                if not f:  # Kept.
+                    trace.append("B")
+                    if not f:  # Removed!
+                        trace.append("C")
+                    trace.append("D")
+                    if f:  # Removed!
+                        trace.append("X")
+                    trace.append("E")
+                trace.append("F")
+                if f:  # Removed!
+                    trace.append("X")
+                trace.append("G")
+            return trace
+
+        trace, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(trace, list("ABCDEFG") * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        # Only one guard remains:
+        self.assertEqual(uops.count("_GUARD_IS_FALSE_POP"), 1)
+        self.assertEqual(uops.count("_GUARD_IS_TRUE_POP"), 0)
+        # But all of the appends we care about are still there:
+        self.assertEqual(uops.count("_CALL_LIST_APPEND"), len("ABCDEFG"))
+
+    # def test_narrow_type_to_constant_bool_true(self):
+    #     def f(n):
+    #         trace = []
+    #         for i in range(n):
+    #             # f is always True, but we can only prove that it's a bool:
+    #             f = i != TIER2_THRESHOLD
+    #             trace.append("A")
+    #             if f:  # Kept.
+    #                 trace.append("B")
+    #                 if not f:  # Removed!
+    #                     trace.append("X")
+    #                 trace.append("C")
+    #                 if f:  # Removed!
+    #                     trace.append("D")
+    #                 trace.append("E")
+    #             trace.append("F")
+    #             if not f:  # Removed!
+    #                 trace.append("X")
+    #             trace.append("G")
+    #         return trace
+
+    #     trace, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+    #     self.assertEqual(trace, list("ABCDEFG") * TIER2_THRESHOLD)
+    #     self.assertIsNotNone(ex)
+    #     uops = get_opnames(ex)
+    #     # Only one guard remains:
+    #     self.assertEqual(uops.count("_GUARD_IS_FALSE_POP"), 0)
+    #     self.assertEqual(uops.count("_GUARD_IS_TRUE_POP"), 1)
+    #     # But all of the appends we care about are still there:
+    #     self.assertEqual(uops.count("_CALL_LIST_APPEND"), len("ABCDEFG"))
 
 def global_identity(x):
     return x

Misc/NEWS.d/next/Core_and_Builtins/2025-02-22-23-03-13.gh-issue-130415.ibOV6B.rst

@@ -0,0 +1 @@
+Improve JIT understanding of integers in boolean context.

Python/optimizer_bytecodes.c

@@ -407,6 +407,23 @@ dummy_func(void) {
             sym_set_type(value, &PyLong_Type);
             res = sym_new_type(ctx, &PyBool_Type);
         }
+        if (!sym_is_const(value)) {
+            assert(sym_matches_type(value, &PyLong_Type));
+            int next_opcode = (this_instr + 1)->opcode;
+            assert(next_opcode == _CHECK_VALIDITY_AND_SET_IP);
+            next_opcode = (this_instr + 2)->opcode;
+            // If the next uop is a guard, we can narrow value. However, we
+            // *can't* narrow res, since that would cause the guard to be
+            // removed and the narrowed value to be invalid:
+            if (next_opcode == _GUARD_IS_FALSE_POP) {
+                sym_set_const(value, Py_GetConstant(Py_CONSTANT_ZERO));
+                res = sym_new_type(ctx, &PyLong_Type);
+            }
+            // else if (next_opcode == _GUARD_IS_TRUE_POP) {
+            //     sym_set_const(value, Py_True);
+            //     res = sym_new_type(ctx, &PyBool_Type);
+            // }
+        }
     }
 
     op(_TO_BOOL_LIST, (value -- res)) {