Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions crates/vm/src/builtins/str.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1563,6 +1563,11 @@ impl Comparable for PyStr {
return Ok(res.into());
}
let other = class_or_notimplemented!(Self, other);
// Equality does not need the ordering, and answers two strings of
// different length without reading either.
if let Some(res) = op.eval_eq(|| zelf.as_wtf8() == other.as_wtf8()) {
return Ok(res.into());
}
Ok(op.eval_ord(zelf.as_wtf8().cmp(other.as_wtf8())).into())
}
}
Expand Down
7 changes: 6 additions & 1 deletion crates/vm/src/bytes_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -345,7 +345,12 @@ impl PyBytesInner {
// but not memoryview, and not equal if compare with unicode str(PyStr)
PyComparisonValue::from_option(
other
.try_bytes_like(vm, |other| op.eval_ord(self.elements.as_slice().cmp(other)))
.try_bytes_like(vm, |other| {
// Equality does not need the ordering, and answers two
// buffers of different length without reading either.
op.eval_eq(|| self.elements.as_slice() == other)
.unwrap_or_else(|| op.eval_ord(self.elements.as_slice().cmp(other)))
})
.ok(),
)
}
Expand Down
9 changes: 6 additions & 3 deletions crates/vm/src/frame.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6894,10 +6894,13 @@ impl ExecutingFrame<'_> {
b.downcast_ref_if_exact::<PyStr>(vm),
) {
let op = self.compare_op_from_arg(arg);
if op != PyComparisonOp::Eq && op != PyComparisonOp::Ne {
// The same two shortcuts the unspecialized comparison takes:
// one object is equal to itself, and equality answers two
// strings of different length without reading either.
let Some(result) = op.eval_eq(|| a.is(b) || a_str.as_wtf8() == b_str.as_wtf8())
else {
return self.execute_compare(vm, arg);
}
let result = op.eval_ord(a_str.as_wtf8().cmp(b_str.as_wtf8()));
};
self.pop_value();
self.pop_value();
self.push_value(vm.ctx.new_bool(result).into());
Expand Down
23 changes: 23 additions & 0 deletions crates/vm/src/types/slot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1995,6 +1995,29 @@ impl PyComparisonOp {
self.map_eq(|| a.borrow().is(b.borrow()))
}

/// The answer to this comparison for two operands that `equal` reports as
/// equal or not, or `None` for an ordering operator, which equality alone
/// cannot settle -- `equal` is not called in that case.
///
/// This is what lets a type answer `==` and `!=` with an equality test
/// rather than with an ordering: the two agree on the answer, but equality
/// can settle a length mismatch without looking at the contents at all.
///
/// The two neighbouring helpers answer different questions: [`Self::map_eq`]
/// answers only where its predicate holds, so a caller still handles the
/// other side, and [`Self::eq_only`] declares the comparison
/// `NotImplemented` for an ordering operator. This one leaves the ordering
/// operators to the caller, which is what a type with a real ordering
/// needs.
#[inline]
pub fn eval_eq(self, equal: impl FnOnce() -> bool) -> Option<bool> {
match self {
Self::Eq => Some(equal()),
Self::Ne => Some(!equal()),
_ => None,
}
}

/// Returns `Some(true)` when self is `Eq` and `f()` returns true. Returns `Some(false)` when self
/// is `Ne` and `f()` returns true. Otherwise returns `None`.
#[inline]
Expand Down
47 changes: 47 additions & 0 deletions extra_tests/snippets/operator_comparison.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,50 @@ def test_type_error(x, y):
assert not math.nan < 123
assert not math.nan >= 123
assert not math.nan <= 123


# str and bytes comparisons, through a function so that the operands are not
# constants the compiler can fold, and in a loop so the specialized comparison
# is reached.
def cmp_all(a, b):
return (a == b, a != b, a < b, a <= b, a > b, a >= b)


def check(a, b, expected):
for _ in range(200):
assert cmp_all(a, b) == expected, (a, b, cmp_all(a, b), expected)


EQ = (True, False, False, True, False, True)
LT = (False, True, True, True, False, False)
GT = (False, True, False, False, True, True)

same = "abc" * 3
check(same, same, EQ) # the very same object
check(same, "abcabcabc", EQ) # equal, distinct objects
check("abc", "abd", LT) # same length, differing content
check("abc", "abcd", LT) # a prefix is less than what extends it
check("abcd", "abc", GT)
check("", "a", LT)
check("", "", EQ)
check("\ud800", "\ud800", EQ) # lone surrogates are compared as themselves
check("\ud800", "\udfff", LT)
check("a\U0001f600", "a\U0001f600", EQ)
check("가나다", "가나다", EQ)
check("가나", "가나다", LT)

# Comparing with a non-string is never an error for == and !=.
assert not "abc" == 3
assert "abc" != 3

bsame = b"abc" * 3
check(bsame, bsame, EQ)
check(bsame, b"abcabcabc", EQ)
check(b"abc", b"abd", LT)
check(b"abc", b"abcd", LT)
check(b"abcd", b"abc", GT)
check(bytearray(b"abc"), bytearray(b"abcd"), LT)
check(bytearray(b"abc"), b"abc", EQ) # bytearray and bytes compare by content
check(b"abc", bytearray(b"abd"), LT)
assert not b"abc" == "abc"
assert b"abc" != "abc"
Loading