diff --git a/crates/vm/src/builtins/set.rs b/crates/vm/src/builtins/set.rs index cd724abc5c1..62bdb0f0da5 100644 --- a/crates/vm/src/builtins/set.rs +++ b/crates/vm/src/builtins/set.rs @@ -195,6 +195,26 @@ impl PySetInner { Ok(set) } + /// Build a set from an arbitrary object, reusing stored hashes when the + /// source is a set/frozenset/dict. + fn from_object(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let set = Self::default(); + set.update_internal(iterable, vm)?; + Ok(set) + } + + /// Elements of `obj` with their stored hashes, or `None` if `obj` keeps + /// none and must be iterated generically. Mirrors the `PyAnySet_Check` / + /// `PyDict_CheckExact` fast paths in CPython's `set_update_internal`. + fn cached_hashes(obj: &PyObject, vm: &VirtualMachine) -> Option> { + if let Some(set) = extract_set(obj) { + Some(set.content.keys_with_hashes()) + } else { + obj.downcast_ref_if_exact::(vm) + .map(|dict| dict._as_dict_inner().keys_with_hashes()) + } + } + fn fold_op( &self, others: impl core::iter::Iterator, @@ -228,6 +248,17 @@ impl PySetInner { Self::wrap_unhashable_error(result, needle, vm) } + /// [`Self::contains`] with a known hash. Such a needle came out of a + /// set/dict, so it is hashable and needs no frozenset retry. + fn contains_known_hash( + &self, + needle: &PyObject, + hash: PyHash, + vm: &VirtualMachine, + ) -> PyResult { + self.content.contains_known_hash(vm, needle, hash) + } + fn compare(&self, other: &Self, op: PyComparisonOp, vm: &VirtualMachine) -> PyResult { if op == PyComparisonOp::Ne { return self.compare(other, PyComparisonOp::Eq, vm).map(|eq| !eq); @@ -251,6 +282,12 @@ impl PySetInner { pub(super) fn union(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { let set = self.clone(); + if let Some(elements) = Self::cached_hashes(other.as_object(), vm) { + for (item, hash) in elements { + set.add_known_hash(item, hash, vm)?; + } + return Ok(set); + } for item in other.iter(vm)? { set.add(item?, vm)?; } @@ -260,6 +297,14 @@ impl PySetInner { pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { let set = Self::default(); + if let Some(elements) = Self::cached_hashes(other.as_object(), vm) { + for (obj, hash) in elements { + if self.contains_known_hash(&obj, hash, vm)? { + set.add_known_hash(obj, hash, vm)?; + } + } + return Ok(set); + } for item in other.iter(vm)? { let obj = item?; if self.contains(&obj, vm)? { @@ -271,6 +316,12 @@ impl PySetInner { pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { let set = self.copy(); + if let Some(elements) = Self::cached_hashes(other.as_object(), vm) { + for (item, hash) in elements { + set.content.delete_if_exists_known_hash(vm, &*item, hash)?; + } + return Ok(set); + } for item in other.iter(vm)? { set.content.delete_if_exists(vm, &*item?)?; } @@ -284,6 +335,16 @@ impl PySetInner { ) -> PyResult { let new_inner = self.clone(); + if let Some(elements) = Self::cached_hashes(other.as_object(), vm) { + // the source is already duplicate-free + for (item, hash) in elements { + new_inner + .content + .delete_or_insert_known_hash(vm, &item, hash, ())?; + } + return Ok(new_inner); + } + // We want to remove duplicates in other let other_set = Self::from_iter(other.iter(vm)?, vm)?; @@ -333,6 +394,12 @@ impl PySetInner { Self::wrap_unhashable_error(result, &item, vm) } + /// [`Self::add`] with a known hash. + fn add_known_hash(&self, item: PyObjectRef, hash: PyHash, vm: &VirtualMachine) -> PyResult<()> { + let result = self.content.insert_known_hash(vm, &*item, hash, ()); + Self::wrap_unhashable_error(result, &item, vm) + } + fn remove(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { let result = self.retry_op_with_frozenset(&item, vm, |item, vm| self.content.delete(vm, item)); @@ -393,15 +460,15 @@ impl PySetInner { } fn merge_set(&self, any_set: AnySet, vm: &VirtualMachine) -> PyResult<()> { - for item in any_set.as_inner().elements() { - self.add(item, vm)?; + for (item, hash) in any_set.as_inner().content.keys_with_hashes() { + self.add_known_hash(item, hash, vm)?; } Ok(()) } fn merge_dict(&self, dict: PyDictRef, vm: &VirtualMachine) -> PyResult<()> { - for (key, _value) in dict { - self.add(key, vm)?; + for (key, hash) in dict._as_dict_inner().keys_with_hashes() { + self.add_known_hash(key, hash, vm)?; } Ok(()) } @@ -413,8 +480,8 @@ impl PySetInner { ) -> PyResult<()> { let temp_inner = self.fold_op(others, Self::intersection, vm)?; self.clear(); - for obj in temp_inner.elements() { - self.add(obj, vm)?; + for (obj, hash) in temp_inner.content.keys_with_hashes() { + self.add_known_hash(obj, hash, vm)?; } Ok(()) } @@ -425,6 +492,12 @@ impl PySetInner { vm: &VirtualMachine, ) -> PyResult<()> { for iterable in others { + if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) { + for (item, hash) in elements { + self.content.delete_if_exists_known_hash(vm, &*item, hash)?; + } + continue; + } let items = iterable.iter(vm)?.collect::, _>>()?; for item in items { self.content.delete_if_exists(vm, &*item)?; @@ -439,6 +512,14 @@ impl PySetInner { vm: &VirtualMachine, ) -> PyResult<()> { for iterable in others { + if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) { + // the source is already duplicate-free + for (item, hash) in elements { + self.content + .delete_or_insert_known_hash(vm, &item, hash, ())?; + } + continue; + } // We want to remove duplicates in iterable let iterable_set = Self::from_iter(iterable.iter(vm)?, vm)?; for item in iterable_set.elements() { @@ -955,7 +1036,7 @@ impl Representable for PySet { } impl Constructor for PyFrozenSet { - type Args = Vec; + type Args = OptionalArg; fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { let is_exact_frozenset = cls.is(vm.ctx.types.frozenset_type); @@ -988,11 +1069,11 @@ impl Constructor for PyFrozenSet { return Ok(input.clone()); } - iterable.into_option() + iterable } else { match &args.args[..] { - [] => None, - [iterable] => Some(iterable.clone()), + [] => OptionalArg::Missing, + [iterable] => OptionalArg::Present(iterable.clone()), slice => { return Err(vm.new_type_error(format!( "frozenset expected at most 1 argument, got {}", @@ -1002,23 +1083,25 @@ impl Constructor for PyFrozenSet { } }; - let elements = if let Some(iterable) = iterable_opt { - iterable.try_to_value(vm)? - } else { - vec![] - }; + let payload = Self::py_new(&cls, iterable_opt, vm)?; // Return empty frozenset singleton - if is_exact_frozenset && elements.is_empty() { + if is_exact_frozenset && payload.inner.len() == 0 { return Ok(vm.ctx.empty_frozenset.clone().into()); } - let payload = Self::py_new(&cls, elements, vm)?; payload.into_ref_with_type(vm, cls).map(Into::into) } - fn py_new(_cls: &Py, elements: Self::Args, vm: &VirtualMachine) -> PyResult { - Self::from_iter(vm, elements) + fn py_new(_cls: &Py, iterable: Self::Args, vm: &VirtualMachine) -> PyResult { + let inner = match iterable { + OptionalArg::Present(iterable) => PySetInner::from_object(iterable, vm)?, + OptionalArg::Missing => PySetInner::default(), + }; + Ok(Self { + inner, + ..Default::default() + }) } } diff --git a/crates/vm/src/dict_inner.rs b/crates/vm/src/dict_inner.rs index 3e75e6f27a6..9dda6194a0c 100644 --- a/crates/vm/src/dict_inner.rs +++ b/crates/vm/src/dict_inner.rs @@ -463,6 +463,24 @@ impl Dict { K: DictKey + ?Sized, { let hash = key.key_hash(vm)?; + self.insert_known_hash(vm, key, hash, value) + } + + /// Store a key whose hash the caller already knows. + /// + /// `hash` must equal `key.key_hash(vm)`; a wrong one lands the entry in a + /// bucket no lookup probes, silently losing the key. Only pass a hash from + /// [`Self::keys_with_hashes`] on a container holding this same key. + pub(crate) fn insert_known_hash( + &self, + vm: &VirtualMachine, + key: &K, + hash: HashValue, + value: T, + ) -> PyResult<()> + where + K: DictKey + ?Sized, + { let _removed = loop { let (entry_index, index_index) = self.lookup(vm, key, hash, None)?; let mut inner = self.write(); @@ -512,7 +530,18 @@ impl Dict { key: &K, ) -> PyResult { let key_hash = key.key_hash(vm)?; - let (entry, _) = self.lookup(vm, key, key_hash, None)?; + self.contains_known_hash(vm, key, key_hash) + } + + /// [`Self::contains`] with a known hash. Same contract as + /// [`Self::insert_known_hash`]. + pub(crate) fn contains_known_hash( + &self, + vm: &VirtualMachine, + key: &K, + hash: HashValue, + ) -> PyResult { + let (entry, _) = self.lookup(vm, key, hash, None)?; Ok(entry.index().is_some()) } @@ -697,6 +726,21 @@ impl Dict { self.remove_if_exists(vm, key).map(|opt| opt.is_some()) } + /// [`Self::delete_if_exists`] with a known hash. Same contract as + /// [`Self::insert_known_hash`]. + pub(crate) fn delete_if_exists_known_hash( + &self, + vm: &VirtualMachine, + key: &K, + hash: HashValue, + ) -> PyResult + where + K: DictKey + ?Sized, + { + self.remove_if_known_hash(vm, key, hash, |_| Ok(true)) + .map(|opt| opt.is_some()) + } + pub(crate) fn delete_if(&self, vm: &VirtualMachine, key: &K, pred: F) -> PyResult where K: DictKey + ?Sized, @@ -725,6 +769,22 @@ impl Dict { F: Fn(&T) -> PyResult, { let hash = key.key_hash(vm)?; + self.remove_if_known_hash(vm, key, hash, pred) + } + + /// [`Self::remove_if`] with a known hash. Same contract as + /// [`Self::insert_known_hash`]. + fn remove_if_known_hash( + &self, + vm: &VirtualMachine, + key: &K, + hash: HashValue, + pred: F, + ) -> PyResult> + where + K: DictKey + ?Sized, + F: Fn(&T) -> PyResult, + { let removed = loop { let lookup = self.lookup(vm, key, hash, None)?; match self.pop_inner_if(lookup, &pred)? { @@ -742,6 +802,18 @@ impl Dict { value: T, ) -> PyResult<()> { let hash = key.key_hash(vm)?; + self.delete_or_insert_known_hash(vm, key, hash, value) + } + + /// [`Self::delete_or_insert`] with a known hash. Same contract as + /// [`Self::insert_known_hash`]. + pub(crate) fn delete_or_insert_known_hash( + &self, + vm: &VirtualMachine, + key: &PyObject, + hash: HashValue, + value: T, + ) -> PyResult<()> { let _removed = loop { let lookup = self.lookup(vm, key, hash, None)?; let (entry, index_index) = lookup; @@ -888,6 +960,16 @@ impl Dict { .collect() } + /// All keys paired with the hash stored in their entry, for feeding + /// [`Self::insert_known_hash`] without re-calling `__hash__`. + pub(crate) fn keys_with_hashes(&self) -> Vec<(PyObjectRef, HashValue)> { + self.read() + .entries + .iter() + .filter_map(|v| v.as_ref().map(|v| (v.key.clone(), v.hash))) + .collect() + } + pub(crate) fn values(&self) -> Vec { self.read() .entries diff --git a/crates/vm/src/function/protocol.rs b/crates/vm/src/function/protocol.rs index 25ef62b458d..d503fabaca8 100644 --- a/crates/vm/src/function/protocol.rs +++ b/crates/vm/src/function/protocol.rs @@ -86,6 +86,11 @@ unsafe impl Traverse for ArgIterable { } impl ArgIterable { + #[must_use] + pub(crate) fn as_object(&self) -> &PyObject { + &self.iterable + } + /// Returns an iterator over this sequence of objects. /// /// This operation may fail if an exception is raised while invoking the