gh-144140: Optimize len for frozen dict/set constants in optimizer (#149969)

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 2f606c2c6eba2d682a87413e0a9de5b131f57afe..790e965d6e5ff2f68751719159105c7afbe9bee9 100644 (file)
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -3154,7 +3154,7 @@ class TestUopsOptimization(unittest.TestCase):
         uops = get_opnames(ex)
         self.assertNotIn("_CHECK_IS_NOT_PY_CALLABLE_KW", uops)
 
-    def test_call_len_string(self):
+    def test_call_len_string_frozen_set_dict(self):
         def testfunc(n):
             for _ in range(n):
                 _ = len("abc")
@@ -3162,12 +3162,14 @@ class TestUopsOptimization(unittest.TestCase):
                 _ = len(d)
                 _ = len(b"def")
                 _ = len(b"")
+                _ = len(FROZEN_SET_CONST)
+                _ = len(FROZEN_DICT_CONST)
 
         _, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertNotIn("_CALL_LEN", uops)
-        self.assertGreaterEqual(count_ops(ex, "_LOAD_CONST_INLINE_BORROW"), 8)
+        self.assertGreaterEqual(count_ops(ex, "_LOAD_CONST_INLINE_BORROW"), 10)
 
     def test_call_len_known_length_small_int(self):
         # Make sure that len(t) is optimized for a tuple of length 5.

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 96dbaea5a5797ef1e8bf794c72e1cc07525d12d2..c968185d77c3317f07d93ad2dce1974626b55afa 100644 (file)
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -2378,7 +2378,7 @@ dummy_func(void) {
         res = sym_new_type(ctx, &PyLong_Type);
         Py_ssize_t length = sym_tuple_length(arg);
 
-        // Not a tuple, check if it's a const string
+        // Not a tuple, check if it's another immutable const with known length
         if (length < 0 && sym_is_const(ctx, arg)) {
             PyObject *const_val = sym_get_const(ctx, arg);
             if (const_val != NULL) {
@@ -2388,6 +2388,12 @@ dummy_func(void) {
                 else if (PyBytes_CheckExact(const_val)) {
                     length = PyBytes_GET_SIZE(const_val);
                 }
+                else if (PyFrozenDict_CheckExact(const_val)) {
+                    length = PyDict_GET_SIZE(const_val);
+                }
+                else if (PyFrozenSet_CheckExact(const_val)) {
+                    length = PySet_GET_SIZE(const_val);
+                }
             }
         }
 

diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index f336549d2ed24407f95c5ade7940f450f474ed0e..d52ebb9804197da993d0da57b05af2041f53f565 100644 (file)
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -4636,12 +4636,13 @@
                         length = PyUnicode_GET_LENGTH(const_val);
                     }
                     else if (PyBytes_CheckExact(const_val)) {
-                        CHECK_STACK_BOUNDS(-2);
-                        stack_pointer[-3] = res;
-                        stack_pointer += -2;
-                        ASSERT_WITHIN_STACK_BOUNDS(__FILE__, __LINE__);
                         length = PyBytes_GET_SIZE(const_val);
-                        stack_pointer += 2;
+                    }
+                    else if (PyFrozenDict_CheckExact(const_val)) {
+                        length = PyDict_GET_SIZE(const_val);
+                    }
+                    else if (PyFrozenSet_CheckExact(const_val)) {
+                        length = PySet_GET_SIZE(const_val);
                     }
                 }
             }

diff --git a/Tools/cases_generator/analyzer.py b/Tools/cases_generator/analyzer.py
index 59bca201a947e355aff227aecc1cc57904660c4d..6f0ddeaeaabf098a104deedc4c6c37ed4b37ac57 100644 (file)
--- a/Tools/cases_generator/analyzer.py
+++ b/Tools/cases_generator/analyzer.py
@@ -616,6 +616,9 @@ NON_ESCAPING_FUNCTIONS = (
     "PyStackRef_RefcountOnObject",
     "PyStackRef_TYPE",
     "PyStackRef_True",
+    "PyBytes_GET_SIZE",
+    "PyDict_GET_SIZE",
+    "PySet_GET_SIZE",
     "PyTuple_GET_ITEM",
     "PyTuple_GET_SIZE",
     "PyType_HasFeature",