draco_core/

decoder_buffer.rs

use crate::status::DracoError;
use crate::version::DEFAULT_MESH_VERSION;
use std::mem;

/// Input buffer for reading compressed Draco data.
///
/// `DecoderBuffer` provides sequential byte and bit-level access to compressed data.
/// It supports both byte-aligned reads (integers, floats, strings) and bit-level
/// reads for entropy-coded data.
///
/// # Example
///
/// ```
/// use draco_core::DecoderBuffer;
///
/// let data = &[0x44, 0x52, 0x41, 0x43, 0x4F]; // "DRACO" header
/// let mut buffer = DecoderBuffer::new(data);
///
/// assert_eq!(buffer.decode_u8().unwrap(), 0x44);
/// assert_eq!(buffer.remaining_size(), 4);
/// ```
pub struct DecoderBuffer<'a> {
    data: &'a [u8],
    pos: usize,
    bit_decoder_active: bool,
    bit_start_pos: usize,
    current_bit_offset: usize,
    bit_stream_end_pos: usize,
    bit_sequence_size_known: bool,
    version_major: u8,
    version_minor: u8,
    /// Bytes this decode has reserved so far, across every buffer it sized.
    ///
    /// The budget belongs to *one decode of one stream*, and that is what this
    /// type is: built once per decode, carrying the stream, and already
    /// threaded to every site that sizes something. Holding the counter here
    /// is what makes the bound cumulative rather than per-allocation -- see
    /// [`charge`](Self::charge).
    spent: usize,
    /// The caller's ceilings on what this decode may produce, and what it has
    /// produced so far against the byte one.
    ///
    /// Here rather than on the decoders because this type is the stream and
    /// already reaches every site that learns a count -- the same reason the
    /// budget above lives here. A decode of a mesh runs a point-cloud decoder
    /// and an attributes decoder under it; one buffer is what all three share.
    limits: crate::decode_limits::DecodeLimits,
    decoded_bytes: u64,
}

impl<'a> DecoderBuffer<'a> {
    /// Creates a new `DecoderBuffer` from a byte slice.
    pub fn new(data: &'a [u8]) -> Self {
        Self {
            data,
            pos: 0,
            bit_decoder_active: false,
            bit_start_pos: 0,
            current_bit_offset: 0,
            bit_stream_end_pos: 0,
            bit_sequence_size_known: false,
            // Default to latest mesh version to match encoder output format
            version_major: DEFAULT_MESH_VERSION.0,
            version_minor: DEFAULT_MESH_VERSION.1,
            spent: 0,
            limits: crate::decode_limits::DecodeLimits::default(),
            decoded_bytes: 0,
        }
    }

    /// Decodes under `limits` rather than under
    /// [`DecodeLimits::default`](crate::DecodeLimits::default).
    ///
    /// ```
    /// use draco_core::{DecodeLimits, DecoderBuffer};
    ///
    /// let stream = [0u8; 0];
    /// let buffer = DecoderBuffer::new(&stream).with_limits(DecodeLimits::permissive());
    /// # let _ = buffer;
    /// ```
    #[must_use]
    pub fn with_limits(mut self, limits: crate::decode_limits::DecodeLimits) -> Self {
        self.limits = limits;
        self
    }

    /// Refuses a decoded point count over the caller's ceiling.
    #[cfg(feature = "point_cloud_decode")]
    pub(crate) fn check_points(&self, points: usize) -> crate::status::Status {
        self.limits.check_points(points as u64)
    }

    /// Refuses a decoded face count over the caller's ceiling.
    pub(crate) fn check_faces(&self, faces: usize) -> crate::status::Status {
        self.limits.check_faces(faces as u64)
    }

    /// Adds one attribute's declared size to this decode's running total and
    /// refuses when the total passes the caller's ceiling.
    ///
    /// Cumulative for the same reason [`charge`](Self::charge) is: the
    /// attribute count is the header's, so a per-attribute check bounds
    /// nothing. Called where the size becomes known and before anything is
    /// allocated for it.
    pub(crate) fn charge_decoded_bytes(&mut self, bytes: usize) -> crate::status::Status {
        let total = self.decoded_bytes.saturating_add(bytes as u64);
        self.limits.check_decoded_bytes(total)?;
        self.decoded_bytes = total;
        Ok(())
    }

    /// Charges `bytes` against what this stream is allowed to make the decoder
    /// allocate, and refuses when the running total outgrows it.
    ///
    /// Cumulative on purpose. Checked per allocation, the same budget is
    /// granted again at every site, and the number of sites is the attacker's:
    /// one attributes decoder carries as many attributes as the header names,
    /// and their buffers coexist. Measured before this existed, on a stream
    /// held at ~32 KB while only the attribute count changed:
    ///
    /// | attributes | reserved |
    /// | ---: | ---: |
    /// | 1 | 19.2 MB |
    /// | 4 | 76.8 MB |
    /// | 16 | 307.2 MB |
    ///
    /// Linear in a count that costs 7.5 bytes of input each, so the per-site
    /// form bounded nothing that mattered. One running total does.
    ///
    /// Nothing is ever credited back. A decode that frees a buffer and sizes
    /// another still pays for both, which is why the ceiling has to cover the
    /// *sum* over a real file rather than its peak -- see
    /// [`MAX_ALLOCATED_BYTES_PER_INPUT_BYTE`], whose value is measured against
    /// exactly that.
    ///
    /// [`MAX_ALLOCATED_BYTES_PER_INPUT_BYTE`]: crate::decode_budget::MAX_ALLOCATED_BYTES_PER_INPUT_BYTE
    pub(crate) fn charge(&mut self, bytes: usize) -> crate::status::Status {
        let total = self.spent.saturating_add(bytes);
        crate::decode_budget::ensure_allocation_is_backed(total, self.data.len())?;
        self.spent = total;
        Ok(())
    }

    /// [`charge`](Self::charge) for a count of `element_size`-byte elements.
    pub(crate) fn charge_elements(
        &mut self,
        count: usize,
        element_size: usize,
    ) -> crate::status::Status {
        self.charge(count.saturating_mul(element_size))
    }

    /// What this decode has reserved so far, for the tests that hold the
    /// budget to being a backstop no legitimate file reaches.
    #[cfg(test)]
    pub(crate) fn spent(&self) -> usize {
        self.spent
    }

    /// Sets the Draco bitstream version for version-dependent decoding.
    pub fn set_version(&mut self, major: u8, minor: u8) {
        self.version_major = major;
        self.version_minor = minor;
    }

    /// Returns the major version number.
    pub fn version_major(&self) -> u8 {
        self.version_major
    }

    /// Returns the minor version number.
    pub fn version_minor(&self) -> u8 {
        self.version_minor
    }

    /// Returns the packed `0xMMmm` bitstream version for ordered comparisons.
    pub fn bitstream_version(&self) -> u16 {
        crate::version::bitstream_version(self.version_major, self.version_minor)
    }

    /// Returns the current read position in bytes.
    pub fn position(&self) -> usize {
        self.pos
    }

    /// Sets the read position.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if:
    /// - Bit decoding is currently active
    /// - Position is beyond the buffer length
    pub fn set_position(&mut self, pos: usize) -> Result<(), DracoError> {
        if self.bit_decoder_active {
            return Err(DracoError::buffer(
                "Cannot set position while bit decoding is active",
            ));
        }
        if pos > self.data.len() {
            return Err(DracoError::buffer(format!(
                "Position {} exceeds buffer length {}",
                pos,
                self.data.len()
            )));
        }
        self.pos = pos;
        Ok(())
    }

    /// Returns the number of bytes remaining in the buffer.
    pub fn remaining_size(&self) -> usize {
        self.data.len().saturating_sub(self.pos)
    }

    /// Returns the total size of the buffer, read or not.
    ///
    /// The decode allocation budget measures against this rather than against
    /// [`remaining_size`](Self::remaining_size): an attribute decoded last has
    /// a legitimately tiny payload left in front of it while its value count is
    /// no smaller than the first attribute's, so charging it only for what
    /// follows would refuse valid streams.
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Peeks at the next `len` bytes without advancing the position.
    pub fn peek_bytes(&self, len: usize) -> Vec<u8> {
        let start = self.pos.min(self.data.len());
        let end = start.saturating_add(len).min(self.data.len());
        self.data[start..end].to_vec()
    }

    /// Starts bit-level decoding mode.
    ///
    /// When `decode_size` is true, reads the bit sequence size from the buffer.
    /// Returns the size in bytes.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if bit decoding is already active.
    pub fn start_bit_decoding(&mut self, decode_size: bool) -> Result<u64, DracoError> {
        if self.bit_decoder_active {
            return Err(DracoError::buffer("Bit decoding already active"));
        }
        let bitstream_version = self.bitstream_version();
        // Draco stores the bit-sequence size in BYTES (not bits) when |decode_size| is true.
        let mut size_bytes: u64 = 0;
        if decode_size {
            if bitstream_version < 0x0202 {
                if !cfg!(feature = "legacy_bitstream_decode") {
                    return Err(DracoError::bitstream_version_unsupported());
                }
                size_bytes = self.decode_u64()?;
            } else {
                size_bytes = self.decode_varint()?;
            }
        }

        self.bit_start_pos = self.pos;
        self.bit_decoder_active = true;
        self.current_bit_offset = 0;
        self.bit_sequence_size_known = decode_size;

        if decode_size {
            let size_bytes = usize::try_from(size_bytes)
                .map_err(|_| DracoError::buffer("Bit stream size too large"))?;
            // Bounded by the buffer, not by the number: the size is a varint
            // out of the stream, and a stream claiming more than it carries is
            // one upstream reads to the end of rather than refuses -- its bit
            // decoder is handed `remaining_size()` and never sees the claim.
            // Left unbounded, the claim lands in `pos` when bit decoding ends,
            // and a position past the buffer is one no later bounds check is
            // written to expect.
            let declared_end = self
                .bit_start_pos
                .checked_add(size_bytes)
                .ok_or_else(|| DracoError::buffer("Bit stream end position overflow"))?;
            self.bit_stream_end_pos = declared_end.min(self.data.len());
        } else {
            // If size is not encoded, assume the rest of the buffer.
            self.bit_stream_end_pos = self.data.len();
        }

        Ok(size_bytes)
    }

    /// Ends bit-level decoding mode and advances the byte position.
    pub fn end_bit_decoding(&mut self) {
        self.bit_decoder_active = false;
        // Draco behavior:
        // - When decoding with size known, the caller typically skips by the stored byte size.
        // - When decoding without size, advance by the number of decoded bits (rounded up).
        if self.bit_sequence_size_known {
            self.pos = self.bit_stream_end_pos;
        } else {
            let bytes_consumed = self.current_bit_offset.div_ceil(8);
            self.pos = self.bit_start_pos + bytes_consumed;
        }
    }

    /// Decodes `nbits` least significant bits as a u32.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if bit decoding is not active or end of stream.
    #[inline(always)]
    pub fn decode_least_significant_bits32(&mut self, nbits: u32) -> Result<u32, DracoError> {
        if !self.bit_decoder_active {
            return Err(DracoError::buffer("Bit decoding not active"));
        }
        self.decode_least_significant_bits32_fast(nbits)
    }

    /// Optimized version for hot paths - reads multiple bytes at once.
    #[inline(always)]
    pub fn decode_least_significant_bits32_fast(&mut self, nbits: u32) -> Result<u32, DracoError> {
        if nbits == 0 {
            return Ok(0);
        }
        // The tagged symbol scheme's per-chunk bit length is a full byte
        // (0..=255) read straight off the wire and checked only against
        // `> 32`, so 32 itself reaches here as a legitimate width -- a
        // 32-bit-wide residual is what the scheme exists for. `nbits` past
        // 32 is not a width this format defines.
        if nbits > 32 {
            return Err(DracoError::buffer("Bit width exceeds 32 bits"));
        }

        let total_bit_offset = self.current_bit_offset;
        let byte_offset = self.bit_start_pos + total_bit_offset / 8;
        let bit_shift = (total_bit_offset % 8) as u32;

        if byte_offset >= self.bit_stream_end_pos || byte_offset >= self.data.len() {
            return Err(DracoError::buffer("Unexpected end of bit stream"));
        }
        let available_end = self.bit_stream_end_pos.min(self.data.len());
        let remaining = available_end - byte_offset;

        // Fast path: read 8 bytes at once when enough data remains (avoids per-byte loop).
        let raw = if remaining >= 8 {
            let mut bytes = [0u8; 8];
            bytes.copy_from_slice(&self.data[byte_offset..byte_offset + 8]);
            u64::from_le_bytes(bytes)
        } else {
            let needed_bytes = (bit_shift + nbits).div_ceil(8) as usize;
            if remaining < needed_bytes {
                return Err(DracoError::buffer("Unexpected end of bit stream"));
            }
            let mut v = 0u64;
            for i in 0..needed_bytes {
                v |= (self.data[byte_offset + i] as u64) << (i * 8);
            }
            v
        };
        // Masking with a u64 keeps `nbits == 32` in range for the shift;
        // `1u32 << 32` is exactly the overflow this guards against.
        let mask = (1u64 << nbits) - 1;
        let value = ((raw >> bit_shift) & mask) as u32;

        self.current_bit_offset += nbits as usize;
        Ok(value)
    }

    #[inline]
    #[allow(dead_code)]
    fn get_bit(&mut self) -> Result<u32, DracoError> {
        let total_bit_offset = self.current_bit_offset;
        let byte_offset = self.bit_start_pos + total_bit_offset / 8;
        let bit_shift = total_bit_offset % 8;

        if byte_offset < self.bit_stream_end_pos && byte_offset < self.data.len() {
            let bit = (self.data[byte_offset] >> bit_shift) & 1;
            self.current_bit_offset += 1;
            Ok(bit as u32)
        } else {
            Err(DracoError::buffer("Unexpected end of bit stream"))
        }
    }

    /// Decodes a value of type T using raw memory copy.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if:
    /// - Bit decoding is active
    /// - Not enough bytes remaining
    pub fn decode<T: Copy + bytemuck::Pod>(&mut self) -> Result<T, DracoError> {
        if self.bit_decoder_active {
            return Err(DracoError::buffer(
                "Cannot decode bytes while bit decoding is active",
            ));
        }
        let size = mem::size_of::<T>();
        if size > self.data.len().saturating_sub(self.pos) {
            return Err(DracoError::buffer(format!(
                "Unexpected end of buffer: need {} bytes, have {}",
                size,
                self.remaining_size()
            )));
        }

        // Safety: bytemuck::Pod guarantees T can be safely read from any bit pattern
        let val = bytemuck::pod_read_unaligned::<T>(&self.data[self.pos..self.pos + size]);
        self.pos += size;
        Ok(val)
    }

    /// Decodes a single byte.
    pub fn decode_u8(&mut self) -> Result<u8, DracoError> {
        self.decode::<u8>()
    }

    /// Decodes a little-endian u16.
    pub fn decode_u16(&mut self) -> Result<u16, DracoError> {
        let mut bytes = [0u8; 2];
        self.decode_bytes(&mut bytes)?;
        Ok(u16::from_le_bytes(bytes))
    }

    /// Decodes a little-endian u32.
    pub fn decode_u32(&mut self) -> Result<u32, DracoError> {
        let mut bytes = [0u8; 4];
        self.decode_bytes(&mut bytes)?;
        Ok(u32::from_le_bytes(bytes))
    }

    /// Decodes a little-endian u64.
    pub fn decode_u64(&mut self) -> Result<u64, DracoError> {
        let mut bytes = [0u8; 8];
        self.decode_bytes(&mut bytes)?;
        Ok(u64::from_le_bytes(bytes))
    }

    /// Decodes a little-endian f32.
    pub fn decode_f32(&mut self) -> Result<f32, DracoError> {
        let mut bytes = [0u8; 4];
        self.decode_bytes(&mut bytes)?;
        Ok(f32::from_le_bytes(bytes))
    }

    /// Decodes a little-endian f64.
    pub fn decode_f64(&mut self) -> Result<f64, DracoError> {
        let mut bytes = [0u8; 8];
        self.decode_bytes(&mut bytes)?;
        Ok(f64::from_le_bytes(bytes))
    }

    /// Decodes a null-terminated string.
    pub fn decode_string(&mut self) -> Result<String, DracoError> {
        let mut bytes = Vec::new();
        loop {
            let b = self.decode_u8()?;
            if b == 0 {
                break;
            }
            bytes.push(b);
        }
        String::from_utf8(bytes)
            .map_err(|e| DracoError::buffer(format!("Invalid UTF-8 string: {}", e)))
    }

    /// Decodes bytes into the provided buffer.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if not enough bytes remaining.
    pub fn decode_bytes(&mut self, out: &mut [u8]) -> Result<(), DracoError> {
        let size = out.len();
        if size > self.data.len().saturating_sub(self.pos) {
            return Err(DracoError::buffer(format!(
                "Unexpected end of buffer: need {} bytes, have {}",
                size,
                self.remaining_size()
            )));
        }
        out.copy_from_slice(&self.data[self.pos..self.pos + size]);
        self.pos += size;
        Ok(())
    }

    /// Decodes a variable-length unsigned integer (varint).
    pub fn decode_varint(&mut self) -> Result<u64, DracoError> {
        let mut val = 0u64;
        let mut shift = 0;
        loop {
            let b = self.decode_u8()?;
            val |= ((b & 0x7F) as u64) << shift;
            if (b & 0x80) == 0 {
                break;
            }
            shift += 7;
            if shift >= 64 {
                return Err(DracoError::buffer("Varint exceeds 64 bits"));
            }
        }
        Ok(val)
    }

    /// Decodes a Draco-compatible signed varint.
    ///
    /// Uses unsigned varint encoding with ConvertSymbolToSignedInt transformation.
    pub fn decode_varint_signed_i32(&mut self) -> Result<i32, DracoError> {
        let symbol = self.decode_varint()? as u32;
        let is_positive = (symbol & 1) == 0;
        let v = symbol >> 1;
        if is_positive {
            Ok(v as i32)
        } else {
            Ok(-(v as i32) - 1)
        }
    }

    /// Returns a slice of the remaining data without advancing.
    pub fn remaining_data(&self) -> &'a [u8] {
        &self.data[self.pos..]
    }

    /// Advances the position by `n` bytes without reading.
    pub fn advance(&mut self, n: usize) {
        self.pos = self.pos.saturating_add(n).min(self.data.len());
    }

    /// Advances the position by `n` bytes without reading.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if the requested advance would move
    /// beyond the end of the input buffer.
    pub fn try_advance(&mut self, n: usize) -> Result<(), DracoError> {
        let new_pos = self
            .pos
            .checked_add(n)
            .ok_or_else(|| DracoError::buffer("Buffer advance overflow"))?;
        if new_pos > self.data.len() {
            return Err(DracoError::buffer(format!(
                "Cannot advance buffer by {} bytes: need position {}, buffer length {}",
                n,
                new_pos,
                self.data.len()
            )));
        }
        self.pos = new_pos;
        Ok(())
    }

    /// Decodes and returns a slice of the specified size.
    ///
    /// # Errors
    ///
    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if not enough bytes remaining.
    pub fn decode_slice(&mut self, size: usize) -> Result<&'a [u8], DracoError> {
        if size > self.data.len().saturating_sub(self.pos) {
            return Err(DracoError::buffer(format!(
                "Unexpected end of buffer: need {} bytes, have {}",
                size,
                self.remaining_size()
            )));
        }
        let slice = &self.data[self.pos..self.pos + size];
        self.pos += size;
        Ok(slice)
    }
}

#[cfg(test)]
mod tests {
    use super::DecoderBuffer;

    #[test]
    fn bit_decode_respects_declared_byte_size() {
        let data = [1, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
        let mut buffer = DecoderBuffer::new(&data);

        assert_eq!(buffer.start_bit_decoding(true).unwrap(), 1);
        assert!(buffer.decode_least_significant_bits32(16).is_err());
    }

    #[test]
    fn try_advance_rejects_out_of_bounds_skip() {
        let data = [0u8; 4];
        let mut buffer = DecoderBuffer::new(&data);

        assert!(buffer.try_advance(5).is_err());
        assert_eq!(buffer.position(), 0);
        assert!(buffer.try_advance(4).is_ok());
        assert_eq!(buffer.position(), 4);
    }

    #[test]
    fn decode_least_significant_bits32_reads_full_width() {
        // `(1u32 << 32) - 1` is the overflow this width used to hit.
        let data = [0xffu8; 5];
        let mut buffer = DecoderBuffer::new(&data);

        buffer.start_bit_decoding(false).unwrap();
        assert_eq!(
            buffer.decode_least_significant_bits32(32).unwrap(),
            u32::MAX
        );
    }

    #[test]
    fn decode_least_significant_bits32_reads_full_width_past_a_bit_shift() {
        // Same width, but starting mid-byte so the fast path's `raw >>
        // bit_shift` also has to keep all 32 bits in range.
        let data = [0xffu8; 6];
        let mut buffer = DecoderBuffer::new(&data);

        buffer.start_bit_decoding(false).unwrap();
        assert_eq!(buffer.decode_least_significant_bits32(3).unwrap(), 0b111);
        assert_eq!(
            buffer.decode_least_significant_bits32(32).unwrap(),
            u32::MAX
        );
    }

    #[test]
    fn decode_least_significant_bits32_rejects_width_above_32() {
        let data = [0xffu8; 5];
        let mut buffer = DecoderBuffer::new(&data);

        buffer.start_bit_decoding(false).unwrap();
        assert!(buffer.decode_least_significant_bits32(33).is_err());
    }
}