atools/
lib.rs

1//! a collection of useful features for working with arrays
2#![cfg_attr(not(test), no_std)]
3#![allow(incomplete_features, internal_features)]
4#![feature(
5    const_destruct,
6    adt_const_params,
7    generic_const_exprs,
8    core_intrinsics,
9    iter_intersperse,
10    const_trait_impl,
11    maybe_uninit_array_assume_init,
12    array_windows,
13    iter_map_windows
14)]
15#![warn(
16    clippy::undocumented_unsafe_blocks,
17    clippy::missing_const_for_fn,
18    clippy::missing_safety_doc,
19    clippy::suboptimal_flops,
20    unsafe_op_in_unsafe_fn,
21    clippy::dbg_macro,
22    clippy::use_self,
23    missing_docs
24)]
25use core::{
26    array::from_fn,
27    intrinsics::transmute_unchecked,
28    marker::Destruct,
29    mem::{offset_of, MaybeUninit as MU},
30};
31pub mod pervasive;
32mod slice;
33mod tuple;
34#[doc(inline)]
35pub use slice::Slice;
36pub use tuple::*;
37
38/// The prelude. You should
39/// ```
40/// use atools::prelude::*;
41/// ```
42pub mod prelude {
43    #[doc(inline)]
44    pub use super::{
45        pervasive::prelude::*, range, slice::r, slice::Slice, splat, Array, ArrayTools, Chunked,
46        CollectArray, Couple, Deconstruct, Deconstruct_, Flatten, Join, Split, Tuple, Zip,
47    };
48    #[doc(inline)]
49    pub use core::array::from_fn;
50}
51
52#[repr(C)]
53struct Pair<X, Y>(X, Y);
54impl<X, Y> Pair<X, Y> {
55    const fn tuple() -> bool {
56        (size_of::<(X, Y)>() == size_of::<Self>())
57            & (offset_of!(Self, 0) == offset_of!((X, Y), 0))
58            & (offset_of!(Self, 1) == offset_of!((X, Y), 1))
59    }
60
61    const fn into(self) -> (X, Y) {
62        if Self::tuple() {
63            // SAFETY: we are a tuple!!!
64            unsafe { transmute_unchecked::<Self, (X, Y)>(self) }
65        } else {
66            // SAFETY: this is safe.
67            let out = unsafe { (core::ptr::read(&self.0), core::ptr::read(&self.1)) };
68            core::mem::forget(self);
69            out
70        }
71    }
72
73    const unsafe fn splat<T>(x: T) -> (X, Y) {
74        assert!(core::mem::size_of::<T>() == core::mem::size_of::<Pair<X, Y>>());
75        // SAFETY: well.
76        unsafe { transmute_unchecked::<_, Self>(x) }.into()
77    }
78}
79
80/// Convenience function for when clonage is required; prefer `[T; N]` if possible. Also useful if `N` should be inferred.
81pub fn splat<T: Clone, const N: usize>(a: T) -> [T; N] {
82    from_fn(|_| a.clone())
83}
84
85/// Creates a array of indices.
86/// ```
87/// # #![feature(generic_const_exprs)]
88/// # use atools::prelude::*;
89/// assert_eq!(range::<5>(), [0, 1, 2, 3, 4]);
90/// ```
91pub const fn range<const N: usize>() -> [usize; N] {
92    let mut out = unsafe { MU::<[MU<usize>; N]>::uninit().assume_init() };
93    let mut i = 0usize;
94    while i < out.len() {
95        out[i] = MU::new(i);
96        i += 1;
97    }
98    unsafe { transmute_unchecked(out) }
99}
100
101/// Collect an iterator into a array.
102pub trait CollectArray<T> {
103    /// Collect an iterator into a array.
104    ///
105    /// # Panics
106    ///
107    /// if the array isn't big enough.
108    fn carr<const N: usize>(&mut self) -> [T; N];
109}
110
111impl<T, I: Iterator<Item = T>> CollectArray<T> for I {
112    fn carr<const N: usize>(&mut self) -> [T; N] {
113        from_fn(|_| self.next().unwrap())
114    }
115}
116
117/// Deconstruct some array.
118/// Use
119/// ```
120/// let [t, arr @ ..] = [1, 2];
121/// ```
122/// when possible. If the length of the array is a const generic, use
123/// ```
124/// # #![feature(generic_const_exprs)]
125/// # use atools::prelude::*;
126/// let (t, arr) = [1, 2].uncons();
127/// ```
128#[const_trait]
129pub trait Deconstruct<T, const N: usize> {
130    /// Gives you the <code>[[head](Deconstruct_::head), [tail](Deconstruct_::tail) @ ..]</code>
131    /// ```
132    /// # #![feature(generic_const_exprs)]
133    /// # use atools::prelude::*;
134    /// let (t, arr) = b"abc".uncons();
135    /// # assert_eq!(t, b'a');
136    /// # assert_eq!(arr, *b"bc");
137    /// ```
138    fn uncons(self) -> (T, [T; N - 1]);
139    /// Gives you the <code>[[init](Deconstruct_::init) @ .., [last](Deconstruct_::last)]</code>
140    /// ```
141    /// # #![feature(generic_const_exprs)]
142    /// # use atools::prelude::*;
143    /// let (arr, t) = [0.1f32, 0.2, 0.3].unsnoc();
144    /// # assert_eq!(arr, [0.1, 0.2]);
145    /// assert_eq!(t, 0.3);
146    /// ```
147    fn unsnoc(self) -> ([T; N - 1], T);
148}
149
150/// Deconstruct some array. (dropping edition).
151///
152/// <img src="https://media.discordapp.net/attachments/273541645579059201/1404772577259294770/listmonster.png?ex=689c67e9&is=689b1669&hm=00525f7bb8ffb2eb096a46d10509ebf8def669ca3175a713df686ff4be7a4e67">
153pub trait Deconstruct_<T, const N: usize> {
154    /// Gives you a <code>[[_](Deconstruct_::init) @ .., last]</code>.
155    /// See also [`unsnoc`](Deconstruct::unsnoc).
156    fn last(self) -> T;
157    /// Gives you a <code>[init @ .., [_](Deconstruct_::last)]</code>
158    /// See also [`unsnoc`](Deconstruct::unsnoc).
159    /// ```
160    /// # #![feature(generic_const_exprs)]
161    /// # use atools::prelude::*;
162    /// let a = *[1u64, 2, 3].init();
163    /// assert_eq!(a, [1, 2]);
164    /// ```
165    fn init(self) -> [T; N - 1];
166    /// Gives you a <code>[head, [_](Deconstruct_::tail) @ ..]</code>.
167    /// See also [`uncons`](Deconstruct::uncons).
168    fn head(self) -> T;
169    /// Gives you a <code>[[_](Deconstruct_::head), tail @ ..]</code>.
170    /// See also [`uncons`](Deconstruct::uncons).
171    /// ```
172    /// # #![feature(generic_const_exprs)]
173    /// # use atools::prelude::*;
174    /// let x = atools::range::<5>();
175    /// assert!(*x.tail() == atools::range::<4>().map(|x| x + 1));
176    /// ```
177    fn tail(self) -> [T; N - 1];
178}
179
180impl<T, const N: usize> const Deconstruct<T, N> for [T; N] {
181    #[doc(alias = "pop_front")]
182    fn uncons(self) -> (T, [T; N - 1]) {
183        // SAFETY: the layout is alright.
184        unsafe { Pair::splat(self) }
185    }
186
187    #[doc(alias = "pop")]
188    fn unsnoc(self) -> ([T; N - 1], T) {
189        // SAFETY: the layout is still alright.
190        unsafe { Pair::splat(self) }
191    }
192}
193
194impl<T, const N: usize> Deconstruct_<T, N> for [T; N]
195where
196    [(); N - 1]:,
197{
198    fn last(self) -> T {
199        self.unsnoc().1
200    }
201    #[doc(alias = "trunc")]
202    fn init(self) -> [T; N - 1] {
203        self.unsnoc().0
204    }
205    fn head(self) -> T {
206        self.uncons().0
207    }
208    fn tail(self) -> [T; N - 1] {
209        self.uncons().1
210    }
211}
212
213/// Join scalars together.
214#[const_trait]
215pub trait Join<T, const N: usize, const O: usize, U> {
216    /// Join a array and an scalar together. For joining two arrays together, see [`Couple`].
217    /// ```
218    /// # #![feature(generic_const_exprs)]
219    /// # use atools::prelude::*;
220    /// let a = [1, 2].join(3);
221    /// let b = 1.join([2, 3]);
222    /// let c = 1.join(2).join(3);
223    /// ```
224    fn join(self, with: U) -> [T; N + O];
225}
226
227/// Couple two arrays together.
228#[const_trait]
229pub trait Couple<T, const N: usize, const O: usize> {
230    /// Couple two arrays together. This could have been [`Join`], but the methods would require disambiguation.
231    /// ```
232    /// # #![feature(generic_const_exprs)]
233    /// # use atools::prelude::*;
234    /// let a = 1.join(2).couple([3, 4]);
235    /// ```
236    fn couple(self, with: [T; O]) -> [T; N + O];
237}
238
239impl<T, const N: usize, const O: usize> const Couple<T, N, O> for [T; N] {
240    fn couple(self, with: [T; O]) -> [T; N + O] {
241        // SAFETY: adjacent
242        unsafe { transmute_unchecked(Pair(self, with)) }
243    }
244}
245
246impl<T, const N: usize> const Join<T, N, 1, T> for [T; N] {
247    fn join(self, with: T) -> [T; N + 1] {
248        self.couple([with])
249    }
250}
251
252impl<T> const Join<T, 1, 1, T> for T {
253    fn join(self, with: T) -> [T; 2] {
254        [self, with]
255    }
256}
257
258impl<T, const O: usize> const Join<T, 1, O, [T; O]> for T {
259    fn join(self, with: [T; O]) -> [T; 1 + O] {
260        [self].couple(with)
261    }
262}
263
264pub(crate) const fn assert_zero(x: usize) -> usize {
265    if x != 0 {
266        panic!("expected zero");
267    } else {
268        0
269    }
270}
271
272/// 🍪
273#[allow(private_bounds)]
274#[const_trait]
275pub trait Chunked<T, const N: usize> {
276    /// Chunks.
277    /// This will compile fail if `N ∤ (does not divide) C`
278    /// ```
279    /// # #![feature(generic_const_exprs)]
280    /// # use atools::prelude::*;
281    /// assert_eq!(range::<6>().chunked::<3>(), [[0, 1, 2], [3, 4, 5]]);
282    /// ```
283    #[allow(private_bounds)]
284    fn chunked<const C: usize>(self) -> [[T; C]; N / C]
285    where
286        // N % C == 0
287        [(); assert_zero(N % C)]:;
288}
289
290impl<const N: usize, T> const Chunked<T, N> for [T; N] {
291    #[allow(private_bounds)]
292    fn chunked<const C: usize>(self) -> [[T; C]; N / C]
293    where
294        [(); assert_zero(N % C)]:,
295    {
296        // SAFETY: N != 0 && wont leak as N % C == 0.
297        let retval = unsafe { self.as_ptr().cast::<[[T; C]; N / C]>().read() };
298        core::mem::forget(self);
299        retval
300    }
301}
302
303/// Flatten arrays.
304#[const_trait]
305pub trait Flatten<T, const N: usize, const N2: usize> {
306    /// Takes a `[[T; N]; N2]`, and flattens it to a `[T; N * N2]`.
307    ///
308    /// # Examples
309    ///
310    /// ```
311    /// # #![feature(generic_const_exprs)]
312    /// # use atools::prelude::*;
313    /// assert_eq!([[1, 2, 3], [4, 5, 6]].flatten(), [1, 2, 3, 4, 5, 6]);
314    ///
315    /// assert_eq!(
316    ///     [[1, 2, 3], [4, 5, 6]].flatten(),
317    ///     [[1, 2], [3, 4], [5, 6]].flatten(),
318    /// );
319    ///
320    /// let array_of_empty_arrays: [[i32; 0]; 5] = [[], [], [], [], []];
321    /// assert!(array_of_empty_arrays.flatten().is_empty());
322    ///
323    /// let empty_array_of_arrays: [[u32; 10]; 0] = [];
324    /// assert!(empty_array_of_arrays.flatten().is_empty());
325    /// ```
326    fn flatten(self) -> [T; N * N2];
327}
328
329impl<T, const N: usize, const N2: usize> const Flatten<T, N, N2> for [[T; N]; N2] {
330    fn flatten(self) -> [T; N * N2] {
331        // SAFETY: layout is the same.
332        unsafe { core::intrinsics::transmute_unchecked(self) }
333    }
334}
335
336#[const_trait]
337/// Splitting arrays up.
338pub trait Split<T, const N: usize> {
339    /// Splits the array into twain.
340    /// ```
341    /// # #![feature(generic_const_exprs)]
342    /// # use atools::prelude::*;
343    /// let x = [1u8, 2, 3];
344    /// let ([x], [y, z]) = x.split::<1>();
345    /// ```
346    fn split<const AT: usize>(self) -> ([T; AT], [T; N - AT]);
347    /// Take `AT` elements, discarding the rest.
348    /// ```
349    /// # #![feature(generic_const_exprs)]
350    /// # use atools::prelude::*;
351    /// assert_eq!(range::<50>().take::<5>(), range::<5>());
352    /// ```
353    fn take<const AT: usize>(self) -> [T; AT]
354    where
355        [(); N - AT]:,
356        T: [const] Destruct;
357    /// Discard `AT` elements, returning the rest.
358    fn drop<const AT: usize>(self) -> [T; N - AT]
359    where
360        T: [const] Destruct;
361}
362
363impl<T, const N: usize> const Split<T, N> for [T; N] {
364    fn split<const AT: usize>(self) -> ([T; AT], [T; N - AT]) {
365        // SAFETY: N - AT overflows when AT > N so the size of the returned "array" is the same.
366        unsafe { Pair::splat(self) }
367    }
368    fn take<const M: usize>(self) -> [T; M]
369    where
370        // M <= N
371        [(); N - M]:,
372        T: [const] Destruct,
373    {
374        self.split::<M>().0
375    }
376    fn drop<const M: usize>(self) -> [T; N - M]
377    where
378        T: [const] Destruct,
379    {
380        self.split::<M>().1
381    }
382}
383#[const_trait]
384/// Zip arrays together.
385pub trait Zip<T, const N: usize> {
386    /// Zip arrays together.
387    fn zip<U>(self, with: [U; N]) -> [(T, U); N];
388}
389
390impl<T, const N: usize> const Zip<T, N> for [T; N] {
391    fn zip<U>(self, with: [U; N]) -> [(T, U); N] {
392        let mut out = unsafe { MU::<[MU<_>; N]>::uninit().assume_init() };
393        let mut i = 0usize;
394        while i < out.len() {
395            out[i] = MU::new(unsafe { (self.as_ptr().add(i).read(), with.as_ptr().add(i).read()) });
396            i += 1;
397        }
398        core::mem::forget((self, with));
399        unsafe { transmute_unchecked(out) }
400    }
401}
402
403/// Array tools.
404pub trait ArrayTools<T, const N: usize> {
405    /// Skip `BY` elements.
406    fn skip<const BY: usize>(self) -> [T; N - BY];
407    /// Skip every `BY` elements.
408    ///
409    /// ```
410    /// # #![feature(generic_const_exprs)]
411    /// # use atools::prelude::*;
412    /// let x = range::<5>().step::<2>();
413    /// assert_eq!(x, [0, 2, 4]);
414    /// let x = range::<20>().step::<5>();
415    /// assert_eq!(x, [0, 5, 10, 15]);
416    /// assert_eq!(range::<50>().step::<3>(), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48]);
417    /// ```
418    fn step<const STEP: usize>(self) -> [T; 1 + (N - 1) / (STEP)];
419    /// Intersperse a element in between items.
420    /// ```
421    /// # #![feature(generic_const_exprs)]
422    /// # use atools::prelude::*;
423    /// let x = range::<3>().intersperse(5);
424    /// assert_eq!(x, [0, 5, 1, 5, 2]);
425    /// ```
426    fn intersperse(self, with: T) -> [T; (N * 2) - 1]
427    where
428        T: Clone;
429    /// Run a function on every element.
430    fn each(self, apply: impl FnMut(T));
431    /// Embed the index.
432    fn enumerate(self) -> [(T, usize); N];
433    /// Get the sliding windows of this array.
434    /// ```
435    /// # #![feature(generic_const_exprs)]
436    /// # use atools::prelude::*;
437    /// assert_eq!(range::<5>().windowed::<2>(), [&[0, 1], &[1, 2], &[2, 3], &[3, 4]]);
438    /// ```
439    fn windowed<const W: usize>(&self) -> [&[T; W]; N - W + 1];
440    /// Inspect every element of this array.
441    fn inspect(self, f: impl FnMut(&T)) -> Self;
442    /// Reverse this array.
443    fn rev(self) -> Self;
444    /// Interleave items from two arrays.
445    /// ```
446    /// # #![feature(generic_const_exprs)]
447    /// # use atools::prelude::*;
448    /// assert_eq!([0u8, 2, 4].interleave([1, 3, 5]), [0, 1, 2, 3, 4, 5]);
449    /// ```
450    fn interleave(self, with: [T; N]) -> [T; N * 2];
451    /// [Cartesian product](https://en.wikipedia.org/wiki/Cartesian_product) (`A  ×  B`) of two arrays.
452    /// ```
453    /// # #![feature(generic_const_exprs)]
454    /// # use atools::prelude::*;
455    /// assert_eq!([1u64, 2].cartesian_product(&["Π", "Σ"]), [(1, "Π"), (1, "Σ"), (2, "Π"), (2, "Σ")]);
456    /// ```
457    fn cartesian_product<U: Clone, const M: usize>(&self, with: &[U; M]) -> [(T, U); N + M]
458    where
459        T: Clone;
460    /// Sorts it. This uses <code>[[T](slice)]::[sort_unstable](slice::sort_unstable)</code>.
461    fn sort(self) -> Self
462    where
463        T: Ord;
464    /// Sum of the array.
465    fn sum(self) -> T
466    where
467        T: core::iter::Sum<T>;
468    /// Product of the array.
469    fn product(self) -> T
470    where
471        T: core::iter::Product<T>;
472}
473
474impl<T, const N: usize> ArrayTools<T, N> for [T; N] {
475    fn skip<const BY: usize>(self) -> [T; N - BY] {
476        self.into_iter().skip(BY).carr()
477    }
478    fn step<const STEP: usize>(self) -> [T; 1 + (N - 1) / (STEP)] {
479        self.into_iter().step_by(STEP).carr()
480    }
481    fn intersperse(self, with: T) -> [T; (N * 2) - 1]
482    where
483        T: Clone,
484    {
485        self.into_iter().intersperse(with).carr()
486    }
487
488    fn each(self, apply: impl FnMut(T)) {
489        self.into_iter().for_each(apply);
490    }
491
492    fn enumerate(self) -> [(T, usize); N] {
493        let mut n = 0;
494        self.map(|x| {
495            let o = n;
496            n += 1;
497            (x, o)
498        })
499    }
500
501    fn windowed<const W: usize>(&self) -> [&[T; W]; N - W + 1] {
502        self.array_windows().carr()
503    }
504
505    fn inspect(self, f: impl FnMut(&T)) -> Self {
506        self.iter().for_each(f);
507        self
508    }
509
510    fn rev(mut self) -> Self {
511        self.reverse();
512        self
513    }
514
515    fn interleave(self, with: [T; N]) -> [T; N * 2] {
516        let mut which = true;
517        let mut a = self.into_iter();
518        let mut b = with.into_iter();
519        from_fn(|_| {
520            which = !which;
521            match which {
522                false => a.next().unwrap(),
523                true => b.next().unwrap(),
524            }
525        })
526    }
527
528    fn cartesian_product<U: Clone, const M: usize>(&self, with: &[U; M]) -> [(T, U); N + M]
529    where
530        T: Clone,
531    {
532        self.iter()
533            .flat_map(|a| with.iter().map(move |b| (a.clone(), b.clone())))
534            .carr()
535    }
536
537    fn sort(mut self) -> Self
538    where
539        T: Ord,
540    {
541        <[T]>::sort_unstable(&mut self);
542        self
543    }
544
545    fn sum(self) -> T
546    where
547        T: core::iter::Sum<T>,
548    {
549        self.into_iter().sum()
550    }
551
552    fn product(self) -> T
553    where
554        T: core::iter::Product<T>,
555    {
556        self.into_iter().product()
557    }
558}