Skip to content

Commit 7d8d90c

Browse files
youknowoneyouknowone
committed
traverse
1 parent e0479fe commit 7d8d90c

File tree

13 files changed

+246
-52
lines changed

13 files changed

+246
-52
lines changed

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

Lines changed: 74 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -574,51 +574,81 @@ 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 HAS_TRAVERSE = has_traverse;
619+
let try_traverse_body = if has_traverse {
620+
quote! {
621+
::rustpython_vm::object::Traverse::traverse(self, tracer_fn);
622+
}
623+
} else {
624+
quote! {
625+
// do nothing
626+
}
627+
};
628+
629+
let try_clear_body = if has_clear {
630+
quote! {
631+
::rustpython_vm::object::Traverse::clear(self, out);
632+
}
633+
} else {
634+
quote! {
635+
// do nothing
636+
}
637+
};
638+
639+
quote! {
640+
impl ::rustpython_vm::object::MaybeTraverse for #ident {
641+
const HAS_TRAVERSE: bool = #HAS_TRAVERSE;
642+
const HAS_CLEAR: bool = #has_clear;
643+
644+
fn try_traverse(&self, tracer_fn: &mut ::rustpython_vm::object::TraverseFn) {
645+
#try_traverse_body
591646
}
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-
);
647+
648+
fn try_clear(&mut self, out: &mut ::std::vec::Vec<::rustpython_vm::PyObjectRef>) {
649+
#try_clear_body
602650
}
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-
)
651+
}
622652
}
623653
};
624654

@@ -675,7 +705,7 @@ pub(crate) fn impl_pyclass(attr: PunctuatedNestedMeta, item: Item) -> Result<Tok
675705
let ret = quote! {
676706
#derive_trace
677707
#item
678-
#maybe_trace_code
708+
#maybe_traverse_code
679709
#class_def
680710
#impl_payload
681711
#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 = false;
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+
// Struct sequences don't need clear
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.pop_all_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: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,50 @@ 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+
// Note: We do NOT clear annotations here.
78+
// Unlike CPython which can set func_annotations to NULL, RustPython always
79+
// has a dict reference. Clearing the dict in-place would affect all functions
80+
// that share the same annotations dict (e.g., via functools.update_wrapper).
81+
// The annotations dict typically doesn't create cycles, so skipping it is safe.
82+
83+
// Replace name and qualname with empty string to break potential str subclass cycles
84+
// This matches CPython's func_clear behavior: "name and qualname could be str
85+
// subclasses, so they could have reference cycles"
86+
if let Some(mut guard) = self.name.try_lock() {
87+
let old_name = std::mem::replace(&mut *guard, Context::genesis().empty_str.to_owned());
88+
out.push(old_name.into());
89+
}
90+
if let Some(mut guard) = self.qualname.try_lock() {
91+
let old_qualname =
92+
std::mem::replace(&mut *guard, Context::genesis().empty_str.to_owned());
93+
out.push(old_qualname.into());
94+
}
95+
96+
// Note: globals, builtins, code are NOT cleared
97+
// as per func_clear behavior (they're required to be non-NULL)
5498
}
5599
}
56100

‎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 pop_all_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)