redstone_ml/util/
index.rs

1use std::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
2
3#[macro_export]
4macro_rules! s {
5    ($($x:expr),*) => {
6        [$($crate::util::index::Index::from($x)),*]
7    };
8}
9
10#[derive(Clone)]
11pub enum Index {
12    Usize(usize),
13    Range(Range<usize>),
14    RangeFrom(RangeFrom<usize>),
15    RangeFull(RangeFull),
16    RangeInclusive(RangeInclusive<usize>),
17    RangeTo(RangeTo<usize>),
18    RangeToInclusive(RangeToInclusive<usize>),
19}
20
21pub(crate) trait Indexer: Clone {
22    /// The resulting dimension of the axis indexed by this indexer.
23    fn indexed_length(&self, axis_length: usize) -> usize;
24    
25    /// The first element along the dimension indexed by this kind of indexer
26    /// For example, 0 for `ndarray [..]` or `ndarray [..2]` but 5 for `ndarray [5..]` or `ndarray [5]`
27    fn index_of_first_element(&self) -> usize;
28
29    /// When indexed with this kind of object, does the dimension of the ndarray collapse?
30    /// Only true for usize since all range-based indexers retain the dimension.
31    fn collapse_dimension(&self) -> bool {
32        false
33    }
34}
35
36impl Indexer for Index {
37    fn indexed_length(&self, len: usize) -> usize {
38        match self {
39            Index::Usize(index) => Indexer::indexed_length(index, len),
40            Index::Range(index) => Indexer::indexed_length(index, len),
41            Index::RangeFrom(index) => Indexer::indexed_length(index, len),
42            Index::RangeFull(index) => Indexer::indexed_length(index, len),
43            Index::RangeInclusive(index) => Indexer::indexed_length(index, len),
44            Index::RangeTo(index) => Indexer::indexed_length(index, len),
45            Index::RangeToInclusive(index) => Indexer::indexed_length(index, len),
46        }
47    }
48
49    fn index_of_first_element(&self) -> usize {
50        match self {
51            Index::Usize(index) => index.index_of_first_element(),
52            Index::Range(index) => index.index_of_first_element(),
53            Index::RangeFrom(index) => index.index_of_first_element(),
54            Index::RangeFull(index) => index.index_of_first_element(),
55            Index::RangeInclusive(index) => index.index_of_first_element(),
56            Index::RangeTo(index) => index.index_of_first_element(),
57            Index::RangeToInclusive(index) => index.index_of_first_element(),
58        }
59    }
60    
61    fn collapse_dimension(&self) -> bool {
62        matches!(self, Index::Usize(_))
63    }
64}
65
66impl Indexer for usize {
67    fn indexed_length(&self, _axis_length: usize) -> usize {
68        1
69    }
70
71    fn index_of_first_element(&self) -> usize {
72        *self
73    }
74
75    fn collapse_dimension(&self) -> bool {
76        true
77    }
78}
79impl Indexer for Range<usize> {
80    fn indexed_length(&self, _axis_length: usize) -> usize {
81        self.end - self.start
82    }
83
84    fn index_of_first_element(&self) -> usize {
85        self.start
86    }
87}
88impl Indexer for RangeFull {
89    fn indexed_length(&self, axis_length: usize) -> usize {
90        axis_length
91    }
92    fn index_of_first_element(&self) -> usize {
93        0
94    }
95}
96impl Indexer for RangeFrom<usize> {
97    fn indexed_length(&self, axis_length: usize) -> usize {
98        axis_length - self.start
99    }
100    fn index_of_first_element(&self) -> usize {
101        self.start
102    }
103}
104impl Indexer for RangeTo<usize> {
105    fn indexed_length(&self, _axis_length: usize) -> usize {
106        self.end
107    }
108    fn index_of_first_element(&self) -> usize {
109        0
110    }
111}
112impl Indexer for RangeToInclusive<usize> {
113    fn indexed_length(&self, _axis_length: usize) -> usize {
114        self.end + 1
115    }
116    fn index_of_first_element(&self) -> usize {
117        0
118    }
119}
120impl Indexer for RangeInclusive<usize> {
121    fn indexed_length(&self, _axis_length: usize) -> usize {
122        self.end() - self.start() + 1
123    }
124    fn index_of_first_element(&self) -> usize {
125        *self.start()
126    }
127}
128
129impl From<usize> for Index {
130    fn from(val: usize) -> Self {
131        Index::Usize(val)
132    }
133}
134
135impl From<Range<usize>> for Index {
136    fn from(val: Range<usize>) -> Self {
137        Index::Range(val)
138    }
139}
140
141impl From<RangeFrom<usize>> for Index {
142    fn from(val: RangeFrom<usize>) -> Self {
143        Index::RangeFrom(val)
144    }
145}
146
147impl From<RangeFull> for Index {
148    fn from(val: RangeFull) -> Self {
149        Index::RangeFull(val)
150    }
151}
152
153impl From<RangeInclusive<usize>> for Index {
154    fn from(val: RangeInclusive<usize>) -> Self {
155        Index::RangeInclusive(val)
156    }
157}
158
159impl From<RangeTo<usize>> for Index {
160    fn from(val: RangeTo<usize>) -> Self {
161        Index::RangeTo(val)
162    }
163}
164
165impl From<RangeToInclusive<usize>> for Index {
166    fn from(val: RangeToInclusive<usize>) -> Self {
167        Index::RangeToInclusive(val)
168    }
169}
redstone_ml/util/index.rs

redstone_ml/util/
index.rs
