Skip to content

Commit ba69646

Browse files
fregataaclaude
fregataa
and
claude
committed
Reuse stored hashes when building a set from a set/frozenset/dict
set and frozenset recomputed __hash__ for every element even when the source object already stored a hash per entry. CPython's set_update_internal branches on PyAnySet_Check / PyDict_CheckExact and feeds set_add_entry the hash read from the source table; RustPython always iterated generically. Split the hash computation out of the Dict entry points so callers can supply a hash they already hold, add keys_with_hashes() to hand out (key, hash) pairs, and take the fast path in the set constructors and in the set operations whose argument is a set/frozenset/exact dict. ArgIterable::as_object() exposes the pre-__iter__ object so the set operations can dispatch on the source type without changing any of their signatures. Closes #8489. dict.fromkeys() is the dict-target counterpart and is tracked in #8490, so test_do_not_rehash_dict_keys keeps its expectedFailure marker until that lands too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 365434b commit ba69646

3 files changed

Lines changed: 198 additions & 12 deletions

File tree

‎crates/vm/src/builtins/set.rs‎

Lines changed: 101 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,29 @@ impl PySetInner {
195195
Ok(set)
196196
}
197197

198+
/// Build a set from an arbitrary object, reusing the source's stored
199+
/// hashes when it is a set/frozenset/dict.
200+
pub(super) fn from_object(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
201+
let set = Self::default();
202+
set.update_internal(iterable, vm)?;
203+
Ok(set)
204+
}
205+
206+
/// Elements of `obj` paired with the hash already stored alongside them,
207+
/// or `None` if `obj` keeps no such hashes and has to be iterated
208+
/// generically (which calls `__hash__` on every element again).
209+
///
210+
/// Mirrors the `PyAnySet_Check` / `PyDict_CheckExact` fast paths in
211+
/// CPython's `set_update_internal`.
212+
fn cached_hashes(obj: &PyObject, vm: &VirtualMachine) -> Option<Vec<(PyObjectRef, PyHash)>> {
213+
if let Some(set) = extract_set(obj) {
214+
Some(set.content.keys_with_hashes())
215+
} else {
216+
obj.downcast_ref_if_exact::<PyDict>(vm)
217+
.map(|dict| dict._as_dict_inner().keys_with_hashes())
218+
}
219+
}
220+
198221
fn fold_op<O>(
199222
&self,
200223
others: impl core::iter::Iterator<Item = O>,
@@ -228,6 +251,18 @@ impl PySetInner {
228251
Self::wrap_unhashable_error(result, needle, vm)
229252
}
230253

254+
/// [`Self::contains`] for a needle whose hash was read from another
255+
/// set/dict. Such a needle is hashable by construction, so there is no
256+
/// set-to-frozenset retry to do.
257+
fn contains_with_hash(
258+
&self,
259+
needle: &PyObject,
260+
hash: PyHash,
261+
vm: &VirtualMachine,
262+
) -> PyResult<bool> {
263+
self.content.contains_with_hash(vm, needle, hash)
264+
}
265+
231266
fn compare(&self, other: &Self, op: PyComparisonOp, vm: &VirtualMachine) -> PyResult<bool> {
232267
if op == PyComparisonOp::Ne {
233268
return self.compare(other, PyComparisonOp::Eq, vm).map(|eq| !eq);
@@ -251,6 +286,12 @@ impl PySetInner {
251286

252287
pub(super) fn union(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
253288
let set = self.clone();
289+
if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
290+
for (item, hash) in elements {
291+
set.add_with_hash(item, hash, vm)?;
292+
}
293+
return Ok(set);
294+
}
254295
for item in other.iter(vm)? {
255296
set.add(item?, vm)?;
256297
}
@@ -260,6 +301,14 @@ impl PySetInner {
260301

261302
pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
262303
let set = Self::default();
304+
if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
305+
for (obj, hash) in elements {
306+
if self.contains_with_hash(&obj, hash, vm)? {
307+
set.add_with_hash(obj, hash, vm)?;
308+
}
309+
}
310+
return Ok(set);
311+
}
263312
for item in other.iter(vm)? {
264313
let obj = item?;
265314
if self.contains(&obj, vm)? {
@@ -271,6 +320,12 @@ impl PySetInner {
271320

272321
pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
273322
let set = self.copy();
323+
if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
324+
for (item, hash) in elements {
325+
set.content.delete_if_exists_with_hash(vm, &*item, hash)?;
326+
}
327+
return Ok(set);
328+
}
274329
for item in other.iter(vm)? {
275330
set.content.delete_if_exists(vm, &*item?)?;
276331
}
@@ -284,6 +339,16 @@ impl PySetInner {
284339
) -> PyResult<Self> {
285340
let new_inner = self.clone();
286341

342+
if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
343+
// the source is already duplicate-free
344+
for (item, hash) in elements {
345+
new_inner
346+
.content
347+
.delete_or_insert_with_hash(vm, &item, hash, ())?;
348+
}
349+
return Ok(new_inner);
350+
}
351+
287352
// We want to remove duplicates in other
288353
let other_set = Self::from_iter(other.iter(vm)?, vm)?;
289354

@@ -333,6 +398,12 @@ impl PySetInner {
333398
Self::wrap_unhashable_error(result, &item, vm)
334399
}
335400

401+
/// [`Self::add`] for an item whose hash was read from another set/dict.
402+
fn add_with_hash(&self, item: PyObjectRef, hash: PyHash, vm: &VirtualMachine) -> PyResult<()> {
403+
let result = self.content.insert_with_hash(vm, &*item, hash, ());
404+
Self::wrap_unhashable_error(result, &item, vm)
405+
}
406+
336407
fn remove(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
337408
let result =
338409
self.retry_op_with_frozenset(&item, vm, |item, vm| self.content.delete(vm, item));
@@ -393,15 +464,15 @@ impl PySetInner {
393464
}
394465

395466
fn merge_set(&self, any_set: AnySet, vm: &VirtualMachine) -> PyResult<()> {
396-
for item in any_set.as_inner().elements() {
397-
self.add(item, vm)?;
467+
for (item, hash) in any_set.as_inner().content.keys_with_hashes() {
468+
self.add_with_hash(item, hash, vm)?;
398469
}
399470
Ok(())
400471
}
401472

402473
fn merge_dict(&self, dict: PyDictRef, vm: &VirtualMachine) -> PyResult<()> {
403-
for (key, _value) in dict {
404-
self.add(key, vm)?;
474+
for (key, hash) in dict._as_dict_inner().keys_with_hashes() {
475+
self.add_with_hash(key, hash, vm)?;
405476
}
406477
Ok(())
407478
}
@@ -413,8 +484,8 @@ impl PySetInner {
413484
) -> PyResult<()> {
414485
let temp_inner = self.fold_op(others, Self::intersection, vm)?;
415486
self.clear();
416-
for obj in temp_inner.elements() {
417-
self.add(obj, vm)?;
487+
for (obj, hash) in temp_inner.content.keys_with_hashes() {
488+
self.add_with_hash(obj, hash, vm)?;
418489
}
419490
Ok(())
420491
}
@@ -425,6 +496,12 @@ impl PySetInner {
425496
vm: &VirtualMachine,
426497
) -> PyResult<()> {
427498
for iterable in others {
499+
if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) {
500+
for (item, hash) in elements {
501+
self.content.delete_if_exists_with_hash(vm, &*item, hash)?;
502+
}
503+
continue;
504+
}
428505
let items = iterable.iter(vm)?.collect::<Result<Vec<_>, _>>()?;
429506
for item in items {
430507
self.content.delete_if_exists(vm, &*item)?;
@@ -439,6 +516,14 @@ impl PySetInner {
439516
vm: &VirtualMachine,
440517
) -> PyResult<()> {
441518
for iterable in others {
519+
if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) {
520+
// the source is already duplicate-free
521+
for (item, hash) in elements {
522+
self.content
523+
.delete_or_insert_with_hash(vm, &item, hash, ())?;
524+
}
525+
continue;
526+
}
442527
// We want to remove duplicates in iterable
443528
let iterable_set = Self::from_iter(iterable.iter(vm)?, vm)?;
444529
for item in iterable_set.elements() {
@@ -1002,18 +1087,23 @@ impl Constructor for PyFrozenSet {
10021087
}
10031088
};
10041089

1005-
let elements = if let Some(iterable) = iterable_opt {
1006-
iterable.try_to_value(vm)?
1090+
// built from the object itself, so a set/frozenset/dict source can
1091+
// hand over its stored hashes instead of being re-hashed
1092+
let inner = if let Some(iterable) = iterable_opt {
1093+
PySetInner::from_object(iterable, vm)?
10071094
} else {
1008-
vec![]
1095+
PySetInner::default()
10091096
};
10101097

10111098
// Return empty frozenset singleton
1012-
if is_exact_frozenset && elements.is_empty() {
1099+
if is_exact_frozenset && inner.len() == 0 {
10131100
return Ok(vm.ctx.empty_frozenset.clone().into());
10141101
}
10151102

1016-
let payload = Self::py_new(&cls, elements, vm)?;
1103+
let payload = Self {
1104+
inner,
1105+
..Default::default()
1106+
};
10171107
payload.into_ref_with_type(vm, cls).map(Into::into)
10181108
}
10191109

‎crates/vm/src/dict_inner.rs‎

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -463,6 +463,26 @@ impl<T: Clone> Dict<T> {
463463
K: DictKey + ?Sized,
464464
{
465465
let hash = key.key_hash(vm)?;
466+
self.insert_with_hash(vm, key, hash, value)
467+
}
468+
469+
/// Store a key whose hash the caller already knows.
470+
///
471+
/// `hash` MUST be the value `key.key_hash(vm)` would return. Passing a
472+
/// different hash corrupts the table: the entry is stored in a bucket no
473+
/// lookup will probe, so the key silently goes missing. Only pass a hash
474+
/// read out of another dict/set entry holding this very key object, e.g.
475+
/// via [`Self::keys_with_hashes`].
476+
pub(crate) fn insert_with_hash<K>(
477+
&self,
478+
vm: &VirtualMachine,
479+
key: &K,
480+
hash: HashValue,
481+
value: T,
482+
) -> PyResult<()>
483+
where
484+
K: DictKey + ?Sized,
485+
{
466486
let _removed = loop {
467487
let (entry_index, index_index) = self.lookup(vm, key, hash, None)?;
468488
let mut inner = self.write();
@@ -512,7 +532,20 @@ impl<T: Clone> Dict<T> {
512532
key: &K,
513533
) -> PyResult<bool> {
514534
let key_hash = key.key_hash(vm)?;
515-
let (entry, _) = self.lookup(vm, key, key_hash, None)?;
535+
self.contains_with_hash(vm, key, key_hash)
536+
}
537+
538+
/// [`Self::contains`] with a caller-supplied hash.
539+
///
540+
/// Same contract as [`Self::insert_with_hash`]: a wrong `hash` makes the
541+
/// lookup probe the wrong bucket and report a present key as missing.
542+
pub(crate) fn contains_with_hash<K: DictKey + ?Sized>(
543+
&self,
544+
vm: &VirtualMachine,
545+
key: &K,
546+
hash: HashValue,
547+
) -> PyResult<bool> {
548+
let (entry, _) = self.lookup(vm, key, hash, None)?;
516549
Ok(entry.index().is_some())
517550
}
518551

@@ -697,6 +730,22 @@ impl<T: Clone> Dict<T> {
697730
self.remove_if_exists(vm, key).map(|opt| opt.is_some())
698731
}
699732

733+
/// [`Self::delete_if_exists`] with a caller-supplied hash.
734+
///
735+
/// Same contract as [`Self::insert_with_hash`].
736+
pub(crate) fn delete_if_exists_with_hash<K>(
737+
&self,
738+
vm: &VirtualMachine,
739+
key: &K,
740+
hash: HashValue,
741+
) -> PyResult<bool>
742+
where
743+
K: DictKey + ?Sized,
744+
{
745+
self.remove_if_with_hash(vm, key, hash, |_| Ok(true))
746+
.map(|opt| opt.is_some())
747+
}
748+
700749
pub(crate) fn delete_if<K, F>(&self, vm: &VirtualMachine, key: &K, pred: F) -> PyResult<bool>
701750
where
702751
K: DictKey + ?Sized,
@@ -725,6 +774,23 @@ impl<T: Clone> Dict<T> {
725774
F: Fn(&T) -> PyResult<bool>,
726775
{
727776
let hash = key.key_hash(vm)?;
777+
self.remove_if_with_hash(vm, key, hash, pred)
778+
}
779+
780+
/// [`Self::remove_if`] with a caller-supplied hash.
781+
///
782+
/// Same contract as [`Self::insert_with_hash`].
783+
pub(crate) fn remove_if_with_hash<K, F>(
784+
&self,
785+
vm: &VirtualMachine,
786+
key: &K,
787+
hash: HashValue,
788+
pred: F,
789+
) -> PyResult<Option<T>>
790+
where
791+
K: DictKey + ?Sized,
792+
F: Fn(&T) -> PyResult<bool>,
793+
{
728794
let removed = loop {
729795
let lookup = self.lookup(vm, key, hash, None)?;
730796
match self.pop_inner_if(lookup, &pred)? {
@@ -742,6 +808,19 @@ impl<T: Clone> Dict<T> {
742808
value: T,
743809
) -> PyResult<()> {
744810
let hash = key.key_hash(vm)?;
811+
self.delete_or_insert_with_hash(vm, key, hash, value)
812+
}
813+
814+
/// [`Self::delete_or_insert`] with a caller-supplied hash.
815+
///
816+
/// Same contract as [`Self::insert_with_hash`].
817+
pub(crate) fn delete_or_insert_with_hash(
818+
&self,
819+
vm: &VirtualMachine,
820+
key: &PyObject,
821+
hash: HashValue,
822+
value: T,
823+
) -> PyResult<()> {
745824
let _removed = loop {
746825
let lookup = self.lookup(vm, key, hash, None)?;
747826
let (entry, index_index) = lookup;
@@ -888,6 +967,18 @@ impl<T: Clone> Dict<T> {
888967
.collect()
889968
}
890969

970+
/// All keys paired with the hash already stored in their entry.
971+
///
972+
/// Lets a caller move keys into another dict/set without calling
973+
/// `__hash__` again; see [`Self::insert_with_hash`].
974+
pub(crate) fn keys_with_hashes(&self) -> Vec<(PyObjectRef, HashValue)> {
975+
self.read()
976+
.entries
977+
.iter()
978+
.filter_map(|v| v.as_ref().map(|v| (v.key.clone(), v.hash)))
979+
.collect()
980+
}
981+
891982
pub(crate) fn values(&self) -> Vec<T> {
892983
self.read()
893984
.entries

‎crates/vm/src/function/protocol.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,11 @@ unsafe impl<T: Traverse> Traverse for ArgIterable<T> {
8686
}
8787

8888
impl<T> ArgIterable<T> {
89+
#[must_use]
90+
pub fn as_object(&self) -> &PyObject {
91+
&self.iterable
92+
}
93+
8994
/// Returns an iterator over this sequence of objects.
9095
///
9196
/// This operation may fail if an exception is raised while invoking the

0 commit comments

Comments
 (0)