From 56eadb29974f7fd6cd3864c8bb0be2d24adafbea Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:13:26 +0300 Subject: [PATCH] Choose smaller difference scans and install completed intersections (#8881) Assisted-by: Codex:GPT-6 --- crates/vm/src/builtins/set.rs | 155 +++++++++++++++++++--------- crates/vm/src/dict_inner.rs | 30 ++++++ extra_tests/snippets/builtin_set.py | 55 ++++++++++ 3 files changed, 189 insertions(+), 51 deletions(-) diff --git a/crates/vm/src/builtins/set.rs b/crates/vm/src/builtins/set.rs index 4993c4e2c..ce4c0c067 100644 --- a/crates/vm/src/builtins/set.rs +++ b/crates/vm/src/builtins/set.rs @@ -247,6 +247,19 @@ impl PySetInner { Ok(result) } + fn difference_multi( + &self, + mut others: impl core::iter::Iterator, + vm: &VirtualMachine, + ) -> PyResult { + let Some(other) = others.next() else { + return Ok(self.copy()); + }; + let result = self.difference_new(other, vm)?; + result.difference_update(others, vm)?; + Ok(result) + } + fn len(&self) -> usize { self.content.len() } @@ -315,28 +328,10 @@ impl PySetInner { } pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { - let (target, elements) = if let Some(other_set) = extract_set(other.as_object()) { - if core::ptr::eq(self, other_set) { - return Ok(self.copy()); - } - let (target, source) = if self.len() < other_set.len() { - (other_set, self) - } else { - (self, other_set) - }; - (target, Some(source.content.keys_with_hashes())) - } else { - (self, None) - }; - let set = Self::default(); - if let Some(elements) = elements { - for (obj, hash) in elements { - if target.contains_known_hash(&obj, hash, vm)? { - set.add_known_hash(&obj, hash, vm)?; - } - } - return Ok(set); + if let Some(other_set) = extract_set(other.as_object()) { + return self.intersection_set(other_set, vm); } + let set = Self::default(); for item in other.iter(vm)? { let obj = item?; let hash = obj.hash(vm)?; @@ -350,20 +345,63 @@ impl PySetInner { Ok(set) } - pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { - let set = self.clone(); - 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); + fn intersection_set(&self, other: &Self, vm: &VirtualMachine) -> PyResult { + if PyRc::ptr_eq(&self.content, &other.content) { + return Ok(self.copy()); } - for item in other.iter(vm)? { - set.content.delete_if_exists(vm, &*item?)?; + let (target, source) = if self.len() < other.len() { + (other, self) + } else { + (self, other) + }; + let set = Self::default(); + for (obj, hash) in source.content.keys_with_hashes() { + if target.contains_known_hash(&obj, hash, vm)? { + set.add_known_hash(&obj, hash, vm)?; + } } Ok(set) } + fn difference_new(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { + // Scanning the left side avoids visiting and retaining a much larger + // exclusion collection. Match CPython's comparison direction as well. + if let Some(other_set) = extract_set(other.as_object()) { + if self.len() >> 2 <= other_set.len() { + return self + .difference_by(|key, hash| other_set.contains_known_hash(key, hash, vm), vm); + } + } else if let Some(dict) = other.as_object().downcast_ref_if_exact::(vm) + && self.len() >> 2 <= dict._as_dict_inner().len() + { + return self.difference_by( + |key, hash| dict._as_dict_inner().contains_known_hash(vm, key, hash), + vm, + ); + } + self.copy().difference(other, vm) + } + + // The dict-view caller supplies a private working table. + pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { + self.difference_update(core::iter::once(other), vm)?; + Ok(self.clone()) + } + + fn difference_by( + &self, + contains: impl Fn(&PyObject, PyHash) -> PyResult, + vm: &VirtualMachine, + ) -> PyResult { + let result = Self::default(); + for (key, hash) in self.content.keys_with_hashes() { + if !contains(&key, hash)? { + result.add_known_hash(&key, hash, vm)?; + } + } + Ok(result) + } + pub(super) fn symmetric_difference( &self, other: ArgIterable, @@ -530,10 +568,8 @@ impl PySetInner { vm: &VirtualMachine, ) -> PyResult<()> { let temp_inner = self.intersection_multi(others, vm)?; - self.clear(); - for (obj, hash) in temp_inner.content.keys_with_hashes() { - self.add_known_hash(&obj, hash, vm)?; - } + let content = PyRc::try_unwrap(temp_inner.content).unwrap_or_else(|table| (*table).clone()); + self.content.replace_contents(content); Ok(()) } @@ -543,7 +579,24 @@ impl PySetInner { vm: &VirtualMachine, ) -> PyResult<()> { for iterable in others { - if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) { + let elements = if let Some(other_set) = extract_set(iterable.as_object()) { + if PyRc::ptr_eq(&self.content, &other_set.content) { + self.clear(); + continue; + } + // Build the intersection first: besides bounding the work by + // our size, this preserves CPython's equality/error ordering. + Some(if other_set.len() >> 3 > self.len() { + self.intersection_set(other_set, vm)? + .content + .keys_with_hashes() + } else { + other_set.content.keys_with_hashes() + }) + } else { + Self::cached_hashes(iterable.as_object(), vm) + }; + if let Some(elements) = elements { for (item, hash) in elements { self.content.delete_if_exists_known_hash(vm, &*item, hash)?; } @@ -748,7 +801,9 @@ impl PySet { others: PosArgs, vm: &VirtualMachine, ) -> PyResult { - self.fold_op(others.into_iter(), PySetInner::difference, vm) + Ok(Self { + inner: self.inner.difference_multi(others.into_iter(), vm)?, + }) } #[pymethod] @@ -807,11 +862,9 @@ impl PySet { vm: &VirtualMachine, ) -> PyResult> { if let Ok(other) = AnySet::try_from_object(vm, other) { - Ok(PyArithmeticValue::Implemented(self.op( - other, - PySetInner::difference, - vm, - )?)) + Ok(PyArithmeticValue::Implemented(Self { + inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?, + })) } else { Ok(PyArithmeticValue::NotImplemented) } @@ -826,8 +879,7 @@ impl PySet { Ok(PyArithmeticValue::Implemented(Self { inner: other .as_inner() - .copy() - .difference(ArgIterable::try_from_object(vm, zelf.into())?, vm)?, + .difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?, })) } else { Ok(PyArithmeticValue::NotImplemented) @@ -1277,7 +1329,10 @@ impl PyFrozenSet { others: PosArgs, vm: &VirtualMachine, ) -> PyResult { - self.fold_op(others.into_iter(), PySetInner::difference, vm) + Ok(Self { + inner: self.inner.difference_multi(others.into_iter(), vm)?, + ..Default::default() + }) } #[pymethod] @@ -1337,11 +1392,10 @@ impl PyFrozenSet { vm: &VirtualMachine, ) -> PyResult> { if let Ok(other) = AnySet::try_from_object(vm, other) { - Ok(PyArithmeticValue::Implemented(self.op( - other, - PySetInner::difference, - vm, - )?)) + Ok(PyArithmeticValue::Implemented(Self { + inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?, + ..Default::default() + })) } else { Ok(PyArithmeticValue::NotImplemented) } @@ -1356,8 +1410,7 @@ impl PyFrozenSet { Ok(PyArithmeticValue::Implemented(Self { inner: other .as_inner() - .copy() - .difference(ArgIterable::try_from_object(vm, zelf.into())?, vm)?, + .difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?, ..Default::default() })) } else { diff --git a/crates/vm/src/dict_inner.rs b/crates/vm/src/dict_inner.rs index 25570e630..c354114ae 100644 --- a/crates/vm/src/dict_inner.rs +++ b/crates/vm/src/dict_inner.rs @@ -860,6 +860,16 @@ impl Dict { }; } + /// Install an owned table, releasing the old entries after unlocking. + pub(crate) fn replace_contents(&self, other: Self) { + let replacement = other.inner.into_inner(); + let _removed = { + let mut inner = self.write(); + self.invalidate_keys_version(); + core::mem::replace(&mut *inner, replacement) + }; + } + /// Delete a key pub(crate) fn delete(&self, vm: &VirtualMachine, key: &K) -> PyResult<()> where @@ -1964,4 +1974,24 @@ mod tests { assert_eq!(hash1, hash2); }) } + + #[test] + fn replace_contents_invalidates_cached_indices() { + Interpreter::without_stdlib(Default::default()).enter(|vm| { + let dict = Dict::default(); + dict.insert(vm, &1usize, ()).unwrap(); + let old_version = dict.assign_keys_version(); + assert_ne!(old_version, 0); + + let replacement = Dict::default(); + replacement.insert(vm, &2usize, ()).unwrap(); + assert_ne!(replacement.assign_keys_version(), 0); + dict.replace_contents(replacement); + + assert_eq!(dict.keys_version(), 0); + assert_eq!(dict.get_index_if_keys_version(old_version, 0), None); + assert!(!dict.contains(vm, &1usize).unwrap()); + assert!(dict.contains(vm, &2usize).unwrap()); + }); + } } diff --git a/extra_tests/snippets/builtin_set.py b/extra_tests/snippets/builtin_set.py index 514809e14..8e96a0959 100644 --- a/extra_tests/snippets/builtin_set.py +++ b/extra_tests/snippets/builtin_set.py @@ -651,3 +651,58 @@ assert_raises(RuntimeError, source.difference_update, iter(source)) assert len(source) == 1 source.difference_update(source) assert source == set() + + +# Difference chooses comparison direction by size, retaining left-side keys. +for left_type in (set, frozenset): + left_key = DirectionalKey(False) + right_key = DirectionalKey(True) + left = left_type([left_key]) + for right_type in (set, frozenset, NativeSetSource, NativeFrozenSetSource, dict): + right = ( + dict.fromkeys([right_key, *range(32)]) + if right_type is dict + else right_type([right_key, *range(32)]) + ) + assert left.difference(right) == set() + assert list(left) == [left_key] + if right_type is not dict: + assert left - right == set() + updated = set(left) + updated.difference_update(right) + assert updated == set() + # Later arguments use in-place difference's comparison direction. + assert list(left.difference(set(), {right_key})) == [left_key] + large_left = left_type([left_key, *range(32)]) + assert left_key in large_left.difference({right_key}) + + +key = StoredHashKey(100) +source = {key} +excluded = {key, *range(32)} +dictionary = dict.fromkeys(excluded) +key.hash_enabled = False +assert not source.difference(excluded) +assert not source.difference(dictionary) +source.difference_update(excluded) +assert not source + + +# Removed objects observe the completed intersection, including their mutations. +intersection_events = [] + + +class RemovedFromIntersection: + def __del__(self): + intersection_events.append(1 in intersection_target) + intersection_target.discard(1) + + +for use_operator in (False, True): + intersection_target = {1, RemovedFromIntersection()} + if use_operator: + intersection_target &= {1} + else: + intersection_target.intersection_update({1}) + assert not intersection_target +assert intersection_events == [True, True]