Skip to content

Commit db48780

Browse files
authored
Collect cycles through sets, frozensets and set iterators (#9343)
Assisted-by: Codex:GPT-6
1 parent 8c0c242 commit db48780

5 files changed

Lines changed: 226 additions & 14 deletions

File tree

‎Lib/test/test_set.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,6 @@ def test_do_not_rehash_dict_keys(self):
352352
self.assertEqual(sum(elem.hash_count for elem in d), n)
353353
self.assertEqual(d3, dict.fromkeys(d, 123))
354354

355-
@unittest.expectedFailure # TODO: RUSTPYTHON
356355
def test_container_iterator(self):
357356
# Bug #3680: tp_traverse was not implemented for set iterator object
358357
class C(object):

‎crates/vm/src/builtins/set.rs‎

Lines changed: 185 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use super::{
88
use crate::{
99
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
1010
atomic_func,
11-
class::{PyClassDef, PyClassImpl},
11+
class::{PyClassDef, PyClassImpl, StaticType},
1212
common::{
1313
ascii,
1414
hash::PyHash,
@@ -22,6 +22,7 @@ use crate::{
2222
ArgIterable, FuncArgs, NameOthers, OptionalArg, PosArgs, PyArithmeticValue,
2323
PyComparisonValue,
2424
},
25+
object::{Traverse, TraverseFn},
2526
protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
2627
recursion::ReprGuard,
2728
types::AsNumber,
@@ -40,12 +41,23 @@ use rustpython_common::{
4041

4142
pub(crate) type SetContentType = dict_inner::Dict<()>;
4243

43-
#[pyclass(module = false, name = "set", unhashable = true, traverse)]
44+
#[pyclass(module = false, name = "set", unhashable = true, traverse = "manual")]
4445
#[derive(Default)]
4546
pub struct PySet {
4647
pub(super) inner: PySetInner,
4748
}
4849

50+
// SAFETY: the inner table visits each exclusively owned reference once.
51+
unsafe impl Traverse for PySet {
52+
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
53+
self.inner.traverse(tracer);
54+
}
55+
56+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
57+
Traverse::clear(&mut self.inner, out);
58+
}
59+
}
60+
4961
impl PySet {
5062
#[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")]
5163
pub fn new_ref(ctx: &Context) -> PyRef<Self> {
@@ -82,12 +94,28 @@ impl PySet {
8294
}
8395
}
8496

85-
#[pyclass(module = false, name = "frozenset", unhashable = true)]
97+
#[pyclass(
98+
module = false,
99+
name = "frozenset",
100+
unhashable = true,
101+
traverse = "manual"
102+
)]
86103
pub struct PyFrozenSet {
87104
inner: PySetInner,
88105
hash: PyAtomic<PyHash>,
89106
}
90107

108+
// SAFETY: the inner table visits each exclusively owned reference once.
109+
unsafe impl Traverse for PyFrozenSet {
110+
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
111+
self.inner.traverse(tracer);
112+
}
113+
114+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
115+
Traverse::clear(&mut self.inner, out);
116+
}
117+
}
118+
91119
impl Default for PyFrozenSet {
92120
fn default() -> Self {
93121
Self {
@@ -161,13 +189,32 @@ impl fmt::Debug for PyFrozenSet {
161189
}
162190

163191
impl PyPayload for PySet {
192+
// Ordinary destruction can release the table directly after untracking.
193+
const CLEAR_ON_DEALLOC: bool = false;
194+
164195
#[inline]
165196
fn class(ctx: &Context) -> &'static Py<PyType> {
166197
ctx.types.set_type
167198
}
168199
}
169200

170201
impl PyPayload for PyFrozenSet {
202+
const CLEAR_ON_DEALLOC: bool = false;
203+
204+
fn needs_gc(&self, class: &Py<PyType>) -> bool {
205+
// Subclasses can acquire instance state. Shared working tables may
206+
// still be mutated by another owner, so keep those tracked as well.
207+
if !class.is(Self::static_type()) || PyRc::strong_count(&self.inner.content) != 1 {
208+
return true;
209+
}
210+
self.inner.content.any_key(|key| {
211+
// Exact untracked frozensets are immutable and were checked at
212+
// allocation. Other untracked objects may later acquire GC edges.
213+
key.gc_may_have_edges()
214+
&& !(key.class().is(Self::static_type()) && !key.is_gc_tracked())
215+
})
216+
}
217+
171218
#[inline]
172219
fn class(ctx: &Context) -> &'static Py<PyType> {
173220
ctx.types.frozenset_type
@@ -179,10 +226,22 @@ pub(super) struct PySetInner {
179226
content: PyRc<SetContentType>,
180227
}
181228

182-
unsafe impl crate::object::Traverse for PySetInner {
183-
fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) {
184-
// FIXME(discord9): Rc means shared ref, so should it be traced?
185-
self.content.traverse(tracer_fn)
229+
// SAFETY: a shared table owns only one reference per entry, regardless of its
230+
// number of owners. Only attribute those edges to an exclusive owner; otherwise
231+
// leave them as external roots rather than subtracting the same edges twice.
232+
unsafe impl Traverse for PySetInner {
233+
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
234+
if PyRc::strong_count(&self.content) == 1 {
235+
self.content.traverse(tracer);
236+
}
237+
}
238+
239+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
240+
if let Some(content) = PyRc::get_mut(&mut self.content) {
241+
// Detach before releasing keys: dropping them may re-enter Python.
242+
out.reserve(content.len());
243+
out.extend(content.drain_entries().map(|(key, ())| key));
244+
}
186245
}
187246
}
188247

@@ -1652,6 +1711,13 @@ pub(crate) struct PyAnySet {
16521711
object: PyObject,
16531712
}
16541713

1714+
// SAFETY: this owner contributes exactly one strong reference to the object.
1715+
unsafe impl Traverse for PyRef<PyAnySet> {
1716+
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
1717+
tracer(self.as_object());
1718+
}
1719+
}
1720+
16551721
impl Borrow<PyObject> for PyAnySet {
16561722
#[inline(always)]
16571723
fn borrow(&self) -> &PyObject {
@@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef<PyAnySet> {
17471813
}
17481814
}
17491815

1750-
#[pyclass(module = false, name = "set_iterator")]
1816+
#[pyclass(module = false, name = "set_iterator", traverse, clear = false)]
17511817
pub(crate) struct PySetIterator {
1818+
#[pytraverse(skip)]
17521819
size: DictSize,
17531820
/// Whether the set was found to have changed, which `setiter_iternext()`
17541821
/// records by writing a size no set can have. Sticky: what it makes the
17551822
/// iterator answer, it answers from then on.
1823+
#[pytraverse(skip)]
17561824
changed: PyAtomic<bool>,
17571825
internal: PyMutex<PositionIterInternal<PyRef<PyAnySet>>>,
17581826
}
@@ -1907,6 +1975,115 @@ pub(crate) fn init(context: &'static Context) {
19071975
mod reference_tests {
19081976
use super::*;
19091977

1978+
#[test]
1979+
fn set_gc_detaches_keys_for_deferred_release() {
1980+
use crate::object::MaybeTraverse;
1981+
1982+
fn check(mut set: impl MaybeTraverse, key: &PyObjectRef) {
1983+
let mut visits = 0;
1984+
set.try_traverse(&mut |child| {
1985+
assert!(core::ptr::eq(child, &**key));
1986+
visits += 1;
1987+
});
1988+
assert_eq!(visits, 1);
1989+
let mut edges = Vec::new();
1990+
set.try_clear(&mut edges);
1991+
set.try_traverse(&mut |_| panic!("cleared set retained a key"));
1992+
set.try_clear(&mut edges);
1993+
assert_eq!(edges.len(), 1);
1994+
assert!(core::ptr::eq(&*edges[0], &**key));
1995+
assert_eq!(key.strong_count(), 2);
1996+
drop(edges);
1997+
assert_eq!(key.strong_count(), 1);
1998+
}
1999+
2000+
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
2001+
for frozen in [false, true] {
2002+
let key: PyObjectRef = vm.ctx.new_int(1000).into();
2003+
let inner = PySetInner::default();
2004+
inner.add(&key, vm).unwrap();
2005+
if frozen {
2006+
check(
2007+
PyFrozenSet {
2008+
inner,
2009+
..Default::default()
2010+
},
2011+
&key,
2012+
);
2013+
} else {
2014+
check(PySet { inner }, &key);
2015+
}
2016+
}
2017+
});
2018+
}
2019+
2020+
#[test]
2021+
fn set_gc_preserves_shared_table_edges() {
2022+
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
2023+
let key: PyObjectRef = vm.ctx.new_int(1000).into();
2024+
let mut inner = PySetInner::default();
2025+
inner.add(&key, vm).unwrap();
2026+
let alias = inner.clone();
2027+
inner.traverse(&mut |_| panic!("shared table edge counted"));
2028+
alias.traverse(&mut |_| panic!("shared table edge counted"));
2029+
let mut edges = Vec::new();
2030+
Traverse::clear(&mut inner, &mut edges);
2031+
assert!(edges.is_empty());
2032+
assert!(alias.contains(&key, vm).unwrap());
2033+
drop(alias);
2034+
Traverse::clear(&mut inner, &mut edges);
2035+
assert_eq!(inner.len(), 0);
2036+
assert_eq!(edges.len(), 1);
2037+
});
2038+
}
2039+
2040+
#[test]
2041+
fn frozenset_gc_tracking_covers_native_construction() {
2042+
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
2043+
let key: PyObjectRef = vm.ctx.new_int(1000).into();
2044+
let atomic = PyRef::new_ref(
2045+
PyFrozenSet::from_iter(vm, [key.clone()]).unwrap(),
2046+
vm.ctx.types.frozenset_type.to_owned(),
2047+
None,
2048+
);
2049+
assert!(!atomic.as_object().is_gc_tracked());
2050+
let nested = PyFrozenSet::from_iter(vm, [atomic.into()])
2051+
.unwrap()
2052+
.into_ref(&vm.ctx);
2053+
assert!(!nested.as_object().is_gc_tracked());
2054+
2055+
// A traversable key must keep the frozen owner visible to GC.
2056+
let tuple: PyObjectRef = vm.ctx.new_tuple(vec![key.clone()]).into();
2057+
let owner = PyFrozenSet::from_iter(vm, [tuple])
2058+
.unwrap()
2059+
.into_ref(&vm.ctx);
2060+
assert!(owner.as_object().is_gc_tracked());
2061+
2062+
let subclass = vm.ctx.new_class(
2063+
None,
2064+
"FrozenSubclass",
2065+
vm.ctx.types.frozenset_type.to_owned(),
2066+
Default::default(),
2067+
);
2068+
let subclass = PyFrozenSet::default()
2069+
.into_ref_with_type(vm, subclass)
2070+
.unwrap();
2071+
assert!(subclass.as_object().is_gc_tracked());
2072+
2073+
// A temporary native alias can still add traversable keys.
2074+
let inner = PySetInner::default();
2075+
inner.add(&key, vm).unwrap();
2076+
let alias = inner.clone();
2077+
let owner = PyFrozenSet {
2078+
inner,
2079+
..Default::default()
2080+
}
2081+
.into_ref(&vm.ctx);
2082+
assert!(owner.as_object().is_gc_tracked());
2083+
drop(alias);
2084+
});
2085+
}
2086+
19102087
#[test]
19112088
fn anyset_reference_layout_and_ownership() {
19122089
static_assertions::assert_not_impl_any!(PyAnySet: PyPayload);

‎crates/vm/src/dict_inner.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1165,6 +1165,16 @@ impl<T: Clone> Dict<T> {
11651165
current != *old
11661166
}
11671167

1168+
/// Inspect keys without allocating or cloning their references.
1169+
/// The predicate must not run Python or mutate this table.
1170+
pub(crate) fn any_key(&self, mut predicate: impl FnMut(&PyObject) -> bool) -> bool {
1171+
self.read()
1172+
.entries
1173+
.iter()
1174+
.flatten()
1175+
.any(|entry| predicate(&entry.key))
1176+
}
1177+
11681178
pub(crate) fn keys(&self) -> Vec<PyObjectRef> {
11691179
self.read()
11701180
.entries

‎crates/vm/src/object/core.rs‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,9 @@ pub(super) unsafe fn default_dealloc<T: PyPayload>(obj: *mut PyObject) {
198198
// them. Some payloads (e.g. FrameObject) drop children in place inside clear_fn
199199
// instead of extracting them, so user code (`__del__`) may run here.
200200
let mut edges = Vec::new();
201-
if let Some(clear_fn) = vtable.clear {
201+
if (T::CLEAR_ON_DEALLOC || T::HAS_FREELIST)
202+
&& let Some(clear_fn) = vtable.clear
203+
{
202204
unsafe { clear_fn(obj, &mut edges) };
203205
}
204206
// Drop extracted child references - may trigger recursive destruction.
@@ -2290,6 +2292,15 @@ impl PyObject {
22902292
GcBits::from_bits_retain(self.0.gc_bits.load(Ordering::Relaxed)).contains(GcBits::TRACKED)
22912293
}
22922294

2295+
/// Whether this object's payload or instance storage can own GC edges.
2296+
/// Unlike is_gc_tracked, this also covers temporarily untracked objects.
2297+
#[inline]
2298+
pub(crate) fn gc_may_have_edges(&self) -> bool {
2299+
self.0.vtable.trace.is_some()
2300+
|| self.class().heaptype_ext.is_some()
2301+
|| self.0.ext_ref().is_some()
2302+
}
2303+
22932304
/// Get the referents (objects directly referenced) of this object.
22942305
/// Uses the full traverse including dict and slots.
22952306
pub fn gc_get_referents(&self) -> Vec<PyObjectRef> {
@@ -2937,6 +2948,8 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
29372948
pub fn new_ref(payload: T, typ: crate::builtins::PyTypeRef, dict: Option<PyDictRef>) -> Self {
29382949
let has_dict = dict.is_some();
29392950
let is_heaptype = typ.heaptype_ext.is_some();
2951+
let track = !T::NEW_REF_UNTRACKED
2952+
&& (has_dict || is_heaptype || (T::HAS_TRAVERSE && payload.needs_gc(&typ)));
29402953

29412954
// Try to reuse from freelist (no dict, no heaptype)
29422955
let cached = if !has_dict && !is_heaptype {
@@ -2969,14 +2982,12 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
29692982
};
29702983

29712984
// Track object if:
2972-
// - HAS_TRAVERSE is true (Rust payload implements Traverse), OR
2985+
// - HAS_TRAVERSE is true and the payload needs cyclic GC, OR
29732986
// - has instance dict (user-defined class instances), OR
29742987
// - heap type (all heap type instances are GC-tracked, like Py_TPFLAGS_HAVE_GC)
29752988
// unless the payload opts out via NEW_REF_UNTRACKED (e.g. call frames,
29762989
// which are tracked lazily only on escape).
2977-
if (<T as crate::object::MaybeTraverse>::HAS_TRAVERSE || has_dict || is_heaptype)
2978-
&& !T::NEW_REF_UNTRACKED
2979-
{
2990+
if track {
29802991
// Tracks under the interpreter running now and collects if this
29812992
// allocation pushed gen0 past its threshold.
29822993
unsafe {

‎crates/vm/src/object/payload.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,21 @@ pub trait PyPayload: MaybeTraverse + PyThreadingConstraint + Sized + 'static {
9191
/// which is created untracked and tracked lazily only on escape.
9292
const NEW_REF_UNTRACKED: bool = false;
9393

94+
/// Whether this payload needs cyclic GC when first wrapped in an object.
95+
/// Returning false promises that this instance cannot participate in a
96+
/// reference cycle for its entire lifetime, not just at allocation time.
97+
/// The check must not run Python. Instance dictionaries and heap types are
98+
/// still tracked, unless NEW_REF_UNTRACKED opts out explicitly.
99+
#[inline]
100+
fn needs_gc(&self, _class: &Py<PyType>) -> bool {
101+
true
102+
}
103+
104+
/// Run tp_clear before ordinary reference-counted destruction. Payloads
105+
/// whose Rust destructor releases their children directly may opt out.
106+
/// Cyclic GC still calls tp_clear; freelists also require a cleared payload.
107+
const CLEAR_ON_DEALLOC: bool = true;
108+
94109
/// Whether this type has a freelist. Types with freelists require
95110
/// immediate (non-deferred) GC untracking during dealloc to prevent
96111
/// race conditions when the object is reused.

0 commit comments

Comments
 (0)