Specialize list.sort() comparisons for homogeneous lists - #8463

crates/vm/src/sorting.rs (3)

232-237: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Rename b_wins in merge_hi; the name states the opposite of the branch action.

The comparison is is_lt(buf[cursor_b], values[cursor_a]), so a true result means the A element is the larger one and the code copies from A. The name b_wins therefore contradicts the branch body. In merge_lo the same name is correct, which makes the mismatch easy to misread during future maintenance. The behavior matches CPython's merge_hi; only the name is wrong.

♻️ Proposed rename

-                let b_wins = match is_lt(&self.buf[cursor_b], &values[cursor_a]) {
+                let a_wins = match is_lt(&self.buf[cursor_b], &values[cursor_a]) {
                     Ok(v) => v,
                     Err(e) => break 'merging Err(e),
                 };
-                if b_wins {
+                if a_wins {

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/vm/src/sorting.rs` around lines 232 - 237, In merge_hi, rename the
boolean local b_wins to a name reflecting that a true comparison selects the A
element, and update its corresponding conditional reference; leave merge_lo’s
b_wins and all merge behavior unchanged.

---

1-3: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove or rewrite the stale TODO at the top of the file.

MERGESTATE_TEMP_SIZE does not exist in this file or in this crate. The comment refers to a CPython implementation detail that this port intentionally replaced with a dynamic Vec. As written, the note describes a symbol that a reader cannot find.

Do you want me to open an issue to track a small-list stack buffer optimization instead?

♻️ Proposed cleanup

-// TODO: MERGESTATE_TEMP_SIZE unused — buf is a dynamic Vec, not a fixed stack array.
 const MIN_GALLOP: usize = 7;
 const MAX_MINRUN: usize = 64;

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/vm/src/sorting.rs` around lines 1 - 3, Remove the stale TODO preceding
MIN_GALLOP and MAX_MINRUN in sorting.rs; do not replace it with a note about
MERGESTATE_TEMP_SIZE, since this implementation intentionally uses a dynamic Vec
and no corresponding symbol exists.

Source: Coding guidelines

---

696-738: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add tests for comparator errors and for stability.

The current tests only use an infallible comparator (Ok::<bool, ()>) and only assert final ordering of i32 values. Two behaviors that this engine guarantees are untested:

1. Error propagation. The Err arms in merge_lo (lines 175-181) and merge_hi (lines 350-354) copy the buffered run back into values before returning the error. No test reaches those arms. list.sort() relies on this path whenever a Python __lt__ raises, and a mistake there would silently corrupt list contents rather than fail loudly.
2. Stability. list.sort() and the specialized comparators in crates/vm/src/builtins/list.rs depend on stable ordering. No test asserts that equal elements keep their input order.

Both need inputs of at least 64 elements to exercise the merge paths rather than the binary-insertion shortcut.

✅ Proposed additional tests

#[test]
fn stability_preserved() {
    // (key, original_index) pairs; compare on key only.
    let mut v: Vec<(i32, usize)> = (0..200).map(|i| ((i % 5) as i32, i)).collect();
    timsort(&mut v, &mut |a, b| Ok::<bool, ()>(a.0 < b.0)).unwrap();
    let mut expected = v.clone();
    expected.sort_by_key(|p| p.0); // std sort_by_key is stable
    assert_eq!(v, expected);
}

#[test]
fn comparator_error_propagates_and_preserves_elements() {
    let mut v: Vec<i32> = (0..200).rev().collect();
    let original: Vec<i32> = v.clone();
    let mut calls = 0usize;
    let res = timsort(&mut v, &mut |a: &i32, b: &i32| {
        calls += 1;
        if calls > 150 { Err(()) } else { Ok(a < b) }
    });
    assert!(res.is_err());
    // The slice must still hold every original element exactly once.
    let mut got = v.clone();
    let mut want = original;
    got.sort();
    want.sort();
    assert_eq!(got, want);
}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/vm/src/sorting.rs` around lines 696 - 738, Add tests in the existing
sorting tests module for stability and comparator-error propagation. Use inputs
of at least 64 elements: sort key/index pairs with a key-only comparator and
verify equal-key elements retain input order; separately use a comparator that
returns Err after several calls, assert timsort returns an error, and compare
sorted copies to confirm the output still contains every original element
exactly once.

crates/vm/src/builtins/list.rs (1)

783-797: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Use the specialized comparator for the first-element equality test.

Line 785 always calls the generic rich_compare_bool(Eq), including for i == 0. The specialization at line 794 then applies only after that generic call has already run. For the common case of two tuples that differ at index 0, every comparison pays one full generic Python-level Eq dispatch, which is the cost the tuple specialization is meant to avoid.

CPython's unsafe_tuple_compare uses the cached specialized rich-comparison for the index-0 equality test as well, then the specialized less-than. Matching that recovers the intended speedup. This is a performance gap, not a correctness bug; results are unchanged.

Add an elem_eq counterpart to elem_lt and use it at index 0.

⚡ Proposed change

 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)? {
+        let eq = if i == 0 {
+            elem_eq(elem, &a[0], &b[0], vm)?
+        } else {
+            a[i].rich_compare_bool(&b[i], PyComparisonOp::Eq, vm)?
+        };
+        if !eq {
             break;
         }
         i += 1;
     }

Add the counterpart helper next to elem_lt:

fn elem_eq(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
    match elem {
        Elem::Str => Ok(a.downcast_ref::<PyStr>().unwrap().as_bytes()
            == b.downcast_ref::<PyStr>().unwrap().as_bytes()),
        Elem::Int => Ok(a.downcast_ref::<PyInt>().unwrap().as_bigint()
            == b.downcast_ref::<PyInt>().unwrap().as_bigint()),
        Elem::Float => Ok(a.downcast_ref::<PyFloat>().unwrap().to_f64()
            == b.downcast_ref::<PyFloat>().unwrap().to_f64()),
        Elem::Object(_) | Elem::Generic => a.rich_compare_bool(b, PyComparisonOp::Eq, vm),
    }
}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/vm/src/builtins/list.rs` around lines 783 - 797, Add an elem_eq helper
alongside elem_lt, using the specialized Str, Int, and Float comparisons and
generic rich_compare_bool for Object and Generic elements. Update the tuple
comparison loop to use elem_eq for the first-element equality check (i == 0),
while preserving generic equality for subsequent elements and the existing
less-than logic.

extra_tests/snippets/builtin_list.py (1)

245-264: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for reverse=True, for key=, and for inputs larger than 64 elements.

The new assertions cover element types well, but they miss three paths that this PR introduces:

1. reverse=True. Each of the six arms in timsort_specialized applies its own operand swap. A missing swap in one arm would pass every assertion here.
2. key=. The keyed path instantiates timsort_specialized over (PyObjectRef, PyObjectRef) and classifies the key results rather than the elements. No assertion exercises it.
3. Input size. Every list here has fewer than 64 elements, so timsort always takes the binary-insertion shortcut. merge_lo, merge_hi, and the galloping code are never reached from Python.

A combined reverse-plus-key case on a list of several hundred elements also verifies stability, which reverse ordering depends on.

✅ Proposed additional assertions

# reverse=True across each specialization
assert sorted([3, 1, 2], reverse=True) == [3, 2, 1]
assert sorted(["b", "a", "c"], reverse=True) == ["c", "b", "a"]
assert sorted([2.5, 1.5, 3.5], reverse=True) == [3.5, 2.5, 1.5]
assert sorted([b"b", b"a"], reverse=True) == [b"b", b"a"]
assert sorted([(1, 2), (1, 1), (0, 9)], reverse=True) == [(1, 2), (1, 1), (0, 9)]

# key= scans key results, not elements
assert sorted(["bbb", "a", "cc"], key=len) == ["a", "cc", "bbb"]
assert sorted([1, 2, 3], key=lambda x: -x) == [3, 2, 1]
assert sorted(["bbb", "a", "cc"], key=len, reverse=True) == ["bbb", "cc", "a"]

# inputs above the 64-element insertion-sort shortcut
big = [(i * 7919) % 500 for i in range(500)]
assert sorted(big) == list(range(500))
assert sorted(big, reverse=True) == list(range(499, -1, -1))
assert sorted([str(i) for i in big]) == sorted([str(i) for i in big])

# stability: equal keys keep input order, including with reverse=True
pairs = [(i % 5, i) for i in range(200)]
assert sorted(pairs, key=lambda p: p[0]) == sorted(pairs, key=lambda p: p[0])
by_key = sorted(pairs, key=lambda p: p[0])
assert [p[1] for p in by_key if p[0] == 0] == list(range(0, 200, 5))
rev = sorted(pairs, key=lambda p: p[0], reverse=True)
assert [p[1] for p in rev if p[0] == 0] == list(range(0, 200, 5))

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@extra_tests/snippets/builtin_list.py` around lines 245 - 264, Expand the
sorted() assertions around the existing builtin sorting coverage to exercise
reverse=True across supported specializations, key= with and without reverse,
and lists larger than 64 elements that reach merge and galloping paths. Add a
combined large reverse-plus-key case that verifies stable ordering of equal
keys, preserving input order within each key group.

Source: Coding guidelines

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes