acme_tensor/actions/iter/
indexed.rs

1/*
2    Appellation: indexed <mod>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5use crate::iter::LayoutIter;
6use crate::tensor::TensorBase;
7
8use super::Position;
9
10pub struct IndexedIter<'a, T: 'a> {
11    inner: LayoutIter,
12    scope: Option<&'a T>,
13    tensor: &'a TensorBase<T>,
14}
15
16impl<'a, T> IndexedIter<'a, T> {
17    pub fn new(tensor: &'a TensorBase<T>) -> Self {
18        Self {
19            inner: tensor.layout().iter(),
20            scope: None,
21            tensor,
22        }
23    }
24}
25
26impl<'a, T> Iterator for IndexedIter<'a, T> {
27    type Item = (&'a T, Position);
28
29    fn next(&mut self) -> Option<Self::Item> {
30        let pos = self.inner.next()?;
31        self.scope = self.tensor.get_by_index(pos.index());
32        self.scope.map(|scope| (scope, pos))
33    }
34}
35
36impl<'a, T> DoubleEndedIterator for IndexedIter<'a, T> {
37    fn next_back(&mut self) -> Option<Self::Item> {
38        let pos = self.inner.next_back()?;
39        self.scope = self.tensor.get_by_index(pos.index());
40        self.scope.map(|scope| (scope, pos))
41    }
42}
43
44impl<'a, T> ExactSizeIterator for IndexedIter<'a, T> {
45    fn len(&self) -> usize {
46        self.tensor.size()
47    }
48}