Skip to content

Commit ad8bfdb

Browse files
authored
Reuse stored hashes when merging dictionaries (#8858)
Assisted-by: Codex:GPT-6
1 parent d31ccae commit ad8bfdb

3 files changed

Lines changed: 91 additions & 7 deletions

File tree

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -278,11 +278,14 @@ impl PyDict {
278278
) -> PyResult<()> {
279279
let dict = &self.entries;
280280
let dict_size = &dict_other.size();
281-
for (key, value) in dict_other {
282-
if !override_existing && dict.contains(vm, &*key)? {
281+
let mut position = 0;
282+
while let Some((next, key, value, hash)) = dict_other.entries.next_entry_with_hash(position)
283+
{
284+
position = next;
285+
if !override_existing && dict.contains_known_hash(vm, &*key, hash)? {
283286
continue;
284287
}
285-
dict.insert(vm, &*key, value)?;
288+
dict.insert_known_hash(vm, &*key, hash, value)?;
286289
}
287290
if dict_other.entries.has_changed_size(dict_size) {
288291
return Err(vm.new_runtime_error("dict mutated during update"));

‎crates/vm/src/dict_inner.rs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -486,8 +486,9 @@ impl<T: Clone> Dict<T> {
486486
/// Store a key whose hash the caller already knows.
487487
///
488488
/// `hash` must be computed for `key` or read from [`Self::keys_with_hashes`]
489-
/// on a container holding this same key. A wrong hash lands the entry in a
490-
/// bucket no lookup probes, silently losing the key.
489+
/// or [`Self::next_entry_with_hash`] on a container holding this same key.
490+
/// A wrong hash lands the entry in a bucket no lookup probes, silently
491+
/// losing the key.
491492
pub(crate) fn insert_known_hash<K>(
492493
&self,
493494
vm: &VirtualMachine,
@@ -1037,13 +1038,21 @@ impl<T: Clone> Dict<T> {
10371038
}
10381039
}
10391040

1040-
pub(crate) fn next_entry(&self, mut position: EntryIndex) -> Option<(usize, PyObjectRef, T)> {
1041+
pub(crate) fn next_entry(&self, position: EntryIndex) -> Option<(usize, PyObjectRef, T)> {
1042+
self.next_entry_with_hash(position)
1043+
.map(|(position, key, value, _)| (position, key, value))
1044+
}
1045+
1046+
pub(crate) fn next_entry_with_hash(
1047+
&self,
1048+
mut position: EntryIndex,
1049+
) -> Option<(usize, PyObjectRef, T, HashValue)> {
10411050
let inner = self.read();
10421051
loop {
10431052
let entry = inner.entries.get(position)?;
10441053
position += 1;
10451054
if let Some(entry) = entry {
1046-
break Some((position, entry.key.clone(), entry.value.clone()));
1055+
break Some((position, entry.key.clone(), entry.value.clone(), entry.hash));
10471056
}
10481057
}
10491058
}

‎extra_tests/snippets/builtin_dict.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -615,3 +615,75 @@ def __eq__(self, other):
615615
lookup = LookupOrderSubclass("lookup", True)
616616
assert {stored: 1}[lookup] == 1
617617
assert LookupOrderKey.calls == ["lookup"]
618+
619+
620+
class MergeHashKey:
621+
hash_disabled = False
622+
623+
def __init__(self, value):
624+
self.value = value
625+
626+
def __hash__(self):
627+
assert not self.hash_disabled, "dictionary merge rehashed a stored key"
628+
return 42
629+
630+
def __eq__(self, other):
631+
if not isinstance(other, MergeHashKey):
632+
return NotImplemented
633+
return self.value == other.value
634+
635+
636+
def merge_with_update(source):
637+
result = {}
638+
result.update(source)
639+
return result
640+
641+
642+
def merge_with_ior(source):
643+
result = {}
644+
result |= source
645+
return result
646+
647+
648+
# Exact dictionary merges reuse hashes, including colliding keys and holes.
649+
merge_keys = [MergeHashKey(i) for i in range(3)]
650+
merge_source = dict(zip(merge_keys, ("first", "removed", "last")))
651+
del merge_source[merge_keys[1]]
652+
existing_merge_key = MergeHashKey(2)
653+
merge_target = {existing_merge_key: "old"}
654+
MergeHashKey.hash_disabled = True
655+
for merge in (
656+
dict,
657+
merge_with_update,
658+
merge_with_ior,
659+
lambda source: {} | source,
660+
lambda source: dict.__ror__(source, {}),
661+
lambda source: {**source},
662+
):
663+
assert list(merge(merge_source).items()) == list(merge_source.items())
664+
665+
# Overwriting a matching key keeps its identity and position.
666+
merge_target.update(merge_source)
667+
assert next(iter(merge_target)) is existing_merge_key
668+
assert list(merge_target.values()) == ["last", "first"]
669+
merge_source.update(merge_source)
670+
merge_source |= merge_source
671+
assert list(merge_source.values()) == ["first", "last"]
672+
MergeHashKey.hash_disabled = False
673+
674+
675+
class MergeMapping(dict):
676+
def __iter__(self):
677+
return iter(("virtual",))
678+
679+
def keys(self):
680+
return ["virtual"]
681+
682+
def __getitem__(self, key):
683+
assert key == "virtual"
684+
return 42
685+
686+
687+
# Generic mappings retain their lookup hooks instead of exposing dict storage.
688+
for merge in (dict, merge_with_update, merge_with_ior):
689+
assert merge(MergeMapping(stored=0)) == {"virtual": 42}

0 commit comments

Comments
 (0)