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

crates/vm/src/builtins/str.rs

@@ -1592,6 +1592,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())
     }
 }

crates/vm/src/bytes_inner.rs

@@ -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(),
         )
     }

crates/vm/src/frame.rs

@@ -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());

crates/vm/src/types/slot.rs

@@ -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]

extra_tests/snippets/operator_comparison.py

@@ -87,3 +87,50 @@ assert not math.nan > 123
 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"