gh-141510: Add can_modify_dict() in dictobject.c (#144955)

can_modify_dict() is stricter than ASSERT_DICT_LOCKED() for
frozendict. It uses PyUnstable_Object_IsUniquelyReferenced() which
matters for free-threaded builds.

Replace anydict_setitem_take2() with setitem_take2_lock_held(). It's
no longer useful to have two functions.

Author: vstinner

Objects/dictobject.c

@@ -282,6 +282,24 @@ load_keys_nentries(PyDictObject *mp)
 
 #endif
 
+#ifndef NDEBUG
+// Check if it's possible to modify a dictionary.
+// Usage: assert(can_modify_dict(mp)).
+static inline int
+can_modify_dict(PyDictObject *mp)
+{
+    if (PyFrozenDict_Check(mp)) {
+        // No locking required to modify a newly created frozendict
+        // since it's only accessible from the current thread.
+        return PyUnstable_Object_IsUniquelyReferenced(_PyObject_CAST(mp));
+    }
+    else {
+        ASSERT_DICT_LOCKED(mp);
+        return 1;
+    }
+}
+#endif
+
 #define _PyAnyDict_CAST(op) \
     (assert(PyAnyDict_Check(op)), _Py_CAST(PyDictObject*, op))
 
@@ -1867,8 +1885,9 @@ insert_split_key(PyDictKeysObject *keys, PyObject *key, Py_hash_t hash)
 static void
 insert_split_value(PyDictObject *mp, PyObject *key, PyObject *value, Py_ssize_t ix)
 {
+    assert(can_modify_dict(mp));
     assert(PyUnicode_CheckExact(key));
-    ASSERT_DICT_LOCKED(mp);
+
     PyObject *old_value = mp->ma_values->values[ix];
     if (old_value == NULL) {
         _PyDict_NotifyEvent(PyDict_EVENT_ADDED, mp, key, value);
@@ -1896,11 +1915,11 @@ static int
 insertdict(PyDictObject *mp,
            PyObject *key, Py_hash_t hash, PyObject *value)
 {
+    assert(can_modify_dict(mp));
+
     PyObject *old_value = NULL;
     Py_ssize_t ix;
 
-    ASSERT_DICT_LOCKED(mp);
-
     if (_PyDict_HasSplitTable(mp) && PyUnicode_CheckExact(key)) {
         ix = insert_split_key(mp->ma_keys, key, hash);
         if (ix != DKIX_EMPTY) {
@@ -1967,8 +1986,8 @@ static int
 insert_to_emptydict(PyDictObject *mp,
                     PyObject *key, Py_hash_t hash, PyObject *value)
 {
+    assert(can_modify_dict(mp));
     assert(mp->ma_keys == Py_EMPTY_KEYS);
-    ASSERT_DICT_LOCKED(mp);
 
     int unicode = PyUnicode_CheckExact(key);
     PyDictKeysObject *newkeys = new_keys_object(PyDict_LOG_MINSIZE, unicode);
@@ -2069,11 +2088,11 @@ static int
 dictresize(PyDictObject *mp,
            uint8_t log2_newsize, int unicode)
 {
+    assert(can_modify_dict(mp));
+
     PyDictKeysObject *oldkeys, *newkeys;
     PyDictValues *oldvalues;
 
-    ASSERT_DICT_LOCKED(mp);
-
     if (log2_newsize >= SIZEOF_SIZE_T*8) {
         PyErr_NoMemory();
         return -1;
@@ -2671,11 +2690,13 @@ _PyDict_LoadBuiltinsFromGlobals(PyObject *globals)
 
 /* Consumes references to key and value */
 static int
-anydict_setitem_take2(PyDictObject *mp, PyObject *key, PyObject *value)
+setitem_take2_lock_held(PyDictObject *mp, PyObject *key, PyObject *value)
 {
+    assert(PyAnyDict_Check(mp));
+    assert(can_modify_dict(mp));
     assert(key);
     assert(value);
-    assert(PyAnyDict_Check(mp));
+
     Py_hash_t hash = _PyObject_HashFast(key);
     if (hash == -1) {
         dict_unhashable_type(key);
@@ -2691,14 +2712,6 @@ anydict_setitem_take2(PyDictObject *mp, PyObject *key, PyObject *value)
     return insertdict(mp, key, hash, value);
 }
 
-/* Consumes references to key and value */
-static int
-setitem_take2_lock_held(PyDictObject *mp, PyObject *key, PyObject *value)
-{
-    ASSERT_DICT_LOCKED(mp);
-    return anydict_setitem_take2(mp, key, value);
-}
-
 int
 _PyDict_SetItem_Take2(PyDictObject *mp, PyObject *key, PyObject *value)
 {
@@ -2800,9 +2813,9 @@ static void
 delitem_common(PyDictObject *mp, Py_hash_t hash, Py_ssize_t ix,
                PyObject *old_value)
 {
-    PyObject *old_key;
+    assert(can_modify_dict(mp));
 
-    ASSERT_DICT_LOCKED(mp);
+    PyObject *old_key;
 
     Py_ssize_t hashpos = lookdict_index(mp->ma_keys, hash, ix);
     assert(hashpos >= 0);
@@ -2856,19 +2869,17 @@ int
 _PyDict_DelItem_KnownHash_LockHeld(PyObject *op, PyObject *key, Py_hash_t hash)
 {
     Py_ssize_t ix;
-    PyDictObject *mp;
     PyObject *old_value;
 
     if (!PyDict_Check(op)) {
         PyErr_BadInternalCall();
         return -1;
     }
-
-    ASSERT_DICT_LOCKED(op);
+    PyDictObject *mp = (PyDictObject *)op;
+    assert(can_modify_dict(mp));
 
     assert(key);
     assert(hash != -1);
-    mp = (PyDictObject *)op;
     ix = _Py_dict_lookup(mp, key, hash, &old_value);
     if (ix == DKIX_ERROR)
         return -1;
@@ -2897,19 +2908,18 @@ delitemif_lock_held(PyObject *op, PyObject *key,
                     int (*predicate)(PyObject *value, void *arg),
                     void *arg)
 {
+    PyDictObject *mp = _PyAnyDict_CAST(op);
+    assert(can_modify_dict(mp));
+
     Py_ssize_t ix;
-    PyDictObject *mp;
     Py_hash_t hash;
     PyObject *old_value;
     int res;
 
-    ASSERT_DICT_LOCKED(op);
-
     assert(key);
     hash = PyObject_Hash(key);
     if (hash == -1)
         return -1;
-    mp = (PyDictObject *)op;
     ix = _Py_dict_lookup(mp, key, hash, &old_value);
     if (ix == DKIX_ERROR) {
         return -1;
@@ -2951,16 +2961,16 @@ _PyDict_DelItemIf(PyObject *op, PyObject *key,
 static void
 clear_lock_held(PyObject *op)
 {
-    PyDictObject *mp;
+    if (!PyDict_Check(op)) {
+        return;
+    }
+    PyDictObject *mp = (PyDictObject *)op;
+    assert(can_modify_dict(mp));
+
     PyDictKeysObject *oldkeys;
     PyDictValues *oldvalues;
     Py_ssize_t i, n;
 
-    ASSERT_DICT_LOCKED(op);
-
-    if (!PyDict_Check(op))
-        return;
-    mp = ((PyDictObject *)op);
     oldkeys = mp->ma_keys;
     oldvalues = mp->ma_values;
     if (oldkeys == Py_EMPTY_KEYS) {
@@ -3106,8 +3116,7 @@ _PyDict_Pop_KnownHash(PyDictObject *mp, PyObject *key, Py_hash_t hash,
                       PyObject **result)
 {
     assert(PyDict_Check(mp));
-
-    ASSERT_DICT_LOCKED(mp);
+    assert(can_modify_dict(mp));
 
     if (mp->ma_used == 0) {
         if (result) {
@@ -3149,8 +3158,6 @@ _PyDict_Pop_KnownHash(PyDictObject *mp, PyObject *key, Py_hash_t hash,
 static int
 pop_lock_held(PyObject *op, PyObject *key, PyObject **result)
 {
-    ASSERT_DICT_LOCKED(op);
-
     if (!PyDict_Check(op)) {
         if (result) {
             *result = NULL;
@@ -3159,6 +3166,7 @@ pop_lock_held(PyObject *op, PyObject *key, PyObject **result)
         return -1;
     }
     PyDictObject *dict = (PyDictObject *)op;
+    assert(can_modify_dict(dict));
 
     if (dict->ma_used == 0) {
         if (result) {
@@ -3361,9 +3369,9 @@ dict_iter_exit:;
     }
     else if (PyFrozenDict_CheckExact(d)) {
         while ((key = PyIter_Next(it)) != NULL) {
-            // anydict_setitem_take2 consumes a reference to key
-            status = anydict_setitem_take2((PyDictObject *)d,
-                                           key, Py_NewRef(value));
+            // setitem_take2_lock_held consumes a reference to key
+            status = setitem_take2_lock_held((PyDictObject *)d,
+                                             key, Py_NewRef(value));
             if (status < 0) {
                 assert(PyErr_Occurred());
                 goto Fail;
@@ -3933,7 +3941,7 @@ PyDict_MergeFromSeq2(PyObject *d, PyObject *seq2, int override)
 static int
 dict_dict_merge(PyDictObject *mp, PyDictObject *other, int override)
 {
-    ASSERT_DICT_LOCKED(mp);
+    assert(can_modify_dict(mp));
     ASSERT_DICT_LOCKED(other);
 
     if (other == mp || other->ma_used == 0)
@@ -4474,15 +4482,14 @@ dict_setdefault_ref_lock_held(PyObject *d, PyObject *key, PyObject *default_valu
     Py_hash_t hash;
     Py_ssize_t ix;
 
-    ASSERT_DICT_LOCKED(d);
-
     if (!PyDict_Check(d)) {
         PyErr_BadInternalCall();
         if (result) {
             *result = NULL;
         }
         return -1;
     }
+    assert(can_modify_dict((PyDictObject*)d));
 
     hash = _PyObject_HashFast(key);
     if (hash == -1) {
@@ -4657,11 +4664,11 @@ static PyObject *
 dict_popitem_impl(PyDictObject *self)
 /*[clinic end generated code: output=e65fcb04420d230d input=ef28b4da5f0f762e]*/
 {
+    assert(can_modify_dict(self));
+
     Py_ssize_t i, j;
     PyObject *res;
 
-    ASSERT_DICT_LOCKED(self);
-
     /* Allocate the result tuple before checking the size.  Believe it
      * or not, this allocation could trigger a garbage collection which
      * could empty the dict, so if we checked the size first and that
@@ -4799,11 +4806,12 @@ _PyDict_SizeOf_LockHeld(PyDictObject *mp)
 }
 
 void
-_PyDict_ClearKeysVersionLockHeld(PyObject *mp)
+_PyDict_ClearKeysVersionLockHeld(PyObject *op)
 {
-    ASSERT_DICT_LOCKED(mp);
+    PyDictObject *mp = _PyAnyDict_CAST(op);
+    assert(can_modify_dict(mp));
 
-    FT_ATOMIC_STORE_UINT32_RELAXED(((PyDictObject *)mp)->ma_keys->dk_version, 0);
+    FT_ATOMIC_STORE_UINT32_RELAXED(mp->ma_keys->dk_version, 0);
 }
 
 Py_ssize_t