add _Py_dict_lookup_unicode_threadsafe_stackref to lookup method

Author: kumaraditya303

Include/internal/pycore_dict.h

@@ -112,6 +112,7 @@ 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 Py_ssize_t _PyDict_LookupIndex(PyDictObject *, PyObject *);
 extern Py_ssize_t _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyObject *key);

Objects/dictobject.c

@@ -1581,33 +1581,41 @@ _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
+unicodekeys_lookup_unicode_threadsafe_stackref(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(PyUnicode_CheckExact(key));
+    assert(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;
+        return DKIX_EMPTY;
+    } 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;
         }
     }
+}
+
+
+Py_ssize_t
+_Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, 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)) {
+        return unicodekeys_lookup_unicode_threadsafe_stackref(dk, key, hash, value_addr);
+    }
 
     PyObject *obj;
     Py_ssize_t ix = _Py_dict_lookup_threadsafe(mp, key, hash, &obj);
@@ -1620,6 +1628,33 @@ _Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t h
     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)
+{
+    assert(PyUnicode_CheckExact(key));
+    PyDictKeysObject *dk = _Py_atomic_load_ptr(&mp->ma_keys);
+    if (dk->dk_kind == DICT_KEYS_UNICODE) {
+        return unicodekeys_lookup_unicode_threadsafe_stackref(dk, key, hash, value_addr);
+    }
+
+    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);
+    }
+    else {
+        *value_addr = PyStackRef_NULL;
+    }
+    return ix;
+}
+
 #else   // Py_GIL_DISABLED
 
 Py_ssize_t
@@ -1644,6 +1679,23 @@ _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)
+{
+    assert(PyUnicode_CheckExact(key));
+    PyObject *val;
+    Py_INCREF(mp);
+    Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, &val);
+    Py_DECREF(mp);
+    if (ix >= 0 && val != NULL) {
+        *value_addr = PyStackRef_FromPyObjectNew(val);
+    }
+    else {
+        *value_addr = PyStackRef_NULL;
+    }
+    return ix;
+}
+
 #endif
 
 int

Objects/object.c

@@ -1759,18 +1759,9 @@ _PyObject_GetMethodStackRef(PyThreadState *ts, PyObject *obj,
         assert(hash != -1);
         // ref is not visible to gc so there should be
         // no escaping calls before assigning it to method
-        PyDictObject *mp = (PyDictObject *)dict;
-        PyDictKeysObject *keys = FT_ATOMIC_LOAD_PTR_ACQUIRE(mp->ma_keys);
-        bool unicode_keys = DK_IS_UNICODE(keys);
         _PyStackRef ref;
-        if (!unicode_keys) {
-            Py_INCREF(mp);
-        }
-        Py_ssize_t ix = _Py_dict_lookup_threadsafe_stackref(mp, name,
-                                                            hash, &ref);
-        if (!unicode_keys) {
-            Py_DECREF(mp);
-        }
+        Py_ssize_t ix = _Py_dict_lookup_unicode_threadsafe_stackref((PyDictObject *)dict,
+                                                                  name, hash, &ref);
         if (ix == DKIX_ERROR) {
             // error
             PyStackRef_CLEAR(*method);