Skip to content

Commit f2b3a06

Browse files
youknowoneyouknowone
committed
traverse and clear
1 parent e0479fe commit f2b3a06

File tree

13 files changed

+250
-53
lines changed

13 files changed

+250
-53
lines changed

‎crates/derive-impl/src/pyclass.rs‎

Lines changed: 73 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -574,51 +574,80 @@ pub(crate) fn impl_pyclass(attr: PunctuatedNestedMeta, item: Item) -> Result<Tok
574574
)?;
575575

576576
const ALLOWED_TRAVERSE_OPTS: &[&str] = &["manual"];
577-
// try to know if it have a `#[pyclass(trace)]` exist on this struct
578-
// TODO(discord9): rethink on auto detect `#[Derive(PyTrace)]`
579-
580-
// 1. no `traverse` at all: generate a dummy try_traverse
581-
// 2. `traverse = "manual"`: generate a try_traverse, but not #[derive(Traverse)]
582-
// 3. `traverse`: generate a try_traverse, and #[derive(Traverse)]
583-
let (maybe_trace_code, derive_trace) = {
584-
if class_meta.inner()._has_key("traverse")? {
585-
let maybe_trace_code = quote! {
586-
impl ::rustpython_vm::object::MaybeTraverse for #ident {
587-
const IS_TRACE: bool = true;
588-
fn try_traverse(&self, tracer_fn: &mut ::rustpython_vm::object::TraverseFn) {
589-
::rustpython_vm::object::Traverse::traverse(self, tracer_fn);
590-
}
577+
// Generate MaybeTraverse impl with both traverse and clear support
578+
//
579+
// For traverse:
580+
// 1. no `traverse` at all: HAS_TRAVERSE = false, try_traverse does nothing
581+
// 2. `traverse = "manual"`: HAS_TRAVERSE = true, but no #[derive(Traverse)]
582+
// 3. `traverse`: HAS_TRAVERSE = true, and #[derive(Traverse)]
583+
//
584+
// For clear (tp_clear):
585+
// 1. no `clear`: HAS_CLEAR = HAS_TRAVERSE (default: same as traverse)
586+
// 2. `clear` or `clear = true`: HAS_CLEAR = true, try_clear calls Traverse::clear
587+
// 3. `clear = false`: HAS_CLEAR = false (rare: traverse without clear)
588+
let has_traverse = class_meta.inner()._has_key("traverse")?;
589+
let has_clear = if class_meta.inner()._has_key("clear")? {
590+
// If clear attribute is present, use its value
591+
class_meta.inner()._bool("clear")?
592+
} else {
593+
// If clear attribute is absent, default to same as traverse
594+
has_traverse
595+
};
596+
597+
let derive_trace = if has_traverse {
598+
// _optional_str returns Err when key exists without value (e.g., `traverse` vs `traverse = "manual"`)
599+
// We want to derive Traverse in that case, so we handle Err as Ok(None)
600+
let value = class_meta.inner()._optional_str("traverse").ok().flatten();
601+
if let Some(s) = value {
602+
if !ALLOWED_TRAVERSE_OPTS.contains(&s.as_str()) {
603+
bail_span!(
604+
item,
605+
"traverse attribute only accept {ALLOWED_TRAVERSE_OPTS:?} as value or no value at all",
606+
);
607+
}
608+
assert_eq!(s, "manual");
609+
quote! {}
610+
} else {
611+
quote! {#[derive(Traverse)]}
612+
}
613+
} else {
614+
quote! {}
615+
};
616+
617+
let maybe_traverse_code = {
618+
let try_traverse_body = if has_traverse {
619+
quote! {
620+
::rustpython_vm::object::Traverse::traverse(self, tracer_fn);
621+
}
622+
} else {
623+
quote! {
624+
// do nothing
625+
}
626+
};
627+
628+
let try_clear_body = if has_clear {
629+
quote! {
630+
::rustpython_vm::object::Traverse::clear(self, out);
631+
}
632+
} else {
633+
quote! {
634+
// do nothing
635+
}
636+
};
637+
638+
quote! {
639+
impl ::rustpython_vm::object::MaybeTraverse for #ident {
640+
const HAS_TRAVERSE: bool = #has_traverse;
641+
const HAS_CLEAR: bool = #has_clear;
642+
643+
fn try_traverse(&self, tracer_fn: &mut ::rustpython_vm::object::TraverseFn) {
644+
#try_traverse_body
591645
}
592-
};
593-
// if the key `traverse` exist but not as key-value, _optional_str return Err(...)
594-
// so we need to check if it is Ok(Some(...))
595-
let value = class_meta.inner()._optional_str("traverse");
596-
let derive_trace = if let Ok(Some(s)) = value {
597-
if !ALLOWED_TRAVERSE_OPTS.contains(&s.as_str()) {
598-
bail_span!(
599-
item,
600-
"traverse attribute only accept {ALLOWED_TRAVERSE_OPTS:?} as value or no value at all",
601-
);
646+
647+
fn try_clear(&mut self, out: &mut ::std::vec::Vec<::rustpython_vm::PyObjectRef>) {
648+
#try_clear_body
602649
}
603-
assert_eq!(s, "manual");
604-
quote! {}
605-
} else {
606-
quote! {#[derive(Traverse)]}
607-
};
608-
(maybe_trace_code, derive_trace)
609-
} else {
610-
(
611-
// a dummy impl, which do nothing
612-
// #attrs
613-
quote! {
614-
impl ::rustpython_vm::object::MaybeTraverse for #ident {
615-
fn try_traverse(&self, tracer_fn: &mut ::rustpython_vm::object::TraverseFn) {
616-
// do nothing
617-
}
618-
}
619-
},
620-
quote! {},
621-
)
650+
}
622651
}
623652
};
624653

@@ -675,7 +704,7 @@ pub(crate) fn impl_pyclass(attr: PunctuatedNestedMeta, item: Item) -> Result<Tok
675704
let ret = quote! {
676705
#derive_trace
677706
#item
678-
#maybe_trace_code
707+
#maybe_traverse_code
679708
#class_def
680709
#impl_payload
681710
#empty_impl

‎crates/derive-impl/src/pystructseq.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,9 +606,16 @@ pub(crate) fn impl_pystruct_sequence(
606606

607607
// MaybeTraverse - delegate to inner PyTuple
608608
impl ::rustpython_vm::object::MaybeTraverse for #pytype_ident {
609+
const HAS_TRAVERSE: bool = true;
610+
const HAS_CLEAR: bool = true;
611+
609612
fn try_traverse(&self, traverse_fn: &mut ::rustpython_vm::object::TraverseFn<'_>) {
610613
self.0.try_traverse(traverse_fn)
611614
}
615+
616+
fn try_clear(&mut self, out: &mut ::std::vec::Vec<::rustpython_vm::PyObjectRef>) {
617+
self.0.try_clear(out)
618+
}
612619
}
613620

614621
// PySubclass for proper inheritance

‎crates/derive-impl/src/util.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,7 @@ impl ItemMeta for ClassItemMeta {
372372
"ctx",
373373
"impl",
374374
"traverse",
375+
"clear", // tp_clear
375376
];
376377

377378
fn from_inner(inner: ItemMetaInner) -> Self {

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use super::{
22
IterStatus, PositionIterInternal, PyBaseExceptionRef, PyGenericAlias, PyMappingProxy, PySet,
33
PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, set::PySetInner,
44
};
5+
use crate::object::{Traverse, TraverseFn};
56
use crate::{
67
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
78
TryFromObject, atomic_func,
@@ -29,13 +30,28 @@ use std::sync::LazyLock;
2930

3031
pub type DictContentType = dict_inner::Dict;
3132

32-
#[pyclass(module = false, name = "dict", unhashable = true, traverse)]
33+
#[pyclass(module = false, name = "dict", unhashable = true, traverse = "manual")]
3334
#[derive(Default)]
3435
pub struct PyDict {
3536
entries: DictContentType,
3637
}
3738
pub type PyDictRef = PyRef<PyDict>;
3839

40+
// SAFETY: Traverse properly visits all owned PyObjectRefs
41+
unsafe impl Traverse for PyDict {
42+
fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
43+
self.entries.traverse(traverse_fn);
44+
}
45+
46+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
47+
// Pop all entries and collect both keys and values
48+
for (key, value) in self.entries.drain_entries() {
49+
out.push(key);
50+
out.push(value);
51+
}
52+
}
53+
}
54+
3955
impl fmt::Debug for PyDict {
4056
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4157
// TODO: implement more detailed, non-recursive Debug formatter

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,70 @@ unsafe impl Traverse for PyFunction {
5151
closure.as_untyped().traverse(tracer_fn);
5252
}
5353
self.defaults_and_kwdefaults.traverse(tracer_fn);
54+
// Traverse additional fields that may contain references
55+
self.type_params.lock().traverse(tracer_fn);
56+
self.annotations.lock().traverse(tracer_fn);
57+
self.module.lock().traverse(tracer_fn);
58+
self.doc.lock().traverse(tracer_fn);
59+
}
60+
61+
fn clear(&mut self, out: &mut Vec<crate::PyObjectRef>) {
62+
// Pop closure if present (equivalent to Py_CLEAR(func_closure))
63+
if let Some(closure) = self.closure.take() {
64+
out.push(closure.into());
65+
}
66+
67+
// Pop defaults and kwdefaults
68+
if let Some(mut guard) = self.defaults_and_kwdefaults.try_lock() {
69+
if let Some(defaults) = guard.0.take() {
70+
out.push(defaults.into());
71+
}
72+
if let Some(kwdefaults) = guard.1.take() {
73+
out.push(kwdefaults.into());
74+
}
75+
}
76+
77+
// Clear annotations and annotate (Py_CLEAR)
78+
if let Some(mut guard) = self.annotations.try_lock()
79+
&& let Some(annotations) = guard.take()
80+
{
81+
out.push(annotations.into());
82+
}
83+
if let Some(mut guard) = self.annotate.try_lock()
84+
&& let Some(annotate) = guard.take()
85+
{
86+
out.push(annotate);
87+
}
88+
89+
// Clear module, doc, and type_params (Py_CLEAR)
90+
if let Some(mut guard) = self.module.try_lock() {
91+
let old_module =
92+
std::mem::replace(&mut *guard, Context::genesis().none.to_owned().into());
93+
out.push(old_module);
94+
}
95+
if let Some(mut guard) = self.doc.try_lock() {
96+
let old_doc = std::mem::replace(&mut *guard, Context::genesis().none.to_owned().into());
97+
out.push(old_doc);
98+
}
99+
if let Some(mut guard) = self.type_params.try_lock() {
100+
let old_type_params =
101+
std::mem::replace(&mut *guard, Context::genesis().empty_tuple.to_owned());
102+
out.push(old_type_params.into());
103+
}
104+
105+
// Replace name and qualname with empty string to break potential str subclass cycles
106+
// name and qualname could be str subclasses, so they could have reference cycles
107+
if let Some(mut guard) = self.name.try_lock() {
108+
let old_name = std::mem::replace(&mut *guard, Context::genesis().empty_str.to_owned());
109+
out.push(old_name.into());
110+
}
111+
if let Some(mut guard) = self.qualname.try_lock() {
112+
let old_qualname =
113+
std::mem::replace(&mut *guard, Context::genesis().empty_str.to_owned());
114+
out.push(old_qualname.into());
115+
}
116+
117+
// Note: globals, builtins, code are NOT cleared (required to be non-NULL)
54118
}
55119
}
56120

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::atomic_func;
33
use crate::common::lock::{
44
PyMappedRwLockReadGuard, PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
55
};
6+
use crate::object::{Traverse, TraverseFn};
67
use crate::{
78
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
89
class::PyClassImpl,
@@ -23,7 +24,7 @@ use crate::{
2324
use alloc::fmt;
2425
use core::ops::DerefMut;
2526

26-
#[pyclass(module = false, name = "list", unhashable = true, traverse)]
27+
#[pyclass(module = false, name = "list", unhashable = true, traverse = "manual")]
2728
#[derive(Default)]
2829
pub struct PyList {
2930
elements: PyRwLock<Vec<PyObjectRef>>,
@@ -50,6 +51,22 @@ impl FromIterator<PyObjectRef> for PyList {
5051
}
5152
}
5253

54+
// SAFETY: Traverse properly visits all owned PyObjectRefs
55+
unsafe impl Traverse for PyList {
56+
fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
57+
self.elements.traverse(traverse_fn);
58+
}
59+
60+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
61+
// During GC, we use interior mutability to access elements.
62+
// This is safe because during GC collection, the object is unreachable
63+
// and no other code should be accessing it.
64+
if let Some(mut guard) = self.elements.try_write() {
65+
out.extend(guard.drain(..));
66+
}
67+
}
68+
}
69+
5370
impl PyPayload for PyList {
5471
#[inline]
5572
fn class(ctx: &Context) -> &'static Py<PyType> {

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1924,9 +1924,16 @@ impl fmt::Display for PyUtf8Str {
19241924
}
19251925

19261926
impl MaybeTraverse for PyUtf8Str {
1927+
const HAS_TRAVERSE: bool = true;
1928+
const HAS_CLEAR: bool = false;
1929+
19271930
fn try_traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
19281931
self.0.try_traverse(traverse_fn);
19291932
}
1933+
1934+
fn try_clear(&mut self, _out: &mut Vec<PyObjectRef>) {
1935+
// No clear needed for PyUtf8Str
1936+
}
19301937
}
19311938

19321939
impl PyPayload for PyUtf8Str {

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::common::{
33
hash::{PyHash, PyUHash},
44
lock::PyMutex,
55
};
6+
use crate::object::{Traverse, TraverseFn};
67
use crate::{
78
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
89
atomic_func,
@@ -24,7 +25,7 @@ use crate::{
2425
use alloc::fmt;
2526
use std::sync::LazyLock;
2627

27-
#[pyclass(module = false, name = "tuple", traverse)]
28+
#[pyclass(module = false, name = "tuple", traverse = "manual")]
2829
pub struct PyTuple<R = PyObjectRef> {
2930
elements: Box<[R]>,
3031
}
@@ -36,6 +37,19 @@ impl<R> fmt::Debug for PyTuple<R> {
3637
}
3738
}
3839

40+
// SAFETY: Traverse properly visits all owned PyObjectRefs
41+
// Note: Only impl for PyTuple<PyObjectRef> (the default)
42+
unsafe impl Traverse for PyTuple {
43+
fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
44+
self.elements.traverse(traverse_fn);
45+
}
46+
47+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
48+
let elements = std::mem::take(&mut self.elements);
49+
out.extend(elements.into_vec());
50+
}
51+
}
52+
3953
impl PyPayload for PyTuple {
4054
#[inline]
4155
fn class(ctx: &Context) -> &'static Py<PyType> {

‎crates/vm/src/dict_inner.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,17 @@ impl<T: Clone> Dict<T> {
724724
+ inner.indices.len() * size_of::<i64>()
725725
+ inner.entries.len() * size_of::<DictEntry<T>>()
726726
}
727+
728+
/// Pop all entries from the dict, returning (key, value) pairs.
729+
/// This is used for circular reference resolution in GC.
730+
/// Requires &mut self to avoid lock contention.
731+
pub fn drain_entries(&mut self) -> impl Iterator<Item = (PyObjectRef, T)> + '_ {
732+
let inner = self.inner.get_mut();
733+
inner.used = 0;
734+
inner.filled = 0;
735+
inner.indices.iter_mut().for_each(|i| *i = IndexEntry::FREE);
736+
inner.entries.drain(..).flatten().map(|e| (e.key, e.value))
737+
}
727738
}
728739

729740
type LookupResult = (IndexEntry, IndexIndex);

0 commit comments

Comments
 (0)