WIP

Author: corona10

Include/internal/pycore_typeobject.h

@@ -278,7 +278,7 @@ 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 int _PyType_CacheGetItemForSpecialization(PyHeapTypeObject *ht, PyObject *descriptor);
+extern int _PyType_CacheGetItemForSpecialization(PyHeapTypeObject *ht, PyObject *descriptor, uint32_t version);
 
 #ifdef __cplusplus
 }

Objects/typeobject.c

@@ -5693,29 +5693,26 @@ _PyType_CacheInitForSpecialization(PyHeapTypeObject *type, PyObject *init,
 }
 
 int
-_PyType_CacheGetItemForSpecialization(PyHeapTypeObject *ht, PyObject *descriptor)
+_PyType_CacheGetItemForSpecialization(PyHeapTypeObject *ht, PyObject *descriptor, uint32_t version)
 {
-    int ret = 0;
+    if (!descriptor || !version) {
+        return 0;
+    }
+    int can_cache;
     BEGIN_TYPE_LOCK();
     // This pointer is invalidated by PyType_Modified (see the comment on
     // struct _specialization_cache):
     PyFunctionObject *func = (PyFunctionObject *)descriptor;
+    can_cache = _PyFunction_GetVersionForCurrentState(func) == version;
 #ifdef Py_GIL_DISABLED
-    if (!_PyObject_HasDeferredRefcount(descriptor)) {
-        ret = -1;
-        goto end;
-    }
+    can_cache = can_cache && _PyObject_HasDeferredRefcount(descriptor);
 #endif
-    uint32_t version = _PyFunction_GetVersionForCurrentState(func);
-    if (!_PyFunction_IsVersionValid(version)) {
-        ret = -1;
-        goto end;
+    if (can_cache) {
+        FT_ATOMIC_STORE_PTR_RELEASE(ht->_spec_cache.getitem, descriptor);
+        FT_ATOMIC_STORE_UINT32_RELAXED(ht->_spec_cache.getitem_version, version);
     }
-    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;
+    return can_cache;
 }
 
 static void

Python/specialize.c

@@ -1763,36 +1763,41 @@ _Py_Specialize_BinarySubscr(
         specialized_op = BINARY_SUBSCR_DICT;
         goto success;
     }
-    PyObject *descriptor = _PyType_Lookup(container_type, &_Py_ID(__getitem__));
+    unsigned int version;
+    PyObject *descriptor = _PyType_LookupRefAndVersion(container_type, &_Py_ID(__getitem__), &version);
     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);
+            Py_DECREF(descriptor);
             goto fail;
         }
         PyFunctionObject *func = (PyFunctionObject *)descriptor;
         PyCodeObject *fcode = (PyCodeObject *)func->func_code;
         int kind = function_kind(fcode);
         if (kind != SIMPLE_FUNCTION) {
             SPECIALIZATION_FAIL(BINARY_SUBSCR, kind);
+            Py_DECREF(descriptor);
             goto fail;
         }
         if (fcode->co_argcount != 2) {
             SPECIALIZATION_FAIL(BINARY_SUBSCR, SPEC_FAIL_WRONG_NUMBER_ARGUMENTS);
+            Py_DECREF(descriptor);
             goto fail;
         }
 
         PyHeapTypeObject *ht = (PyHeapTypeObject *)container_type;
-        if (_PyType_CacheGetItemForSpecialization(ht, 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);
+            Py_DECREF(descriptor);
             goto fail;
         }
-        specialized_op = BINARY_SUBSCR_GETITEM;
-        goto success;
+        if (_PyType_CacheGetItemForSpecialization(ht, descriptor, (uint32_t)version)) {
+            specialized_op = BINARY_SUBSCR_GETITEM;
+            Py_DECREF(descriptor);
+            goto success;
+        }
+        Py_DECREF(descriptor);
     }
     SPECIALIZATION_FAIL(BINARY_SUBSCR,
                         binary_subscr_fail_kind(container_type, sub));