gh-151722: Defer GC tracking in frozendict.copy() (#152230)

Fix _PyDict_Or() and frozendict.copy(): only track the frozendict by
the GC once the dictionary is fully initialized.

Functions modifying frozendict now ensures that the object is not
tracked by the GC (in debug mode).

* can_modify_dict() checks that _PyObject_GC_IS_TRACKED() is false
  for frozendicts.
* dict_merge_api() makes sure that the dictionary is tracked by the
  GC.

Author: vstinner

Objects/dictobject.c

@@ -305,16 +305,20 @@ static inline int
 can_modify_dict(PyDictObject *mp)
 {
     if (PyFrozenDict_Check(mp)) {
+        // gh-151722: A frozendict must not be tracked by the GC
+        // when it's being modified.
+        assert(!_PyObject_GC_IS_TRACKED(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));
+        assert(PyUnstable_Object_IsUniquelyReferenced(_PyObject_CAST(mp)));
     }
     else {
         // Locking is only required if the dictionary is not
         // uniquely referenced.
         ASSERT_DICT_LOCKED(mp);
-        return 1;
     }
+    return 1;
 }
 #endif
 
@@ -937,7 +941,7 @@ free_values(PyDictValues *values, bool use_qsbr)
 static inline PyObject *
 new_dict_impl(PyDictObject *mp, PyDictKeysObject *keys,
               PyDictValues *values, Py_ssize_t used,
-              int free_values_on_failure, int frozendict)
+              int free_values_on_failure, int frozendict, int gc_track)
 {
     assert(keys != NULL);
     if (mp == NULL) {
@@ -956,12 +960,14 @@ new_dict_impl(PyDictObject *mp, PyDictKeysObject *keys,
         ((PyFrozenDictObject *)mp)->ma_hash = -1;
     }
     ASSERT_CONSISTENT(mp);
-    _PyObject_GC_TRACK(mp);
+    if (gc_track) {
+        _PyObject_GC_TRACK(mp);
+    }
     return (PyObject *)mp;
 }
 
 /* Consumes a reference to the keys object */
-static PyObject *
+static PyObject*
 new_dict(PyDictKeysObject *keys, PyDictValues *values,
          Py_ssize_t used, int free_values_on_failure)
 {
@@ -971,16 +977,30 @@ new_dict(PyDictKeysObject *keys, PyDictValues *values,
     }
     assert(mp == NULL || Py_IS_TYPE(mp, &PyDict_Type));
 
-    return new_dict_impl(mp, keys, values, used, free_values_on_failure, 0);
+    return new_dict_impl(mp, keys, values, used, free_values_on_failure, 0, 1);
 }
 
 /* Consumes a reference to the keys object */
-static PyObject *
-new_frozendict(PyDictKeysObject *keys, PyDictValues *values,
-               Py_ssize_t used, int free_values_on_failure)
+static PyObject*
+new_dict_untracked(PyDictKeysObject *keys, PyDictValues *values,
+                   Py_ssize_t used, int free_values_on_failure)
+{
+    PyDictObject *mp = _Py_FREELIST_POP(PyDictObject, dicts);
+    if (mp == NULL) {
+        mp = PyObject_GC_New(PyDictObject, &PyDict_Type);
+    }
+    assert(mp == NULL || Py_IS_TYPE(mp, &PyDict_Type));
+
+    return new_dict_impl(mp, keys, values, used, free_values_on_failure, 0, 0);
+}
+
+/* Consumes a reference to the keys object */
+static PyObject*
+new_frozendict_untracked(PyDictKeysObject *keys, PyDictValues *values,
+                         Py_ssize_t used, int free_values_on_failure)
 {
     PyDictObject *mp = PyObject_GC_New(PyDictObject, &PyFrozenDict_Type);
-    return new_dict_impl(mp, keys, values, used, free_values_on_failure, 1);
+    return new_dict_impl(mp, keys, values, used, free_values_on_failure, 1, 0);
 }
 
 static PyObject *
@@ -3471,7 +3491,6 @@ _PyDict_FromKeys(PyObject *cls, PyObject *iterable, PyObject *value)
     }
     if (PyFrozenDict_Check(d)) {
         assert(can_modify_dict((PyDictObject*)d));
-        assert(!_PyObject_GC_IS_TRACKED(d));
     }
 
     if (PyDict_CheckExact(d)) {
@@ -4076,6 +4095,7 @@ merge_from_seq2_lock_held(PyObject *d, PyObject *seq2, int override)
     assert(d != NULL);
     assert(PyAnyDict_Check(d));
     assert(seq2 != NULL);
+    assert(can_modify_dict((PyDictObject*)d));
 
     it = PyObject_GetIter(seq2);
     if (it == NULL)
@@ -4277,11 +4297,13 @@ dict_merge(PyObject *a, PyObject *b, int override, PyObject **dupkey)
     assert(0 <= override && override <= 2);
 
     PyDictObject *mp = _PyAnyDict_CAST(a);
+
     int res = 0;
     if (PyAnyDict_Check(b) && (Py_TYPE(b)->tp_iter == dict_iter)) {
         PyDictObject *other = (PyDictObject*)b;
         int res;
         Py_BEGIN_CRITICAL_SECTION2(a, b);
+        assert(can_modify_dict(mp));
         res = dict_dict_merge((PyDictObject *)a, other, override, dupkey);
         ASSERT_CONSISTENT(a);
         Py_END_CRITICAL_SECTION2();
@@ -4290,6 +4312,8 @@ dict_merge(PyObject *a, PyObject *b, int override, PyObject **dupkey)
     else {
         /* Do it the generic, slower way */
         Py_BEGIN_CRITICAL_SECTION(a);
+        assert(can_modify_dict(mp));
+
         PyObject *keys = PyMapping_Keys(b);
         PyObject *iter;
         PyObject *key, *value;
@@ -4382,9 +4406,7 @@ dict_merge_api(PyObject *a, PyObject *b, int override, PyObject **dupkey)
     }
 
     int res = dict_merge(a, b, override, dupkey);
-    if (PyDict_Check(a)) {
-        assert(_PyObject_GC_IS_TRACKED(a));
-    }
+    assert(_PyObject_GC_IS_TRACKED(a));
     return res;
 }
 
@@ -4442,7 +4464,7 @@ copy_values(PyDictValues *values)
 }
 
 static PyObject *
-copy_lock_held(PyObject *o, int as_frozendict)
+copy_lock_held_untracked(PyObject *o, int as_frozendict)
 {
     PyObject *copy;
     PyDictObject *mp;
@@ -4455,12 +4477,15 @@ copy_lock_held(PyObject *o, int as_frozendict)
     mp = (PyDictObject *)o;
     if (mp->ma_used == 0) {
         /* The dict is empty; just return a new dict. */
+        PyObject *d;
         if (as_frozendict) {
-            return PyFrozenDict_New(NULL);
+            d = frozendict_new_untracked(&PyFrozenDict_Type);
         }
         else {
-            return PyDict_New();
+            d = dict_new_untracked(&PyDict_Type);
         }
+        assert(!_PyObject_GC_IS_TRACKED(d));
+        return d;
     }
 
     if (_PyDict_HasSplitTable(mp)) {
@@ -4492,7 +4517,7 @@ copy_lock_held(PyObject *o, int as_frozendict)
             PyFrozenDictObject *frozen = (PyFrozenDictObject *)split_copy;
             frozen->ma_hash = -1;
         }
-        _PyObject_GC_TRACK(split_copy);
+        assert(!_PyObject_GC_IS_TRACKED(split_copy));
         return (PyObject *)split_copy;
     }
 
@@ -4520,10 +4545,10 @@ copy_lock_held(PyObject *o, int as_frozendict)
         }
         PyDictObject *new;
         if (as_frozendict) {
-            new = (PyDictObject *)new_frozendict(keys, NULL, 0, 0);
+            new = (PyDictObject *)new_frozendict_untracked(keys, NULL, 0, 0);
         }
         else {
-            new = (PyDictObject *)new_dict(keys, NULL, 0, 0);
+            new = (PyDictObject *)new_dict_untracked(keys, NULL, 0, 0);
         }
         if (new == NULL) {
             /* In case of an error, new_dict()/new_frozendict() takes care of
@@ -4533,14 +4558,15 @@ copy_lock_held(PyObject *o, int as_frozendict)
 
         new->ma_used = mp->ma_used;
         ASSERT_CONSISTENT(new);
+        assert(!_PyObject_GC_IS_TRACKED(new));
         return (PyObject *)new;
     }
 
     if (as_frozendict) {
-        copy = PyFrozenDict_New(NULL);
+        copy = frozendict_new_untracked(&PyFrozenDict_Type);
     }
     else {
-        copy = PyDict_New();
+        copy = dict_new_untracked(&PyDict_Type);
     }
     if (copy == NULL)
         return NULL;
@@ -4549,9 +4575,7 @@ copy_lock_held(PyObject *o, int as_frozendict)
         return NULL;
     }
 
-    if (PyDict_Check(copy)) {
-        assert(_PyObject_GC_IS_TRACKED(copy));
-    }
+    assert(!_PyObject_GC_IS_TRACKED(copy));
     return copy;
 }
 
@@ -4565,25 +4589,28 @@ PyDict_Copy(PyObject *o)
 
     PyObject *res;
     Py_BEGIN_CRITICAL_SECTION(o);
-    res = copy_lock_held(o, 0);
+    res = copy_lock_held_untracked(o, 0);
     Py_END_CRITICAL_SECTION();
+    if (res != NULL) {
+        _PyObject_GC_TRACK(res);
+    }
     return res;
 }
 
 // Similar to PyDict_Copy(), but return a frozendict if the argument
 // is a frozendict.
 static PyObject *
-anydict_copy(PyObject *o)
+anydict_copy_untracked(PyObject *o)
 {
     assert(PyAnyDict_Check(o));
 
     PyObject *res;
     if (PyFrozenDict_Check(o)) {
-        res = copy_lock_held(o, 1);
+        res = copy_lock_held_untracked(o, 1);
     }
     else {
         Py_BEGIN_CRITICAL_SECTION(o);
-        res = copy_lock_held(o, 0);
+        res = copy_lock_held_untracked(o, 0);
         Py_END_CRITICAL_SECTION();
     }
     return res;
@@ -4598,13 +4625,16 @@ _PyDict_CopyAsDict(PyObject *o)
 
     PyObject *res;
     if (PyFrozenDict_Check(o)) {
-        res = copy_lock_held(o, 0);
+        res = copy_lock_held_untracked(o, 0);
     }
     else {
         Py_BEGIN_CRITICAL_SECTION(o);
-        res = copy_lock_held(o, 0);
+        res = copy_lock_held_untracked(o, 0);
         Py_END_CRITICAL_SECTION();
     }
+    if (res != NULL) {
+        _PyObject_GC_TRACK(res);
+    }
     return res;
 }
 
@@ -5157,14 +5187,15 @@ _PyDict_Or(PyObject *self, PyObject *other)
     if (!PyAnyDict_Check(self) || !PyAnyDict_Check(other)) {
         Py_RETURN_NOTIMPLEMENTED;
     }
-    PyObject *new = anydict_copy(self);
+    PyObject *new = anydict_copy_untracked(self);
     if (new == NULL) {
         return NULL;
     }
     if (dict_update_arg(new, other)) {
         Py_DECREF(new);
         return NULL;
     }
+    _PyObject_GC_TRACK(new);
     return new;
 }
 
@@ -5352,7 +5383,6 @@ dict_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwds
     if (self == NULL) {
         return NULL;
     }
-    assert(!_PyObject_GC_IS_TRACKED(self));
     _PyObject_GC_TRACK(self);
     return self;
 }
@@ -5434,7 +5464,7 @@ frozendict_vectorcall(PyObject *type, PyObject * const*args,
             }
         }
     }
-    assert(!_PyObject_GC_IS_TRACKED(self));
+
     _PyObject_GC_TRACK(self);
     return self;
 }
@@ -6753,10 +6783,12 @@ dictitems_xor_lock_held(PyObject *d1, PyObject *d2)
     ASSERT_DICT_LOCKED(d1);
     ASSERT_DICT_LOCKED(d2);
 
-    PyObject *temp_dict = copy_lock_held(d1, 0);
+    PyObject *temp_dict = copy_lock_held_untracked(d1, 0);
     if (temp_dict == NULL) {
         return NULL;
     }
+    _PyObject_GC_TRACK(temp_dict);
+
     PyObject *result_set = PySet_New(NULL);
     if (result_set == NULL) {
         Py_CLEAR(temp_dict);
@@ -8488,7 +8520,6 @@ frozendict_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         assert(kwds == NULL);
     }
 
-    assert(!_PyObject_GC_IS_TRACKED(d));
     _PyObject_GC_TRACK(d);
     return d;
 }
@@ -8533,7 +8564,11 @@ frozendict_copy_impl(PyFrozenDictObject *self)
         return Py_NewRef(self);
     }
 
-    return anydict_copy((PyObject*)self);
+    PyObject *copy = anydict_copy_untracked((PyObject*)self);
+    if (copy != NULL) {
+        _PyObject_GC_TRACK(copy);
+    }
+    return copy;
 }