Skip to content

Commit db5de5e

Browse files
fregataaclaude
andauthored
Reuse stored hashes in dict.fromkeys() (#8503)
* Reuse stored hashes in dict.fromkeys() Closes #8490. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Make fromkeys_known_hashes a PyDict associated fn --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 3a98ef7 commit db5de5e

3 files changed

Lines changed: 43 additions & 5 deletions

File tree

Lib/test/test_set.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,6 @@ def test_cyclical_repr(self):
330330
name = repr(s).partition('(')[0] # strip class name
331331
self.assertEqual(repr(s), '%s({%s(...)})' % (name, name))
332332

333-
@unittest.expectedFailure # TODO: RUSTPYTHON
334333
def test_do_not_rehash_dict_keys(self):
335334
n = 10
336335
d = dict.fromkeys(map(HashCountingInt, range(n)))

crates/vm/src/builtins/dict.rs

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use super::{
22
IterStatus, PositionIterInternal, PyBaseExceptionRef, PyGenericAlias, PyMappingProxy, PySet,
3-
PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, set::PySetInner,
3+
PyStr, PyStrRef, PyTupleRef, PyType, PyTypeRef, set, set::PySetInner,
44
};
55
use crate::common::lock::LazyLock;
66
use crate::object::{Traverse, TraverseFn};
@@ -9,7 +9,7 @@ use crate::{
99
TryFromObject, atomic_func,
1010
builtins::{PyList, PyTuple, iter::builtins_iter, type_::PyAttributes},
1111
class::{PyClassDef, PyClassImpl},
12-
common::ascii,
12+
common::{ascii, hash::PyHash},
1313
dict_inner::{self, DictKey},
1414
function::{ArgIterable, FuncArgs, KwArgs, OptionalArg, PyArithmeticValue, PyComparisonValue},
1515
iter::PyExactSizeIterator,
@@ -354,6 +354,20 @@ impl PyDict {
354354
) -> PyResult<Option<PyObjectRef>> {
355355
self.entries.get(vm, key)
356356
}
357+
358+
/// Keys of `obj` with their stored hashes, or `None` if it must be iterated
359+
/// generically. Only exact dicts and sets qualify, as in CPython's
360+
/// `_PyDict_FromKeys`: a subclass may override `__iter__`.
361+
fn fromkeys_known_hashes(
362+
obj: &PyObject,
363+
vm: &VirtualMachine,
364+
) -> Option<Vec<(PyObjectRef, PyHash)>> {
365+
if let Some(dict) = obj.downcast_ref_if_exact::<Self>(vm) {
366+
Some(dict.entries.keys_with_hashes())
367+
} else {
368+
set::exact_set_keys_with_hashes(obj, vm)
369+
}
370+
}
357371
}
358372

359373
// Python dict methods:
@@ -384,8 +398,16 @@ impl PyDict {
384398
let d = PyType::call(&class, ().into(), vm)?;
385399
match d.downcast_exact::<Self>(vm) {
386400
Ok(pydict) => {
387-
for key in iterable.iter(vm)? {
388-
pydict.__setitem__(key?, value.clone(), vm)?;
401+
if let Some(keys) = Self::fromkeys_known_hashes(iterable.as_object(), vm) {
402+
for (key, hash) in keys {
403+
pydict
404+
.entries
405+
.insert_known_hash(vm, &*key, hash, value.clone())?;
406+
}
407+
} else {
408+
for key in iterable.iter(vm)? {
409+
pydict.__setitem__(key?, value.clone(), vm)?;
410+
}
389411
}
390412
Ok(pydict.into_pyref().into())
391413
}

crates/vm/src/builtins/set.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -605,6 +605,23 @@ fn extract_set(obj: &PyObject) -> Option<&PySetInner> {
605605
})
606606
}
607607

608+
/// Elements of `obj` with their stored hashes, or `None` unless `obj` is exactly
609+
/// a `set` or `frozenset` — `PyAnySet_CheckExact`, where [`extract_set`] is the
610+
/// subclass-inclusive `PyAnySet_Check`.
611+
pub(super) fn exact_set_keys_with_hashes(
612+
obj: &PyObject,
613+
vm: &VirtualMachine,
614+
) -> Option<Vec<(PyObjectRef, PyHash)>> {
615+
let inner = obj
616+
.downcast_ref_if_exact::<PySet>(vm)
617+
.map(|set| &set.inner)
618+
.or_else(|| {
619+
obj.downcast_ref_if_exact::<PyFrozenSet>(vm)
620+
.map(|frozen| &frozen.inner)
621+
})?;
622+
Some(inner.content.keys_with_hashes())
623+
}
624+
608625
fn reduce_set(zelf: &PyObject, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
609626
(
610627
zelf.class().to_owned(),

0 commit comments

Comments
 (0)