python · corona10 · Dec 19, 2024 · Dec 8, 2024 · Dec 8, 2024 · Dec 8, 2024
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
@@ -278,6 +278,9 @@ typedef int (*_py_validate_type)(PyTypeObject *);
 // and if the validation is passed, it will set the ``tp_version`` as valid
 // tp_version_tag from the ``ty``.
 extern int _PyType_Validate(PyTypeObject *ty, _py_validate_type validate, unsigned int *tp_version);
+extern PyObject *_PyType_GetItemFromCache(PyTypeObject *type);
+extern PyObject *_PyType_GetItemFromCacheWithVersion(PyTypeObject *type, uint32_t *version);
+extern int _PyType_CacheGetItemForSpecialization(PyTypeObject *type, PyObject *descriptor);
 
 #ifdef __cplusplus
 }

diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -609,7 +609,7 @@ def assert_races_do_not_crash(
             for writer in writers:
                 writer.join()
 
-    @requires_specialization
+    @requires_specialization_ft
     def test_binary_subscr_getitem(self):
         def get_items():
             class C:
@@ -1069,7 +1069,7 @@ def write(items):
         opname = "STORE_SUBSCR_LIST_INT"
         self.assert_races_do_not_crash(opname, get_items, read, write)
 
-    @requires_specialization
+    @requires_specialization_ft
     def test_unpack_sequence_list(self):
         def get_items():
             items = []
@@ -1520,6 +1520,20 @@ def binary_subscr_str_int():
         self.assert_specialized(binary_subscr_str_int, "BINARY_SUBSCR_STR_INT")
         self.assert_no_opcode(binary_subscr_str_int, "BINARY_SUBSCR")
 
+        def binary_subscr_getitems():
+            class C:
+                def __init__(self, val):
+                    self.val = val
+                def __getitem__(self, item):
+                    return self.val
+            items = [C(i) for i in range(100)]
+            for i in range(100):
+                self.assertEqual(items[i][i], i)
+
+        binary_subscr_getitems()
+        self.assert_specialized(binary_subscr_getitems, "BINARY_SUBSCR_GETITEM")
+        self.assert_no_opcode(binary_subscr_getitems, "BINARY_SUBSCR")
+
 
 if __name__ == "__main__":
     unittest.main()
@@ -5692,6 +5692,64 @@ _PyType_CacheInitForSpecialization(PyHeapTypeObject *type, PyObject *init,
     return can_cache;
 }
 
+int
+_PyType_CacheGetItemForSpecialization(PyTypeObject *type, PyObject *descriptor)
+{
+    int ret = 0;
+    if (!type) {
+        return -1;
+    }
+    BEGIN_TYPE_LOCK();
+    // This pointer is invalidated by PyType_Modified (see the comment on
+    // struct _specialization_cache):
+    PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
+    PyFunctionObject *func = (PyFunctionObject *)descriptor;
+    uint32_t version = _PyFunction_GetVersionForCurrentState(func);
+    if (!_PyFunction_IsVersionValid(version)) {
+        ret = -1;
+        goto end;
+    }
+    FT_ATOMIC_STORE_PTR_RELEASE(ht->_spec_cache.getitem, descriptor);
+    FT_ATOMIC_STORE_UINT32_RELAXED(ht->_spec_cache.getitem_version, version);
+end:
+    END_TYPE_LOCK();
+    return ret;
+}
+
+PyObject *
+_PyType_GetItemFromCache(PyTypeObject *type)
+{
+    PyObject *res = NULL;
+    BEGIN_TYPE_LOCK();
+    PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
+    res = ht->_spec_cache.getitem;
+    if (res == NULL || !PyFunction_Check(res)) {
+        res = NULL;
+    }
+    END_TYPE_LOCK();
+    return res;
+}
+
+PyObject *
+_PyType_GetItemFromCacheWithVersion(PyTypeObject *type, uint32_t *version)
+{
+    PyObject *res = NULL;
+    BEGIN_TYPE_LOCK();
+    PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
+    res = ht->_spec_cache.getitem;
+    if (res == NULL) {
+        goto end;
+    }
+    if (!PyFunction_Check(res)) {
+        res = NULL;
+        goto end;
+    }
+    *version = ht->_spec_cache.getitem_version;
+end:
+    END_TYPE_LOCK();
+    return res;
+}
+
 static void
 set_flags(PyTypeObject *self, unsigned long mask, unsigned long flags)
 {

@@ -867,11 +867,10 @@ dummy_func(
         op(_BINARY_SUBSCR_CHECK_FUNC, (container, unused -- container, unused)) {
             PyTypeObject *tp = Py_TYPE(PyStackRef_AsPyObjectBorrow(container));
             DEOPT_IF(!PyType_HasFeature(tp, Py_TPFLAGS_HEAPTYPE));
-            PyHeapTypeObject *ht = (PyHeapTypeObject *)tp;
-            PyObject *getitem = ht->_spec_cache.getitem;
+            uint32_t cached_version;
+            PyObject *getitem = _PyType_GetItemFromCacheWithVersion(tp, &cached_version);
             DEOPT_IF(getitem == NULL);
             assert(PyFunction_Check(getitem));
-            uint32_t cached_version = ht->_spec_cache.getitem_version;
             DEOPT_IF(((PyFunctionObject *)getitem)->func_version != cached_version);
             PyCodeObject *code = (PyCodeObject *)PyFunction_GET_CODE(getitem);
             assert(code->co_argcount == 2);
@@ -881,8 +880,8 @@ dummy_func(
 
         op(_BINARY_SUBSCR_INIT_CALL, (container, sub -- new_frame: _PyInterpreterFrame* )) {
             PyTypeObject *tp = Py_TYPE(PyStackRef_AsPyObjectBorrow(container));
-            PyHeapTypeObject *ht = (PyHeapTypeObject *)tp;
-            PyObject *getitem = ht->_spec_cache.getitem;
+            PyObject *getitem = _PyType_GetItemFromCache(tp);
+            DEOPT_IF(getitem == NULL);
             new_frame = _PyFrame_PushUnchecked(tstate, PyStackRef_FromPyObjectNew(getitem), 2, frame);
             new_frame->localsplus[0] = container;
             new_frame->localsplus[1] = sub;

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -1085,6 +1085,7 @@ specialize_instance_load_attr(PyObject* owner, _Py_CODEUNIT* instr, PyObject* na
                 SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_ATTR_METHOD);
                 return -1;
             }
+            /* Don't specialize if PEP 523 is active */
             if (_PyInterpreterState_GET()->eval_frame) {
                 SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_OTHER);
                 return -1;
@@ -1154,6 +1155,7 @@ specialize_instance_load_attr(PyObject* owner, _Py_CODEUNIT* instr, PyObject* na
             if (version == 0) {
                 return -1;
             }
+            /* Don't specialize if PEP 523 is active */
             if (_PyInterpreterState_GET()->eval_frame) {
                 SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_OTHER);
                 return -1;
@@ -1761,9 +1763,7 @@ _Py_Specialize_BinarySubscr(
         specialized_op = BINARY_SUBSCR_DICT;
         goto success;
     }
-#ifndef Py_GIL_DISABLED
-    PyTypeObject *cls = Py_TYPE(container);
-    PyObject *descriptor = _PyType_Lookup(cls, &_Py_ID(__getitem__));
+    PyObject *descriptor = _PyType_Lookup(container_type, &_Py_ID(__getitem__));
     if (descriptor && Py_TYPE(descriptor) == &PyFunction_Type) {
         if (!(container_type->tp_flags & Py_TPFLAGS_HEAPTYPE)) {
             SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_SUBSCR_NOT_HEAP_TYPE);
@@ -1780,24 +1780,18 @@ _Py_Specialize_BinarySubscr(
             SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_WRONG_NUMBER_ARGUMENTS);
             goto fail;
         }
-        uint32_t version = _PyFunction_GetVersionForCurrentState(func);
-        if (!_PyFunction_IsVersionValid(version)) {
+        if (_PyType_CacheGetItemForSpecialization(container_type, descriptor) < 0) {
             SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_OUT_OF_VERSIONS);
             goto fail;
         }
+        /* Don't specialize if PEP 523 is active */
         if (_PyInterpreterState_GET()->eval_frame) {
             SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_OTHER);
             goto fail;
         }
-        PyHeapTypeObject *ht = (PyHeapTypeObject *)container_type;
-        // This pointer is invalidated by PyType_Modified (see the comment on
-        // struct _specialization_cache):
-        ht->_spec_cache.getitem = descriptor;
-        ht->_spec_cache.getitem_version = version;
         specialized_op = BINARY_SUBSCR_GETITEM;
         goto success;
     }
-#endif   // Py_GIL_DISABLED
     SPECIALIZATION_FAIL(BINARY_SUBSCR,
                         binary_subscr_fail_kind(container_type, sub));
 fail:
@@ -2606,6 +2600,7 @@ _Py_Specialize_ForIter(_PyStackRef iter, _Py_CODEUNIT *instr, int oparg)
         assert(instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == END_FOR  ||
             instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == INSTRUMENTED_END_FOR
         );
+        /* Don't specialize if PEP 523 is active */
         if (_PyInterpreterState_GET()->eval_frame) {
             SPECIALIZATION_FAIL(FOR_ITER, SPEC_FAIL_OTHER);
             goto failure;
@@ -2634,6 +2629,7 @@ _Py_Specialize_Send(_PyStackRef receiver_st, _Py_CODEUNIT *instr)
     assert(_PyOpcode_Caches[SEND] == INLINE_CACHE_ENTRIES_SEND);
     PyTypeObject *tp = Py_TYPE(receiver);
     if (tp == &PyGen_Type || tp == &PyCoro_Type) {
+        /* Don't specialize if PEP 523 is active */
         if (_PyInterpreterState_GET()->eval_frame) {
             SPECIALIZATION_FAIL(SEND, SPEC_FAIL_OTHER);
             goto failure;