Skip to content

Commit 56eadb2

Browse files
1ndahous31ndahous3
authored
Choose smaller difference scans and install completed intersections (#8881)
Assisted-by: Codex:GPT-6
1 parent fd10823 commit 56eadb2

3 files changed

Lines changed: 189 additions & 51 deletions

File tree

‎crates/vm/src/builtins/set.rs‎

Lines changed: 104 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,19 @@ impl PySetInner {
247247
Ok(result)
248248
}
249249

250+
fn difference_multi(
251+
&self,
252+
mut others: impl core::iter::Iterator<Item = ArgIterable>,
253+
vm: &VirtualMachine,
254+
) -> PyResult<Self> {
255+
let Some(other) = others.next() else {
256+
return Ok(self.copy());
257+
};
258+
let result = self.difference_new(other, vm)?;
259+
result.difference_update(others, vm)?;
260+
Ok(result)
261+
}
262+
250263
fn len(&self) -> usize {
251264
self.content.len()
252265
}
@@ -315,28 +328,10 @@ impl PySetInner {
315328
}
316329

317330
pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
318-
let (target, elements) = if let Some(other_set) = extract_set(other.as_object()) {
319-
if core::ptr::eq(self, other_set) {
320-
return Ok(self.copy());
321-
}
322-
let (target, source) = if self.len() < other_set.len() {
323-
(other_set, self)
324-
} else {
325-
(self, other_set)
326-
};
327-
(target, Some(source.content.keys_with_hashes()))
328-
} else {
329-
(self, None)
330-
};
331-
let set = Self::default();
332-
if let Some(elements) = elements {
333-
for (obj, hash) in elements {
334-
if target.contains_known_hash(&obj, hash, vm)? {
335-
set.add_known_hash(&obj, hash, vm)?;
336-
}
337-
}
338-
return Ok(set);
331+
if let Some(other_set) = extract_set(other.as_object()) {
332+
return self.intersection_set(other_set, vm);
339333
}
334+
let set = Self::default();
340335
for item in other.iter(vm)? {
341336
let obj = item?;
342337
let hash = obj.hash(vm)?;
@@ -350,20 +345,63 @@ impl PySetInner {
350345
Ok(set)
351346
}
352347

353-
pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
354-
let set = self.clone();
355-
if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
356-
for (item, hash) in elements {
357-
set.content.delete_if_exists_known_hash(vm, &*item, hash)?;
358-
}
359-
return Ok(set);
348+
fn intersection_set(&self, other: &Self, vm: &VirtualMachine) -> PyResult<Self> {
349+
if PyRc::ptr_eq(&self.content, &other.content) {
350+
return Ok(self.copy());
360351
}
361-
for item in other.iter(vm)? {
362-
set.content.delete_if_exists(vm, &*item?)?;
352+
let (target, source) = if self.len() < other.len() {
353+
(other, self)
354+
} else {
355+
(self, other)
356+
};
357+
let set = Self::default();
358+
for (obj, hash) in source.content.keys_with_hashes() {
359+
if target.contains_known_hash(&obj, hash, vm)? {
360+
set.add_known_hash(&obj, hash, vm)?;
361+
}
363362
}
364363
Ok(set)
365364
}
366365

366+
fn difference_new(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
367+
// Scanning the left side avoids visiting and retaining a much larger
368+
// exclusion collection. Match CPython's comparison direction as well.
369+
if let Some(other_set) = extract_set(other.as_object()) {
370+
if self.len() >> 2 <= other_set.len() {
371+
return self
372+
.difference_by(|key, hash| other_set.contains_known_hash(key, hash, vm), vm);
373+
}
374+
} else if let Some(dict) = other.as_object().downcast_ref_if_exact::<PyDict>(vm)
375+
&& self.len() >> 2 <= dict._as_dict_inner().len()
376+
{
377+
return self.difference_by(
378+
|key, hash| dict._as_dict_inner().contains_known_hash(vm, key, hash),
379+
vm,
380+
);
381+
}
382+
self.copy().difference(other, vm)
383+
}
384+
385+
// The dict-view caller supplies a private working table.
386+
pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
387+
self.difference_update(core::iter::once(other), vm)?;
388+
Ok(self.clone())
389+
}
390+
391+
fn difference_by(
392+
&self,
393+
contains: impl Fn(&PyObject, PyHash) -> PyResult<bool>,
394+
vm: &VirtualMachine,
395+
) -> PyResult<Self> {
396+
let result = Self::default();
397+
for (key, hash) in self.content.keys_with_hashes() {
398+
if !contains(&key, hash)? {
399+
result.add_known_hash(&key, hash, vm)?;
400+
}
401+
}
402+
Ok(result)
403+
}
404+
367405
pub(super) fn symmetric_difference(
368406
&self,
369407
other: ArgIterable,
@@ -530,10 +568,8 @@ impl PySetInner {
530568
vm: &VirtualMachine,
531569
) -> PyResult<()> {
532570
let temp_inner = self.intersection_multi(others, vm)?;
533-
self.clear();
534-
for (obj, hash) in temp_inner.content.keys_with_hashes() {
535-
self.add_known_hash(&obj, hash, vm)?;
536-
}
571+
let content = PyRc::try_unwrap(temp_inner.content).unwrap_or_else(|table| (*table).clone());
572+
self.content.replace_contents(content);
537573
Ok(())
538574
}
539575

@@ -543,7 +579,24 @@ impl PySetInner {
543579
vm: &VirtualMachine,
544580
) -> PyResult<()> {
545581
for iterable in others {
546-
if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) {
582+
let elements = if let Some(other_set) = extract_set(iterable.as_object()) {
583+
if PyRc::ptr_eq(&self.content, &other_set.content) {
584+
self.clear();
585+
continue;
586+
}
587+
// Build the intersection first: besides bounding the work by
588+
// our size, this preserves CPython's equality/error ordering.
589+
Some(if other_set.len() >> 3 > self.len() {
590+
self.intersection_set(other_set, vm)?
591+
.content
592+
.keys_with_hashes()
593+
} else {
594+
other_set.content.keys_with_hashes()
595+
})
596+
} else {
597+
Self::cached_hashes(iterable.as_object(), vm)
598+
};
599+
if let Some(elements) = elements {
547600
for (item, hash) in elements {
548601
self.content.delete_if_exists_known_hash(vm, &*item, hash)?;
549602
}
@@ -748,7 +801,9 @@ impl PySet {
748801
others: PosArgs<ArgIterable, NameOthers>,
749802
vm: &VirtualMachine,
750803
) -> PyResult<Self> {
751-
self.fold_op(others.into_iter(), PySetInner::difference, vm)
804+
Ok(Self {
805+
inner: self.inner.difference_multi(others.into_iter(), vm)?,
806+
})
752807
}
753808

754809
#[pymethod]
@@ -807,11 +862,9 @@ impl PySet {
807862
vm: &VirtualMachine,
808863
) -> PyResult<PyArithmeticValue<Self>> {
809864
if let Ok(other) = AnySet::try_from_object(vm, other) {
810-
Ok(PyArithmeticValue::Implemented(self.op(
811-
other,
812-
PySetInner::difference,
813-
vm,
814-
)?))
865+
Ok(PyArithmeticValue::Implemented(Self {
866+
inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?,
867+
}))
815868
} else {
816869
Ok(PyArithmeticValue::NotImplemented)
817870
}
@@ -826,8 +879,7 @@ impl PySet {
826879
Ok(PyArithmeticValue::Implemented(Self {
827880
inner: other
828881
.as_inner()
829-
.copy()
830-
.difference(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
882+
.difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
831883
}))
832884
} else {
833885
Ok(PyArithmeticValue::NotImplemented)
@@ -1277,7 +1329,10 @@ impl PyFrozenSet {
12771329
others: PosArgs<ArgIterable, NameOthers>,
12781330
vm: &VirtualMachine,
12791331
) -> PyResult<Self> {
1280-
self.fold_op(others.into_iter(), PySetInner::difference, vm)
1332+
Ok(Self {
1333+
inner: self.inner.difference_multi(others.into_iter(), vm)?,
1334+
..Default::default()
1335+
})
12811336
}
12821337

12831338
#[pymethod]
@@ -1337,11 +1392,10 @@ impl PyFrozenSet {
13371392
vm: &VirtualMachine,
13381393
) -> PyResult<PyArithmeticValue<Self>> {
13391394
if let Ok(other) = AnySet::try_from_object(vm, other) {
1340-
Ok(PyArithmeticValue::Implemented(self.op(
1341-
other,
1342-
PySetInner::difference,
1343-
vm,
1344-
)?))
1395+
Ok(PyArithmeticValue::Implemented(Self {
1396+
inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?,
1397+
..Default::default()
1398+
}))
13451399
} else {
13461400
Ok(PyArithmeticValue::NotImplemented)
13471401
}
@@ -1356,8 +1410,7 @@ impl PyFrozenSet {
13561410
Ok(PyArithmeticValue::Implemented(Self {
13571411
inner: other
13581412
.as_inner()
1359-
.copy()
1360-
.difference(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
1413+
.difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
13611414
..Default::default()
13621415
}))
13631416
} else {

‎crates/vm/src/dict_inner.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -860,6 +860,16 @@ impl<T: Clone> Dict<T> {
860860
};
861861
}
862862

863+
/// Install an owned table, releasing the old entries after unlocking.
864+
pub(crate) fn replace_contents(&self, other: Self) {
865+
let replacement = other.inner.into_inner();
866+
let _removed = {
867+
let mut inner = self.write();
868+
self.invalidate_keys_version();
869+
core::mem::replace(&mut *inner, replacement)
870+
};
871+
}
872+
863873
/// Delete a key
864874
pub(crate) fn delete<K>(&self, vm: &VirtualMachine, key: &K) -> PyResult<()>
865875
where
@@ -1964,4 +1974,24 @@ mod tests {
19641974
assert_eq!(hash1, hash2);
19651975
})
19661976
}
1977+
1978+
#[test]
1979+
fn replace_contents_invalidates_cached_indices() {
1980+
Interpreter::without_stdlib(Default::default()).enter(|vm| {
1981+
let dict = Dict::default();
1982+
dict.insert(vm, &1usize, ()).unwrap();
1983+
let old_version = dict.assign_keys_version();
1984+
assert_ne!(old_version, 0);
1985+
1986+
let replacement = Dict::default();
1987+
replacement.insert(vm, &2usize, ()).unwrap();
1988+
assert_ne!(replacement.assign_keys_version(), 0);
1989+
dict.replace_contents(replacement);
1990+
1991+
assert_eq!(dict.keys_version(), 0);
1992+
assert_eq!(dict.get_index_if_keys_version(old_version, 0), None);
1993+
assert!(!dict.contains(vm, &1usize).unwrap());
1994+
assert!(dict.contains(vm, &2usize).unwrap());
1995+
});
1996+
}
19671997
}

‎extra_tests/snippets/builtin_set.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -651,3 +651,58 @@ def __iter__(self):
651651
assert len(source) == 1
652652
source.difference_update(source)
653653
assert source == set()
654+
655+
656+
# Difference chooses comparison direction by size, retaining left-side keys.
657+
for left_type in (set, frozenset):
658+
left_key = DirectionalKey(False)
659+
right_key = DirectionalKey(True)
660+
left = left_type([left_key])
661+
for right_type in (set, frozenset, NativeSetSource, NativeFrozenSetSource, dict):
662+
right = (
663+
dict.fromkeys([right_key, *range(32)])
664+
if right_type is dict
665+
else right_type([right_key, *range(32)])
666+
)
667+
assert left.difference(right) == set()
668+
assert list(left) == [left_key]
669+
if right_type is not dict:
670+
assert left - right == set()
671+
updated = set(left)
672+
updated.difference_update(right)
673+
assert updated == set()
674+
# Later arguments use in-place difference's comparison direction.
675+
assert list(left.difference(set(), {right_key})) == [left_key]
676+
large_left = left_type([left_key, *range(32)])
677+
assert left_key in large_left.difference({right_key})
678+
679+
680+
key = StoredHashKey(100)
681+
source = {key}
682+
excluded = {key, *range(32)}
683+
dictionary = dict.fromkeys(excluded)
684+
key.hash_enabled = False
685+
assert not source.difference(excluded)
686+
assert not source.difference(dictionary)
687+
source.difference_update(excluded)
688+
assert not source
689+
690+
691+
# Removed objects observe the completed intersection, including their mutations.
692+
intersection_events = []
693+
694+
695+
class RemovedFromIntersection:
696+
def __del__(self):
697+
intersection_events.append(1 in intersection_target)
698+
intersection_target.discard(1)
699+
700+
701+
for use_operator in (False, True):
702+
intersection_target = {1, RemovedFromIntersection()}
703+
if use_operator:
704+
intersection_target &= {1}
705+
else:
706+
intersection_target.intersection_update({1})
707+
assert not intersection_target
708+
assert intersection_events == [True, True]

0 commit comments

Comments
 (0)