numrst/core/ndarray/
iter.rs

1use crate::{StorageIndices, StorageRef, WithDType};
2
3use super::NdArray;
4
5pub struct NdArrayIter<'a, T> {
6    indexes: StorageIndices<'a>,
7    storage: StorageRef<'a, T>,
8}
9
10impl<'a, T: WithDType> Iterator for NdArrayIter<'a, T> {
11    type Item = T;
12
13    fn next(&mut self) -> Option<T> {
14        let index = self.indexes.next()?;
15        return self.storage.get(index)
16    }
17}
18
19impl<T: WithDType> NdArray<T> {
20    pub fn iter(&self) -> NdArrayIter<T> {
21        NdArrayIter {
22            indexes: self.0.layout.storage_indices(),
23            storage: self.storage_ref(0)
24        }
25    }
26}
27
28pub trait ResettableIterator: Iterator {
29    fn reset(&mut self);
30}
31
32impl<'a, T: WithDType> ResettableIterator for NdArrayIter<'a, T> {
33    fn reset(&mut self) {
34        self.indexes.reset();
35    }
36}
37
38impl<'a, T: WithDType> ExactSizeIterator for NdArrayIter<'a, T> {
39    fn len(&self) -> usize {
40        self.indexes.len()
41    }
42}
numrst/core/ndarray/iter.rs

numrst/core/ndarray/
iter.rs
