From db4878080a7a9cb75ee7889cf632248dd72098af Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Sun, 25 Oct 2026 09:10:43 +0300 Subject: [PATCH] Collect cycles through sets, frozensets and set iterators (#9343) Assisted-by: Codex:GPT-6 --- Lib/test/test_set.py | 1 - crates/vm/src/builtins/set.rs | 193 ++++++++++++++++++++++++++++++-- crates/vm/src/dict_inner.rs | 10 ++ crates/vm/src/object/core.rs | 21 +++- crates/vm/src/object/payload.rs | 15 +++ 5 files changed, 226 insertions(+), 14 deletions(-) diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py index b615af508..7079eb1b9 100644 --- a/Lib/test/test_set.py +++ b/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): diff --git a/crates/vm/src/builtins/set.rs b/crates/vm/src/builtins/set.rs index bd204fea1..ce251efca 100644 --- a/crates/vm/src/builtins/set.rs +++ b/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) { + Traverse::clear(&mut self.inner, out); + } +} + impl PySet { #[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")] pub fn new_ref(ctx: &Context) -> PyRef { @@ -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, } +// 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) { + 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 { 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) -> 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 { ctx.types.frozenset_type @@ -179,10 +226,22 @@ pub(super) struct PySetInner { content: PyRc, } -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) { + 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 { + fn traverse(&self, tracer: &mut TraverseFn<'_>) { + tracer(self.as_object()); + } +} + impl Borrow for PyAnySet { #[inline(always)] fn borrow(&self) -> &PyObject { @@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef { } } -#[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, internal: PyMutex>>, } @@ -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); diff --git a/crates/vm/src/dict_inner.rs b/crates/vm/src/dict_inner.rs index 8b4a8eb6d..6388f3ed3 100644 --- a/crates/vm/src/dict_inner.rs +++ b/crates/vm/src/dict_inner.rs @@ -1165,6 +1165,16 @@ impl Dict { 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 { self.read() .entries diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs index 72c15855c..fa0e6be1a 100644 --- a/crates/vm/src/object/core.rs +++ b/crates/vm/src/object/core.rs @@ -198,7 +198,9 @@ pub(super) unsafe fn default_dealloc(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 { @@ -2937,6 +2948,8 @@ impl PyRef { pub fn new_ref(payload: T, typ: crate::builtins::PyTypeRef, dict: Option) -> 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 PyRef { }; // 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 (::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 { diff --git a/crates/vm/src/object/payload.rs b/crates/vm/src/object/payload.rs index a5a4bf35b..36de55fc3 100644 --- a/crates/vm/src/object/payload.rs +++ b/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) -> 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.