From 6e1c04ed9999511234f4bc9f5faf5d6897a84f4d Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 01:42:43 +0900 Subject: [PATCH 01/11] Add powersort implementation for list sorting Add crates/vm/src/sorting.rs implementing Tim Peters' timsort with powersort's merge-ordering policy (CPython 3.11+): run detection, binary insertion for short runs, galloping merge (merge_lo/merge_hi), and power-based merge ordering (powerloop). Comparison is passed in as a fallible `is_lt` closure, so the algorithm stays generic over the element type and free of interpreter details. Not yet wired into list.sort(); replaces rust-timsort in a follow-up. --- crates/vm/src/lib.rs | 1 + crates/vm/src/sorting.rs | 736 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 737 insertions(+) create mode 100644 crates/vm/src/sorting.rs diff --git a/crates/vm/src/lib.rs b/crates/vm/src/lib.rs index c353f8dfc46..df67c979739 100644 --- a/crates/vm/src/lib.rs +++ b/crates/vm/src/lib.rs @@ -92,6 +92,7 @@ pub mod scope; pub mod sequence; pub mod signal; pub mod sliceable; +pub mod sorting; pub mod stdlib; pub mod suggestion; pub mod types; diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs new file mode 100644 index 00000000000..8c8e8985e9e --- /dev/null +++ b/crates/vm/src/sorting.rs @@ -0,0 +1,736 @@ +// 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; + +enum Breakout { + Succeed, + CopyA, + CopyB, +} + +#[derive(Clone, Copy)] +struct Run { + base: usize, + len: usize, + power: u32, +} + +struct MergeState { + buf: Vec, + min_gallop: usize, + pending: Vec, +} + +impl MergeState { + fn merge_lo( + &mut self, + values: &mut [T], + is_lt: &mut F, + start_a: usize, + mut len_a: usize, + start_b: usize, + mut len_b: usize, + ) -> Result<(), E> + where + F: FnMut(&T, &T) -> Result, + { + debug_assert!(len_a > 0); + debug_assert!(len_b > 0); + debug_assert!(start_a + len_a == start_b); + + self.buf.clear(); + self.buf + .extend_from_slice(&values[start_a..start_a + len_a]); + + let mut cursor_a = 0; + let mut cursor_b = start_b; + let mut dest = start_a; + + values[dest] = values[cursor_b].clone(); + dest += 1; + cursor_b += 1; + len_b -= 1; + + if len_b == 0 { + values[dest..dest + len_a].clone_from_slice(&self.buf[cursor_a..cursor_a + len_a]); + return Ok(()); + } + if len_a == 1 { + copy_within_clone(values, cursor_b, dest, len_b); + values[dest + len_b] = self.buf[cursor_a].clone(); + return Ok(()); + } + + let mut min_gallop = self.min_gallop; + let mut breakout: Option = None; + + loop { + let mut a_count = 0; + let mut b_count = 0; + + loop { + if is_lt(&values[cursor_b], &self.buf[cursor_a])? { + values[dest] = values[cursor_b].clone(); + dest += 1; + cursor_b += 1; + len_b -= 1; + b_count += 1; + a_count = 0; + if len_b == 0 { + breakout = Some(Breakout::Succeed); + break; + } + if b_count >= min_gallop { + break; + } + } else { + values[dest] = self.buf[cursor_a].clone(); + dest += 1; + cursor_a += 1; + len_a -= 1; + a_count += 1; + b_count = 0; + if len_a == 1 { + breakout = Some(Breakout::CopyB); + break; + } + if a_count >= min_gallop { + break; + } + } + } + + if breakout.is_some() { + break; + } + + min_gallop += 1; + loop { + if min_gallop > 1 { + min_gallop -= 1; + } + self.min_gallop = min_gallop; + let mut k = gallop_right(&self.buf, is_lt, &values[cursor_b], cursor_a, len_a, 0)?; + a_count = k; + if k > 0 { + values[dest..dest + k].clone_from_slice(&self.buf[cursor_a..cursor_a + k]); + dest += k; + cursor_a += k; + len_a -= k; + if len_a == 1 { + breakout = Some(Breakout::CopyB); + break; + } + if len_a == 0 { + breakout = Some(Breakout::Succeed); + break; + } + } + values[dest] = values[cursor_b].clone(); + dest += 1; + cursor_b += 1; + len_b -= 1; + if len_b == 0 { + breakout = Some(Breakout::Succeed); + break; + } + k = gallop_left(&values, is_lt, &self.buf[cursor_a], cursor_b, len_b, 0)?; + b_count = k; + if k > 0 { + copy_within_clone(values, cursor_b, dest, k); + dest += k; + cursor_b += k; + len_b -= k; + if len_b == 0 { + breakout = Some(Breakout::Succeed); + break; + } + } + if len_a == 1 { + breakout = Some(Breakout::CopyB); + break; + } + if a_count < MIN_GALLOP && b_count < MIN_GALLOP { + break; + } + } + if breakout.is_some() { + break; + } + min_gallop += 1; + self.min_gallop = min_gallop; + } + + match breakout { + Some(Breakout::Succeed) => { + if len_a > 0 { + values[dest..dest + len_a] + .clone_from_slice(&self.buf[cursor_a..cursor_a + len_a]); + } + Ok(()) + } + Some(Breakout::CopyB) => { + copy_within_clone(values, cursor_b, dest, len_b); + values[dest + len_b] = self.buf[cursor_a].clone(); + Ok(()) + } + _ => unreachable!(), + } + } + + fn merge_hi( + &mut self, + values: &mut [T], + is_lt: &mut F, + start_a: usize, + mut len_a: usize, + start_b: usize, + mut len_b: usize, + ) -> Result<(), E> + where + F: FnMut(&T, &T) -> Result, + { + debug_assert!(len_a > 0); + debug_assert!(len_b > 0); + debug_assert!(start_a + len_a == start_b); + + self.buf.clear(); + self.buf + .extend_from_slice(&values[start_b..start_b + len_b]); + + let mut dest = start_b + len_b - 1; + let mut cursor_a = start_a + len_a - 1; + let mut cursor_b = len_b - 1; + + values[dest] = values[cursor_a].clone(); + dest -= 1; + cursor_a -= 1; + len_a -= 1; + + if len_a == 0 { + values[dest - len_b + 1..dest + 1].clone_from_slice(&self.buf[0..len_b]); + return Ok(()); + } + if len_b == 1 { + let src = cursor_a + 1 - len_a; + let dst = dest + 1 - len_a; + copy_within_clone(values, src, dst, len_a); + values[dst - 1] = self.buf[cursor_b].clone(); + return Ok(()); + } + + let mut min_gallop = self.min_gallop; + let mut breakout: Option = None; + + loop { + let mut a_count = 0; + let mut b_count = 0; + + loop { + if is_lt(&self.buf[cursor_b], &values[cursor_a])? { + values[dest] = values[cursor_a].clone(); + dest -= 1; + len_a -= 1; + + if len_a == 0 { + breakout = Some(Breakout::Succeed); + break; + } + + cursor_a -= 1; + a_count += 1; + b_count = 0; + + if a_count >= min_gallop { + break; + } + } else { + values[dest] = self.buf[cursor_b].clone(); + dest -= 1; + cursor_b -= 1; + len_b -= 1; + b_count += 1; + a_count = 0; + if len_b == 1 { + breakout = Some(Breakout::CopyA); + break; + } + if b_count >= min_gallop { + break; + } + } + } + + if breakout.is_some() { + break; + } + + min_gallop += 1; + loop { + if min_gallop > 1 { + min_gallop -= 1; + } + self.min_gallop = min_gallop; + let mut k = gallop_right( + &values, + is_lt, + &self.buf[cursor_b], + start_a, + len_a, + len_a - 1, + )?; + k = len_a - k; + a_count = k; + if k > 0 { + copy_within_clone(values, cursor_a + 1 - k, dest + 1 - k, k); + dest -= k; + len_a -= k; + if len_a == 0 { + breakout = Some(Breakout::Succeed); + break; + } + cursor_a -= k; + } + values[dest] = self.buf[cursor_b].clone(); + dest -= 1; + cursor_b -= 1; + len_b -= 1; + if len_b == 1 { + breakout = Some(Breakout::CopyA); + break; + } + k = gallop_left(&self.buf, is_lt, &values[cursor_a], 0, len_b, len_b - 1)?; + k = len_b - k; + b_count = k; + if k > 0 { + values[dest + 1 - k..dest + 1] + .clone_from_slice(&self.buf[cursor_b + 1 - k..cursor_b + 1]); + dest -= k; + len_b -= k; + + if len_b == 0 { + breakout = Some(Breakout::Succeed); + break; + } + cursor_b -= k; + + if len_b == 1 { + breakout = Some(Breakout::CopyA); + break; + } + } + values[dest] = values[cursor_a].clone(); + dest -= 1; + len_a -= 1; + + if len_a == 0 { + breakout = Some(Breakout::Succeed); + break; + } + + cursor_a -= 1; + + if a_count < MIN_GALLOP && b_count < MIN_GALLOP { + break; + } + } + if breakout.is_some() { + break; + } + min_gallop += 1; + self.min_gallop = min_gallop; + } + + match breakout { + Some(Breakout::Succeed) => { + if len_b > 0 { + values[dest - len_b + 1..dest + 1].clone_from_slice(&self.buf[0..len_b]); + } + Ok(()) + } + Some(Breakout::CopyA) => { + let src = cursor_a + 1 - len_a; + let dst = dest + 1 - len_a; + copy_within_clone(values, src, dst, len_a); + values[dst - 1] = self.buf[cursor_b].clone(); + Ok(()) + } + _ => unreachable!(), + } + } + + fn merge_at(&mut self, values: &mut [T], is_lt: &mut F, i: usize) -> Result<(), E> + where + F: FnMut(&T, &T) -> Result, + { + debug_assert!(self.pending.len() >= 2); + debug_assert!(i == self.pending.len() - 2 || i == self.pending.len() - 3); + + let mut start_a = self.pending[i].base; + let mut len_a = self.pending[i].len; + let start_b = self.pending[i + 1].base; + let mut len_b = self.pending[i + 1].len; + + debug_assert!(len_a > 0); + debug_assert!(len_b > 0); + debug_assert!(start_a + len_a == start_b); + + self.pending[i].len = len_a + len_b; + self.pending.remove(i + 1); + + let k = gallop_right(values, is_lt, &values[start_b], start_a, len_a, 0)?; + start_a += k; + len_a -= k; + + if len_a == 0 { + return Ok(()); + } + + len_b = gallop_left( + values, + is_lt, + &values[start_a + len_a - 1], + start_b, + len_b, + len_b - 1, + )?; + + if len_b == 0 { + return Ok(()); + } + + if len_a <= len_b { + self.merge_lo(values, is_lt, start_a, len_a, start_b, len_b)?; + } else { + self.merge_hi(values, is_lt, start_a, len_a, start_b, len_b)?; + } + Ok(()) + } + + fn found_new_run(&mut self, new_run_len: usize, values: &mut [T], is_lt: &mut F) -> Result<(), E> + where + F: FnMut(&T, &T) -> Result, + { + if !self.pending.is_empty() { + let last = self.pending.len() - 1; + let s1 = self.pending[last].base; + let n1 = self.pending[last].len; + let power = powerloop(s1, n1, new_run_len, values.len()); + + while self.pending.len() > 1 && self.pending[self.pending.len() - 2].power > power { + self.merge_at(values, is_lt, self.pending.len() - 2)?; + } + + debug_assert!( + self.pending.len() < 2 || self.pending[self.pending.len() - 2].power < power + ); + let last = self.pending.len() - 1; + self.pending[last].power = power; + } + Ok(()) + } + + fn push_run(&mut self, base: usize, len: usize) { + self.pending.push(Run { + base, + len, + power: 0, + }) + } + + fn merge_force_collapse(&mut self, values: &mut [T], is_lt: &mut F) -> Result<(), E> + where + F: FnMut(&T, &T) -> Result, + { + while self.pending.len() > 1 { + let mut n = self.pending.len() - 2; + if n > 0 && self.pending[n - 1].len < self.pending[n + 1].len { + n -= 1; + } + self.merge_at(values, is_lt, n)?; + } + Ok(()) + } +} + +fn binary_insertion_sort(values: &mut [T], is_lt: &mut F, start: usize) -> Result<(), E> +where + F: FnMut(&T, &T) -> Result, +{ + for i in start..values.len() { + let mut l = 0; + let mut r = i; + + while l < r { + let m = (l + r) / 2; + if is_lt(&values[i], &values[m])? { + r = m; + } else { + l = m + 1; + } + } + values[l..=i].rotate_right(1); + } + Ok(()) +} + +fn copy_within_clone(values: &mut [T], src: usize, dest: usize, n: usize) { + if dest <= src { + for k in 0..n { + values[dest + k] = values[src + k].clone(); + } + } else { + for k in (0..n).rev() { + values[dest + k] = values[src + k].clone(); + } + } +} + +fn count_run(values: &[T], is_lt: &mut F) -> Result<(usize, bool), E> +where + F: FnMut(&T, &T) -> Result, +{ + let n = values.len(); + if n == 1 { + return Ok((1, false)); + } + let mut i = 2; + let descending = is_lt(&values[1], &values[0])?; + if descending { + while i < n && is_lt(&values[i], &values[i - 1])? { + i += 1; + } + } else { + while i < n && !is_lt(&values[i], &values[i - 1])? { + i += 1; + } + } + Ok((i, descending)) +} + +fn gallop_left( + values: &[T], + is_lt: &mut F, + key: &T, + base: usize, + len: usize, + hint: usize, +) -> Result +where + F: FnMut(&T, &T) -> Result, +{ + debug_assert!(hint < len); + let mut lastofs: isize = 0; + let mut ofs: isize = 1; + let hint_i = hint as isize; + let len_i = len as isize; + + if is_lt(&values[base + hint], key)? { + let maxofs = len_i - hint_i; + while ofs < maxofs && is_lt(&values[base + hint + ofs as usize], key)? { + lastofs = ofs; + ofs = (ofs * 2) + 1; + } + if ofs > maxofs { + ofs = maxofs; + } + lastofs += hint_i; + ofs += hint_i; + } else { + let maxofs = hint_i + 1; + while ofs < maxofs && !is_lt(&values[base + (hint_i - ofs) as usize], key)? { + lastofs = ofs; + ofs = (ofs * 2) + 1; + } + if ofs > maxofs { + ofs = maxofs; + } + (lastofs, ofs) = (hint_i - ofs, hint_i - lastofs); + } + lastofs += 1; + while lastofs < ofs { + let m = lastofs + ((ofs - lastofs) / 2); + if is_lt(&values[base + m as usize], key)? { + lastofs = m + 1; + } else { + ofs = m; + } + } + Ok(ofs as usize) +} + +fn gallop_right( + values: &[T], + is_lt: &mut F, + key: &T, + base: usize, + len: usize, + hint: usize, +) -> Result +where + F: FnMut(&T, &T) -> Result, +{ + debug_assert!(hint < len); + let mut lastofs: isize = 0; + let mut ofs: isize = 1; + let hint_i = hint as isize; + let len_i = len as isize; + + if is_lt(key, &values[base + hint])? { + let maxofs = hint_i + 1; + while ofs < maxofs && is_lt(key, &values[base + (hint_i - ofs) as usize])? { + lastofs = ofs; + ofs = (ofs * 2) + 1; + } + if ofs > maxofs { + ofs = maxofs; + } + (lastofs, ofs) = (hint_i - ofs, hint_i - lastofs); + } else { + let maxofs = len_i - hint_i; + while ofs < maxofs && !is_lt(key, &values[base + hint + ofs as usize])? { + lastofs = ofs; + ofs = (ofs * 2) + 1; + } + if ofs > maxofs { + ofs = maxofs; + } + lastofs += hint_i; + ofs += hint_i; + } + lastofs += 1; + while lastofs < ofs { + let m = lastofs + ((ofs - lastofs) / 2); + if is_lt(key, &values[base + m as usize])? { + ofs = m; + } else { + lastofs = m + 1; + } + } + Ok(ofs as usize) +} + +// TODO: consider CPython 3.12+'s incremental minrun (mr_current/mr_e/mr_mask) +// for a more precise minrun; current bit-shift version is the classic one. +fn merge_compute_minrun(mut n: usize) -> usize { + let mut r = 0; + while n >= MAX_MINRUN { + r |= n & 1; + n >>= 1; + } + n + r +} + +fn powerloop(s1: usize, n1: usize, n2: usize, n: usize) -> u32 { + let mut result: u32 = 0; + let mut a = 2 * s1 + n1; + let mut b = a + n1 + n2; + + loop { + result += 1; + if a >= n { + debug_assert!(b >= a); + a -= n; + b -= n; + } else if b >= n { + break; + } + debug_assert!(a < b && b < n); + a <<= 1; + b <<= 1; + } + result +} + +/// Stable adaptive mergesort (Tim Peters' timsort with powersort's +/// merge-ordering policy, matching CPython 3.11+). `is_lt` provides comparison. +pub(crate) fn timsort(values: &mut [T], is_lt: &mut F) -> Result<(), E> +where + T: Clone, + F: FnMut(&T, &T) -> Result, +{ + let n = values.len(); + let mut ms = MergeState { + buf: Vec::new(), + min_gallop: MIN_GALLOP, + pending: Vec::new(), + }; + + if n < 2 { + return Ok(()); + } + + if n < MAX_MINRUN { + let (l, desc) = count_run(&values, is_lt)?; + if desc { + values[0..l].reverse(); + } + binary_insertion_sort(values, is_lt, l)?; + return Ok(()); + } + + let minrun = merge_compute_minrun(n); + let mut lo = 0; + + while lo < n { + let (mut l, desc) = count_run(&values[lo..n], is_lt)?; + if desc { + values[lo..lo + l].reverse(); + } + if l < minrun { + let force = minrun.min(n - lo); + binary_insertion_sort(&mut values[lo..lo + force], is_lt, l)?; + l = force; + } + ms.found_new_run(l, values, is_lt)?; + ms.push_run(lo, l); + lo += l; + } + ms.merge_force_collapse(values, is_lt)?; + debug_assert!(ms.pending.len() == 1 && ms.pending[0].len == n); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sort(mut v: Vec) -> Vec { + timsort(&mut v, &mut |a: &i32, b: &i32| Ok::(a < b)).unwrap(); + v + } + + #[test] + fn test_basic() { + assert_eq!(sort(vec![3, 1, 2]), vec![1, 2, 3]); + assert_eq!(sort(Vec::::new()), Vec::::new()); + assert_eq!(sort(vec![1]), vec![1]); + assert_eq!(sort(vec![2, 1]), vec![1, 2]); + } + + #[test] + fn test_ordered() { + assert_eq!(sort(vec![1, 2, 3, 4, 5]), vec![1, 2, 3, 4, 5]); + assert_eq!(sort(vec![5, 4, 3, 2, 1]), vec![1, 2, 3, 4, 5]); + } + + #[test] + fn test_duplicates() { + assert_eq!(sort(vec![3, 1, 3, 1, 2, 2]), vec![1, 1, 2, 2, 3, 3]); + } + + #[test] + fn test_large() { + let mut v: Vec = (0..1000).rev().collect(); // 999..0 + let sorted: Vec = (0..1000).collect(); + assert_eq!(sort(v), sorted); + } + + #[test] + fn test_random_ish() { + let mut v: Vec = (0..500).map(|i| (i * 7919) % 500).collect(); + let mut expected = v.clone(); + expected.sort(); + assert_eq!(sort(v), expected); + } +} From 5b52036263f9d6a5b9b6c22d01216ecd5eb4b732 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 02:06:06 +0900 Subject: [PATCH 02/11] Use powersort for list.sort() and drop rust-timsort Wire list sorting through crate::sorting::timsort and remove the rust-timsort dependency. Fixes the O(N^2) behavior on random input (1M random floats: ~16min -> ~0.8s), now within ~4x of CPython. --- Cargo.lock | 7 ------- Cargo.toml | 1 - crates/vm/Cargo.toml | 1 - crates/vm/src/builtins/list.rs | 13 ++++++------- 4 files changed, 6 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0acd840c079..5b9870a0b0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3716,7 +3716,6 @@ dependencies = [ "strum_macros", "thin-vec", "thiserror", - "timsort", "wasm-bindgen", "widestring", ] @@ -4285,12 +4284,6 @@ dependencies = [ "time-core", ] -[[package]] -name = "timsort" -version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "639ce8ef6d2ba56be0383a94dd13b92138d58de44c62618303bb798fa92bdc00" - [[package]] name = "tinystr" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index ac1e35c5314..6bb22e36f04 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -305,7 +305,6 @@ textwrap = { version = "0.16.2", default-features = false } termios = "0.3.3" thiserror = "2.0" thin-vec = "0.2.14" -timsort = "0.1.2" tk-sys = { git = "https://github.com/arihant2math/tkinter.git", tag = "v0.2.0" } icu_casemap = "2" icu_locale = "2" diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml index ca92b9dfd76..6601fb03dc2 100644 --- a/crates/vm/Cargo.toml +++ b/crates/vm/Cargo.toml @@ -80,7 +80,6 @@ half = { workspace = true } psm = { workspace = true } optional = { workspace = true } result-like = { workspace = true } -timsort = { workspace = true } [target.'cfg(unix)'.dependencies] exitcode = { workspace = true } diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index 2bdbaf63cde..baa7f9bd46d 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -18,6 +18,7 @@ use crate::{ recursion::ReprGuard, sequence::{MutObjectSequenceOp, OptionalRangeArgs, SequenceExt, SequenceMutExt}, sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp}, + sorting::timsort, types::{ AsMapping, AsSequence, Comparable, Constructor, Initializer, IterNext, Iterable, PyComparisonOp, Representable, SelfIter, @@ -642,13 +643,11 @@ fn do_sort( ) -> PyResult<()> { // CPython uses __lt__ for all comparisons in sort. // try_sort_by_gt expects is_gt(a, b) = true when a should come AFTER b. - let cmp = |a: &PyObjectRef, b: &PyObjectRef| { + let mut is_lt = |a: &PyObjectRef, b: &PyObjectRef| { if reverse { - // Descending: a comes after b when a < b - a.rich_compare_bool(b, PyComparisonOp::Lt, vm) - } else { - // Ascending: a comes after b when b < a b.rich_compare_bool(a, PyComparisonOp::Lt, vm) + } else { + a.rich_compare_bool(b, PyComparisonOp::Lt, vm) } }; @@ -657,10 +656,10 @@ fn do_sort( .iter() .map(|x| Ok((x.clone(), key_func.call((x.clone(),), vm)?))) .collect::, _>>()?; - timsort::try_sort_by_gt(&mut items, |a, b| cmp(&a.1, &b.1))?; + timsort(&mut items, &mut |a, b| is_lt(&a.1, &b.1))?; *values = items.into_iter().map(|(val, _)| val).collect(); } else { - timsort::try_sort_by_gt(values, cmp)?; + timsort(values, &mut is_lt)?; } Ok(()) From e5a5247776e934a3cc0482335f3352d5e3ca5f31 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 02:23:49 +0900 Subject: [PATCH 03/11] Fix usize underflow in merge_hi Succeed path --- crates/vm/src/sorting.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs index 8c8e8985e9e..f2819109ab7 100644 --- a/crates/vm/src/sorting.rs +++ b/crates/vm/src/sorting.rs @@ -344,7 +344,7 @@ impl MergeState { match breakout { Some(Breakout::Succeed) => { if len_b > 0 { - values[dest - len_b + 1..dest + 1].clone_from_slice(&self.buf[0..len_b]); + values[(dest + 1) - len_b..dest + 1].clone_from_slice(&self.buf[0..len_b]); } Ok(()) } From 189e6e7a3e2185798be85bd52e66f489e20effb4 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 02:48:47 +0900 Subject: [PATCH 04/11] Fix clippy warnings and formatting in sorting.rs --- crates/vm/src/sorting.rs | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs index f2819109ab7..e35c1062543 100644 --- a/crates/vm/src/sorting.rs +++ b/crates/vm/src/sorting.rs @@ -134,7 +134,7 @@ impl MergeState { breakout = Some(Breakout::Succeed); break; } - k = gallop_left(&values, is_lt, &self.buf[cursor_a], cursor_b, len_b, 0)?; + k = gallop_left(values, is_lt, &self.buf[cursor_a], cursor_b, len_b, 0)?; b_count = k; if k > 0 { copy_within_clone(values, cursor_b, dest, k); @@ -208,7 +208,7 @@ impl MergeState { len_a -= 1; if len_a == 0 { - values[dest - len_b + 1..dest + 1].clone_from_slice(&self.buf[0..len_b]); + values[(dest - len_b + 1)..=dest].clone_from_slice(&self.buf[0..len_b]); return Ok(()); } if len_b == 1 { @@ -272,7 +272,7 @@ impl MergeState { } self.min_gallop = min_gallop; let mut k = gallop_right( - &values, + values, is_lt, &self.buf[cursor_b], start_a, @@ -303,8 +303,8 @@ impl MergeState { k = len_b - k; b_count = k; if k > 0 { - values[dest + 1 - k..dest + 1] - .clone_from_slice(&self.buf[cursor_b + 1 - k..cursor_b + 1]); + values[dest + 1 - k..=dest] + .clone_from_slice(&self.buf[cursor_b + 1 - k..=cursor_b]); dest -= k; len_b -= k; @@ -344,7 +344,7 @@ impl MergeState { match breakout { Some(Breakout::Succeed) => { if len_b > 0 { - values[(dest + 1) - len_b..dest + 1].clone_from_slice(&self.buf[0..len_b]); + values[(dest + 1) - len_b..=dest].clone_from_slice(&self.buf[0..len_b]); } Ok(()) } @@ -407,7 +407,12 @@ impl MergeState { Ok(()) } - fn found_new_run(&mut self, new_run_len: usize, values: &mut [T], is_lt: &mut F) -> Result<(), E> + fn found_new_run( + &mut self, + new_run_len: usize, + values: &mut [T], + is_lt: &mut F, + ) -> Result<(), E> where F: FnMut(&T, &T) -> Result, { @@ -661,7 +666,7 @@ where } if n < MAX_MINRUN { - let (l, desc) = count_run(&values, is_lt)?; + let (l, desc) = count_run(values, is_lt)?; if desc { values[0..l].reverse(); } @@ -721,14 +726,14 @@ mod tests { #[test] fn test_large() { - let mut v: Vec = (0..1000).rev().collect(); // 999..0 + let v: Vec = (0..1000).rev().collect(); // 999..0 let sorted: Vec = (0..1000).collect(); assert_eq!(sort(v), sorted); } #[test] fn test_random_ish() { - let mut v: Vec = (0..500).map(|i| (i * 7919) % 500).collect(); + let v: Vec = (0..500).map(|i| (i * 7919) % 500).collect(); let mut expected = v.clone(); expected.sort(); assert_eq!(sort(v), expected); From 526820be67197421a4f43a42f43426d1f06654e1 Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 13:22:52 +0900 Subject: [PATCH 05/11] Drop redundant test_ prefixes in sorting tests --- crates/vm/src/sorting.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs index e35c1062543..5062ceadd37 100644 --- a/crates/vm/src/sorting.rs +++ b/crates/vm/src/sorting.rs @@ -706,7 +706,7 @@ mod tests { } #[test] - fn test_basic() { + fn basic_examples() { assert_eq!(sort(vec![3, 1, 2]), vec![1, 2, 3]); assert_eq!(sort(Vec::::new()), Vec::::new()); assert_eq!(sort(vec![1]), vec![1]); @@ -714,25 +714,25 @@ mod tests { } #[test] - fn test_ordered() { + fn five_elements_forwards_and_backwards() { assert_eq!(sort(vec![1, 2, 3, 4, 5]), vec![1, 2, 3, 4, 5]); assert_eq!(sort(vec![5, 4, 3, 2, 1]), vec![1, 2, 3, 4, 5]); } #[test] - fn test_duplicates() { + fn six_elements_with_duplicates() { assert_eq!(sort(vec![3, 1, 3, 1, 2, 2]), vec![1, 1, 2, 2, 3, 3]); } #[test] - fn test_large() { + fn one_thousand_elements() { let v: Vec = (0..1000).rev().collect(); // 999..0 let sorted: Vec = (0..1000).collect(); assert_eq!(sort(v), sorted); } #[test] - fn test_random_ish() { + fn pseudorandom_collection() { let v: Vec = (0..500).map(|i| (i * 7919) % 500).collect(); let mut expected = v.clone(); expected.sort(); From ecbda8e5a52eaa84d078b04a389c04b738def05c Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 16:00:14 +0900 Subject: [PATCH 06/11] Restore buffered elements when a comparison fails mid-merge --- crates/vm/src/builtins/list.rs | 5 +- crates/vm/src/sorting.rs | 141 ++++++++++++++++----------------- 2 files changed, 71 insertions(+), 75 deletions(-) diff --git a/crates/vm/src/builtins/list.rs b/crates/vm/src/builtins/list.rs index baa7f9bd46d..8a426685ad2 100644 --- a/crates/vm/src/builtins/list.rs +++ b/crates/vm/src/builtins/list.rs @@ -642,7 +642,10 @@ fn do_sort( reverse: bool, ) -> PyResult<()> { // CPython uses __lt__ for all comparisons in sort. - // try_sort_by_gt expects is_gt(a, b) = true when a should come AFTER b. + // `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) diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs index 5062ceadd37..ddc2d6f800a 100644 --- a/crates/vm/src/sorting.rs +++ b/crates/vm/src/sorting.rs @@ -2,12 +2,16 @@ const MIN_GALLOP: usize = 7; const MAX_MINRUN: usize = 64; -enum Breakout { +enum LoBreakout { Succeed, - CopyA, CopyB, } +enum HiBreakout { + Succeed, + CopyA, +} + #[derive(Clone, Copy)] struct Run { base: usize, @@ -62,14 +66,17 @@ impl MergeState { } let mut min_gallop = self.min_gallop; - let mut breakout: Option = None; - loop { + let breakout: Result = 'merging: loop { let mut a_count = 0; let mut b_count = 0; loop { - if is_lt(&values[cursor_b], &self.buf[cursor_a])? { + let b_wins = match is_lt(&values[cursor_b], &self.buf[cursor_a]) { + Ok(v) => v, + Err(e) => break 'merging Err(e), + }; + if b_wins { values[dest] = values[cursor_b].clone(); dest += 1; cursor_b += 1; @@ -77,8 +84,7 @@ impl MergeState { b_count += 1; a_count = 0; if len_b == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(LoBreakout::Succeed); } if b_count >= min_gallop { break; @@ -91,8 +97,7 @@ impl MergeState { a_count += 1; b_count = 0; if len_a == 1 { - breakout = Some(Breakout::CopyB); - break; + break 'merging Ok(LoBreakout::CopyB); } if a_count >= min_gallop { break; @@ -100,17 +105,17 @@ impl MergeState { } } - if breakout.is_some() { - break; - } - min_gallop += 1; loop { if min_gallop > 1 { min_gallop -= 1; } self.min_gallop = min_gallop; - let mut k = gallop_right(&self.buf, is_lt, &values[cursor_b], cursor_a, len_a, 0)?; + let mut k = + match gallop_right(&self.buf, is_lt, &values[cursor_b], cursor_a, len_a, 0) { + Ok(k) => k, + Err(e) => break 'merging Err(e), + }; a_count = k; if k > 0 { values[dest..dest + k].clone_from_slice(&self.buf[cursor_a..cursor_a + k]); @@ -118,12 +123,10 @@ impl MergeState { cursor_a += k; len_a -= k; if len_a == 1 { - breakout = Some(Breakout::CopyB); - break; + break 'merging Ok(LoBreakout::CopyB); } if len_a == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(LoBreakout::Succeed); } } values[dest] = values[cursor_b].clone(); @@ -131,10 +134,12 @@ impl MergeState { cursor_b += 1; len_b -= 1; if len_b == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(LoBreakout::Succeed); } - k = gallop_left(values, is_lt, &self.buf[cursor_a], cursor_b, len_b, 0)?; + k = match gallop_left(values, is_lt, &self.buf[cursor_a], cursor_b, len_b, 0) { + Ok(k) => k, + Err(e) => break 'merging Err(e), + }; b_count = k; if k > 0 { copy_within_clone(values, cursor_b, dest, k); @@ -142,39 +147,34 @@ impl MergeState { cursor_b += k; len_b -= k; if len_b == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(LoBreakout::Succeed); } } if len_a == 1 { - breakout = Some(Breakout::CopyB); - break; + break 'merging Ok(LoBreakout::CopyB); } if a_count < MIN_GALLOP && b_count < MIN_GALLOP { break; } } - if breakout.is_some() { - break; - } + min_gallop += 1; self.min_gallop = min_gallop; - } + }; match breakout { - Some(Breakout::Succeed) => { + Ok(LoBreakout::CopyB) => { + copy_within_clone(values, cursor_b, dest, len_b); + values[dest + len_b] = self.buf[cursor_a].clone(); + Ok(()) + } + other => { if len_a > 0 { values[dest..dest + len_a] .clone_from_slice(&self.buf[cursor_a..cursor_a + len_a]); } - Ok(()) + other.map(|_| ()) } - Some(Breakout::CopyB) => { - copy_within_clone(values, cursor_b, dest, len_b); - values[dest + len_b] = self.buf[cursor_a].clone(); - Ok(()) - } - _ => unreachable!(), } } @@ -220,21 +220,22 @@ impl MergeState { } let mut min_gallop = self.min_gallop; - let mut breakout: Option = None; - - loop { + let breakout: Result = 'merging: loop { let mut a_count = 0; let mut b_count = 0; loop { - if is_lt(&self.buf[cursor_b], &values[cursor_a])? { + let b_wins = match is_lt(&self.buf[cursor_b], &values[cursor_a]) { + Ok(v) => v, + Err(e) => break 'merging Err(e), + }; + if b_wins { values[dest] = values[cursor_a].clone(); dest -= 1; len_a -= 1; if len_a == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(HiBreakout::Succeed); } cursor_a -= 1; @@ -252,8 +253,7 @@ impl MergeState { b_count += 1; a_count = 0; if len_b == 1 { - breakout = Some(Breakout::CopyA); - break; + break 'merging Ok(HiBreakout::CopyA); } if b_count >= min_gallop { break; @@ -261,24 +261,23 @@ impl MergeState { } } - if breakout.is_some() { - break; - } - min_gallop += 1; loop { if min_gallop > 1 { min_gallop -= 1; } self.min_gallop = min_gallop; - let mut k = gallop_right( + let mut k = match gallop_right( values, is_lt, &self.buf[cursor_b], start_a, len_a, len_a - 1, - )?; + ) { + Ok(k) => k, + Err(e) => break 'merging Err(e), + }; k = len_a - k; a_count = k; if k > 0 { @@ -286,8 +285,7 @@ impl MergeState { dest -= k; len_a -= k; if len_a == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(HiBreakout::Succeed); } cursor_a -= k; } @@ -296,10 +294,12 @@ impl MergeState { cursor_b -= 1; len_b -= 1; if len_b == 1 { - breakout = Some(Breakout::CopyA); - break; + break 'merging Ok(HiBreakout::CopyA); } - k = gallop_left(&self.buf, is_lt, &values[cursor_a], 0, len_b, len_b - 1)?; + k = match gallop_left(&self.buf, is_lt, &values[cursor_a], 0, len_b, len_b - 1) { + Ok(k) => k, + Err(e) => break 'merging Err(e), + }; k = len_b - k; b_count = k; if k > 0 { @@ -309,14 +309,12 @@ impl MergeState { len_b -= k; if len_b == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(HiBreakout::Succeed); } cursor_b -= k; if len_b == 1 { - breakout = Some(Breakout::CopyA); - break; + break 'merging Ok(HiBreakout::CopyA); } } values[dest] = values[cursor_a].clone(); @@ -324,8 +322,7 @@ impl MergeState { len_a -= 1; if len_a == 0 { - breakout = Some(Breakout::Succeed); - break; + break 'merging Ok(HiBreakout::Succeed); } cursor_a -= 1; @@ -334,28 +331,24 @@ impl MergeState { break; } } - if breakout.is_some() { - break; - } min_gallop += 1; self.min_gallop = min_gallop; - } + }; match breakout { - Some(Breakout::Succeed) => { - if len_b > 0 { - values[(dest + 1) - len_b..=dest].clone_from_slice(&self.buf[0..len_b]); - } - Ok(()) - } - Some(Breakout::CopyA) => { + Ok(HiBreakout::CopyA) => { let src = cursor_a + 1 - len_a; let dst = dest + 1 - len_a; copy_within_clone(values, src, dst, len_a); values[dst - 1] = self.buf[cursor_b].clone(); Ok(()) } - _ => unreachable!(), + other => { + if len_b > 0 { + values[(dest + 1) - len_b..=dest].clone_from_slice(&self.buf[0..len_b]); + } + other.map(|_| ()) + } } } From b4fda72fa412bda498d897a1d11f700b5064a57d Mon Sep 17 00:00:00 2001 From: KangDora Date: Sat, 1 Aug 2026 16:01:18 +0900 Subject: [PATCH 07/11] Copy one run-A element per gallop round in merge_lo --- crates/vm/src/sorting.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/vm/src/sorting.rs b/crates/vm/src/sorting.rs index ddc2d6f800a..5f61d6624e1 100644 --- a/crates/vm/src/sorting.rs +++ b/crates/vm/src/sorting.rs @@ -150,6 +150,10 @@ impl MergeState { break 'merging Ok(LoBreakout::Succeed); } } + values[dest] = self.buf[cursor_a].clone(); + dest += 1; + cursor_a += 1; + len_a -= 1; if len_a == 1 { break 'merging Ok(LoBreakout::CopyB); } From 6a153dd833de7a244af5f427e32b5d21f50e2c85 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 09:50:31 +0900 Subject: [PATCH 08/11] 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 8a426685ad2..ef6a0771dd2 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}, @@ -635,34 +635,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 b92659d36f467bac3419d3bf19c9a2b4858b5018 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:26:33 +0900 Subject: [PATCH 09/11] 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 3010dd499d4e5474acffcf06ce1cf6a37eff6758 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:27:31 +0900 Subject: [PATCH 10/11] 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 ef6a0771dd2..cb8cc66c918 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, }; @@ -639,6 +639,7 @@ enum PreSort { Str, Int, Float, + Object(RichCompareFunc), Generic, } @@ -659,6 +660,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 } @@ -676,6 +679,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], @@ -711,6 +739,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 fc58bdfff6fd524cbefbfd88f1cb5f3eb95dba79 Mon Sep 17 00:00:00 2001 From: KangDora Date: Wed, 5 Aug 2026 11:29:19 +0900 Subject: [PATCH 11/11] 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 cb8cc66c918..fed39010a3b 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}, @@ -635,14 +635,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, @@ -650,23 +685,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() } @@ -704,6 +766,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], @@ -747,6 +840,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]