traverse for generator (#6760)

crates/vm/src/builtins/asyncgenerator.rs

@@ -7,13 +7,14 @@ use crate::{
     coroutine::{Coro, warn_deprecated_throw_signature},
     frame::FrameRef,
     function::OptionalArg,
+    object::{Traverse, TraverseFn},
     protocol::PyIterReturn,
     types::{Destructor, IterNext, Iterable, Representable, SelfIter},
 };
 
 use crossbeam_utils::atomic::AtomicCell;
 
-#[pyclass(name = "async_generator", module = false)]
+#[pyclass(name = "async_generator", module = false, traverse = "manual")]
 #[derive(Debug)]
 pub struct PyAsyncGen {
     inner: Coro,
@@ -23,6 +24,13 @@ pub struct PyAsyncGen {
     // ag_origin_or_finalizer - stores the finalizer callback
     ag_finalizer: PyMutex<Option<PyObjectRef>>,
 }
+
+unsafe impl Traverse for PyAsyncGen {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.inner.traverse(tracer_fn);
+        self.ag_finalizer.traverse(tracer_fn);
+    }
+}
 type PyAsyncGenRef = PyRef<PyAsyncGen>;
 
 impl PyPayload for PyAsyncGen {
@@ -199,9 +207,20 @@ impl Representable for PyAsyncGen {
     }
 }
 
-#[pyclass(module = false, name = "async_generator_wrapped_value")]
+#[pyclass(
+    module = false,
+    name = "async_generator_wrapped_value",
+    traverse = "manual"
+)]
 #[derive(Debug)]
 pub(crate) struct PyAsyncGenWrappedValue(pub PyObjectRef);
+
+unsafe impl Traverse for PyAsyncGenWrappedValue {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.0.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyAsyncGenWrappedValue {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -244,7 +263,7 @@ enum AwaitableState {
     Closed,
 }
 
-#[pyclass(module = false, name = "async_generator_asend")]
+#[pyclass(module = false, name = "async_generator_asend", traverse = "manual")]
 #[derive(Debug)]
 pub(crate) struct PyAsyncGenASend {
     ag: PyAsyncGenRef,
@@ -252,6 +271,13 @@ pub(crate) struct PyAsyncGenASend {
     value: PyObjectRef,
 }
 
+unsafe impl Traverse for PyAsyncGenASend {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.ag.traverse(tracer_fn);
+        self.value.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyAsyncGenASend {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -338,7 +364,7 @@ impl IterNext for PyAsyncGenASend {
     }
 }
 
-#[pyclass(module = false, name = "async_generator_athrow")]
+#[pyclass(module = false, name = "async_generator_athrow", traverse = "manual")]
 #[derive(Debug)]
 pub(crate) struct PyAsyncGenAThrow {
     ag: PyAsyncGenRef,
@@ -347,6 +373,13 @@ pub(crate) struct PyAsyncGenAThrow {
     value: (PyObjectRef, PyObjectRef, PyObjectRef),
 }
 
+unsafe impl Traverse for PyAsyncGenAThrow {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.ag.traverse(tracer_fn);
+        self.value.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyAsyncGenAThrow {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -489,7 +522,7 @@ impl IterNext for PyAsyncGenAThrow {
 
 /// Awaitable wrapper for anext() builtin with default value.
 /// When StopAsyncIteration is raised, it converts it to StopIteration(default).
-#[pyclass(module = false, name = "anext_awaitable")]
+#[pyclass(module = false, name = "anext_awaitable", traverse = "manual")]
 #[derive(Debug)]
 pub struct PyAnextAwaitable {
     wrapped: PyObjectRef,
@@ -497,6 +530,13 @@ pub struct PyAnextAwaitable {
     state: AtomicCell<AwaitableState>,
 }
 
+unsafe impl Traverse for PyAnextAwaitable {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.wrapped.traverse(tracer_fn);
+        self.default_value.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyAnextAwaitable {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {

crates/vm/src/builtins/coroutine.rs

@@ -5,18 +5,25 @@ use crate::{
     coroutine::{Coro, warn_deprecated_throw_signature},
     frame::FrameRef,
     function::OptionalArg,
+    object::{Traverse, TraverseFn},
     protocol::PyIterReturn,
     types::{IterNext, Iterable, Representable, SelfIter},
 };
 use crossbeam_utils::atomic::AtomicCell;
 
-#[pyclass(module = false, name = "coroutine")]
+#[pyclass(module = false, name = "coroutine", traverse = "manual")]
 #[derive(Debug)]
 // PyCoro_Type in CPython
 pub struct PyCoroutine {
     inner: Coro,
 }
 
+unsafe impl Traverse for PyCoroutine {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.inner.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyCoroutine {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
@@ -138,7 +145,7 @@ impl IterNext for PyCoroutine {
     }
 }
 
-#[pyclass(module = false, name = "coroutine_wrapper")]
+#[pyclass(module = false, name = "coroutine_wrapper", traverse = "manual")]
 #[derive(Debug)]
 // PyCoroWrapper_Type in CPython
 pub struct PyCoroutineWrapper {
@@ -146,6 +153,12 @@ pub struct PyCoroutineWrapper {
     closed: AtomicCell<bool>,
 }
 
+unsafe impl Traverse for PyCoroutineWrapper {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.coro.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyCoroutineWrapper {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {

crates/vm/src/builtins/generator.rs

@@ -9,16 +9,23 @@ use crate::{
     coroutine::{Coro, warn_deprecated_throw_signature},
     frame::FrameRef,
     function::OptionalArg,
+    object::{Traverse, TraverseFn},
     protocol::PyIterReturn,
     types::{IterNext, Iterable, Representable, SelfIter},
 };
 
-#[pyclass(module = false, name = "generator")]
+#[pyclass(module = false, name = "generator", traverse = "manual")]
 #[derive(Debug)]
 pub struct PyGenerator {
     inner: Coro,
 }
 
+unsafe impl Traverse for PyGenerator {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.inner.traverse(tracer_fn);
+    }
+}
+
 impl PyPayload for PyGenerator {
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {

crates/vm/src/coroutine.rs

@@ -5,6 +5,7 @@ use crate::{
     exceptions::types::PyBaseException,
     frame::{ExecutionResult, FrameRef},
     function::OptionalArg,
+    object::{Traverse, TraverseFn},
     protocol::PyIterReturn,
 };
 use crossbeam_utils::atomic::AtomicCell;
@@ -38,6 +39,15 @@ pub struct Coro {
     exception: PyMutex<Option<PyBaseExceptionRef>>, // exc_state
 }
 
+unsafe impl Traverse for Coro {
+    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
+        self.frame.traverse(tracer_fn);
+        self.name.traverse(tracer_fn);
+        self.qualname.traverse(tracer_fn);
+        self.exception.traverse(tracer_fn);
+    }
+}
+
 fn gen_name(jen: &PyObject, vm: &VirtualMachine) -> &'static str {
     let typ = jen.class();
     if typ.is(vm.ctx.types.coroutine_type) {