Collect cycles through sets, frozensets and set iterators (#9343)

Assisted-by: Codex:GPT-6

Lib/test/test_set.py

@@ -352,7 +352,6 @@ class TestJointOps:
         self.assertEqual(sum(elem.hash_count for elem in d), n)
         self.assertEqual(d3, dict.fromkeys(d, 123))
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_container_iterator(self):
         # Bug #3680: tp_traverse was not implemented for set iterator object
         class C(object):

crates/vm/src/builtins/set.rs

@@ -8,7 +8,7 @@ use super::{
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     atomic_func,
-    class::{PyClassDef, PyClassImpl},
+    class::{PyClassDef, PyClassImpl, StaticType},
     common::{
         ascii,
         hash::PyHash,
@@ -22,6 +22,7 @@ use crate::{
         ArgIterable, FuncArgs, NameOthers, OptionalArg, PosArgs, PyArithmeticValue,
         PyComparisonValue,
     },
+    object::{Traverse, TraverseFn},
     protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
     recursion::ReprGuard,
     types::AsNumber,
@@ -40,12 +41,23 @@ use rustpython_common::{
 
 pub(crate) type SetContentType = dict_inner::Dict<()>;
 
-#[pyclass(module = false, name = "set", unhashable = true, traverse)]
+#[pyclass(module = false, name = "set", unhashable = true, traverse = "manual")]
 #[derive(Default)]
 pub struct PySet {
     pub(super) inner: PySetInner,
 }
 
+// SAFETY: the inner table visits each exclusively owned reference once.
+unsafe impl Traverse for PySet {
+    fn traverse(&self, tracer: &mut TraverseFn<'_>) {
+        self.inner.traverse(tracer);
+    }
+
+    fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
+        Traverse::clear(&mut self.inner, out);
+    }
+}
+
 impl PySet {
     #[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")]
     pub fn new_ref(ctx: &Context) -> PyRef<Self> {
@@ -82,12 +94,28 @@ impl PySet {
     }
 }
 
-#[pyclass(module = false, name = "frozenset", unhashable = true)]
+#[pyclass(
+    module = false,
+    name = "frozenset",
+    unhashable = true,
+    traverse = "manual"
+)]
 pub struct PyFrozenSet {
     inner: PySetInner,
     hash: PyAtomic<PyHash>,
 }
 
+// SAFETY: the inner table visits each exclusively owned reference once.
+unsafe impl Traverse for PyFrozenSet {
+    fn traverse(&self, tracer: &mut TraverseFn<'_>) {
+        self.inner.traverse(tracer);
+    }
+
+    fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
+        Traverse::clear(&mut self.inner, out);
+    }
+}
+
 impl Default for PyFrozenSet {
     fn default() -> Self {
         Self {
@@ -161,6 +189,9 @@ impl fmt::Debug for PyFrozenSet {
 }
 
 impl PyPayload for PySet {
+    // Ordinary destruction can release the table directly after untracking.
+    const CLEAR_ON_DEALLOC: bool = false;
+
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
         ctx.types.set_type
@@ -168,6 +199,22 @@ impl PyPayload for PySet {
 }
 
 impl PyPayload for PyFrozenSet {
+    const CLEAR_ON_DEALLOC: bool = false;
+
+    fn needs_gc(&self, class: &Py<PyType>) -> bool {
+        // Subclasses can acquire instance state. Shared working tables may
+        // still be mutated by another owner, so keep those tracked as well.
+        if !class.is(Self::static_type()) || PyRc::strong_count(&self.inner.content) != 1 {
+            return true;
+        }
+        self.inner.content.any_key(|key| {
+            // Exact untracked frozensets are immutable and were checked at
+            // allocation. Other untracked objects may later acquire GC edges.
+            key.gc_may_have_edges()
+                && !(key.class().is(Self::static_type()) && !key.is_gc_tracked())
+        })
+    }
+
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
         ctx.types.frozenset_type
@@ -179,10 +226,22 @@ pub(super) struct PySetInner {
     content: PyRc<SetContentType>,
 }
 
-unsafe impl crate::object::Traverse for PySetInner {
-    fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) {
-        // FIXME(discord9): Rc means shared ref, so should it be traced?
-        self.content.traverse(tracer_fn)
+// SAFETY: a shared table owns only one reference per entry, regardless of its
+// number of owners. Only attribute those edges to an exclusive owner; otherwise
+// leave them as external roots rather than subtracting the same edges twice.
+unsafe impl Traverse for PySetInner {
+    fn traverse(&self, tracer: &mut TraverseFn<'_>) {
+        if PyRc::strong_count(&self.content) == 1 {
+            self.content.traverse(tracer);
+        }
+    }
+
+    fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
+        if let Some(content) = PyRc::get_mut(&mut self.content) {
+            // Detach before releasing keys: dropping them may re-enter Python.
+            out.reserve(content.len());
+            out.extend(content.drain_entries().map(|(key, ())| key));
+        }
     }
 }
 
@@ -1652,6 +1711,13 @@ pub(crate) struct PyAnySet {
     object: PyObject,
 }
 
+// SAFETY: this owner contributes exactly one strong reference to the object.
+unsafe impl Traverse for PyRef<PyAnySet> {
+    fn traverse(&self, tracer: &mut TraverseFn<'_>) {
+        tracer(self.as_object());
+    }
+}
+
 impl Borrow<PyObject> for PyAnySet {
     #[inline(always)]
     fn borrow(&self) -> &PyObject {
@@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef<PyAnySet> {
     }
 }
 
-#[pyclass(module = false, name = "set_iterator")]
+#[pyclass(module = false, name = "set_iterator", traverse, clear = false)]
 pub(crate) struct PySetIterator {
+    #[pytraverse(skip)]
     size: DictSize,
     /// Whether the set was found to have changed, which `setiter_iternext()`
     /// records by writing a size no set can have. Sticky: what it makes the
     /// iterator answer, it answers from then on.
+    #[pytraverse(skip)]
     changed: PyAtomic<bool>,
     internal: PyMutex<PositionIterInternal<PyRef<PyAnySet>>>,
 }
@@ -1907,6 +1975,115 @@ pub(crate) fn init(context: &'static Context) {
 mod reference_tests {
     use super::*;
 
+    #[test]
+    fn set_gc_detaches_keys_for_deferred_release() {
+        use crate::object::MaybeTraverse;
+
+        fn check(mut set: impl MaybeTraverse, key: &PyObjectRef) {
+            let mut visits = 0;
+            set.try_traverse(&mut |child| {
+                assert!(core::ptr::eq(child, &**key));
+                visits += 1;
+            });
+            assert_eq!(visits, 1);
+            let mut edges = Vec::new();
+            set.try_clear(&mut edges);
+            set.try_traverse(&mut |_| panic!("cleared set retained a key"));
+            set.try_clear(&mut edges);
+            assert_eq!(edges.len(), 1);
+            assert!(core::ptr::eq(&*edges[0], &**key));
+            assert_eq!(key.strong_count(), 2);
+            drop(edges);
+            assert_eq!(key.strong_count(), 1);
+        }
+
+        crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            for frozen in [false, true] {
+                let key: PyObjectRef = vm.ctx.new_int(1000).into();
+                let inner = PySetInner::default();
+                inner.add(&key, vm).unwrap();
+                if frozen {
+                    check(
+                        PyFrozenSet {
+                            inner,
+                            ..Default::default()
+                        },
+                        &key,
+                    );
+                } else {
+                    check(PySet { inner }, &key);
+                }
+            }
+        });
+    }
+
+    #[test]
+    fn set_gc_preserves_shared_table_edges() {
+        crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            let key: PyObjectRef = vm.ctx.new_int(1000).into();
+            let mut inner = PySetInner::default();
+            inner.add(&key, vm).unwrap();
+            let alias = inner.clone();
+            inner.traverse(&mut |_| panic!("shared table edge counted"));
+            alias.traverse(&mut |_| panic!("shared table edge counted"));
+            let mut edges = Vec::new();
+            Traverse::clear(&mut inner, &mut edges);
+            assert!(edges.is_empty());
+            assert!(alias.contains(&key, vm).unwrap());
+            drop(alias);
+            Traverse::clear(&mut inner, &mut edges);
+            assert_eq!(inner.len(), 0);
+            assert_eq!(edges.len(), 1);
+        });
+    }
+
+    #[test]
+    fn frozenset_gc_tracking_covers_native_construction() {
+        crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
+            let key: PyObjectRef = vm.ctx.new_int(1000).into();
+            let atomic = PyRef::new_ref(
+                PyFrozenSet::from_iter(vm, [key.clone()]).unwrap(),
+                vm.ctx.types.frozenset_type.to_owned(),
+                None,
+            );
+            assert!(!atomic.as_object().is_gc_tracked());
+            let nested = PyFrozenSet::from_iter(vm, [atomic.into()])
+                .unwrap()
+                .into_ref(&vm.ctx);
+            assert!(!nested.as_object().is_gc_tracked());
+
+            // A traversable key must keep the frozen owner visible to GC.
+            let tuple: PyObjectRef = vm.ctx.new_tuple(vec![key.clone()]).into();
+            let owner = PyFrozenSet::from_iter(vm, [tuple])
+                .unwrap()
+                .into_ref(&vm.ctx);
+            assert!(owner.as_object().is_gc_tracked());
+
+            let subclass = vm.ctx.new_class(
+                None,
+                "FrozenSubclass",
+                vm.ctx.types.frozenset_type.to_owned(),
+                Default::default(),
+            );
+            let subclass = PyFrozenSet::default()
+                .into_ref_with_type(vm, subclass)
+                .unwrap();
+            assert!(subclass.as_object().is_gc_tracked());
+
+            // A temporary native alias can still add traversable keys.
+            let inner = PySetInner::default();
+            inner.add(&key, vm).unwrap();
+            let alias = inner.clone();
+            let owner = PyFrozenSet {
+                inner,
+                ..Default::default()
+            }
+            .into_ref(&vm.ctx);
+            assert!(owner.as_object().is_gc_tracked());
+            drop(alias);
+        });
+    }
+
     #[test]
     fn anyset_reference_layout_and_ownership() {
         static_assertions::assert_not_impl_any!(PyAnySet: PyPayload);

crates/vm/src/dict_inner.rs

@@ -1165,6 +1165,16 @@ impl<T: Clone> Dict<T> {
         current != *old
     }
 
+    /// Inspect keys without allocating or cloning their references.
+    /// The predicate must not run Python or mutate this table.
+    pub(crate) fn any_key(&self, mut predicate: impl FnMut(&PyObject) -> bool) -> bool {
+        self.read()
+            .entries
+            .iter()
+            .flatten()
+            .any(|entry| predicate(&entry.key))
+    }
+
     pub(crate) fn keys(&self) -> Vec<PyObjectRef> {
         self.read()
             .entries

crates/vm/src/object/core.rs

@@ -198,7 +198,9 @@ pub(super) unsafe fn default_dealloc<T: PyPayload>(obj: *mut PyObject) {
     // them. Some payloads (e.g. FrameObject) drop children in place inside clear_fn
     // instead of extracting them, so user code (`__del__`) may run here.
     let mut edges = Vec::new();
-    if let Some(clear_fn) = vtable.clear {
+    if (T::CLEAR_ON_DEALLOC || T::HAS_FREELIST)
+        && let Some(clear_fn) = vtable.clear
+    {
         unsafe { clear_fn(obj, &mut edges) };
     }
     // Drop extracted child references - may trigger recursive destruction.
@@ -2290,6 +2292,15 @@ impl PyObject {
         GcBits::from_bits_retain(self.0.gc_bits.load(Ordering::Relaxed)).contains(GcBits::TRACKED)
     }
 
+    /// Whether this object's payload or instance storage can own GC edges.
+    /// Unlike is_gc_tracked, this also covers temporarily untracked objects.
+    #[inline]
+    pub(crate) fn gc_may_have_edges(&self) -> bool {
+        self.0.vtable.trace.is_some()
+            || self.class().heaptype_ext.is_some()
+            || self.0.ext_ref().is_some()
+    }
+
     /// Get the referents (objects directly referenced) of this object.
     /// Uses the full traverse including dict and slots.
     pub fn gc_get_referents(&self) -> Vec<PyObjectRef> {
@@ -2937,6 +2948,8 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
     pub fn new_ref(payload: T, typ: crate::builtins::PyTypeRef, dict: Option<PyDictRef>) -> Self {
         let has_dict = dict.is_some();
         let is_heaptype = typ.heaptype_ext.is_some();
+        let track = !T::NEW_REF_UNTRACKED
+            && (has_dict || is_heaptype || (T::HAS_TRAVERSE && payload.needs_gc(&typ)));
 
         // Try to reuse from freelist (no dict, no heaptype)
         let cached = if !has_dict && !is_heaptype {
@@ -2969,14 +2982,12 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
         };
 
         // Track object if:
-        // - HAS_TRAVERSE is true (Rust payload implements Traverse), OR
+        // - HAS_TRAVERSE is true and the payload needs cyclic GC, OR
         // - has instance dict (user-defined class instances), OR
         // - heap type (all heap type instances are GC-tracked, like Py_TPFLAGS_HAVE_GC)
         // unless the payload opts out via NEW_REF_UNTRACKED (e.g. call frames,
         // which are tracked lazily only on escape).
-        if (<T as crate::object::MaybeTraverse>::HAS_TRAVERSE || has_dict || is_heaptype)
-            && !T::NEW_REF_UNTRACKED
-        {
+        if track {
             // Tracks under the interpreter running now and collects if this
             // allocation pushed gen0 past its threshold.
             unsafe {

crates/vm/src/object/payload.rs

@@ -91,6 +91,21 @@ pub trait PyPayload: MaybeTraverse + PyThreadingConstraint + Sized + 'static {
     /// which is created untracked and tracked lazily only on escape.
     const NEW_REF_UNTRACKED: bool = false;
 
+    /// Whether this payload needs cyclic GC when first wrapped in an object.
+    /// Returning false promises that this instance cannot participate in a
+    /// reference cycle for its entire lifetime, not just at allocation time.
+    /// The check must not run Python. Instance dictionaries and heap types are
+    /// still tracked, unless NEW_REF_UNTRACKED opts out explicitly.
+    #[inline]
+    fn needs_gc(&self, _class: &Py<PyType>) -> bool {
+        true
+    }
+
+    /// Run tp_clear before ordinary reference-counted destruction. Payloads
+    /// whose Rust destructor releases their children directly may opt out.
+    /// Cyclic GC still calls tp_clear; freelists also require a cleared payload.
+    const CLEAR_ON_DEALLOC: bool = true;
+
     /// Whether this type has a freelist. Types with freelists require
     /// immediate (non-deferred) GC untracking during dealloc to prevent
     /// race conditions when the object is reused.