Reuse stored hashes for frozenset hashing and in-place set union (#8859)

Assisted-by: Codex:GPT-6

crates/vm/src/builtins/set.rs

@@ -491,20 +491,6 @@ impl PySetInner {
         }
     }
 
-    fn update(
-        &self,
-        others: impl core::iter::Iterator<Item = ArgIterable>,
-        vm: &VirtualMachine,
-    ) -> PyResult<()> {
-        for iterable in others {
-            for item in iterable.iter(vm)? {
-                let item = item?;
-                self.add(&item, vm)?;
-            }
-        }
-        Ok(())
-    }
-
     fn update_internal(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
         // check AnySet
         if let Ok(any_set) = AnySet::try_from_object(vm, iterable.to_owned()) {
@@ -592,15 +578,15 @@ impl PySetInner {
         Ok(())
     }
 
-    fn hash(&self, vm: &VirtualMachine) -> PyResult<PyHash> {
-        let hasher = self.content.try_fold_keys(
+    fn hash(&self) -> PyHash {
+        let hasher = self.content.fold_hashes(
             hash::FrozenSetHash::new(self.len()),
-            |mut hasher, element| {
-                hasher.add(element.hash(vm)?);
-                Ok(hasher)
+            |mut hasher, element_hash| {
+                hasher.add(element_hash);
+                hasher
             },
-        )?;
-        Ok(hasher.finish())
+        );
+        hasher.finish()
     }
 
     // Run operation, on failure, if item is a set/set subclass, convert it
@@ -887,7 +873,7 @@ impl PySet {
     }
 
     fn __ior__(zelf: PyRef<Self>, set: AnySet, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-        zelf.inner.update(set.into_iterable_iter(vm)?, vm)?;
+        zelf.inner.merge_set(set, vm)?;
         Ok(zelf)
     }
 
@@ -1424,10 +1410,10 @@ impl AsSequence for PyFrozenSet {
 
 impl Hashable for PyFrozenSet {
     #[inline]
-    fn hash(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
+    fn hash(zelf: &crate::Py<Self>, _vm: &VirtualMachine) -> PyResult<PyHash> {
         let hash = match zelf.hash.load(Ordering::Relaxed) {
             hash::SENTINEL => {
-                let hash = zelf.inner.hash(vm)?;
+                let hash = zelf.inner.hash();
                 match Radium::compare_exchange(
                     &zelf.hash,
                     hash::SENTINEL,

crates/vm/src/dict_inner.rs

@@ -1114,15 +1114,16 @@ impl<T: Clone> Dict<T> {
             .collect()
     }
 
-    pub(crate) fn try_fold_keys<Acc, Fold>(&self, init: Acc, f: Fold) -> PyResult<Acc>
+    /// Fold stored hashes under the read lock; the callback must not run Python.
+    pub(crate) fn fold_hashes<Acc, Fold>(&self, init: Acc, f: Fold) -> Acc
     where
-        Fold: FnMut(Acc, &PyObject) -> PyResult<Acc>,
+        Fold: FnMut(Acc, HashValue) -> Acc,
     {
         self.read()
             .entries
             .iter()
-            .filter_map(|v| v.as_ref().map(|v| v.key.as_object()))
-            .try_fold(init, f)
+            .filter_map(|v| v.as_ref().map(|v| v.hash))
+            .fold(init, f)
     }
 
     /// Lookup the index for the given key.

extra_tests/snippets/builtin_set.py

@@ -566,3 +566,49 @@ assert key.hash_count == 1
 dictionary = {key: None}
 key.hash_enabled = False
 assert_raises(AssertionError, source.intersection, dictionary)
+
+
+# Native operations reuse hashes stored when an element enters the collection.
+class RemainingHashKey:
+    blocked = False
+
+    def __init__(self, value):
+        self.value = value
+
+    def __hash__(self):
+        assert not type(self).blocked, "stored key was hashed again"
+        return 7
+
+    def __eq__(self, other):
+        return isinstance(other, RemainingHashKey) and self.value == other.value
+
+
+class NativeSetSource(set):
+    def __iter__(self):
+        raise AssertionError("native set operation called __iter__")
+
+
+class NativeFrozenSetSource(frozenset):
+    def __iter__(self):
+        raise AssertionError("native set operation called __iter__")
+
+
+stored_keys = [RemainingHashKey(1), RemainingHashKey(2)]
+stored_frozen = frozenset(stored_keys)
+stored_frozen_reversed = frozenset(reversed(stored_keys))
+stored_sources = [
+    source_type(stored_keys)
+    for source_type in (set, frozenset, NativeSetSource, NativeFrozenSetSource)
+]
+RemainingHashKey.blocked = True
+stored_hash = hash(stored_frozen)
+assert stored_hash == hash(stored_frozen_reversed)
+assert stored_hash == hash(stored_frozen)
+for stored_source in stored_sources:
+    stored_target = set()
+    assert stored_target.__ior__(stored_source) is stored_target
+    assert len(stored_target) == 2
+    assert {item.value for item in stored_target} == {1, 2}
+    stored_target |= stored_target
+    assert len(stored_target) == 2
+RemainingHashKey.blocked = False