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

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

/// Elements of `obj` with their stored hashes, or `None` if `obj` keeps
/// none and must be iterated generically. 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 +248,17 @@ impl PySetInner {
Self::wrap_unhashable_error(result, needle, vm)
}

/// [`Self::contains`] with a known hash. Such a needle came out of a
/// set/dict, so it is hashable and needs no frozenset retry.
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 +282,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 +297,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 +316,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 +335,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 +394,12 @@ impl PySetInner {
Self::wrap_unhashable_error(result, &item, vm)
}

/// [`Self::add`] with a known hash.
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 +460,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 +463 to +471

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 +480,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 +492,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 +512,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 +1036,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 +1069,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 +1083,25 @@ 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> {
let inner = match iterable {
OptionalArg::Present(iterable) => PySetInner::from_object(iterable, vm)?,
OptionalArg::Missing => PySetInner::default(),
};
Ok(Self {
inner,
..Default::default()
})
}
}

Expand Down
84 changes: 83 additions & 1 deletion crates/vm/src/dict_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,24 @@ 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 equal `key.key_hash(vm)`; a wrong one lands the entry in a
/// bucket no lookup probes, silently losing the key. Only pass a hash from
/// [`Self::keys_with_hashes`] on a container holding this same key.
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 +530,18 @@ 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 known hash. Same contract as
/// [`Self::insert_known_hash`].
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 +726,21 @@ impl<T: Clone> Dict<T> {
self.remove_if_exists(vm, key).map(|opt| opt.is_some())
}

/// [`Self::delete_if_exists`] with a known 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 +769,22 @@ 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 known 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 +802,18 @@ 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 known 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 +960,16 @@ impl<T: Clone> Dict<T> {
.collect()
}

/// All keys paired with the hash stored in their entry, for feeding
/// [`Self::insert_known_hash`] without re-calling `__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(crate) 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
Loading