1use crate::error::Result;
16use crate::header::DataFormat;
17use crate::header::PcdHeader;
18use 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 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 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 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 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 })?; 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 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 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}