From f9d9a1e15e5304cc3d3d2e007ccf73bee2e77afd Mon Sep 17 00:00:00 2001 From: Noa <33094578+coolreader18@users.noreply.github.com> Date: Fri, 5 Nov 2021 19:38:59 -0500 Subject: [PATCH] Implement weakref._remove_dead_weakref --- Lib/test/test_copy.py | 4 --- Lib/test/test_weakref.py | 4 --- vm/src/builtins/dict.rs | 6 ++++ vm/src/dictdatatype.rs | 73 ++++++++++++++++++++++++++-------------- vm/src/pyobjectrc.rs | 5 +++ vm/src/stdlib/weakref.rs | 18 ++++++++-- 6 files changed, 74 insertions(+), 36 deletions(-) diff --git a/Lib/test/test_copy.py b/Lib/test/test_copy.py index 5c28a1eeb..c0988f953 100644 --- a/Lib/test/test_copy.py +++ b/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): diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py index f91fc9a1d..67a17d341 100644 --- a/Lib/test/test_weakref.py +++ b/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() diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index 0ee51973d..ba5ad5319 100644 --- a/vm/src/builtins/dict.rs +++ b/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) diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs index dbbcfddde..eed94e622 100644 --- a/vm/src/dictdatatype.rs +++ b/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 DictInner { } } +type PopInnerResult = ControlFlow>>; + impl Dict { fn read(&self) -> PyRwLockReadGuard<'_, DictInner> { self.inner.read() @@ -348,18 +351,23 @@ impl Dict { pub fn delete_if_exists(&self, vm: &VirtualMachine, key: &K) -> PyResult 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(&self, vm: &VirtualMachine, key: &K, pred: F) -> PyResult + where + K: DictKey, + F: Fn(&T) -> PyResult, { 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 Dict { 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 Dict { } // returns Err(()) if changed since lookup - fn pop_inner(&self, lookup: LookupResult) -> Result>, ()> { + fn pop_inner(&self, lookup: LookupResult) -> PopInnerResult { + self.pop_inner_if(lookup, |_| Ok::<_, std::convert::Infallible>(true)) + .unwrap_or_else(|x| match x {}) + } + + fn pop_inner_if( + &self, + lookup: LookupResult, + pred: impl Fn(&T) -> Result, + ) -> Result, 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 Dict { 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) diff --git a/vm/src/pyobjectrc.rs b/vm/src/pyobjectrc.rs index 96cfc243f..750651c75 100644 --- a/vm/src/pyobjectrc.rs +++ b/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 = { diff --git a/vm/src/stdlib/weakref.rs b/vm/src/stdlib/weakref.rs index e4a1e9cac..ffeb03628 100644 --- a/vm/src/stdlib/weakref.rs +++ b/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::() + .ok_or_else(|| vm.new_type_error("not a weakref".to_owned()))?; + Ok(wr.is_dead()) + }) + .map(drop) } }