diff --git a/crates/vm/src/builtins/asyncgenerator.rs b/crates/vm/src/builtins/asyncgenerator.rs index 891083f3e6b..8f7518f0eae 100644 --- a/crates/vm/src/builtins/asyncgenerator.rs +++ b/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>, } + +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; 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 { @@ -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 { @@ -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 { @@ -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, } +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 { diff --git a/crates/vm/src/builtins/coroutine.rs b/crates/vm/src/builtins/coroutine.rs index 9e8d5d534f1..961c352f8df 100644 --- a/crates/vm/src/builtins/coroutine.rs +++ b/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 { @@ -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, } +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 { diff --git a/crates/vm/src/builtins/generator.rs b/crates/vm/src/builtins/generator.rs index 9a1e737500b..ceae2e61c3b 100644 --- a/crates/vm/src/builtins/generator.rs +++ b/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 { diff --git a/crates/vm/src/coroutine.rs b/crates/vm/src/coroutine.rs index 19830496e6e..236388ecf8b 100644 --- a/crates/vm/src/coroutine.rs +++ b/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>, // 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) {