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 if self.geometry_type != EncodedGeometryType::PointCloud {
246 return Err(DracoError::DracoError(
247 "PointCloudDecoder cannot decode mesh bitstreams".to_string(),
248 ));
249 }
250
251 self.method = buffer.decode_u8()?;
252
253 self.flags = buffer
255 .decode_u16()
256 .map_err(|_| DracoError::DracoError("Failed to decode flags".to_string()))?;
257
258 Ok(())
259 }
260
261 #[cfg(feature = "point_cloud_decode")]
262 fn decode_geometry_data(&mut self, buffer: &mut DecoderBuffer, pc: &mut PointCloud) -> Status {
263 let bitstream_version: u16 =
264 crate::version::bitstream_version(self.version_major, self.version_minor);
265 let num_points: usize = buffer.decode_u32()? as usize;
270 if num_points > buffer.remaining_size().saturating_mul(8) {
279 return Err(DracoError::count_exceeds_bitstream(
280 num_points,
281 buffer.remaining_size(),
282 ));
283 }
284 pc.set_num_points(num_points);
285
286 let num_attributes_decoders = buffer.decode_u8()? as usize;
287
288 if self.method == 1 {
289 for _ in 0..num_attributes_decoders {
291 let mut att_decoder = KdTreeAttributesDecoder::new(0);
292 if !att_decoder.decode_attributes_decoder_data(pc, buffer) {
293 return Err(DracoError::DracoError(
294 "Failed to decode attribute metadata".to_string(),
295 ));
296 }
297 if !att_decoder.decode_attributes(pc, buffer) {
298 return Err(DracoError::DracoError(
299 "Failed to decode attributes".to_string(),
300 ));
301 }
302 }
303 } else {
304 struct PendingQuant {
306 att_id: i32,
307 portable: PointAttribute,
308 transform: AttributeQuantizationTransform,
309 }
310
311 struct PendingNormal {
312 att_id: i32,
313 portable: PointAttribute,
314 quantization_bits: u8,
315 }
316
317 struct AttributeSpec {
318 att_type: GeometryAttributeType,
319 data_type: DataType,
320 num_components: u8,
321 normalized: bool,
322 unique_id: u32,
323 }
324
325 for _ in 0..num_attributes_decoders {
326 let num_attributes_in_decoder: usize = if bitstream_version < 0x0200 {
327 buffer.decode_u32()? as usize
328 } else {
329 buffer.decode_varint()? as usize
330 };
331 if num_attributes_in_decoder == 0 {
332 return Err(DracoError::DracoError(
333 "Invalid number of attributes".to_string(),
334 ));
335 }
336 validate_num_attributes_in_decoder(
337 num_attributes_in_decoder,
338 buffer.remaining_size(),
339 )?;
340
341 let mut attribute_specs: Vec<AttributeSpec> =
342 Vec::with_capacity(num_attributes_in_decoder);
343 let mut att_ids: Vec<i32> = Vec::with_capacity(num_attributes_in_decoder);
344 let mut decoder_types: Vec<u8> = Vec::with_capacity(num_attributes_in_decoder);
345 let mut pending_quant: Vec<PendingQuant> = Vec::new();
346 let mut pending_normals: Vec<PendingNormal> = Vec::new();
347
348 for _ in 0..num_attributes_in_decoder {
349 let att_type_val = buffer.decode_u8()?;
350 let att_type = GeometryAttributeType::try_from(att_type_val)?;
351
352 let data_type_val = buffer.decode_u8()?;
353 let data_type = DataType::try_from(data_type_val)?;
354
355 let num_components = buffer.decode_u8()?;
356 validate_num_components(num_components)?;
357 let normalized = buffer.decode_u8()? != 0;
358 let unique_id: u32 = if bitstream_version < 0x0103 {
359 buffer.decode_u16()? as u32
360 } else {
361 buffer.decode_varint()? as u32
362 };
363
364 attribute_specs.push(AttributeSpec {
365 att_type,
366 data_type,
367 num_components,
368 normalized,
369 unique_id,
370 });
371 }
372
373 for _ in 0..num_attributes_in_decoder {
374 decoder_types.push(buffer.decode_u8()?);
375 }
376
377 for (local_i, spec) in attribute_specs.iter().enumerate() {
378 if decoder_types[local_i] == 0 {
379 let entry_size =
380 spec.num_components as usize * spec.data_type.byte_length();
381 let bytes_needed = entry_size.checked_mul(num_points).ok_or_else(|| {
382 DracoError::DracoError(
383 "Raw point cloud attribute byte count overflow".to_string(),
384 )
385 })?;
386 if buffer.remaining_size() < bytes_needed {
387 return Err(DracoError::DracoError(
388 "Not enough data for raw point cloud attribute values".to_string(),
389 ));
390 }
391 }
392
393 let mut att = PointAttribute::new();
394 att.try_init(
395 spec.att_type,
396 spec.num_components,
397 spec.data_type,
398 spec.normalized,
399 num_points,
400 )?;
401 att.set_unique_id(spec.unique_id);
402 let att_id = pc.add_attribute_preserve_unique_id(att);
403 att_ids.push(att_id);
404 }
405
406 let point_ids = if decoder_types.iter().any(|&decoder_type| decoder_type != 0) {
407 Some(make_point_ids(num_points)?)
408 } else {
409 None
410 };
411
412 for (local_i, &att_id) in att_ids.iter().enumerate() {
413 let decoder_type = decoder_types[local_i];
414 match decoder_type {
415 1 => {
416 let point_ids = point_ids.as_ref().ok_or_else(|| {
417 DracoError::DracoError(
418 "Point ids missing for integer attribute decoder".to_string(),
419 )
420 })?;
421 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
422 att_decoder.init(self, att_id);
423 if !att_decoder.decode_values(
424 pc, point_ids, buffer, None, None, None, None, None, None,
425 ) {
426 return Err(DracoError::DracoError(
427 "Failed to decode integer attribute".to_string(),
428 ));
429 }
430 }
431 2 => {
432 let original = pc.try_attribute(att_id)?;
433 let (original_type, original_num_components) =
434 (original.attribute_type(), original.num_components());
435 let mut portable = PointAttribute::default();
436 portable.try_init(
437 original_type,
438 original_num_components,
439 DataType::Uint32,
440 false,
441 num_points,
442 )?;
443 let mut transform = AttributeQuantizationTransform::new();
444
445 let quant_skip_bytes = if bitstream_version < 0x0102 {
449 let saved_pos = buffer.position();
450 let method_byte = buffer.decode_u8().map_err(|_| {
451 DracoError::DracoError("read pred method".to_string())
452 })?;
453 if method_byte != 0xFF {
454 let _transform_byte = buffer.decode_u8().map_err(|_| {
455 DracoError::DracoError("read transform".to_string())
456 })?;
457 }
458 let original = pc.try_attribute(att_id)?;
459 if !transform.decode_parameters(original, buffer) {
460 return Err(DracoError::DracoError(
461 "Failed to decode quantization parameters (v<2.0)"
462 .to_string(),
463 ));
464 }
465 let bytes_consumed = buffer.position() - saved_pos;
466 let pred_header_bytes = if method_byte != 0xFF { 2 } else { 1 };
467 let skip = bytes_consumed - pred_header_bytes;
468 buffer
469 .set_position(saved_pos)
470 .map_err(|_| DracoError::DracoError("buf reset".to_string()))?;
471 skip
472 } else {
473 0
474 };
475 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
476 att_decoder.init(self, att_id);
477 let mut skip_fn =
478 move |buf: &mut crate::decoder_buffer::DecoderBuffer<'_>| -> bool {
479 if quant_skip_bytes > 0
480 && buf.try_advance(quant_skip_bytes).is_err()
481 {
482 return false;
483 }
484 true
485 };
486 let hook: Option<
487 &mut dyn FnMut(
488 &mut crate::decoder_buffer::DecoderBuffer<'_>,
489 ) -> bool,
490 > = if quant_skip_bytes > 0 {
491 Some(&mut skip_fn)
492 } else {
493 None
494 };
495 if !att_decoder.decode_values(
496 pc,
497 point_ids.as_ref().ok_or_else(|| {
498 DracoError::DracoError(
499 "Point ids missing for quantized attribute decoder"
500 .to_string(),
501 )
502 })?,
503 buffer,
504 None,
505 None,
506 None,
507 Some(&mut portable),
508 None,
509 hook,
510 ) {
511 return Err(DracoError::DracoError(
512 "Failed to decode quantized portable values".to_string(),
513 ));
514 }
515 pending_quant.push(PendingQuant {
516 att_id,
517 portable,
518 transform,
519 });
520 }
521 3 => {
522 let mut portable = PointAttribute::default();
523 portable.try_init(
524 GeometryAttributeType::Generic,
525 2,
526 DataType::Uint32,
527 false,
528 num_points,
529 )?;
530 let mut quant_bits: u8 = 0;
534 let normal_skip_bytes = if bitstream_version < 0x0102 {
535 let saved_pos = buffer.position();
536 let method_byte = buffer.decode_u8().map_err(|_| {
537 DracoError::DracoError("read pred method".to_string())
538 })?;
539 if method_byte != 0xFF {
540 let _transform_byte = buffer.decode_u8().map_err(|_| {
541 DracoError::DracoError("read transform".to_string())
542 })?;
543 }
544 quant_bits = buffer.decode_u8().map_err(|_| {
545 DracoError::DracoError("read normal quant_bits".to_string())
546 })?;
547 if !AttributeOctahedronTransform::is_valid_quantization_bits(
548 quant_bits as i32,
549 ) {
550 return Err(DracoError::DracoError(
551 "Invalid normal quantization bits".to_string(),
552 ));
553 }
554 let bytes_consumed = buffer.position() - saved_pos;
555 let pred_header_bytes = if method_byte != 0xFF { 2 } else { 1 };
556 let skip = bytes_consumed - pred_header_bytes;
557 buffer
558 .set_position(saved_pos)
559 .map_err(|_| DracoError::DracoError("buf reset".to_string()))?;
560 skip
561 } else {
562 0
563 };
564 let mut att_decoder = SequentialIntegerAttributeDecoder::new();
565 att_decoder.init(self, att_id);
566 let mut skip_fn =
567 move |buf: &mut crate::decoder_buffer::DecoderBuffer<'_>| -> bool {
568 if normal_skip_bytes > 0
569 && buf.try_advance(normal_skip_bytes).is_err()
570 {
571 return false;
572 }
573 true
574 };
575 let hook: Option<
576 &mut dyn FnMut(
577 &mut crate::decoder_buffer::DecoderBuffer<'_>,
578 ) -> bool,
579 > = if normal_skip_bytes > 0 {
580 Some(&mut skip_fn)
581 } else {
582 None
583 };
584 if !att_decoder.decode_values(
585 pc,
586 point_ids.as_ref().ok_or_else(|| {
587 DracoError::DracoError(
588 "Point ids missing for normal attribute decoder"
589 .to_string(),
590 )
591 })?,
592 buffer,
593 None,
594 None,
595 None,
596 Some(&mut portable),
597 None,
598 hook,
599 ) {
600 return Err(DracoError::DracoError(
601 "Failed to decode normal portable values".to_string(),
602 ));
603 }
604 pending_normals.push(PendingNormal {
605 att_id,
606 portable,
607 quantization_bits: quant_bits,
608 });
609 }
610 0 => {
611 decode_raw_attribute_values(
614 buffer,
615 pc.try_attribute_mut(att_id)?,
616 num_points,
617 )?;
618 }
619 _ => {
620 return Err(DracoError::DracoError(format!(
621 "Unsupported sequential decoder type: {}",
622 decoder_type
623 )));
624 }
625 }
626 }
627
628 for (local_i, &att_id) in att_ids.iter().enumerate() {
629 match decoder_types[local_i] {
630 2 if bitstream_version >= 0x0102 => {
631 let idx = pending_quant
632 .iter()
633 .position(|p| p.att_id == att_id)
634 .ok_or_else(|| {
635 DracoError::DracoError(
636 "Missing pending quantized attribute transform".to_string(),
637 )
638 })?;
639 let original = pc.try_attribute(att_id)?;
640 if !pending_quant[idx]
641 .transform
642 .decode_parameters(original, buffer)
643 {
644 return Err(DracoError::DracoError(
645 "Failed to decode quantization parameters".to_string(),
646 ));
647 }
648 }
649 3 if bitstream_version >= 0x0102 => {
650 let idx = pending_normals
651 .iter()
652 .position(|p| p.att_id == att_id)
653 .ok_or_else(|| {
654 DracoError::DracoError(
655 "Missing pending normal attribute transform".to_string(),
656 )
657 })?;
658 let quantization_bits = buffer.decode_u8()?;
659 if !AttributeOctahedronTransform::is_valid_quantization_bits(
660 quantization_bits as i32,
661 ) {
662 return Err(DracoError::DracoError(
663 "Invalid normal quantization bits".to_string(),
664 ));
665 }
666 pending_normals[idx].quantization_bits = quantization_bits;
667 }
668 _ => {}
669 }
670 }
671
672 for q in pending_quant {
673 let dst = pc.try_attribute_mut(q.att_id)?;
674 if !q.transform.inverse_transform_attribute(&q.portable, dst) {
675 return Err(DracoError::DracoError(
676 "Failed to dequantize attribute".to_string(),
677 ));
678 }
679 }
680 for n in pending_normals {
681 let mut oct = AttributeOctahedronTransform::new(-1);
682 if !oct.set_parameters(n.quantization_bits as i32) {
683 return Err(DracoError::DracoError(
684 "Invalid normal quantization bits".to_string(),
685 ));
686 }
687 let dst = pc.try_attribute_mut(n.att_id)?;
688 if !oct.inverse_transform_attribute_with_legacy_octahedron(
689 &n.portable,
690 dst,
691 bitstream_version < 0x0102,
692 ) {
693 return Err(DracoError::DracoError(
694 "Failed to decode normals".to_string(),
695 ));
696 }
697 }
698 }
699 }
700
701 Ok(())
702 }
703
704 pub fn get_geometry_type(&self) -> EncodedGeometryType {
706 self.geometry_type
707 }
708}
709
710#[cfg(all(test, feature = "point_cloud_decode"))]
711mod tests {
712 use super::*;
713
714 #[test]
715 fn decode_raw_attribute_values_rejects_required_size_overflow() {
716 let bytes = [];
717 let mut buffer = DecoderBuffer::new(&bytes);
718 let mut attribute = PointAttribute::new();
719 attribute.init(
720 GeometryAttributeType::Generic,
721 1,
722 DataType::Uint32,
723 false,
724 1,
725 );
726
727 let status = decode_raw_attribute_values(&mut buffer, &mut attribute, usize::MAX);
728
729 assert!(status.is_err());
730 }
731
732 #[test]
733 fn decode_raw_attribute_values_rejects_truncated_input() {
734 let bytes = [1u8, 2, 3];
735 let mut buffer = DecoderBuffer::new(&bytes);
736 let mut attribute = PointAttribute::new();
737 attribute.init(
738 GeometryAttributeType::Generic,
739 1,
740 DataType::Uint32,
741 false,
742 1,
743 );
744
745 let status = decode_raw_attribute_values(&mut buffer, &mut attribute, 1);
746
747 assert!(status.is_err());
748 }
749}