rs_pcd/io/
writer.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::error::Result;
16use crate::header::DataFormat;
17use crate::header::PcdHeader;
18// use crate::header::ValueType;
19use crate::error::PcdError;
20use crate::storage::PointBlock;
21use byteorder::{LittleEndian, WriteBytesExt};
22use lzf;
23use std::io::Write;
24
25pub struct PcdWriter<W: Write> {
26    writer: W,
27}
28
29impl<W: Write> PcdWriter<W> {
30    pub fn new(writer: W) -> Self {
31        Self { writer }
32    }
33
34    pub fn write_pcd(&mut self, header: &PcdHeader, data: &PointBlock) -> Result<()> {
35        self.write_header(header)?;
36        match header.data {
37            DataFormat::Binary => self.write_binary(header, data)?,
38            DataFormat::Ascii => self.write_ascii(header, data)?,
39            DataFormat::BinaryCompressed => self.write_compressed_binary(header, data)?,
40        }
41        Ok(())
42    }
43
44    fn write_header(&mut self, header: &PcdHeader) -> Result<()> {
45        writeln!(self.writer, "VERSION {}", header.version)?;
46        writeln!(self.writer, "FIELDS {}", header.fields.join(" "))?;
47
48        let sizes_str: Vec<String> = header.sizes.iter().map(|s| s.to_string()).collect();
49        writeln!(self.writer, "SIZE {}", sizes_str.join(" "))?;
50
51        let types_str: Vec<String> = header.types.iter().map(|t| t.to_string()).collect();
52        writeln!(self.writer, "TYPE {}", types_str.join(" "))?;
53
54        let counts_str: Vec<String> = header.counts.iter().map(|c| c.to_string()).collect();
55        writeln!(self.writer, "COUNT {}", counts_str.join(" "))?;
56
57        writeln!(self.writer, "WIDTH {}", header.width)?;
58        writeln!(self.writer, "HEIGHT {}", header.height)?;
59
60        // VIEWPOINT
61        writeln!(
62            self.writer,
63            "VIEWPOINT {} {} {} {} {} {} {}",
64            header.viewpoint[0],
65            header.viewpoint[1],
66            header.viewpoint[2],
67            header.viewpoint[3],
68            header.viewpoint[4],
69            header.viewpoint[5],
70            header.viewpoint[6]
71        )?;
72
73        writeln!(self.writer, "POINTS {}", header.points)?;
74
75        match header.data {
76            DataFormat::Ascii => writeln!(self.writer, "DATA ascii")?,
77            DataFormat::Binary => writeln!(self.writer, "DATA binary")?,
78            DataFormat::BinaryCompressed => writeln!(self.writer, "DATA binary_compressed")?,
79        }
80
81        Ok(())
82    }
83
84    fn write_binary(&mut self, header: &PcdHeader, data: &PointBlock) -> Result<()> {
85        // Optimization: Collect column references once
86        let mut columns = Vec::with_capacity(header.fields.len());
87        for name in &header.fields {
88            columns.push(
89                data.get_column(name).ok_or_else(|| {
90                    PcdError::InvalidDataFormat(format!("Missing column {}", name))
91                })?,
92            );
93        }
94
95        // Loop points, then fields (AoS)
96        for i in 0..header.points {
97            for (field_idx, _name) in header.fields.iter().enumerate() {
98                let col = columns[field_idx];
99                let count = header.counts[field_idx];
100                let start = i * count;
101
102                match header.types[field_idx] {
103                    'F' => {
104                        // Check sizes: 4 bytes -> F32, 8 bytes -> F64
105                        match header.sizes[field_idx] {
106                            4 => {
107                                let vec = col.as_f32().ok_or_else(|| PcdError::LayoutMismatch {
108                                    expected: 0,
109                                    got: 0,
110                                })?; // Todo better error
111                                for k in 0..count {
112                                    self.writer.write_f32::<LittleEndian>(vec[start + k])?;
113                                }
114                            }
115                            8 => {
116                                let vec = col.as_f64().ok_or_else(|| PcdError::LayoutMismatch {
117                                    expected: 0,
118                                    got: 0,
119                                })?;
120                                for k in 0..count {
121                                    self.writer.write_f64::<LittleEndian>(vec[start + k])?;
122                                }
123                            }
124                            _ => {
125                                return Err(PcdError::UnsupportedType(format!(
126                                    "F{}",
127                                    header.sizes[field_idx]
128                                )));
129                            }
130                        }
131                    }
132                    'U' => match header.sizes[field_idx] {
133                        1 => {
134                            let vec = col.as_u8().ok_or(PcdError::LayoutMismatch {
135                                expected: 0,
136                                got: 0,
137                            })?;
138                            for k in 0..count {
139                                self.writer.write_u8(vec[start + k])?;
140                            }
141                        }
142                        2 => {
143                            let vec = col.as_u16().ok_or(PcdError::LayoutMismatch {
144                                expected: 0,
145                                got: 0,
146                            })?;
147                            for k in 0..count {
148                                self.writer.write_u16::<LittleEndian>(vec[start + k])?;
149                            }
150                        }
151                        4 => {
152                            let vec = col.as_u32().ok_or(PcdError::LayoutMismatch {
153                                expected: 0,
154                                got: 0,
155                            })?;
156                            for k in 0..count {
157                                self.writer.write_u32::<LittleEndian>(vec[start + k])?;
158                            }
159                        }
160                        _ => {
161                            return Err(PcdError::UnsupportedType(format!(
162                                "U{}",
163                                header.sizes[field_idx]
164                            )));
165                        }
166                    },
167                    'I' => match header.sizes[field_idx] {
168                        1 => {
169                            let vec = col.as_i8().ok_or(PcdError::LayoutMismatch {
170                                expected: 0,
171                                got: 0,
172                            })?;
173                            for k in 0..count {
174                                self.writer.write_i8(vec[start + k])?;
175                            }
176                        }
177                        2 => {
178                            let vec = col.as_i16().ok_or(PcdError::LayoutMismatch {
179                                expected: 0,
180                                got: 0,
181                            })?;
182                            for k in 0..count {
183                                self.writer.write_i16::<LittleEndian>(vec[start + k])?;
184                            }
185                        }
186                        4 => {
187                            let vec = col.as_i32().ok_or(PcdError::LayoutMismatch {
188                                expected: 0,
189                                got: 0,
190                            })?;
191                            for k in 0..count {
192                                self.writer.write_i32::<LittleEndian>(vec[start + k])?;
193                            }
194                        }
195                        _ => {
196                            return Err(PcdError::UnsupportedType(format!(
197                                "I{}",
198                                header.sizes[field_idx]
199                            )));
200                        }
201                    },
202                    _ => {
203                        return Err(PcdError::UnsupportedType(
204                            header.types[field_idx].to_string(),
205                        ));
206                    }
207                }
208            }
209        }
210        Ok(())
211    }
212
213    fn write_ascii(&mut self, header: &PcdHeader, data: &PointBlock) -> Result<()> {
214        // Optimization: Collect column references once
215        let mut columns = Vec::with_capacity(header.fields.len());
216        for name in &header.fields {
217            columns.push(
218                data.get_column(name).ok_or_else(|| {
219                    PcdError::InvalidDataFormat(format!("Missing column {}", name))
220                })?,
221            );
222        }
223
224        for i in 0..header.points {
225            let mut line_tokens = Vec::with_capacity(header.fields.len());
226            for (field_idx, _name) in header.fields.iter().enumerate() {
227                let col = columns[field_idx];
228                let count = header.counts[field_idx];
229                let start = i * count;
230
231                match header.types[field_idx] {
232                    'F' => match header.sizes[field_idx] {
233                        4 => {
234                            let vec = col.as_f32().ok_or(PcdError::LayoutMismatch {
235                                expected: 0,
236                                got: 0,
237                            })?;
238                            for k in 0..count {
239                                line_tokens.push(format!("{:.6}", vec[start + k]));
240                            }
241                        }
242                        8 => {
243                            let vec = col.as_f64().ok_or(PcdError::LayoutMismatch {
244                                expected: 0,
245                                got: 0,
246                            })?;
247                            for k in 0..count {
248                                line_tokens.push(format!("{:.6}", vec[start + k]));
249                            }
250                        }
251                        _ => {}
252                    },
253                    'U' => match header.sizes[field_idx] {
254                        1 => {
255                            let vec = col.as_u8().ok_or(PcdError::LayoutMismatch {
256                                expected: 0,
257                                got: 0,
258                            })?;
259                            for k in 0..count {
260                                line_tokens.push(format!("{}", vec[start + k]));
261                            }
262                        }
263                        2 => {
264                            let vec = col.as_u16().ok_or(PcdError::LayoutMismatch {
265                                expected: 0,
266                                got: 0,
267                            })?;
268                            for k in 0..count {
269                                line_tokens.push(format!("{}", vec[start + k]));
270                            }
271                        }
272                        4 => {
273                            let vec = col.as_u32().ok_or(PcdError::LayoutMismatch {
274                                expected: 0,
275                                got: 0,
276                            })?;
277                            for k in 0..count {
278                                line_tokens.push(format!("{}", vec[start + k]));
279                            }
280                        }
281                        _ => {}
282                    },
283                    'I' => match header.sizes[field_idx] {
284                        1 => {
285                            let vec = col.as_i8().ok_or(PcdError::LayoutMismatch {
286                                expected: 0,
287                                got: 0,
288                            })?;
289                            for k in 0..count {
290                                line_tokens.push(format!("{}", vec[start + k]));
291                            }
292                        }
293                        2 => {
294                            let vec = col.as_i16().ok_or(PcdError::LayoutMismatch {
295                                expected: 0,
296                                got: 0,
297                            })?;
298                            for k in 0..count {
299                                line_tokens.push(format!("{}", vec[start + k]));
300                            }
301                        }
302                        4 => {
303                            let vec = col.as_i32().ok_or(PcdError::LayoutMismatch {
304                                expected: 0,
305                                got: 0,
306                            })?;
307                            for k in 0..count {
308                                line_tokens.push(format!("{}", vec[start + k]));
309                            }
310                        }
311                        _ => {}
312                    },
313                    _ => {}
314                }
315            }
316            writeln!(self.writer, "{}", line_tokens.join(" "))?;
317        }
318        Ok(())
319    }
320    fn write_compressed_binary(&mut self, header: &PcdHeader, data: &PointBlock) -> Result<()> {
321        let mut uncompressed_data = Vec::new();
322
323        // Binary Compressed is SoA in the buffer
324        for (field_idx, name) in header.fields.iter().enumerate() {
325            let col = data
326                .get_column(name)
327                .ok_or_else(|| PcdError::InvalidDataFormat(format!("Missing column {}", name)))?;
328            let _count = header.counts[field_idx];
329
330            match header.types[field_idx] {
331                'F' => {
332                    if header.sizes[field_idx] == 4 {
333                        let vec = col.as_f32().unwrap();
334                        for val in vec {
335                            uncompressed_data.write_f32::<LittleEndian>(*val)?;
336                        }
337                    } else {
338                        let vec = col.as_f64().unwrap();
339                        for val in vec {
340                            uncompressed_data.write_f64::<LittleEndian>(*val)?;
341                        }
342                    }
343                }
344                'U' => match header.sizes[field_idx] {
345                    1 => uncompressed_data.write_all(col.as_u8().unwrap())?,
346                    2 => {
347                        let vec = col.as_u16().unwrap();
348                        for val in vec {
349                            uncompressed_data.write_u16::<LittleEndian>(*val)?;
350                        }
351                    }
352                    4 => {
353                        let vec = col.as_u32().unwrap();
354                        for val in vec {
355                            uncompressed_data.write_u32::<LittleEndian>(*val)?;
356                        }
357                    }
358                    _ => {}
359                },
360                'I' => match header.sizes[field_idx] {
361                    1 => {
362                        let vec = col.as_i8().unwrap();
363                        for val in vec {
364                            uncompressed_data.write_i8(*val)?;
365                        }
366                    }
367                    2 => {
368                        let vec = col.as_i16().unwrap();
369                        for val in vec {
370                            uncompressed_data.write_i16::<LittleEndian>(*val)?;
371                        }
372                    }
373                    4 => {
374                        let vec = col.as_i32().unwrap();
375                        for val in vec {
376                            uncompressed_data.write_i32::<LittleEndian>(*val)?;
377                        }
378                    }
379                    _ => {}
380                },
381                _ => {}
382            }
383        }
384
385        let uncompressed_size = uncompressed_data.len();
386        let compressed_result = lzf::compress(&uncompressed_data);
387
388        let (final_compressed_size, final_data) = match compressed_result {
389            Ok(data) => (data.len(), data),
390            Err(lzf::LzfError::NoCompressionPossible) => (uncompressed_size, uncompressed_data),
391            Err(e) => return Err(PcdError::Other(format!("Compression failed: {:?}", e))),
392        };
393
394        self.writer
395            .write_u32::<LittleEndian>(final_compressed_size as u32)?;
396        self.writer
397            .write_u32::<LittleEndian>(uncompressed_size as u32)?;
398        self.writer.write_all(&final_data)?;
399
400        Ok(())
401    }
402}