Skip to content

Commit 2a02904

Browse files
committed
Specialize tuple sorts on their first elements
1 parent f49aa68 commit 2a02904

2 files changed

Lines changed: 120 additions & 14 deletions

File tree

‎crates/vm/src/builtins/list.rs‎

Lines changed: 115 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use crate::common::lock::{
99
use crate::object::{Traverse, TraverseFn};
1010
use crate::{
1111
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
12-
builtins::{PyFloat, PyInt, PyStr},
12+
builtins::{PyFloat, PyInt, PyStr, PyTuple},
1313
class::PyClassImpl,
1414
convert::ToPyObject,
1515
function::{ArgSize, Either, FuncArgs, OptionalArg, PyComparisonValue},
@@ -638,38 +638,100 @@ impl Representable for PyList {
638638
}
639639
}
640640

641+
enum Elem {
642+
Str,
643+
Int,
644+
Float,
645+
Object(RichCompareFunc),
646+
Generic,
647+
}
648+
641649
enum PreSort {
642650
Str,
643651
Int,
644652
Float,
645653
Object(RichCompareFunc),
654+
Tuple(Elem),
646655
Generic,
647656
}
648657

658+
impl From<Elem> for PreSort {
659+
fn from(e: Elem) -> Self {
660+
match e {
661+
Elem::Str => Self::Str,
662+
Elem::Int => Self::Int,
663+
Elem::Float => Self::Float,
664+
Elem::Object(f) => Self::Object(f),
665+
Elem::Generic => Self::Generic,
666+
}
667+
}
668+
}
669+
670+
fn classify(class: &Py<PyType>, vm: &VirtualMachine) -> Elem {
671+
if class.is(vm.ctx.types.str_type) {
672+
Elem::Str
673+
} else if class.is(vm.ctx.types.int_type) {
674+
Elem::Int
675+
} else if class.is(vm.ctx.types.float_type) {
676+
Elem::Float
677+
} else if let Some(f) = class.slots.richcompare.load() {
678+
Elem::Object(f)
679+
} else {
680+
Elem::Generic
681+
}
682+
}
683+
649684
fn pre_sort_check<'a>(
650685
mut keys: impl Iterator<Item = &'a PyObjectRef>,
651686
vm: &VirtualMachine,
652687
) -> PreSort {
653688
let Some(first) = keys.next() else {
654689
return PreSort::Generic;
655690
};
656-
let class = first.class();
657-
if !keys.all(|o| o.class().is(class)) {
658-
return PreSort::Generic;
659-
}
660-
if class.is(vm.ctx.types.str_type) {
661-
PreSort::Str
662-
} else if class.is(vm.ctx.types.int_type) {
663-
PreSort::Int
664-
} else if class.is(vm.ctx.types.float_type) {
665-
PreSort::Float
666-
} else if let Some(f) = class.slots.richcompare.load() {
667-
PreSort::Object(f)
691+
692+
if let Some(t) = first
693+
.downcast_ref_if_exact::<PyTuple>(vm)
694+
.filter(|t| !t.as_slice().is_empty())
695+
{
696+
pre_sort_check_tuples(&t.as_slice()[0], keys, vm)
668697
} else {
669-
PreSort::Generic
698+
let class = first.class();
699+
if keys.all(|k| k.class().is(class)) {
700+
classify(class, vm).into()
701+
} else {
702+
PreSort::Generic
703+
}
670704
}
671705
}
672706

707+
fn pre_sort_check_tuples<'a>(
708+
first_elem: &PyObjectRef,
709+
keys: impl Iterator<Item = &'a PyObjectRef>,
710+
vm: &VirtualMachine,
711+
) -> PreSort {
712+
let class = first_elem.class();
713+
let mut all_same_type = true;
714+
715+
for k in keys {
716+
let Some(t) = k
717+
.downcast_ref_if_exact::<PyTuple>(vm)
718+
.filter(|t| !t.as_slice().is_empty())
719+
else {
720+
return PreSort::Generic;
721+
};
722+
if all_same_type && !t.as_slice()[0].class().is(class) {
723+
all_same_type = false;
724+
}
725+
}
726+
727+
let elem = if !all_same_type || class.is(vm.ctx.types.tuple_type) {
728+
Elem::Generic
729+
} else {
730+
classify(class, vm)
731+
};
732+
PreSort::Tuple(elem)
733+
}
734+
673735
fn str_lt(a: &PyObjectRef, b: &PyObjectRef) -> bool {
674736
a.downcast_ref::<PyStr>().unwrap().as_bytes() < b.downcast_ref::<PyStr>().unwrap().as_bytes()
675737
}
@@ -707,6 +769,37 @@ fn object_lt(
707769
}
708770
}
709771

772+
fn elem_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
773+
match elem {
774+
Elem::Str => Ok(str_lt(a, b)),
775+
Elem::Int => Ok(int_lt(a, b)),
776+
Elem::Float => Ok(float_lt(a, b)),
777+
Elem::Object(f) => object_lt(*f, a, b, vm),
778+
Elem::Generic => a.rich_compare_bool(b, PyComparisonOp::Lt, vm),
779+
}
780+
}
781+
782+
fn tuple_lt(elem: &Elem, a: &PyObjectRef, b: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
783+
let a = a.downcast_ref::<PyTuple>().unwrap().as_slice();
784+
let b = b.downcast_ref::<PyTuple>().unwrap().as_slice();
785+
786+
let mut i = 0;
787+
while i < a.len() && i < b.len() {
788+
if !a[i].rich_compare_bool(&b[i], PyComparisonOp::Eq, vm)? {
789+
break;
790+
}
791+
i += 1;
792+
}
793+
if i >= a.len() || i >= b.len() {
794+
return Ok(a.len() < b.len());
795+
}
796+
if i == 0 {
797+
elem_lt(elem, &a[0], &b[0], vm)
798+
} else {
799+
a[i].rich_compare_bool(&b[i], PyComparisonOp::Lt, vm)
800+
}
801+
}
802+
710803
fn timsort_specialized<T, K>(
711804
vm: &VirtualMachine,
712805
items: &mut [T],
@@ -750,6 +843,14 @@ where
750843
};
751844
object_lt(cmp, a, b, vm)
752845
}),
846+
PreSort::Tuple(elem) => timsort(items, &mut |a, b| {
847+
let (a, b) = if reverse {
848+
(key(b), key(a))
849+
} else {
850+
(key(a), key(b))
851+
};
852+
tuple_lt(&elem, a, b, vm)
853+
}),
753854
PreSort::Generic => timsort(items, &mut |a, b| {
754855
let (a, b) = if reverse {
755856
(key(b), key(a))

‎extra_tests/snippets/builtin_list.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,11 @@ class IntSub(int):
257257
nan = float("nan")
258258
assert repr(sorted([nan, 1.0, 2.0])) == "[nan, 1.0, 2.0]"
259259
assert sorted([b"b", b"a", b"c"]) == [b"a", b"b", b"c"]
260+
assert sorted([(2, 9), (1, 5), (2, 1)]) == [(1, 5), (2, 1), (2, 9)]
261+
assert sorted([(1, "b"), (1, "a")]) == [(1, "a"), (1, "b")]
262+
assert sorted([(1,), (1, 2), ()]) == [(), (1,), (1, 2)]
263+
assert sorted([((2,), "x"), ((1,), "y")]) == [((1,), "y"), ((2,), "x")]
264+
assert sorted([(1, "a"), (2.5, "b"), (0, "c")]) == [(0, "c"), (1, "a"), (2.5, "b")]
260265

261266
lst = [3, 1, 5, 2, 4]
262267

0 commit comments

Comments
 (0)