1#[macro_export]
2macro_rules! make_index_fn {
3 ($name:ident; $dim:literal) => {
4 impl<T: tensor_macros::traits::TensorTrait> std::ops::Index<usize> for $name<T> {
5 type Output = T;
6
7 fn index(&self, i: usize) -> &Self::Output {
8 &self.0[i]
9 }
10 }
11
12 impl<T: tensor_macros::traits::TensorTrait> std::ops::IndexMut<usize> for $name<T> {
13 fn index_mut(&mut self, i: usize) -> &mut T {
14 &mut self.0[i]
15 }
16 }
17 };
18
19 ($name:ident; $($dims:literal),+) => {
20 impl<T: tensor_macros::traits::TensorTrait> std::ops::Index<usize> for $name<T> {
21 type Output = T;
22
23 fn index(&self, i: usize) -> &Self::Output {
24 &self.0[i]
25 }
26 }
27
28 impl<T: tensor_macros::traits::TensorTrait> std::ops::IndexMut<usize> for $name<T> {
29 fn index_mut(&mut self, i: usize) -> &mut T {
30 &mut self.0[i]
31 }
32 }
33
34 make_index_fn!($name; $($dims),*;;;);
35 };
36
37 ($name:ident; $dim:literal $(,$dims:literal)*; $($i:ident),*; $($t:ty),*; $($dims_bk:literal),*) => {
38 make_index_fn!($name; $($dims),*; $($i,)* i; $($t,)* usize; $($dims_bk,)* $dim);
39 };
40 ($name:ident; ; $($i:ident),*; $($t:ty),*; $($dims:literal),*) => {
41 impl<T: tensor_macros::traits::TensorTrait> std::ops::Index<( $($t),* )> for $name<T> {
42 type Output = T;
43
44 fn index(&self, ( $($i),* ): ( $($t),* )) -> &Self::Output {
45 &self.0[
46 make_index_val!($($dims),*; $($i),*)
47 ]
48 }
49 }
50
51 impl<T: tensor_macros::traits::TensorTrait> std::ops::IndexMut<( $($t),* )> for $name<T> {
52 fn index_mut(&mut self, ( $($i),* ): ( $($t),* )) -> &mut T {
53 &mut self.0[
54 make_index_val!($($dims),*; $($i),*)
55 ]
56 }
57 }
58 };
59}
60
61#[macro_export]
62macro_rules! make_index_val {
63 ($dim:literal $(,$dims:literal)*; $i:expr $(,$is:expr)* ) => (
64 $i * mul!($($dims),*) + make_index_val!($($dims),*; $($is),*)
65 );
66 (;) => (0)
67}