Implement weakref._remove_dead_weakref

Lib/test/test_copy.py

@@ -829,8 +829,6 @@ class TestCopy(unittest.TestCase):
     def test_copy_weakkeydict(self):
         self._check_copy_weakdict(weakref.WeakKeyDictionary)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_copy_weakvaluedict(self):
         self._check_copy_weakdict(weakref.WeakValueDictionary)
 
@@ -853,8 +851,6 @@ class TestCopy(unittest.TestCase):
         del c
         self.assertEqual(len(v), 1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_deepcopy_weakvaluedict(self):
         class C(object):
             def __init__(self, i):

Lib/test/test_weakref.py

@@ -1267,8 +1267,6 @@ class MappingTestCase(TestBase):
     def test_weak_valued_len_race(self):
         self.check_len_race(weakref.WeakValueDictionary, lambda k: (1, k))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_weak_values(self):
         #
         #  This exercises d.copy(), d.items(), d[], del d[], len(d).
@@ -1501,8 +1499,6 @@ class MappingTestCase(TestBase):
         dict, objects = self.make_weak_keyed_dict()
         self.check_weak_del_and_len_while_iterating(dict, testcontext)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_weak_values_destroy_while_iterating(self):
         # Issue #7105: iterators shouldn't crash when a key is implicitly removed
         dict, objects = self.make_weak_valued_dict()

vm/src/builtins/dict.rs

@@ -63,6 +63,12 @@ impl PyDict {
 #[allow(clippy::len_without_is_empty)]
 #[pyimpl(with(AsMapping, Hashable, Comparable, Iterable), flags(BASETYPE))]
 impl PyDict {
+    /// escape hatch to access the underlying data structure directly. prefer adding a method on
+    /// PyDict instead of using this
+    pub(crate) fn _as_dict_inner(&self) -> &DictContentType {
+        &self.entries
+    }
+
     #[pyslot]
     fn slot_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult {
         PyDict::default().into_pyresult_with_type(vm, cls)

vm/src/dictdatatype.rs

@@ -13,6 +13,7 @@ use crate::{
 };
 use std::fmt;
 use std::mem::size_of;
+use std::ops::ControlFlow;
 
 // HashIndex is intended to be same size with hash::PyHash
 // but it doesn't mean the values are compatible with actual pyhash value
@@ -223,6 +224,8 @@ impl<T> DictInner<T> {
     }
 }
 
+type PopInnerResult<T> = ControlFlow<Option<DictEntry<T>>>;
+
 impl<T: Clone> Dict<T> {
     fn read(&self) -> PyRwLockReadGuard<'_, DictInner<T>> {
         self.inner.read()
@@ -348,18 +351,23 @@ impl<T: Clone> Dict<T> {
     pub fn delete_if_exists<K>(&self, vm: &VirtualMachine, key: &K) -> PyResult<bool>
     where
         K: DictKey,
+    {
+        self.delete_if(vm, key, |_| Ok(true))
+    }
+
+    /// pred should be VERY CAREFUL about what it does as it is called while
+    /// the dict's internal mutex is held
+    pub(crate) fn delete_if<K, F>(&self, vm: &VirtualMachine, key: &K, pred: F) -> PyResult<bool>
+    where
+        K: DictKey,
+        F: Fn(&T) -> PyResult<bool>,
     {
         let hash = key.key_hash(vm)?;
         let deleted = loop {
             let lookup = self.lookup(vm, key, hash, None)?;
-            if let IndexEntry::Index(_) = lookup.0 {
-                if let Ok(Some(entry)) = self.pop_inner(lookup) {
-                    break Some(entry);
-                } else {
-                    // The dict was changed since we did lookup. Let's try again.
-                }
-            } else {
-                break None;
+            match self.pop_inner_if(lookup, &pred)? {
+                ControlFlow::Break(entry) => break entry,
+                ControlFlow::Continue(()) => continue,
             }
         };
         Ok(deleted.is_some())
@@ -371,10 +379,9 @@ impl<T: Clone> Dict<T> {
             let lookup = self.lookup(vm, key, hash, None)?;
             let (entry, index_index) = lookup;
             if let IndexEntry::Index(_) = entry {
-                if let Ok(Some(entry)) = self.pop_inner(lookup) {
-                    break Some(entry);
-                } else {
-                    // The dict was changed since we did lookup. Let's try again.
+                match self.pop_inner(lookup) {
+                    ControlFlow::Break(Some(entry)) => break Some(entry),
+                    _ => continue,
                 }
             } else {
                 let mut inner = self.write();
@@ -573,25 +580,42 @@ impl<T: Clone> Dict<T> {
     }
 
     // returns Err(()) if changed since lookup
-    fn pop_inner(&self, lookup: LookupResult) -> Result<Option<DictEntry<T>>, ()> {
+    fn pop_inner(&self, lookup: LookupResult) -> PopInnerResult<T> {
+        self.pop_inner_if(lookup, |_| Ok::<_, std::convert::Infallible>(true))
+            .unwrap_or_else(|x| match x {})
+    }
+
+    fn pop_inner_if<E>(
+        &self,
+        lookup: LookupResult,
+        pred: impl Fn(&T) -> Result<bool, E>,
+    ) -> Result<PopInnerResult<T>, E> {
         let (entry_index, index_index) = lookup;
         let entry_index = if let IndexEntry::Index(entry_index) = entry_index {
             entry_index
         } else {
-            return Ok(None);
+            return Ok(ControlFlow::Break(None));
         };
-        let mut inner = self.write();
-        if matches!(inner.entries.get(entry_index), Some(Some(entry)) if entry.index == index_index)
-        {
-            // all good
+        let inner = &mut *self.write();
+        let slot = if let Some(slot) = inner.entries.get_mut(entry_index) {
+            slot
         } else {
             // The dict was changed since we did lookup. Let's try again.
-            return Err(());
+            return Ok(ControlFlow::Continue(()));
         };
+        match slot {
+            Some(entry) if entry.index == index_index => {
+                if pred(&entry.value)? == false {
+                    return Ok(ControlFlow::Break(None));
+                }
+            }
+            // The dict was changed since we did lookup. Let's try again.
+            _ => return Ok(ControlFlow::Continue(())),
+        }
         inner.indices[index_index] = IndexEntry::DUMMY;
         inner.used -= 1;
-        let removed = std::mem::take(&mut inner.entries[entry_index]);
-        Ok(removed)
+        let removed = slot.take();
+        Ok(ControlFlow::Break(removed))
     }
 
     /// Retrieve and delete a key
@@ -599,10 +623,9 @@ impl<T: Clone> Dict<T> {
         let hash_value = key.key_hash(vm)?;
         let removed = loop {
             let lookup = self.lookup(vm, key, hash_value, None)?;
-            if let Ok(ret) = self.pop_inner(lookup) {
-                break ret.map(|e| e.value);
-            } else {
-                // changed since lookup, loop again
+            match self.pop_inner(lookup) {
+                ControlFlow::Break(entry) => break entry.map(|e| e.value),
+                ControlFlow::Continue(()) => continue,
             }
         };
         Ok(removed)

vm/src/pyobjectrc.rs

@@ -290,6 +290,11 @@ impl PyWeak {
         }
     }
 
+    pub(crate) fn is_dead(&self) -> bool {
+        let guard = unsafe { self.parent.as_ref().lock() };
+        guard.obj.is_none()
+    }
+
     #[inline(always)]
     fn drop_inner(&self) {
         let dealloc = {

vm/src/stdlib/weakref.rs

@@ -8,7 +8,8 @@ pub(crate) use _weakref::make_module;
 
 #[pymodule]
 mod _weakref {
-    use crate::{builtins::PyTypeRef, PyObjectRef, VirtualMachine};
+    use crate::builtins::{PyDictRef, PyTypeRef, PyWeak};
+    use crate::{PyObjectRef, PyResult, VirtualMachine};
 
     #[pyattr(name = "ref")]
     fn ref_(vm: &VirtualMachine) -> PyTypeRef {
@@ -45,7 +46,18 @@ mod _weakref {
     }
 
     #[pyfunction]
-    fn _remove_dead_weakref(_obj: PyObjectRef, _key: PyObjectRef) {
-        // TODO
+    fn _remove_dead_weakref(
+        dict: PyDictRef,
+        key: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
+        dict._as_dict_inner()
+            .delete_if(vm, &key, |wr| {
+                let wr = wr
+                    .payload::<PyWeak>()
+                    .ok_or_else(|| vm.new_type_error("not a weakref".to_owned()))?;
+                Ok(wr.is_dead())
+            })
+            .map(drop)
     }
 }