From 2a87b893fbdfa648e401a9ae0704461ed4e41fbe Mon Sep 17 00:00:00 2001 From: Noa <33094578+coolreader18@users.noreply.github.com> Date: Tue, 9 Nov 2021 15:41:03 -0600 Subject: [PATCH] Reuse standard no-cb weakref --- Lib/test/test_weakref.py | 2 - vm/src/builtins/pytype.rs | 2 +- vm/src/builtins/weakproxy.rs | 16 ++++++- vm/src/builtins/weakref.rs | 2 +- vm/src/pyobjectrc.rs | 86 ++++++++++++++++++++++++------------ 5 files changed, 75 insertions(+), 33 deletions(-) diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py index 67a17d341..3b783c456 100644 --- a/Lib/test/test_weakref.py +++ b/Lib/test/test_weakref.py @@ -198,8 +198,6 @@ class ReferencesTestCase(TestBase): self.assertIsNone(ref(), "ref2 should be dead after deleting object reference") - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_ref_reuse(self): o = C() ref1 = weakref.ref(o) diff --git a/vm/src/builtins/pytype.rs b/vm/src/builtins/pytype.rs index b6e249595..3ed6f77bc 100644 --- a/vm/src/builtins/pytype.rs +++ b/vm/src/builtins/pytype.rs @@ -121,7 +121,7 @@ impl PyType { base.subclasses.write().push( new_type .as_object() - .downgrade_with_typ_opt(None, weakref_type.clone()) + .downgrade_with_weakref_typ_opt(None, weakref_type.clone()) .unwrap(), ); } diff --git a/vm/src/builtins/weakproxy.rs b/vm/src/builtins/weakproxy.rs index f658a3893..8cf4ab404 100644 --- a/vm/src/builtins/weakproxy.rs +++ b/vm/src/builtins/weakproxy.rs @@ -33,14 +33,28 @@ impl Constructor for PyWeakProxy { Self::Args { referent, callback }: Self::Args, vm: &VirtualMachine, ) -> PyResult { + // 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, + super::PyWeak::make_slots(), + ) + }); // TODO: PyWeakProxy should use the same payload as PyWeak PyWeakProxy { - weak: referent.downgrade(callback.into_option(), vm)?, + weak: referent.downgrade_with_typ(callback.into_option(), weak_cls.clone(), vm)?, } .into_pyresult_with_type(vm, cls) } } +crate::common::static_cell! { + static WEAK_SUBCLASS: PyTypeRef; +} + #[pyimpl(with(SetAttr, Constructor))] impl PyWeakProxy { // TODO: callbacks diff --git a/vm/src/builtins/weakref.rs b/vm/src/builtins/weakref.rs index df600ce64..841681444 100644 --- a/vm/src/builtins/weakref.rs +++ b/vm/src/builtins/weakref.rs @@ -105,7 +105,7 @@ impl Comparable for PyWeak { let both = zelf.upgrade().and_then(|s| other.upgrade().map(|o| (s, o))); let eq = match both { Some((a, b)) => vm.bool_eq(&a, &b)?, - None => false, + None => zelf.is(other), }; Ok(eq.into()) }) diff --git a/vm/src/pyobjectrc.rs b/vm/src/pyobjectrc.rs index 750651c75..5e844fad6 100644 --- a/vm/src/pyobjectrc.rs +++ b/vm/src/pyobjectrc.rs @@ -114,21 +114,12 @@ impl fmt::Debug for WeakRefList { struct WeakListInner { list: LinkedList>, + generic_weakref: Option>>, obj: Option>, // one for each live PyWeak with a reference to this, + 1 for the referent object if it's not dead ref_count: usize, } -impl Default for WeakListInner { - fn default() -> Self { - Self { - list: LinkedList::default(), - obj: None, - ref_count: 1, - } - } -} - cfg_if::cfg_if! { if #[cfg(feature = "threading")] { unsafe impl Send for WeakListInner {} @@ -152,13 +143,29 @@ impl WeakRefList { &self, obj: &PyObject, cls: PyTypeRef, + cls_is_weakref: bool, callback: Option, dict: Option, ) -> PyObjectWeak { - let inner_ptr = self.inner.get_or_init(Default::default); + let is_generic = cls_is_weakref && callback.is_none(); + let inner_ptr = self.inner.get_or_init(|| { + Box::new(PyMutex::new(WeakListInner { + list: LinkedList::default(), + generic_weakref: None, + obj: Some(NonNull::from(obj)), + ref_count: 1, + })) + }); let mut inner = unsafe { inner_ptr.as_ref().lock() }; - if inner.obj.is_none() { - inner.obj = Some(NonNull::from(obj)); + if is_generic { + if let Some(generic_weakref) = inner.generic_weakref { + let generic_weakref = unsafe { generic_weakref.as_ref() }; + if generic_weakref.0.refcount.get() != 0 { + return PyObjectWeak { + weak: generic_weakref.to_owned(), + }; + } + } } let obj = PyWeak { pointers: Pointers::new(), @@ -171,6 +178,9 @@ impl WeakRefList { // one out of the list we immediately wrap it in ManuallyDrop or forget it inner.list.push_front(unsafe { ptr::read(&weak) }); inner.ref_count += 1; + if is_generic { + inner.generic_weakref = Some(NonNull::from(&*weak)); + } PyObjectWeak { weak } } @@ -185,7 +195,8 @@ impl WeakRefList { // TODO: can be an arrayvec let mut v = Vec::with_capacity(16); loop { - let iter = inner + let inner2 = &mut *inner; + let iter = inner2 .list .drain_filter(|_| true) .filter_map(|wr| { @@ -195,6 +206,10 @@ impl WeakRefList { // this but before we reach the loop body below let wr = ManuallyDrop::new(wr); + if Some(NonNull::from(&**wr)) == inner2.generic_weakref { + inner2.generic_weakref = None + } + // if strong_count == 0 there's some reentrancy going on. we don't // want to call the callback (wr.as_object().strong_count() > 0).then(|| (*wr).clone()) @@ -235,6 +250,24 @@ impl WeakRefList { unsafe fn dealloc(ptr: NonNull>) { Box::from_raw(ptr.as_ptr()); } + + fn get_weak_references(&self) -> Vec { + let inner = match self.try_lock() { + Some(inner) => inner, + None => return vec![], + }; + let mut v = Vec::with_capacity(inner.ref_count - 1); + v.extend(inner.iter().map(|wr| PyObjectWeak { + weak: wr.to_owned(), + })); + v + } +} + +impl WeakListInner { + fn iter<'a>(&'a self) -> impl Iterator> { + self.list.iter().filter(|wr| wr.0.refcount.get() > 0) + } } impl Default for WeakRefList { @@ -306,6 +339,9 @@ impl PyWeak { // right now so that should be obvious!! std::mem::forget(unsafe { guard.list.remove(node_ptr) }); guard.ref_count -= 1; + if Some(node_ptr) == guard.generic_weakref { + guard.generic_weakref = None; + } guard.ref_count == 0 }; if dealloc { @@ -537,13 +573,14 @@ impl PyObject { Some(&self.0.weaklist) } - pub(crate) fn downgrade_with_typ_opt( + pub(crate) fn downgrade_with_weakref_typ_opt( &self, callback: Option, + // a reference to weakref_type **specifically** typ: PyTypeRef, ) -> Option { self.weakreflist() - .map(|wr| wr.add(self, typ, callback, None)) + .map(|wr| wr.add(self, typ, true, callback, None)) } pub(crate) fn downgrade_with_typ( @@ -561,8 +598,9 @@ impl PyObject { } else { None }; + let cls_is_weakref = typ.is(&vm.ctx.types.weakref_type); self.weakreflist() - .map(|wr| wr.add(self, typ, callback, dict)) + .map(|wr| wr.add(self, typ, cls_is_weakref, callback, dict)) .ok_or_else(|| { vm.new_type_error(format!( "cannot create weak reference to '{}' object", @@ -580,15 +618,7 @@ impl PyObject { } pub fn get_weak_references(&self) -> Option> { - self.weakreflist().map(|wrl| { - wrl.try_lock() - .iter() - .flat_map(|inner| inner.list.iter()) - .map(|wr| PyObjectWeak { - weak: wr.to_owned(), - }) - .collect() - }) + self.weakreflist().map(|wrl| wrl.get_weak_references()) } pub fn payload_is(&self) -> bool { @@ -1099,14 +1129,14 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { object_type.subclasses.write().push( type_type .as_object() - .downgrade_with_typ_opt(None, weakref_type.clone()) + .downgrade_with_weakref_typ_opt(None, weakref_type.clone()) .unwrap(), ); object_type.subclasses.write().push( weakref_type .as_object() - .downgrade_with_typ_opt(None, weakref_type.clone()) + .downgrade_with_weakref_typ_opt(None, weakref_type.clone()) .unwrap(), );