diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h index 25bb224921aa8b..6ab569393e5ce1 100644 --- a/Include/internal/pycore_dict.h +++ b/Include/internal/pycore_dict.h @@ -113,6 +113,8 @@ extern Py_ssize_t _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t has 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 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); diff --git a/Include/internal/pycore_stackref.h b/Include/internal/pycore_stackref.h index 6bf82d8322f508..c4e8f10fe05276 100644 --- a/Include/internal/pycore_stackref.h +++ b/Include/internal/pycore_stackref.h @@ -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 { \ diff --git a/Objects/dictobject.c b/Objects/dictobject.c index 0ed52ac5e87b6e..928b905aaedb1b 100644 --- a/Objects/dictobject.c +++ b/Objects/dictobject.c @@ -1581,32 +1581,47 @@ _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyOb return ix; } +static Py_ssize_t +lookup_threadsafe_unicode(PyDictKeysObject *dk, PyObject *key, Py_hash_t hash, _PyStackRef *value_addr) +{ + 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; + } + if (_Py_TryIncrefCompare(addr_of_value, value)) { + *value_addr = PyStackRef_FromPyObjectSteal(value); + return ix; + } + return DKIX_KEY_CHANGED; + } + assert(ix == DKIX_KEY_CHANGED); + 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)) { - Py_ssize_t ix = unicodekeys_lookup_unicode_threadsafe(dk, key, hash); - if (ix == DKIX_EMPTY) { - *value_addr = PyStackRef_NULL; + 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; @@ -1646,6 +1661,46 @@ _Py_dict_lookup_threadsafe_stackref(PyDictObject *mp, PyObject *key, Py_hash_t h #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)); + 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_threadsafe(mp, key, hash, &obj); + Py_DECREF(mp); + if (ix == DKIX_ERROR) { + PyStackRef_CLEAR(*method); + return -1; + } + else if (ix >= 0 && obj != NULL) { + PyStackRef_XSETREF(*method, PyStackRef_FromPyObjectSteal(obj)); + return 1; + } + return 0; // not found +} + int _PyDict_HasOnlyStringKeys(PyObject *dict) { diff --git a/Objects/object.c b/Objects/object.c index 3ed7d55593dffa..479f4176a46039 100644 --- a/Objects/object.c +++ b/Objects/object.c @@ -1753,20 +1753,15 @@ _PyObject_GetMethodStackRef(PyThreadState *ts, PyObject *obj, } } if (dict != NULL) { - // TODO: use _Py_dict_lookup_threadsafe_stackref - Py_INCREF(dict); - PyObject *value; - if (PyDict_GetItemRef(dict, name, &value) != 0) { - // found or error - Py_DECREF(dict); - PyStackRef_CLEAR(*method); - if (value != NULL) { - *method = PyStackRef_FromPyObjectSteal(value); - } + assert(PyUnicode_CheckExact(name)); + int found = _PyDict_GetMethodStackRef((PyDictObject *)dict, name, method); + if (found < 0) { + assert(PyStackRef_IsNull(*method)); + return -1; + } + else if (found) { return 0; } - // not found - Py_DECREF(dict); } if (meth_found) {