Skip to content
Open
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
127 changes: 108 additions & 19 deletions crates/vm/src/builtins/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,29 @@ impl PySetInner {
Ok(set)
}

/// Build a set from an arbitrary object, reusing the source's stored
/// hashes when it is a set/frozenset/dict.
pub(super) fn from_object(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
let set = Self::default();
set.update_internal(iterable, vm)?;
Ok(set)
}

/// Elements of `obj` paired with the hash already stored alongside them,
/// or `None` if `obj` keeps no such hashes and has to be iterated
/// generically (which calls `__hash__` on every element again).
///
/// Mirrors the `PyAnySet_Check` / `PyDict_CheckExact` fast paths in
/// CPython's `set_update_internal`.
fn cached_hashes(obj: &PyObject, vm: &VirtualMachine) -> Option<Vec<(PyObjectRef, PyHash)>> {
if let Some(set) = extract_set(obj) {
Some(set.content.keys_with_hashes())
} else {
obj.downcast_ref_if_exact::<PyDict>(vm)
.map(|dict| dict._as_dict_inner().keys_with_hashes())
}
}

fn fold_op<O>(
&self,
others: impl core::iter::Iterator<Item = O>,
Expand Down Expand Up @@ -228,6 +251,18 @@ impl PySetInner {
Self::wrap_unhashable_error(result, needle, vm)
}

/// [`Self::contains`] for a needle whose hash was read from another
/// set/dict. Such a needle is hashable by construction, so there is no
/// set-to-frozenset retry to do.
fn contains_known_hash(
&self,
needle: &PyObject,
hash: PyHash,
vm: &VirtualMachine,
) -> PyResult<bool> {
self.content.contains_known_hash(vm, needle, hash)
}

fn compare(&self, other: &Self, op: PyComparisonOp, vm: &VirtualMachine) -> PyResult<bool> {
if op == PyComparisonOp::Ne {
return self.compare(other, PyComparisonOp::Eq, vm).map(|eq| !eq);
Expand All @@ -251,6 +286,12 @@ impl PySetInner {

pub(super) fn union(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
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)?;
}
Expand All @@ -260,6 +301,14 @@ impl PySetInner {

pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
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)? {
Expand All @@ -271,6 +320,12 @@ impl PySetInner {

pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
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?)?;
}
Expand All @@ -284,6 +339,16 @@ impl PySetInner {
) -> PyResult<Self> {
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)?;

Expand Down Expand Up @@ -333,6 +398,12 @@ impl PySetInner {
Self::wrap_unhashable_error(result, &item, vm)
}

/// [`Self::add`] for an item whose hash was read from another set/dict.
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));
Expand Down Expand Up @@ -393,15 +464,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)?;
Comment on lines +467 to +475

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

ast-grep outline crates/vm/src/builtins/set.rs --items all --type function --match 'update|update_internal|merge_set|merge_dict'
rg -n -C 5 --type rust 'fn update\b|\.update\(' crates/vm/src/builtins/set.rs

Repository: RustPython/RustPython

Length of output: 2076


🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '400,490p' crates/vm/src/builtins/set.rs
sed -n '780,815p' crates/vm/src/builtins/set.rs
rg -n -C 4 --type rust 'enum ArgIterable|struct ArgIterable|impl ArgIterable|fn as_object|into_iterable_iter|keys_with_hashes|add_known_hash' crates/vm/src
rg -n -C 6 --type rust 'update_internal|\.inner\.update\(|inner\.update\(' crates/vm/src/builtins/set.rs crates/vm/src | head -n 240
rg -n --glob '*.rs' 'hash.*raises|raises.*hash|__hash__|Hash' Lib tests crates/vm/src/builtins/set.rs | head -n 240

Repository: RustPython/RustPython

Length of output: 25269


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

p = Path("crates/vm/src/builtins/set.rs")
lines = p.read_text().splitlines()
for start, end in [(430, 485), (790, 810), (870, 890)]:
    print(f"--- {p}:{start}-{end} ---")
    for n in range(start, min(end, len(lines)) + 1):
        print(f"{n}:{lines[n-1]}")
PY

rg -n -C 8 --type rust 'ArgIterable' crates/vm/src | head -n 320
rg -n -C 8 --type rust 'keys_with_hashes|add_known_hash' crates/vm/src/builtins | head -n 320

Repository: RustPython/RustPython

Length of output: 32348


🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '60,115p' crates/vm/src/function/protocol.rs
sed -n '1410,1450p' crates/vm/src/builtins/set.rs
rg -n -C 10 --type rust 'enum AnySet|struct AnySet|impl.*AnySet|TryFromObject.*AnySet|into_iterable_iter' crates/vm/src/builtins/set.rs crates/vm/src
rg -n --type rust 'inner\.update\(' crates/vm/src/builtins/set.rs

Repository: RustPython/RustPython

Length of output: 19080


🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '70,115p' crates/vm/src/function/protocol.rs
sed -n '1380,1460p' crates/vm/src/builtins/set.rs
rg -n -C 10 --type rust 'AnySet' crates/vm/src/builtins/set.rs | head -n 260
rg -n --type rust 'inner\.update\(' crates/vm/src/builtins/set.rs

Repository: RustPython/RustPython

Length of output: 13457


Route set.__ior__() through the cached-hash path.

set.update() already calls update_internal, but set.__ior__() calls PySetInner::update, which recomputes __hash__ for elements from an existing set. Delegate PySetInner::update to update_internal and add a regression test for set |= existing_set.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/vm/src/builtins/set.rs` around lines 467 - 475, Update
PySetInner::update to delegate to update_internal so set.__ior__() reuses cached
hashes when merging an existing set instead of recomputing __hash__. Preserve
set.update() behavior and add a regression test covering set |= existing_set
with hash invocation tracking.

}
Ok(())
}
Expand All @@ -413,8 +484,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(())
}
Expand All @@ -425,6 +496,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::<Result<Vec<_>, _>>()?;
for item in items {
self.content.delete_if_exists(vm, &*item)?;
Expand All @@ -439,6 +516,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() {
Expand Down Expand Up @@ -955,7 +1040,7 @@ impl Representable for PySet {
}

impl Constructor for PyFrozenSet {
type Args = Vec<PyObjectRef>;
type Args = OptionalArg<PyObjectRef>;

fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
let is_exact_frozenset = cls.is(vm.ctx.types.frozenset_type);
Expand Down Expand Up @@ -988,11 +1073,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 {}",
Expand All @@ -1002,23 +1087,27 @@ 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<PyType>, elements: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
Self::from_iter(vm, elements)
fn py_new(_cls: &Py<PyType>, iterable: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
// built from the object itself, so a set/frozenset/dict source can
// hand over its stored hashes instead of being re-hashed
let inner = match iterable {
OptionalArg::Present(iterable) => PySetInner::from_object(iterable, vm)?,
OptionalArg::Missing => PySetInner::default(),
};
Ok(Self {
inner,
..Default::default()
})
}
}

Expand Down
93 changes: 92 additions & 1 deletion crates/vm/src/dict_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,26 @@ impl<T: Clone> Dict<T> {
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 be the value `key.key_hash(vm)` would return. Passing a
/// different hash corrupts the table: the entry is stored in a bucket no
/// lookup will probe, so the key silently goes missing. Only pass a hash
/// read out of another dict/set entry holding this very key object, e.g.
/// via [`Self::keys_with_hashes`].
pub(crate) fn insert_known_hash<K>(
&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();
Expand Down Expand Up @@ -512,7 +532,20 @@ impl<T: Clone> Dict<T> {
key: &K,
) -> PyResult<bool> {
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 caller-supplied hash.
///
/// Same contract as [`Self::insert_known_hash`]: a wrong `hash` makes the
/// lookup probe the wrong bucket and report a present key as missing.
pub(crate) fn contains_known_hash<K: DictKey + ?Sized>(
&self,
vm: &VirtualMachine,
key: &K,
hash: HashValue,
) -> PyResult<bool> {
let (entry, _) = self.lookup(vm, key, hash, None)?;
Ok(entry.index().is_some())
}

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

/// [`Self::delete_if_exists`] with a caller-supplied hash.
///
/// Same contract as [`Self::insert_known_hash`].
pub(crate) fn delete_if_exists_known_hash<K>(
&self,
vm: &VirtualMachine,
key: &K,
hash: HashValue,
) -> PyResult<bool>
where
K: DictKey + ?Sized,
{
self.remove_if_known_hash(vm, key, hash, |_| Ok(true))
.map(|opt| opt.is_some())
}

pub(crate) fn delete_if<K, F>(&self, vm: &VirtualMachine, key: &K, pred: F) -> PyResult<bool>
where
K: DictKey + ?Sized,
Expand Down Expand Up @@ -725,6 +774,23 @@ impl<T: Clone> Dict<T> {
F: Fn(&T) -> PyResult<bool>,
{
let hash = key.key_hash(vm)?;
self.remove_if_known_hash(vm, key, hash, pred)
}

/// [`Self::remove_if`] with a caller-supplied hash.
///
/// Same contract as [`Self::insert_known_hash`].
fn remove_if_known_hash<K, F>(
&self,
vm: &VirtualMachine,
key: &K,
hash: HashValue,
pred: F,
) -> PyResult<Option<T>>
where
K: DictKey + ?Sized,
F: Fn(&T) -> PyResult<bool>,
{
let removed = loop {
let lookup = self.lookup(vm, key, hash, None)?;
match self.pop_inner_if(lookup, &pred)? {
Expand All @@ -742,6 +808,19 @@ impl<T: Clone> Dict<T> {
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 caller-supplied 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;
Expand Down Expand Up @@ -888,6 +967,18 @@ impl<T: Clone> Dict<T> {
.collect()
}

/// All keys paired with the hash already stored in their entry.
///
/// Lets a caller move keys into another dict/set without calling
/// `__hash__` again; see [`Self::insert_known_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<T> {
self.read()
.entries
Expand Down
5 changes: 5 additions & 0 deletions crates/vm/src/function/protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ unsafe impl<T: Traverse> Traverse for ArgIterable<T> {
}

impl<T> ArgIterable<T> {
#[must_use]
pub 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
Expand Down