Specialize list.sort() comparisons for homogeneous lists (#8464)

* Use type-specialized comparators for homogeneous list sorts

Scan the sort keys once before sorting; when every element is exactly str, int, or float, compare wtf8 bytes / BigInt / f64 directly instead of going through rich_compare_bool dispatch, mirroring CPython's pre-sort check in listsort.c (unsafe_latin_compare and friends). Subclasses and mixed-type lists keep the generic __lt__ path.

* Add edge-case tests for specialized list sorts

* Cache the richcompare slot for homogeneous list sorts

* Specialize tuple sorts on their first elements

* Apply ruff formatting to the new list sort tests

* Deduplicate the reverse swap across sort dispatch arms

crates/vm/src/builtins/list.rs

@@ -9,10 +9,10 @@ use crate::common::lock::{
 use crate::object::{Traverse, TraverseFn};
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
-    builtins::PyStr,
+    builtins::{PyFloat, PyInt, PyStr, PyTuple},
     class::PyClassImpl,
     convert::ToPyObject,
-    function::{ArgSize, FuncArgs, OptionalArg, PyComparisonValue},
+    function::{ArgSize, Either, FuncArgs, OptionalArg, PyComparisonValue},
     iter::PyExactSizeIterator,
     protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
     recursion::ReprGuard,
@@ -21,7 +21,7 @@ use crate::{
     sorting::timsort,
     types::{
         AsMapping, AsSequence, Comparable, Constructor, Initializer, IterNext, Iterable,
-        PyComparisonOp, Representable, SelfIter,
+        PyComparisonOp, Representable, RichCompareFunc, SelfIter,
     },
     vm::VirtualMachine,
 };
@@ -638,34 +638,226 @@ impl Representable for PyList {
     }
 }
 
+enum Elem {
+    Str,
+    Int,
+    Float,
+    Object(RichCompareFunc),
+    Generic,
+}
+
+enum PreSort {
+    Str,
+    Int,
+    Float,
+    Object(RichCompareFunc),
+    Tuple(Elem),
+    Generic,
+}
+
+impl From<Elem> for PreSort {
+    fn from(e: Elem) -> Self {
+        match e {
+            Elem::Str => Self::Str,
+            Elem::Int => Self::Int,
+            Elem::Float => Self::Float,
+            Elem::Object(f) => Self::Object(f),
+            Elem::Generic => Self::Generic,
+        }
+    }
+}
+
+fn classify(class: &Py<PyType>, vm: &VirtualMachine) -> Elem {
+    if class.is(vm.ctx.types.str_type) {
+        Elem::Str
+    } else if class.is(vm.ctx.types.int_type) {
+        Elem::Int
+    } else if class.is(vm.ctx.types.float_type) {
+        Elem::Float
+    } else if let Some(f) = class.slots.richcompare.load() {
+        Elem::Object(f)
+    } else {
+        Elem::Generic
+    }
+}
+
+fn pre_sort_check<'a>(
+    mut keys: impl Iterator<Item = &'a PyObjectRef>,
+    vm: &VirtualMachine,
+) -> PreSort {
+    let Some(first) = keys.next() else {
+        return PreSort::Generic;
+    };
+
+    if let Some(t) = first
+        .downcast_ref_if_exact::<PyTuple>(vm)
+        .filter(|t| !t.as_slice().is_empty())
+    {
+        pre_sort_check_tuples(&t.as_slice()[0], keys, vm)
+    } else {
+        let class = first.class();
+        if keys.all(|k| k.class().is(class)) {
+            classify(class, vm).into()
+        } else {
+            PreSort::Generic
+        }
+    }
+}
+
+fn pre_sort_check_tuples<'a>(
+    first_elem: &PyObjectRef,
+    keys: impl Iterator<Item = &'a PyObjectRef>,
+    vm: &VirtualMachine,
+) -> PreSort {
+    let class = first_elem.class();
+    let mut all_same_type = true;
+
+    for k in keys {
+        let Some(t) = k
+            .downcast_ref_if_exact::<PyTuple>(vm)
+            .filter(|t| !t.as_slice().is_empty())
+        else {
+            return PreSort::Generic;
+        };
+        if all_same_type && !t.as_slice()[0].class().is(class) {
+            all_same_type = false;
+        }
+    }
+
+    let elem = if !all_same_type || class.is(vm.ctx.types.tuple_type) {
+        Elem::Generic
+    } else {
+        classify(class, vm)
+    };
+    PreSort::Tuple(elem)
+}
+
+fn str_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool {
+    a.downcast_ref::<PyStr>().unwrap().as_bytes() < b.downcast_ref::<PyStr>().unwrap().as_bytes()
+}
+
+fn int_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool {
+    a.downcast_ref::<PyInt>().unwrap().as_bigint() < b.downcast_ref::<PyInt>().unwrap().as_bigint()
+}
+
+fn float_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool {
+    a.downcast_ref::<PyFloat>().unwrap().to_f64() < b.downcast_ref::<PyFloat>().unwrap().to_f64()
+}
+
+fn object_lt(
+    cmp: RichCompareFunc,
+    a: &PyObjectRef,
+    b: &PyObjectRef,
+    vm: &VirtualMachine,
+) -> PyResult<bool> {
+    #[allow(unpredictable_function_pointer_comparisons)]
+    if a.class().slots.richcompare.load() != Some(cmp) {
+        return a.rich_compare_bool(b, PyComparisonOp::Lt, vm);
+    }
+    match cmp(a, b, PyComparisonOp::Lt, vm)? {
+        Either::B(PyComparisonValue::Implemented(v)) => Ok(v),
+        Either::B(PyComparisonValue::NotImplemented) => {
+            a.rich_compare_bool(b, PyComparisonOp::Lt, vm)
+        }
+        Either::A(obj) => {
+            if obj.is(&vm.ctx.not_implemented) {
+                a.rich_compare_bool(b, PyComparisonOp::Lt, vm)
+            } else {
+                obj.try_to_bool(vm)
+            }
+        }
+    }
+}
+
+fn elem_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
+    match elem {
+        Elem::Str => Ok(str_lt(a, b)),
+        Elem::Int => Ok(int_lt(a, b)),
+        Elem::Float => Ok(float_lt(a, b)),
+        Elem::Object(f) => object_lt(*f, a, b, vm),
+        Elem::Generic => a.rich_compare_bool(b, PyComparisonOp::Lt, vm),
+    }
+}
+
+fn tuple_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
+    let a = a.downcast_ref::<PyTuple>().unwrap().as_slice();
+    let b = b.downcast_ref::<PyTuple>().unwrap().as_slice();
+
+    let mut i = 0;
+    while i < a.len() && i < b.len() {
+        if !a[i].rich_compare_bool(&b[i], PyComparisonOp::Eq, vm)? {
+            break;
+        }
+        i += 1;
+    }
+    if i >= a.len() || i >= b.len() {
+        return Ok(a.len() < b.len());
+    }
+    if i == 0 {
+        elem_lt(elem, &a[0], &b[0], vm)
+    } else {
+        a[i].rich_compare_bool(&b[i], PyComparisonOp::Lt, vm)
+    }
+}
+
+fn timsort_by<T, K, L>(items: &mut [T], reverse: bool, key: &K, mut lt: L) -> PyResult<()>
+where
+    T: Clone,
+    K: Fn(&T) -> &PyObjectRef,
+    L: FnMut(&PyObjectRef, &PyObjectRef) -> PyResult<bool>,
+{
+    timsort(items, &mut |a, b| {
+        let (a, b) = if reverse {
+            (key(b), key(a))
+        } else {
+            (key(a), key(b))
+        };
+        lt(a, b)
+    })
+}
+
+fn timsort_specialized<T, K>(
+    vm: &VirtualMachine,
+    items: &mut [T],
+    reverse: bool,
+    key: K,
+) -> PyResult<()>
+where
+    T: Clone,
+    K: Fn(&T) -> &PyObjectRef,
+{
+    match pre_sort_check(items.iter().map(&key), vm) {
+        PreSort::Str => timsort_by(items, reverse, &key, |a, b| Ok(str_lt(a, b))),
+        PreSort::Int => timsort_by(items, reverse, &key, |a, b| Ok(int_lt(a, b))),
+        PreSort::Float => timsort_by(items, reverse, &key, |a, b| Ok(float_lt(a, b))),
+        PreSort::Object(cmp) => timsort_by(items, reverse, &key, |a, b| object_lt(cmp, a, b, vm)),
+        PreSort::Tuple(elem) => timsort_by(items, reverse, &key, |a, b| tuple_lt(&elem, a, b, vm)),
+        PreSort::Generic => timsort_by(items, reverse, &key, |a, b| {
+            a.rich_compare_bool(b, PyComparisonOp::Lt, vm)
+        }),
+    }
+}
+
 fn do_sort(
     vm: &VirtualMachine,
     values: &mut Vec<PyObjectRef>,
     key_func: Option<PyObjectRef>,
     reverse: bool,
 ) -> PyResult<()> {
-    // CPython uses __lt__ for all comparisons in sort.
-    // `timsort` expects is_lt(a, b) = true when a must be placed BEFORE b.
-    // For reverse=True, swapping the operands yields a descending order that is
-    // still stable in the original relative order, matching CPython's
-    // reverse-sort-reverse approach.
-    let mut is_lt = |a: &PyObjectRef, b: &PyObjectRef| {
-        if reverse {
-            b.rich_compare_bool(a, PyComparisonOp::Lt, vm)
-        } else {
-            a.rich_compare_bool(b, PyComparisonOp::Lt, vm)
-        }
-    };
-
     if let Some(ref key_func) = key_func {
         let mut items = values
             .iter()
             .map(|x| Ok((x.clone(), key_func.call((x.clone(),), vm)?)))
             .collect::<Result<Vec<_>, _>>()?;
-        timsort(&mut items, &mut |a, b| is_lt(&a.1, &b.1))?;
+        timsort_specialized(
+            vm,
+            &mut items,
+            reverse,
+            |item: &(PyObjectRef, PyObjectRef)| &item.1,
+        )?;
         *values = items.into_iter().map(|(val, _)| val).collect();
     } else {
-        timsort(values, &mut is_lt)?;
+        timsort_specialized(vm, values, reverse, |x: &PyObjectRef| x)?
     }
 
     Ok(())

extra_tests/snippets/builtin_list.py

@@ -242,6 +242,33 @@ assert sorted([(1, 2, 3), (0, 3, 6)], key=lambda x: x[0]) == [(0, 3, 6), (1, 2,
 assert sorted([(1, 2, 3), (0, 3, 6)], key=lambda x: x[1]) == [(1, 2, 3), (0, 3, 6)]
 assert sorted([(1, 2), (), (5,)], key=len) == [(), (5,), (1, 2)]
 
+assert sorted(["b", "a", "é", "z\U0001f600", "z"]) == [
+    "a",
+    "b",
+    "z",
+    "z\U0001f600",
+    "é",
+]
+assert sorted([10**30, -(10**30), 5, 0]) == [-(10**30), 0, 5, 10**30]
+assert sorted([True, False, True]) == [False, True, True]
+
+
+class IntSub(int):
+    pass
+
+
+assert sorted([IntSub(2), 3, IntSub(1)]) == [1, 2, 3]
+assert sorted([2.5, 1, 3.0, 2]) == [1, 2, 2.5, 3.0]
+assert_raises(TypeError, sorted, [1, "a"])
+nan = float("nan")
+assert repr(sorted([nan, 1.0, 2.0])) == "[nan, 1.0, 2.0]"
+assert sorted([b"b", b"a", b"c"]) == [b"a", b"b", b"c"]
+assert sorted([(2, 9), (1, 5), (2, 1)]) == [(1, 5), (2, 1), (2, 9)]
+assert sorted([(1, "b"), (1, "a")]) == [(1, "a"), (1, "b")]
+assert sorted([(1,), (1, 2), ()]) == [(), (1,), (1, 2)]
+assert sorted([((2,), "x"), ((1,), "y")]) == [((1,), "y"), ((2,), "x")]
+assert sorted([(1, "a"), (2.5, "b"), (0, "c")]) == [(0, "c"), (1, "a"), (2.5, "b")]
+
 lst = [3, 1, 5, 2, 4]