Skip to content

Commit 833a2ba

Browse files
youknowoneyouknowone
authored
str, bytes: answer equality with equality rather than with an ordering (#8531)
PyStr's comparison, PyBytesInner's, and the specialized CompareOpStr instruction all answered == and != by taking Ord::cmp of the two buffers and asking whether the result was Equal. An ordering has to read the bytes: it memcmps the common prefix even where the lengths already settle the question. CompareOpStr bypasses the Comparable slot, so it had also lost the identity shortcut that slot takes, and a string compared with itself was read end to end. Add PyComparisonOp::eval_eq, which settles Eq and Ne from an equality test and leaves an ordering operator to the caller, and answer through it in the three places: slice equality checks the length first, and CompareOpStr answers an object compared with itself the way the slot it specializes does. n=1,000,000, per comparison: before after s == s (the very same object) 23.21us 0.16us s == a string one shorter 24.63us 0.17us b == bytes one shorter 25.28us 0.20us ba == bytearray one shorter 27.68us 0.23us s == an equal, distinct string 23.78us 24.50us s < an equal string 29.87us 25.07us Assisted-by: Claude
1 parent 5919be9 commit 833a2ba

5 files changed

Lines changed: 87 additions & 4 deletions

File tree

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1592,6 +1592,11 @@ impl Comparable for PyStr {
15921592
return Ok(res.into());
15931593
}
15941594
let other = class_or_notimplemented!(Self, other);
1595+
// Equality does not need the ordering, and answers two strings of
1596+
// different length without reading either.
1597+
if let Some(res) = op.eval_eq(|| zelf.as_wtf8() == other.as_wtf8()) {
1598+
return Ok(res.into());
1599+
}
15951600
Ok(op.eval_ord(zelf.as_wtf8().cmp(other.as_wtf8())).into())
15961601
}
15971602
}

‎crates/vm/src/bytes_inner.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -345,7 +345,12 @@ impl PyBytesInner {
345345
// but not memoryview, and not equal if compare with unicode str(PyStr)
346346
PyComparisonValue::from_option(
347347
other
348-
.try_bytes_like(vm, |other| op.eval_ord(self.elements.as_slice().cmp(other)))
348+
.try_bytes_like(vm, |other| {
349+
// Equality does not need the ordering, and answers two
350+
// buffers of different length without reading either.
351+
op.eval_eq(|| self.elements.as_slice() == other)
352+
.unwrap_or_else(|| op.eval_ord(self.elements.as_slice().cmp(other)))
353+
})
349354
.ok(),
350355
)
351356
}

‎crates/vm/src/frame.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6894,10 +6894,13 @@ impl ExecutingFrame<'_> {
68946894
b.downcast_ref_if_exact::<PyStr>(vm),
68956895
) {
68966896
let op = self.compare_op_from_arg(arg);
6897-
if op != PyComparisonOp::Eq && op != PyComparisonOp::Ne {
6897+
// The same two shortcuts the unspecialized comparison takes:
6898+
// one object is equal to itself, and equality answers two
6899+
// strings of different length without reading either.
6900+
let Some(result) = op.eval_eq(|| a.is(b) || a_str.as_wtf8() == b_str.as_wtf8())
6901+
else {
68986902
return self.execute_compare(vm, arg);
6899-
}
6900-
let result = op.eval_ord(a_str.as_wtf8().cmp(b_str.as_wtf8()));
6903+
};
69016904
self.pop_value();
69026905
self.pop_value();
69036906
self.push_value(vm.ctx.new_bool(result).into());

‎crates/vm/src/types/slot.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1995,6 +1995,29 @@ impl PyComparisonOp {
19951995
self.map_eq(|| a.borrow().is(b.borrow()))
19961996
}
19971997

1998+
/// The answer to this comparison for two operands that `equal` reports as
1999+
/// equal or not, or `None` for an ordering operator, which equality alone
2000+
/// cannot settle -- `equal` is not called in that case.
2001+
///
2002+
/// This is what lets a type answer `==` and `!=` with an equality test
2003+
/// rather than with an ordering: the two agree on the answer, but equality
2004+
/// can settle a length mismatch without looking at the contents at all.
2005+
///
2006+
/// The two neighbouring helpers answer different questions: [`Self::map_eq`]
2007+
/// answers only where its predicate holds, so a caller still handles the
2008+
/// other side, and [`Self::eq_only`] declares the comparison
2009+
/// `NotImplemented` for an ordering operator. This one leaves the ordering
2010+
/// operators to the caller, which is what a type with a real ordering
2011+
/// needs.
2012+
#[inline]
2013+
pub fn eval_eq(self, equal: impl FnOnce() -> bool) -> Option<bool> {
2014+
match self {
2015+
Self::Eq => Some(equal()),
2016+
Self::Ne => Some(!equal()),
2017+
_ => None,
2018+
}
2019+
}
2020+
19982021
/// Returns `Some(true)` when self is `Eq` and `f()` returns true. Returns `Some(false)` when self
19992022
/// is `Ne` and `f()` returns true. Otherwise returns `None`.
20002023
#[inline]

‎extra_tests/snippets/operator_comparison.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,3 +87,50 @@ def test_type_error(x, y):
8787
assert not math.nan < 123
8888
assert not math.nan >= 123
8989
assert not math.nan <= 123
90+
91+
92+
# str and bytes comparisons, through a function so that the operands are not
93+
# constants the compiler can fold, and in a loop so the specialized comparison
94+
# is reached.
95+
def cmp_all(a, b):
96+
return (a == b, a != b, a < b, a <= b, a > b, a >= b)
97+
98+
99+
def check(a, b, expected):
100+
for _ in range(200):
101+
assert cmp_all(a, b) == expected, (a, b, cmp_all(a, b), expected)
102+
103+
104+
EQ = (True, False, False, True, False, True)
105+
LT = (False, True, True, True, False, False)
106+
GT = (False, True, False, False, True, True)
107+
108+
same = "abc" * 3
109+
check(same, same, EQ) # the very same object
110+
check(same, "abcabcabc", EQ) # equal, distinct objects
111+
check("abc", "abd", LT) # same length, differing content
112+
check("abc", "abcd", LT) # a prefix is less than what extends it
113+
check("abcd", "abc", GT)
114+
check("", "a", LT)
115+
check("", "", EQ)
116+
check("\ud800", "\ud800", EQ) # lone surrogates are compared as themselves
117+
check("\ud800", "\udfff", LT)
118+
check("a\U0001f600", "a\U0001f600", EQ)
119+
check("가나다", "가나다", EQ)
120+
check("가나", "가나다", LT)
121+
122+
# Comparing with a non-string is never an error for == and !=.
123+
assert not "abc" == 3
124+
assert "abc" != 3
125+
126+
bsame = b"abc" * 3
127+
check(bsame, bsame, EQ)
128+
check(bsame, b"abcabcabc", EQ)
129+
check(b"abc", b"abd", LT)
130+
check(b"abc", b"abcd", LT)
131+
check(b"abcd", b"abc", GT)
132+
check(bytearray(b"abc"), bytearray(b"abcd"), LT)
133+
check(bytearray(b"abc"), b"abc", EQ) # bytearray and bytes compare by content
134+
check(b"abc", bytearray(b"abd"), LT)
135+
assert not b"abc" == "abc"
136+
assert b"abc" != "abc"

0 commit comments

Comments
 (0)