Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
228 changes: 210 additions & 18 deletions crates/vm/src/builtins/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
};
Expand Down Expand Up @@ -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)
}),
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

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(())
Expand Down
27 changes: 27 additions & 0 deletions extra_tests/snippets/builtin_list.py
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,33 @@ 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]"
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]


Expand Down
Loading