Skip to content

Commit a997f2b

Browse files
committed
Own generated native closures and retain their method definitions
Assisted-by: Codex:gpt-6
1 parent 7536730 commit a997f2b

4 files changed

Lines changed: 159 additions & 25 deletions

File tree

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

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,19 @@ use crate::{
55
common::wtf8::Wtf8,
66
convert::TryFromObject,
77
function::{Callee, FuncArgs, PyComparisonValue, PyMethodDef, PyMethodFlags, PyNativeFn},
8+
object::{Traverse, TraverseFn},
89
types::{Callable, Comparable, PyComparisonOp, Representable},
910
};
1011
use alloc::fmt;
1112

1213
// PyCFunctionObject in CPython
1314
#[repr(C)]
14-
#[pyclass(name = "builtin_function_or_method", module = false, traverse)]
15+
#[pyclass(
16+
name = "builtin_function_or_method",
17+
module = false,
18+
traverse = "manual"
19+
)]
1520
pub struct PyNativeFunction {
16-
#[pytraverse(skip)]
1721
pub(crate) value: &'static PyMethodDef,
1822
pub(crate) zelf: Option<PyObjectRef>,
1923
// Module that owns this function. Not passed as a call argument.
@@ -24,6 +28,23 @@ pub struct PyNativeFunction {
2428
pub(crate) _method_def_owner: Option<PyObjectRef>,
2529
}
2630

31+
unsafe impl Traverse for PyNativeFunction {
32+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
33+
self.zelf.traverse(tracer_fn);
34+
self.module_object.traverse(tracer_fn);
35+
self.module.traverse(tracer_fn);
36+
self._method_def_owner.traverse(tracer_fn);
37+
}
38+
39+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
40+
out.extend(self.zelf.take());
41+
out.extend(self.module_object.take());
42+
// GC has exclusive access while clearing this unreachable object.
43+
out.extend(unsafe { self.module.swap(None) });
44+
// Keep the definition owner until deallocation: `value` borrows it.
45+
}
46+
}
47+
2748
impl PyPayload for PyNativeFunction {
2849
fn class(ctx: &Context) -> &'static Py<PyType> {
2950
ctx.types.builtin_function_or_method_type
@@ -78,7 +99,7 @@ impl PyNativeFunction {
7899
self.zelf.as_deref().or(self.module_object.as_deref())
79100
}
80101

81-
pub const fn as_func(&self) -> &'static dyn PyNativeFn {
102+
pub const fn as_func(&self) -> &dyn PyNativeFn {
82103
self.value.func
83104
}
84105
}
@@ -234,12 +255,22 @@ impl PyNativeFunction {
234255
}
235256

236257
// PyCMethodObject in CPython
237-
#[pyclass(name = "builtin_function_or_method", module = false, base = PyNativeFunction, ctx = "builtin_function_or_method_type")]
258+
#[pyclass(name = "builtin_function_or_method", module = false, base = PyNativeFunction, ctx = "builtin_function_or_method_type", traverse = "manual")]
238259
pub struct PyNativeMethod {
239260
pub(crate) func: PyNativeFunction,
240261
pub(crate) class: &'static Py<PyType>, // TODO: the actual life is &'self
241262
}
242263

264+
unsafe impl Traverse for PyNativeMethod {
265+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
266+
self.func.traverse(tracer_fn);
267+
}
268+
269+
fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
270+
self.func.clear(out);
271+
}
272+
}
273+
243274
// All Python-visible behavior (getters, slots) is registered by PyNativeFunction::extend_class.
244275
// PyNativeMethod only extends the Rust-side struct with the defining class reference.
245276
// The func field at offset 0 (#[repr(C)]) allows NativeFunctionOrMethod to read it safely.

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

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::{
88
function::{
99
Callee, FuncArgs, ItemDoc, PyMethodDef, PyMethodFlags, PySetterValue, PySsize, plain_doc,
1010
},
11+
object::{Traverse, TraverseFn},
1112
protocol::{PyNumberBinaryFunc, PyNumberTernaryFunc, PyNumberUnaryFunc},
1213
types::{
1314
Callable, Comparable, DelFunc, DescrGetFunc, DescrSetFunc, GenericMethod, GetDescriptor,
@@ -34,17 +35,37 @@ pub struct PyDescriptorOwned {
3435
pub qualname: PyRwLock<Option<String>>,
3536
}
3637

37-
#[pyclass(name = "method_descriptor", module = false)]
38+
impl PyDescriptor {
39+
fn bind(
40+
&self,
41+
method: &'static PyMethodDef,
42+
owner: Option<PyObjectRef>,
43+
obj: PyObjectRef,
44+
ctx: &Context,
45+
) -> PyRef<PyNativeMethod> {
46+
let mut bound = method.to_bound_method(obj, self.typ);
47+
bound.func._method_def_owner = owner;
48+
bound.into_ref(ctx)
49+
}
50+
}
51+
52+
#[pyclass(name = "method_descriptor", module = false, traverse = "manual")]
3853
pub struct PyMethodDescriptor {
3954
#[pymember(name = "__objclass__", path = "typ")]
4055
#[pymember(name = "__name__", path = "name")]
4156
pub common: PyDescriptor,
42-
pub method: &'static PyMethodDef,
57+
pub(crate) method: &'static PyMethodDef,
4358
// vectorcall: vector_call_func,
4459
/// Prevent HeapMethodDef from being freed while this descriptor references it
4560
pub(crate) _method_def_owner: Option<PyObjectRef>,
4661
}
4762

63+
unsafe impl Traverse for PyMethodDescriptor {
64+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
65+
self._method_def_owner.traverse(tracer_fn);
66+
}
67+
}
68+
4869
impl PyMethodDescriptor {
4970
pub fn new(method: &'static PyMethodDef, typ: &'static Py<PyType>, ctx: &Context) -> Self {
5071
Self {
@@ -125,7 +146,8 @@ impl Callable for PyMethodDescriptor {
125146

126147
impl PyMethodDescriptor {
127148
pub fn bind(&self, obj: PyObjectRef, ctx: &Context) -> PyRef<PyNativeMethod> {
128-
self.method.build_bound_method(ctx, obj, self.common.typ)
149+
self.common
150+
.bind(self.method, self._method_def_owner.clone(), obj, ctx)
129151
}
130152
}
131153

@@ -173,15 +195,21 @@ impl Representable for PyMethodDescriptor {
173195
}
174196

175197
// METH_CLASS descriptors. Same layout as method_descriptor; a distinct type.
176-
#[pyclass(name = "classmethod_descriptor", module = false)]
198+
#[pyclass(name = "classmethod_descriptor", module = false, traverse = "manual")]
177199
pub struct PyClassMethodDescriptor {
178200
#[pymember(name = "__objclass__", path = "typ")]
179201
#[pymember(name = "__name__", path = "name")]
180202
pub common: PyDescriptor,
181-
pub method: &'static PyMethodDef,
203+
pub(crate) method: &'static PyMethodDef,
182204
pub(crate) _method_def_owner: Option<PyObjectRef>,
183205
}
184206

207+
unsafe impl Traverse for PyClassMethodDescriptor {
208+
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
209+
self._method_def_owner.traverse(tracer_fn);
210+
}
211+
}
212+
185213
impl PyClassMethodDescriptor {
186214
pub fn new(method: &'static PyMethodDef, typ: &'static Py<PyType>, ctx: &Context) -> Self {
187215
Self {
@@ -196,7 +224,8 @@ impl PyClassMethodDescriptor {
196224
}
197225

198226
pub fn bind(&self, obj: PyObjectRef, ctx: &Context) -> PyRef<PyNativeMethod> {
199-
self.method.build_bound_method(ctx, obj, self.common.typ)
227+
self.common
228+
.bind(self.method, self._method_def_owner.clone(), obj, ctx)
200229
}
201230
}
202231

‎crates/vm/src/function/method.rs‎

Lines changed: 88 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -375,15 +375,51 @@ impl core::fmt::Debug for PyMethodDef {
375375
// This is not a part of CPython API.
376376
// But useful to support dynamically generated methods
377377
#[pyclass(name, module = false, ctx = "method_def")]
378-
#[derive(Debug)]
379378
pub struct HeapMethodDef {
380379
method: PyMethodDef,
380+
_function_owner: Option<Box<dyn PyNativeFn>>,
381+
}
382+
383+
impl core::fmt::Debug for HeapMethodDef {
384+
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
385+
self.method.fmt(f)
386+
}
381387
}
382388

383389
impl HeapMethodDef {
384390
#[must_use]
385391
pub const fn new(method: PyMethodDef) -> Self {
386-
Self { method }
392+
Self {
393+
method,
394+
_function_owner: None,
395+
}
396+
}
397+
398+
pub(crate) fn with_owned_function(
399+
name: &'static str,
400+
function: Box<dyn PyNativeFn>,
401+
flags: PyMethodFlags,
402+
doc: super::ItemDoc,
403+
) -> Self {
404+
// SAFETY: the box has a stable address and is retained by this payload.
405+
// The definition is private; all functions/descriptors referring to it
406+
// retain this HeapMethodDef and expose no static reference to callers.
407+
let func = unsafe { &*core::ptr::from_ref(&*function) };
408+
Self {
409+
method: PyMethodDef {
410+
name,
411+
func,
412+
flags,
413+
#[cfg(feature = "doc")]
414+
doc_off: doc.offset,
415+
#[cfg(feature = "doc")]
416+
doc_len: doc.len,
417+
#[cfg(feature = "doc")]
418+
doc_body_pending: false,
419+
doc: doc.text,
420+
},
421+
_function_owner: Some(function),
422+
}
387423
}
388424
}
389425

@@ -431,3 +467,53 @@ pub(crate) fn init(ctx: &'static Context) {
431467

432468
let _ = ctx.intern_str(HeapMethodDef::NAME);
433469
}
470+
471+
#[cfg(test)]
472+
mod tests {
473+
use super::*;
474+
use crate::{AsObject, function::FuncArgs};
475+
476+
#[test]
477+
fn generated_functions_own_and_release_their_native_captures() {
478+
use alloc::sync::Arc;
479+
let captured = Arc::new(17);
480+
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
481+
let capture = captured.clone();
482+
let definition = vm.ctx.new_method_def(
483+
"captured",
484+
move |_args: FuncArgs| *capture,
485+
PyMethodFlags::empty(),
486+
super::super::ItemDoc::NONE,
487+
);
488+
let function = definition.build_function(vm, None);
489+
let descriptor = definition.build_method(vm.ctx.types.object_type, vm);
490+
drop(definition);
491+
assert_eq!(Arc::strong_count(&captured), 2);
492+
assert_eq!(
493+
function
494+
.as_object()
495+
.call((), vm)
496+
.unwrap()
497+
.try_into_value::<i32>(vm)
498+
.unwrap(),
499+
17
500+
);
501+
let method = descriptor.bind(vm.ctx.none(), &vm.ctx);
502+
assert!(method.func._method_def_owner.is_some());
503+
drop(function);
504+
drop(descriptor);
505+
assert_eq!(Arc::strong_count(&captured), 2);
506+
assert_eq!(
507+
method
508+
.as_object()
509+
.call((), vm)
510+
.unwrap()
511+
.try_into_value::<i32>(vm)
512+
.unwrap(),
513+
17
514+
);
515+
drop(method);
516+
assert_eq!(Arc::strong_count(&captured), 1);
517+
});
518+
}
519+
}

‎crates/vm/src/vm/context.rs‎

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -659,19 +659,7 @@ impl Context {
659659
where
660660
F: IntoPyNativeFn<FKind>,
661661
{
662-
let def = PyMethodDef {
663-
name,
664-
func: Box::leak(Box::new(f.into_func())),
665-
flags,
666-
#[cfg(feature = "doc")]
667-
doc_off: doc.offset,
668-
#[cfg(feature = "doc")]
669-
doc_len: doc.len,
670-
#[cfg(feature = "doc")]
671-
doc_body_pending: false,
672-
doc: doc.text,
673-
};
674-
let payload = HeapMethodDef::new(def);
662+
let payload = HeapMethodDef::with_owned_function(name, Box::new(f.into_func()), flags, doc);
675663
PyRef::new_ref(payload, self.types.method_def.to_owned(), None)
676664
}
677665

0 commit comments

Comments
 (0)