sort_const/lib.rs
1#![no_std]
2
3//! A macro for sorting arrays and slices at compile time.
4//!
5//! ## Usage
6//!
7//! Just use the [`const_quicksort`] or [`const_shellsort`] macros.
8//!
9//! ```
10//! use sort_const::const_quicksort;
11//!
12//! const fn lt(a: &u8, b: &u8) -> bool {
13//! *a < *b
14//! }
15//!
16//! const A: &[u8] = &const_quicksort!([3, 1, 2]);
17//! const B: &[u8] = &const_quicksort!([3, 1, 2], |a, b| *a < *b);
18//! const C: &[u8] = &const_quicksort!([3, 1, 2], lt);
19//!
20//! assert_eq!(A, [1, 2, 3]);
21//! assert_eq!(B, [1, 2, 3]);
22//! assert_eq!(C, [1, 2, 3]);
23//! ```
24
25use core::mem::forget;
26
27/// Used for our `call stack`
28#[doc(hidden)]
29pub use arrayvec_const::ArrayVec;
30
31/// A helper to turn everything into mutable slices
32#[doc(hidden)]
33pub struct Wrapper<T>(pub T);
34impl<T, const N: usize> Wrapper<[T; N]> {
35 #[inline(always)]
36 pub const fn as_mut_slice(&mut self) -> &mut [T] {
37 &mut self.0
38 }
39}
40impl<T> Wrapper<&'_ mut [T]> {
41 #[inline(always)]
42 pub const fn as_mut_slice(&mut self) -> &mut [T] {
43 self.0
44 }
45}
46impl<T, const N: usize> Wrapper<&'_ mut [T; N]> {
47 #[inline(always)]
48 pub const fn as_mut_slice(&mut self) -> &mut [T] {
49 self.0
50 }
51}
52impl<T, const N: usize> Wrapper<&'_ [T; N]> {
53 #[inline(always)]
54 pub const fn as_mut_slice(&mut self) -> &mut [T] {
55 panic!("Cannot sort non-mut `&`. Did you mean to use `&mut`?")
56 }
57}
58
59#[doc(hidden)]
60#[track_caller]
61#[inline]
62pub const fn expect_push<T, const LEN: usize>(v: &mut ArrayVec<T, LEN>, element: T) {
63 let res = v.try_push(element);
64 if res.is_err() {
65 panic!("ran out of capacity");
66 }
67 forget(res);
68}
69
70/// Some nice gaps for shellsort
71#[doc(hidden)]
72pub const A366726: [usize; 32] = [
73 1,
74 4,
75 9,
76 20,
77 45,
78 102,
79 230,
80 516,
81 1158,
82 2599,
83 5831,
84 13082,
85 29351,
86 65853,
87 147748,
88 331490,
89 743735,
90 1668650,
91 3743800,
92 8399623,
93 18845471,
94 42281871,
95 94863989,
96 212837706,
97 477524607,
98 1071378536,
99 2403754591,
100 5393085583,
101 12099975682,
102 27147615084,
103 60908635199,
104 136655165852,
105];
106
107/// Sorts a `const` array or mutable slice as a `const` expression using quicksort.
108///
109/// Can be called with just the data to be sorted, in which case elements are compared by `a < b` resulting in the smallest element at
110/// the head of the data and the largest at the tail.
111///
112/// ```
113/// use sort_const::const_quicksort;
114///
115/// const U8S: &[u8] = &const_quicksort!([3, 1, 2]);
116///
117/// assert_eq!(U8S, [1, 2, 3]);
118/// ```
119///
120/// Can also be called with a lambda-like expression (e.g. `|a, b| a < b`) where `true` identifies that `a` should come before `b`.
121///
122/// ```
123/// use sort_const::const_quicksort;
124///
125/// #[derive(Debug, PartialEq)]
126/// struct Foo(u8);
127///
128/// const FOOS_MUT_REF: &[Foo] = &{
129/// let mut foo = [Foo(1), Foo(2), Foo(4), Foo(3)];
130/// const_quicksort!(foo.split_last_mut().expect("failed").1, |a, b| a.0 > b.0);
131/// foo
132/// };
133///
134/// assert_eq!(FOOS_MUT_REF, [4, 2, 1, 3].map(Foo));
135/// ```
136///
137/// Can also be called with the name of a `const` function which must return a boolean and which will be evaluated over `&data[a], &data[b]`.
138///
139/// ```
140/// use sort_const::const_quicksort;
141///
142/// #[derive(Debug, PartialEq)]
143/// struct Foo(u8);
144///
145/// const fn gt(lhs: &Foo, rhs: &Foo) -> bool {
146/// lhs.0 > rhs.0
147/// }
148/// const FOOS_MUT_REF: &[Foo] = &{
149/// let mut foo = [Foo(1), Foo(2), Foo(4), Foo(3)];
150/// const_quicksort!(foo.split_last_mut().expect("failed").1, gt);
151/// foo
152/// };
153///
154/// assert_eq!(FOOS_MUT_REF, [4, 2, 1, 3].map(Foo));
155/// ```
156///
157/// The `@depth` parameter should only be used if you encounter a scenario where "stack overflows" start occuring.
158/// ```compile_fail
159/// use sort_const::const_quicksort;
160///
161/// const SORTED_ARRAY: &[u32] = &{
162/// let mut data = [0_u32; 1000];
163/// let mut i = 0;
164/// while i < data.len() {
165/// if i & 1 == 0 {
166/// data[i] = i as u32;
167/// }
168/// i += 1;
169/// }
170/// const_quicksort!(@8, &mut data);
171/// data
172/// };
173/// ```
174///
175#[macro_export]
176macro_rules! const_quicksort {
177 ($(@$depth:literal,)? $data:expr) => {{
178 macro_rules! cmp {
179 ($a:expr, $b:expr) => {{ $a < $b }};
180 }
181 $crate::const_quicksort_adv!($(@$depth,)? $data, cmp)
182 }};
183 ($(@$depth:literal,)? $data:expr, |$a:ident, $b:ident| $cmp:expr) => {{
184 macro_rules! cmp {
185 ($a_:expr, $b_:expr) => {{
186 let ($a, $b) = (&$a_, &$b_);
187 $cmp
188 }};
189 }
190 $crate::const_quicksort_adv!($(@$depth,)? $data, cmp)
191 }};
192 ($(@$depth:literal,)? $data:expr, $cmp:expr) => {{
193 macro_rules! cmp {
194 ($a:expr, $b:expr) => {{ $cmp(&$a, &$b) }};
195 }
196 $crate::const_quicksort_adv!($(@$depth,)? $data, cmp)
197 }};
198}
199
200/// Sorts a `const` array or mutable slice as a `const` expression using quicksort.
201///
202/// The optional `@` prefix parameter is the max stack size for recursion.
203///
204/// The first parameter is the data to be sorted.
205///
206/// The second parameter is the path identifying the macro which should be invoked to compare elements `cmp!(a, b)`
207#[macro_export]
208macro_rules! const_quicksort_adv {
209 ($data:expr, $cmp:path) => { $crate::const_quicksort_adv!(@1024, $data, $cmp) };
210 (@$depth:literal, $data:expr, $cmp:path) => {{
211 let mut data = $crate::Wrapper($data);
212 let mut slices = $crate::ArrayVec::<&mut [_], $depth>::new();
213 $crate::expect_push(&mut slices, data.as_mut_slice());
214 while let Some(data) = slices.pop() {
215 match data.len() {
216 0 | 1 => continue,
217 2 => {
218 if !{$cmp!(data[0], data[1])} {
219 data.swap(0, 1);
220 }
221 continue;
222 },
223 _ => {}
224 }
225
226
227 let (pivot, rest) = data
228 .split_first_mut()
229 .expect("slice is not empty, as verified above");
230
231 let mut left = 0;
232 let mut right = rest.len() - 1;
233 while left <= right {
234 if {$cmp!(rest[left], *pivot)} {
235 left += 1;
236 } else if !{$cmp!(rest[right], *pivot)} {
237 if right == 0 {
238 break;
239 }
240 right -= 1;
241 } else {
242 rest.swap(left, right);
243 left += 1;
244 if right == 0 {
245 break;
246 }
247 right -= 1;
248 }
249 }
250
251 data.swap(0, left);
252 let (left, right) = data.split_at_mut(left);
253 match right.split_first_mut() {
254 None => {
255 $crate::expect_push(&mut slices, left);
256 }
257 Some((_pivot, right)) if left.len() >= right.len() => {
258 $crate::expect_push(&mut slices, left);
259 $crate::expect_push(&mut slices, right);
260 }
261 Some((_pivot, right)) => {
262 $crate::expect_push(&mut slices, right);
263 $crate::expect_push(&mut slices, left);
264 }
265 }
266 }
267 ::core::mem::forget(slices);
268
269 data.0
270 }};
271}
272
273/// Sorts a `const` array or mutable slice as a `const` expression using quicksort.
274///
275/// Can be called with just the data to be sorted, in which case elements are compared by `a < b` resulting in the smallest element at
276/// the head of the data and the largest at the tail.
277///
278/// ```
279/// use sort_const::const_shellsort;
280///
281/// const U8S: &[u8] = &const_shellsort!([3, 1, 2]);
282///
283/// assert_eq!(U8S, [1, 2, 3]);
284/// ```
285///
286/// Can also be called with a lambda-like expression (e.g. `|a, b| a < b`) where `true` identifies that `a` should come before `b`.
287///
288/// ```
289/// use sort_const::const_shellsort;
290///
291/// #[derive(Debug, PartialEq)]
292/// struct Foo(u8);
293///
294/// const FOOS_MUT_REF: &[Foo] = &{
295/// let mut foo = [Foo(1), Foo(2), Foo(4), Foo(3)];
296/// const_shellsort!(foo.split_last_mut().expect("failed").1, |a, b| a.0 > b.0);
297/// foo
298/// };
299///
300/// assert_eq!(FOOS_MUT_REF, [4, 2, 1, 3].map(Foo));
301/// ```
302///
303/// Can also be called with the name of a `const` function which must return a boolean and which will be evaluated over `&data[a], &data[b]`.
304#[macro_export]
305macro_rules! const_shellsort {
306 ($(@$seq:expr,)? $data:expr) => {{
307 macro_rules! cmp {
308 ($a:expr, $b:expr) => {{ $a < $b }};
309 }
310 $crate::const_shellsort_adv!($(@$seq,)? $data, cmp)
311 }};
312 ($(@$seq:expr,)? $data:expr, |$a:ident, $b:ident| $cmp:expr) => {{
313 macro_rules! cmp {
314 ($a_:expr, $b_:expr) => {{
315 let ($a, $b) = (&$a_, &$b_);
316 $cmp
317 }};
318 }
319 $crate::const_shellsort_adv!($(@$seq,)? $data, cmp)
320 }};
321 ($(@$seq:expr,)? $data:expr, $cmp:expr) => {{
322 macro_rules! cmp {
323 ($a:expr, $b:expr) => {{ $cmp(&$a, &$b) }};
324 }
325 $crate::const_shellsort_adv!($(@$seq,)? $data, cmp)
326 }};
327}
328
329/// Sorts a `const` array or mutable slice as a `const` expression using shellsort.
330///
331/// The optional `@` prefix parameter is the sort.
332///
333/// The first parameter is the data to be sorted.
334///
335/// The second parameter is the path identifying the macro which should be invoked to compare elements `cmp!(a, b)`
336#[macro_export]
337macro_rules! const_shellsort_adv {
338 ($data:expr, $cmp:path) => { $crate::const_shellsort_adv!(@$crate::A366726, $data, $cmp) };
339 (@$seq:expr, $data:expr, $cmp:path) => {{
340 const GAPS: &[usize] = &$seq;
341 const GAPS_LEN: usize = GAPS.len();
342
343 let mut data = $crate::Wrapper($data);
344
345 let arr = data.as_mut_slice();
346 let n = arr.len();
347 let mut gaps = {
348 let mut gaps = $crate::ArrayVec::<usize, GAPS_LEN>::new();
349 let mut i = 0;
350 while i < GAPS_LEN && GAPS[i] < arr.len() {
351 $crate::expect_push(&mut gaps, GAPS[i]);
352 i += 1;
353 }
354 gaps
355 };
356
357 // Start with the largest gap and work down to a gap of 1
358 while let Some(gap) = gaps.pop() {
359 // Do a gapped insertion sort for every element in gaps
360 // Each loop leaves a[0..gap-1] in gapped order
361 let mut i = gap;
362 while i < arr.len() {
363 // save a[i] in temp and make a hole at position i
364 // This is safe because we are holding a temporary over a series of swaps.
365 let temp = unsafe { ::core::ptr::read(&arr[i]) };
366
367 // shift earlier gap-sorted elements up until the correct location for a[i] is found
368 let mut j = i;
369 while j >= gap && {$cmp!(temp, arr[j - gap])} {
370 // This is safe because the previous item is either saved away in `temp` or
371 // was previously copied to it's new location
372 unsafe { ::core::ptr::copy(&arr[j - gap], &mut arr[j], 1); }
373 j -= gap
374 }
375
376 // put temp (the original a[i]) in its correct location
377 // This is safe because we writing the last swap.
378 unsafe { ::core::ptr::write(&mut arr[j], temp); }
379 i += 1;
380 }
381 }
382 core::mem::forget(gaps);
383 data.0
384 }};
385}
386
387#[cfg(test)]
388mod test {
389 #[derive(Debug)]
390 struct Foo(u8);
391
392 const fn lt(a: &u8, b: &u8) -> bool {
393 *a < *b
394 }
395
396 const U8S: &[u8] = &const_quicksort!([3, 2, 1]);
397 const U8S_FN: &[u8] = &const_quicksort!([3, 2, 1], lt);
398 const FOOS: &[Foo] = &const_quicksort!([Foo(1), Foo(2), Foo(4), Foo(3)], |a, b| a.0 > b.0);
399 const FOOS_MUT_REF: &[Foo] = &{
400 let mut foo = [Foo(1), Foo(2), Foo(4), Foo(3)];
401 const_quicksort!(foo.split_last_mut().expect("failed").1, |a, b| a.0 > b.0);
402 foo
403 };
404
405 #[test]
406 fn test_u8() {
407 assert_eq!(U8S, [1, 2, 3]);
408 }
409
410 #[test]
411 fn test_u8_fn() {
412 assert_eq!(U8S_FN, [1, 2, 3]);
413 }
414
415 #[test]
416 fn test_foo() {
417 assert!(FOOS.iter().map(|v| v.0).eq([4, 3, 2, 1]));
418 assert!(FOOS_MUT_REF.iter().map(|v| v.0).eq([4, 2, 1, 3]));
419 }
420}