winsfs_core/saf/iter/
sites.rs

1use rayon::iter::{
2    plumbing::{bridge, Consumer, Producer, ProducerCallback, UnindexedConsumer},
3    IndexedParallelIterator, IntoParallelIterator, ParallelIterator,
4};
5
6use crate::saf::{AsSiteView, Saf, SafView, SiteView};
7
8/// A type that can be turned into an iterator over SAF sites.
9pub trait IntoSiteIterator<const N: usize> {
10    /// The type of each individual site.
11    type Item: AsSiteView<N>;
12    /// The type of iterator.
13    type Iter: ExactSizeIterator<Item = Self::Item>;
14
15    /// Convert this type into a SAF site iterator.
16    fn into_site_iter(self) -> Self::Iter;
17}
18
19impl<'a, const N: usize> IntoSiteIterator<N> for &'a Saf<N> {
20    type Item = SiteView<'a, N>;
21    type Iter = SiteIter<'a, N>;
22
23    fn into_site_iter(self) -> Self::Iter {
24        SiteIter::new(self.view())
25    }
26}
27
28impl<'a, const N: usize> IntoSiteIterator<N> for SafView<'a, N> {
29    type Item = SiteView<'a, N>;
30    type Iter = SiteIter<'a, N>;
31
32    fn into_site_iter(self) -> Self::Iter {
33        SiteIter::new(self)
34    }
35}
36
37impl<'a, 'b, const N: usize> IntoSiteIterator<N> for &'b SafView<'a, N> {
38    type Item = SiteView<'a, N>;
39    type Iter = SiteIter<'a, N>;
40
41    fn into_site_iter(self) -> Self::Iter {
42        SiteIter::new(*self)
43    }
44}
45
46impl<const N: usize, T> IntoSiteIterator<N> for T
47where
48    T: IntoIterator,
49    T::IntoIter: ExactSizeIterator,
50    T::Item: AsSiteView<N>,
51{
52    type Item = T::Item;
53    type Iter = T::IntoIter;
54
55    fn into_site_iter(self) -> Self::Iter {
56        self.into_iter()
57    }
58}
59
60/// A type that can be turned into a parallel iterator over SAF sites.
61pub trait IntoParallelSiteIterator<const N: usize> {
62    /// The type of each individual site.
63    type Item: AsSiteView<N>;
64    /// The type of iterator.
65    type Iter: IndexedParallelIterator<Item = Self::Item>;
66
67    /// Convert this type into a parallel SAF site iterator.
68    fn into_par_site_iter(self) -> Self::Iter;
69}
70
71impl<'a, const N: usize> IntoParallelSiteIterator<N> for &'a Saf<N> {
72    type Item = SiteView<'a, N>;
73    type Iter = ParSiteIter<'a, N>;
74
75    fn into_par_site_iter(self) -> Self::Iter {
76        ParSiteIter::new(self.view())
77    }
78}
79
80impl<'a, const N: usize> IntoParallelSiteIterator<N> for SafView<'a, N> {
81    type Item = SiteView<'a, N>;
82    type Iter = ParSiteIter<'a, N>;
83
84    fn into_par_site_iter(self) -> Self::Iter {
85        ParSiteIter::new(self)
86    }
87}
88
89impl<'a, 'b, const N: usize> IntoParallelSiteIterator<N> for &'b SafView<'a, N> {
90    type Item = SiteView<'a, N>;
91    type Iter = ParSiteIter<'a, N>;
92
93    fn into_par_site_iter(self) -> Self::Iter {
94        ParSiteIter::new(*self)
95    }
96}
97
98impl<const N: usize, T> IntoParallelSiteIterator<N> for T
99where
100    T: IntoParallelIterator,
101    T::Iter: IndexedParallelIterator,
102    T::Item: AsSiteView<N>,
103{
104    type Item = T::Item;
105    type Iter = T::Iter;
106
107    fn into_par_site_iter(self) -> Self::Iter {
108        self.into_par_iter()
109    }
110}
111
112/// An iterator over SAF sites.
113#[derive(Debug)]
114pub struct SiteIter<'a, const N: usize> {
115    iter: ::std::slice::Chunks<'a, f32>,
116    shape: [usize; N],
117}
118
119impl<'a, const N: usize> SiteIter<'a, N> {
120    pub(in crate::saf) fn new(saf: SafView<'a, N>) -> Self {
121        let iter = saf.values.chunks(saf.width());
122
123        Self {
124            iter,
125            shape: saf.shape,
126        }
127    }
128}
129
130impl<'a, const N: usize> Iterator for SiteIter<'a, N> {
131    type Item = SiteView<'a, N>;
132
133    #[inline]
134    fn next(&mut self) -> Option<Self::Item> {
135        self.iter
136            .next()
137            .map(|item| SiteView::new_unchecked(item, self.shape))
138    }
139
140    fn size_hint(&self) -> (usize, Option<usize>) {
141        self.iter.size_hint()
142    }
143}
144
145impl<'a, const N: usize> ExactSizeIterator for SiteIter<'a, N> {
146    fn len(&self) -> usize {
147        self.iter.len()
148    }
149}
150
151impl<'a, const N: usize> DoubleEndedIterator for SiteIter<'a, N> {
152    fn next_back(&mut self) -> Option<Self::Item> {
153        self.iter
154            .next_back()
155            .map(|item| SiteView::new_unchecked(item, self.shape))
156    }
157}
158
159/// A parallel iterator over SAF sites.
160#[derive(Debug)]
161pub struct ParSiteIter<'a, const N: usize> {
162    values: &'a [f32],
163    shape: [usize; N],
164    chunk_size: usize,
165}
166
167impl<'a, const N: usize> ParSiteIter<'a, N> {
168    pub(in crate::saf) fn new(saf: SafView<'a, N>) -> Self {
169        Self {
170            values: saf.values,
171            shape: saf.shape,
172            chunk_size: saf.width(),
173        }
174    }
175}
176
177/*
178    All the trait impls below are mainly boilerplate, and largely adapted from the
179    implementation of rayon::slice::Chunks,
180    see https://docs.rs/rayon/latest/src/rayon/slice/chunks.rs.html#8-11
181*/
182
183impl<'a, const N: usize> ParallelIterator for ParSiteIter<'a, N> {
184    type Item = SiteView<'a, N>;
185
186    fn drive_unindexed<C>(self, consumer: C) -> C::Result
187    where
188        C: UnindexedConsumer<Self::Item>,
189    {
190        bridge(self, consumer)
191    }
192
193    fn opt_len(&self) -> Option<usize> {
194        Some(self.len())
195    }
196}
197
198impl<'a, const N: usize> IndexedParallelIterator for ParSiteIter<'a, N> {
199    fn drive<C>(self, consumer: C) -> C::Result
200    where
201        C: Consumer<Self::Item>,
202    {
203        bridge(self, consumer)
204    }
205
206    fn len(&self) -> usize {
207        let n = self.values.len();
208        if n == 0 {
209            0
210        } else {
211            (n - 1) / self.chunk_size + 1
212        }
213    }
214
215    fn with_producer<CB>(self, callback: CB) -> CB::Output
216    where
217        CB: ProducerCallback<Self::Item>,
218    {
219        callback.callback(SiteProducer {
220            values: self.values,
221            shape: self.shape,
222            chunk_size: self.chunk_size,
223        })
224    }
225}
226
227struct SiteProducer<'a, const N: usize> {
228    values: &'a [f32],
229    shape: [usize; N],
230    chunk_size: usize,
231}
232
233impl<'a, const N: usize> Producer for SiteProducer<'a, N> {
234    type Item = SiteView<'a, N>;
235    type IntoIter = SiteIter<'a, N>;
236
237    fn into_iter(self) -> Self::IntoIter {
238        SiteIter {
239            iter: self.values.chunks(self.chunk_size),
240            shape: self.shape,
241        }
242    }
243
244    fn split_at(self, index: usize) -> (Self, Self) {
245        let elem_index = self.values.len().min(index * self.chunk_size);
246        let (left, right) = self.values.split_at(elem_index);
247
248        (
249            Self {
250                values: left,
251                shape: self.shape,
252                chunk_size: self.chunk_size,
253            },
254            Self {
255                values: right,
256                shape: self.shape,
257                chunk_size: self.chunk_size,
258            },
259        )
260    }
261}
winsfs_core/saf/iter/sites.rs

winsfs_core/saf/iter/
sites.rs
