1use crate::compression_config::EncodedGeometryType;
2use crate::corner_table::CornerTable;
3#[cfg(feature = "point_cloud_decode")]
4use crate::decoder_buffer::DecoderBuffer;
5#[cfg(feature = "point_cloud_decode")]
6use crate::draco_types::DataType;
7#[cfg(feature = "point_cloud_decode")]
8use crate::geometry_attribute::{GeometryAttributeType, PointAttribute};
9#[cfg(feature = "point_cloud_decode")]
10use crate::geometry_indices::PointIndex;
11#[cfg(feature = "point_cloud_decode")]
12use crate::kd_tree_attributes_decoder::KdTreeAttributesDecoder;
13use crate::mesh::Mesh;
14use crate::point_cloud::PointCloud;
15#[cfg(feature = "point_cloud_decode")]
16use crate::sequential_integer_attribute_decoder::SequentialIntegerAttributeDecoder;
17#[cfg(feature = "point_cloud_decode")]
18use crate::status::{DracoError, Status};
19
20#[cfg(feature = "point_cloud_decode")]
21use crate::attribute_octahedron_transform::AttributeOctahedronTransform;
22#[cfg(feature = "point_cloud_decode")]
23use crate::attribute_quantization_transform::AttributeQuantizationTransform;
24#[cfg(feature = "point_cloud_decode")]
25use crate::attribute_transform::AttributeTransform;
26#[cfg(feature = "point_cloud_decode")]
27use crate::version::{version_at_least, VERSION_FLAGS_INTRODUCED};
28
29pub trait GeometryDecoder {
31 fn point_cloud(&self) -> Option<&PointCloud>;
33 fn mesh(&self) -> Option<&Mesh>;
35 fn corner_table(&self) -> Option<&CornerTable>;
37 fn get_geometry_type(&self) -> EncodedGeometryType;
39 fn get_attribute_encoding_method(&self, _att_id: i32) -> Option<i32> {
41 None
42 }
43}
44
45pub struct PointCloudDecoder {
54 geometry_type: EncodedGeometryType,
55 #[cfg(feature = "point_cloud_decode")]
56 method: u8,
57 #[cfg(feature = "point_cloud_decode")]
58 flags: u16,
59 #[cfg(feature = "point_cloud_decode")]
60 version_major: u8,
61 #[cfg(feature = "point_cloud_decode")]
62 version_minor: u8,
63}
64
65impl GeometryDecoder for PointCloudDecoder {
66 fn point_cloud(&self) -> Option<&PointCloud> {
67 None }
73
74 fn mesh(&self) -> Option<&Mesh> {
75 None
76 }
77
78 fn corner_table(&self) -> Option<&CornerTable> {
79 None
80 }
81
82 fn get_geometry_type(&self) -> EncodedGeometryType {
83 self.geometry_type
84 }
85}
86
87impl Default for PointCloudDecoder {
88 fn default() -> Self {
89 Self::new()
90 }
91}
92
93#[cfg(feature = "point_cloud_decode")]
94fn make_point_ids(num_points: usize) -> Result<Vec<PointIndex>, DracoError> {
95 let mut point_ids = Vec::new();
96 point_ids
97 .try_reserve_exact(num_points)
98 .map_err(|_| DracoError::DracoError("Failed to allocate point ids".to_string()))?;
99 for i in 0..num_points {
100 point_ids.push(PointIndex(i as u32));
101 }
102 Ok(point_ids)
103}
104
105#[cfg(feature = "point_cloud_decode")]
106fn validate_num_attributes_in_decoder(
107 num_attributes_in_decoder: usize,
108 remaining_bytes: usize,
109) -> Result<(), DracoError> {
110 const MIN_ATTRIBUTE_BYTES: usize = 6;
114 if num_attributes_in_decoder == 0
115 || num_attributes_in_decoder > remaining_bytes / MIN_ATTRIBUTE_BYTES
116 {
117 return Err(DracoError::DracoError(
118 "Invalid number of attributes".to_string(),
119 ));
120 }
121 Ok(())
122}
123
124#[cfg(feature = "point_cloud_decode")]
125fn validate_num_components(num_components: u8) -> Result<(), DracoError> {
126 if num_components == 0 {
127 return Err(DracoError::DracoError(
128 "Invalid attribute component count".to_string(),
129 ));
130 }
131 Ok(())
132}
133
134#[cfg(feature = "point_cloud_decode")]
135fn decode_raw_attribute_values(
136 buffer: &mut DecoderBuffer<'_>,
137 attribute: &mut PointAttribute,
138 num_points: usize,
139) -> Result<(), DracoError> {
140 let entry_size = attribute.byte_stride() as usize;
141 if entry_size == 0 {
142 return Err(DracoError::DracoError(
143 "Invalid point cloud attribute entry size".to_string(),
144 ));
145 }
146 let required_size = entry_size.checked_mul(num_points).ok_or_else(|| {
147 DracoError::DracoError("Point cloud raw attribute byte count overflow".to_string())
148 })?;
149
150 let dst = attribute.buffer_mut().data_mut();
151 if dst.len() < required_size {
152 return Err(DracoError::DracoError(
153 "Point cloud attribute buffer too small".to_string(),
154 ));
155 }
156
157 for chunk in dst[..required_size].chunks_exact_mut(entry_size) {
158 buffer.decode_bytes(chunk).map_err(|_| {
159 DracoError::DracoError("Failed to decode raw point cloud attribute values".to_string())
160 })?;
161 }
162
163 Ok(())
164}
165
166impl PointCloudDecoder {
167 pub fn new() -> Self {
169 Self {
170 geometry_type: EncodedGeometryType::PointCloud,
171 #[cfg(feature = "point_cloud_decode")]
172 method: 0,
173 #[cfg(feature = "point_cloud_decode")]
174 flags: 0,
175 #[cfg(feature = "point_cloud_decode")]
176 version_major: 0,
177 #[cfg(feature = "point_cloud_decode")]
178 version_minor: 0,
179 }
180 }
181
182 #[cfg(feature = "point_cloud_decode")]
183 pub fn decode(&mut self, in_buffer: &mut DecoderBuffer, out_pc: &mut PointCloud) -> Status {
190 self.decode_header(in_buffer)?;
192
193 if version_at_least(
194 self.version_major,
195 self.version_minor,
196 VERSION_FLAGS_INTRODUCED,
197 ) && (self.flags & crate::metadata::METADATA_FLAG_MASK) != 0
198 {
199 let metadata = crate::metadata::GeometryMetadata::decode(in_buffer)
200 .map_err(|_| DracoError::DracoError("Failed to decode metadata".to_string()))?;
201 out_pc.set_metadata(Some(metadata));
202 }
203
204 self.decode_geometry_data(in_buffer, out_pc)
206 }
207
208 #[cfg(feature = "point_cloud_decode")]
211 pub fn decode_after_header(
212 &mut self,
213 version_major: u8,
214 version_minor: u8,
215 method: u8,
216 buffer: &mut DecoderBuffer,
217 out_pc: &mut PointCloud,
218 ) -> Status {
219 self.version_major = version_major;
220 self.version_minor = version_minor;
221 self.method = method;
222 self.flags = 0;
223 self.geometry_type = EncodedGeometryType::PointCloud;
224 self.decode_geometry_data(buffer, out_pc)
225 }
226
227 #[cfg(feature = "point_cloud_decode")]
228 fn decode_header(&mut self, buffer: &mut DecoderBuffer) -> Status {
229 let mut magic = [0u8; 5];
230 buffer.decode_bytes(&mut magic)?;
231 if &magic != b"DRACO" {
232 return Err(DracoError::DracoError("Invalid magic".to_string()));
233 }
234
235 self.version_major = buffer.decode_u8()?;
236 self.version_minor = buffer.decode_u8()?;
237 buffer.set_version(self.version_major, self.version_minor);
238
239 let g_type = buffer.decode_u8()?;
240 self.geometry_type = match g_type {
241 0 => EncodedGeometryType::PointCloud,
242 1 => EncodedGeometryType::TriangularMesh,
243 _ => return Err(DracoError::DracoError("Invalid geometry type".to_string())),
244 };
245
246 self.method = buffer.decode_u8()?;
247
248 self.flags = buffer
250 .decode_u16()
251 .map_err(|_| DracoError::DracoError("Failed to decode flags".to_string()))?;
252
253 Ok(())
254 }
255
256 #[cfg(feature = "point_cloud_decode")]
257 fn decode_geometry_data(&mut self, buffer: &mut DecoderBuffer, pc: &mut PointCloud) -> Status {
258 let bitstream_version: u16 =
259 crate::version::bitstream_version(self.version_major, self.version_minor);
260 let num_points: usize = buffer.decode_u32()? as usize;
265 if num_points > buffer.remaining_size().saturating_mul(8) {
274 return Err(DracoError::DracoError(
275 "Point count exceeds remaining bitstream size".to_string(),
276 ));
277 }
278 pc.set_num_points(num_points);
279
280 let num_attributes_decoders = buffer.decode_u8()? as usize;
281
282 if self.method == 1 {
283 for _ in 0..num_attributes_decoders {
285 let mut att_decoder = KdTreeAttributesDecoder::new(0);
286 if !att_decoder.decode_attributes_decoder_data(pc, buffer) {
287 return Err(DracoError::DracoError(
288 "Failed to decode attribute metadata".to_string(),
289 ));
290 }
291 if !att_decoder.decode_attributes(pc, buffer) {
292 return Err(DracoError::DracoError(
293 "Failed to decode attributes".to_string(),
294 ));
295 }
296 }
297 } else {
298 struct PendingQuant {
300 att_id: i32,
301 portable: PointAttribute,
302 transform: AttributeQuantizationTransform,
303 }
304
305 struct PendingNormal {
306 att_id: i32,
307 portable: PointAttribute,
308 quantization_bits: u8,
309 }
310
311 struct AttributeSpec {
312 att_type: GeometryAttributeType,
313 data_type: DataType,
314 num_components: u8,
315 normalized: bool,
316 unique_id: u32,
317 }
318
319 for _ in 0..num_attributes_decoders {
320 let num_attributes_in_decoder: usize = if bitstream_version < 0x0200 {
321 buffer.decode_u32()? as usize
322 } else {
323 buffer.decode_varint()? as usize
324 };
325 if num_attributes_in_decoder == 0 {
326 return Err(DracoError::DracoError(
327 "Invalid number of attributes".to_string(),
328 ));
329 }
330 validate_num_attributes_in_decoder(
331 num_attributes_in_decoder,
332 buffer.remaining_size(),
333 )?;
334
335 let mut attribute_specs: Vec<AttributeSpec> =
336 Vec::with_capacity(num_attributes_in_decoder);
337 let mut att_ids: Vec<i32> = Vec::with_capacity(num_attributes_in_decoder);
338 let mut decoder_types: Vec<u8> = Vec::with_capacity(num_attributes_in_decoder);
339 let mut pending_quant: Vec<PendingQuant> = Vec::new();
340 let mut pending_normals: Vec<PendingNormal> = Vec::new();
341
342 for _ in 0..num_attributes_in_decoder {
343 let att_type_val = buffer.decode_u8()?;
344 let att_type = GeometryAttributeType::try_from(att_type_val)?;
345
346 let data_type_val = buffer.decode_u8()?;
347 let data_type = DataType::try_from(data_type_val)?;
348
349 let num_components = buffer.decode_u8()?;
350 validate_num_components(num_components)?;
351 let normalized = buffer.decode_u8()? != 0;
352 let unique_id: u32 = if bitstream_version < 0x0103 {
353 buffer.decode_u16()? as u32
354 } else {
355 buffer.decode_varint()? as u32
356 };
357
358 attribute_specs.push(AttributeSpec {
359 att_type,
360 data_type,
361 num_components,
362 normalized,
363 unique_id,
364 });
365 }
366
367 for _ in 0..num_attributes_in_decoder {
368 decoder_types.push(buffer.decode_u8()?);
369 }
370
371 for (local_i, spec) in attribute_specs.iter().enumerate() {
372 if decoder_types[local_i] == 0 {
373 let entry_size =
374 spec.num_components as usize * spec.data_type.byte_length();
375 let bytes_needed = entry_size.checked_mul(num_points).ok_or_else(|| {
376 DracoError::DracoError(
377 "Raw point cloud attribute byte count overflow".to_string(),
378 )
379 })?;
380 if buffer.remaining_size() < bytes_needed {
381 return Err(DracoError::DracoError(
382 "Not enough data for raw point cloud attribute values".to_string(),
383 ));
384 }
385 }
386
387 let mut att = PointAttribute::new();
388 att.try_init(
389 spec.att_type,
390 spec.num_components,
391 spec.data_type,
392 spec.normalized,
393 num_points,
394 )?;
395 att.set_unique_id(spec.unique_id);
396 let att_id = pc.add_attribute_preserve_unique_id(att);
397 att_ids.push(att_id);
398 }
399
400 let point_ids = if decoder_types.iter().any(|&decoder_type| decoder_type != 0) {
401 Some(make_point_ids(num_points)?)
402 } else {
403 None
404 };
405
406 for (local_i, &att_id) in att_ids.iter().enumerate() {
407 let decoder_type = decoder_types[local_i];
408 match decoder_type {
409 1 => {
410 let point_ids = point_ids.as_ref().ok_or_else(|| {
411 DracoError::DracoError(
412 "Point ids missing for integer attribute decoder".to_string(),
413 )
414 })?;
415 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
416 att_decoder.init(self, att_id);
417 if !att_decoder.decode_values(
418 pc, point_ids, buffer, None, None, None, None, None, None,
419 ) {
420 return Err(DracoError::DracoError(
421 "Failed to decode integer attribute".to_string(),
422 ));
423 }
424 }
425 2 => {
426 let original = pc.try_attribute(att_id)?;
427 let (original_type, original_num_components) =
428 (original.attribute_type(), original.num_components());
429 let mut portable = PointAttribute::default();
430 portable.try_init(
431 original_type,
432 original_num_components,
433 DataType::Uint32,
434 false,
435 num_points,
436 )?;
437 let mut transform = AttributeQuantizationTransform::new();
438
439 let quant_skip_bytes = if bitstream_version < 0x0102 {
443 let saved_pos = buffer.position();
444 let method_byte = buffer.decode_u8().map_err(|_| {
445 DracoError::DracoError("read pred method".to_string())
446 })?;
447 if method_byte != 0xFF {
448 let _transform_byte = buffer.decode_u8().map_err(|_| {
449 DracoError::DracoError("read transform".to_string())
450 })?;
451 }
452 let original = pc.try_attribute(att_id)?;
453 if !transform.decode_parameters(original, buffer) {
454 return Err(DracoError::DracoError(
455 "Failed to decode quantization parameters (v<2.0)"
456 .to_string(),
457 ));
458 }
459 let bytes_consumed = buffer.position() - saved_pos;
460 let pred_header_bytes = if method_byte != 0xFF { 2 } else { 1 };
461 let skip = bytes_consumed - pred_header_bytes;
462 buffer
463 .set_position(saved_pos)
464 .map_err(|_| DracoError::DracoError("buf reset".to_string()))?;
465 skip
466 } else {
467 0
468 };
469 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
470 att_decoder.init(self, att_id);
471 let mut skip_fn =
472 move |buf: &mut crate::decoder_buffer::DecoderBuffer<'_>| -> bool {
473 if quant_skip_bytes > 0
474 && buf.try_advance(quant_skip_bytes).is_err()
475 {
476 return false;
477 }
478 true
479 };
480 let hook: Option<
481 &mut dyn FnMut(
482 &mut crate::decoder_buffer::DecoderBuffer<'_>,
483 ) -> bool,
484 > = if quant_skip_bytes > 0 {
485 Some(&mut skip_fn)
486 } else {
487 None
488 };
489 if !att_decoder.decode_values(
490 pc,
491 point_ids.as_ref().ok_or_else(|| {
492 DracoError::DracoError(
493 "Point ids missing for quantized attribute decoder"
494 .to_string(),
495 )
496 })?,
497 buffer,
498 None,
499 None,
500 None,
501 Some(&mut portable),
502 None,
503 hook,
504 ) {
505 return Err(DracoError::DracoError(
506 "Failed to decode quantized portable values".to_string(),
507 ));
508 }
509 pending_quant.push(PendingQuant {
510 att_id,
511 portable,
512 transform,
513 });
514 }
515 3 => {
516 let mut portable = PointAttribute::default();
517 portable.try_init(
518 GeometryAttributeType::Generic,
519 2,
520 DataType::Uint32,
521 false,
522 num_points,
523 )?;
524 let mut quant_bits: u8 = 0;
528 let normal_skip_bytes = if bitstream_version < 0x0102 {
529 let saved_pos = buffer.position();
530 let method_byte = buffer.decode_u8().map_err(|_| {
531 DracoError::DracoError("read pred method".to_string())
532 })?;
533 if method_byte != 0xFF {
534 let _transform_byte = buffer.decode_u8().map_err(|_| {
535 DracoError::DracoError("read transform".to_string())
536 })?;
537 }
538 quant_bits = buffer.decode_u8().map_err(|_| {
539 DracoError::DracoError("read normal quant_bits".to_string())
540 })?;
541 if !AttributeOctahedronTransform::is_valid_quantization_bits(
542 quant_bits as i32,
543 ) {
544 return Err(DracoError::DracoError(
545 "Invalid normal quantization bits".to_string(),
546 ));
547 }
548 let bytes_consumed = buffer.position() - saved_pos;
549 let pred_header_bytes = if method_byte != 0xFF { 2 } else { 1 };
550 let skip = bytes_consumed - pred_header_bytes;
551 buffer
552 .set_position(saved_pos)
553 .map_err(|_| DracoError::DracoError("buf reset".to_string()))?;
554 skip
555 } else {
556 0
557 };
558 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
559 att_decoder.init(self, att_id);
560 let mut skip_fn =
561 move |buf: &mut crate::decoder_buffer::DecoderBuffer<'_>| -> bool {
562 if normal_skip_bytes > 0
563 && buf.try_advance(normal_skip_bytes).is_err()
564 {
565 return false;
566 }
567 true
568 };
569 let hook: Option<
570 &mut dyn FnMut(
571 &mut crate::decoder_buffer::DecoderBuffer<'_>,
572 ) -> bool,
573 > = if normal_skip_bytes > 0 {
574 Some(&mut skip_fn)
575 } else {
576 None
577 };
578 if !att_decoder.decode_values(
579 pc,
580 point_ids.as_ref().ok_or_else(|| {
581 DracoError::DracoError(
582 "Point ids missing for normal attribute decoder"
583 .to_string(),
584 )
585 })?,
586 buffer,
587 None,
588 None,
589 None,
590 Some(&mut portable),
591 None,
592 hook,
593 ) {
594 return Err(DracoError::DracoError(
595 "Failed to decode normal portable values".to_string(),
596 ));
597 }
598 pending_normals.push(PendingNormal {
599 att_id,
600 portable,
601 quantization_bits: quant_bits,
602 });
603 }
604 0 => {
605 decode_raw_attribute_values(
608 buffer,
609 pc.try_attribute_mut(att_id)?,
610 num_points,
611 )?;
612 }
613 _ => {
614 return Err(DracoError::DracoError(format!(
615 "Unsupported sequential decoder type: {}",
616 decoder_type
617 )));
618 }
619 }
620 }
621
622 for (local_i, &att_id) in att_ids.iter().enumerate() {
623 match decoder_types[local_i] {
624 2 if bitstream_version >= 0x0102 => {
625 let idx = pending_quant
626 .iter()
627 .position(|p| p.att_id == att_id)
628 .ok_or_else(|| {
629 DracoError::DracoError(
630 "Missing pending quantized attribute transform".to_string(),
631 )
632 })?;
633 let original = pc.try_attribute(att_id)?;
634 if !pending_quant[idx]
635 .transform
636 .decode_parameters(original, buffer)
637 {
638 return Err(DracoError::DracoError(
639 "Failed to decode quantization parameters".to_string(),
640 ));
641 }
642 }
643 3 if bitstream_version >= 0x0102 => {
644 let idx = pending_normals
645 .iter()
646 .position(|p| p.att_id == att_id)
647 .ok_or_else(|| {
648 DracoError::DracoError(
649 "Missing pending normal attribute transform".to_string(),
650 )
651 })?;
652 let quantization_bits = buffer.decode_u8()?;
653 if !AttributeOctahedronTransform::is_valid_quantization_bits(
654 quantization_bits as i32,
655 ) {
656 return Err(DracoError::DracoError(
657 "Invalid normal quantization bits".to_string(),
658 ));
659 }
660 pending_normals[idx].quantization_bits = quantization_bits;
661 }
662 _ => {}
663 }
664 }
665
666 for q in pending_quant {
667 let dst = pc.try_attribute_mut(q.att_id)?;
668 if !q.transform.inverse_transform_attribute(&q.portable, dst) {
669 return Err(DracoError::DracoError(
670 "Failed to dequantize attribute".to_string(),
671 ));
672 }
673 }
674 for n in pending_normals {
675 let mut oct = AttributeOctahedronTransform::new(-1);
676 if !oct.set_parameters(n.quantization_bits as i32) {
677 return Err(DracoError::DracoError(
678 "Invalid normal quantization bits".to_string(),
679 ));
680 }
681 let dst = pc.try_attribute_mut(n.att_id)?;
682 if !oct.inverse_transform_attribute_with_legacy_octahedron(
683 &n.portable,
684 dst,
685 bitstream_version < 0x0102,
686 ) {
687 return Err(DracoError::DracoError(
688 "Failed to decode normals".to_string(),
689 ));
690 }
691 }
692 }
693 }
694
695 Ok(())
696 }
697
698 pub fn get_geometry_type(&self) -> EncodedGeometryType {
700 self.geometry_type
701 }
702}
703
704#[cfg(all(test, feature = "point_cloud_decode"))]
705mod tests {
706 use super::*;
707
708 #[test]
709 fn decode_raw_attribute_values_rejects_required_size_overflow() {
710 let bytes = [];
711 let mut buffer = DecoderBuffer::new(&bytes);
712 let mut attribute = PointAttribute::new();
713 attribute.init(
714 GeometryAttributeType::Generic,
715 1,
716 DataType::Uint32,
717 false,
718 1,
719 );
720
721 let status = decode_raw_attribute_values(&mut buffer, &mut attribute, usize::MAX);
722
723 assert!(status.is_err());
724 }
725
726 #[test]
727 fn decode_raw_attribute_values_rejects_truncated_input() {
728 let bytes = [1u8, 2, 3];
729 let mut buffer = DecoderBuffer::new(&bytes);
730 let mut attribute = PointAttribute::new();
731 attribute.init(
732 GeometryAttributeType::Generic,
733 1,
734 DataType::Uint32,
735 false,
736 1,
737 );
738
739 let status = decode_raw_attribute_values(&mut buffer, &mut attribute, 1);
740
741 assert!(status.is_err());
742 }
743}