gh-146381: Constant-fold frozendict subscript lookups via REPLACE_OPCODE_IF_EVALUATES...

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 126bcedef99993da037733870f9692a0e523e0e3..60f5cb6edbce43bc60cc432fccafe59a17963f82 100644 (file)
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -20,6 +20,10 @@ from _testinternalcapi import _PY_NSMALLPOSINTS, TIER2_THRESHOLD, TIER2_RESUME_T
 #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 @@ class TestUopsOptimization(unittest.TestCase):
         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

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index f5d059dea3ef0bb73c0d3885fc3a1a56585d71a1..6702f3c56651331f443abe00518fdb6a43832e85 100644 (file)
--- a/Python/optimizer_bytecodes.c
+++ b/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)) {

diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index ba6817cfdf2c0b9228a533ccc53171824af51e59..dd1d2813a215ca4ee7bb164dce41567adcf9eb2d 100644 (file)
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -1327,6 +1327,54 @@
             res = sym_new_not_null(ctx);
             ds = dict_st;
             ss = sub_st;
+            if (sym_matches_type(dict_st, &PyFrozenDict_Type)) {
+                if (
+                    sym_is_safe_const(ctx, dict_st) &&
+                    sym_is_safe_const(ctx, sub_st)
+                ) {
+                    JitOptRef dict_st_sym = dict_st;
+                    JitOptRef sub_st_sym = sub_st;
+                    _PyStackRef dict_st = sym_get_const_as_stackref(ctx, dict_st_sym);
+                    _PyStackRef sub_st = sym_get_const_as_stackref(ctx, sub_st_sym);
+                    _PyStackRef res_stackref;
+                    _PyStackRef ds_stackref;
+                    _PyStackRef ss_stackref;
+                    /* Start of uop copied from bytecodes for constant evaluation */
+                    PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
+                    PyObject *dict = PyStackRef_AsPyObjectBorrow(dict_st);
+                    assert(PyAnyDict_CheckExact(dict));
+                    STAT_INC(BINARY_OP, hit);
+                    PyObject *res_o;
+                    int rc = PyDict_GetItemRef(dict, sub, &res_o);
+                    if (rc == 0) {
+                        _PyErr_SetKeyError(sub);
+                    }
+                    if (rc <= 0) {
+                        JUMP_TO_LABEL(error);
+                    }
+                    res_stackref = PyStackRef_FromPyObjectSteal(res_o);
+                    ds_stackref = dict_st;
+                    ss_stackref = sub_st;
+                    /* End of uop copied from bytecodes for constant evaluation */
+                    (void)ds_stackref;
+                    (void)ss_stackref;
+                    res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref));
+                    if (sym_is_const(ctx, res)) {
+                        PyObject *result = sym_get_const(ctx, res);
+                        if (_Py_IsImmortal(result)) {
+                            // Replace with _INSERT_2_LOAD_CONST_INLINE_BORROW since we have two inputs and an immortal result
+                            ADD_OP(_INSERT_2_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)result);
+                        }
+                    }
+                    CHECK_STACK_BOUNDS(1);
+                    stack_pointer[-2] = res;
+                    stack_pointer[-1] = ds;
+                    stack_pointer[0] = ss;
+                    stack_pointer += 1;
+                    ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
+                    break;
+                }
+            }
             CHECK_STACK_BOUNDS(1);
             stack_pointer[-2] = res;
             stack_pointer[-1] = ds;

diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
index 0bc3c5055812c4703ed3e6c325a0eede562a5410..d6f1c09490aac9d7398087d91a8a5c54f98c305f 100644 (file)
--- a/Python/optimizer_symbols.c
+++ b/Python/optimizer_symbols.c
@@ -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