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, 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/// <img src="https://media.discordapp.net/attachments/1190100233233895585/1430294602144813167/listmonster.png?ex=68f94126&is=68f7efa6&hm=1e2f1d83a8348d1369eb33c5f4f62aa5787eb05eec3d782bf49d93f84b04d94a&=&width=710&height=355">
129#[const_trait]
130pub trait Deconstruct<T, const N: usize> {
131    /// Gives you the <code>[[head](Deconstruct_::head), [tail](Deconstruct_::tail) @ ..]</code>
132    /// ```
133    /// # #![feature(generic_const_exprs)]
134    /// # use atools::prelude::*;
135    /// let (t, arr) = b"abc".uncons();
136    /// # assert_eq!(t, b'a');
137    /// # assert_eq!(arr, *b"bc");
138    /// ```
139    fn uncons(self) -> (T, [T; N - 1]);
140    /// Gives you the <code>[[init](Deconstruct_::init) @ .., [last](Deconstruct_::last)]</code>
141    /// ```
142    /// # #![feature(generic_const_exprs)]
143    /// # use atools::prelude::*;
144    /// let (arr, t) = [0.1f32, 0.2, 0.3].unsnoc();
145    /// # assert_eq!(arr, [0.1, 0.2]);
146    /// assert_eq!(t, 0.3);
147    /// ```
148    fn unsnoc(self) -> ([T; N - 1], T);
149    /// Gives you a <code>[init @ .., [_](Deconstruct_::last)]</code>
150    /// See also [`unsnoc`](Deconstruct::unsnoc).
151    /// ```
152    /// # #![feature(generic_const_exprs)]
153    /// # use atools::prelude::*;
154    /// let a = *[1u64, 2, 3].init();
155    /// assert_eq!(a, [1, 2]);
156    /// ```
157    fn init(self) -> [T; N - 1];
158    /// Gives you a <code>[[_](Deconstruct_::head), tail @ ..]</code>.
159    /// See also [`uncons`](Deconstruct::uncons).
160    /// ```
161    /// # #![feature(generic_const_exprs)]
162    /// # use atools::prelude::*;
163    /// let x = atools::range::<5>();
164    /// assert!(*x.tail() == atools::range::<4>().map(|x| x + 1));
165    /// ```
166    fn tail(self) -> [T; N - 1];
167    /// Gives you a <code>[[_](Deconstruct_::init) @ .., last]</code>.
168    /// See also [`unsnoc`](Deconstruct::unsnoc).
169    fn last(self) -> T
170    where
171        [(); N - 1]:;
172
173    /// Gives you a <code>[head, [_](Deconstruct_::tail) @ ..]</code>.
174    /// See also [`uncons`](Deconstruct::uncons).
175    fn head(self) -> T
176    where
177        [(); N - 1]:;
178}
179
180impl<T, const N: usize> const Deconstruct<T, N> for [T; N]
181where
182    T: [const] Destruct,
183{
184    #[doc(alias = "pop_front")]
185    fn uncons(self) -> (T, [T; N - 1]) {
186        // SAFETY: the layout is alright.
187        unsafe { Pair::splat(self) }
188    }
189
190    #[doc(alias = "pop")]
191    fn unsnoc(self) -> ([T; N - 1], T) {
192        // SAFETY: the layout is still alright.
193        unsafe { Pair::splat(self) }
194    }
195    fn tail(self) -> [T; N - 1]
196    where
197        T: [const] Destruct,
198    {
199        self.uncons().1
200    }
201    #[doc(alias = "trunc")]
202    fn init(self) -> [T; N - 1] {
203        self.unsnoc().0
204    }
205    fn last(self) -> T
206    where
207        [(); N - 1]:,
208    {
209        self.unsnoc().1
210    }
211
212    fn head(self) -> T
213    where
214        [(); N - 1]:,
215    {
216        self.uncons().0
217    }
218}
219
220/// Join scalars together.
221#[const_trait]
222pub trait Join<T, const N: usize, const O: usize, U> {
223    /// Join a array and an scalar together. For joining two arrays together, see [`Couple`].
224    /// ```
225    /// # #![feature(generic_const_exprs)]
226    /// # use atools::prelude::*;
227    /// let a = [1, 2].join(3);
228    /// let b = 1.join([2, 3]);
229    /// let c = 1.join(2).join(3);
230    /// ```
231    fn join(self, with: U) -> [T; N + O];
232}
233
234/// Couple two arrays together.
235#[const_trait]
236pub trait Couple<T, const N: usize, const O: usize> {
237    /// Couple two arrays together. This could have been [`Join`], but the methods would require disambiguation.
238    /// ```
239    /// # #![feature(generic_const_exprs)]
240    /// # use atools::prelude::*;
241    /// let a = 1.join(2).couple([3, 4]);
242    /// ```
243    fn couple(self, with: [T; O]) -> [T; N + O];
244}
245
246impl<T, const N: usize, const O: usize> const Couple<T, N, O> for [T; N] {
247    fn couple(self, with: [T; O]) -> [T; N + O] {
248        // SAFETY: adjacent
249        unsafe { transmute_unchecked(Pair(self, with)) }
250    }
251}
252
253impl<T, const N: usize> const Join<T, N, 1, T> for [T; N] {
254    fn join(self, with: T) -> [T; N + 1] {
255        self.couple([with])
256    }
257}
258
259impl<T> const Join<T, 1, 1, T> for T {
260    fn join(self, with: T) -> [T; 2] {
261        [self, with]
262    }
263}
264
265impl<T, const O: usize> const Join<T, 1, O, [T; O]> for T {
266    fn join(self, with: [T; O]) -> [T; 1 + O] {
267        [self].couple(with)
268    }
269}
270
271pub(crate) const fn assert_zero(x: usize) -> usize {
272    if x != 0 {
273        panic!("expected zero");
274    } else {
275        0
276    }
277}
278
279/// 🍪
280#[allow(private_bounds)]
281#[const_trait]
282pub trait Chunked<T, const N: usize> {
283    /// Chunks.
284    /// This will compile fail if `N ∤ (does not divide) C`
285    /// ```
286    /// # #![feature(generic_const_exprs)]
287    /// # use atools::prelude::*;
288    /// assert_eq!(range::<6>().chunked::<3>(), [[0, 1, 2], [3, 4, 5]]);
289    /// ```
290    #[allow(private_bounds)]
291    fn chunked<const C: usize>(self) -> [[T; C]; N / C]
292    where
293        // N % C == 0
294        [(); assert_zero(N % C)]:;
295}
296
297impl<const N: usize, T> const Chunked<T, N> for [T; N] {
298    #[allow(private_bounds)]
299    fn chunked<const C: usize>(self) -> [[T; C]; N / C]
300    where
301        [(); assert_zero(N % C)]:,
302    {
303        // SAFETY: N != 0 && wont leak as N % C == 0.
304        let retval = unsafe { self.as_ptr().cast::<[[T; C]; N / C]>().read() };
305        core::mem::forget(self);
306        retval
307    }
308}
309
310/// Flatten arrays.
311#[const_trait]
312pub trait Flatten<T, const N: usize, const N2: usize> {
313    /// Takes a `[[T; N]; N2]`, and flattens it to a `[T; N * N2]`.
314    ///
315    /// # Examples
316    ///
317    /// ```
318    /// # #![feature(generic_const_exprs)]
319    /// # use atools::prelude::*;
320    /// assert_eq!([[1, 2, 3], [4, 5, 6]].flatten(), [1, 2, 3, 4, 5, 6]);
321    ///
322    /// assert_eq!(
323    ///     [[1, 2, 3], [4, 5, 6]].flatten(),
324    ///     [[1, 2], [3, 4], [5, 6]].flatten(),
325    /// );
326    ///
327    /// let array_of_empty_arrays: [[i32; 0]; 5] = [[], [], [], [], []];
328    /// assert!(array_of_empty_arrays.flatten().is_empty());
329    ///
330    /// let empty_array_of_arrays: [[u32; 10]; 0] = [];
331    /// assert!(empty_array_of_arrays.flatten().is_empty());
332    /// ```
333    fn flatten(self) -> [T; N * N2];
334}
335
336impl<T, const N: usize, const N2: usize> const Flatten<T, N, N2> for [[T; N]; N2] {
337    fn flatten(self) -> [T; N * N2] {
338        // SAFETY: layout is the same.
339        unsafe { core::intrinsics::transmute_unchecked(self) }
340    }
341}
342
343#[const_trait]
344/// Splitting arrays up.
345pub trait Split<T, const N: usize> {
346    /// Splits the array into twain.
347    /// ```
348    /// # #![feature(generic_const_exprs)]
349    /// # use atools::prelude::*;
350    /// let x = [1u8, 2, 3];
351    /// let ([x], [y, z]) = x.split::<1>();
352    /// ```
353    fn split<const AT: usize>(self) -> ([T; AT], [T; N - AT]);
354    /// Take `AT` elements, discarding the rest.
355    /// ```
356    /// # #![feature(generic_const_exprs)]
357    /// # use atools::prelude::*;
358    /// assert_eq!(range::<50>().take::<5>(), range::<5>());
359    /// ```
360    fn take<const AT: usize>(self) -> [T; AT]
361    where
362        [(); N - AT]:,
363        T: [const] Destruct;
364    /// Discard `AT` elements, returning the rest.
365    fn drop<const AT: usize>(self) -> [T; N - AT]
366    where
367        T: [const] Destruct;
368}
369
370impl<T, const N: usize> const Split<T, N> for [T; N] {
371    fn split<const AT: usize>(self) -> ([T; AT], [T; N - AT]) {
372        // SAFETY: N - AT overflows when AT > N so the size of the returned "array" is the same.
373        unsafe { Pair::splat(self) }
374    }
375    fn take<const M: usize>(self) -> [T; M]
376    where
377        // M <= N
378        [(); N - M]:,
379        T: [const] Destruct,
380    {
381        self.split::<M>().0
382    }
383    fn drop<const M: usize>(self) -> [T; N - M]
384    where
385        T: [const] Destruct,
386    {
387        self.split::<M>().1
388    }
389}
390#[const_trait]
391/// Zip arrays together.
392pub trait Zip<T, const N: usize> {
393    /// Zip arrays together.
394    fn zip<U>(self, with: [U; N]) -> [(T, U); N];
395}
396
397impl<T, const N: usize> const Zip<T, N> for [T; N] {
398    fn zip<U>(self, with: [U; N]) -> [(T, U); N] {
399        let mut out = unsafe { MU::<[MU<_>; N]>::uninit().assume_init() };
400        let mut i = 0usize;
401        while i < out.len() {
402            out[i] = MU::new(unsafe { (self.as_ptr().add(i).read(), with.as_ptr().add(i).read()) });
403            i += 1;
404        }
405        core::mem::forget((self, with));
406        unsafe { transmute_unchecked(out) }
407    }
408}
409
410/// Array tools.
411pub trait ArrayTools<T, const N: usize> {
412    /// Skip `BY` elements.
413    fn skip<const BY: usize>(self) -> [T; N - BY];
414    /// Skip every `BY` elements.
415    ///
416    /// ```
417    /// # #![feature(generic_const_exprs)]
418    /// # use atools::prelude::*;
419    /// let x = range::<5>().step::<2>();
420    /// assert_eq!(x, [0, 2, 4]);
421    /// let x = range::<20>().step::<5>();
422    /// assert_eq!(x, [0, 5, 10, 15]);
423    /// assert_eq!(range::<50>().step::<3>(), [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48]);
424    /// ```
425    fn step<const STEP: usize>(self) -> [T; 1 + (N - 1) / (STEP)];
426    /// Intersperse a element in between items.
427    /// ```
428    /// # #![feature(generic_const_exprs)]
429    /// # use atools::prelude::*;
430    /// let x = range::<3>().intersperse(5);
431    /// assert_eq!(x, [0, 5, 1, 5, 2]);
432    /// ```
433    fn intersperse(self, with: T) -> [T; (N * 2) - 1]
434    where
435        T: Clone;
436    /// Run a function on every element.
437    fn each(self, apply: impl FnMut(T));
438    /// Embed the index.
439    fn enumerate(self) -> [(T, usize); N];
440    /// Get the sliding windows of this array.
441    /// ```
442    /// # #![feature(generic_const_exprs)]
443    /// # use atools::prelude::*;
444    /// assert_eq!(range::<5>().windowed::<2>(), [&[0, 1], &[1, 2], &[2, 3], &[3, 4]]);
445    /// ```
446    fn windowed<const W: usize>(&self) -> [&[T; W]; N - W + 1];
447    /// Inspect every element of this array.
448    fn inspect(self, f: impl FnMut(&T)) -> Self;
449    /// Reverse this array.
450    fn rev(self) -> Self;
451    /// Interleave items from two arrays.
452    /// ```
453    /// # #![feature(generic_const_exprs)]
454    /// # use atools::prelude::*;
455    /// assert_eq!([0u8, 2, 4].interleave([1, 3, 5]), [0, 1, 2, 3, 4, 5]);
456    /// ```
457    fn interleave(self, with: [T; N]) -> [T; N * 2];
458    /// [Cartesian product](https://en.wikipedia.org/wiki/Cartesian_product) (`A  ×  B`) of two arrays.
459    /// ```
460    /// # #![feature(generic_const_exprs)]
461    /// # use atools::prelude::*;
462    /// assert_eq!([1u64, 2].cartesian_product(&["Π", "Σ"]), [(1, "Π"), (1, "Σ"), (2, "Π"), (2, "Σ")]);
463    /// ```
464    fn cartesian_product<U: Clone, const M: usize>(&self, with: &[U; M]) -> [(T, U); N + M]
465    where
466        T: Clone;
467    /// Sorts it. This uses <code>[[T](slice)]::[sort_unstable](slice::sort_unstable)</code>.
468    fn sort(self) -> Self
469    where
470        T: Ord;
471    /// Sum of the array.
472    fn sum(self) -> T
473    where
474        T: core::iter::Sum<T>;
475    /// Product of the array.
476    fn product(self) -> T
477    where
478        T: core::iter::Product<T>;
479}
480
481impl<T, const N: usize> ArrayTools<T, N> for [T; N] {
482    fn skip<const BY: usize>(self) -> [T; N - BY] {
483        self.into_iter().skip(BY).carr()
484    }
485    fn step<const STEP: usize>(self) -> [T; 1 + (N - 1) / (STEP)] {
486        self.into_iter().step_by(STEP).carr()
487    }
488    fn intersperse(self, with: T) -> [T; (N * 2) - 1]
489    where
490        T: Clone,
491    {
492        self.into_iter().intersperse(with).carr()
493    }
494
495    fn each(self, apply: impl FnMut(T)) {
496        self.into_iter().for_each(apply);
497    }
498
499    fn enumerate(self) -> [(T, usize); N] {
500        let mut n = 0;
501        self.map(|x| {
502            let o = n;
503            n += 1;
504            (x, o)
505        })
506    }
507
508    fn windowed<const W: usize>(&self) -> [&[T; W]; N - W + 1] {
509        self.array_windows().carr()
510    }
511
512    fn inspect(self, f: impl FnMut(&T)) -> Self {
513        self.iter().for_each(f);
514        self
515    }
516
517    fn rev(mut self) -> Self {
518        self.reverse();
519        self
520    }
521
522    fn interleave(self, with: [T; N]) -> [T; N * 2] {
523        let mut which = true;
524        let mut a = self.into_iter();
525        let mut b = with.into_iter();
526        from_fn(|_| {
527            which = !which;
528            match which {
529                false => a.next().unwrap(),
530                true => b.next().unwrap(),
531            }
532        })
533    }
534
535    fn cartesian_product<U: Clone, const M: usize>(&self, with: &[U; M]) -> [(T, U); N + M]
536    where
537        T: Clone,
538    {
539        self.iter()
540            .flat_map(|a| with.iter().map(move |b| (a.clone(), b.clone())))
541            .carr()
542    }
543
544    fn sort(mut self) -> Self
545    where
546        T: Ord,
547    {
548        <[T]>::sort_unstable(&mut self);
549        self
550    }
551
552    fn sum(self) -> T
553    where
554        T: core::iter::Sum<T>,
555    {
556        self.into_iter().sum()
557    }
558
559    fn product(self) -> T
560    where
561        T: core::iter::Product<T>,
562    {
563        self.into_iter().product()
564    }
565}