From d1bb889e829316bf16cc62d025464bada43acd86 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 09:50:31 +0900 Subject: [PATCH 1/6] 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. --- crates/vm/src/builtins/list.rs | 111 ++++++++++++++++++++++++++++----- 1 file changed, 95 insertions(+), 16 deletions(-) diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index 2b83a5c1007..a396bc263da 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -9,7 +9,7 @@ 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}, class::PyClassImpl, convert::ToPyObject, function::{ArgSize, FuncArgs, OptionalArg, PyComparisonValue}, @@ -638,34 +638,113 @@ impl Representable for PyList { } } +enum PreSort { + Str, + Int, + Float, + Generic, +} + +fn pre_sort_check<'a>( + mut keys: impl Iterator, + vm: &VirtualMachine, +) -> PreSort { + let Some(first) = keys.next() else { + return PreSort::Generic; + }; + let class = first.class(); + if !keys.all(|o| o.class().is(class)) { + return PreSort::Generic; + } + if class.is(vm.ctx.types.str_type) { + PreSort::Str + } else if class.is(vm.ctx.types.int_type) { + PreSort::Int + } else if class.is(vm.ctx.types.float_type) { + PreSort::Float + } else { + PreSort::Generic + } +} + +fn str_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool { + a.downcast_ref::().unwrap().as_bytes() < b.downcast_ref::().unwrap().as_bytes() +} + +fn int_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool { + a.downcast_ref::().unwrap().as_bigint() < b.downcast_ref::().unwrap().as_bigint() +} + +fn float_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool { + a.downcast_ref::().unwrap().to_f64() < b.downcast_ref::().unwrap().to_f64() +} + +fn timsort_specialized( + 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(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + Ok(str_lt(a, b)) + }), + PreSort::Int => timsort(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + Ok(int_lt(a, b)) + }), + PreSort::Float => timsort(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + Ok(float_lt(a, b)) + }), + PreSort::Generic => timsort(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + a.rich_compare_bool(b, PyComparisonOp::Lt, vm) + }), + } +} + fn do_sort( vm: &VirtualMachine, values: &mut Vec, key_func: Option, 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::, _>>()?; - 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(()) From 03924c290e504b56d8906ffe0448860270839e3c Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:26:33 +0900 Subject: [PATCH 2/6] Add edge-case tests for specialized list sorts --- extra_tests/snippets/builtin_list.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/extra_tests/snippets/builtin_list.py b/extra_tests/snippets/builtin_list.py index d4afbffa1cb..5c47b2969ef 100644 --- a/extra_tests/snippets/builtin_list.py +++ b/extra_tests/snippets/builtin_list.py @@ -242,6 +242,21 @@ def __eq__(self, x): 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]" + lst = [3, 1, 5, 2, 4] From f49aa6810a7fc7004303577b978038397e303e6d Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:27:31 +0900 Subject: [PATCH 3/6] Cache the richcompare slot for homogeneous list sorts --- crates/vm/src/builtins/list.rs | 40 ++++++++++++++++++++++++++-- extra_tests/snippets/builtin_list.py | 1 + 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index a396bc263da..cc7e0ea5cd6 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -12,7 +12,7 @@ use crate::{ builtins::{PyFloat, PyInt, PyStr}, 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, }; @@ -642,6 +642,7 @@ enum PreSort { Str, Int, Float, + Object(RichCompareFunc), Generic, } @@ -662,6 +663,8 @@ fn pre_sort_check<'a>( PreSort::Int } else if class.is(vm.ctx.types.float_type) { PreSort::Float + } else if let Some(f) = class.slots.richcompare.load() { + PreSort::Object(f) } else { PreSort::Generic } @@ -679,6 +682,31 @@ fn float_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool { a.downcast_ref::().unwrap().to_f64() < b.downcast_ref::().unwrap().to_f64() } +fn object_lt( + cmp: RichCompareFunc, + a: &PyObjectRef, + b: &PyObjectRef, + vm: &VirtualMachine, +) -> PyResult { + #[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 timsort_specialized( vm: &VirtualMachine, items: &mut [T], @@ -714,6 +742,14 @@ where }; Ok(float_lt(a, b)) }), + PreSort::Object(cmp) => timsort(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + object_lt(cmp, a, b, vm) + }), PreSort::Generic => timsort(items, &mut |a, b| { let (a, b) = if reverse { (key(b), key(a)) diff --git a/extra_tests/snippets/builtin_list.py b/extra_tests/snippets/builtin_list.py index 5c47b2969ef..908d982cbba 100644 --- a/extra_tests/snippets/builtin_list.py +++ b/extra_tests/snippets/builtin_list.py @@ -256,6 +256,7 @@ class IntSub(int): 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"] lst = [3, 1, 5, 2, 4] From 2a0290410c4c0226244d90e858e0b992f9f6a933 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:29:19 +0900 Subject: [PATCH 4/6] Specialize tuple sorts on their first elements --- crates/vm/src/builtins/list.rs | 129 ++++++++++++++++++++++++--- extra_tests/snippets/builtin_list.py | 5 ++ 2 files changed, 120 insertions(+), 14 deletions(-) diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index cc7e0ea5cd6..e69a661061a 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -9,7 +9,7 @@ use crate::common::lock::{ use crate::object::{Traverse, TraverseFn}; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, - builtins::{PyFloat, PyInt, PyStr}, + builtins::{PyFloat, PyInt, PyStr, PyTuple}, class::PyClassImpl, convert::ToPyObject, function::{ArgSize, Either, FuncArgs, OptionalArg, PyComparisonValue}, @@ -638,14 +638,49 @@ impl Representable for PyList { } } +enum Elem { + Str, + Int, + Float, + Object(RichCompareFunc), + Generic, +} + enum PreSort { Str, Int, Float, Object(RichCompareFunc), + Tuple(Elem), Generic, } +impl From 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, 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, vm: &VirtualMachine, @@ -653,23 +688,50 @@ fn pre_sort_check<'a>( let Some(first) = keys.next() else { return PreSort::Generic; }; - let class = first.class(); - if !keys.all(|o| o.class().is(class)) { - return PreSort::Generic; - } - if class.is(vm.ctx.types.str_type) { - PreSort::Str - } else if class.is(vm.ctx.types.int_type) { - PreSort::Int - } else if class.is(vm.ctx.types.float_type) { - PreSort::Float - } else if let Some(f) = class.slots.richcompare.load() { - PreSort::Object(f) + + if let Some(t) = first + .downcast_ref_if_exact::(vm) + .filter(|t| !t.as_slice().is_empty()) + { + pre_sort_check_tuples(&t.as_slice()[0], keys, vm) } else { - PreSort::Generic + 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, + 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::(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::().unwrap().as_bytes() < b.downcast_ref::().unwrap().as_bytes() } @@ -707,6 +769,37 @@ fn object_lt( } } +fn elem_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + 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 { + let a = a.downcast_ref::().unwrap().as_slice(); + let b = b.downcast_ref::().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_specialized( vm: &VirtualMachine, items: &mut [T], @@ -750,6 +843,14 @@ where }; object_lt(cmp, a, b, vm) }), + PreSort::Tuple(elem) => timsort(items, &mut |a, b| { + let (a, b) = if reverse { + (key(b), key(a)) + } else { + (key(a), key(b)) + }; + tuple_lt(&elem, a, b, vm) + }), PreSort::Generic => timsort(items, &mut |a, b| { let (a, b) = if reverse { (key(b), key(a)) diff --git a/extra_tests/snippets/builtin_list.py b/extra_tests/snippets/builtin_list.py index 908d982cbba..12a78b53731 100644 --- a/extra_tests/snippets/builtin_list.py +++ b/extra_tests/snippets/builtin_list.py @@ -257,6 +257,11 @@ class IntSub(int): 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] From f76f442b8663649b85684ffd856ccb53da9f9136 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 8 Aug 2026 08:00:12 +0900 Subject: [PATCH 5/6] Apply ruff formatting to the new list sort tests --- extra_tests/snippets/builtin_list.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/extra_tests/snippets/builtin_list.py b/extra_tests/snippets/builtin_list.py index 12a78b53731..d62cae03b50 100644 --- a/extra_tests/snippets/builtin_list.py +++ b/extra_tests/snippets/builtin_list.py @@ -242,7 +242,13 @@ def __eq__(self, x): 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(["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] From d8e530761392d838041382ae42719e7581fd88d1 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 8 Aug 2026 08:00:12 +0900 Subject: [PATCH 6/6] Deduplicate the reverse swap across sort dispatch arms --- crates/vm/src/builtins/list.rs | 68 +++++++++++----------------------- 1 file changed, 22 insertions(+), 46 deletions(-) diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index e69a661061a..c2059e28806 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -800,6 +800,22 @@ fn tuple_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) } } +fn timsort_by(items: &mut [T], reverse: bool, key: &K, mut lt: L) -> PyResult<()> +where + T: Clone, + K: Fn(&T) -> &PyObjectRef, + L: FnMut(&PyObjectRef, &PyObjectRef) -> PyResult, +{ + 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( vm: &VirtualMachine, items: &mut [T], @@ -811,52 +827,12 @@ where K: Fn(&T) -> &PyObjectRef, { match pre_sort_check(items.iter().map(&key), vm) { - PreSort::Str => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; - Ok(str_lt(a, b)) - }), - PreSort::Int => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; - Ok(int_lt(a, b)) - }), - PreSort::Float => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; - Ok(float_lt(a, b)) - }), - PreSort::Object(cmp) => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; - object_lt(cmp, a, b, vm) - }), - PreSort::Tuple(elem) => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; - tuple_lt(&elem, a, b, vm) - }), - PreSort::Generic => timsort(items, &mut |a, b| { - let (a, b) = if reverse { - (key(b), key(a)) - } else { - (key(a), key(b)) - }; + 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) }), }