gh-146381: Constant-fold frozendict subscript lookups via REPLACE_OPCODE_IF_EVALUATES_PURE (gh-146490)

Author: corona10

Lib/test/test_capi/test_opt.py

@@ -20,6 +20,10 @@
 #For test of issue 136154
 GLOBAL_136154 = 42
 
+# For frozendict JIT tests
+FROZEN_DICT_CONST = frozendict(x=1, y=2)
+
+
 @contextlib.contextmanager
 def clear_executors(func):
     # Clear executors in func before and after running a block
@@ -4155,6 +4159,35 @@ def testfunc(n):
         self.assertLessEqual(count_ops(ex, "_POP_TOP_INT"), 1)
         self.assertIn("_POP_TOP_NOP", uops)
 
+    def test_binary_subscr_frozendict_lowering(self):
+        def testfunc(n):
+            x = 0
+            for _ in range(n):
+                x += FROZEN_DICT_CONST['x']
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_INSERT_2_LOAD_CONST_INLINE_BORROW", uops)
+        self.assertNotIn("_BINARY_OP_SUBSCR_DICT", uops)
+
+    def test_binary_subscr_frozendict_const_fold(self):
+        def testfunc(n):
+            x = 0
+            for _ in range(n):
+                if FROZEN_DICT_CONST['x'] == 1:
+                    x += 1
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        # lookup result is folded to constant 1, so comparison is optimized away
+        self.assertNotIn("_COMPARE_OP_INT", uops)
+
     def test_binary_subscr_list_slice(self):
         def testfunc(n):
             x = 0

Python/optimizer_bytecodes.c

@@ -485,6 +485,9 @@ dummy_func(void) {
         res = sym_new_not_null(ctx);
         ds = dict_st;
         ss = sub_st;
+        if (sym_matches_type(dict_st, &PyFrozenDict_Type)) {
+            REPLACE_OPCODE_IF_EVALUATES_PURE(dict_st, sub_st, res);
+        }
     }
 
     op(_BINARY_OP_SUBSCR_LIST_SLICE, (list_st, sub_st -- res, ls, ss)) {

Python/optimizer_cases.c.h

@@ -282,7 +282,8 @@ _Py_uop_sym_is_safe_const(JitOptContext *ctx, JitOptRef sym)
     return (typ == &PyUnicode_Type) ||
            (typ == &PyFloat_Type) ||
            (typ == &_PyNone_Type) ||
-           (typ == &PyBool_Type);
+           (typ == &PyBool_Type) ||
+           (typ == &PyFrozenDict_Type);
 }
 
 void