Skip to content

Commit abea6bd

Browse files
youknowoneyouknowone
authored
traverse for generator (#6760)
1 parent ca76cb7 commit abea6bd

File tree

4 files changed

+78
-8
lines changed

4 files changed

+78
-8
lines changed

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

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,14 @@ use crate::{
77
coroutine::{Coro, warn_deprecated_throw_signature},
88
frame::FrameRef,
99
function::OptionalArg,
10+
object::{Traverse, TraverseFn},
1011
protocol::PyIterReturn,
1112
types::{Destructor, IterNext, Iterable, Representable, SelfIter},
1213
};
1314

1415
use crossbeam_utils::atomic::AtomicCell;
1516

16-
#[pyclass(name = "async_generator", module = false)]
17+
#[pyclass(name = "async_generator", module = false, traverse = "manual")]
1718
#[derive(Debug)]
1819
pub struct PyAsyncGen {
1920
inner: Coro,
@@ -23,6 +24,13 @@ pub struct PyAsyncGen {
2324
// ag_origin_or_finalizer - stores the finalizer callback
2425
ag_finalizer: PyMutex<Option<PyObjectRef>>,
2526
}
27+
28+
unsafe impl Traverse for PyAsyncGen {
29+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
30+
self.inner.traverse(tracer_fn);
31+
self.ag_finalizer.traverse(tracer_fn);
32+
}
33+
}
2634
type PyAsyncGenRef = PyRef<PyAsyncGen>;
2735

2836
impl PyPayload for PyAsyncGen {
@@ -199,9 +207,20 @@ impl Representable for PyAsyncGen {
199207
}
200208
}
201209

202-
#[pyclass(module = false, name = "async_generator_wrapped_value")]
210+
#[pyclass(
211+
module = false,
212+
name = "async_generator_wrapped_value",
213+
traverse = "manual"
214+
)]
203215
#[derive(Debug)]
204216
pub(crate) struct PyAsyncGenWrappedValue(pub PyObjectRef);
217+
218+
unsafe impl Traverse for PyAsyncGenWrappedValue {
219+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
220+
self.0.traverse(tracer_fn);
221+
}
222+
}
223+
205224
impl PyPayload for PyAsyncGenWrappedValue {
206225
#[inline]
207226
fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -244,14 +263,21 @@ enum AwaitableState {
244263
Closed,
245264
}
246265

247-
#[pyclass(module = false, name = "async_generator_asend")]
266+
#[pyclass(module = false, name = "async_generator_asend", traverse = "manual")]
248267
#[derive(Debug)]
249268
pub(crate) struct PyAsyncGenASend {
250269
ag: PyAsyncGenRef,
251270
state: AtomicCell<AwaitableState>,
252271
value: PyObjectRef,
253272
}
254273

274+
unsafe impl Traverse for PyAsyncGenASend {
275+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
276+
self.ag.traverse(tracer_fn);
277+
self.value.traverse(tracer_fn);
278+
}
279+
}
280+
255281
impl PyPayload for PyAsyncGenASend {
256282
#[inline]
257283
fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -338,7 +364,7 @@ impl IterNext for PyAsyncGenASend {
338364
}
339365
}
340366

341-
#[pyclass(module = false, name = "async_generator_athrow")]
367+
#[pyclass(module = false, name = "async_generator_athrow", traverse = "manual")]
342368
#[derive(Debug)]
343369
pub(crate) struct PyAsyncGenAThrow {
344370
ag: PyAsyncGenRef,
@@ -347,6 +373,13 @@ pub(crate) struct PyAsyncGenAThrow {
347373
value: (PyObjectRef, PyObjectRef, PyObjectRef),
348374
}
349375

376+
unsafe impl Traverse for PyAsyncGenAThrow {
377+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
378+
self.ag.traverse(tracer_fn);
379+
self.value.traverse(tracer_fn);
380+
}
381+
}
382+
350383
impl PyPayload for PyAsyncGenAThrow {
351384
#[inline]
352385
fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -489,14 +522,21 @@ impl IterNext for PyAsyncGenAThrow {
489522

490523
/// Awaitable wrapper for anext() builtin with default value.
491524
/// When StopAsyncIteration is raised, it converts it to StopIteration(default).
492-
#[pyclass(module = false, name = "anext_awaitable")]
525+
#[pyclass(module = false, name = "anext_awaitable", traverse = "manual")]
493526
#[derive(Debug)]
494527
pub struct PyAnextAwaitable {
495528
wrapped: PyObjectRef,
496529
default_value: PyObjectRef,
497530
state: AtomicCell<AwaitableState>,
498531
}
499532

533+
unsafe impl Traverse for PyAnextAwaitable {
534+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
535+
self.wrapped.traverse(tracer_fn);
536+
self.default_value.traverse(tracer_fn);
537+
}
538+
}
539+
500540
impl PyPayload for PyAnextAwaitable {
501541
#[inline]
502542
fn class(ctx: &Context) -> &'static Py<PyType> {

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,25 @@ use crate::{
55
coroutine::{Coro, warn_deprecated_throw_signature},
66
frame::FrameRef,
77
function::OptionalArg,
8+
object::{Traverse, TraverseFn},
89
protocol::PyIterReturn,
910
types::{IterNext, Iterable, Representable, SelfIter},
1011
};
1112
use crossbeam_utils::atomic::AtomicCell;
1213

13-
#[pyclass(module = false, name = "coroutine")]
14+
#[pyclass(module = false, name = "coroutine", traverse = "manual")]
1415
#[derive(Debug)]
1516
// PyCoro_Type in CPython
1617
pub struct PyCoroutine {
1718
inner: Coro,
1819
}
1920

21+
unsafe impl Traverse for PyCoroutine {
22+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
23+
self.inner.traverse(tracer_fn);
24+
}
25+
}
26+
2027
impl PyPayload for PyCoroutine {
2128
#[inline]
2229
fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -138,14 +145,20 @@ impl IterNext for PyCoroutine {
138145
}
139146
}
140147

141-
#[pyclass(module = false, name = "coroutine_wrapper")]
148+
#[pyclass(module = false, name = "coroutine_wrapper", traverse = "manual")]
142149
#[derive(Debug)]
143150
// PyCoroWrapper_Type in CPython
144151
pub struct PyCoroutineWrapper {
145152
coro: PyRef<PyCoroutine>,
146153
closed: AtomicCell<bool>,
147154
}
148155

156+
unsafe impl Traverse for PyCoroutineWrapper {
157+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
158+
self.coro.traverse(tracer_fn);
159+
}
160+
}
161+
149162
impl PyPayload for PyCoroutineWrapper {
150163
#[inline]
151164
fn class(ctx: &Context) -> &'static Py<PyType> {

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,23 @@ use crate::{
99
coroutine::{Coro, warn_deprecated_throw_signature},
1010
frame::FrameRef,
1111
function::OptionalArg,
12+
object::{Traverse, TraverseFn},
1213
protocol::PyIterReturn,
1314
types::{IterNext, Iterable, Representable, SelfIter},
1415
};
1516

16-
#[pyclass(module = false, name = "generator")]
17+
#[pyclass(module = false, name = "generator", traverse = "manual")]
1718
#[derive(Debug)]
1819
pub struct PyGenerator {
1920
inner: Coro,
2021
}
2122

23+
unsafe impl Traverse for PyGenerator {
24+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
25+
self.inner.traverse(tracer_fn);
26+
}
27+
}
28+
2229
impl PyPayload for PyGenerator {
2330
#[inline]
2431
fn class(ctx: &Context) -> &'static Py<PyType> {

‎crates/vm/src/coroutine.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use crate::{
55
exceptions::types::PyBaseException,
66
frame::{ExecutionResult, FrameRef},
77
function::OptionalArg,
8+
object::{Traverse, TraverseFn},
89
protocol::PyIterReturn,
910
};
1011
use crossbeam_utils::atomic::AtomicCell;
@@ -38,6 +39,15 @@ pub struct Coro {
3839
exception: PyMutex<Option<PyBaseExceptionRef>>, // exc_state
3940
}
4041

42+
unsafe impl Traverse for Coro {
43+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
44+
self.frame.traverse(tracer_fn);
45+
self.name.traverse(tracer_fn);
46+
self.qualname.traverse(tracer_fn);
47+
self.exception.traverse(tracer_fn);
48+
}
49+
}
50+
4151
fn gen_name(jen: &PyObject, vm: &VirtualMachine) -> &'static str {
4252
let typ = jen.class();
4353
if typ.is(vm.ctx.types.coroutine_type) {

0 commit comments

Comments
 (0)