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}