draco-core 2.3.0

//! KD-tree point-cloud attribute encoder.
//!
//! [`KdTreeAttributesEncoder`] compresses point-cloud attributes by recursively
//! partitioning points in space and coding positions (and other attributes)
//! against the KD-tree structure. Encode-side counterpart of
//! `KdTreeAttributesDecoder`. Port of Draco's `kd_tree_attributes_encoder.h`.

use crate::attribute_quantization_transform::AttributeQuantizationTransform;
use crate::attribute_transform::AttributeTransform;
use crate::draco_types::DataType;
use crate::dynamic_integer_points_kd_tree::{DynamicIntegerPointsKdTreeEncoder, PointDVector};
use crate::encoder_buffer::EncoderBuffer;
use crate::encoder_options::EncoderOptions;
use crate::geometry_indices::PointIndex;
use crate::point_cloud::PointCloud;
use crate::prediction_scheme::EntryToPointIdMap;
use crate::status::{DracoError, Status};

pub struct KdTreeAttributesEncoder {
    attribute_ids: Vec<i32>,
    num_components: u32,
    attribute_quantization_transforms: Vec<AttributeQuantizationTransform>,
    quantized_portable_attributes: Vec<crate::geometry_attribute::PointAttribute>,
    min_signed_values: Vec<i32>,
}

impl KdTreeAttributesEncoder {
    pub fn new(first_att_id: i32) -> Self {
        Self {
            attribute_ids: vec![first_att_id],
            num_components: 0,
            attribute_quantization_transforms: Vec::new(),
            quantized_portable_attributes: Vec::new(),
            min_signed_values: Vec::new(),
        }
    }

    pub fn add_attribute_id(&mut self, att_id: i32) {
        self.attribute_ids.push(att_id);
    }

    pub fn encode_attributes_encoder_data(
        &self,
        point_cloud: &PointCloud,
        out_buffer: &mut EncoderBuffer,
    ) -> Status {
        // Encode number of attributes
        out_buffer.encode_varint(self.attribute_ids.len() as u64);

        for &att_id in &self.attribute_ids {
            let att = point_cloud.attribute(att_id);
            // Encode attribute metadata
            out_buffer.encode_u8(att.attribute_type() as u8);
            out_buffer.encode_u8(att.data_type() as u8);
            out_buffer.encode_u8(att.num_components());
            out_buffer.encode_u8(if att.normalized() { 1 } else { 0 });
            out_buffer.encode_varint(att.unique_id() as u64);
        }
        Ok(())
    }

    pub fn attribute_ids(&self) -> &[i32] {
        &self.attribute_ids
    }

    pub fn transform_attributes_to_portable_format(
        &mut self,
        point_cloud: &PointCloud,
        options: &EncoderOptions,
    ) -> Status {
        self.attribute_quantization_transforms.clear();
        self.quantized_portable_attributes.clear();
        self.min_signed_values.clear();

        let num_points = point_cloud.num_points();
        let point_ids: Vec<PointIndex> = (0..num_points).map(|i| PointIndex(i as u32)).collect();

        let mut total_components: u32 = 0;
        for &att_id in &self.attribute_ids {
            let att = point_cloud.attribute(att_id);
            total_components += att.num_components() as u32;
        }
        self.num_components = total_components;

        for &att_id in &self.attribute_ids {
            let att = point_cloud.attribute(att_id);
            match att.data_type() {
                DataType::Float32 => {
                    let quantization_bits =
                        options.get_attribute_int(att_id, "quantization_bits", -1);
                    if quantization_bits < 1 {
                        return Err(DracoError::invalid_parameter(format!(
                            "Attribute {att_id} is float and needs quantization_bits,                              which is {quantization_bits}"
                        )));
                    }
                    let mut transform = AttributeQuantizationTransform::new();
                    transform.compute_parameters(att, quantization_bits)?;

                    let mut portable_att = crate::geometry_attribute::PointAttribute::default();
                    portable_att.init(
                        att.attribute_type(),
                        att.num_components(),
                        DataType::Uint32,
                        false,
                        num_points,
                    );
                    portable_att.set_identity_mapping();

                    transform.transform_attribute(
                        att,
                        EntryToPointIdMap::from_point_indices(&point_ids),
                        &mut portable_att,
                    )?;

                    self.attribute_quantization_transforms.push(transform);
                    self.quantized_portable_attributes.push(portable_att);
                }
                DataType::Int32 | DataType::Int16 | DataType::Int8 => {
                    // Determine per-component minimum.
                    let num_components = att.num_components() as usize;
                    let mut min_vals = vec![i32::MAX; num_components];
                    let stride = att.byte_stride() as usize;
                    let component_size = att.data_type().byte_length();
                    for i in 0..num_points {
                        let avi = att.mapped_index(PointIndex(i as u32));
                        let base = avi.0 as usize * stride;
                        for c in 0..num_components {
                            let v = read_as_i32(
                                att.buffer(),
                                base + c * component_size,
                                att.data_type(),
                            );
                            if v < min_vals[c] {
                                min_vals[c] = v;
                            }
                        }
                    }
                    self.min_signed_values.extend_from_slice(&min_vals);
                }
                _ => {}
            }
        }

        Ok(())
    }

    pub fn encode_attributes(
        &mut self,
        point_cloud: &PointCloud,
        options: &EncoderOptions,
        out_buffer: &mut EncoderBuffer,
    ) -> Status {
        // Draco C++: compression_level = min(10 - GetSpeed(), 6), and GetSpeed
        // is the larger of the encoding and decoding speeds.
        let speed = options.get_speed();
        // Saturating because the speed is a caller-set option with no declared
        // range and `10 - i32::MIN` overflows.
        //
        // The lower bound is this crate's, and deliberately so. Upstream writes
        // `std::min(10 - GetSpeed(), 6)` into a `uint8_t`, so a speed above 10
        // makes the subtraction negative and the narrowing turns it into 255 or
        // thereabouts; upstream then encodes that byte and only afterwards hits
        // the `default:` arm of its level switch and fails the encode, leaving a
        // stray level byte in the buffer. Its own `DRACO_DCHECK_LE` on the same
        // line says that state is not intended, so clamping keeps the invariant
        // upstream asserts rather than reproducing what its release build does
        // when the assert is compiled out. Verified against C++ 1.5.7: every
        // speed at or below 10 is byte-identical either way, and at 11 and above
        // upstream's encode fails outright rather than producing a stream this
        // could differ from.
        let mut compression_level: u8 = 10i32.saturating_sub(speed).clamp(0, 6) as u8;
        if compression_level == 6 && self.num_components > 15 {
            compression_level = 5;
        }

        out_buffer.encode_u8(compression_level);

        let num_points = point_cloud.num_points();
        let mut point_vector = PointDVector::new(num_points, self.num_components as usize);

        let mut num_processed_components: usize = 0;
        let mut num_processed_quantized_attributes: usize = 0;
        let mut num_processed_signed_components: usize = 0;

        for &att_id in &self.attribute_ids {
            let att = point_cloud.attribute(att_id);
            let use_quantized;

            let src: &crate::geometry_attribute::PointAttribute = match att.data_type() {
                DataType::Uint32
                | DataType::Uint16
                | DataType::Uint8
                | DataType::Int32
                | DataType::Int16
                | DataType::Int8 => {
                    use_quantized = false;
                    att
                }
                DataType::Float32 => {
                    use_quantized = true;
                    let pa =
                        &self.quantized_portable_attributes[num_processed_quantized_attributes];
                    num_processed_quantized_attributes += 1;
                    pa
                }
                other => {
                    return Err(DracoError::unsupported_feature(format!(
                        "The KD-tree coder does not carry {other:?} attributes"
                    )));
                }
            };

            let num_att_components = src.num_components() as usize;
            let stride = src.byte_stride() as usize;
            let component_size = src.data_type().byte_length();

            match src.data_type() {
                DataType::Uint32 => {
                    for p in 0..num_points {
                        let avi = src.mapped_index(PointIndex(p as u32));
                        let base = avi.0 as usize * stride;
                        let dst = point_vector.point_mut(p);
                        for c in 0..num_att_components {
                            let v = read_as_u32(
                                src.buffer(),
                                base + c * component_size,
                                DataType::Uint32,
                            );
                            dst[num_processed_components + c] = v;
                        }
                    }
                }
                DataType::Int32 | DataType::Int16 | DataType::Int8 => {
                    for p in 0..num_points {
                        let avi = src.mapped_index(PointIndex(p as u32));
                        let base = avi.0 as usize * stride;
                        let dst = point_vector.point_mut(p);
                        for c in 0..num_att_components {
                            let signed = read_as_i32(
                                src.buffer(),
                                base + c * component_size,
                                src.data_type(),
                            );
                            let minv = self.min_signed_values[num_processed_signed_components + c];
                            // Wrapping, as C++ `signed_point[c] - min_signed_values_[..]`
                            // assigned to a uint32_t is: the difference of two
                            // i32 does not fit i32 when the attribute spans the
                            // full range, but two's complement makes the
                            // wrapped bits the exact u32 difference, which is
                            // what the decoder adds `minv` back to. Plain
                            // subtraction only differed by panicking in debug.
                            dst[num_processed_components + c] = signed.wrapping_sub(minv) as u32;
                        }
                    }
                    num_processed_signed_components += num_att_components;
                }
                DataType::Uint16 | DataType::Uint8 => {
                    for p in 0..num_points {
                        let avi = src.mapped_index(PointIndex(p as u32));
                        let base = avi.0 as usize * stride;
                        let dst = point_vector.point_mut(p);
                        for c in 0..num_att_components {
                            let v = read_as_u32(
                                src.buffer(),
                                base + c * component_size,
                                src.data_type(),
                            );
                            dst[num_processed_components + c] = v;
                        }
                    }
                }
                other => {
                    // Float32 reaches here only as its Uint32 portable copy, so
                    // this arm means the portable conversion did not happen.
                    return Err(DracoError::general(format!(
                        "Attribute {att_id} reached the KD-tree coder as {other:?}                         {}",
                        if use_quantized {
                            ", with a portable copy that is not Uint32"
                        } else {
                            ""
                        }
                    )));
                }
            }

            num_processed_components += num_att_components;
        }

        // Compute maximum bit length.
        let mut num_bits: u32 = 0;
        for &v in point_vector.as_slice() {
            if v != 0 {
                let msb = 32 - v.leading_zeros();
                if msb > num_bits {
                    num_bits = msb;
                }
            }
        }

        let mut encoder =
            DynamicIntegerPointsKdTreeEncoder::new(compression_level, self.num_components);
        encoder.encode_points(&mut point_vector, num_bits, out_buffer);
        Ok(())
    }

    /// Writes the pre-2.3 layout, as Draco 1.0 to 1.2 wrote it in
    /// `KdTreeAttributesEncoder::EncodeDataNeededByPortableTransforms`.
    ///
    /// Those encoders took one three-component attribute and nothing else. A
    /// `Float32` one goes through `FloatPointsTreeEncoder`: a header of its own
    /// with one range for all three axes, every coordinate quantized against
    /// it by that era's `Quantizer` and offset by the largest quantized value to
    /// make it unsigned, and the tree told the bit length `bits + 1` rather than
    /// measuring it. A `Uint32` one is written as it is, the tree told 32 bits.
    /// Everything that decides a byte follows those encoders, so the output is
    /// theirs: compared against Draco 1.0.0, 1.1.0 and 1.2.5 in the C++ bridge.
    #[cfg(feature = "legacy_bitstream_encode")]
    pub fn encode_pre_2_3(
        &self,
        point_cloud: &PointCloud,
        options: &EncoderOptions,
        out_buffer: &mut EncoderBuffer,
    ) -> Status {
        let att_id = match self.attribute_ids.as_slice() {
            &[att_id] => att_id,
            ids => {
                return Err(DracoError::unsupported_feature(format!(
                    "A pre-2.3 KD-tree carries one attribute, not {}",
                    ids.len()
                )))
            }
        };
        let att = point_cloud.attribute(att_id);
        if att.num_components() != 3 {
            return Err(DracoError::unsupported_feature(format!(
                "A pre-2.3 KD-tree carries three components, not {}",
                att.num_components()
            )));
        }
        let compression_level = 10i32.saturating_sub(options.get_speed()).clamp(0, 6) as u8;
        let num_points = point_cloud.num_points();
        let read = |p: usize| -> [u32; 3] {
            let base =
                att.mapped_index(PointIndex(p as u32)).0 as usize * att.byte_stride() as usize;
            let mut bytes = [0u8; 12];
            att.buffer().read(base, &mut bytes);
            std::array::from_fn(|c| u32::from_le_bytes(bytes[c * 4..][..4].try_into().unwrap()))
        };
        let mut points = PointDVector::new(num_points, 3);

        let bit_length = match att.data_type() {
            DataType::Float32 => {
                let bits = options.get_attribute_int(att_id, "quantization_bits", -1);
                if !(1..=30).contains(&bits) {
                    return Err(DracoError::invalid_parameter(format!(
                        "A pre-2.3 KD-tree quantizes positions to 1..=30 bits, not {bits}"
                    )));
                }
                let bits = bits as u32;
                let coordinates: Vec<[f32; 3]> = (0..num_points)
                    .map(|p| read(p).map(f32::from_bits))
                    .collect();
                // `std::max(std::fabs(c), max_range)`, argument order kept: a
                // NaN coordinate becomes the range, as it does upstream.
                let range = coordinates.iter().flatten().fold(0f32, |max, c| {
                    if c.abs() < max {
                        max
                    } else {
                        c.abs()
                    }
                });
                let max_quantized_value = (1u32 << bits) - 1;

                out_buffer.encode_u8(0); // kKdTreeQuantizationEncoding
                out_buffer.encode_u8(compression_level);
                out_buffer.encode_u32(num_points as u32);
                out_buffer.encode_u32(3); // FloatPointsTreeEncoder::version_
                out_buffer.encode_u8(1); // KDTREE
                out_buffer.encode_u32(bits);
                out_buffer.encode_u32(range.to_bits());
                out_buffer.encode_u32(num_points as u32);
                out_buffer.encode_u32(compression_level as u32);
                if num_points == 0 {
                    return Ok(());
                }
                for (p, point) in coordinates.iter().enumerate() {
                    let dst = points.point_mut(p);
                    for (c, &value) in point.iter().enumerate() {
                        dst[c] = (quantize_pre_2_3(value, range, max_quantized_value) as u32)
                            .wrapping_add(max_quantized_value);
                    }
                }
                bits + 1
            }
            DataType::Uint32 => {
                out_buffer.encode_u8(1); // kKdTreeIntegerEncoding
                out_buffer.encode_u8(compression_level);
                out_buffer.encode_u32(num_points as u32);
                for p in 0..num_points {
                    points.point_mut(p).copy_from_slice(&read(p));
                }
                32
            }
            other => {
                return Err(DracoError::unsupported_feature(format!(
                    "A pre-2.3 KD-tree carries Float32 or Uint32, not {other:?}"
                )))
            }
        };
        DynamicIntegerPointsKdTreeEncoder::new(compression_level, 3).encode_points(
            &mut points,
            bit_length,
            out_buffer,
        );
        Ok(())
    }

    pub fn encode_data_needed_by_portable_transforms(
        &self,
        out_buffer: &mut EncoderBuffer,
    ) -> Status {
        for t in &self.attribute_quantization_transforms {
            t.encode_parameters(out_buffer)?;
        }

        for &minv in &self.min_signed_values {
            out_buffer.encode_varint_signed_i32(minv);
        }

        Ok(())
    }
}

/// The `Quantizer::QuantizeFloat` of Draco 1.0 to 1.2: the magnitude divided by
/// the range, scaled, rounded half up, and the sign put back.
///
/// The float-to-integer conversion is `static_cast<int32_t>` as an x86 build
/// performs it, which gives `i32::MIN` for a NaN or a value out of range -- the
/// case for every coordinate of a cloud whose range is zero, where the
/// division is `0 / 0`. Rust's `as` saturates instead, and would write other
/// bytes for the same cloud.
#[cfg(feature = "legacy_bitstream_encode")]
fn quantize_pre_2_3(value: f32, range: f32, max_quantized_value: u32) -> i32 {
    let negative = value < 0.0;
    let magnitude = if negative { -value } else { value } / range;
    let rounded = (magnitude * max_quantized_value as f32 + 0.5).floor();
    let truncated = if rounded.is_nan() || !(-2_147_483_648.0..2_147_483_648.0).contains(&rounded) {
        i32::MIN
    } else {
        rounded as i32
    };
    if negative {
        truncated.wrapping_neg()
    } else {
        truncated
    }
}

fn read_as_i32(buffer: &crate::data_buffer::DataBuffer, offset: usize, data_type: DataType) -> i32 {
    match data_type {
        DataType::Int8 => {
            let mut bytes = [0u8; 1];
            buffer.read(offset, &mut bytes);
            bytes[0] as i8 as i32
        }
        DataType::Int16 => {
            let mut bytes = [0u8; 2];
            buffer.read(offset, &mut bytes);
            i16::from_le_bytes(bytes) as i32
        }
        DataType::Int32 => {
            let mut bytes = [0u8; 4];
            buffer.read(offset, &mut bytes);
            i32::from_le_bytes(bytes)
        }
        _ => 0,
    }
}

fn read_as_u32(buffer: &crate::data_buffer::DataBuffer, offset: usize, data_type: DataType) -> u32 {
    match data_type {
        DataType::Uint8 => {
            let mut bytes = [0u8; 1];
            buffer.read(offset, &mut bytes);
            bytes[0] as u32
        }
        DataType::Uint16 => {
            let mut bytes = [0u8; 2];
            buffer.read(offset, &mut bytes);
            u16::from_le_bytes(bytes) as u32
        }
        DataType::Uint32 => {
            let mut bytes = [0u8; 4];
            buffer.read(offset, &mut bytes);
            u32::from_le_bytes(bytes)
        }
        _ => 0,
    }
}