Skip to main content

draco_core/
decoder_buffer.rs

1use crate::status::DracoError;
2use crate::version::DEFAULT_MESH_VERSION;
3use std::mem;
4
5/// Input buffer for reading compressed Draco data.
6///
7/// `DecoderBuffer` provides sequential byte and bit-level access to compressed data.
8/// It supports both byte-aligned reads (integers, floats, strings) and bit-level
9/// reads for entropy-coded data.
10///
11/// # Example
12///
13/// ```
14/// use draco_core::DecoderBuffer;
15///
16/// let data = &[0x44, 0x52, 0x41, 0x43, 0x4F]; // "DRACO" header
17/// let mut buffer = DecoderBuffer::new(data);
18///
19/// assert_eq!(buffer.decode_u8().unwrap(), 0x44);
20/// assert_eq!(buffer.remaining_size(), 4);
21/// ```
22pub struct DecoderBuffer<'a> {
23    data: &'a [u8],
24    pos: usize,
25    bit_decoder_active: bool,
26    bit_start_pos: usize,
27    current_bit_offset: usize,
28    bit_stream_end_pos: usize,
29    bit_sequence_size_known: bool,
30    version_major: u8,
31    version_minor: u8,
32    /// Bytes this decode has reserved so far, across every buffer it sized.
33    ///
34    /// The budget belongs to *one decode of one stream*, and that is what this
35    /// type is: built once per decode, carrying the stream, and already
36    /// threaded to every site that sizes something. Holding the counter here
37    /// is what makes the bound cumulative rather than per-allocation -- see
38    /// [`charge`](Self::charge).
39    spent: usize,
40    /// The length of the stream the budget's ratio is taken against. The whole
41    /// stream's, not this buffer's: a buffer over one attribute's bytes is still
42    /// one decode of one stream, and judging its reservations against its own
43    /// few kilobytes would refuse a constant attribute the whole file backs.
44    budget_len: usize,
45    /// The running total shared with the buffers over other pieces of the same
46    /// stream, where this is one of them: they charge one counter, so the bound
47    /// stays cumulative over the whole decode whichever piece spends.
48    shared_spent: Option<std::sync::Arc<std::sync::atomic::AtomicUsize>>,
49    /// The caller's ceilings on what this decode may produce, and what it has
50    /// produced so far against the byte one.
51    ///
52    /// Here rather than on the decoders because this type is the stream and
53    /// already reaches every site that learns a count -- the same reason the
54    /// budget above lives here. A decode of a mesh runs a point-cloud decoder
55    /// and an attributes decoder under it; one buffer is what all three share.
56    limits: crate::decode_limits::DecodeLimits,
57    decoded_bytes: u64,
58    /// Attribute values the caller's limits admitted that no unbacked run has
59    /// drawn on yet -- see [`charge_unbacked`](Self::charge_unbacked).
60    admitted_values: usize,
61    /// The allowance shared with the buffers over other pieces of the same
62    /// stream while they decode beside this one, where it is in place of
63    /// `admitted_values`: one allowance for the decode, so an admitted value
64    /// excuses one unbacked value once whichever piece draws it.
65    shared_admitted: Option<std::sync::Arc<std::sync::atomic::AtomicUsize>>,
66}
67
68/// A decode's standing against its budget and limits, as
69/// [`DecoderBuffer::budget`] takes it.
70#[cfg(feature = "point_cloud_decode")]
71#[derive(Clone, Copy, Debug, PartialEq, Eq)]
72pub(crate) struct Budget {
73    spent: usize,
74    decoded_bytes: u64,
75    admitted_values: usize,
76}
77
78/// The version, ceilings and budget of one decode, set aside to open buffers over
79/// pieces of its stream. See [`DecoderBuffer::child_template`].
80#[cfg(feature = "point_cloud_decode")]
81pub(crate) struct ChildTemplate {
82    version_major: u8,
83    version_minor: u8,
84    budget_len: usize,
85    limits: crate::decode_limits::DecodeLimits,
86    decoded_bytes: u64,
87    /// What the stream had spent when this was taken, and what it and every
88    /// piece opened from this have spent since.
89    spent: std::sync::Arc<std::sync::atomic::AtomicUsize>,
90    spent_at_start: usize,
91    /// The stream's admitted-value allowance while the pieces are out, and what
92    /// it held when this was taken.
93    admitted: std::sync::Arc<std::sync::atomic::AtomicUsize>,
94    admitted_at_start: usize,
95}
96
97#[cfg(feature = "point_cloud_decode")]
98impl ChildTemplate {
99    /// A buffer over `data`, a stretch of the stream this was taken from. It
100    /// charges the counter every other such buffer, and the stream itself,
101    /// charges, so what the pieces spend side by side is bounded the way what
102    /// one buffer spends in order is.
103    pub(crate) fn open<'a>(&self, data: &'a [u8]) -> DecoderBuffer<'a> {
104        let mut buffer = DecoderBuffer::new(data).with_limits(self.limits);
105        buffer.version_major = self.version_major;
106        buffer.version_minor = self.version_minor;
107        buffer.spent = self.spent_at_start;
108        buffer.budget_len = self.budget_len;
109        buffer.decoded_bytes = self.decoded_bytes;
110        buffer.shared_spent = Some(self.spent.clone());
111        buffer.shared_admitted = Some(self.admitted.clone());
112        buffer
113    }
114
115    /// What the pieces opened from this, and the stream beside them, have
116    /// spent since it was taken.
117    #[cfg(test)]
118    pub(crate) fn spent_since(&self) -> usize {
119        self.spent
120            .load(std::sync::atomic::Ordering::Relaxed)
121            .saturating_sub(self.spent_at_start)
122    }
123}
124
125impl<'a> DecoderBuffer<'a> {
126    /// Creates a new `DecoderBuffer` from a byte slice.
127    pub fn new(data: &'a [u8]) -> Self {
128        Self {
129            data,
130            pos: 0,
131            bit_decoder_active: false,
132            bit_start_pos: 0,
133            current_bit_offset: 0,
134            bit_stream_end_pos: 0,
135            bit_sequence_size_known: false,
136            // Default to latest mesh version to match encoder output format
137            version_major: DEFAULT_MESH_VERSION.0,
138            version_minor: DEFAULT_MESH_VERSION.1,
139            spent: 0,
140            budget_len: data.len(),
141            shared_spent: None,
142            limits: crate::decode_limits::DecodeLimits::default(),
143            decoded_bytes: 0,
144            admitted_values: 0,
145            shared_admitted: None,
146        }
147    }
148
149    /// What a buffer over a stretch of this one's stream needs from it, taken
150    /// now so that the stretch can be opened later, on another thread, while this
151    /// buffer goes on being read.
152    ///
153    /// This buffer's admitted-value allowance and what it has spent move into
154    /// counters the pieces share, and this buffer draws on and charges them too
155    /// until [`rejoin`](Self::rejoin): the pieces and whatever this buffer
156    /// decodes in the meantime are one decode, every charge checked against
157    /// what all of them have spent.
158    #[cfg(feature = "point_cloud_decode")]
159    pub(crate) fn child_template(&mut self) -> ChildTemplate {
160        let admitted =
161            std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(self.admitted_values));
162        let spent = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(self.spent));
163        self.shared_admitted = Some(admitted.clone());
164        self.shared_spent = Some(spent.clone());
165        ChildTemplate {
166            version_major: self.version_major,
167            version_minor: self.version_minor,
168            budget_len: self.budget_len,
169            limits: self.limits,
170            decoded_bytes: self.decoded_bytes,
171            spent,
172            spent_at_start: self.spent,
173            admitted,
174            admitted_at_start: self.admitted_values,
175        }
176    }
177
178    /// Takes the allowance and the budget back from the pieces opened from
179    /// `template`: as they and this buffer left them where their work is kept,
180    /// and as they stood when it was taken where it is thrown away to be
181    /// decoded again in order.
182    #[cfg(feature = "point_cloud_decode")]
183    pub(crate) fn rejoin(&mut self, template: &ChildTemplate, kept: bool) {
184        use std::sync::atomic::Ordering;
185        (self.admitted_values, self.spent) = if kept {
186            (
187                template.admitted.load(Ordering::Relaxed),
188                template.spent.load(Ordering::Relaxed),
189            )
190        } else {
191            (template.admitted_at_start, template.spent_at_start)
192        };
193        self.shared_admitted = None;
194        self.shared_spent = None;
195    }
196
197    /// Decodes under `limits` rather than under
198    /// [`DecodeLimits::default`](crate::DecodeLimits::default).
199    ///
200    /// ```
201    /// use draco_core::{DecodeLimits, DecoderBuffer};
202    ///
203    /// let stream = [0u8; 0];
204    /// let buffer = DecoderBuffer::new(&stream).with_limits(DecodeLimits::permissive());
205    /// # let _ = buffer;
206    /// ```
207    #[must_use]
208    pub fn with_limits(mut self, limits: crate::decode_limits::DecodeLimits) -> Self {
209        self.limits = limits;
210        self
211    }
212
213    /// Refuses a decoded point count over the caller's ceiling.
214    #[cfg(feature = "point_cloud_decode")]
215    pub(crate) fn check_points(&self, points: usize) -> crate::status::Status {
216        self.limits.check_points(points as u64)
217    }
218
219    /// Refuses a decoded face count over the caller's ceiling.
220    pub(crate) fn check_faces(&self, faces: usize) -> crate::status::Status {
221        self.limits.check_faces(faces as u64)
222    }
223
224    /// Adds one attribute's declared size to this decode's running total and
225    /// refuses when the total passes the caller's ceiling.
226    ///
227    /// Cumulative for the same reason [`charge`](Self::charge) is: the
228    /// attribute count is the header's, so a per-attribute check bounds
229    /// nothing. Called where the size becomes known and before anything is
230    /// allocated for it.
231    ///
232    /// An admitted attribute's values are also credited to the allowance
233    /// [`charge_unbacked`](Self::charge_unbacked) draws on, since the caller
234    /// has now said a decode this large is one it is willing to make.
235    pub(crate) fn admit_attribute(
236        &mut self,
237        num_points: usize,
238        num_components: usize,
239        component_bytes: usize,
240    ) -> crate::status::Status {
241        let num_values = num_components.saturating_mul(num_points);
242        let total = self
243            .decoded_bytes
244            .saturating_add(num_values.saturating_mul(component_bytes) as u64);
245        self.limits.check_decoded_bytes(total)?;
246        self.decoded_bytes = total;
247        self.admitted_values = self.admitted_values.saturating_add(num_values);
248        Ok(())
249    }
250
251    /// Charges `bytes` against what this stream is allowed to make the decoder
252    /// allocate, and refuses when the running total outgrows it.
253    ///
254    /// Cumulative on purpose. Checked per allocation, the same budget is
255    /// granted again at every site, and the number of sites is the attacker's:
256    /// one attributes decoder carries as many attributes as the header names,
257    /// and their buffers coexist. Measured before this existed, on a stream
258    /// held at ~32 KB while only the attribute count changed:
259    ///
260    /// | attributes | reserved |
261    /// | ---: | ---: |
262    /// | 1 | 19.2 MB |
263    /// | 4 | 76.8 MB |
264    /// | 16 | 307.2 MB |
265    ///
266    /// Linear in a count that costs 7.5 bytes of input each, so the per-site
267    /// form bounded nothing that mattered. One running total does.
268    ///
269    /// Nothing is ever credited back. A decode that frees a buffer and sizes
270    /// another still pays for both, which is why the ceiling has to cover the
271    /// *sum* over a real file rather than its peak -- see
272    /// [`MAX_ALLOCATED_BYTES_PER_INPUT_BYTE`], whose value is measured against
273    /// exactly that.
274    ///
275    /// [`MAX_ALLOCATED_BYTES_PER_INPUT_BYTE`]: crate::decode_budget::MAX_ALLOCATED_BYTES_PER_INPUT_BYTE
276    pub(crate) fn charge(&mut self, bytes: usize) -> crate::status::Status {
277        let total = match &self.shared_spent {
278            Some(shared) => shared
279                .fetch_add(bytes, std::sync::atomic::Ordering::Relaxed)
280                .saturating_add(bytes),
281            None => self.spent.saturating_add(bytes),
282        };
283        crate::decode_budget::ensure_allocation_is_backed(total, self.budget_len)?;
284        self.spent = total;
285        Ok(())
286    }
287
288    /// [`charge`](Self::charge) for a count of `element_size`-byte elements.
289    pub(crate) fn charge_elements(
290        &mut self,
291        count: usize,
292        element_size: usize,
293    ) -> crate::status::Status {
294        self.charge(count.saturating_mul(element_size))
295    }
296
297    /// Charges `count` values of `element_size` bytes that nothing in the
298    /// stream backs, drawing first on the attribute values the caller's limits
299    /// admitted and charging only the rest to [`charge`](Self::charge).
300    ///
301    /// The case this exists for is a run over an alphabet of one -- or raw
302    /// corrections declared zero bytes wide -- which is what a constant
303    /// attribute codes to. It carries no payload, so its count is unbacked by
304    /// construction, and the stream is as short for a billion values as for a
305    /// thousand. No bound derived from the stream separates such a file from a
306    /// hostile claim, so charging it to the budget would only move the refusal
307    /// of legitimate files to a larger point count: any ceiling the budget
308    /// could carry is one real scans walk past.
309    ///
310    /// What does separate them is the caller's ceiling. These values *are* the
311    /// attribute's own, whose declared size [`admit_attribute`] already put to
312    /// [`DecodeLimits`], so they are bounded there -- with the defaults, a
313    /// stream naming billions of values is refused at the header before
314    /// anything is decoded. Charging them here as well would count one
315    /// attribute twice, against two ceilings, the second of which only exists
316    /// to catch what the first cannot see.
317    ///
318    /// The allowance is drawn down rather than consulted, so each admitted
319    /// value excuses one unbacked value once, and the sum the budget would
320    /// have refused is bounded by the sum the limits admitted. That bound is
321    /// in values, while the limits count bytes of the declared type: the
322    /// decoder's working buffer holds four bytes per value, so a one-byte
323    /// attribute is decoded through up to four times the bytes the caller
324    /// admitted. That is what decoding it costs whether or not it is constant;
325    /// a caller whose limits are [`DecodeLimits::permissive`] has said it
326    /// trusts its input, and is now taken at its word here too.
327    ///
328    /// A count with no admission ahead of it -- connectivity, which is decoded
329    /// before any attribute is declared, or a run longer than the attributes
330    /// declared -- still pays the budget in full.
331    ///
332    /// [`admit_attribute`]: Self::admit_attribute
333    /// [`DecodeLimits`]: crate::decode_limits::DecodeLimits
334    /// [`DecodeLimits::permissive`]: crate::decode_limits::DecodeLimits::permissive
335    pub(crate) fn charge_unbacked(
336        &mut self,
337        count: usize,
338        element_size: usize,
339    ) -> crate::status::Status {
340        use std::sync::atomic::Ordering;
341        let Some(shared) = &self.shared_admitted else {
342            let admitted = count.min(self.admitted_values);
343            self.charge_elements(count - admitted, element_size)?;
344            self.admitted_values -= admitted;
345            return Ok(());
346        };
347        // Drawn before the charge so no other piece can draw the same values,
348        // and handed back if the charge refuses the rest. A compare-exchange
349        // loop rather than `fetch_update`, deprecated since 1.99, or the
350        // `try_update` it became, which is past this crate's 1.88.
351        let mut left = shared.load(Ordering::Relaxed);
352        let admitted = loop {
353            let admitted = count.min(left);
354            match shared.compare_exchange(
355                left,
356                left - admitted,
357                Ordering::Relaxed,
358                Ordering::Relaxed,
359            ) {
360                Ok(_) => break admitted,
361                Err(now) => left = now,
362            }
363        };
364        let shared = shared.clone();
365        self.charge_elements(count - admitted, element_size)
366            .inspect_err(|_| {
367                shared.fetch_add(admitted, Ordering::Relaxed);
368            })
369    }
370
371    /// Where this decode stands against its budget and its limits, for a
372    /// decode that tries a faster way and gives it up to put back as it was.
373    /// A buffer read on one thread; one whose pieces share their counters is
374    /// not one this is taken of.
375    #[cfg(feature = "point_cloud_decode")]
376    pub(crate) fn budget(&self) -> Budget {
377        debug_assert!(self.shared_spent.is_none() && self.shared_admitted.is_none());
378        Budget {
379            spent: self.spent,
380            decoded_bytes: self.decoded_bytes,
381            admitted_values: self.admitted_values,
382        }
383    }
384
385    /// Puts this decode's budget and limits back to `budget`.
386    #[cfg(feature = "point_cloud_decode")]
387    pub(crate) fn restore_budget(&mut self, budget: Budget) {
388        self.spent = budget.spent;
389        self.decoded_bytes = budget.decoded_bytes;
390        self.admitted_values = budget.admitted_values;
391    }
392
393    /// A buffer over the same stream at `pos`, carrying this one's version,
394    /// limits and the length its budget is judged against, and its budget as
395    /// it stands: what it charges is checked against everything this decode
396    /// has charged before it, and [`adopt_budget`] makes it this buffer's. For
397    /// reading a second stream of the same decode beside the first, on the
398    /// same thread.
399    ///
400    /// [`adopt_budget`]: Self::adopt_budget
401    #[cfg(feature = "point_cloud_decode")]
402    pub(crate) fn fork_at(&self, pos: usize) -> Result<DecoderBuffer<'a>, DracoError> {
403        let mut fork = DecoderBuffer::new(self.data).with_limits(self.limits);
404        fork.set_version(self.version_major, self.version_minor);
405        fork.budget_len = self.budget_len;
406        fork.restore_budget(self.budget());
407        fork.set_position(pos)?;
408        Ok(fork)
409    }
410
411    /// Takes `fork`'s budget as this decode's: what was charged through it, on
412    /// top of what this buffer had charged when it was forked.
413    #[cfg(feature = "point_cloud_decode")]
414    pub(crate) fn adopt_budget(&mut self, fork: &DecoderBuffer) {
415        self.restore_budget(fork.budget());
416    }
417
418    /// What this decode has reserved so far, for the tests that hold the
419    /// budget to being a backstop no legitimate file reaches.
420    #[cfg(test)]
421    pub(crate) fn spent(&self) -> usize {
422        self.spent
423    }
424
425    /// Sets the Draco bitstream version for version-dependent decoding.
426    pub fn set_version(&mut self, major: u8, minor: u8) {
427        self.version_major = major;
428        self.version_minor = minor;
429    }
430
431    /// Returns the major version number.
432    pub fn version_major(&self) -> u8 {
433        self.version_major
434    }
435
436    /// Returns the minor version number.
437    pub fn version_minor(&self) -> u8 {
438        self.version_minor
439    }
440
441    /// Returns the packed `0xMMmm` bitstream version for ordered comparisons.
442    pub fn bitstream_version(&self) -> u16 {
443        crate::version::bitstream_version(self.version_major, self.version_minor)
444    }
445
446    /// Returns the current read position in bytes.
447    pub fn position(&self) -> usize {
448        self.pos
449    }
450
451    /// Sets the read position.
452    ///
453    /// # Errors
454    ///
455    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if:
456    /// - Bit decoding is currently active
457    /// - Position is beyond the buffer length
458    pub fn set_position(&mut self, pos: usize) -> Result<(), DracoError> {
459        if self.bit_decoder_active {
460            return Err(DracoError::buffer(
461                "Cannot set position while bit decoding is active",
462            ));
463        }
464        if pos > self.data.len() {
465            return Err(DracoError::buffer(format!(
466                "Position {} exceeds buffer length {}",
467                pos,
468                self.data.len()
469            )));
470        }
471        self.pos = pos;
472        Ok(())
473    }
474
475    /// Returns the number of bytes remaining in the buffer.
476    pub fn remaining_size(&self) -> usize {
477        self.data.len().saturating_sub(self.pos)
478    }
479
480    /// Returns the total size of the buffer, read or not.
481    ///
482    /// The decode allocation budget measures against this rather than against
483    /// [`remaining_size`](Self::remaining_size): an attribute decoded last has
484    /// a legitimately tiny payload left in front of it while its value count is
485    /// no smaller than the first attribute's, so charging it only for what
486    /// follows would refuse valid streams.
487    pub fn size(&self) -> usize {
488        self.data.len()
489    }
490
491    /// Peeks at the next `len` bytes without advancing the position.
492    pub fn peek_bytes(&self, len: usize) -> Vec<u8> {
493        let start = self.pos.min(self.data.len());
494        let end = start.saturating_add(len).min(self.data.len());
495        self.data[start..end].to_vec()
496    }
497
498    /// Starts bit-level decoding mode.
499    ///
500    /// When `decode_size` is true, reads the bit sequence size from the buffer.
501    /// Returns the size in bytes.
502    ///
503    /// # Errors
504    ///
505    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if bit decoding is already active.
506    pub fn start_bit_decoding(&mut self, decode_size: bool) -> Result<u64, DracoError> {
507        if self.bit_decoder_active {
508            return Err(DracoError::buffer("Bit decoding already active"));
509        }
510        let bitstream_version = self.bitstream_version();
511        // Draco stores the bit-sequence size in BYTES (not bits) when |decode_size| is true.
512        let mut size_bytes: u64 = 0;
513        if decode_size {
514            if bitstream_version < 0x0202 {
515                if !cfg!(feature = "legacy_bitstream_decode") {
516                    return Err(DracoError::bitstream_version_unsupported());
517                }
518                size_bytes = self.decode_u64()?;
519            } else {
520                size_bytes = self.decode_varint()?;
521            }
522        }
523
524        self.bit_start_pos = self.pos;
525        self.bit_decoder_active = true;
526        self.current_bit_offset = 0;
527        self.bit_sequence_size_known = decode_size;
528
529        if decode_size {
530            let size_bytes = usize::try_from(size_bytes)
531                .map_err(|_| DracoError::buffer("Bit stream size too large"))?;
532            // Bounded by the buffer, not by the number: the size is a varint
533            // out of the stream, and a stream claiming more than it carries is
534            // one upstream reads to the end of rather than refuses -- its bit
535            // decoder is handed `remaining_size()` and never sees the claim.
536            // Left unbounded, the claim lands in `pos` when bit decoding ends,
537            // and a position past the buffer is one no later bounds check is
538            // written to expect.
539            let declared_end = self
540                .bit_start_pos
541                .checked_add(size_bytes)
542                .ok_or_else(|| DracoError::buffer("Bit stream end position overflow"))?;
543            self.bit_stream_end_pos = declared_end.min(self.data.len());
544        } else {
545            // If size is not encoded, assume the rest of the buffer.
546            self.bit_stream_end_pos = self.data.len();
547        }
548
549        Ok(size_bytes)
550    }
551
552    /// Ends bit-level decoding mode and advances the byte position.
553    pub fn end_bit_decoding(&mut self) {
554        self.bit_decoder_active = false;
555        // Draco behavior:
556        // - When decoding with size known, the caller typically skips by the stored byte size.
557        // - When decoding without size, advance by the number of decoded bits (rounded up).
558        if self.bit_sequence_size_known {
559            self.pos = self.bit_stream_end_pos;
560        } else {
561            let bytes_consumed = self.current_bit_offset.div_ceil(8);
562            self.pos = self.bit_start_pos + bytes_consumed;
563        }
564    }
565
566    /// Decodes `nbits` least significant bits as a u32.
567    ///
568    /// # Errors
569    ///
570    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if bit decoding is not active or end of stream.
571    #[inline(always)]
572    pub fn decode_least_significant_bits32(&mut self, nbits: u32) -> Result<u32, DracoError> {
573        if !self.bit_decoder_active {
574            return Err(DracoError::buffer("Bit decoding not active"));
575        }
576        self.decode_least_significant_bits32_fast(nbits)
577    }
578
579    /// Optimized version for hot paths - reads multiple bytes at once.
580    #[inline(always)]
581    pub fn decode_least_significant_bits32_fast(&mut self, nbits: u32) -> Result<u32, DracoError> {
582        if nbits == 0 {
583            return Ok(0);
584        }
585        // The tagged symbol scheme's per-chunk bit length is a full byte
586        // (0..=255) read straight off the wire and checked only against
587        // `> 32`, so 32 itself reaches here as a legitimate width -- a
588        // 32-bit-wide residual is what the scheme exists for. `nbits` past
589        // 32 is not a width this format defines.
590        if nbits > 32 {
591            return Err(DracoError::buffer("Bit width exceeds 32 bits"));
592        }
593
594        let total_bit_offset = self.current_bit_offset;
595        let byte_offset = self.bit_start_pos + total_bit_offset / 8;
596        let bit_shift = (total_bit_offset % 8) as u32;
597
598        if byte_offset >= self.bit_stream_end_pos || byte_offset >= self.data.len() {
599            return Err(DracoError::buffer("Unexpected end of bit stream"));
600        }
601        let available_end = self.bit_stream_end_pos.min(self.data.len());
602        let remaining = available_end - byte_offset;
603
604        // Fast path: read 8 bytes at once when enough data remains (avoids per-byte loop).
605        let raw = if remaining >= 8 {
606            let mut bytes = [0u8; 8];
607            bytes.copy_from_slice(&self.data[byte_offset..byte_offset + 8]);
608            u64::from_le_bytes(bytes)
609        } else {
610            let needed_bytes = (bit_shift + nbits).div_ceil(8) as usize;
611            if remaining < needed_bytes {
612                return Err(DracoError::buffer("Unexpected end of bit stream"));
613            }
614            let mut v = 0u64;
615            for i in 0..needed_bytes {
616                v |= (self.data[byte_offset + i] as u64) << (i * 8);
617            }
618            v
619        };
620        // Masking with a u64 keeps `nbits == 32` in range for the shift;
621        // `1u32 << 32` is exactly the overflow this guards against.
622        let mask = (1u64 << nbits) - 1;
623        let value = ((raw >> bit_shift) & mask) as u32;
624
625        self.current_bit_offset += nbits as usize;
626        Ok(value)
627    }
628
629    #[inline]
630    #[allow(dead_code)]
631    fn get_bit(&mut self) -> Result<u32, DracoError> {
632        let total_bit_offset = self.current_bit_offset;
633        let byte_offset = self.bit_start_pos + total_bit_offset / 8;
634        let bit_shift = total_bit_offset % 8;
635
636        if byte_offset < self.bit_stream_end_pos && byte_offset < self.data.len() {
637            let bit = (self.data[byte_offset] >> bit_shift) & 1;
638            self.current_bit_offset += 1;
639            Ok(bit as u32)
640        } else {
641            Err(DracoError::buffer("Unexpected end of bit stream"))
642        }
643    }
644
645    /// Decodes a value of type T using raw memory copy.
646    ///
647    /// # Errors
648    ///
649    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if:
650    /// - Bit decoding is active
651    /// - Not enough bytes remaining
652    pub fn decode<T: Copy + bytemuck::Pod>(&mut self) -> Result<T, DracoError> {
653        if self.bit_decoder_active {
654            return Err(DracoError::buffer(
655                "Cannot decode bytes while bit decoding is active",
656            ));
657        }
658        let size = mem::size_of::<T>();
659        if size > self.data.len().saturating_sub(self.pos) {
660            return Err(DracoError::buffer(format!(
661                "Unexpected end of buffer: need {} bytes, have {}",
662                size,
663                self.remaining_size()
664            )));
665        }
666
667        // Safety: bytemuck::Pod guarantees T can be safely read from any bit pattern
668        let val = bytemuck::pod_read_unaligned::<T>(&self.data[self.pos..self.pos + size]);
669        self.pos += size;
670        Ok(val)
671    }
672
673    /// Decodes a single byte.
674    pub fn decode_u8(&mut self) -> Result<u8, DracoError> {
675        self.decode::<u8>()
676    }
677
678    /// Decodes a little-endian u16.
679    pub fn decode_u16(&mut self) -> Result<u16, DracoError> {
680        let mut bytes = [0u8; 2];
681        self.decode_bytes(&mut bytes)?;
682        Ok(u16::from_le_bytes(bytes))
683    }
684
685    /// Decodes a little-endian u32.
686    pub fn decode_u32(&mut self) -> Result<u32, DracoError> {
687        let mut bytes = [0u8; 4];
688        self.decode_bytes(&mut bytes)?;
689        Ok(u32::from_le_bytes(bytes))
690    }
691
692    /// Decodes a little-endian u64.
693    pub fn decode_u64(&mut self) -> Result<u64, DracoError> {
694        let mut bytes = [0u8; 8];
695        self.decode_bytes(&mut bytes)?;
696        Ok(u64::from_le_bytes(bytes))
697    }
698
699    /// Decodes a little-endian f32.
700    pub fn decode_f32(&mut self) -> Result<f32, DracoError> {
701        let mut bytes = [0u8; 4];
702        self.decode_bytes(&mut bytes)?;
703        Ok(f32::from_le_bytes(bytes))
704    }
705
706    /// Decodes a little-endian f64.
707    pub fn decode_f64(&mut self) -> Result<f64, DracoError> {
708        let mut bytes = [0u8; 8];
709        self.decode_bytes(&mut bytes)?;
710        Ok(f64::from_le_bytes(bytes))
711    }
712
713    /// Decodes a null-terminated string.
714    pub fn decode_string(&mut self) -> Result<String, DracoError> {
715        let mut bytes = Vec::new();
716        loop {
717            let b = self.decode_u8()?;
718            if b == 0 {
719                break;
720            }
721            bytes.push(b);
722        }
723        String::from_utf8(bytes)
724            .map_err(|e| DracoError::buffer(format!("Invalid UTF-8 string: {}", e)))
725    }
726
727    /// Decodes bytes into the provided buffer.
728    ///
729    /// # Errors
730    ///
731    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if not enough bytes remaining.
732    pub fn decode_bytes(&mut self, out: &mut [u8]) -> Result<(), DracoError> {
733        let size = out.len();
734        if size > self.data.len().saturating_sub(self.pos) {
735            return Err(DracoError::buffer(format!(
736                "Unexpected end of buffer: need {} bytes, have {}",
737                size,
738                self.remaining_size()
739            )));
740        }
741        out.copy_from_slice(&self.data[self.pos..self.pos + size]);
742        self.pos += size;
743        Ok(())
744    }
745
746    /// Decodes a variable-length unsigned integer (varint).
747    pub fn decode_varint(&mut self) -> Result<u64, DracoError> {
748        let mut val = 0u64;
749        let mut shift = 0;
750        loop {
751            let b = self.decode_u8()?;
752            val |= ((b & 0x7F) as u64) << shift;
753            if (b & 0x80) == 0 {
754                break;
755            }
756            shift += 7;
757            if shift >= 64 {
758                return Err(DracoError::buffer("Varint exceeds 64 bits"));
759            }
760        }
761        Ok(val)
762    }
763
764    /// Decodes a Draco-compatible signed varint.
765    ///
766    /// Uses unsigned varint encoding with ConvertSymbolToSignedInt transformation.
767    pub fn decode_varint_signed_i32(&mut self) -> Result<i32, DracoError> {
768        let symbol = self.decode_varint()? as u32;
769        let is_positive = (symbol & 1) == 0;
770        let v = symbol >> 1;
771        if is_positive {
772            Ok(v as i32)
773        } else {
774            Ok(-(v as i32) - 1)
775        }
776    }
777
778    /// Returns a slice of the remaining data without advancing.
779    pub fn remaining_data(&self) -> &'a [u8] {
780        &self.data[self.pos..]
781    }
782
783    /// Advances the position by `n` bytes without reading.
784    pub fn advance(&mut self, n: usize) {
785        self.pos = self.pos.saturating_add(n).min(self.data.len());
786    }
787
788    /// Advances the position by `n` bytes without reading.
789    ///
790    /// # Errors
791    ///
792    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if the requested advance would move
793    /// beyond the end of the input buffer.
794    pub fn try_advance(&mut self, n: usize) -> Result<(), DracoError> {
795        let new_pos = self
796            .pos
797            .checked_add(n)
798            .ok_or_else(|| DracoError::buffer("Buffer advance overflow"))?;
799        if new_pos > self.data.len() {
800            return Err(DracoError::buffer(format!(
801                "Cannot advance buffer by {} bytes: need position {}, buffer length {}",
802                n,
803                new_pos,
804                self.data.len()
805            )));
806        }
807        self.pos = new_pos;
808        Ok(())
809    }
810
811    /// Decodes and returns a slice of the specified size.
812    ///
813    /// # Errors
814    ///
815    /// Returns an [`ErrorKind::Buffer`](crate::ErrorKind::Buffer) error if not enough bytes remaining.
816    pub fn decode_slice(&mut self, size: usize) -> Result<&'a [u8], DracoError> {
817        if size > self.data.len().saturating_sub(self.pos) {
818            return Err(DracoError::buffer(format!(
819                "Unexpected end of buffer: need {} bytes, have {}",
820                size,
821                self.remaining_size()
822            )));
823        }
824        let slice = &self.data[self.pos..self.pos + size];
825        self.pos += size;
826        Ok(slice)
827    }
828}
829
830#[cfg(test)]
831mod tests {
832    use super::DecoderBuffer;
833
834    #[test]
835    fn bit_decode_respects_declared_byte_size() {
836        let data = [1, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
837        let mut buffer = DecoderBuffer::new(&data);
838
839        assert_eq!(buffer.start_bit_decoding(true).unwrap(), 1);
840        assert!(buffer.decode_least_significant_bits32(16).is_err());
841    }
842
843    #[test]
844    fn try_advance_rejects_out_of_bounds_skip() {
845        let data = [0u8; 4];
846        let mut buffer = DecoderBuffer::new(&data);
847
848        assert!(buffer.try_advance(5).is_err());
849        assert_eq!(buffer.position(), 0);
850        assert!(buffer.try_advance(4).is_ok());
851        assert_eq!(buffer.position(), 4);
852    }
853
854    #[test]
855    fn decode_least_significant_bits32_reads_full_width() {
856        // `(1u32 << 32) - 1` is the overflow this width used to hit.
857        let data = [0xffu8; 5];
858        let mut buffer = DecoderBuffer::new(&data);
859
860        buffer.start_bit_decoding(false).unwrap();
861        assert_eq!(
862            buffer.decode_least_significant_bits32(32).unwrap(),
863            u32::MAX
864        );
865    }
866
867    #[test]
868    fn decode_least_significant_bits32_reads_full_width_past_a_bit_shift() {
869        // Same width, but starting mid-byte so the fast path's `raw >>
870        // bit_shift` also has to keep all 32 bits in range.
871        let data = [0xffu8; 6];
872        let mut buffer = DecoderBuffer::new(&data);
873
874        buffer.start_bit_decoding(false).unwrap();
875        assert_eq!(buffer.decode_least_significant_bits32(3).unwrap(), 0b111);
876        assert_eq!(
877            buffer.decode_least_significant_bits32(32).unwrap(),
878            u32::MAX
879        );
880    }
881
882    #[test]
883    fn decode_least_significant_bits32_rejects_width_above_32() {
884        let data = [0xffu8; 5];
885        let mut buffer = DecoderBuffer::new(&data);
886
887        buffer.start_bit_decoding(false).unwrap();
888        assert!(buffer.decode_least_significant_bits32(33).is_err());
889    }
890
891    /// The count a constant RGB colour over 29,351,662 points codes to: past
892    /// what the budget alone permits, from a stream of any length.
893    const CONSTANT_COLOUR_VALUES: usize = 29_351_662 * 3;
894
895    #[test]
896    fn an_unbacked_run_the_limits_admitted_charges_nothing() {
897        let data = [0u8; 113];
898        let mut buffer = DecoderBuffer::new(&data);
899        assert!(buffer
900            .charge_unbacked(CONSTANT_COLOUR_VALUES, size_of::<u32>())
901            .is_err());
902
903        let mut buffer = DecoderBuffer::new(&data);
904        buffer
905            .admit_attribute(CONSTANT_COLOUR_VALUES / 3, 3, 1)
906            .unwrap();
907        buffer
908            .charge_unbacked(CONSTANT_COLOUR_VALUES, size_of::<u32>())
909            .unwrap();
910        assert_eq!(buffer.spent(), 0);
911    }
912
913    /// Each admitted value excuses one unbacked value once: a second run over
914    /// the same admission pays the budget, so the sum the budget lets through
915    /// stays bounded by the sum the limits admitted.
916    #[test]
917    fn an_admission_is_drawn_down_rather_than_reused() {
918        let data = [0u8; 113];
919        let mut buffer = DecoderBuffer::new(&data);
920        buffer.admit_attribute(1_000, 1, 1).unwrap();
921        buffer.charge_unbacked(1_000, size_of::<u32>()).unwrap();
922        assert_eq!(buffer.spent(), 0);
923        buffer.charge_unbacked(10, size_of::<u32>()).unwrap();
924        assert_eq!(buffer.spent(), 40);
925        assert!(buffer
926            .charge_unbacked(CONSTANT_COLOUR_VALUES, size_of::<u32>())
927            .is_err());
928    }
929
930    /// The protection the budget exists for is kept, by the limits: the
931    /// header of a tiny stream naming billions of values is refused before
932    /// anything is admitted, so nothing is there to draw on.
933    #[test]
934    fn a_claim_past_the_default_limits_is_refused_at_admission() {
935        let data = [0u8; 170];
936        let mut buffer = DecoderBuffer::new(&data);
937        let error = buffer
938            .admit_attribute(2_270_000_000, 3, 1)
939            .expect_err("6.8 billion values admitted under the defaults");
940        assert_eq!(error.kind(), crate::status::ErrorKind::LimitExceeded);
941        assert!(buffer
942            .charge_unbacked(u32::MAX as usize, size_of::<u32>())
943            .is_err());
944    }
945
946    /// Pieces of one stream decoded side by side draw on one allowance: what
947    /// two of them draw together is what one buffer could, the rest pays the
948    /// budget they share, and the stream takes back what is left -- or all of
949    /// it, where their work is thrown away to be decoded again in order.
950    #[cfg(feature = "point_cloud_decode")]
951    #[test]
952    fn pieces_decoded_side_by_side_share_one_allowance() {
953        let data = [0u8; 113];
954        let mut stream = DecoderBuffer::new(&data);
955        stream.admit_attribute(1_000, 1, 1).unwrap();
956        let template = stream.child_template();
957        let mut a = template.open(&data);
958        let mut b = template.open(&data);
959        a.charge_unbacked(600, size_of::<u32>()).unwrap();
960        b.charge_unbacked(600, size_of::<u32>()).unwrap();
961        stream.charge_unbacked(100, size_of::<u32>()).unwrap();
962        assert_eq!(template.spent_since(), 300 * size_of::<u32>());
963
964        stream.rejoin(&template, true);
965        assert_eq!(stream.admitted_values, 0);
966        assert_eq!(stream.spent(), 300 * size_of::<u32>());
967        stream.rejoin(&template, false);
968        assert_eq!(stream.admitted_values, 1_000);
969        assert_eq!(stream.spent(), 0);
970    }
971
972    /// What the stream charges while its pieces are out counts against the
973    /// ceiling they charge: two charges, each under it alone, are refused
974    /// together whichever of them made it.
975    #[cfg(feature = "point_cloud_decode")]
976    #[test]
977    fn the_stream_and_its_pieces_charge_one_budget() {
978        let data = [0u8; 113];
979        let two_thirds = (113 << 20) * 2 / 3;
980        for stream_first in [true, false] {
981            let mut stream = DecoderBuffer::new(&data);
982            let template = stream.child_template();
983            let mut piece = template.open(&data);
984            let (first, second) = if stream_first {
985                (stream.charge(two_thirds), piece.charge(two_thirds))
986            } else {
987                (piece.charge(two_thirds), stream.charge(two_thirds))
988            };
989            assert!(
990                first.is_ok() && second.is_err(),
991                "stream first: {stream_first}"
992            );
993        }
994    }
995}