weakref: reuse proxy objects (#8307)

* weakref: reuse callback-less proxy objects like CPython

Make PyWeakProxy a #[repr(transparent)] view sharing PyWeak's payload
(manual PyPayload with shared PAYLOAD_TYPE_ID + class check), so a proxy
is the weakref-list node itself and add()'s reuse applies to it
directly. Port CPython's insert_weakref/get_basic_refs slot discipline:
genericref at head, generic proxy right after it, with callback/class
re-verificationbefore reuse. Constructor moves to a slot_new override
since reuse returns an existing object. Drops the hidden __weakproxy
marker subclass and the proxy branch in PyWeakref_GetRef.

Assisted-by: Claude Code:claude-opus-4-8

* weakref: always validate class in basic-proxy slot lookup

* weakref: rustfmt find_generic_proxy_ptr

* weakref: unify reuse check shape

Lib/test/test_weakref.py

@@ -289,7 +289,6 @@ class ReferencesTestCase(TestBase):
         self.assertEqual(weakref.getweakrefcount(o), 1,
                      "wrong weak ref count for object after deleting proxy")
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_proxy_reuse(self):
         o = C()
         proxy1 = weakref.proxy(o)
@@ -383,7 +382,6 @@ class ReferencesTestCase(TestBase):
     def test_shared_ref_without_callback(self):
         self.check_shared_without_callback(weakref.ref)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_shared_proxy_without_callback(self):
         self.check_shared_without_callback(weakref.proxy)
 

crates/capi/src/weakrefobject.rs

@@ -20,8 +20,6 @@ pub unsafe extern "C" fn PyWeakref_GetRef(
         let reference = unsafe { &*reference };
         let upgraded = if let Some(weak) = reference.downcast_ref::<PyWeak>() {
             weak.upgrade()
-        } else if let Some(proxy) = reference.downcast_ref::<PyWeakProxy>() {
-            proxy.get_weak().upgrade()
         } else {
             return Err(vm.new_type_error("expected a weakref"));
         };

crates/vm/src/builtins/weakproxy.rs

@@ -4,7 +4,7 @@ use crate::{
     Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
     class::PyClassImpl,
     common::hash::PyHash,
-    function::{OptionalArg, PyArithmeticValue, PyComparisonValue, PySetterValue},
+    function::{FuncArgs, OptionalArg, PyArithmeticValue, PyComparisonValue, PySetterValue},
     protocol::{PyIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     stdlib::builtins::reversed,
     types::{
@@ -13,13 +13,22 @@ use crate::{
     },
 };
 
-#[pyclass(module = false, name = "weakproxy", unhashable = true, traverse)]
+#[pyclass(module = false, name = "weakproxy", unhashable = true)]
 #[derive(Debug)]
-pub struct PyWeakProxy {
-    weak: PyRef<PyWeak>,
-}
+#[repr(transparent)]
+pub struct PyWeakProxy(PyWeak);
 
 impl PyPayload for PyWeakProxy {
+    const PAYLOAD_TYPE_ID: core::any::TypeId = <PyWeak as PyPayload>::PAYLOAD_TYPE_ID;
+
+    #[inline]
+    unsafe fn validate_downcastable_from(obj: &PyObject) -> bool {
+        <Self as ::rustpython_vm::class::PyClassDef>::BASICSIZE <= obj.class().slots.basicsize
+            && obj
+                .class()
+                .fast_issubclass(<Self as ::rustpython_vm::class::StaticType>::static_type())
+    }
+
     #[inline]
     fn class(ctx: &Context) -> &'static Py<PyType> {
         ctx.types.weakproxy_type
@@ -37,52 +46,34 @@ pub struct WeakProxyNewArgs {
 impl Constructor for PyWeakProxy {
     type Args = WeakProxyNewArgs;
 
-    fn py_new(
-        _cls: &Py<PyType>,
-        Self::Args { referent, callback }: Self::Args,
-        vm: &VirtualMachine,
-    ) -> PyResult<Self> {
-        let weak = Self::new_weak(referent.as_ref(), callback.into_option(), vm)?;
-        // TODO: PyWeakProxy should use the same payload as PyWeak
-        Ok(Self { weak })
+    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
+        let _ = cls;
+        let Self::Args { referent, callback } = args.bind(vm)?;
+        let callback = callback
+            .into_option()
+            .filter(|callback| !vm.is_none(callback));
+        let proxy = Self::new_weakproxy(referent.as_ref(), callback, vm)?;
+        Ok(proxy.into())
     }
-}
 
-crate::common::static_cell! {
-    static WEAK_SUBCLASS: PyTypeRef;
+    fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
+        unimplemented!("use slot_new")
+    }
 }
 
 impl PyWeakProxy {
-    fn new_weak(
-        referent: &PyObject,
-        callback: Option<PyObjectRef>,
-        vm: &VirtualMachine,
-    ) -> PyResult<PyRef<PyWeak>> {
-        // using an internal subclass as the class prevents us from getting the generic weakref,
-        // which would mess up the weakref count
-        let weak_cls = WEAK_SUBCLASS.get_or_init(|| {
-            vm.ctx.new_class(
-                None,
-                "__weakproxy",
-                vm.ctx.types.weakref_type.to_owned(),
-                super::PyWeak::make_slots(),
-            )
-        });
-        referent.downgrade_with_typ(callback, weak_cls.clone(), vm)
-    }
-
     pub fn new_weakproxy(
         referent: &PyObject,
         callback: Option<PyObjectRef>,
         vm: &VirtualMachine,
-    ) -> PyResult<PyRef<Self>> {
-        let weak = Self::new_weak(referent, callback, vm)?;
-        Ok(Self { weak }.into_ref(&vm.ctx))
+    ) -> PyResult<PyRef<PyWeak>> {
+        let typ = vm.ctx.types.weakproxy_type.to_owned();
+        referent.downgrade_with_typ(callback, typ, vm)
     }
 
     #[must_use]
-    pub fn get_weak(&self) -> &PyRef<PyWeak> {
-        &self.weak
+    pub fn get_weak(&self) -> &PyWeak {
+        &self.0
     }
 }
 
@@ -99,7 +90,7 @@ impl PyWeakProxy {
 ))]
 impl PyWeakProxy {
     fn try_upgrade(&self, vm: &VirtualMachine) -> PyResult {
-        self.weak.upgrade().ok_or_else(|| new_reference_error(vm))
+        self.0.upgrade().ok_or_else(|| new_reference_error(vm))
     }
 
     #[pymethod]

crates/vm/src/object/core.rs

@@ -582,6 +582,18 @@ unsafe fn unlink_weakref(wrl: &WeakRefList, node: NonNull<Py<PyWeak>>) {
     }
 }
 
+// try_reuse_basic_ref
+unsafe fn try_reuse_weakref(ptr: *mut Py<PyWeak>) -> Option<PyRef<PyWeak>> {
+    if ptr.is_null() {
+        return None;
+    }
+    let node = unsafe { &*ptr };
+    node.0
+        .ref_count
+        .safe_inc()
+        .then(|| unsafe { PyRef::from_raw(ptr) })
+}
+
 impl WeakRefList {
     pub(super) fn new() -> Self {
         Self {
@@ -596,22 +608,25 @@ impl WeakRefList {
         obj: &PyObject,
         cls: PyTypeRef,
         cls_is_weakref: bool,
+        cls_is_weakproxy: bool,
         callback: Option<PyObjectRef>,
         dict: Option<PyDictRef>,
     ) -> PyRef<PyWeak> {
         let is_generic = cls_is_weakref && callback.is_none();
+        let is_generic_proxy = cls_is_weakproxy && callback.is_none();
 
         // Try reuse under lock first (fast path, no allocation)
         {
             let _lock = weakref_lock::lock(obj as *const PyObject as usize);
-            if is_generic {
-                let generic_ptr = self.generic.load(Ordering::Relaxed);
-                if !generic_ptr.is_null() {
-                    let generic = unsafe { &*generic_ptr };
-                    if generic.0.ref_count.safe_inc() {
-                        return unsafe { PyRef::from_raw(generic_ptr) };
-                    }
-                }
+            let existing = if is_generic {
+                unsafe { try_reuse_weakref(self.generic.load(Ordering::Relaxed)) }
+            } else if is_generic_proxy {
+                unsafe { try_reuse_weakref(self.find_generic_proxy_ptr()) }
+            } else {
+                None
+            };
+            if let Some(existing) = existing {
+                return existing;
             }
         }
 
@@ -629,66 +644,100 @@ impl WeakRefList {
         // Re-acquire lock for linked list insertion
         let _lock = weakref_lock::lock(obj as *const PyObject as usize);
 
-        // Re-check: another thread may have inserted a generic ref while we
-        // were allocating outside the lock. If so, reuse it and drop ours.
-        if is_generic {
-            let generic_ptr = self.generic.load(Ordering::Relaxed);
-            if !generic_ptr.is_null() {
-                let generic = unsafe { &*generic_ptr };
-                if generic.0.ref_count.safe_inc() {
-                    // Nullify wr_object so drop_inner won't unlink an
-                    // un-inserted node (which would corrupt the list head).
-                    weak.wr_object.store(ptr::null_mut(), Ordering::Relaxed);
-                    return unsafe { PyRef::from_raw(generic_ptr) };
-                }
-            }
+        // Re-check: another thread may have inserted a generic ref/proxy
+        // while we were allocating outside the lock. If so, reuse it and
+        // drop ours.
+        let existing = if is_generic {
+            unsafe { try_reuse_weakref(self.generic.load(Ordering::Relaxed)) }
+        } else if is_generic_proxy {
+            unsafe { try_reuse_weakref(self.find_generic_proxy_ptr()) }
+        } else {
+            None
+        };
+        if let Some(existing) = existing {
+            // Nullify wr_object so drop_inner won't unlink an
+            // un-inserted node (which would corrupt the list head).
+            weak.wr_object.store(ptr::null_mut(), Ordering::Relaxed);
+            return existing;
         }
 
         // Insert into linked list under stripe lock
+        // (insert_weakref: generic ref at head, generic proxy right after it)
         let node_ptr = NonNull::from(&*weak);
-        unsafe {
-            let mut ptrs = WeakLink::pointers(node_ptr);
-            if is_generic {
-                // Generic ref goes to head (insert_head for basic ref)
-                let old_head = self.head.load(Ordering::Relaxed);
-                ptrs.as_mut().set_next(NonNull::new(old_head));
-                ptrs.as_mut().set_prev(None);
-                if let Some(old_head) = NonNull::new(old_head) {
-                    WeakLink::pointers(old_head)
-                        .as_mut()
-                        .set_prev(Some(node_ptr));
-                }
-                self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
-                self.generic.store(node_ptr.as_ptr(), Ordering::Relaxed);
-            } else {
-                // Non-generic refs go after generic ref (insert_after)
-                let generic_ptr = self.generic.load(Ordering::Relaxed);
-                if let Some(after) = NonNull::new(generic_ptr) {
-                    let after_next = WeakLink::pointers(after).as_ref().get_next();
-                    ptrs.as_mut().set_prev(Some(after));
-                    ptrs.as_mut().set_next(after_next);
-                    WeakLink::pointers(after).as_mut().set_next(Some(node_ptr));
-                    if let Some(next) = after_next {
-                        WeakLink::pointers(next).as_mut().set_prev(Some(node_ptr));
-                    }
-                } else {
-                    // No generic ref; insert at head
-                    let old_head = self.head.load(Ordering::Relaxed);
-                    ptrs.as_mut().set_next(NonNull::new(old_head));
-                    ptrs.as_mut().set_prev(None);
-                    if let Some(old_head) = NonNull::new(old_head) {
-                        WeakLink::pointers(old_head)
-                            .as_mut()
-                            .set_prev(Some(node_ptr));
-                    }
-                    self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
-                }
-            }
+        let after = if is_generic {
+            None
+        } else if is_generic_proxy {
+            NonNull::new(self.generic.load(Ordering::Relaxed))
+        } else {
+            NonNull::new(self.find_generic_proxy_ptr())
+                .or_else(|| NonNull::new(self.generic.load(Ordering::Relaxed)))
+        };
+        match after {
+            Some(after) => unsafe { self.insert_after(after, node_ptr) },
+            None => unsafe { self.insert_at_head(node_ptr) },
+        }
+        if is_generic {
+            self.generic.store(node_ptr.as_ptr(), Ordering::Relaxed);
         }
 
         weak
     }
 
+    unsafe fn insert_at_head(&self, node_ptr: NonNull<Py<PyWeak>>) {
+        unsafe {
+            let mut ptrs = WeakLink::pointers(node_ptr);
+            let old_head = self.head.load(Ordering::Relaxed);
+            ptrs.as_mut().set_next(NonNull::new(old_head));
+            ptrs.as_mut().set_prev(None);
+            if let Some(old_head) = NonNull::new(old_head) {
+                WeakLink::pointers(old_head)
+                    .as_mut()
+                    .set_prev(Some(node_ptr));
+            }
+            self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
+        }
+    }
+
+    unsafe fn insert_after(&self, after: NonNull<Py<PyWeak>>, node_ptr: NonNull<Py<PyWeak>>) {
+        unsafe {
+            let mut ptrs = WeakLink::pointers(node_ptr);
+            let after_next = WeakLink::pointers(after).as_ref().get_next();
+            ptrs.as_mut().set_prev(Some(after));
+            ptrs.as_mut().set_next(after_next);
+            WeakLink::pointers(after).as_mut().set_next(Some(node_ptr));
+            if let Some(next) = after_next {
+                WeakLink::pointers(next).as_mut().set_prev(Some(node_ptr));
+            }
+        }
+    }
+
+    // get_basic_refs
+    fn find_generic_proxy_ptr(&self) -> *mut Py<PyWeak> {
+        let generic_ptr = self.generic.load(Ordering::Relaxed);
+        let candidate_ptr = if let Some(generic_node) = NonNull::new(generic_ptr) {
+            unsafe { WeakLink::pointers(generic_node).as_ref().get_next() }
+                .map_or(ptr::null_mut(), |n| n.as_ptr())
+        } else {
+            self.head.load(Ordering::Relaxed)
+        };
+        match NonNull::new(candidate_ptr) {
+            Some(candidate) => {
+                let node = unsafe { candidate.as_ref() };
+                let has_callback = unsafe { (&*node.0.payload.callback.get()).is_some() };
+                // PyWeakref_CheckProxy: the basic-proxy slot is reserved for
+                // the canonical proxy type; subclasses and callback-less ref
+                // subclasses must not be mistaken for it.
+                let is_proxy = node.class().is(crate::builtins::PyWeakProxy::static_type());
+                if has_callback || !is_proxy {
+                    ptr::null_mut()
+                } else {
+                    candidate_ptr
+                }
+            }
+            None => ptr::null_mut(),
+        }
+    }
+
     /// Clear all weakrefs and call their callbacks.
     /// Called when the owner object is being dropped.
     // PyObject_ClearWeakRefs
@@ -1445,7 +1494,7 @@ impl PyObject {
         typ: PyTypeRef,
     ) -> Option<PyRef<PyWeak>> {
         self.weak_ref_list()
-            .map(|wrl| wrl.add(self, typ, true, callback, None))
+            .map(|wrl| wrl.add(self, typ, true, false, callback, None))
     }
 
     pub(crate) fn downgrade_with_typ(
@@ -1476,13 +1525,14 @@ impl PyObject {
             None
         };
         let cls_is_weakref = typ.is(vm.ctx.types.weakref_type);
+        let cls_is_weakproxy = typ.is(vm.ctx.types.weakproxy_type);
         let wrl = self.weak_ref_list().ok_or_else(|| {
             vm.new_type_error(format!(
                 "cannot create weak reference to '{}' object",
                 self.class().name()
             ))
         })?;
-        Ok(wrl.add(self, typ, cls_is_weakref, callback, dict))
+        Ok(wrl.add(self, typ, cls_is_weakref, cls_is_weakproxy, callback, dict))
     }
 
     pub fn downgrade(