Add _PyDictKeys_StringLookupAndVersion

Look up a unicode key in an all unicode keys object along with the
keys version, assigning one if not present.

We need a keys version that is consistent with presence of the key
for use in the guards.

Author: mpage

Include/internal/pycore_dict.h

@@ -113,7 +113,18 @@ extern Py_ssize_t _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py
 extern Py_ssize_t _Py_dict_lookup_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);
+
+/* Look up a string key in an all unicode dict keys, assign the keys object a version, and
+ * stored it in version.
+ *
+ * Returns DKIX_ERROR if key is not a string or if the keys object is not all
+ * strings.
+ *
+ * Returns DKIX_EMPTY if the key is not present.
+ */
+extern Py_ssize_t _PyDictKeys_StringLookupAndVersion(PyDictKeysObject* dictkeys, PyObject *key, uint32_t *version);
 PyAPI_FUNC(PyObject *)_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObject *);
 PyAPI_FUNC(void) _PyDict_LoadGlobalStackRef(PyDictObject *, PyDictObject *, PyObject *, _PyStackRef *);
 

Objects/dictobject.c

@@ -1129,6 +1129,23 @@ dictkeys_generic_lookup(PyDictObject *mp, PyDictKeysObject* dk, PyObject *key, P
     return do_lookup(mp, dk, key, hash, compare_generic);
 }
 
+static Py_hash_t
+check_keys_and_hash(PyDictKeysObject *dk, PyObject *key) {
+    DictKeysKind kind = dk->dk_kind;
+    if (!PyUnicode_CheckExact(key) || kind == DICT_KEYS_GENERAL) {
+        return -1;
+    }
+    Py_hash_t hash = unicode_get_hash(key);
+    if (hash == -1) {
+        hash = PyUnicode_Type.tp_hash(key);
+        if (hash == -1) {
+            PyErr_Clear();
+            return -1;
+        }
+    }
+    return hash;
+}
+
 /* Lookup a string in a (all unicode) dict keys.
  * Returns DKIX_ERROR if key is not a string,
  * or if the dict keys is not all strings.
@@ -1138,19 +1155,27 @@ dictkeys_generic_lookup(PyDictObject *mp, PyDictKeysObject* dk, PyObject *key, P
 Py_ssize_t
 _PyDictKeys_StringLookup(PyDictKeysObject* dk, PyObject *key)
 {
-    DictKeysKind kind = dk->dk_kind;
-    if (!PyUnicode_CheckExact(key) || kind == DICT_KEYS_GENERAL) {
+    Py_hash_t hash = check_keys_and_hash(dk, key);
+    if (hash == -1) {
         return DKIX_ERROR;
     }
-    Py_hash_t hash = unicode_get_hash(key);
+    return unicodekeys_lookup_unicode(dk, key, hash);
+}
+
+Py_ssize_t
+_PyDictKeys_StringLookupAndVersion(PyDictKeysObject* dk, PyObject *key, uint32_t *version)
+{
+    Py_hash_t hash = check_keys_and_hash(dk, key);
     if (hash == -1) {
-        hash = PyUnicode_Type.tp_hash(key);
-        if (hash == -1) {
-            PyErr_Clear();
-            return DKIX_ERROR;
-        }
+        return DKIX_ERROR;
     }
-    return unicodekeys_lookup_unicode(dk, key, hash);
+
+    Py_ssize_t ix;
+    LOCK_KEYS(dk);
+    ix = unicodekeys_lookup_unicode(dk, key, hash);
+    *version = _PyDictKeys_GetVersionForCurrentState(_PyInterpreterState_GET(), dk);
+    UNLOCK_KEYS(dk);
+    return ix;
 }
 
 #ifdef Py_GIL_DISABLED