rustpython_vm/object/
traverse.rs

1use core::ptr::NonNull;
2
3use rustpython_common::lock::{PyMutex, PyRwLock};
4
5use crate::{
6    AsObject, PyObject, PyObjectRef, PyRef, PyStackRef, function::Either, object::PyObjectPayload,
7};
8
9pub type TraverseFn<'a> = dyn FnMut(&PyObject) + 'a;
10
11/// This trait is used as a "Optional Trait"(I 'd like to use `Trace?` but it's not allowed yet) for PyObjectPayload type
12///
13/// impl for PyObjectPayload, `pyclass` proc macro will handle the actual dispatch if type impl `Trace`
14/// Every PyObjectPayload impl `MaybeTrace`, which may or may not be traceable
15pub trait MaybeTraverse {
16    /// if is traceable, will be used by vtable to determine
17    const HAS_TRAVERSE: bool = false;
18    /// if has clear implementation for circular reference resolution (tp_clear)
19    const HAS_CLEAR: bool = false;
20    // if this type is traceable, then call with tracer_fn, default to do nothing
21    fn try_traverse(&self, traverse_fn: &mut TraverseFn<'_>);
22    // if this type has clear, extract child refs for circular reference resolution (tp_clear)
23    fn try_clear(&mut self, _out: &mut Vec<PyObjectRef>) {}
24}
25
26/// Type that need traverse it's children should impl [`Traverse`] (not [`MaybeTraverse`])
27/// # Safety
28/// Please carefully read [`Traverse::traverse()`] and follow the guideline
29pub unsafe trait Traverse {
30    /// impl `traverse()` with caution! Following those guideline so traverse doesn't cause memory error!:
31    /// - Make sure that every owned object(Every PyObjectRef/PyRef) is called with traverse_fn **at most once**.
32    ///   If some field is not called, the worst results is just memory leak,
33    ///   but if some field is called repeatedly, panic and deadlock can happen.
34    ///
35    /// - _**DO NOT**_ clone a [`PyObjectRef`] or [`PyRef<T>`] in [`Traverse::traverse()`]
36    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>);
37
38    /// Extract all owned child PyObjectRefs for circular reference resolution (tp_clear).
39    /// Called just before object deallocation to break circular references.
40    /// Default implementation does nothing.
41    fn clear(&mut self, _out: &mut Vec<PyObjectRef>) {}
42}
43
44unsafe impl Traverse for PyObjectRef {
45    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
46        traverse_fn(self)
47    }
48}
49
50unsafe impl Traverse for PyStackRef {
51    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
52        traverse_fn(self.as_object())
53    }
54}
55
56unsafe impl<T: PyObjectPayload> Traverse for PyRef<T> {
57    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
58        traverse_fn(self.as_object())
59    }
60}
61
62unsafe impl Traverse for () {
63    fn traverse(&self, _traverse_fn: &mut TraverseFn<'_>) {}
64}
65
66unsafe impl<T: Traverse> Traverse for Option<T> {
67    #[inline]
68    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
69        if let Some(v) = self {
70            v.traverse(traverse_fn);
71        }
72    }
73}
74
75unsafe impl<T> Traverse for [T]
76where
77    T: Traverse,
78{
79    #[inline]
80    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
81        for elem in self {
82            elem.traverse(traverse_fn);
83        }
84    }
85}
86
87unsafe impl<T> Traverse for Box<[T]>
88where
89    T: Traverse,
90{
91    #[inline]
92    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
93        for elem in &**self {
94            elem.traverse(traverse_fn);
95        }
96    }
97}
98
99unsafe impl<T> Traverse for Vec<T>
100where
101    T: Traverse,
102{
103    #[inline]
104    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
105        for elem in self {
106            elem.traverse(traverse_fn);
107        }
108    }
109}
110
111unsafe impl<T: Traverse> Traverse for PyRwLock<T> {
112    #[inline]
113    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
114        // if can't get a lock, this means something else is holding the lock,
115        // but since gc stopped the world, during gc the lock is always held
116        // so it is safe to ignore those in gc
117        if let Some(inner) = self.try_read_recursive() {
118            inner.traverse(traverse_fn)
119        }
120    }
121}
122
123/// Safety: We can't hold lock during traverse it's child because it may cause deadlock.
124/// TODO(discord9): check if this is thread-safe to do
125/// (Outside of gc phase, only incref/decref will call trace,
126/// and refcnt is atomic, so it should be fine?)
127unsafe impl<T: Traverse> Traverse for PyMutex<T> {
128    #[inline]
129    fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
130        let mut chs: Vec<NonNull<PyObject>> = Vec::new();
131        if let Some(obj) = self.try_lock() {
132            obj.traverse(&mut |ch| {
133                chs.push(NonNull::from(ch));
134            })
135        }
136        chs.iter()
137            .map(|ch| {
138                // Safety: during gc, this should be fine, because nothing should write during gc's tracing?
139                let ch = unsafe { ch.as_ref() };
140                traverse_fn(ch);
141            })
142            .count();
143    }
144}
145
146macro_rules! trace_tuple {
147    ($(($NAME: ident, $NUM: tt)),*) => {
148        unsafe impl<$($NAME: Traverse),*> Traverse for ($($NAME),*) {
149            #[inline]
150            fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {
151                $(
152                    self.$NUM.traverse(traverse_fn);
153                )*
154            }
155        }
156
157    };
158}
159
160unsafe impl<A: Traverse, B: Traverse> Traverse for Either<A, B> {
161    #[inline]
162    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
163        match self {
164            Self::A(a) => a.traverse(tracer_fn),
165            Self::B(b) => b.traverse(tracer_fn),
166        }
167    }
168}
169
170// only tuple with 12 elements or less is supported,
171// because long tuple is extremely rare in almost every case
172unsafe impl<A: Traverse> Traverse for (A,) {
173    #[inline]
174    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
175        self.0.traverse(tracer_fn);
176    }
177}
178
179trace_tuple!((A, 0), (B, 1));
180trace_tuple!((A, 0), (B, 1), (C, 2));
181trace_tuple!((A, 0), (B, 1), (C, 2), (D, 3));
182trace_tuple!((A, 0), (B, 1), (C, 2), (D, 3), (E, 4));
183trace_tuple!((A, 0), (B, 1), (C, 2), (D, 3), (E, 4), (F, 5));
184trace_tuple!((A, 0), (B, 1), (C, 2), (D, 3), (E, 4), (F, 5), (G, 6));
185trace_tuple!(
186    (A, 0),
187    (B, 1),
188    (C, 2),
189    (D, 3),
190    (E, 4),
191    (F, 5),
192    (G, 6),
193    (H, 7)
194);
195trace_tuple!(
196    (A, 0),
197    (B, 1),
198    (C, 2),
199    (D, 3),
200    (E, 4),
201    (F, 5),
202    (G, 6),
203    (H, 7),
204    (I, 8)
205);
206trace_tuple!(
207    (A, 0),
208    (B, 1),
209    (C, 2),
210    (D, 3),
211    (E, 4),
212    (F, 5),
213    (G, 6),
214    (H, 7),
215    (I, 8),
216    (J, 9)
217);
218trace_tuple!(
219    (A, 0),
220    (B, 1),
221    (C, 2),
222    (D, 3),
223    (E, 4),
224    (F, 5),
225    (G, 6),
226    (H, 7),
227    (I, 8),
228    (J, 9),
229    (K, 10)
230);
231trace_tuple!(
232    (A, 0),
233    (B, 1),
234    (C, 2),
235    (D, 3),
236    (E, 4),
237    (F, 5),
238    (G, 6),
239    (H, 7),
240    (I, 8),
241    (J, 9),
242    (K, 10),
243    (L, 11)
244);
rustpython_vm/object/traverse.rs

rustpython_vm/object/
traverse.rs
