kir/
flatindex.rs

1pub trait FlatIndex<T> {
2  fn flat_len(&self) -> usize;
3  fn flat_index(&self, i: usize) -> &T;
4  fn flat_index_mut(&mut self, i: usize) -> &mut T;
5}
6impl<T> FlatIndex<T> for T {
7  fn flat_len(&self) -> usize {
8    1
9  }
10  fn flat_index(&self, i: usize) -> &T {
11    assert!(i == 0);
12    self
13  }
14  fn flat_index_mut(&mut self, i: usize) -> &mut T {
15    assert!(i == 0);
16    self
17  }
18}
19impl<T> FlatIndex<T> for [T] {
20  fn flat_len(&self) -> usize {
21    self.len()
22  }
23  fn flat_index(&self, i: usize) -> &T {
24    &self[i]
25  }
26  fn flat_index_mut(&mut self, i: usize) -> &mut T {
27    &mut self[i]
28  }
29}
30impl<T> FlatIndex<T> for Option<T> {
31  fn flat_len(&self) -> usize {
32    self.is_some() as usize
33  }
34  fn flat_index(&self, i: usize) -> &T {
35    assert!(i == 0);
36    self.as_ref().unwrap()
37  }
38  fn flat_index_mut(&mut self, i: usize) -> &mut T {
39    assert!(i == 0);
40    self.as_mut().unwrap()
41  }
42}
43impl<T> FlatIndex<T> for Vec<T> {
44  fn flat_len(&self) -> usize {
45    self.len()
46  }
47  fn flat_index(&self, i: usize) -> &T {
48    &self[i]
49  }
50  fn flat_index_mut(&mut self, i: usize) -> &mut T {
51    &mut self[i]
52  }
53}
54impl<T, const N: usize> FlatIndex<T> for [T; N] {
55  fn flat_len(&self) -> usize {
56    N
57  }
58  fn flat_index(&self, i: usize) -> &T {
59    &self[i]
60  }
61  fn flat_index_mut(&mut self, i: usize) -> &mut T {
62    &mut self[i]
63  }
64}
65
66#[macro_export]
67macro_rules! flat_combine_len {
68    (<$t:ty>) => { 0 };
69    (<$t:ty> <$e:expr, $le:ident, $re:ident> $(, $($tt:tt)*)?) => {
70        $e.$le() + $crate::flat_combine_len!(<$t> $($($tt)*)?)
71    };
72    (<$t:ty> $e:expr $(, $($tt:tt)*)?) => {
73        $crate::flatindex::FlatIndex::<$t>::flat_len($e) + $crate::flat_combine_len!(<$t> $($($tt)*)?)
74    };
75    // (<$t:ty> $e:expr $(, $rest_e:expr)* $(,)?) => {
76    //     $crate::flatindex::FlatIndex::<$t>::flat_len($e) + $crate::flat_combine_len!(<$t> $($rest_e),*)
77    // };
78}
79
80#[macro_export]
81macro_rules! flat_combine {
82    (<$t:ty> $idx:expr $(,)?) => { panic!("Index out of bounds") };
83    (<$t:ty> $idx:expr, <$e:expr, $le:ident, $re:ident> $(, $($tt:tt)*)?) => {{
84        let __idx: usize = $idx;
85        let __len = $e.$le();
86        if __idx < __len {
87            $e.$re(__idx)
88        } else {
89            $crate::flat_combine!(<$t> __idx - __len $(, $($tt)*)?)
90        }
91    }};
92    (<$t:ty> $idx:expr, $e:expr $(, $($tt:tt)*)?) => {{
93        let __idx: usize = $idx;
94        let __len = $crate::flatindex::FlatIndex::<$t>::flat_len($e);
95        if __idx < __len {
96            $crate::flatindex::FlatIndex::<$t>::flat_index($e, __idx)
97        } else {
98            $crate::flat_combine!(<$t> __idx - __len $(, $($tt)*)?)
99        }
100    }};
101    // (<$t:ty> $idx:expr, $e:expr $(, $rest_e:expr)* $(,)?) => {{
102    //     let __idx: usize = $idx;
103    //     let __len = $crate::flatindex::FlatIndex::<$t>::flat_len($e);
104    //     if __idx < __len {
105    //         $crate::flatindex::FlatIndex::<$t>::flat_index($e, __idx)
106    //     } else {
107    //         $crate::flat_combine!(<$t> __idx - __len $(, $rest_e)*)
108    //     }
109    // }};
110}
111
112#[macro_export]
113macro_rules! flat_combine_mut {
114    (<$t:ty> $idx:expr $(,)?) => { panic!("Index out of bounds") };
115    (<$t:ty> $idx:expr, <$e:expr, $le:ident, $re:ident> $(, $($tt:tt)*)?) => {{
116        let __idx: usize = $idx;
117        let __len = $e.$le();
118        if __idx < __len {
119            $e.$re(__idx)
120        } else {
121            $crate::flat_combine_mut!(<$t> __idx - __len $(, $($tt)*)?)
122        }
123    }};
124    (<$t:ty> $idx:expr, $e:expr $(, $($tt:tt)*)?) => {{
125        let __idx: usize = $idx;
126        let __len = $crate::flatindex::FlatIndex::<$t>::flat_len($e);
127        if __idx < __len {
128            $crate::flatindex::FlatIndex::<$t>::flat_index_mut($e, __idx)
129        } else {
130            $crate::flat_combine_mut!(<$t> __idx - __len $(, $($tt)*)?)
131        }
132    }};
133    // (<$t:ty> $idx:expr, $e:expr $(, $rest_e:expr)* $(,)?) => {{
134    //     let __idx: usize = $idx;
135    //     let __len = $crate::flatindex::FlatIndex::<$t>::flat_len($e);
136    //     if __idx < __len {
137    //         $crate::flatindex::FlatIndex::<$t>::flat_index_mut($e, __idx)
138    //     } else {
139    //         $crate::flat_combine_mut!(<$t> __idx - __len $(, $rest_e)*)
140    //     }
141    // }};
142}
143
144#[macro_export]
145macro_rules! flat_combine_deref {
146    (<$t:ty> $idx:expr $(,)?) => { panic!("Index out of bounds") };
147    (<$t:ty> $idx:expr, <$e:expr, $le:ident, $re:ident> $(, $($tt:tt)*)?) => {{
148        let __idx: usize = $idx;
149        let __len = $e.$le();
150        if __idx < __len {
151            $e.$re(__idx)
152        } else {
153            $crate::flat_combine_deref!(<$t> __idx - __len $(, $($tt)*)?)
154        }
155    }};
156    (<$t:ty> $idx:expr, $e:expr $(, $($tt:tt)*)?) => {{
157        let __idx: usize = $idx;
158        let __len = $crate::flatindex::FlatIndex::<$t>::flat_len($e);
159        if __idx < __len {
160            *$crate::flatindex::FlatIndex::<$t>::flat_index($e, __idx)
161        } else {
162            $crate::flat_combine_deref!(<$t> __idx - __len $(, $($tt)*)?)
163        }
164    }};
165}
166
167#[cfg(test)]
168mod tests {
169  use std::ops::Index;
170  #[test]
171  fn combine() {
172    let mut a: Vec<usize> = vec![0, 1, 2, 3, 4, 5];
173    let mut b: Vec<usize> = vec![6, 7, 8, 9, 10, 11];
174    let mut c: Option<usize> = Some(12);
175    let mut d: Option<usize> = None;
176    let mut e: [usize; 3] = [13, 14, 15];
177    assert_eq!(flat_combine_len!(<usize> &a, &b), 12);
178    assert_eq!(flat_combine_len!(<usize> &a, &b, &c, &d, &e), 16);
179    for i in 0..16 {
180      let combined = flat_combine!(<usize> i, &a, <&b, len, index>, &c, &d, &e);
181      assert_eq!(combined, &i);
182      assert_eq!(
183        *flat_combine_mut!(<usize> i, &mut a, <&mut b, len, index>, &mut c, &mut d, &mut e),
184        i
185      );
186      assert_eq!(flat_combine_deref!(<usize> i, &a, &b, &c, &d, &e), i);
187    }
188  }
189
190  #[test]
191  #[should_panic]
192  fn out_of_range() {
193    let a: Vec<usize> = vec![0, 1, 2, 3, 4, 5];
194    let b: Vec<usize> = vec![6, 7, 8, 9, 10, 11];
195    let c: Option<usize> = Some(12);
196    let d: Option<usize> = None;
197    let e: [usize; 3] = [13, 14, 15];
198    flat_combine!(<usize> 16, &a, <&b, len, index>, &c, &d, &e);
199  }
200}
kir/flatindex.rs

kir/
flatindex.rs
