Skip to main content

spatialrust_core/
tensor_aoso.rs

1//! AoSoA chunk packing: interleaved `[x,y,z, …]` buffers for SIMD/GPU staging.
2//!
3//! Enabled by the `tensor-aoso` feature. The underlying cloud stays schema-SoA;
4//! packing copies one chunk at a time on demand.
5
6use crate::{
7    HasIntensity, HasNormals3, HasPositions3, PointCloud, SpatialResult, SpatialTensorChunk,
8};
9
10/// Explicit interleaved field layout for an [`AoSoAAttributeChunk`].
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct AoSoAAttributeLayout {
13    stride_f32: usize,
14    position_offsets: [usize; 3],
15    intensity_offset: Option<usize>,
16    normal_offsets: Option<[usize; 3]>,
17}
18
19impl AoSoAAttributeLayout {
20    /// Interleaved `[x, y, z, intensity]` layout.
21    pub const XYZ_INTENSITY: Self = Self {
22        stride_f32: 4,
23        position_offsets: [0, 1, 2],
24        intensity_offset: Some(3),
25        normal_offsets: None,
26    };
27
28    /// Interleaved `[x, y, z, nx, ny, nz]` layout.
29    pub const XYZ_NORMALS: Self = Self {
30        stride_f32: 6,
31        position_offsets: [0, 1, 2],
32        intensity_offset: None,
33        normal_offsets: Some([3, 4, 5]),
34    };
35
36    /// Interleaved `[x, y, z, intensity, nx, ny, nz]` layout.
37    pub const XYZ_INTENSITY_NORMALS: Self = Self {
38        stride_f32: 7,
39        position_offsets: [0, 1, 2],
40        intensity_offset: Some(3),
41        normal_offsets: Some([4, 5, 6]),
42    };
43
44    /// Returns the distance between adjacent points in `f32` elements.
45    #[must_use]
46    pub const fn stride_f32(self) -> usize {
47        self.stride_f32
48    }
49
50    /// Returns the x/y/z offsets within one point record.
51    #[must_use]
52    pub const fn position_offsets(self) -> [usize; 3] {
53        self.position_offsets
54    }
55
56    /// Returns the intensity offset when present.
57    #[must_use]
58    pub const fn intensity_offset(self) -> Option<usize> {
59        self.intensity_offset
60    }
61
62    /// Returns the normal x/y/z offsets when present.
63    #[must_use]
64    pub const fn normal_offsets(self) -> Option<[usize; 3]> {
65        self.normal_offsets
66    }
67}
68
69/// Owned interleaved positions and optional capability attributes for one chunk.
70#[derive(Clone, Debug, PartialEq)]
71pub struct AoSoAAttributeChunk {
72    data: Vec<f32>,
73    point_count: usize,
74    layout: AoSoAAttributeLayout,
75}
76
77impl AoSoAAttributeChunk {
78    fn pack(
79        chunk: &SpatialTensorChunk,
80        cloud: &PointCloud,
81        layout: AoSoAAttributeLayout,
82    ) -> SpatialResult<Self> {
83        let (x, y, z) = cloud.positions3()?;
84        let intensity = layout.intensity_offset().map(|_| cloud.intensity()).transpose()?;
85        let normals = layout.normal_offsets().map(|_| cloud.normals3()).transpose()?;
86        let range = chunk.range();
87        let point_count = range.len();
88        let mut data = Vec::with_capacity(point_count * layout.stride_f32());
89        for index in range {
90            data.extend_from_slice(&[x[index], y[index], z[index]]);
91            if let Some(intensity) = intensity {
92                data.push(intensity[index]);
93            }
94            if let Some((nx, ny, nz)) = normals {
95                data.extend_from_slice(&[nx[index], ny[index], nz[index]]);
96            }
97        }
98        Ok(Self { data, point_count, layout })
99    }
100
101    /// Packs `[x, y, z, intensity]` records using [`HasIntensity`].
102    pub fn pack_xyz_intensity(
103        chunk: &SpatialTensorChunk,
104        cloud: &PointCloud,
105    ) -> SpatialResult<Self> {
106        Self::pack(chunk, cloud, AoSoAAttributeLayout::XYZ_INTENSITY)
107    }
108
109    /// Packs `[x, y, z, nx, ny, nz]` records using [`HasNormals3`].
110    pub fn pack_xyz_normals(chunk: &SpatialTensorChunk, cloud: &PointCloud) -> SpatialResult<Self> {
111        Self::pack(chunk, cloud, AoSoAAttributeLayout::XYZ_NORMALS)
112    }
113
114    /// Packs `[x, y, z, intensity, nx, ny, nz]` capability records.
115    pub fn pack_xyz_intensity_normals(
116        chunk: &SpatialTensorChunk,
117        cloud: &PointCloud,
118    ) -> SpatialResult<Self> {
119        Self::pack(chunk, cloud, AoSoAAttributeLayout::XYZ_INTENSITY_NORMALS)
120    }
121
122    /// Returns the number of packed points.
123    #[must_use]
124    pub const fn len(&self) -> usize {
125        self.point_count
126    }
127
128    /// Returns whether this chunk contains no points.
129    #[must_use]
130    pub const fn is_empty(&self) -> bool {
131        self.point_count == 0
132    }
133
134    /// Returns the explicit stride and field-offset metadata.
135    #[must_use]
136    pub const fn layout(&self) -> AoSoAAttributeLayout {
137        self.layout
138    }
139
140    /// Returns the packed `f32` records.
141    #[must_use]
142    pub fn as_slice(&self) -> &[f32] {
143        &self.data
144    }
145}
146
147/// Interleaved XYZ layout for one [`SpatialTensorChunk`] (`3 * point_count` floats).
148#[derive(Clone, Debug, PartialEq)]
149pub struct AoSoAXyzChunk {
150    data: Vec<f32>,
151    point_count: usize,
152}
153
154impl AoSoAXyzChunk {
155    /// Packs chunk points into interleaved `[x,y,z, …]` order.
156    pub fn pack(chunk: &SpatialTensorChunk, cloud: &PointCloud) -> SpatialResult<Self> {
157        let (x, y, z) = cloud.positions3()?;
158        let range = chunk.range();
159        let point_count = range.len();
160        let mut data = Vec::with_capacity(point_count * 3);
161        for index in range {
162            data.push(x[index]);
163            data.push(y[index]);
164            data.push(z[index]);
165        }
166        Ok(Self { data, point_count })
167    }
168
169    /// Returns the number of points in this chunk.
170    #[must_use]
171    pub const fn len(&self) -> usize {
172        self.point_count
173    }
174
175    /// Returns whether this chunk contains zero points.
176    #[must_use]
177    pub fn is_empty(&self) -> bool {
178        self.point_count == 0
179    }
180
181    /// Returns the interleaved `[x,y,z, …]` slice (`len == 3 * point_count`).
182    #[must_use]
183    pub fn as_slice(&self) -> &[f32] {
184        &self.data
185    }
186
187    /// Returns one point by local chunk index.
188    #[must_use]
189    pub fn point(&self, local_index: usize) -> [f32; 3] {
190        let base = local_index * 3;
191        [self.data[base], self.data[base + 1], self.data[base + 2]]
192    }
193}
194
195impl SpatialTensorChunk {
196    /// Packs this chunk into an owned interleaved XYZ buffer.
197    pub fn pack_xyz(&self, cloud: &PointCloud) -> SpatialResult<AoSoAXyzChunk> {
198        AoSoAXyzChunk::pack(self, cloud)
199    }
200
201    /// Packs interleaved XYZ and intensity records.
202    pub fn pack_xyz_intensity(&self, cloud: &PointCloud) -> SpatialResult<AoSoAAttributeChunk> {
203        AoSoAAttributeChunk::pack_xyz_intensity(self, cloud)
204    }
205
206    /// Packs interleaved XYZ and normal records.
207    pub fn pack_xyz_normals(&self, cloud: &PointCloud) -> SpatialResult<AoSoAAttributeChunk> {
208        AoSoAAttributeChunk::pack_xyz_normals(self, cloud)
209    }
210
211    /// Packs interleaved XYZ, intensity, and normal records.
212    pub fn pack_xyz_intensity_normals(
213        &self,
214        cloud: &PointCloud,
215    ) -> SpatialResult<AoSoAAttributeChunk> {
216        AoSoAAttributeChunk::pack_xyz_intensity_normals(self, cloud)
217    }
218
219    /// Packs this chunk into `out`, returning the number of points written.
220    ///
221    /// `out` is cleared first and resized to `3 * chunk.len()` floats.
222    pub fn pack_xyz_into(&self, cloud: &PointCloud, out: &mut Vec<f32>) -> SpatialResult<usize> {
223        let (x, y, z) = cloud.positions3()?;
224        let range = self.range();
225        out.clear();
226        out.reserve(range.len() * 3);
227        for index in range {
228            out.push(x[index]);
229            out.push(y[index]);
230            out.push(z[index]);
231        }
232        Ok(out.len() / 3)
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use crate::{AoSoAAttributeLayout, PointCloudBuilder, SpatialTensor, StandardSchemas};
239
240    #[test]
241    fn pack_matches_column_slices() {
242        let mut builder = PointCloudBuilder::xyz();
243        builder.push_point([1.0, 2.0, 3.0]).unwrap();
244        builder.push_point([4.0, 5.0, 6.0]).unwrap();
245        builder.push_point([7.0, 8.0, 9.0]).unwrap();
246        let cloud = builder.build().unwrap();
247        let chunk = SpatialTensor::new(&cloud, 2).unwrap().chunks().next().unwrap();
248
249        let packed = chunk.pack_xyz(&cloud).unwrap();
250        assert_eq!(packed.len(), 2);
251        assert_eq!(packed.as_slice(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
252        assert_eq!(packed.point(1), [4.0, 5.0, 6.0]);
253    }
254
255    #[test]
256    fn pack_into_reuses_buffer() {
257        let mut builder = PointCloudBuilder::xyz();
258        builder.push_point([0.0, 1.0, 2.0]).unwrap();
259        let cloud = builder.build().unwrap();
260        let chunk = SpatialTensor::new(&cloud, 4).unwrap().chunks().next().unwrap();
261
262        let mut buffer = vec![f32::NAN; 99];
263        let count = chunk.pack_xyz_into(&cloud, &mut buffer).unwrap();
264        assert_eq!(count, 1);
265        assert_eq!(buffer, vec![0.0, 1.0, 2.0]);
266    }
267
268    #[test]
269    fn packs_composite_capabilities_with_explicit_layout() {
270        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyzinormal());
271        builder.push_point([1.0, 2.0, 3.0, 0.5, 0.0, 1.0, 0.0]).unwrap();
272        builder.push_point([4.0, 5.0, 6.0, 0.8, 1.0, 0.0, 0.0]).unwrap();
273        let cloud = builder.build().unwrap();
274        let chunk = SpatialTensor::new(&cloud, 8).unwrap().chunks().next().unwrap();
275
276        let packed = chunk.pack_xyz_intensity_normals(&cloud).unwrap();
277        assert_eq!(packed.layout(), AoSoAAttributeLayout::XYZ_INTENSITY_NORMALS);
278        assert_eq!(packed.layout().stride_f32(), 7);
279        assert_eq!(packed.layout().intensity_offset(), Some(3));
280        assert_eq!(packed.layout().normal_offsets(), Some([4, 5, 6]));
281        assert_eq!(
282            packed.as_slice(),
283            &[1.0, 2.0, 3.0, 0.5, 0.0, 1.0, 0.0, 4.0, 5.0, 6.0, 0.8, 1.0, 0.0, 0.0]
284        );
285    }
286
287    #[test]
288    fn attribute_packers_reject_missing_capabilities() {
289        let mut builder = PointCloudBuilder::xyz();
290        builder.push_point([1.0, 2.0, 3.0]).unwrap();
291        let cloud = builder.build().unwrap();
292        let chunk = SpatialTensor::new(&cloud, 4).unwrap().chunks().next().unwrap();
293
294        assert!(chunk.pack_xyz_intensity(&cloud).is_err());
295        assert!(chunk.pack_xyz_normals(&cloud).is_err());
296    }
297}