Refactor

python · kumaraditya303 · Jul 28, 2025 · Jul 7, 2025 · Jul 8, 2025 · Jul 9, 2025
commit bbc7fb836f0d2349a44aff4bed85d0b377bcecef
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -112,7 +112,8 @@ extern void _PyDictKeys_DecRef(PyDictKeysObject *keys);
 extern Py_ssize_t _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
 extern Py_ssize_t _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
 extern Py_ssize_t _Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr);
-extern Py_ssize_t _Py_dict_lookup_unicode_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr);
+
+extern int _PyDict_GetMethodStackRef(PyDictObject *dict, PyObject *name, _PyStackRef *method);
 
 extern Py_ssize_t _PyDict_LookupIndex(PyDictObject *, PyObject *);
 extern Py_ssize_t _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyObject *key);

@@ -839,6 +839,13 @@ _Py_TryXGetStackRef(PyObject **src, _PyStackRef *out)
 
 #endif
 
+#define PyStackRef_XSETREF(dst, src) \
+    do { \
+        _PyStackRef _tmp_dst_ref = (dst); \
+        (dst) = (src); \
+        PyStackRef_XCLOSE(_tmp_dst_ref); \
+    } while(0)
+
 // Like Py_VISIT but for _PyStackRef fields
 #define _Py_VISIT_STACKREF(ref)                                         \
     do {                                                                \

@@ -1581,84 +1581,51 @@ _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyOb
     return ix;
 }
 
-
-Py_ssize_t
-_Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
+static Py_ssize_t
+lookup_threadsafe_unicode(PyDictKeysObject *dk, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
 {
-    PyDictKeysObject *dk = _Py_atomic_load_ptr(&mp->ma_keys);
-    if (dk->dk_kind == DICT_KEYS_UNICODE && PyUnicode_CheckExact(key)) {
-        Py_ssize_t ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash);
-        if (ix == DKIX_EMPTY) {
+    assert(dk->dk_kind == DICT_KEYS_UNICODE);
+    assert(PyUnicode_CheckExact(key));
+
+    Py_ssize_t ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash);
+    if (ix == DKIX_EMPTY) {
+        *value_addr = PyStackRef_NULL;
+        return ix;
+    }
+    else if (ix >= 0) {
+        PyObject **addr_of_value = &DK_UNICODE_ENTRIES(dk)[ix].me_value;
+        PyObject *value = _Py_atomic_load_ptr(addr_of_value);
+        if (value == NULL) {
             *value_addr = PyStackRef_NULL;
+            return DKIX_EMPTY;
+        }
+        if (_PyObject_HasDeferredRefcount(value)) {
+            *value_addr =  (_PyStackRef){ .bits = (uintptr_t)value | Py_TAG_DEFERRED };
             return ix;
         }
-        else if (ix >= 0) {
-            PyObject **addr_of_value = &DK_UNICODE_ENTRIES(dk)[ix].me_value;
-            PyObject *value = _Py_atomic_load_ptr(addr_of_value);
-            if (value == NULL) {
-                *value_addr = PyStackRef_NULL;
-                return DKIX_EMPTY;
-            }
-            if (_PyObject_HasDeferredRefcount(value)) {
-                *value_addr =  (_PyStackRef){ .bits = (uintptr_t)value | Py_TAG_DEFERRED };
-                return ix;
-            }
-            if (_Py_TryIncrefCompare(addr_of_value, value)) {
-                *value_addr = PyStackRef_FromPyObjectSteal(value);
-                return ix;
-            }
+        if (_Py_TryIncrefCompare(addr_of_value, value)) {
+            *value_addr = PyStackRef_FromPyObjectSteal(value);
+            return ix;
         }
+        return DKIX_KEY_CHANGED;
     }
-
-    PyObject *obj;
-    Py_ssize_t ix = _Py_dict_lookup_threadsafe(mp, key, hash, &obj);
-    if (ix >= 0 && obj != NULL) {
-        *value_addr = PyStackRef_FromPyObjectSteal(obj);
-    }
-    else {
-        *value_addr = PyStackRef_NULL;
-    }
+    assert(ix == DKIX_KEY_CHANGED);
     return ix;
 }
 
-// This is similar to _Py_dict_lookup_threadsafe_stackref() but
-// it is used when dict is borrowed reference and key is known to be unicode.
-// It avoids increfing the dict if dict only has unicode keys in which case
-// the lookup is safe, otherwise it increfs the dict and lookups the key.
-
 Py_ssize_t
-_Py_dict_lookup_unicode_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
+_Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
 {
-    assert(PyUnicode_CheckExact(key));
     PyDictKeysObject *dk = _Py_atomic_load_ptr(&mp->ma_keys);
-    if (dk->dk_kind == DICT_KEYS_UNICODE) {
-        Py_ssize_t ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash);
-        if (ix == DKIX_EMPTY) {
-            *value_addr = PyStackRef_NULL;
+    if (dk->dk_kind == DICT_KEYS_UNICODE && PyUnicode_CheckExact(key)) {
+        Py_ssize_t ix = lookup_threadsafe_unicode(dk, key, hash, value_addr);
+        if (ix != DKIX_KEY_CHANGED) {
             return ix;
         }
-        else if (ix >= 0) {
-            PyObject **addr_of_value = &DK_UNICODE_ENTRIES(dk)[ix].me_value;
-            PyObject *value = _Py_atomic_load_ptr(addr_of_value);
-            if (value == NULL) {
-                *value_addr = PyStackRef_NULL;
-                return DKIX_EMPTY;
-            }
-            if (_PyObject_HasDeferredRefcount(value)) {
-                *value_addr =  (_PyStackRef){ .bits = (uintptr_t)value | Py_TAG_DEFERRED };
-                return ix;
-            }
-            if (_Py_TryIncrefCompare(addr_of_value, value)) {
-                *value_addr = PyStackRef_FromPyObjectSteal(value);
-                return ix;
-            }
-        }
     }
 
     PyObject *obj;
-    Py_INCREF(mp);
     Py_ssize_t ix = _Py_dict_lookup_threadsafe(mp, key, hash, &obj);
-    Py_DECREF(mp);
     if (ix >= 0 && obj != NULL) {
         *value_addr = PyStackRef_FromPyObjectSteal(obj);
     }
@@ -1692,25 +1659,48 @@ _Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t h
     return ix;
 }
 
-Py_ssize_t
-_Py_dict_lookup_unicode_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr)
+#endif
+
+// Looks up the unicode key `key` in the dictionary. Note that `*method` may
+// already contain a valid value! See _PyObject_GetMethodStackRef().
+int
+_PyDict_GetMethodStackRef(PyDictObject *mp, PyObject *key, _PyStackRef *method)
 {
     assert(PyUnicode_CheckExact(key));
-    PyObject *val;
+    Py_hash_t hash = hash_unicode_key(key);
+
+#ifdef Py_GIL_DISABLED
+    PyDictKeysObject *dk = _Py_atomic_load_ptr_acquire(&mp->ma_keys);
+    if (dk->dk_kind == DICT_KEYS_UNICODE) {
+        _PyStackRef ref;  //
+        Py_ssize_t ix = lookup_threadsafe_unicode(dk, key, hash, &ref);
+        if (ix >= 0) {
+            assert(!PyStackRef_IsNull(ref));
+            PyStackRef_XSETREF(*method, ref);
+            return 1;
+        }
+        else if (ix == DKIX_EMPTY) {
+            return 0;
+        }
+        assert(ix == DKIX_KEY_CHANGED);
+    }
+#endif
+
+    PyObject *obj;
     Py_INCREF(mp);
-    Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, &val);
+    Py_ssize_t ix = _Py_dict_lookup_threadsafe(mp, key, hash, &obj);
     Py_DECREF(mp);
-    if (ix >= 0 && val != NULL) {
-        *value_addr = PyStackRef_FromPyObjectNew(val);
+    if (ix == DKIX_ERROR) {
+        PyStackRef_CLEAR(*method);
+        return -1;
     }
-    else {
-        *value_addr = PyStackRef_NULL;
+    else if (ix >= 0 && obj != NULL) {
+        PyStackRef_XSETREF(*method, PyStackRef_FromPyObjectSteal(obj));
+        return 1;
     }
-    return ix;
+    return 0;  // not found
 }
 
-#endif
-
 int
 _PyDict_HasOnlyStringKeys(PyObject *dict)
 {

@@ -1754,24 +1754,12 @@ _PyObject_GetMethodStackRef(PyThreadState *ts, PyObject *obj,
     }
     if (dict != NULL) {
         assert(PyUnicode_CheckExact(name));
-        Py_hash_t hash = _PyObject_HashFast(name);
-        // cannot fail for exact unicode
-        assert(hash != -1);
-        // ref is not visible to gc so there should be
-        // no escaping calls before assigning it to method
-        _PyStackRef ref;
-        Py_ssize_t ix = _Py_dict_lookup_unicode_threadsafe_stackref((PyDictObject *)dict,
-                                                                  name, hash, &ref);
-        if (ix == DKIX_ERROR) {
-            // error
-            PyStackRef_CLEAR(*method);
+        int found = _PyDict_GetMethodStackRef(dict, name, method);
+        if (found < 0) {
+            assert(PyStackRef_IsNull(*method));
             return -1;
         }
-        else if (!PyStackRef_IsNull(ref)) {
-            // found
-            _PyStackRef tmp = *method;
-            *method = ref;
-            PyStackRef_XCLOSE(tmp);
+        else if (found) {
             return 0;
         }
     }