aeon_tensor/
compound.rs

1use crate::indices::{Gen, IndexIterator, Sym, TensorIndex};
2use std::marker::PhantomData;
3
4/// Singleton for implementing `CompondSum`
5pub struct CSum;
6
7/// Seal CompoundSum to CSum.
8mod private {
9    pub trait Sealed {}
10}
11
12impl private::Sealed for CSum {}
13
14pub trait CompoundSum<const T: usize, const L: usize, const R: usize>: private::Sealed {
15    fn join(idx: ([usize; L], [usize; R])) -> [usize; T];
16    fn split(idx: [usize; T]) -> ([usize; L], [usize; R]);
17}
18
19// DRY macro for Implementing Compound sum
20macro_rules! impl_compound_sum {
21    ($T:literal, $L:literal, $R:literal, [$($left:ident),+] [$($right:ident),+]) => {
22        impl CompoundSum<$T, $L, $R> for CSum {
23            fn join(([$($left),*], [$($right),*]): ([usize; $L], [usize; $R])) -> [usize; $T] {
24                [$($left),* , $($right),*]
25            }
26
27            fn split([$($left),* , $($right),*]: [usize; $T]) -> ([usize; $L], [usize; $R]) {
28                ([$($left),*], [$($right),*])
29            }
30        }
31    };
32}
33
34impl_compound_sum!(3, 1, 2, [a][b, c]);
35impl_compound_sum!(3, 2, 1, [a, b][c]);
36impl_compound_sum!(4, 2, 2, [a, b][c, d]);
37impl_compound_sum!(4, 1, 3, [a][b, c, d]);
38impl_compound_sum!(4, 3, 1, [a, b, c][d]);
39
40/// Represents compound combinations of fundemental indices.
41pub struct Compound<const T: usize, const L: usize, const R: usize, Left, Right>
42where
43    CSum: CompoundSum<T, L, R>,
44{
45    _marker: PhantomData<(Left, Right)>,
46}
47
48impl<const N: usize, const T: usize, const L: usize, const R: usize, Left, Right> TensorIndex<N, T>
49    for Compound<T, L, R, Left, Right>
50where
51    Left: TensorIndex<N, L>,
52    Right: TensorIndex<N, R>,
53    CSum: CompoundSum<T, L, R>,
54{
55    type Indices = CompoundIndices<N, T, L, R, Left::Indices, Right::Indices>;
56
57    fn offset_from_index(index: [usize; T]) -> usize {
58        let (a, b) = CSum::split(index);
59
60        let stride = Right::count();
61        let most_sig = Left::offset_from_index(a);
62        let least_sig = Right::offset_from_index(b);
63
64        most_sig * stride + least_sig
65    }
66
67    fn indices() -> Self::Indices {
68        CompoundIndices::zero()
69    }
70
71    fn count() -> usize {
72        Left::count() * Right::count()
73    }
74
75    fn for_each_index(mut f: impl FnMut([usize; T])) {
76        Left::for_each_index(|a| {
77            Right::for_each_index(|b| {
78                let idx = CSum::join((a, b));
79                f(idx)
80            })
81        });
82    }
83}
84
85/// Indexes over compound indices.
86#[derive(Default, Clone, Debug)]
87pub struct CompoundIndices<
88    const N: usize,
89    const T: usize,
90    const L: usize,
91    const R: usize,
92    Left,
93    Right,
94> {
95    first: Left,
96    first_cur: Option<Option<[usize; L]>>,
97    second: Right,
98}
99
100impl<const N: usize, const T: usize, const L: usize, const R: usize, Left, Right> Iterator
101    for CompoundIndices<N, T, L, R, Left, Right>
102where
103    Left: IndexIterator<L>,
104    Right: IndexIterator<R>,
105    CSum: CompoundSum<T, L, R>,
106{
107    type Item = [usize; T];
108
109    fn next(&mut self) -> Option<Self::Item> {
110        let b = match self.second.next() {
111            Some(b) => b,
112            None => {
113                // Advance left iterator
114                self.first_cur = Some(self.first.next());
115                // Reset right-most iterator
116                self.second = Right::zero();
117                // Get current value of left-most iterator
118                self.second.next()?
119            }
120        };
121
122        self.first_cur
123            .get_or_insert_with(|| self.first.next())
124            .map(|a| CSum::join((a, b)))
125    }
126}
127
128impl<const N: usize, const T: usize, const L: usize, const R: usize, Left, Right> IndexIterator<T>
129    for CompoundIndices<N, T, L, R, Left, Right>
130where
131    Left: IndexIterator<L>,
132    Right: IndexIterator<R>,
133    CSum: CompoundSum<T, L, R>,
134{
135    fn zero() -> Self {
136        Self {
137            first: Left::zero(),
138            first_cur: None,
139            second: Right::zero(),
140        }
141    }
142}
143
144/// A tensor of the form Tᵢ₍ⱼₖ₎
145pub type VecSym = Compound<3, 1, 2, Gen, Sym>;
146pub type SymVec = Compound<3, 2, 1, Sym, Gen>;
147pub type SymSym = Compound<4, 2, 2, Sym, Sym>;
148pub type VecSymVec = Compound<4, 3, 1, VecSym, Gen>;
aeon_tensor/compound.rs

aeon_tensor/
compound.rs
