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