pasture_core/containers/
point_iterators.rs

1use std::marker::PhantomData;
2
3use crate::layout::PointType;
4
5use super::point_buffer::{BorrowedBuffer, InterleavedBuffer, InterleavedBufferMut};
6
7/// Iterator over strongly typed points by value
8pub struct PointIteratorByValue<'a, 'b, T: PointType, B: BorrowedBuffer<'a> + ?Sized>
9where
10    'a: 'b,
11{
12    buffer: &'b B,
13    current_index: usize,
14    _phantom: PhantomData<&'a T>,
15}
16
17impl<'a, 'b, T: PointType, B: BorrowedBuffer<'a> + ?Sized> From<&'b B>
18    for PointIteratorByValue<'a, 'b, T, B>
19{
20    fn from(value: &'b B) -> Self {
21        Self {
22            buffer: value,
23            current_index: 0,
24            _phantom: Default::default(),
25        }
26    }
27}
28
29impl<'a, 'b, T: PointType, B: BorrowedBuffer<'a> + ?Sized> Iterator
30    for PointIteratorByValue<'a, 'b, T, B>
31{
32    type Item = T;
33
34    fn next(&mut self) -> Option<Self::Item> {
35        if self.current_index == self.buffer.len() {
36            None
37        } else {
38            let mut point = T::zeroed();
39            let point_bytes = bytemuck::bytes_of_mut(&mut point);
40            self.buffer.get_point(self.current_index, point_bytes);
41            self.current_index += 1;
42            Some(point)
43        }
44    }
45
46    fn nth(&mut self, n: usize) -> Option<Self::Item> {
47        self.current_index += n;
48        self.next()
49    }
50}
51
52/// Iterator over strongly typed points by immutable borrow
53pub struct PointIteratorByRef<'a, T: PointType> {
54    point_data: &'a [T],
55    current_index: usize,
56}
57
58impl<'a, 'b, T: PointType, B: InterleavedBuffer<'b> + ?Sized> From<&'a B>
59    for PointIteratorByRef<'a, T>
60where
61    'b: 'a,
62{
63    fn from(value: &'a B) -> Self {
64        let points_memory = value.get_point_range_ref(0..value.len());
65        Self {
66            point_data: bytemuck::cast_slice(points_memory),
67            current_index: 0,
68        }
69    }
70}
71
72impl<'a, T: PointType> Iterator for PointIteratorByRef<'a, T> {
73    type Item = &'a T;
74
75    fn next(&mut self) -> Option<Self::Item> {
76        if self.current_index == self.point_data.len() {
77            None
78        } else {
79            let point = &self.point_data[self.current_index];
80            self.current_index += 1;
81            Some(point)
82        }
83    }
84
85    fn nth(&mut self, n: usize) -> Option<Self::Item> {
86        self.current_index += n;
87        self.next()
88    }
89}
90
91/// Iterator over strongly typed points by mutable borrow
92pub struct PointIteratorByMut<'a, T: PointType> {
93    point_data: &'a mut [T],
94    current_index: usize,
95    _phantom: PhantomData<T>,
96}
97
98impl<'a, 'b, T: PointType, B: InterleavedBufferMut<'b> + ?Sized> From<&'a mut B>
99    for PointIteratorByMut<'a, T>
100where
101    'b: 'a,
102{
103    fn from(value: &'a mut B) -> Self {
104        let memory_for_all_points = value.get_point_range_mut(0..value.len());
105        Self {
106            point_data: bytemuck::cast_slice_mut(memory_for_all_points),
107            current_index: 0,
108            _phantom: Default::default(),
109        }
110    }
111}
112
113impl<'a, T: PointType> Iterator for PointIteratorByMut<'a, T> {
114    type Item = &'a mut T;
115
116    fn next(&mut self) -> Option<Self::Item> {
117        if self.current_index == self.point_data.len() {
118            None
119        } else {
120            unsafe {
121                let memory = self.point_data.as_mut_ptr().add(self.current_index);
122                self.current_index += 1;
123                Some(&mut *memory)
124            }
125        }
126    }
127
128    fn nth(&mut self, n: usize) -> Option<Self::Item> {
129        self.current_index += n;
130        self.next()
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use rand::{thread_rng, Rng};
137
138    use crate::{
139        containers::{BorrowedBufferExt, BorrowedMutBufferExt, VectorBuffer},
140        test_utils::{CustomPointTypeSmall, DefaultPointDistribution},
141    };
142
143    #[test]
144    #[allow(clippy::iter_nth_zero)]
145    fn point_iterator_nth() {
146        const COUNT: usize = 16;
147        let mut points = thread_rng()
148            .sample_iter::<CustomPointTypeSmall, _>(DefaultPointDistribution)
149            .take(COUNT)
150            .collect::<VectorBuffer>();
151        let mut points_clone = points.clone();
152
153        let points_view = points.view::<CustomPointTypeSmall>();
154        {
155            let mut points_iter = points_view.clone().into_iter();
156            assert_eq!(points_iter.nth(0), Some(points_view.at(0)));
157        }
158        {
159            let mut points_iter = points_view.clone().into_iter();
160            assert_eq!(points_iter.nth(6), Some(points_view.at(6)));
161        }
162        {
163            let mut points_iter = points_view.clone().into_iter();
164            assert_eq!(points_iter.nth(COUNT), None);
165        }
166        {
167            let mut points_iter = points_view.clone().into_iter();
168            points_iter.nth(0);
169            assert_eq!(points_iter.nth(0), Some(points_view.at(1)));
170        }
171
172        {
173            let mut points_iter = points_view.iter();
174            assert_eq!(points_iter.nth(0), Some(points_view.at_ref(0)));
175        }
176        {
177            let mut points_iter = points_view.iter();
178            assert_eq!(points_iter.nth(6), Some(points_view.at_ref(6)));
179        }
180        {
181            let mut points_iter = points_view.iter();
182            assert_eq!(points_iter.nth(COUNT), None);
183        }
184        {
185            let mut points_iter = points_view.iter();
186            points_iter.nth(0);
187            assert_eq!(points_iter.nth(0), Some(points_view.at_ref(1)));
188        }
189
190        let mut points_view_mut = points.view_mut::<CustomPointTypeSmall>();
191        let mut cloned_points_view_mut = points_clone.view_mut::<CustomPointTypeSmall>();
192        {
193            let mut points_iter = points_view_mut.iter_mut();
194            assert_eq!(points_iter.nth(0), Some(cloned_points_view_mut.at_mut(0)));
195        }
196        {
197            let mut points_iter = points_view_mut.iter_mut();
198            assert_eq!(points_iter.nth(6), Some(cloned_points_view_mut.at_mut(6)));
199        }
200        {
201            let mut points_iter = points_view_mut.iter_mut();
202            assert_eq!(points_iter.nth(COUNT), None);
203        }
204        {
205            let mut points_iter = points_view_mut.iter_mut();
206            points_iter.nth(0);
207            assert_eq!(points_iter.nth(0), Some(cloned_points_view_mut.at_mut(1)));
208        }
209    }
210}