acme_tensor/actions/index/
slice.rs

1/*
2    Appellation: slice <mod>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5//! # Slice
6//!
7//!
8use super::Ixs;
9
10#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
11#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
12pub struct Slice {
13    pub start: Ixs,
14    pub end: Option<Ixs>,
15    pub step: Ixs,
16}
17
18impl Slice {
19    pub fn new(start: Ixs, end: Option<Ixs>, step: Ixs) -> Self {
20        debug_assert_ne!(step, 0, "step must be non-zero");
21        Self { start, end, step }
22    }
23
24    pub fn start(&self) -> Ixs {
25        self.start
26    }
27
28    pub fn end(&self) -> Option<Ixs> {
29        self.end
30    }
31
32    pub fn step(&self) -> Ixs {
33        self.step
34    }
35}
36
37#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
38#[cfg_attr(
39    feature = "serde",
40    derive(serde::Deserialize, serde::Serialize),
41    serde(rename_all = "snake_case")
42)]
43pub enum SliceType {
44    Index(Ixs),
45    Slice(Slice),
46    NewAxis,
47}
48
49impl SliceType {
50    /// Create a new axis
51    pub fn new_axis() -> Self {
52        Self::NewAxis
53    }
54    /// Create a new index
55    pub fn index(index: Ixs) -> Self {
56        Self::Index(index)
57    }
58    /// Create a new slice
59    pub fn slice(start: Ixs, end: Option<Ixs>, step: Ixs) -> Self {
60        Self::Slice(Slice::new(start, end, step))
61    }
62}
63
64pub struct Slices {
65    slices: Vec<SliceType>,
66}
67
68impl Slices {
69    pub fn new(slices: impl IntoIterator<Item = SliceType>) -> Self {
70        Self {
71            slices: Vec::from_iter(slices),
72        }
73    }
74
75    pub fn iter(&self) -> core::slice::Iter<'_, SliceType> {
76        self.slices.iter()
77    }
78}
79
80macro_rules! impl_from_range {
81    ($self:ty, $constructor:expr, $idx:ty, [$($index:ty),*]) => {
82        $(
83            impl_from_range!($self, $constructor, $idx, $index);
84        )*
85    };
86    ($self:ty, $constructor:expr, $idx:ty, $index:ty) => {
87        impl From<core::ops::Range<$index>> for $self {
88            #[inline]
89            fn from(r: core::ops::Range<$index>) -> $self {
90                $constructor(r.start as $idx, Some(r.end as $idx), 1)
91            }
92        }
93
94        impl From<core::ops::RangeFrom<$index>> for $self {
95            #[inline]
96            fn from(r: core::ops::RangeFrom<$index>) -> $self {
97                $constructor(r.start as $idx, None, 1)
98            }
99        }
100
101        impl From<core::ops::RangeInclusive<$index>> for $self {
102            #[inline]
103            fn from(r: core::ops::RangeInclusive<$index>) -> $self {
104                let end = *r.end() as $idx;
105                $constructor(*r.start() as $idx, if end == -1 { None } else { Some(end + 1) }, 1)
106            }
107        }
108
109        impl From<core::ops::RangeTo<$index>> for $self {
110            #[inline]
111            fn from(r: core::ops::RangeTo<$index>) -> $self {
112                $constructor(0, Some(r.end as $idx), 1)
113            }
114        }
115
116        impl From<core::ops::RangeToInclusive<$index>> for $self {
117            #[inline]
118            fn from(r: core::ops::RangeToInclusive<$index>) -> $self {
119                let end = r.end as $idx;
120                $constructor(0, if end == -1 { None } else { Some(end + 1) }, 1)
121            }
122        }
123    };
124}
125
126impl_from_range!(Slice, Slice::new, Ixs, [i32, isize, usize]);
127impl_from_range!(SliceType, SliceType::slice, Ixs, [i32, isize, usize]);