Add failing unit tests for str comparison

Author: reidenong

Lib/test/test_capi/test_opt.py

@@ -978,6 +978,50 @@ def return_tenth():
         # Constant narrowing allows constant folding for second comparison
         self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1)
 
+    def test_compare_str_eq_narrows_to_constant(self):
+        def f(n):
+            def return_hello():
+                return "hello"
+
+            hits = 0
+            v = return_hello()
+            for _ in range(n):
+                if v == "hello":
+                    if v == "hello":
+                        hits += 1
+            return hits
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # Constant narrowing allows constant folding for second comparison
+        self.assertLessEqual(count_ops(ex, "_COMPARE_OP_STR"), 1)
+
+    def test_compare_str_ne_narrows_to_constant(self):
+        def f(n):
+            def return_hello():
+                return "hello"
+
+            hits = 0
+            v = return_hello()
+            for _ in range(n):
+                if v != "hello":
+                    hits += 1000
+                else:
+                    if v == "hello":
+                        hits += 1
+            return hits
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # Constant narrowing allows constant folding for second comparison
+        self.assertLessEqual(count_ops(ex, "_COMPARE_OP_STR"), 1)
+
     @unittest.skip("gh-139109 WIP")
     def test_combine_stack_space_checks_sequential(self):
         def dummy12(x):

Python/optimizer_symbols.c

@@ -1419,6 +1419,58 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored))
     _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true);
     TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)");
 
+    // Test narrowing subject to constant from EQ predicate for str
+    subject = _Py_uop_sym_new_unknown(ctx);
+    PyObject *str_hello_obj = PyUnicode_FromString("hello");
+    JitOptRef const_str_hello = _Py_uop_sym_new_const(ctx, str_hello_obj);
+    if (PyJitRef_IsNull(subject) || str_hello_obj == NULL || PyJitRef_IsNull(const_str_hello)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_EQ);
+    if (PyJitRef_IsNull(ref)) {
+        goto fail;
+    }
+    _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true);
+    TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (str)");
+    TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == str_hello_obj, "predicate narrowing did not narrow subject to hello");
+
+    // Resolving EQ predicate to False should not narrow subject for str
+    subject = _Py_uop_sym_new_unknown(ctx);
+    if (PyJitRef_IsNull(subject)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_EQ);
+    if (PyJitRef_IsNull(ref)) {
+        goto fail;
+    }
+    _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false);
+    TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)");
+
+    // Test narrowing subject to constant from NE predicate for str
+    subject = _Py_uop_sym_new_unknown(ctx);
+    if (PyJitRef_IsNull(subject)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_NE);
+    if (PyJitRef_IsNull(ref)) {
+        goto fail;
+    }
+    _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false);
+    TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (str)");
+    TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == str_hello_obj, "predicate narrowing did not narrow subject to hello");
+
+    // Resolving NE predicate to true should not narrow subject for str
+    subject = _Py_uop_sym_new_unknown(ctx);
+    if (PyJitRef_IsNull(subject)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_NE);
+    if (PyJitRef_IsNull(ref)) {
+        goto fail;
+    }
+    _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true);
+    TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)");
+
     val_big = PyNumber_Lshift(_PyLong_GetOne(), PyLong_FromLong(66));
     if (val_big == NULL) {
         goto fail;