gh-130415: Narrowing to constants in branches involving is comparisons with a constant

Include/internal/pycore_optimizer.h

@@ -205,6 +205,8 @@ extern JitOptRef _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef value,
 extern bool _Py_uop_sym_is_compact_int(JitOptRef sym);
 extern JitOptRef _Py_uop_sym_new_compact_int(JitOptContext *ctx);
 extern void _Py_uop_sym_set_compact_int(JitOptContext *ctx,  JitOptRef sym);
+extern JitOptRef _Py_uop_sym_new_predicate(JitOptContext *ctx, JitOptRef lhs_ref, JitOptRef rhs_ref, JitOptPredicateKind kind);
+extern void _Py_uop_sym_apply_predicate_narrowing(JitOptContext *ctx, JitOptRef sym, bool branch_is_true);
 
 extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);

Include/internal/pycore_optimizer_types.h

@@ -40,6 +40,7 @@ typedef enum _JitSymType {
     JIT_SYM_TUPLE_TAG = 8,
     JIT_SYM_TRUTHINESS_TAG = 9,
     JIT_SYM_COMPACT_INT = 10,
+    JIT_SYM_PREDICATE_TAG = 11,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -72,6 +73,18 @@ typedef struct {
     uint16_t value;
 } JitOptTruthiness;
 
+typedef enum {
+    JIT_PRED_IS,
+    JIT_PRED_IS_NOT,
+} JitOptPredicateKind;
+
+typedef struct {
+    uint8_t tag;
+    uint8_t kind;
+    uint16_t lhs;
+    uint16_t rhs;
+} JitOptPredicate;
+
 typedef struct {
     uint8_t tag;
 } JitOptCompactInt;
@@ -84,6 +97,7 @@ typedef union _jit_opt_symbol {
     JitOptTuple tuple;
     JitOptTruthiness truthiness;
     JitOptCompactInt compact;
+    JitOptPredicate predicate;
 } JitOptSymbol;
 
 // This mimics the _PyStackRef API

Lib/test/test_capi/test_opt.py

@@ -3551,6 +3551,46 @@ def test_is_none(n):
         self.assertIn("_POP_TOP_NOP", uops)
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
+    def test_is_true_narrows_to_constant(self):
+        def f(n):
+            def return_true():
+                return True
+
+            hits = 0
+            v = return_true()
+            for i in range(n):
+                if v is True:
+                    hits += v + 1
+            return hits
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD * 2)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # v + 1 should be constant folded
+        self.assertNotIn("_BINARY_OP", uops)
+
+    def test_is_false_narrows_to_constant(self):
+        def f(n):
+            def return_false():
+                return False
+
+            hits = 0
+            v = return_false()
+            for i in range(n):
+                if v is False:
+                    hits += v + 1
+            return hits
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # v + 1 should be constant folded
+        self.assertNotIn("_BINARY_OP", uops)
+
     def test_for_iter_gen_frame(self):
         def f(n):
             for i in range(n):

Python/optimizer_analysis.c

@@ -247,6 +247,8 @@ add_op(JitOptContext *ctx, _PyUOpInstruction *this_instr,
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_predicate _Py_uop_sym_new_predicate
+#define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
 #define JUMP_TO_LABEL(label) goto label;
 

Python/optimizer_bytecodes.c

@@ -38,6 +38,8 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_predicate _Py_uop_sym_new_predicate
+#define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
 extern int
 optimize_to_bool(
@@ -533,7 +535,7 @@ dummy_func(void) {
     }
 
     op(_IS_OP, (left, right -- b, l, r)) {
-        b = sym_new_type(ctx, &PyBool_Type);
+        b = sym_new_predicate(ctx, left, right, (oparg ? JIT_PRED_IS_NOT : JIT_PRED_IS));
         l = left;
         r = right;
     }
@@ -1173,6 +1175,8 @@ dummy_func(void) {
     }
 
     op(_GUARD_IS_TRUE_POP, (flag -- )) {
+        sym_apply_predicate_narrowing(ctx, flag, true);
+
         if (sym_is_const(ctx, flag)) {
             PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
@@ -1191,6 +1195,8 @@ dummy_func(void) {
     }
 
     op(_GUARD_IS_FALSE_POP, (flag -- )) {
+        sym_apply_predicate_narrowing(ctx, flag, false);
+
         if (sym_is_const(ctx, flag)) {
             PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);