Add _PyObject_SetMaybeWeakref

mpage · mpage · commit 9bf3820977b7 · 2024-03-22T12:29:13.000-07:00
We need to tag weakrefs and their referents
diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
@@ -507,6 +507,25 @@ _Py_XNewRefWithLock(PyObject *obj)
     return _Py_NewRefWithLock(obj);
 }
 
+static inline void
+_PyObject_SetMaybeWeakref(PyObject *op)
+{
+    if (_Py_IsImmortal(op)) {
+        return;
+    }
+    for (;;) {
+        Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
+        if ((shared & _Py_REF_SHARED_FLAG_MASK) != 0) {
+            // Nothing to do if it's in WEAKREFS, QUEUED, or MERGED states.
+            return;
+        }
+        if (_Py_atomic_compare_exchange_ssize(
+                &op->ob_ref_shared, &shared, shared | _Py_REF_MAYBE_WEAKREF)) {
+            return;
+        }
+    }
+}
+
 #endif
 
 #ifdef Py_REF_DEBUG
