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}