pasture_core/containers/
attribute_iterators.rs

1use std::marker::PhantomData;
2
3use crate::layout::{PointAttributeDefinition, PointAttributeMember, PrimitiveType};
4
5use super::point_buffer::{BorrowedBuffer, ColumnarBuffer, ColumnarBufferMut};
6
7/// An iterator over strongly typed attribute data in a point buffer. Returns attribute data
8/// by value and makes assumptions about the memory layout of the underlying buffer
9pub struct AttributeIteratorByValue<'a, 'b, T: PrimitiveType, B: BorrowedBuffer<'a> + ?Sized>
10where
11    'a: 'b,
12{
13    buffer: &'b B,
14    attribute_member: &'b PointAttributeMember,
15    current_index: usize,
16    _phantom: PhantomData<&'a T>,
17}
18
19impl<'a, 'b, T: PrimitiveType, B: BorrowedBuffer<'a> + ?Sized>
20    AttributeIteratorByValue<'a, 'b, T, B>
21{
22    pub(crate) fn new(buffer: &'b B, attribute: &PointAttributeDefinition) -> Self {
23        Self {
24            attribute_member: buffer
25                .point_layout()
26                .get_attribute(attribute)
27                .expect("Attribute not found in PointLayout of buffer"),
28            buffer,
29            current_index: 0,
30            _phantom: Default::default(),
31        }
32    }
33}
34
35impl<'a, 'b, T: PrimitiveType, B: BorrowedBuffer<'a> + ?Sized> Iterator
36    for AttributeIteratorByValue<'a, 'b, T, B>
37{
38    type Item = T;
39
40    fn next(&mut self) -> Option<Self::Item> {
41        if self.current_index == self.buffer.len() {
42            None
43        } else {
44            let mut attribute = T::zeroed();
45            let attribute_bytes = bytemuck::bytes_of_mut(&mut attribute);
46            // This is safe because in `new` we obtain the `attribute_member` from the point layout of the buffer
47            unsafe {
48                self.buffer.get_attribute_unchecked(
49                    self.attribute_member,
50                    self.current_index,
51                    attribute_bytes,
52                );
53            }
54            self.current_index += 1;
55            Some(attribute)
56        }
57    }
58
59    fn nth(&mut self, n: usize) -> Option<Self::Item> {
60        self.current_index += n;
61        self.next()
62    }
63}
64
65/// Like [`AttributeIteratorByValue`], but returns attribute data by immutable reference. Can only be
66/// constructed from a buffer that implements [`ColumnarBuffer`]
67pub struct AttributeIteratorByRef<'a, T: PrimitiveType> {
68    attribute_data: &'a [T],
69    current_index: usize,
70}
71
72impl<'a, T: PrimitiveType> AttributeIteratorByRef<'a, T> {
73    pub(crate) fn new<'b, B: ColumnarBuffer<'b> + ?Sized>(
74        buffer: &'a B,
75        attribute: &PointAttributeDefinition,
76    ) -> Self
77    where
78        'b: 'a,
79    {
80        let attribute_memory = buffer.get_attribute_range_ref(attribute, 0..buffer.len());
81        Self {
82            attribute_data: bytemuck::cast_slice(attribute_memory),
83            current_index: 0,
84        }
85    }
86}
87
88impl<'a, T: PrimitiveType> Iterator for AttributeIteratorByRef<'a, T> {
89    type Item = &'a T;
90
91    fn next(&mut self) -> Option<Self::Item> {
92        if self.current_index == self.attribute_data.len() {
93            None
94        } else {
95            let ret = &self.attribute_data[self.current_index];
96            self.current_index += 1;
97            Some(ret)
98        }
99    }
100
101    fn nth(&mut self, n: usize) -> Option<Self::Item> {
102        self.current_index += n;
103        self.next()
104    }
105}
106
107/// Like [`AttributeIteratorByRef`], but returns attribute data by mutable reference, allowing mutation
108/// of the attribute data in-place. Can only be constructed from a buffer that implements [`ColumnarBufferMut`]
109pub struct AttributeIteratorByMut<'a, T: PrimitiveType> {
110    attribute_data: &'a mut [T],
111    current_index: usize,
112}
113
114impl<'a, T: PrimitiveType> AttributeIteratorByMut<'a, T> {
115    pub(crate) fn new<'b, B: ColumnarBufferMut<'b> + ?Sized>(
116        buffer: &'a mut B,
117        attribute: &PointAttributeDefinition,
118    ) -> Self
119    where
120        'b: 'a,
121    {
122        let attribute_memory = buffer.get_attribute_range_mut(attribute, 0..buffer.len());
123        Self {
124            attribute_data: bytemuck::cast_slice_mut(attribute_memory),
125            current_index: 0,
126        }
127    }
128}
129
130impl<'a, T: PrimitiveType> Iterator for AttributeIteratorByMut<'a, T> {
131    type Item = &'a mut T;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        if self.current_index == self.attribute_data.len() {
135            None
136        } else {
137            unsafe {
138                let attribute_ptr = self.attribute_data.as_mut_ptr().add(self.current_index);
139                self.current_index += 1;
140                Some(&mut *attribute_ptr)
141            }
142        }
143    }
144
145    fn nth(&mut self, n: usize) -> Option<Self::Item> {
146        self.current_index += n;
147        self.next()
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use nalgebra::Vector3;
154    use rand::{thread_rng, Rng};
155
156    use crate::{
157        containers::{BorrowedBufferExt, BorrowedMutBufferExt, HashMapBuffer},
158        layout::attributes::POSITION_3D,
159        test_utils::{CustomPointTypeSmall, DefaultPointDistribution},
160    };
161
162    #[test]
163    #[allow(clippy::iter_nth_zero)]
164    fn attribute_iterator_nth() {
165        const COUNT: usize = 16;
166        let mut points = thread_rng()
167            .sample_iter::<CustomPointTypeSmall, _>(DefaultPointDistribution)
168            .take(COUNT)
169            .collect::<HashMapBuffer>();
170        let mut points_clone = points.clone();
171
172        let positions_view = points.view_attribute::<Vector3<f64>>(&POSITION_3D);
173        {
174            let mut positions_iter = positions_view.clone().into_iter();
175            assert_eq!(positions_iter.nth(0), Some(positions_view.at(0)));
176        }
177        {
178            let mut positions_iter = positions_view.clone().into_iter();
179            assert_eq!(positions_iter.nth(6), Some(positions_view.at(6)));
180        }
181        {
182            let mut positions_iter = positions_view.clone().into_iter();
183            assert_eq!(positions_iter.nth(COUNT), None);
184        }
185        {
186            let mut positions_iter = positions_view.clone().into_iter();
187            positions_iter.nth(0);
188            assert_eq!(positions_iter.nth(0), Some(positions_view.at(1)));
189        }
190
191        {
192            let mut positions_iter = positions_view.iter();
193            assert_eq!(positions_iter.nth(0), Some(positions_view.at_ref(0)));
194        }
195        {
196            let mut positions_iter = positions_view.iter();
197            assert_eq!(positions_iter.nth(6), Some(positions_view.at_ref(6)));
198        }
199        {
200            let mut positions_iter = positions_view.iter();
201            assert_eq!(positions_iter.nth(COUNT), None);
202        }
203        {
204            let mut positions_iter = positions_view.iter();
205            positions_iter.nth(0);
206            assert_eq!(positions_iter.nth(0), Some(positions_view.at_ref(1)));
207        }
208
209        {
210            let mut positions_view_mut = points.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
211            let mut cloned_positions_view_mut =
212                points_clone.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
213            let mut positions_iter = positions_view_mut.iter_mut();
214            assert_eq!(
215                positions_iter.nth(0),
216                Some(cloned_positions_view_mut.at_mut(0))
217            );
218        }
219        {
220            let mut positions_view_mut = points.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
221            let mut cloned_positions_view_mut =
222                points_clone.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
223            let mut positions_iter = positions_view_mut.iter_mut();
224            assert_eq!(
225                positions_iter.nth(6),
226                Some(cloned_positions_view_mut.at_mut(6))
227            );
228        }
229        {
230            let mut positions_view_mut = points.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
231            let mut positions_iter = positions_view_mut.iter_mut();
232            assert_eq!(positions_iter.nth(COUNT), None);
233        }
234        {
235            let mut positions_view_mut = points.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
236            let mut cloned_positions_view_mut =
237                points_clone.view_attribute_mut::<Vector3<f64>>(&POSITION_3D);
238            let mut positions_iter = positions_view_mut.iter_mut();
239            positions_iter.nth(0);
240            assert_eq!(
241                positions_iter.nth(0),
242                Some(cloned_positions_view_mut.at_mut(1))
243            );
244        }
245    }
246}
pasture_core/containers/attribute_iterators.rs

pasture_core/containers/
attribute_iterators.rs
