Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions crates/vm/src/builtins/dict.rs
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,9 @@ impl PyDict {
match d.downcast_exact::<Self>(vm) {
Ok(pydict) => {
if let Some(keys) = Self::fromkeys_known_hashes(iterable.as_object(), vm) {
if class.is(vm.ctx.types.dict_type) {
pydict.entries.reserve_for_empty(keys.len());
}
for (key, hash) in keys {
pydict
.entries
Expand Down
68 changes: 68 additions & 0 deletions crates/vm/src/dict_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,35 @@ impl<T: Clone> Dict<T> {
self.inner.write()
}

/// Reserve space for known keys without changing an existing key layout.
/// Only use on an unpublished dictionary: resizing must not race with a
/// lookup that has already probed the old index table.
pub(crate) fn reserve_for_empty(&self, capacity: usize) {
if capacity == 0 {
return;
}
let Some(index_capacity) = capacity
.checked_add(capacity.div_ceil(2))
.and_then(usize::checked_next_power_of_two)
else {
return;
};
let mut inner = self.write();
if inner.filled != 0 || !inner.entries.is_empty() {
return;
}
// Keep the usual incremental growth path if preallocation fails.
if inner.entries.try_reserve_exact(capacity).is_err() {
return;
}
let additional = index_capacity.saturating_sub(inner.indices.len());
if additional != 0 && inner.indices.try_reserve_exact(additional).is_ok() {
// Every existing index is FREE, so extending needs no rehash and
// leaves the empty key layout and its version unchanged.
inner.indices.resize(index_capacity, IndexEntry::FREE);
}
}

/// Store a key
pub(crate) fn insert<K>(&self, vm: &VirtualMachine, key: &K, value: T) -> PyResult<()>
where
Expand Down Expand Up @@ -1758,6 +1787,45 @@ mod tests {
})
}

#[test]
fn reserve_empty_avoids_growth() {
Interpreter::without_stdlib(Default::default()).enter(|vm| {
let dict = Dict::default();
dict.reserve_for_empty(usize::MAX);
assert_eq!(dict.read().indices.len(), 8);
assert_eq!(dict.read().entries.capacity(), 0);

let version = dict.assign_keys_version();
let count = 1000;
dict.reserve_for_empty(count);
assert_eq!(dict.keys_version(), version);
let capacity = {
let inner = dict.read();
(inner.indices.len(), inner.entries.capacity())
};
for key in 0..count {
dict.insert(vm, &key, vm.ctx.none()).unwrap();
let inner = dict.read();
assert_eq!((inner.indices.len(), inner.entries.capacity()), capacity);
}
assert_eq!(dict.len(), count);
assert!(dict.contains(vm, &(count - 1)).unwrap());

// Existing and tombstoned layouts are left intact.
dict.reserve_for_empty(count * 4);
{
let inner = dict.read();
assert_eq!((inner.indices.len(), inner.entries.capacity()), capacity);
}
for key in 0..count {
dict.delete(vm, &key).unwrap();
}
let old_size = dict.read().size();
dict.reserve_for_empty(count * 4);
assert_eq!(dict.read().size(), old_size);
})
}

macro_rules! hash_tests {
($($name:ident: $example_hash:expr,)*) => {
$(
Expand Down
19 changes: 19 additions & 0 deletions extra_tests/snippets/builtin_dict.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,25 @@ class MyClass:
assert {1: None, "b": None} == dict.fromkeys([1, "b"])
assert {1: 0, "b": 0} == dict.fromkeys([1, "b"], 0)

for source in ({i: -1 for i in range(128)}, set(range(128)), frozenset(range(128))):
shared_value = []
result = dict.fromkeys(source, shared_value)
assert list(result) == list(source)
assert all(value is shared_value for value in result.values())


class UnsizedFromKeys:
def __iter__(self):
return iter((1, 1, 2))

def __len__(self):
raise AssertionError("fromkeys must not request a length hint")

__length_hint__ = __len__


assert dict.fromkeys(UnsizedFromKeys()) == {1: None, 2: None}

x = {"a": 1, "b": 1, "c": 1}
y = {"b": 2, "c": 2, "d": 2}
z = {"c": 3, "d": 3, "e": 3}
Expand Down
Loading