diff --git a/crates/vm/src/builtins/set.rs b/crates/vm/src/builtins/set.rs index 8638e18efd3..664cd24262d 100644 --- a/crates/vm/src/builtins/set.rs +++ b/crates/vm/src/builtins/set.rs @@ -232,6 +232,21 @@ impl PySetInner { Ok(res) } + fn intersection_multi( + &self, + mut others: impl core::iter::Iterator, + vm: &VirtualMachine, + ) -> PyResult { + let Some(other) = others.next() else { + return Ok(self.copy()); + }; + let mut result = self.intersection(other, vm)?; + for other in others { + result = result.intersection(other, vm)?; + } + Ok(result) + } + fn len(&self) -> usize { self.content.len() } @@ -252,8 +267,7 @@ 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. + /// Look up a key whose hash is already known, without a frozenset retry. fn contains_known_hash( &self, needle: &PyObject, @@ -272,12 +286,12 @@ impl PySetInner { } let (superset, subset) = match op { - PyComparisonOp::Lt | PyComparisonOp::Le => (other, self), + PyComparisonOp::Lt | PyComparisonOp::Le | PyComparisonOp::Eq => (other, self), _ => (self, other), }; - for key in subset.elements() { - if !superset.contains(&key, vm)? { + for (key, hash) in subset.content.keys_with_hashes() { + if !superset.contains_known_hash(&key, hash, vm)? { return Ok(false); } } @@ -301,10 +315,23 @@ 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) = Self::cached_hashes(other.as_object(), vm) { + if let Some(elements) = elements { for (obj, hash) in elements { - if self.contains_known_hash(&obj, hash, vm)? { + if target.contains_known_hash(&obj, hash, vm)? { set.add_known_hash(&obj, hash, vm)?; } } @@ -312,8 +339,12 @@ impl PySetInner { } for item in other.iter(vm)? { let obj = item?; - if self.contains(&obj, vm)? { - set.add(&obj, vm)?; + let hash = obj.hash(vm)?; + if self.contains_known_hash(&obj, hash, vm)? { + set.add_known_hash(&obj, hash, vm)?; + if set.len() >= self.len() { + break; + } } } Ok(set) @@ -361,6 +392,9 @@ impl PySetInner { } fn issuperset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { + if let Some(other_set) = extract_set(other.as_object()) { + return self.compare(other_set, PyComparisonOp::Ge, vm); + } for item in other.iter(vm)? { if !self.contains(&*item?, vm)? { return Ok(false); @@ -370,11 +404,32 @@ impl PySetInner { } fn issubset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { - let other_set = Self::from_iter(other.iter(vm)?, vm)?; - self.compare(&other_set, PyComparisonOp::Le, vm) + if let Some(other_set) = extract_set(other.as_object()) { + return self.compare(other_set, PyComparisonOp::Le, vm); + } + Ok(self.intersection(other, vm)?.len() == self.len()) } pub(super) fn isdisjoint(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult { + if let Some(other_set) = extract_set(other.as_object()) { + if core::ptr::eq(self, other_set) { + return Ok(self.len() == 0); + } + let other_type = other.as_object().class(); + if other_type.is(vm.ctx.types.set_type) || other_type.is(vm.ctx.types.frozenset_type) { + let (target, source) = if self.len() < other_set.len() { + (other_set, self) + } else { + (self, other_set) + }; + for (key, hash) in source.content.keys_with_hashes() { + if target.contains_known_hash(&key, hash, vm)? { + return Ok(false); + } + } + return Ok(true); + } + } for item in other.iter(vm)? { if self.contains(&*item?, vm)? { return Ok(false); @@ -486,7 +541,7 @@ impl PySetInner { others: impl core::iter::Iterator, vm: &VirtualMachine, ) -> PyResult<()> { - let temp_inner = self.fold_op(others, Self::intersection, vm)?; + 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)?; @@ -694,7 +749,9 @@ impl PySet { others: PosArgs, vm: &VirtualMachine, ) -> PyResult { - self.fold_op(others.into_iter(), PySetInner::intersection, vm) + Ok(Self { + inner: self.inner.intersection_multi(others.into_iter(), vm)?, + }) } #[pymethod] @@ -748,11 +805,9 @@ impl PySet { vm: &VirtualMachine, ) -> PyResult> { if let Ok(other) = AnySet::try_from_object(vm, other) { - Ok(PyArithmeticValue::Implemented(self.op( - other, - PySetInner::intersection, - vm, - )?)) + Ok(PyArithmeticValue::Implemented(Self { + inner: self.inner.intersection(other.into_iterable(vm)?, vm)?, + })) } else { Ok(PyArithmeticValue::NotImplemented) } @@ -1221,7 +1276,10 @@ impl PyFrozenSet { others: PosArgs, vm: &VirtualMachine, ) -> PyResult { - self.fold_op(others.into_iter(), PySetInner::intersection, vm) + Ok(Self { + inner: self.inner.intersection_multi(others.into_iter(), vm)?, + ..Default::default() + }) } #[pymethod] @@ -1275,11 +1333,10 @@ impl PyFrozenSet { vm: &VirtualMachine, ) -> PyResult> { if let Ok(other) = AnySet::try_from_object(vm, other) { - Ok(PyArithmeticValue::Implemented(self.op( - other, - PySetInner::intersection, - vm, - )?)) + Ok(PyArithmeticValue::Implemented(Self { + inner: self.inner.intersection(other.into_iterable(vm)?, vm)?, + ..Default::default() + })) } else { Ok(PyArithmeticValue::NotImplemented) } diff --git a/crates/vm/src/dict_inner.rs b/crates/vm/src/dict_inner.rs index ea857878946..b2cf0788609 100644 --- a/crates/vm/src/dict_inner.rs +++ b/crates/vm/src/dict_inner.rs @@ -485,9 +485,9 @@ impl Dict { /// 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. + /// `hash` must be computed for `key` or read from [`Self::keys_with_hashes`] + /// on a container holding this same key. A wrong hash lands the entry in a + /// bucket no lookup probes, silently losing the key. pub(crate) fn insert_known_hash( &self, vm: &VirtualMachine, @@ -1330,6 +1330,7 @@ pub trait DictKey { } fn key_hash(&self, vm: &VirtualMachine) -> PyResult; fn key_is(&self, other: &PyObject) -> bool; + /// Compare the stored entry (`other_key`) with this lookup key, in that order. fn key_eq(&self, vm: &VirtualMachine, other_key: &PyObject) -> PyResult; fn key_as_isize(&self, vm: &VirtualMachine) -> PyResult; } @@ -1369,7 +1370,7 @@ impl DictKey for PyObject { let b = unsafe { other_key.downcast_unchecked_ref::() }; return Ok(a.as_bigint() == b.as_bigint()); } - vm.bool_eq(self, other_key) + vm.bool_eq(other_key, self) } #[inline] @@ -1401,7 +1402,7 @@ impl DictKey for Py { } else if let Some(pystr) = other_key.downcast_ref_if_exact::(vm) { Ok(self.as_wtf8() == pystr.as_wtf8()) } else { - vm.bool_eq(self.as_object(), other_key) + vm.bool_eq(other_key, self.as_object()) } } @@ -1709,7 +1710,7 @@ impl DictKey for usize { } } else { let int = vm.ctx.new_int(*self); - vm.bool_eq(int.as_ref(), other_key) + vm.bool_eq(other_key, int.as_ref()) } } diff --git a/extra_tests/snippets/builtin_dict.py b/extra_tests/snippets/builtin_dict.py index bfda8a08b1c..41ab8522754 100644 --- a/extra_tests/snippets/builtin_dict.py +++ b/extra_tests/snippets/builtin_dict.py @@ -568,3 +568,50 @@ def __missing__(self, key): assert ("ports", [80, 443]) in lookup_items assert ("ports", [8080]) not in lookup_items assert ("missing", []) not in lookup_items + + +class LookupOrderKey: + calls = [] + + def __init__(self, name, equal): + self.name = name + self.equal = equal + + def __hash__(self): + return 17 + + def __eq__(self, other): + self.calls.append(self.name) + return self.equal + + +stored = LookupOrderKey("stored", True) +lookup = LookupOrderKey("lookup", False) +mapping = {stored: 1} +assert mapping[lookup] == 1 +mapping[lookup] = 2 +assert len(mapping) == 1 +assert next(iter(mapping)) is stored +del mapping[lookup] +assert not mapping +assert LookupOrderKey.calls == ["stored", "stored", "stored"] + +LookupOrderKey.calls.clear() +stored = LookupOrderKey("stored", NotImplemented) +lookup = LookupOrderKey("lookup", True) +assert {stored: 1}[lookup] == 1 +assert LookupOrderKey.calls == ["stored", "lookup"] + + +class LookupOrderSubclass(LookupOrderKey): + __hash__ = LookupOrderKey.__hash__ + + def __eq__(self, other): + return super().__eq__(other) + + +LookupOrderKey.calls.clear() +stored = LookupOrderKey("stored", False) +lookup = LookupOrderSubclass("lookup", True) +assert {stored: 1}[lookup] == 1 +assert LookupOrderKey.calls == ["lookup"] diff --git a/extra_tests/snippets/builtin_set.py b/extra_tests/snippets/builtin_set.py index 950875ea09a..a43cd5405aa 100644 --- a/extra_tests/snippets/builtin_set.py +++ b/extra_tests/snippets/builtin_set.py @@ -445,3 +445,124 @@ class FS(frozenset): assert repr(FS()) == "FS()" assert repr(FS([1, 2, 3])) == "FS({1, 2, 3})" + + +class StoredHashKey: + def __init__(self, value): + self.value = value + self.hash_enabled = True + self.hash_count = 0 + + def __hash__(self): + self.hash_count += 1 + assert self.hash_enabled, "set operation recomputed a stored hash" + return self.value + + +for left_type in (set, frozenset): + for right_type in (set, frozenset): + keys = [StoredHashKey(i) for i in range(3)] + small = left_type(keys[:1]) + equal = right_type(keys[:1]) + large = right_type(keys[:2]) + separate = right_type(keys[2:]) + for key in keys: + key.hash_enabled = False + assert small == equal + assert not (small != equal) + assert small < large + assert small <= large + assert large > small + assert large >= small + assert small.issubset(large) + assert large.issuperset(small) + assert small.isdisjoint(separate) + assert not small.isdisjoint(large) + + +for set_type in (set, frozenset): + + class IteratorOverride(set_type): + def __iter__(self): + raise RuntimeError("overridden iterator") + + subclass = IteratorOverride([1, 2]) + assert set_type([1]).issubset(subclass) + assert set_type([1, 2, 3]).issuperset(subclass) + assert set_type([1]).intersection(subclass) == {1} + assert not subclass.isdisjoint(subclass) + empty_subclass = IteratorOverride() + assert empty_subclass.isdisjoint(empty_subclass) + assert_raises(RuntimeError, set_type([3]).isdisjoint, subclass) + + +# Intersections retain the key from the smaller operand, or RHS on a tie. +small_key = float("1") +small = {small_key} +large = {1, 2} +assert next(iter(small & large)) is small_key +assert next(iter(large & small)) is small_key +assert next(iter({1} & small)) is small_key +result = large.intersection(small) +result.clear() +assert large == {1, 2} +assert small == {small_key} + + +def clear_during_intersection(): + live_source.clear() + yield 1 + + +live_source = {1, 2} +assert live_source.intersection(clear_during_intersection()) == set() + + +class DirectionalKey: + def __init__(self, equal): + self.equal = equal + + def __hash__(self): + return 17 + + def __eq__(self, other): + return self.equal + + +small_key = DirectionalKey(False) +large_key = DirectionalKey(True) +small = {small_key} +large = {large_key, 12345} +equal_size = {large_key} +assert next(iter(small & large)) is small_key +assert next(iter(large & small)) is small_key +assert not small.isdisjoint(large) +assert not large.isdisjoint(small) +assert small.issubset(large) +assert large.issuperset(small) +assert small == equal_size +assert equal_size != small +assert not small.issubset([large_key]) + + +def matched_then_error(): + yield 1 + raise ValueError("iterator consumed after matching") + + +for set_type in (set, frozenset): + assert set_type([1]).intersection(matched_then_error()) == {1} + assert set_type([1]).issubset(matched_then_error()) + assert_raises(ValueError, set_type([1, 2]).intersection, matched_then_error()) + assert_raises(ValueError, set_type([1, 2]).issubset, matched_then_error()) + assert_raises(ValueError, set_type().intersection, matched_then_error()) + assert_raises(TypeError, set_type().intersection, [set()]) + +key = StoredHashKey(4) +source = {key} +key.hash_count = 0 +assert len(source.intersection([key])) == 1 +assert key.hash_count == 1 +dictionary = {key: None} +key.hash_enabled = False +assert_raises(AssertionError, source.intersection, dictionary)