atools - const generic expr based fixed length array manipulation crate


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use core::ops::{RangeFrom, RangeInclusive, RangeTo, RangeToInclusive};
#[derive(Eq, PartialEq, Copy, Clone, core::marker::ConstParamTy)]
pub enum Range {
    Range(usize, usize),
    RangeFrom(usize),
    RangeTo(usize),
}

impl Range {
    pub const fn range<const N: usize>(self) -> (usize, usize) {
        match self {
            Range::Range(begin, end) => (begin, end),
            Range::RangeFrom(begin) => (begin, N),
            Range::RangeTo(end) => (0, end),
        }
    }
    pub const fn length<const N: usize>(self) -> usize {
        let (begin, end) = self.range::<N>();
        end - begin
    }
    pub const fn valid<const N: usize>(self) -> usize {
        (self.range::<N>().1 <= N) as usize
    }
}

#[diagnostic::on_unimplemented(
    message = "{Self} is not a valid range type",
    label = "use a correct range type, such as {{Range(x..y), RangeInclusive(x..=y), RangeTo(..x), RangeToInclusive(..=x)}}"
)]
pub const trait Ranged {
    fn range(self) -> Range;
}
impl const Ranged for core::ops::Range<usize> {
    fn range(self) -> Range {
        Range::Range(self.start, self.end)
    }
}
impl const Ranged for RangeInclusive<usize> {
    fn range(self) -> Range {
        Range::Range(*self.start(), *self.end() + 1)
    }
}

impl const Ranged for RangeFrom<usize> {
    fn range(self) -> Range {
        Range::RangeFrom(self.start)
    }
}
impl const Ranged for RangeTo<usize> {
    fn range(self) -> Range {
        Range::RangeTo(self.end)
    }
}
impl const Ranged for RangeToInclusive<usize> {
    fn range(self) -> Range {
        Range::RangeTo(self.end + 1)
    }
}
/// Constifies a range. For use with [`slice`](Slice::slice).
///
/// Takes a type in the form {[`Range`], [`RangeInclusive`], [`RangeTo`], [`RangeToInclusive`]}.
#[allow(private_bounds)]
pub const fn r<T: [const] Ranged>(x: T) -> Range {
    Ranged::range(x)
}

/// Slicing arrays up.
pub const trait Slice<T, const N: usize> {
    /// Slices the array.
    /// Compile time checked.
    /// ```
    /// # #![feature(generic_const_exprs, const_trait_impl)]
    /// # use atools::prelude::*;
    /// let x = atools::range::<5>();
    /// assert_eq!(*x.slice::<{ r(2..=4) }>(), [2, 3, 4]);
    /// // x.slice::<{ r(..10) }>(); // ERROR
    /// ```
    fn slice<const RANGE: Range>(&self) -> &[T; RANGE.length::<N>()]
    where
        // comptime length check
        [(); RANGE.valid::<N>() - 1]:;
}

const unsafe fn slice<T, const N: usize, const TO: usize>(x: &[T; N], begin: usize) -> &[T; TO] {
    // SAFETY: up to caller
    unsafe { &*x.as_ptr().add(begin).cast::<[T; TO]>() }
}

impl<T, const N: usize> const Slice<T, N> for [T; N] {
    fn slice<const RANGE: Range>(&self) -> &[T; RANGE.length::<N>()]
    where
        [(); RANGE.valid::<N>() - 1]:,
    {
        // SAFETY: the validity check ensures that the array will be in bounds.
        unsafe { slice::<T, N, { RANGE.length::<N>() }>(self, RANGE.range::<N>().0) }
    }
}

#[test]
fn slicing() {
    let x = [1, 2, 3];
    let &[y] = x.slice::<{ r(2..) }>();
    assert_eq!(y, 3);
    let &[y, z] = x.slice::<{ r(1..=2) }>();
    assert_eq!([y, z], [2, 3]);
    let &y = x.slice::<{ r(..=2) }>();
    assert_eq!(x, y);
    let &y = x.slice::<{ r(..2) }>();
    assert_eq!(y, [1, 2]);
}