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