rs_pcd/storage/
mod.rs

1// Copyright 2025 bigpear0201
2
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6
7//     http://www.apache.org/licenses/LICENSE-2.0
8
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::header::ValueType;
16use std::collections::HashMap;
17
18pub mod view;
19pub use view::{ColumnView, PointView};
20
21#[derive(Debug, Clone)]
22pub enum Column {
23    U8(Vec<u8>),
24    U16(Vec<u16>),
25    U32(Vec<u32>),
26    I8(Vec<i8>),
27    I16(Vec<i16>),
28    I32(Vec<i32>),
29    F32(Vec<f32>),
30    F64(Vec<f64>),
31}
32
33impl Column {
34    pub fn new(value_type: ValueType, capacity: usize) -> Self {
35        match value_type {
36            ValueType::U8 => Column::U8(vec![0; capacity]),
37            ValueType::U16 => Column::U16(vec![0; capacity]),
38            ValueType::U32 => Column::U32(vec![0; capacity]),
39            ValueType::I8 => Column::I8(vec![0; capacity]),
40            ValueType::I16 => Column::I16(vec![0; capacity]),
41            ValueType::I32 => Column::I32(vec![0; capacity]),
42            ValueType::F32 => Column::F32(vec![0.0; capacity]),
43            ValueType::F64 => Column::F64(vec![0.0; capacity]),
44        }
45    }
46
47    pub fn resize(&mut self, new_len: usize) {
48        match self {
49            Column::U8(v) => v.resize(new_len, 0),
50            Column::U16(v) => v.resize(new_len, 0),
51            Column::U32(v) => v.resize(new_len, 0),
52            Column::I8(v) => v.resize(new_len, 0),
53            Column::I16(v) => v.resize(new_len, 0),
54            Column::I32(v) => v.resize(new_len, 0),
55            Column::F32(v) => v.resize(new_len, 0.0),
56            Column::F64(v) => v.resize(new_len, 0.0),
57        }
58    }
59
60    #[must_use]
61    pub fn len(&self) -> usize {
62        match self {
63            Column::U8(v) => v.len(),
64            Column::U16(v) => v.len(),
65            Column::U32(v) => v.len(),
66            Column::I8(v) => v.len(),
67            Column::I16(v) => v.len(),
68            Column::I32(v) => v.len(),
69            Column::F32(v) => v.len(),
70            Column::F64(v) => v.len(),
71        }
72    }
73
74    #[must_use]
75    pub fn is_empty(&self) -> bool {
76        self.len() == 0
77    }
78
79    pub fn as_f32_slice(&self) -> Option<&[f32]> {
80        if let Column::F32(v) = self {
81            Some(v)
82        } else {
83            None
84        }
85    }
86
87    // Mutable access for decoders
88    // Safe internal mutable access
89    pub fn as_u8_mut(&mut self) -> Option<&mut Vec<u8>> {
90        if let Column::U8(v) = self {
91            Some(v)
92        } else {
93            None
94        }
95    }
96    pub fn as_u16_mut(&mut self) -> Option<&mut Vec<u16>> {
97        if let Column::U16(v) = self {
98            Some(v)
99        } else {
100            None
101        }
102    }
103    pub fn as_u32_mut(&mut self) -> Option<&mut Vec<u32>> {
104        if let Column::U32(v) = self {
105            Some(v)
106        } else {
107            None
108        }
109    }
110    pub fn as_i8_mut(&mut self) -> Option<&mut Vec<i8>> {
111        if let Column::I8(v) = self {
112            Some(v)
113        } else {
114            None
115        }
116    }
117    pub fn as_i16_mut(&mut self) -> Option<&mut Vec<i16>> {
118        if let Column::I16(v) = self {
119            Some(v)
120        } else {
121            None
122        }
123    }
124    pub fn as_i32_mut(&mut self) -> Option<&mut Vec<i32>> {
125        if let Column::I32(v) = self {
126            Some(v)
127        } else {
128            None
129        }
130    }
131    pub fn as_f32_mut(&mut self) -> Option<&mut Vec<f32>> {
132        if let Column::F32(v) = self {
133            Some(v)
134        } else {
135            None
136        }
137    }
138    pub fn as_f64_mut(&mut self) -> Option<&mut Vec<f64>> {
139        if let Column::F64(v) = self {
140            Some(v)
141        } else {
142            None
143        }
144    }
145
146    // Read access variants (useful for Writer)
147    pub fn as_u8(&self) -> Option<&[u8]> {
148        if let Column::U8(v) = self {
149            Some(v)
150        } else {
151            None
152        }
153    }
154    pub fn as_u16(&self) -> Option<&[u16]> {
155        if let Column::U16(v) = self {
156            Some(v)
157        } else {
158            None
159        }
160    }
161    pub fn as_u32(&self) -> Option<&[u32]> {
162        if let Column::U32(v) = self {
163            Some(v)
164        } else {
165            None
166        }
167    }
168    pub fn as_i8(&self) -> Option<&[i8]> {
169        if let Column::I8(v) = self {
170            Some(v)
171        } else {
172            None
173        }
174    }
175    pub fn as_i16(&self) -> Option<&[i16]> {
176        if let Column::I16(v) = self {
177            Some(v)
178        } else {
179            None
180        }
181    }
182    pub fn as_i32(&self) -> Option<&[i32]> {
183        if let Column::I32(v) = self {
184            Some(v)
185        } else {
186            None
187        }
188    }
189    pub fn as_f32(&self) -> Option<&[f32]> {
190        if let Column::F32(v) = self {
191            Some(v)
192        } else {
193            None
194        }
195    }
196    pub fn as_f64(&self) -> Option<&[f64]> {
197        if let Column::F64(v) = self {
198            Some(v)
199        } else {
200            None
201        }
202    }
203
204    // Unsafe methods to get mutable slice for parallel writing.
205    // Safety: Caller must ensure exclusive access to the slice regions if writing in parallel.
206    pub unsafe fn as_ptr_mut(&mut self) -> (*mut u8, usize) {
207        match self {
208            Column::U8(v) => (v.as_mut_ptr() as *mut u8, v.len() * 1),
209            Column::U16(v) => (v.as_mut_ptr() as *mut u8, v.len() * 2),
210            Column::U32(v) => (v.as_mut_ptr() as *mut u8, v.len() * 4),
211            Column::I8(v) => (v.as_mut_ptr() as *mut u8, v.len() * 1),
212            Column::I16(v) => (v.as_mut_ptr() as *mut u8, v.len() * 2),
213            Column::I32(v) => (v.as_mut_ptr() as *mut u8, v.len() * 4),
214            Column::F32(v) => (v.as_mut_ptr() as *mut u8, v.len() * 4),
215            Column::F64(v) => (v.as_mut_ptr() as *mut u8, v.len() * 8),
216        }
217    }
218}
219
220/// SoA (Structure of Arrays) storage for point cloud data.
221/// 
222/// Internally uses Vec<Column> for O(1) index-based access, with a HashMap
223/// for name-based lookups. This provides efficient iteration while maintaining
224/// backwards-compatible named access.
225#[derive(Debug)]
226pub struct PointBlock {
227    /// Column data stored in schema order for O(1) indexed access
228    columns: Vec<Column>,
229    /// Field names in schema order
230    schema: Vec<String>,
231    /// Name to index mapping for backwards-compatible get_column(name) API
232    name_to_index: HashMap<String, usize>,
233    /// Number of points
234    pub len: usize,
235}
236
237impl Default for PointBlock {
238    fn default() -> Self {
239        Self {
240            columns: Vec::new(),
241            schema: Vec::new(),
242            name_to_index: HashMap::new(),
243            len: 0,
244        }
245    }
246}
247
248impl PointBlock {
249    pub fn new(schema: &Vec<(String, ValueType)>, capacity: usize) -> Self {
250        let mut columns = Vec::with_capacity(schema.len());
251        let mut names = Vec::with_capacity(schema.len());
252        let mut name_to_index = HashMap::with_capacity(schema.len());
253
254        for (i, (name, dtype)) in schema.iter().enumerate() {
255            columns.push(Column::new(*dtype, capacity));
256            names.push(name.clone());
257            name_to_index.insert(name.clone(), i);
258        }
259
260        PointBlock {
261            columns,
262            schema: names,
263            name_to_index,
264            len: capacity,
265        }
266    }
267
268    pub fn resize(&mut self, new_len: usize) {
269        for col in &mut self.columns {
270            col.resize(new_len);
271        }
272        self.len = new_len;
273    }
274
275    /// Get a column by name (backwards-compatible API).
276    /// For performance-critical code, prefer `get_column_by_index`.
277    #[must_use]
278    pub fn get_column(&self, name: &str) -> Option<&Column> {
279        self.name_to_index.get(name).map(|&idx| &self.columns[idx])
280    }
281
282    /// Get a mutable column by name (backwards-compatible API).
283    /// For performance-critical code, prefer `get_column_mut_by_index`.
284    pub fn get_column_mut(&mut self, name: &str) -> Option<&mut Column> {
285        if let Some(&idx) = self.name_to_index.get(name) {
286            Some(&mut self.columns[idx])
287        } else {
288            None
289        }
290    }
291
292    /// O(1) indexed access to a column.
293    #[inline]
294    #[must_use]
295    pub fn get_column_by_index(&self, index: usize) -> Option<&Column> {
296        self.columns.get(index)
297    }
298
299    /// O(1) mutable indexed access to a column.
300    #[inline]
301    pub fn get_column_mut_by_index(&mut self, index: usize) -> Option<&mut Column> {
302        self.columns.get_mut(index)
303    }
304
305    /// Get the index of a column by name.
306    #[inline]
307    #[must_use]
308    pub fn get_column_index(&self, name: &str) -> Option<usize> {
309        self.name_to_index.get(name).copied()
310    }
311
312    /// Get the schema (field names in order).
313    #[must_use]
314    pub fn schema(&self) -> &[String] {
315        &self.schema
316    }
317
318    /// Number of columns.
319    #[must_use]
320    pub fn num_columns(&self) -> usize {
321        self.columns.len()
322    }
323
324    /// Optimized: Get multiple mutable columns simultaneously.
325    /// Returns None if any column is missing or if names contain duplicates.
326    /// This avoids O(N*M) lookup inside tight loops.
327    pub fn get_columns_mut(&mut self, names: &[String]) -> Option<Vec<&mut Column>> {
328        // Simple check for duplicates (O(M^2) but M is small, e.g. < 10)
329        for i in 0..names.len() {
330            for j in i + 1..names.len() {
331                if names[i] == names[j] {
332                    return None; // Duplicate requested
333                }
334            }
335        }
336
337        // Get indices for all requested names
338        let mut indices = Vec::with_capacity(names.len());
339        for name in names {
340            if let Some(&idx) = self.name_to_index.get(name) {
341                indices.push(idx);
342            } else {
343                return None; // Missing column
344            }
345        }
346
347        // Use raw pointers to bypass borrow checker for multiple mutable references
348        // Safety: We verified all indices are unique above, so all pointers point to disjoint memory.
349        let mut results = Vec::with_capacity(names.len());
350        let base_ptr = self.columns.as_mut_ptr();
351        for idx in indices {
352            unsafe {
353                results.push(&mut *base_ptr.add(idx));
354            }
355        }
356        Some(results)
357    }
358
359    /// Access underlying columns slice (for iteration).
360    #[must_use]
361    pub fn columns(&self) -> &[Column] {
362        &self.columns
363    }
364
365    /// Access underlying columns mutably.
366    pub fn columns_mut(&mut self) -> &mut [Column] {
367        &mut self.columns
368    }
369
370    // ========================
371    // Typed Convenience Accessors
372    // ========================
373
374    /// Get XYZ coordinates as f32 slices.
375    /// Returns None if any of x, y, z columns are missing or not F32.
376    #[must_use]
377    pub fn xyz(&self) -> Option<(&[f32], &[f32], &[f32])> {
378        let x = self.get_column("x")?.as_f32()?;
379        let y = self.get_column("y")?.as_f32()?;
380        let z = self.get_column("z")?.as_f32()?;
381        Some((x, y, z))
382    }
383
384    /// Get XYZ + intensity as f32 slices.
385    /// Returns None if any column is missing or has wrong type.
386    #[must_use]
387    pub fn xyzi(&self) -> Option<(&[f32], &[f32], &[f32], &[f32])> {
388        let x = self.get_column("x")?.as_f32()?;
389        let y = self.get_column("y")?.as_f32()?;
390        let z = self.get_column("z")?.as_f32()?;
391        let i = self.get_column("intensity")?.as_f32()?;
392        Some((x, y, z, i))
393    }
394
395    /// Get XYZ + RGB (packed as u32) slices.
396    /// Returns None if any column is missing or has wrong type.
397    #[must_use]
398    pub fn xyzrgb(&self) -> Option<(&[f32], &[f32], &[f32], &[u32])> {
399        let x = self.get_column("x")?.as_f32()?;
400        let y = self.get_column("y")?.as_f32()?;
401        let z = self.get_column("z")?.as_f32()?;
402        let rgb = self.get_column("rgb")?.as_u32()?;
403        Some((x, y, z, rgb))
404    }
405
406    /// Get XYZ + intensity + ring (common LiDAR format).
407    /// Returns None if any column is missing or has wrong type.
408    /// - intensity: F32
409    /// - ring: U16
410    #[must_use]
411    pub fn xyzir(&self) -> Option<(&[f32], &[f32], &[f32], &[f32], &[u16])> {
412        let x = self.get_column("x")?.as_f32()?;
413        let y = self.get_column("y")?.as_f32()?;
414        let z = self.get_column("z")?.as_f32()?;
415        let intensity = self.get_column("intensity")?.as_f32()?;
416        let ring = self.get_column("ring")?.as_u16()?;
417        Some((x, y, z, intensity, ring))
418    }
419
420    /// Get XYZ + intensity + ring + timestamp (full LiDAR format).
421    /// Returns None if any column is missing or has wrong type.
422    /// - intensity: F32
423    /// - ring: U16
424    /// - timestamp: F64
425    #[must_use]
426    pub fn xyzirt(&self) -> Option<(&[f32], &[f32], &[f32], &[f32], &[u16], &[f64])> {
427        let x = self.get_column("x")?.as_f32()?;
428        let y = self.get_column("y")?.as_f32()?;
429        let z = self.get_column("z")?.as_f32()?;
430        let intensity = self.get_column("intensity")?.as_f32()?;
431        let ring = self.get_column("ring")?.as_u16()?;
432        let timestamp = self.get_column("timestamp")?.as_f64()?;
433        Some((x, y, z, intensity, ring, timestamp))
434    }
435
436    /// Get XYZIRT + id (LiDAR format with point ID/label).
437    /// Returns None if any column is missing or has wrong type.
438    /// - intensity: F32
439    /// - ring: U16
440    /// - timestamp: F64
441    /// - id: U32
442    #[must_use]
443    pub fn xyzirt_id(&self) -> Option<(&[f32], &[f32], &[f32], &[f32], &[u16], &[f64], &[u32])> {
444        let x = self.get_column("x")?.as_f32()?;
445        let y = self.get_column("y")?.as_f32()?;
446        let z = self.get_column("z")?.as_f32()?;
447        let intensity = self.get_column("intensity")?.as_f32()?;
448        let ring = self.get_column("ring")?.as_u16()?;
449        let timestamp = self.get_column("timestamp")?.as_f64()?;
450        let id = self.get_column("id")?.as_u32()?;
451        Some((x, y, z, intensity, ring, timestamp, id))
452    }
453}