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

Author: lgeiger

Lib/test/test_capi/test_opt.py

@@ -3154,20 +3154,22 @@ def testfunc(n):
         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")
                 d = ''
                 _ = 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.

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);
+                }
             }
         }
 

Python/optimizer_cases.c.h

@@ -616,6 +616,9 @@ def has_error_without_pop(op: parser.CodeDef) -> bool:
     "PyStackRef_RefcountOnObject",
     "PyStackRef_TYPE",
     "PyStackRef_True",
+    "PyBytes_GET_SIZE",
+    "PyDict_GET_SIZE",
+    "PySet_GET_SIZE",
     "PyTuple_GET_ITEM",
     "PyTuple_GET_SIZE",
     "PyType_HasFeature",