Skip to main content

structured_zstd/decoding/
workspace.rs

1//! Streaming decoder that runs entirely in caller-owned memory.
2//!
3//! [`WorkspaceDecoder`] decodes a stream of zstd frames chunk by chunk without
4//! touching the heap: the window, the entropy tables, the literals and the
5//! staging for blocks split across input chunks all live in one byte slice the
6//! caller supplies. [`workspace_bytes`] says how large that slice must be for a
7//! given window limit, so the caller can admit the memory in its own accounting
8//! before decoding starts, and keep the charge for as long as the slice lives.
9//!
10//! Every step is bounded three ways, all chosen by the caller: the window a
11//! frame may declare and the total output ([`DecodeLimits`]), and the work a
12//! step may do ([`WorkBudget`]). A refusal is reported apart from malformed
13//! input, an unsupported feature, and input that ends mid-frame, and it ends
14//! the stream: every later step fails with [`WorkspaceDecodeError::Failed`]
15//! until [`WorkspaceDecoder::reset`].
16//!
17//! Frames that name a dictionary are refused with
18//! [`Unsupported::Dictionary`]; this decoder keeps no dictionary registry.
19//!
20//! The decoding itself is the same code the other decoders run, so output is
21//! identical to theirs.
22//!
23//! # Examples
24//!
25//! ```
26//! use structured_zstd::decoding::workspace::{
27//!     workspace_bytes, DecodeLimits, StepState, WorkBudget, WorkspaceDecoder,
28//! };
29//! use structured_zstd::encoding::{compress_slice_to_vec, CompressionLevel};
30//!
31//! let data = b"the same words, the same words, the same words".repeat(100);
32//! let compressed = compress_slice_to_vec(&data, CompressionLevel::Default);
33//!
34//! let limits = DecodeLimits::new(1 << 20).with_max_output(1 << 20);
35//! let mut workspace = vec![0u8; workspace_bytes(limits.max_window_size).unwrap()];
36//! let mut decoder = WorkspaceDecoder::new(&mut workspace, limits).unwrap();
37//! let mut work = WorkBudget::new(1 << 24);
38//!
39//! let mut decoded = Vec::new();
40//! let mut out = [0u8; 512];
41//! let mut input = &compressed[..];
42//! loop {
43//!     // Feed at most 100 bytes at a time; say when nothing more will come.
44//!     let chunk = &input[..input.len().min(100)];
45//!     let last = chunk.len() == input.len();
46//!     let progress = decoder.step(chunk, last, &mut out, &mut work).unwrap();
47//!     decoded.extend_from_slice(&out[..progress.produced]);
48//!     input = &input[progress.consumed..];
49//!     if progress.state == StepState::Done {
50//!         break;
51//!     }
52//! }
53//! assert_eq!(decoded, data);
54//! ```
55
56use super::block_decoder::{self, BlockDecoder};
57use super::buffer_backend::{BufferBackend, WILDCOPY_OVERLENGTH};
58use super::decode_buffer::DecodeBuffer;
59use super::errors::{
60    BlockHeaderReadError, DecodeBlockContentError, FrameHeaderError, ReadFrameHeaderError,
61};
62use super::frame;
63use super::ringbuffer::RingBuffer;
64use super::scratch::{ByteScratch, FSEScratch, HuffmanScratch, Workspace, WorkspaceRef};
65use super::user_slice_buf::UserSliceBackend;
66use crate::blocks::block::{BlockHeader, BlockType};
67use crate::common::{MAGIC_NUM, MAX_DECODER_WINDOW_SIZE, MIN_WINDOW_SIZE};
68use core::marker::PhantomData;
69use core::ptr::NonNull;
70
71/// Work units charged before a compressed block decodes, on top of its input
72/// bytes and its output bound: the most its entropy-table builds can cost
73/// (one Huffman table of `1 << 12` entries and the literal-length,
74/// match-length and offset tables at `1 << 9`, `1 << 9` and `1 << 8`).
75pub const COMPRESSED_BLOCK_UNITS: u64 = (1 << crate::huff0::MAX_MAX_NUM_BITS)
76    + (1 << super::sequence_section_decoder::LL_MAX_LOG)
77    + (1 << super::sequence_section_decoder::ML_MAX_LOG)
78    + (1 << super::sequence_section_decoder::OF_MAX_LOG);
79
80/// Largest frame header: magic, descriptor, window byte, 4-byte dictionary id
81/// and 8-byte content size (RFC 8878 3.1.1.1).
82const MAX_FRAME_HEADER: usize = 18;
83
84/// Alignment the entropy tables need at the start of the workspace.
85const TABLE_ALIGN: usize = {
86    let fse = core::mem::align_of::<FSEScratch>();
87    let huf = core::mem::align_of::<HuffmanScratch>();
88    if fse > huf { fse } else { huf }
89};
90
91/// Bytes of workspace a [`WorkspaceDecoder`] needs to decode frames whose
92/// window is at most `max_window_size`.
93///
94/// The figure covers every byte the decoder writes: the window with room for
95/// the block being decoded, the entropy tables, the literals, the staging for a
96/// block split across input chunks, and alignment slack, so any sub-slice of
97/// that length works whatever its address.
98///
99/// # Errors
100///
101/// [`WorkspaceConfigError::WindowLimitOutOfRange`] when `max_window_size` lies
102/// outside 1 KiB ..= [`MAX_DECODER_WINDOW_SIZE`],
103/// and [`WorkspaceConfigError::SizeOverflow`] when the size does not fit in
104/// `usize` on this target.
105///
106/// # Examples
107///
108/// ```
109/// use structured_zstd::decoding::workspace::workspace_bytes;
110/// let small = workspace_bytes(64 * 1024).unwrap();
111/// let large = workspace_bytes(8 << 20).unwrap();
112/// assert!(small < large);
113/// assert!(workspace_bytes(100).is_err());
114/// ```
115pub fn workspace_bytes(max_window_size: u64) -> Result<usize, WorkspaceConfigError> {
116    Layout::for_window(max_window_size).map(|layout| layout.total)
117}
118
119/// Where each part of the workspace sits, relative to its aligned start.
120struct Layout {
121    huf_offset: usize,
122    tables_len: usize,
123    ring_cap: usize,
124    literals_len: usize,
125    staging_len: usize,
126    total: usize,
127}
128
129impl Layout {
130    fn for_window(max_window_size: u64) -> Result<Layout, WorkspaceConfigError> {
131        use WorkspaceConfigError::SizeOverflow;
132        if !(MIN_WINDOW_SIZE..=MAX_DECODER_WINDOW_SIZE).contains(&max_window_size) {
133            return Err(WorkspaceConfigError::WindowLimitOutOfRange {
134                requested: max_window_size,
135            });
136        }
137        let window = usize::try_from(max_window_size).map_err(|_| SizeOverflow)?;
138        let block_max = block_decoder::block_maximum(window);
139        let huf_offset = core::mem::size_of::<FSEScratch>()
140            .next_multiple_of(core::mem::align_of::<HuffmanScratch>());
141        let tables_len = huf_offset + core::mem::size_of::<HuffmanScratch>();
142        // The window, plus two block maxima: a block's matches stop at one
143        // block past the window (the per-block ceiling), and literals copied
144        // after its last match can add up to one more before the block-size
145        // check rejects it. One more slot is the ring's empty/full sentinel.
146        let ring_cap = window
147            .checked_add(2 * block_max)
148            .and_then(|n| n.checked_add(1))
149            .ok_or(SizeOverflow)?;
150        let literals_len = block_max + WILDCOPY_OVERLENGTH;
151        let staging_len = block_max;
152        let total = [
153            tables_len,
154            ring_cap,
155            WILDCOPY_OVERLENGTH,
156            literals_len,
157            staging_len,
158        ]
159        .iter()
160        .try_fold(TABLE_ALIGN - 1, |sum, &part| sum.checked_add(part))
161        .ok_or(SizeOverflow)?;
162        Ok(Layout {
163            huf_offset,
164            tables_len,
165            ring_cap,
166            literals_len,
167            staging_len,
168            total,
169        })
170    }
171}
172
173/// Limits a [`WorkspaceDecoder`] holds every frame of the stream to.
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub struct DecodeLimits {
176    /// Largest window a frame may declare. Also sizes the workspace; see
177    /// [`workspace_bytes`].
178    pub max_window_size: u64,
179    /// Most decompressed bytes the whole stream may produce.
180    pub max_output: u64,
181}
182
183impl DecodeLimits {
184    /// Limits with the given window ceiling and no output limit.
185    pub const fn new(max_window_size: u64) -> Self {
186        Self {
187            max_window_size,
188            max_output: u64::MAX,
189        }
190    }
191
192    /// The same limits with the total output capped at `max_output` bytes.
193    pub const fn with_max_output(self, max_output: u64) -> Self {
194        Self { max_output, ..self }
195    }
196}
197
198/// Work a caller allows its decode steps, in units charged before the work
199/// they pay for runs.
200///
201/// One unit is charged per input byte consumed (frame and block headers,
202/// block contents, skippable frames, checksums), and one per byte a block may
203/// produce, before it is produced: a Raw block's bytes as each part of it is
204/// copied, an RLE block's stated size, and a compressed block's frame block
205/// maximum plus [`COMPRESSED_BLOCK_UNITS`] for its table builds. A charge the
206/// budget cannot cover refuses the step before that work starts.
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
208pub struct WorkBudget {
209    remaining: u64,
210}
211
212impl WorkBudget {
213    /// A budget of `units`.
214    pub const fn new(units: u64) -> Self {
215        Self { remaining: units }
216    }
217
218    /// Units still available.
219    pub const fn remaining(&self) -> u64 {
220        self.remaining
221    }
222
223    /// Add `units` to the budget, stopping at `u64::MAX`, which is already
224    /// more work than any stream can ask for.
225    pub fn add(&mut self, units: u64) {
226        self.remaining = self.remaining.saturating_add(units);
227    }
228
229    fn charge(&mut self, units: u64) -> Result<(), WorkspaceDecodeError> {
230        match self.remaining.checked_sub(units) {
231            Some(left) => {
232                self.remaining = left;
233                Ok(())
234            }
235            None => Err(WorkspaceDecodeError::Limit(LimitExceeded::Work {
236                needed: units,
237                remaining: self.remaining,
238            })),
239        }
240    }
241}
242
243/// What one [`WorkspaceDecoder::step`] did.
244#[derive(Debug, Clone, Copy, PartialEq, Eq)]
245pub struct Progress {
246    /// Input bytes consumed; the next step starts after them.
247    pub consumed: usize,
248    /// Bytes written to the front of the output slice.
249    pub produced: usize,
250    /// What the decoder needs next.
251    pub state: StepState,
252}
253
254/// Why a [`WorkspaceDecoder::step`] returned.
255#[derive(Debug, Clone, Copy, PartialEq, Eq)]
256#[non_exhaustive]
257pub enum StepState {
258    /// Every input byte was consumed; call again with more, or with the
259    /// final-input flag set once there is no more.
260    NeedInput,
261    /// The output slice is full; call again with room.
262    NeedOutput,
263    /// The input was final and every frame in it is decoded and delivered.
264    Done,
265}
266
267/// A workspace size or limit [`workspace_bytes`] or [`WorkspaceDecoder::new`]
268/// refused.
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270#[non_exhaustive]
271pub enum WorkspaceConfigError {
272    /// The window limit lies outside 1 KiB ..=
273    /// [`MAX_DECODER_WINDOW_SIZE`].
274    WindowLimitOutOfRange {
275        /// The limit asked for.
276        requested: u64,
277    },
278    /// The workspace size does not fit in `usize` on this target.
279    SizeOverflow,
280    /// The workspace slice is shorter than [`workspace_bytes`] asks for.
281    WorkspaceTooSmall {
282        /// Bytes the limits require.
283        needed: usize,
284        /// Bytes supplied.
285        provided: usize,
286    },
287}
288
289impl core::fmt::Display for WorkspaceConfigError {
290    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
291        match self {
292            Self::WindowLimitOutOfRange { requested } => write!(
293                f,
294                "window limit {requested} lies outside {MIN_WINDOW_SIZE}..={MAX_DECODER_WINDOW_SIZE}"
295            ),
296            Self::SizeOverflow => write!(f, "workspace size does not fit in usize"),
297            Self::WorkspaceTooSmall { needed, provided } => write!(
298                f,
299                "workspace of {provided} bytes is smaller than the {needed} the limits require"
300            ),
301        }
302    }
303}
304
305#[cfg(feature = "std")]
306impl std::error::Error for WorkspaceConfigError {}
307
308/// Why a [`WorkspaceDecoder::step`] failed. Every variant ends the stream:
309/// later steps return [`Self::Failed`] until [`WorkspaceDecoder::reset`].
310#[derive(Debug)]
311#[non_exhaustive]
312pub enum WorkspaceDecodeError {
313    /// A caller-chosen limit refused the stream.
314    Limit(LimitExceeded),
315    /// The input is not a valid zstd stream.
316    Malformed(MalformedInput),
317    /// The input is valid but needs something this decoder does not do.
318    Unsupported(Unsupported),
319    /// The input was marked final in the middle of a frame.
320    Truncated,
321    /// An earlier step failed; call [`WorkspaceDecoder::reset`] first.
322    Failed,
323}
324
325/// A limit from [`DecodeLimits`] or [`WorkBudget`] that refused the stream.
326#[derive(Debug, Clone, Copy, PartialEq, Eq)]
327#[non_exhaustive]
328pub enum LimitExceeded {
329    /// A frame declared a window above [`DecodeLimits::max_window_size`].
330    Window {
331        /// The window the frame declared.
332        requested: u64,
333        /// The limit.
334        limit: u64,
335    },
336    /// The stream would produce more than [`DecodeLimits::max_output`].
337    Output {
338        /// The limit.
339        limit: u64,
340    },
341    /// The next unit of work costs more than the budget has left.
342    Work {
343        /// Units the refused work costs.
344        needed: u64,
345        /// Units the budget had.
346        remaining: u64,
347    },
348}
349
350/// Something this decoder does not do, met in otherwise valid input.
351#[derive(Debug, Clone, Copy, PartialEq, Eq)]
352#[non_exhaustive]
353pub enum Unsupported {
354    /// The frame needs dictionary `id`, and this decoder takes none.
355    Dictionary {
356        /// The dictionary id the frame header names.
357        id: u32,
358    },
359}
360
361/// How the input breaks the zstd format.
362#[derive(Debug)]
363#[non_exhaustive]
364pub enum MalformedInput {
365    /// A frame header, or the magic number before it, is invalid.
366    FrameHeader(ReadFrameHeaderError),
367    /// The frame header declares an invalid window.
368    Window(FrameHeaderError),
369    /// A block header is invalid.
370    BlockHeader(BlockHeaderReadError),
371    /// A block states a size above its frame's block maximum
372    /// (RFC 8878 3.1.1.2.4).
373    BlockTooLarge {
374        /// The size the block header states.
375        size: usize,
376        /// The frame's block maximum.
377        maximum: usize,
378    },
379    /// A block's content does not decode.
380    Block(DecodeBlockContentError),
381    /// The frame's blocks do not add up to its declared content size.
382    ContentSize {
383        /// The size the frame header declares.
384        declared: u64,
385        /// The size its blocks produce (so far, when they overrun it).
386        produced: u64,
387    },
388    /// The content checksum does not match the decoded output.
389    Checksum {
390        /// The checksum the frame carries.
391        expected: u32,
392        /// The low 32 bits of the XXH64 of the decoded frame.
393        calculated: u32,
394    },
395}
396
397impl core::fmt::Display for WorkspaceDecodeError {
398    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
399        match self {
400            Self::Limit(limit) => match limit {
401                LimitExceeded::Window { requested, limit } => write!(
402                    f,
403                    "frame window of {requested} bytes exceeds the {limit}-byte limit"
404                ),
405                LimitExceeded::Output { limit } => {
406                    write!(f, "output exceeds the {limit}-byte limit")
407                }
408                LimitExceeded::Work { needed, remaining } => write!(
409                    f,
410                    "work of {needed} units exceeds the {remaining} remaining"
411                ),
412            },
413            Self::Malformed(malformed) => match malformed {
414                MalformedInput::FrameHeader(e) => write!(f, "invalid frame header: {e}"),
415                MalformedInput::Window(e) => write!(f, "invalid frame window: {e}"),
416                MalformedInput::BlockHeader(e) => write!(f, "invalid block header: {e}"),
417                MalformedInput::BlockTooLarge { size, maximum } => write!(
418                    f,
419                    "block of {size} bytes exceeds the frame's block maximum of {maximum}"
420                ),
421                MalformedInput::Block(e) => write!(f, "invalid block: {e}"),
422                MalformedInput::ContentSize { declared, produced } => write!(
423                    f,
424                    "frame declares {declared} bytes but its blocks produce {produced}"
425                ),
426                MalformedInput::Checksum {
427                    expected,
428                    calculated,
429                } => write!(
430                    f,
431                    "content checksum {calculated:#010x} does not match the stored {expected:#010x}"
432                ),
433            },
434            Self::Unsupported(Unsupported::Dictionary { id }) => {
435                write!(
436                    f,
437                    "frame needs dictionary {id}, which this decoder does not take"
438                )
439            }
440            Self::Truncated => write!(f, "input ended in the middle of a frame"),
441            Self::Failed => write!(f, "an earlier step failed; reset the decoder"),
442        }
443    }
444}
445
446#[cfg(feature = "std")]
447impl std::error::Error for WorkspaceDecodeError {}
448
449/// A fixed region of the workspace standing in for the `Vec` the owned
450/// decoders keep for literals.
451///
452/// A push past the end is dropped but still counted, so the literals decoder's
453/// own count check rejects the block as it would with a growable buffer. Every
454/// other write is sized by the caller below the region's length; the asserts
455/// are what keeps a broken size argument from writing outside it.
456struct SliceScratch<'w> {
457    buf: &'w mut [u8],
458    len: usize,
459}
460
461impl ByteScratch for SliceScratch<'_> {
462    #[inline(always)]
463    fn len(&self) -> usize {
464        self.len
465    }
466    #[inline(always)]
467    fn clear(&mut self) {
468        self.len = 0;
469    }
470    #[inline(always)]
471    fn truncate(&mut self, len: usize) {
472        self.len = self.len.min(len);
473    }
474    #[inline(always)]
475    fn reserve(&mut self, additional: usize) {
476        assert!(
477            self.len
478                .checked_add(additional)
479                .is_some_and(|end| end <= self.buf.len()),
480            "literals exceed the workspace region sized for a block"
481        );
482    }
483    #[inline(always)]
484    fn push(&mut self, byte: u8) {
485        if let Some(slot) = self.buf.get_mut(self.len) {
486            *slot = byte;
487        }
488        self.len += 1;
489    }
490    fn resize(&mut self, new_len: usize, value: u8) {
491        if new_len > self.len {
492            self.buf[self.len..new_len].fill(value);
493        }
494        self.len = new_len;
495    }
496    fn extend_from_slice(&mut self, bytes: &[u8]) {
497        let end = self.len + bytes.len();
498        self.buf[self.len..end].copy_from_slice(bytes);
499        self.len = end;
500    }
501    #[inline(always)]
502    fn as_mut_ptr(&mut self) -> *mut u8 {
503        self.buf.as_mut_ptr()
504    }
505    #[inline(always)]
506    unsafe fn set_len(&mut self, new_len: usize) {
507        assert!(new_len <= self.buf.len());
508        self.len = new_len;
509    }
510    #[inline(always)]
511    fn as_slice(&self) -> &[u8] {
512        &self.buf[..self.len.min(self.buf.len())]
513    }
514    #[inline(always)]
515    fn as_mut_slice(&mut self) -> &mut [u8] {
516        let len = self.len.min(self.buf.len());
517        &mut self.buf[..len]
518    }
519}
520
521/// The current frame's window over the workspace's window region, chosen per
522/// frame as `FrameDecoder` chooses its backend: a single-segment frame's whole
523/// content is its window, so it decodes linearly into the region, the shape
524/// `decode_all` writes into a caller's slice; any other frame decodes into a
525/// ring.
526enum Window<'w> {
527    /// No frame has started since construction or reset.
528    Idle,
529    Ring(DecodeBuffer<RingBuffer>),
530    Linear(DecodeBuffer<UserSliceBackend<'w>>),
531}
532
533/// Run `$body` with `$buffer` bound to the active window's buffer, or yield
534/// `$idle` when no frame has started.
535macro_rules! with_window {
536    ($window:expr, $buffer:ident => $body:expr, idle => $idle:expr) => {
537        match $window {
538            Window::Ring($buffer) => $body,
539            Window::Linear($buffer) => $body,
540            Window::Idle => $idle,
541        }
542    };
543}
544
545/// Output bytes `buffer` has produced this frame. The linear backend's inline
546/// sequence path skips the running counter, so its write cursor counts there;
547/// the ring keeps the counter, its cursor wrapping.
548fn produced_so_far<B: BufferBackend>(buffer: &DecodeBuffer<B>) -> u64 {
549    if B::INLINE_EXEC_MAINTAINS_OUTPUT_COUNTER {
550        buffer.total_output()
551    } else {
552        buffer.buffer_ref().tail() as u64
553    }
554}
555
556/// The state the shared block decoder works on for one block: the tables and
557/// literals in the workspace, and the active window's buffer.
558struct BlockScratch<'a, 'w, B: BufferBackend> {
559    huf: &'a mut HuffmanScratch,
560    fse: &'a mut FSEScratch,
561    buffer: &'a mut DecodeBuffer<B>,
562    offset_hist: &'a mut [u32; 3],
563    literals: &'a mut SliceScratch<'w>,
564    /// The Read-sourced block path's buffer; this decoder stages blocks
565    /// itself and hands them over as slices, so it stays empty.
566    block_content: &'a mut SliceScratch<'w>,
567}
568
569impl<'w, B: BufferBackend> Workspace for BlockScratch<'_, 'w, B> {
570    type Backend = B;
571    type Bytes = SliceScratch<'w>;
572    fn split(&mut self) -> WorkspaceRef<'_, B, SliceScratch<'w>> {
573        WorkspaceRef {
574            huf: &mut *self.huf,
575            fse: &mut *self.fse,
576            buffer: &mut *self.buffer,
577            offset_hist: &mut *self.offset_hist,
578            literals_buffer: &mut *self.literals,
579            block_content_buffer: &mut *self.block_content,
580        }
581    }
582}
583
584/// Decode one block whose whole content is `content`, returning the bytes it
585/// produced.
586fn run_block<B: BufferBackend>(
587    decoder: &mut BlockDecoder,
588    scratch: &mut BlockScratch<'_, '_, B>,
589    header: &BlockHeader,
590    mut content: &[u8],
591) -> Result<u64, DecodeBlockContentError> {
592    let before = produced_so_far(scratch.buffer);
593    decoder.decode_block_content_from_slice(header, scratch, None, &mut content)?;
594    Ok(produced_so_far(scratch.buffer) - before)
595}
596
597/// What the decoder reads next.
598#[derive(Clone, Copy, PartialEq, Eq, Debug)]
599enum Phase {
600    /// The 4-byte magic number of the next frame.
601    Magic,
602    /// The descriptor byte, which gives the header's length.
603    Descriptor,
604    /// The rest of the frame header.
605    FrameHeader,
606    /// The 4-byte length of a skippable frame.
607    SkippableSize,
608    /// The payload of a skippable frame.
609    Skipping { remaining: u64 },
610    /// A 3-byte block header.
611    BlockHeader,
612    /// A block's content.
613    BlockBody(PendingBlock),
614    /// The 4-byte content checksum.
615    Checksum,
616}
617
618/// A block whose header was read, carried while its content arrives.
619#[derive(Clone, Copy, PartialEq, Eq, Debug)]
620struct PendingBlock {
621    last_block: bool,
622    block_type: BlockType,
623    decompressed_size: u32,
624    content_size: u32,
625}
626
627impl PendingBlock {
628    fn header(self) -> BlockHeader {
629        BlockHeader {
630            last_block: self.last_block,
631            block_type: self.block_type,
632            decompressed_size: self.decompressed_size,
633            content_size: self.content_size,
634        }
635    }
636}
637
638/// The frame being decoded.
639#[derive(Clone, Copy)]
640struct FrameState {
641    declared_size: Option<u64>,
642    has_checksum: bool,
643    block_maximum: usize,
644    produced: u64,
645}
646
647/// Allocation-free streaming zstd decoder over a caller-owned workspace.
648///
649/// See the [module documentation](self) for the contract and an example.
650pub struct WorkspaceDecoder<'w> {
651    huf: &'w mut HuffmanScratch,
652    fse: &'w mut FSEScratch,
653    offset_hist: [u32; 3],
654    literals: SliceScratch<'w>,
655    block_content: SliceScratch<'w>,
656    window: Window<'w>,
657    /// The workspace region every frame's window is laid over, held as a
658    /// pointer because frames alternate between a ring (which addresses it by
659    /// pointer) and a linear buffer (which borrows it), never both at once.
660    region: NonNull<u8>,
661    /// `region`'s length: the ring's capacity plus the copiers' slack.
662    region_len: usize,
663    ring_cap: usize,
664    _region: PhantomData<&'w mut [u8]>,
665    staging: &'w mut [u8],
666    staged: usize,
667    block_decoder: BlockDecoder,
668    kernel: crate::cpu_kernel::CpuKernelTag,
669    limits: DecodeLimits,
670    phase: Phase,
671    header: [u8; MAX_FRAME_HEADER],
672    header_len: usize,
673    header_target: usize,
674    frame: FrameState,
675    /// Every block of the current frame is decoded, so all of the window is
676    /// output to deliver rather than history to keep.
677    blocks_done: bool,
678    stream_produced: u64,
679    failed: bool,
680}
681
682// SAFETY: `region` points into the `&'w mut [u8]` workspace this decoder
683// borrows exclusively for `'w` (`_region` records that borrow), and every other
684// field is `Send`. Moving the decoder moves that exclusive borrow with it, as
685// moving the `&mut [u8]` itself would.
686unsafe impl Send for WorkspaceDecoder<'_> {}
687
688impl<'w> WorkspaceDecoder<'w> {
689    /// Build a decoder in `workspace`, which must hold at least
690    /// [`workspace_bytes`]`(limits.max_window_size)` bytes. Nothing is
691    /// allocated, now or later; the workspace stays borrowed for the decoder's
692    /// life.
693    ///
694    /// # Errors
695    ///
696    /// [`WorkspaceConfigError`] when the window limit is out of range or the
697    /// workspace is too small for it.
698    pub fn new(
699        workspace: &'w mut [u8],
700        limits: DecodeLimits,
701    ) -> Result<Self, WorkspaceConfigError> {
702        let layout = Layout::for_window(limits.max_window_size)?;
703        let provided = workspace.len();
704        if provided < layout.total {
705            return Err(WorkspaceConfigError::WorkspaceTooSmall {
706                needed: layout.total,
707                provided,
708            });
709        }
710        // `layout.total` holds `TABLE_ALIGN - 1` bytes of slack, and the
711        // offset a byte pointer needs to reach an alignment is below it.
712        let skip = workspace.as_ptr().align_offset(TABLE_ALIGN);
713        debug_assert!(skip < TABLE_ALIGN);
714        let rest = &mut workspace[skip..];
715        let (tables, rest) = rest.split_at_mut(layout.tables_len);
716        let (ring, rest) = rest.split_at_mut(layout.ring_cap + WILDCOPY_OVERLENGTH);
717        let (literals, rest) = rest.split_at_mut(layout.literals_len);
718        let staging = &mut rest[..layout.staging_len];
719
720        let tables = tables.as_mut_ptr();
721        let fse_ptr = tables.cast::<FSEScratch>();
722        // SAFETY: `tables` starts at a `TABLE_ALIGN` boundary and spans
723        // `tables_len` bytes: the FSE scratch at its start and the Huffman
724        // scratch at `huf_offset`, a multiple of its alignment, both in bounds.
725        // Each is fully initialised in place before a reference to it exists,
726        // and the bytes are reached through nothing else while `workspace`
727        // stays borrowed for `'w`. Neither type owns heap memory under this
728        // feature, so never running their destructors leaks nothing.
729        let (fse, huf) = unsafe {
730            let huf_ptr = tables.add(layout.huf_offset).cast::<HuffmanScratch>();
731            FSEScratch::init_in_place(fse_ptr);
732            HuffmanScratch::init_in_place(huf_ptr);
733            (&mut *fse_ptr, &mut *huf_ptr)
734        };
735        let region_len = ring.len();
736        let region = NonNull::from(ring).cast::<u8>();
737
738        let kernel = crate::cpu_kernel::detect_cpu_kernel();
739        Ok(WorkspaceDecoder {
740            huf,
741            fse,
742            offset_hist: [1, 4, 8],
743            literals: SliceScratch {
744                buf: literals,
745                len: 0,
746            },
747            block_content: SliceScratch {
748                buf: &mut [],
749                len: 0,
750            },
751            window: Window::Idle,
752            region,
753            region_len,
754            ring_cap: layout.ring_cap,
755            _region: PhantomData,
756            staging,
757            staged: 0,
758            block_decoder: block_decoder::with_kernel(kernel),
759            kernel,
760            limits,
761            phase: Phase::Magic,
762            header: [0; MAX_FRAME_HEADER],
763            header_len: 0,
764            header_target: 4,
765            frame: FrameState {
766                declared_size: None,
767                has_checksum: false,
768                block_maximum: 0,
769                produced: 0,
770            },
771            blocks_done: false,
772            stream_produced: 0,
773            failed: false,
774        })
775    }
776
777    /// Start a new stream: forget the current one, including a failure, and
778    /// expect a frame's magic number next. The workspace is reused as it is.
779    pub fn reset(&mut self) {
780        self.window = Window::Idle;
781        self.literals.clear();
782        self.staged = 0;
783        self.block_decoder = block_decoder::with_kernel(self.kernel);
784        self.phase = Phase::Magic;
785        self.header_len = 0;
786        self.header_target = 4;
787        self.blocks_done = false;
788        self.stream_produced = 0;
789        self.failed = false;
790    }
791
792    /// Decompressed bytes the stream has produced so far, delivered or not.
793    pub fn total_output(&self) -> u64 {
794        self.stream_produced
795    }
796
797    /// Decode from `input` into `output` until one of them runs out or the
798    /// stream ends.
799    ///
800    /// `input_is_final` says no input follows this slice. Without it, running
801    /// out of input asks for more ([`StepState::NeedInput`]); with it, input
802    /// that ends between frames finishes the stream ([`StepState::Done`]) and
803    /// input that ends inside one is [`WorkspaceDecodeError::Truncated`].
804    /// Unconsumed input must be passed again, first, on the next call.
805    ///
806    /// Concatenated frames decode one after another, and skippable frames are
807    /// consumed without output.
808    ///
809    /// # Errors
810    ///
811    /// See [`WorkspaceDecodeError`]. Any error ends the stream; the output
812    /// already written by this call is then not counted in a [`Progress`].
813    pub fn step(
814        &mut self,
815        input: &[u8],
816        input_is_final: bool,
817        output: &mut [u8],
818        work: &mut WorkBudget,
819    ) -> Result<Progress, WorkspaceDecodeError> {
820        if self.failed {
821            return Err(WorkspaceDecodeError::Failed);
822        }
823        let result = self.step_inner(input, input_is_final, output, work);
824        if result.is_err() {
825            self.failed = true;
826        }
827        result
828    }
829
830    fn step_inner(
831        &mut self,
832        input: &[u8],
833        input_is_final: bool,
834        output: &mut [u8],
835        work: &mut WorkBudget,
836    ) -> Result<Progress, WorkspaceDecodeError> {
837        let mut consumed = 0;
838        let mut produced = 0;
839        loop {
840            produced += self.deliver(&mut output[produced..]);
841            if self.pending() != 0 {
842                return Ok(Progress {
843                    consumed,
844                    produced,
845                    state: StepState::NeedOutput,
846                });
847            }
848            let rest = &input[consumed..];
849            match self.phase {
850                Phase::Skipping { remaining } => {
851                    if remaining == 0 {
852                        self.expect_frame();
853                        continue;
854                    }
855                    if rest.is_empty() {
856                        return self.out_of_input(input_is_final, consumed, produced);
857                    }
858                    let take = rest
859                        .len()
860                        .min(usize::try_from(remaining).unwrap_or(usize::MAX));
861                    work.charge(take as u64)?;
862                    consumed += take;
863                    self.phase = Phase::Skipping {
864                        remaining: remaining - take as u64,
865                    };
866                }
867                Phase::BlockBody(block) if block.block_type == BlockType::Raw => {
868                    // A Raw block's content is its output: it goes straight
869                    // into the window as it arrives, never through staging.
870                    // `staged` counts the bytes written so far.
871                    let remaining = block.content_size as usize - self.staged;
872                    if remaining == 0 {
873                        self.finish_block(block)?;
874                        continue;
875                    }
876                    if rest.is_empty() {
877                        return self.out_of_input(input_is_final, consumed, produced);
878                    }
879                    let take = rest.len().min(remaining);
880                    // One unit for the input byte, one for the output byte.
881                    work.charge(2 * take as u64)?;
882                    self.push_raw(&rest[..take])?;
883                    self.staged += take;
884                    consumed += take;
885                }
886                Phase::BlockBody(block) => {
887                    let size = block.content_size as usize;
888                    if self.staged == 0 && rest.len() >= size {
889                        work.charge(size as u64)?;
890                        consumed += size;
891                        self.decode_block(block, &rest[..size], work)?;
892                        continue;
893                    }
894                    if rest.is_empty() {
895                        return self.out_of_input(input_is_final, consumed, produced);
896                    }
897                    let take = rest.len().min(size - self.staged);
898                    work.charge(take as u64)?;
899                    self.staging[self.staged..self.staged + take].copy_from_slice(&rest[..take]);
900                    self.staged += take;
901                    consumed += take;
902                    if self.staged == size {
903                        self.staged = 0;
904                        // The staging region is a separate field from the
905                        // state the block decoder writes, so it lends the
906                        // content in place.
907                        let staging = core::mem::take(&mut self.staging);
908                        let decoded = self.decode_block(block, &staging[..size], work);
909                        self.staging = staging;
910                        decoded?;
911                    }
912                }
913                _ => {
914                    if self.header_len < self.header_target {
915                        if rest.is_empty() {
916                            return self.out_of_input(input_is_final, consumed, produced);
917                        }
918                        let take = rest.len().min(self.header_target - self.header_len);
919                        work.charge(take as u64)?;
920                        self.header[self.header_len..self.header_len + take]
921                            .copy_from_slice(&rest[..take]);
922                        self.header_len += take;
923                        consumed += take;
924                    }
925                    if self.header_len == self.header_target {
926                        self.header_complete()?;
927                    }
928                }
929            }
930        }
931    }
932
933    /// Input ran out. Only between frames does final input end the stream.
934    fn out_of_input(
935        &self,
936        input_is_final: bool,
937        consumed: usize,
938        produced: usize,
939    ) -> Result<Progress, WorkspaceDecodeError> {
940        if !input_is_final {
941            return Ok(Progress {
942                consumed,
943                produced,
944                state: StepState::NeedInput,
945            });
946        }
947        if self.phase == Phase::Magic && self.header_len == 0 {
948            return Ok(Progress {
949                consumed,
950                produced,
951                state: StepState::Done,
952            });
953        }
954        Err(WorkspaceDecodeError::Truncated)
955    }
956
957    /// Decoded bytes not delivered yet, whether or not the frame is finished.
958    fn pending(&self) -> usize {
959        with_window!(&self.window, buffer => buffer.undelivered(), idle => 0)
960    }
961
962    fn deliver(&mut self, output: &mut [u8]) -> usize {
963        if output.is_empty() {
964            return 0;
965        }
966        let blocks_done = self.blocks_done;
967        let drained = with_window!(&mut self.window, buffer => if blocks_done {
968            buffer.read_all(output)
969        } else {
970            buffer.read_reporting_pending(output).map(|(read, _)| read)
971        }, idle => Ok(0));
972        // Draining into a slice cannot fail: the copy reports no error.
973        drained.unwrap_or(0)
974    }
975
976    fn expect_frame(&mut self) {
977        self.phase = Phase::Magic;
978        self.header_len = 0;
979        self.header_target = 4;
980    }
981
982    fn expect_header(&mut self, phase: Phase, target: usize) {
983        self.phase = phase;
984        self.header_target = target;
985    }
986
987    /// The bytes the current phase collects are all in `header`.
988    fn header_complete(&mut self) -> Result<(), WorkspaceDecodeError> {
989        use WorkspaceDecodeError::Malformed;
990        match self.phase {
991            Phase::Magic => {
992                let magic = u32::from_le_bytes([
993                    self.header[0],
994                    self.header[1],
995                    self.header[2],
996                    self.header[3],
997                ]);
998                if (0x184D2A50..=0x184D2A5F).contains(&magic) {
999                    self.expect_header(Phase::SkippableSize, 8);
1000                } else if magic == MAGIC_NUM {
1001                    self.expect_header(Phase::Descriptor, 5);
1002                } else {
1003                    return Err(Malformed(MalformedInput::FrameHeader(
1004                        ReadFrameHeaderError::BadMagicNumber(magic),
1005                    )));
1006                }
1007            }
1008            Phase::SkippableSize => {
1009                let length = u32::from_le_bytes([
1010                    self.header[4],
1011                    self.header[5],
1012                    self.header[6],
1013                    self.header[7],
1014                ]);
1015                self.phase = Phase::Skipping {
1016                    remaining: u64::from(length),
1017                };
1018            }
1019            Phase::Descriptor => {
1020                let length = frame::measure_header_from_slice(&self.header[..5])
1021                    .map_err(|e| Malformed(MalformedInput::FrameHeader(e)))?;
1022                self.expect_header(Phase::FrameHeader, length);
1023            }
1024            Phase::FrameHeader => self.start_frame()?,
1025            Phase::BlockHeader => self.start_block()?,
1026            Phase::Checksum => {
1027                self.check_checksum()?;
1028                self.expect_frame();
1029            }
1030            Phase::Skipping { .. } | Phase::BlockBody(_) => {
1031                unreachable!("phases without a header buffer are handled by the step loop")
1032            }
1033        }
1034        Ok(())
1035    }
1036
1037    fn start_frame(&mut self) -> Result<(), WorkspaceDecodeError> {
1038        use WorkspaceDecodeError::{Limit, Malformed};
1039        let mut source = &self.header[..self.header_target];
1040        let (header, _) = frame::read_frame_header_from_slice(&mut source, false)
1041            .map_err(|e| Malformed(MalformedInput::FrameHeader(e)))?;
1042        if let Some(id) = header.dictionary_id() {
1043            return Err(WorkspaceDecodeError::Unsupported(Unsupported::Dictionary {
1044                id,
1045            }));
1046        }
1047        let window = header
1048            .window_size()
1049            .map_err(|e| Malformed(MalformedInput::Window(e)))?;
1050        if window > self.limits.max_window_size {
1051            return Err(Limit(LimitExceeded::Window {
1052                requested: window,
1053                limit: self.limits.max_window_size,
1054            }));
1055        }
1056        let declared_size = header.fcs_declared().then(|| header.frame_content_size());
1057        if let Some(declared) = declared_size
1058            && self
1059                .stream_produced
1060                .checked_add(declared)
1061                .is_none_or(|total| total > self.limits.max_output)
1062        {
1063            return Err(Limit(LimitExceeded::Output {
1064                limit: self.limits.max_output,
1065            }));
1066        }
1067        // Fits `usize`: it is at most the limit, which `Layout` converted.
1068        let window = window as usize;
1069        let has_checksum = header.descriptor.content_checksum_flag();
1070
1071        // End the previous frame's view of the region before laying out the
1072        // next one, so the two never coexist.
1073        self.window = Window::Idle;
1074        self.window = if header.descriptor.single_segment_flag() {
1075            // SAFETY: `region` is `region_len` initialised bytes borrowed for
1076            // `'w`, and no other view of it is alive: the previous window was
1077            // dropped just above.
1078            let region =
1079                unsafe { core::slice::from_raw_parts_mut(self.region.as_ptr(), self.region_len) };
1080            Window::Linear(DecodeBuffer::from_backend(
1081                UserSliceBackend::from_slice(region),
1082                window,
1083            ))
1084        } else {
1085            // SAFETY: as above; the ring addresses `ring_cap` bytes of the
1086            // region plus the copiers' slack, `region_len` in all, and
1087            // `ring_cap >= 1`.
1088            let ring = unsafe { RingBuffer::from_borrowed(self.region, self.ring_cap) };
1089            Window::Ring(DecodeBuffer::from_backend(ring, window))
1090        };
1091        with_window!(&mut self.window, buffer => {
1092            buffer.set_declared_content(declared_size);
1093            #[cfg(feature = "hash")]
1094            buffer.set_compute_hash(has_checksum);
1095        }, idle => {});
1096        self.offset_hist = [1, 4, 8];
1097        self.literals.clear();
1098        super::scratch::reset_entropy(self.huf, self.fse);
1099
1100        self.frame = FrameState {
1101            declared_size,
1102            has_checksum,
1103            block_maximum: block_decoder::block_maximum(window),
1104            produced: 0,
1105        };
1106        self.blocks_done = false;
1107        self.header_len = 0;
1108        self.expect_header(Phase::BlockHeader, 3);
1109        Ok(())
1110    }
1111
1112    fn start_block(&mut self) -> Result<(), WorkspaceDecodeError> {
1113        use WorkspaceDecodeError::Malformed;
1114        let (header, _) = self
1115            .block_decoder
1116            .read_block_header(&self.header[..3])
1117            .map_err(|e| Malformed(MalformedInput::BlockHeader(e)))?;
1118        // RFC 8878 3.1.1.2.4: Block_Size is at most Block_Maximum_Size for
1119        // every block type; for RLE it is the regenerated size.
1120        let size = match header.block_type {
1121            BlockType::RLE => header.decompressed_size,
1122            _ => header.content_size,
1123        } as usize;
1124        let maximum = self.frame.block_maximum;
1125        if size > maximum {
1126            return Err(Malformed(MalformedInput::BlockTooLarge { size, maximum }));
1127        }
1128        self.header_len = 0;
1129        self.staged = 0;
1130        self.phase = Phase::BlockBody(PendingBlock {
1131            last_block: header.last_block,
1132            block_type: header.block_type,
1133            decompressed_size: header.decompressed_size,
1134            content_size: header.content_size,
1135        });
1136        Ok(())
1137    }
1138
1139    /// Decode one block whose whole content is `content`, charging its work
1140    /// first and holding its output to the limits.
1141    fn decode_block(
1142        &mut self,
1143        block: PendingBlock,
1144        content: &[u8],
1145        work: &mut WorkBudget,
1146    ) -> Result<(), WorkspaceDecodeError> {
1147        use WorkspaceDecodeError::{Limit, Malformed};
1148        let output_bound = match block.block_type {
1149            BlockType::Compressed => self.frame.block_maximum as u64 + COMPRESSED_BLOCK_UNITS,
1150            _ => u64::from(block.decompressed_size),
1151        };
1152        work.charge(output_bound)?;
1153        let output_limit = Limit(LimitExceeded::Output {
1154            limit: self.limits.max_output,
1155        });
1156        if block.block_type != BlockType::Compressed
1157            && self
1158                .stream_produced
1159                .checked_add(u64::from(block.decompressed_size))
1160                .is_none_or(|total| total > self.limits.max_output)
1161        {
1162            return Err(output_limit);
1163        }
1164
1165        let header = block.header();
1166        let produced = with_window!(&mut self.window, buffer => run_block(
1167            &mut self.block_decoder,
1168            &mut BlockScratch {
1169                huf: &mut *self.huf,
1170                fse: &mut *self.fse,
1171                buffer,
1172                offset_hist: &mut self.offset_hist,
1173                literals: &mut self.literals,
1174                block_content: &mut self.block_content,
1175            },
1176            &header,
1177            content,
1178        ), idle => unreachable!("a block is decoded only inside a started frame"))
1179        .map_err(|e| Malformed(MalformedInput::Block(e)))?;
1180
1181        self.stream_produced += produced;
1182        self.frame.produced += produced;
1183        if self.stream_produced > self.limits.max_output {
1184            return Err(output_limit);
1185        }
1186        if let Some(declared) = self.frame.declared_size
1187            && self.frame.produced > declared
1188        {
1189            return Err(Malformed(MalformedInput::ContentSize {
1190                declared,
1191                produced: self.frame.produced,
1192            }));
1193        }
1194        self.finish_block(block)
1195    }
1196
1197    /// Append the next part of a Raw block to the window, after holding it to
1198    /// the output limit and to the frame's declared size.
1199    fn push_raw(&mut self, bytes: &[u8]) -> Result<(), WorkspaceDecodeError> {
1200        use crate::cpu_kernel::{BaselineKernel, CpuKernelTag, ScalarKernel};
1201        use WorkspaceDecodeError::{Limit, Malformed};
1202        let len = bytes.len() as u64;
1203        let stream_total = self.stream_produced + len;
1204        if stream_total > self.limits.max_output {
1205            return Err(Limit(LimitExceeded::Output {
1206                limit: self.limits.max_output,
1207            }));
1208        }
1209        let frame_total = self.frame.produced + len;
1210        if let Some(declared) = self.frame.declared_size
1211            && frame_total > declared
1212        {
1213            return Err(Malformed(MalformedInput::ContentSize {
1214                declared,
1215                produced: frame_total,
1216            }));
1217        }
1218        // One copy per piece, under the kernel a Raw block takes elsewhere:
1219        // the scalar tier stays portable.
1220        let scalar = self.kernel == CpuKernelTag::Scalar;
1221        let pushed = with_window!(&mut self.window, buffer => if scalar {
1222            buffer.try_push::<ScalarKernel>(bytes)
1223        } else {
1224            buffer.try_push::<BaselineKernel>(bytes)
1225        }, idle => unreachable!("a block is decoded only inside a started frame"));
1226        // The block was held to the block maximum and the window region holds
1227        // a window and two block maxima, so only a frame that outruns its own
1228        // window can reach here.
1229        pushed.map_err(|_| {
1230            Malformed(MalformedInput::Block(
1231                DecodeBlockContentError::BackendOverflow {
1232                    step: BlockType::Raw,
1233                },
1234            ))
1235        })?;
1236        self.stream_produced = stream_total;
1237        self.frame.produced = frame_total;
1238        Ok(())
1239    }
1240
1241    /// The block's content is all decoded: check a last block against the
1242    /// declared size and read what follows it.
1243    fn finish_block(&mut self, block: PendingBlock) -> Result<(), WorkspaceDecodeError> {
1244        if let Some(declared) = self.frame.declared_size
1245            && block.last_block
1246            && self.frame.produced != declared
1247        {
1248            return Err(WorkspaceDecodeError::Malformed(
1249                MalformedInput::ContentSize {
1250                    declared,
1251                    produced: self.frame.produced,
1252                },
1253            ));
1254        }
1255        self.staged = 0;
1256        self.header_len = 0;
1257        if !block.last_block {
1258            self.expect_header(Phase::BlockHeader, 3);
1259        } else {
1260            self.blocks_done = true;
1261            if self.frame.has_checksum {
1262                self.expect_header(Phase::Checksum, 4);
1263            } else {
1264                self.expect_frame();
1265            }
1266        }
1267        Ok(())
1268    }
1269
1270    /// Compare the stored checksum with the decoded frame's. The step loop
1271    /// reaches this only once the whole frame was delivered, which is when the
1272    /// running digest covers all of it. Without the `hash` feature the
1273    /// checksum is read and not checked.
1274    fn check_checksum(&mut self) -> Result<(), WorkspaceDecodeError> {
1275        #[cfg(feature = "hash")]
1276        {
1277            use core::hash::Hasher;
1278            let expected = u32::from_le_bytes([
1279                self.header[0],
1280                self.header[1],
1281                self.header[2],
1282                self.header[3],
1283            ]);
1284            let calculated =
1285                with_window!(&self.window, buffer => buffer.hash.finish(), idle => 0) as u32;
1286            if expected != calculated {
1287                return Err(WorkspaceDecodeError::Malformed(MalformedInput::Checksum {
1288                    expected,
1289                    calculated,
1290                }));
1291            }
1292        }
1293        Ok(())
1294    }
1295}
1296
1297#[cfg(test)]
1298mod tests;