@@ -102,6 +102,7 @@ use core::mem::size_of;
102
102
use core:: mem;
103
103
use core:: ptr;
104
104
use core:: slice as core_slice;
105
+ use core:: { u8, u16, u32} ;
105
106
106
107
use borrow:: { Borrow , BorrowMut , ToOwned } ;
107
108
use boxed:: Box ;
@@ -1329,19 +1330,31 @@ impl<T> [T] {
1329
1330
pub fn sort_by_key < K , F > ( & mut self , f : F )
1330
1331
where F : FnMut ( & T ) -> K , K : Ord
1331
1332
{
1332
- let mut indices: Vec < _ > = self . iter ( ) . map ( f) . enumerate ( ) . map ( |( i, k) | ( k, i) ) . collect ( ) ;
1333
- // The elements of `indices` are unique, as they are indexed, so any sort will be stable
1334
- // with respect to the original slice. We use `sort_unstable` here because it requires less
1335
- // memory allocation.
1336
- indices. sort_unstable ( ) ;
1337
- for i in 0 ..self . len ( ) {
1338
- let mut index = indices[ i] . 1 ;
1339
- while index < i {
1340
- index = indices[ index] . 1 ;
1341
- }
1342
- indices[ i] . 1 = index;
1343
- self . swap ( i, index) ;
1333
+ // Helper macro for indexing our vector by the smallest possible type, to reduce allocation.
1334
+ macro_rules! sort_by_key {
1335
+ ( $t: ty, $slice: ident, $f: ident) => ( {
1336
+ let mut indices: Vec <_> =
1337
+ $slice. iter( ) . map( $f) . enumerate( ) . map( |( i, k) | ( k, i as $t) ) . collect( ) ;
1338
+ // The elements of `indices` are unique, as they are indexed, so any sort will be
1339
+ // stable with respect to the original slice. We use `sort_unstable` here because it
1340
+ // requires less memory allocation.
1341
+ indices. sort_unstable( ) ;
1342
+ for i in 0 ..$slice. len( ) {
1343
+ let mut index = indices[ i] . 1 ;
1344
+ while ( index as usize ) < i {
1345
+ index = indices[ index as usize ] . 1 ;
1346
+ }
1347
+ indices[ i] . 1 = index;
1348
+ $slice. swap( i, index as usize ) ;
1349
+ }
1350
+ } )
1344
1351
}
1352
+
1353
+ let len = self . len ( ) ;
1354
+ if len <= ( u8:: MAX as usize ) { return sort_by_key ! ( u8 , self , f) }
1355
+ if len <= ( u16:: MAX as usize ) { return sort_by_key ! ( u16 , self , f) }
1356
+ if len <= ( u32:: MAX as usize ) { return sort_by_key ! ( u32 , self , f) }
1357
+ sort_by_key ! ( usize , self , f)
1345
1358
}
1346
1359
1347
1360
/// Sorts the slice, but may not preserve the order of equal elements.
0 commit comments