Reuse stored hashes when building a set from a set/frozenset/dict

fregataa · claude · fregataa · commit 2924f86ff184 · 2026-08-11T21:09:56.000+09:00
set and frozenset recomputed __hash__ for every element even when the source object already stored a hash per entry. CPython's set_update_internal branches on PyAnySet_Check / PyDict_CheckExact and feeds set_add_entry the hash read from the source table; RustPython always iterated generically. Split the hash computation out of the Dict entry points so callers can supply a hash they already hold, add keys_with_hashes() to hand out (key, hash) pairs, and take the fast path in the set constructors and in the set operations whose argument is a set/frozenset/exact dict. ArgIterable::as_object() exposes the pre-__iter__ object so the set operations can dispatch on the source type without changing any of their signatures. Closes #8489. dict.fromkeys() is the dict-target counterpart and is tracked in #8490, so test_do_not_rehash_dict_keys keeps its expectedFailure marker until that lands too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
diff --git a/crates/vm/src/builtins/set.rs b/crates/vm/src/builtins/set.rs
@@ -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>,
@@ -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_with_hash(
+        &self,
+        needle: &PyObject,
+        hash: PyHash,
+        vm: &VirtualMachine,
+    ) -> PyResult<bool> {
+        self.content.contains_with_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);
@@ -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_with_hash(item, hash, vm)?;
+            }
+            return Ok(set);
+        }
         for item in other.iter(vm)? {
             set.add(item?, vm)?;
         }
@@ -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_with_hash(&obj, hash, vm)? {
+                    set.add_with_hash(obj, hash, vm)?;
+                }
+            }
+            return Ok(set);
+        }
         for item in other.iter(vm)? {
             let obj = item?;
             if self.contains(&obj, vm)? {
@@ -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_with_hash(vm, &*item, hash)?;
+            }
+            return Ok(set);
+        }
         for item in other.iter(vm)? {
             set.content.delete_if_exists(vm, &*item?)?;
         }
@@ -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_with_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)?;
 
@@ -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_with_hash(&self, item: PyObjectRef, hash: PyHash, vm: &VirtualMachine) -> PyResult<()> {
+        let result = self.content.insert_with_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));
@@ -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_with_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_with_hash(key, hash, vm)?;
         }
         Ok(())
     }
@@ -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_with_hash(obj, hash, vm)?;
         }
         Ok(())
     }
@@ -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_with_hash(vm, &*item, hash)?;
+                }
+                continue;
+            }
             let items = iterable.iter(vm)?.collect::<Result<Vec<_>, _>>()?;
             for item in items {
                 self.content.delete_if_exists(vm, &*item)?;
@@ -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_with_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() {
@@ -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);
@@ -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 {}",
@@ -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()
+        })
     }
 }
 
diff --git a/crates/vm/src/dict_inner.rs b/crates/vm/src/dict_inner.rs
@@ -463,6 +463,26 @@ impl<T: Clone> Dict<T> {
         K: DictKey + ?Sized,
     {
         let hash = key.key_hash(vm)?;
+        self.insert_with_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_with_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();
@@ -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_with_hash(vm, key, key_hash)
+    }
+
+    /// [`Self::contains`] with a caller-supplied hash.
+    ///
+    /// Same contract as [`Self::insert_with_hash`]: a wrong `hash` makes the
+    /// lookup probe the wrong bucket and report a present key as missing.
+    pub(crate) fn contains_with_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())
     }
 
@@ -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_with_hash`].
+    pub(crate) fn delete_if_exists_with_hash<K>(
+        &self,
+        vm: &VirtualMachine,
+        key: &K,
+        hash: HashValue,
+    ) -> PyResult<bool>
+    where
+        K: DictKey + ?Sized,
+    {
+        self.remove_if_with_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,
@@ -725,6 +774,23 @@ impl<T: Clone> Dict<T> {
         F: Fn(&T) -> PyResult<bool>,
     {
         let hash = key.key_hash(vm)?;
+        self.remove_if_with_hash(vm, key, hash, pred)
+    }
+
+    /// [`Self::remove_if`] with a caller-supplied hash.
+    ///
+    /// Same contract as [`Self::insert_with_hash`].
+    fn remove_if_with_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)? {
@@ -742,6 +808,19 @@ impl<T: Clone> Dict<T> {
         value: T,
     ) -> PyResult<()> {
         let hash = key.key_hash(vm)?;
+        self.delete_or_insert_with_hash(vm, key, hash, value)
+    }
+
+    /// [`Self::delete_or_insert`] with a caller-supplied hash.
+    ///
+    /// Same contract as [`Self::insert_with_hash`].
+    pub(crate) fn delete_or_insert_with_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;
@@ -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_with_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
diff --git a/crates/vm/src/function/protocol.rs b/crates/vm/src/function/protocol.rs
@@ -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
