@@ -9,7 +9,7 @@ use crate::common::lock::{
99use crate :: object:: { Traverse , TraverseFn } ;
1010use 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+
641649enum 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+
649684fn 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+
673735fn 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+
710803fn 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) )
0 commit comments