From 21b3ea15099c866b0d6b760db607142f1d6b3857 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 15 Sep 2026 22:23:06 +0300 Subject: [PATCH 01/29] perf(decoding): hold a streamed frame to one window plus a block decode_from_to decoded every block its input held before draining any of it, so a caller handing over a whole frame grew the ring to the frame's content size, doubling and copying its way there. It now drains into the target before each block and decodes no further while the target is full, as upstream ZSTD_decompressStream flushes each block before the next. The ring's amortized growth stops at the frame's window plus one block instead of doubling past it, so a full power-of-two window asking room for a block no longer becomes two windows; a caller that holds more than that still grows. Frames that reserve up front reserve the window plus a block, content-capped, so filling the window costs no copy; frames of unknown size still grow lazily. Tests: a 4 MiB frame with a 1 MiB window streamed in 128 KiB steps keeps the workspace under 1.5 windows (it held the whole content), a 1 KiB target still receives every byte, and the ring's limit cuts the doubling step, leaves small growth alone and still grows for a need past it. Part of #508 --- zstd/src/decoding/buffer_backend.rs | 6 ++ zstd/src/decoding/decode_buffer.rs | 17 ++++- zstd/src/decoding/frame_decoder.rs | 93 +++++++++++++++++------ zstd/src/decoding/frame_decoder/tests.rs | 96 ++++++++++++++++++++++++ zstd/src/decoding/ringbuffer.rs | 35 ++++++++- zstd/src/decoding/ringbuffer/tests.rs | 45 +++++++++++ 6 files changed, 263 insertions(+), 29 deletions(-) diff --git a/zstd/src/decoding/buffer_backend.rs b/zstd/src/decoding/buffer_backend.rs index d083a2c12..63dc764e2 100644 --- a/zstd/src/decoding/buffer_backend.rs +++ b/zstd/src/decoding/buffer_backend.rs @@ -357,6 +357,12 @@ pub(crate) trait BufferBackend: Sized { /// fixed-capacity backends (`UserSliceBackend`), which are already bounded. fn set_max_capacity(&mut self, _max_capacity: usize) {} + /// Live byte count the frame's decode tops out at (window plus one + /// block), which amortized growth stops at instead of doubling past it. + /// Only `RingBuffer` grows by doubling across a whole window; the flat + /// backends are sized once per frame and take this no-op. + fn set_growth_limit(&mut self, _growth_limit: usize) {} + /// Live byte count: bytes between the logical head and tail. fn len(&self) -> usize; diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index d21dd006a..c370aea19 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -87,10 +87,23 @@ impl Read for DecodeBuffer { } } +/// Live bytes a frame with `window_size` holds at most while its output is +/// drained as it is produced: the window plus the block being decoded into it. +/// Upstream sizes its stream buffer the same way (`ZSTD_decodingBufferSize_min`). +/// A window so large the sum does not fit sets no limit at all. +fn peak_buffered_len(window_size: usize) -> usize { + let block = window_size.min(crate::common::MAX_BLOCK_SIZE as usize); + // Saturating on purpose: `usize::MAX` is the "no limit" value of the + // growth limit, which is exactly what a sum past it should mean. + window_size.saturating_add(block) +} + impl DecodeBuffer { pub fn new(window_size: usize) -> DecodeBuffer { + let mut buffer = B::new(); + buffer.set_growth_limit(peak_buffered_len(window_size)); DecodeBuffer { - buffer: B::new(), + buffer, window_size, total_output_counter: 0, #[cfg(feature = "hash")] @@ -116,6 +129,7 @@ impl DecodeBuffer { /// it issues vanish in the per-frame reset noise. pub fn from_backend(mut buffer: B, window_size: usize) -> DecodeBuffer { buffer.clear(); + buffer.set_growth_limit(peak_buffered_len(window_size)); DecodeBuffer { buffer, window_size, @@ -174,6 +188,7 @@ impl DecodeBuffer { pub fn reset(&mut self, window_size: usize) { self.window_size = window_size; self.buffer.clear(); + self.buffer.set_growth_limit(peak_buffered_len(window_size)); // No reserve here: capacity decisions are pushed up to the frame // layer. Direct-decode frames (`run_direct_decode`) write through // `UserSliceBackend` and never touch this buffer, so a long-lived diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index b12268f47..2a03bd9e5 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -513,8 +513,9 @@ impl DecoderScratchKind { /// frame writes only through `UserSliceBackend` and leaves this /// buffer empty. /// - /// `window_size` is the TARGET visible-window capacity: callers pass - /// the full window, and the method itself computes the shortfall past + /// `window_size` is the TARGET buffer capacity: callers pass the + /// frame's window plus a block (`decoding_buffer_size`), and the + /// method itself computes the shortfall past /// the bytes already buffered before calling the backend's /// ADDITIONAL-semantics `reserve_exact`. That keeps re-entries (the /// decode_all fallback loop runs `decode_blocks` once per strategy @@ -844,6 +845,24 @@ impl FrameDecoderState { } } + /// The up-front reservation for a frame that decodes through the buffer: + /// the window plus the block decoded into it, or the whole content when + /// that is smaller (upstream `ZSTD_decodingBufferSize_min`). Reserving the + /// window alone left the first block past a full window to grow the buffer + /// and copy the window across. + fn decoding_buffer_size(&self) -> usize { + let window_size = self.frame_header.window_size().unwrap_or(0); + let block = window_size.min(u64::from(crate::common::MAX_BLOCK_SIZE)); + // No overflow: the window was checked against + // `MAXIMUM_ALLOWED_WINDOW_SIZE` when the header was taken. + let needed = window_size + block; + if self.frame_header.fcs_declared() { + needed.min(self.frame_header.frame_content_size()) as usize + } else { + needed as usize + } + } + /// Construct a new frame decoder state, reading the frame header /// from `source`. When `magicless` is `true`, the 4-byte magic /// number prefix is NOT consumed (upstream zstd `ZSTD_f_zstd1_magicless`). @@ -1751,15 +1770,15 @@ impl FrameDecoder { } // Streaming entry point: pre-reserve the backing buffer to - // the FCS-capped window so multi-block frames don't pay repeated - // `reserve_amortized` grow steps (128 KiB → 256 KiB → ... → - // window) as blocks accumulate. `decode_all` does the same up - // front in `decode_all_impl`; this mirrors it for callers - // driving `decode_blocks` directly. Idempotent — the - // backend's `reserve` early-returns when capacity is already - // sufficient. - let useful_window = state.useful_window_size(); - state.decoder_scratch.reserve_buffer(useful_window); + // the FCS-capped window plus a block so multi-block frames don't pay + // repeated `reserve_amortized` grow steps (128 KiB → 256 KiB → ... → + // window) as blocks accumulate, nor a copy of the window when it + // fills. `decode_all` does the same up front in `decode_all_impl`; + // this mirrors it for callers driving `decode_blocks` directly. + // Idempotent — the backend's `reserve` early-returns when capacity + // is already sufficient. + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); let mut block_dec = decoding::block_decoder::new(); @@ -1993,13 +2012,13 @@ impl FrameDecoder { } // Mirror `decode_blocks`: pre-reserve the backing buffer to the - // FCS-capped window so multi-block frames don't pay repeated grow - // steps. The RAW frame window stays separately bound — the resume - // logic below bounds match reach by the frame's window semantics, - // not by the (possibly smaller) reservation cap. + // FCS-capped window plus a block so multi-block frames don't pay + // repeated grow steps. The RAW frame window stays separately bound — + // the resume logic below bounds match reach by the frame's window + // semantics, not by the (possibly smaller) reservation cap. let window_size = state.frame_header.window_size().unwrap_or(0) as usize; - let useful_window = state.useful_window_size(); - state.decoder_scratch.reserve_buffer(useful_window); + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); // Cold resume: prime the match window + restore entropy/repcode state + // advance the block cursor BEFORE the loop, so the first in-range block @@ -2307,7 +2326,10 @@ impl FrameDecoder { } } - /// Decodes as many blocks as possible from the source slice and reads from the decodebuffer into the target slice + /// Decodes blocks from the source slice and reads from the decodebuffer into the target slice, one block at a + /// time: output is handed to `target` as each block completes, and no further block is decoded while `target` + /// cannot take what is already decoded. The decode buffer so holds one window plus one block however much input + /// is supplied; call again with the unread input to continue. /// The source slice may contain only parts of a frame but must contain at least one full block to make progress /// /// By all means use decode_blocks if you have a io.Reader available. This is just for compatibility with other decompressors @@ -2330,6 +2352,8 @@ impl FrameDecoder { Some(s) => s.bytes_read_counter, None => 0, }; + // Bytes already handed to `target` by the per-block drain below. + let mut written = 0usize; if !self.is_finished() || self.state.is_none() { let mut mt_source = source; @@ -2383,6 +2407,22 @@ impl FrameDecoder { if state.frame_finished { break; } + // Hand what the window no longer needs to `target` before + // decoding more, and decode no further while `target` cannot + // take it: the buffer then holds one window plus the block + // being decoded, whatever the caller supplies, as upstream + // `ZSTD_decompressStream` flushes each block before the next. + written += state + .decoder_scratch + .buffer_read(&mut target[written..]) + .map_err(err::FailedToDrainDecodebuffer)?; + if state + .decoder_scratch + .buffer_can_drain_to_window_size() + .is_some_and(|pending| pending > 0) + { + break; + } //check if there are enough bytes for the next header if mt_source.len() < 3 { break; @@ -2450,7 +2490,10 @@ impl FrameDecoder { } } - let result_len = self.read(target).map_err(err::FailedToDrainDecodebuffer)?; + let result_len = written + + self + .read(&mut target[written..]) + .map_err(err::FailedToDrainDecodebuffer)?; // Once the frame is fully decoded and drained, the running digest is // final: validate it in `Verify` mode (no-op otherwise). Same finish // point as the streaming reader. @@ -2861,12 +2904,12 @@ impl FrameDecoder { // > 128 KiB otherwise grows through several intermediate // sizes with `alloc_zeroed + memcpy` each time). if let Some(state) = self.state.as_mut() { - // FCS-capped via `useful_window_size` — the same cap + // FCS-capped via `decoding_buffer_size` — the same cap // `decode_blocks` applies, so its per-iteration reserve in // the loop below cannot grow the buffer back to the raw // frame window. - let useful_window = state.useful_window_size(); - state.decoder_scratch.reserve_buffer(useful_window); + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); } let frame_start_total = total_bytes_written; loop { @@ -2995,12 +3038,12 @@ impl FrameDecoder { // `window_size` once so the per-block growth cycle is // skipped (see same comment on the no-lsm path above). if let Some(state) = self.state.as_mut() { - // FCS-capped via `useful_window_size` — the same cap + // FCS-capped via `decoding_buffer_size` — the same cap // `decode_blocks` applies, so its per-iteration reserve in // the loop below cannot grow the buffer back to the raw // frame window. - let useful_window = state.useful_window_size(); - state.decoder_scratch.reserve_buffer(useful_window); + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); } let frame_start_total = total_bytes_written; loop { diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index ab80907cc..27771cc82 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -862,6 +862,102 @@ fn reserve_buffer_reserves_the_shortfall_not_the_full_window_again() { ); } +/// A frame longer than its window, streamed through `decode_from_to` with the +/// whole frame as input, keeps one window plus one block of ring, as upstream +/// sizes its stream buffer (`ZSTD_decodingBufferSize_min`). The decode used to +/// run every buffered block before draining any, so the ring grew to the +/// frame's content size, doubling (and copying) its way there. +#[test] +fn a_streamed_frame_longer_than_its_window_keeps_one_window_of_ring() { + use crate::encoding::CompressionParameters; + let window_log = 20u32; + let window = 1usize << window_log; + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let payload: Vec = (0..4 * window) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + b"abcdefgh"[(state >> 61) as usize] + }) + .collect(); + let params = CompressionParameters::builder(CompressionLevel::Level(1)) + .window_log(window_log) + .build() + .expect("window_log within bounds"); + let mut compressor = FrameCompressor::new(CompressionLevel::Level(1)); + compressor.set_parameters(¶ms); + compressor.set_source(payload.as_slice()); + let mut compressed = Vec::new(); + compressor.set_drain(&mut compressed); + compressor.compress(); + // The frame must declare the 1 MiB window rather than be single-segment, + // or the decode takes the flat buffer and the ring is never exercised. + let header = crate::decoding::read_frame_header_info(&compressed, false).expect("header"); + assert_eq!(header.window_size, window as u64); + + let mut decoder = FrameDecoder::new(); + let mut source = compressed.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut decoded = Vec::with_capacity(payload.len()); + let mut chunk = alloc::vec![0u8; 128 * 1024]; + while !(decoder.is_finished() && decoder.can_collect() == 0) { + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + source = &source[read..]; + decoded.extend_from_slice(&chunk[..written]); + assert!(read > 0 || written > 0, "decode made no progress"); + } + assert_eq!(decoded, payload); + let workspace = decoder.workspace_size(); + assert!( + workspace < window + window / 2, + "ring grew past one window plus a block: workspace {workspace} bytes \ + for a {window}-byte window" + ); +} + +/// The per-block drain must not change what a caller with a buffer too small +/// for one block receives: a 1 KiB target still gets every byte, in order. +#[test] +fn a_streamed_frame_drains_through_a_target_smaller_than_a_block() { + let mut state = 0x2545_F491_4F6C_DD1Du64; + let payload: Vec = (0..600 * 1024) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + b"0123456789abcdef"[(state & 15) as usize] + }) + .collect(); + let params = crate::encoding::CompressionParameters::builder(CompressionLevel::Level(3)) + .window_log(17) + .build() + .expect("window_log within bounds"); + let mut compressor = FrameCompressor::new(CompressionLevel::Level(3)); + compressor.set_parameters(¶ms); + compressor.set_source(payload.as_slice()); + let mut compressed = Vec::new(); + compressor.set_drain(&mut compressed); + compressor.compress(); + + let mut decoder = FrameDecoder::new(); + let mut source = compressed.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut decoded = Vec::with_capacity(payload.len()); + let mut chunk = [0u8; 1024]; + while !(decoder.is_finished() && decoder.can_collect() == 0) { + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + source = &source[read..]; + decoded.extend_from_slice(&chunk[..written]); + assert!(read > 0 || written > 0, "decode made no progress"); + } + assert_eq!(decoded, payload); +} + #[test] fn dict_frame_decodes_through_direct_path() { // A dictionary frame decoded via `decode_all_with_dict_handle` diff --git a/zstd/src/decoding/ringbuffer.rs b/zstd/src/decoding/ringbuffer.rs index a8d0f8b9d..4916daa87 100644 --- a/zstd/src/decoding/ringbuffer.rs +++ b/zstd/src/decoding/ringbuffer.rs @@ -40,6 +40,12 @@ pub struct RingBuffer { /// grow, so well-formed blocks (covered by the upfront /// `reserve(MAX_BLOCK_SIZE)`) never pay for the check. max_capacity: usize, + /// Live byte count the frame's decode is expected to top out at (its + /// window plus one block). Amortized growth doubles, and the step that + /// would carry the ring past this is cut to it, so a window-sized ring is + /// not doubled for one block. A need beyond it still grows as before. + /// `usize::MAX` (the default) leaves the doubling alone. + growth_limit: usize, } // SAFETY: RingBuffer does not hold any thread specific values -> it can be sent to another thread -> RingBuffer is Send @@ -58,6 +64,7 @@ impl RingBuffer { head: 0, tail: 0, max_capacity: usize::MAX, + growth_limit: usize::MAX, } } @@ -192,6 +199,13 @@ impl RingBuffer { self.max_capacity = max_capacity; } + /// Set the live byte count growth should stop short at (see + /// [`Self::growth_limit`]). `usize::MAX` restores plain doubling. + #[inline] + pub fn set_growth_limit(&mut self, growth_limit: usize) { + self.growth_limit = growth_limit; + } + /// Fallible [`Self::reserve`]: identical fast path, but when the /// reserve would have to *grow* the ring it first rejects any target /// `len() + amount` past [`Self::max_capacity`]. This is where the @@ -250,9 +264,20 @@ impl RingBuffer { .cap .checked_add(amount) .expect("ringbuffer capacity overflow"); - let new_cap = usize::max(self.cap.next_power_of_two(), needed.next_power_of_two()) - .checked_add(1) - .expect("ringbuffer capacity overflow"); + let doubled = usize::max(self.cap.next_power_of_two(), needed.next_power_of_two()); + // Stop at the frame's expected peak rather than doubling past it: a + // full power-of-two window asking room for one more block would + // otherwise become two windows. Only when the need itself fits under + // the limit; a caller that holds more than the limit grows on. + // `needed` counts the sentinel slot (`cap + shortfall = len + amount + + // 1`) and the limit counts live bytes, hence the `- 1`; `needed >= 1` + // since growth only runs for a positive shortfall. + let target = if needed - 1 <= self.growth_limit { + doubled.min(self.growth_limit) + } else { + doubled + }; + let new_cap = target.checked_add(1).expect("ringbuffer capacity overflow"); // Check that the capacity isn't bigger than isize::MAX, which is the max allowed by LLVM, or that // we are on a >= 64 bit system which will never allow that much memory to be allocated @@ -1180,6 +1205,10 @@ impl super::buffer_backend::BufferBackend for RingBuffer { Self::set_max_capacity(self, max_capacity); } #[inline] + fn set_growth_limit(&mut self, growth_limit: usize) { + Self::set_growth_limit(self, growth_limit); + } + #[inline] fn len(&self) -> usize { Self::len(self) } diff --git a/zstd/src/decoding/ringbuffer/tests.rs b/zstd/src/decoding/ringbuffer/tests.rs index 722fb6148..f2f0fd92a 100644 --- a/zstd/src/decoding/ringbuffer/tests.rs +++ b/zstd/src/decoding/ringbuffer/tests.rs @@ -28,6 +28,51 @@ fn assert_branchless_matches_checked( assert_buffers_equal(&checked, &branchless); } +/// A full power-of-two window asking room for one more block grows to the +/// frame's limit (window plus block), not to twice the window. +#[test] +fn growth_stops_at_the_limit_instead_of_doubling_past_it() { + let window = 1usize << 16; + let block = 1usize << 12; + let mut rb = RingBuffer::new(); + rb.set_growth_limit(window + block); + rb.reserve(window); + assert_eq!(rb.cap, window + 1, "the window itself is a power of two"); + rb.extend(&alloc::vec![7u8; window]); + rb.reserve(block); + assert_eq!(rb.cap, window + block + 1); + assert!(rb.free() >= block); + assert_eq!(rb.len(), window); +} + +/// Growth below the limit keeps doubling, so small frames grow as before. +#[test] +fn growth_below_the_limit_keeps_doubling() { + let mut rb = RingBuffer::new(); + rb.set_growth_limit(1 << 20); + rb.reserve(1000); + assert_eq!(rb.cap, 1024 + 1); +} + +/// A caller holding more than the limit (one that does not drain) still gets +/// the room it asks for. +#[test] +fn a_need_past_the_limit_still_grows() { + let mut rb = RingBuffer::new(); + rb.set_growth_limit(4096); + rb.extend(&alloc::vec![1u8; 4000]); + rb.reserve(10_000); + assert!(rb.free() >= 10_000); + assert_eq!(rb.len(), 4000); + assert!( + rb.as_slices() + .0 + .iter() + .chain(rb.as_slices().1) + .all(|&b| b == 1) + ); +} + #[test] fn inline_exec_ok_respects_block_output_ceiling() { // The inline sequence-exec path bypasses `try_reserve`, so it must From 984c3a3aaefd4c4d8cf937d63caf118cab443102 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 15 Sep 2026 22:32:00 +0300 Subject: [PATCH 02/29] perf(decoding): allocate a declared-size stream frame's buffer once A fresh decoder streaming a frame grew its ring from nothing through every doubling, reallocating, copying and faulting its pages in per frame; bounding the growth at a window plus a block changed the sizes glibc sees and made the per-frame cost worse (an 8 MiB level-19 frame through a new ZSTD_DStream: 9.37 -> 16.56 ms, page faults 16.5K -> 122.8K over 20 frames). A frame that declares its size now gets min(window + block, content) in one allocation on its first decode_from_to call, as upstream allocates its stream buffer at the frame header. A frame of unknown size keeps growing lazily under the limit, so a small one is not charged its whole declared window. Part of #508 --- zstd/src/decoding/frame_decoder.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 2a03bd9e5..c67de1c3d 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -2368,6 +2368,17 @@ impl FrameDecoder { Some(s) => s, None => panic!("Bug in library"), }; + // A frame that declares its size gets its buffer in one + // allocation on its first call, as upstream allocates its + // stream buffer at the frame header: growing it block by block + // cost a fresh decoder several reallocations, copies and + // page-fault passes per frame. The size is content-capped, so a + // small frame gets a small buffer; a frame of unknown size keeps + // growing lazily rather than paying for its whole window. + if state.block_counter == 0 && state.frame_header.fcs_declared() { + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); + } let mut block_dec = decoding::block_decoder::new(); // Honour the content-checksum mode on this hand-rolled decode From 80990962fc08202ae5fe559e5ba97988918824c3 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 15 Sep 2026 22:41:41 +0300 Subject: [PATCH 03/29] perf(decoding): decode a frame of unknown size straight into the caller's slice decode_all took the direct path only for a frame declaring its size; one that declares none went through the drain path, which reserves the frame's declared window. A streamed producer's frame declares no size and, at level 19, an 8 MiB window, so a one-shot decode of 4 KiB allocated and zeroed megabytes per call (ZSTD_decompress: 186 us against libzstd's 4.3 us). Such a frame now decodes into the caller's slice as well, the slice being its limit, as upstream ZSTD_decompressDCtx decodes into dst; output past the slice is TargetTooSmall, as it was on the drain path. Frames declaring a size keep every content-size check. Tests: a 4 KiB level-19 frame of unknown size decodes into its slice without reserving its window (it reserved 8.5 MB), and one byte short is still TargetTooSmall. Part of #508 --- zstd/src/decoding/frame_decoder.rs | 128 +++++++++++++---------- zstd/src/decoding/frame_decoder/tests.rs | 64 ++++++++++++ 2 files changed, 137 insertions(+), 55 deletions(-) diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index c67de1c3d..50020b7c3 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -2771,15 +2771,17 @@ impl FrameDecoder { output.resize(frame_end, 0); // On error, drop the just-grown (zeroed) tail before propagating so // callers never observe bytes that were never decoded. - let written = - match self.run_direct_decode(&mut *input, &mut output[frame_start..], content_size) - { - Ok(n) => n, - Err(e) => { - output.truncate(frame_start); - return Err(e); - } - }; + let written = match self.run_direct_decode( + &mut *input, + &mut output[frame_start..], + Some(content_size), + ) { + Ok(n) => n, + Err(e) => { + output.truncate(frame_start); + return Err(e); + } + }; output.truncate(frame_start + written); #[cfg(feature = "hash")] self.verify_content_checksum()?; @@ -2896,9 +2898,18 @@ impl FrameDecoder { // that the spec relies on for `offset <= window_size` // validation. Path choice no longer alters checksum // semantics. - let direct_eligible = content_size > 0 && (output.len() as u64) >= content_size; + // A frame that declares no size decodes straight into the + // caller's slice too, the slice being its limit, as upstream + // `ZSTD_decompressDCtx` decodes into `dst`. The drain path + // reserved the frame's whole declared window for it, which for a + // streamed producer's small frame is megabytes for kilobytes. + let declared_size = fcs_declared.then_some(content_size); + let direct_eligible = match declared_size { + Some(declared) => declared > 0 && (output.len() as u64) >= declared, + None => true, + }; if direct_eligible { - let written = self.run_direct_decode(&mut input, output, content_size)?; + let written = self.run_direct_decode(&mut input, output, declared_size)?; output = &mut output[written..]; total_bytes_written += written; // Per-frame content-checksum verification (no-op unless the @@ -3034,9 +3045,18 @@ impl FrameDecoder { // `UserSliceBackend::exec_sequence_bounded`, so no // `WILDCOPY_OVERLENGTH` trailing slack is required (see the // no-lsm path above). - let direct_eligible = content_size > 0 && (output.len() as u64) >= content_size; + // A frame that declares no size decodes straight into the + // caller's slice too, the slice being its limit, as upstream + // `ZSTD_decompressDCtx` decodes into `dst`. The drain path + // reserved the frame's whole declared window for it, which for a + // streamed producer's small frame is megabytes for kilobytes. + let declared_size = fcs_declared.then_some(content_size); + let direct_eligible = match declared_size { + Some(declared) => declared > 0 && (output.len() as u64) >= declared, + None => true, + }; if direct_eligible { - let written = self.run_direct_decode(&mut input, output, content_size)?; + let written = self.run_direct_decode(&mut input, output, declared_size)?; output = &mut output[written..]; total_bytes_written += written; // Per-frame content-checksum verification (no-op unless the @@ -3160,12 +3180,13 @@ impl FrameDecoder { /// /// - `self.init` (or `init_with_dict_handle`) was called for /// this frame so `self.state` is populated. - /// - `content_size` matches `self.state.frame_header - /// .frame_content_size()` and is `> 0` (caller already passed - /// the eligibility gate). - /// - `output.len() >= content_size`. No `WILDCOPY_OVERLENGTH` - /// trailing slack is required: the trailing sequence(s) take the - /// bounded (non-overshooting) copy in + /// - `declared_size` is the frame's declared content size, `> 0`, with + /// `output.len() >= declared_size` (the eligibility gate), or `None` + /// for a frame that declares none. Then `output` itself is the limit, + /// as upstream `ZSTD_decompressDCtx` decodes into `dst`, and a frame + /// that does not fit is `TargetTooSmall` rather than a size mismatch. + /// No `WILDCOPY_OVERLENGTH` trailing slack is required: the trailing + /// sequence(s) take the bounded (non-overshooting) copy in /// [`UserSliceBackend::exec_sequence_bounded`]. /// /// Dictionary frames are supported: the scratch buffer's shared @@ -3181,7 +3202,7 @@ impl FrameDecoder { &mut self, input: &mut &[u8], output: &mut [u8], - content_size: u64, + declared_size: Option, ) -> Result { #[cfg(test)] { @@ -3194,6 +3215,16 @@ impl FrameDecoder { use crate::io::Read; use FrameDecoderError as err; + // The most the frame may write: its declared size, or the caller's + // slice for a frame that declares none. + let limit = declared_size.unwrap_or(output.len() as u64); + // Output past `limit`: the frame lied about its size, or it does not + // fit the caller's slice. + let overflow = |produced: u64| match declared_size { + Some(declared) => err::FrameContentSizeMismatch { declared, produced }, + None => err::TargetTooSmall, + }; + let state = self .state .as_mut() @@ -3215,7 +3246,7 @@ impl FrameDecoder { let n = bh.decompressed_size as usize; if bh.last_block && matches!(bh.block_type, crate::blocks::block::BlockType::Raw) - && n as u64 == content_size + && declared_size == Some(n as u64) && probe.len() >= n && output.len() >= n { @@ -3381,13 +3412,10 @@ impl FrameDecoder { // post-decode check below catches overflow via the // backend's actual write counter delta. let block_upper = u64::from(block_header.decompressed_size); - if block_upper > 0 && produced + block_upper > content_size { - // Frame is corrupt — Raw/RLE block headers claim - // more output than the FCS allows. - return Err(err::FrameContentSizeMismatch { - declared: content_size, - produced: produced + block_upper, - }); + if block_upper > 0 && produced + block_upper > limit { + // Raw/RLE block headers claim more output than the FCS + // allows (a corrupt frame) or the caller's slice holds. + return Err(overflow(produced + block_upper)); } // Slice-source fast path: consume the block body // straight from `input` without copying into the @@ -3412,11 +3440,9 @@ impl FrameDecoder { // accumulated `produced` can grow toward // u64::MAX across adversarial frames. Saturating // avoids a panic on the error path itself. - return Err(err::FrameContentSizeMismatch { - declared: content_size, - produced: produced - .saturating_add(u64::from(block_header.decompressed_size)), - }); + return Err(overflow( + produced.saturating_add(u64::from(block_header.decompressed_size)), + )); } // Compressed-block in-block overshoot: the sequence // executor (upstream zstd-inline path) or the match-repeat @@ -3426,12 +3452,11 @@ impl FrameDecoder { // from the partial fill: `tail` bytes were written before // the failing op, and `requested` is what overflowed — // their sum is a strict lower bound on the frame's true - // expanded size and is always > `content_size` (the - // direct path is only entered when the slice is sized to - // `content_size + WILDCOPY_OVERLENGTH`, so any overflow - // means the frame exceeded the declared FCS, never a - // caller-undersized buffer). Folds into the same - // `FrameContentSizeMismatch` contract as Raw/RLE. + // expanded size and is always > `limit`. With a declared + // size the slice holds at least that much, so any overflow + // means the frame exceeded its FCS, never a caller-undersized + // buffer, and folds into the same `FrameContentSizeMismatch` + // contract as Raw/RLE; without one the slice is the limit. Err(crate::decoding::errors::DecodeBlockContentError::DecompressBlockError( crate::decoding::errors::DecompressBlockError::ExecuteSequencesError(ref e), )) if e.output_overflow_requested().is_some() => { @@ -3439,10 +3464,7 @@ impl FrameDecoder { .output_overflow_requested() .expect("guard guarantees Some") as u64; let tail = direct.buffer.buffer_ref().tail() as u64; - return Err(err::FrameContentSizeMismatch { - declared: content_size, - produced: tail.saturating_add(requested), - }); + return Err(overflow(tail.saturating_add(requested))); } Err(e) => { return Err(block_body_decode_error( @@ -3465,11 +3487,8 @@ impl FrameDecoder { } produced = direct.buffer.buffer_ref().tail() as u64; // Post-decode FCS overflow check. - if produced > content_size { - return Err(err::FrameContentSizeMismatch { - declared: content_size, - produced, - }); + if produced > limit { + return Err(overflow(produced)); } state.bytes_read_counter += body_consumed; state.block_counter += 1; @@ -3493,15 +3512,14 @@ impl FrameDecoder { break; } } - // Final sanity: blocks summed to exactly `content_size`. - if produced != content_size { - return Err(err::FrameContentSizeMismatch { - declared: content_size, - produced, - }); + // Final sanity: blocks summed to exactly the declared size. + if let Some(declared) = declared_size + && produced != declared + { + return Err(err::FrameContentSizeMismatch { declared, produced }); } - let written = content_size as usize; + let written = produced as usize; state.frame_finished = true; // `direct`'s last use is in the decode loop above; NLL therefore // releases its `&mut output` borrow before here, freeing `output` for diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 27771cc82..e5b185f6d 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -958,6 +958,70 @@ fn a_streamed_frame_drains_through_a_target_smaller_than_a_block() { assert_eq!(decoded, payload); } +/// A streamed level-19 frame of 4 KiB that declares no content size, and so an +/// 8 MiB window, the shape a `ZSTD_compressStream2` / `zstd -` producer emits. +#[cfg(feature = "std")] +fn small_frame_of_unknown_size() -> (Vec, Vec) { + use crate::encoding::StreamingEncoder; + use std::io::Write as _; + let payload: Vec = (0..4096u32) + .map(|i| b"GET /index.html 200\n"[(i % 20) as usize] ^ (i / 97) as u8) + .collect(); + let mut encoder = StreamingEncoder::new(Vec::new(), CompressionLevel::Level(19)); + encoder.write_all(&payload).unwrap(); + let frame = encoder.finish().unwrap(); + let header = crate::decoding::read_frame_header_info(&frame, false).expect("header"); + assert!( + matches!( + header.content_size, + crate::decoding::FrameContentSize::Unknown + ), + "the fixture must declare no content size" + ); + assert!( + header.window_size >= 1 << 20, + "the fixture must declare a window far past its content" + ); + (payload, frame) +} + +/// Decoding a frame of unknown size into the caller's slice writes straight +/// into it, as upstream `ZSTD_decompressDCtx` decodes into `dst`: the frame's +/// declared window is never reserved. The drain path used to allocate (and +/// zero) the whole window for every such frame, a few kilobytes of content. +#[cfg(feature = "std")] +#[test] +fn a_frame_of_unknown_size_decodes_into_the_slice_without_its_window() { + let (payload, frame) = small_frame_of_unknown_size(); + let mut decoder = FrameDecoder::new(); + let mut out = alloc::vec![0u8; payload.len()]; + let written = decoder.decode_all(&frame, &mut out).expect("frame decodes"); + assert_eq!(written, payload.len()); + assert_eq!(out, payload); + let workspace = decoder.workspace_size(); + assert!( + workspace < 1 << 20, + "decoding into the caller's slice reserved {workspace} bytes of window" + ); +} + +/// A slice too small for a frame of unknown size is the caller's error, as +/// before, not a content-size mismatch. +#[cfg(feature = "std")] +#[test] +fn a_frame_of_unknown_size_into_a_short_slice_is_target_too_small() { + let (payload, frame) = small_frame_of_unknown_size(); + let mut decoder = FrameDecoder::new(); + let mut out = alloc::vec![0u8; payload.len() - 1]; + let err = decoder + .decode_all(&frame, &mut out) + .expect_err("one byte short must fail"); + assert!( + matches!(err, super::FrameDecoderError::TargetTooSmall), + "expected TargetTooSmall, got {err:?}" + ); +} + #[test] fn dict_frame_decodes_through_direct_path() { // A dictionary frame decoded via `decode_all_with_dict_handle` From d3ecc38ed0950a54f78757c40c2d7531158ad42a Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 15 Sep 2026 23:16:58 +0300 Subject: [PATCH 04/29] fix(decoding): hold direct-path blocks to the block maximum, reserve a stream buffer per block in hand - Direct decoding into the caller's slice never enforced MAX_BLOCK_SIZE per block: UserSliceBackend took the per-block ceiling as a no-op, so a malformed compressed block expanding past 128 KiB was accepted whenever the slice (or the declared content size) had room. It now keeps the ceiling on the live byte count and checks it where RingBuffer does, in the inline gate and the match reservation; an overflow within the slice is reported as a malformed block rather than TargetTooSmall. - decode_from_to reserved a declared-size frame's buffer on its first call, before knowing a block was there, so a header followed by nothing reserved the whole declared window. The reservation now waits for the first complete block. - A content-capped reservation below the window rounded up to the next power of two in the ring (600 KiB content, 1 MiB window: 1 MiB of ring). reserve_buffer now lowers the ring's growth limit to its target, so the doubling lands on the content size. Tests: a hand-built block of two RLE sequences writing 131,080 bytes is rejected on the direct path with and without a declared size (both were accepted) and on the ring; a 64 MiB header with no block reserves nothing (it reserved 67 MB); a 600 KiB frame with a 1 MiB window keeps a ring of its content (it held 1 MiB). Part of #508 --- zstd/src/decoding/decode_buffer.rs | 8 ++ zstd/src/decoding/frame_decoder.rs | 37 ++++-- zstd/src/decoding/frame_decoder/tests.rs | 150 +++++++++++++++++++++++ zstd/src/decoding/user_slice_buf.rs | 71 ++++++++--- 4 files changed, 237 insertions(+), 29 deletions(-) diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index c370aea19..8ddf42847 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -323,6 +323,14 @@ impl DecodeBuffer { self.buffer.reserve_exact(amount); } + /// Lower the live byte count growth stops at, for a frame that knows it + /// holds less than a window plus a block (its declared content is + /// smaller). `reset` sets the window-derived limit for every frame. + #[inline] + pub(crate) fn set_growth_limit(&mut self, growth_limit: usize) { + self.buffer.set_growth_limit(growth_limit); + } + /// Mutable backend handle. Lets the inline sequence executor /// write straight into the backend's physical storage; the /// `tail()` cursor on the backend is the authoritative output diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 50020b7c3..c812836a9 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -537,6 +537,10 @@ impl DecoderScratchKind { // window-sized buffer toward 2x window. match self { Self::Ring(s) => { + // The target is the most the frame holds, so growth stops + // there: a content-capped size below the window would + // otherwise round up to the next power of two. + s.buffer.set_growth_limit(window_size); let additional = window_size.saturating_sub(s.buffer.len()); s.buffer.reserve_exact(additional); } @@ -2368,17 +2372,6 @@ impl FrameDecoder { Some(s) => s, None => panic!("Bug in library"), }; - // A frame that declares its size gets its buffer in one - // allocation on its first call, as upstream allocates its - // stream buffer at the frame header: growing it block by block - // cost a fresh decoder several reallocations, copies and - // page-fault passes per frame. The size is content-capped, so a - // small frame gets a small buffer; a frame of unknown size keeps - // growing lazily rather than paying for its whole window. - if state.block_counter == 0 && state.frame_header.fcs_declared() { - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); - } let mut block_dec = decoding::block_decoder::new(); // Honour the content-checksum mode on this hand-rolled decode @@ -2452,6 +2445,19 @@ impl FrameDecoder { break; } state.bytes_read_counter += u64::from(block_header_size); + // A frame that declares its size gets its buffer in one + // allocation once its first block is in hand, as upstream + // allocates its stream buffer per frame: growing it block by + // block cost a fresh decoder several reallocations, copies + // and page-fault passes per frame. Not on the header alone, + // which would let a header followed by nothing reserve its + // whole declared window. The size is content-capped, so a + // small frame gets a small buffer; a frame of unknown size + // keeps growing lazily rather than paying for its window. + if state.block_counter == 0 && state.frame_header.fcs_declared() { + let buffer_size = state.decoding_buffer_size(); + state.decoder_scratch.reserve_buffer(buffer_size); + } // Only expose the held dictionary while THIS frame is dict-backed // (`using_dict` is set per dict-apply, cleared on reset). A reused @@ -3457,9 +3463,16 @@ impl FrameDecoder { // means the frame exceeded its FCS, never a caller-undersized // buffer, and folds into the same `FrameContentSizeMismatch` // contract as Raw/RLE; without one the slice is the limit. + // An overflow that stays within `limit` was refused by the + // per-block output ceiling instead: a malformed block, which + // takes the generic arm below. Err(crate::decoding::errors::DecodeBlockContentError::DecompressBlockError( crate::decoding::errors::DecompressBlockError::ExecuteSequencesError(ref e), - )) if e.output_overflow_requested().is_some() => { + )) if e.output_overflow_requested().is_some_and(|requested| { + (direct.buffer.buffer_ref().tail() as u64).saturating_add(requested as u64) + > limit + }) => + { let requested = e .output_overflow_requested() .expect("guard guarantees Some") as u64; diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index e5b185f6d..2e8b32371 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1022,6 +1022,156 @@ fn a_frame_of_unknown_size_into_a_short_slice_is_target_too_small() { ); } +/// A multi-segment frame whose declared content is smaller than its window +/// gets a ring of its content, not of the window rounded up: the header +/// carries both, and the ring's amortized growth would otherwise round the +/// content-capped reservation up to the next power of two. +#[test] +fn a_streamed_frame_smaller_than_its_window_gets_a_ring_of_its_content() { + let content = 600 * 1024u32; + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + frame.extend_from_slice(&content.to_le_bytes()); + let mut payload = Vec::with_capacity(content as usize); + let mut left = content; + while left > 0 { + let size = left.min(128 * 1024); + left -= size; + // Raw block header: last flag, type 0, size. + let header = size << 3 | u32::from(left == 0); + frame.extend_from_slice(&header.to_le_bytes()[..3]); + let body: Vec = (0..size).map(|i| (i * 7 + left) as u8).collect(); + payload.extend_from_slice(&body); + frame.extend_from_slice(&body); + } + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut decoded = Vec::with_capacity(payload.len()); + let mut chunk = alloc::vec![0u8; 128 * 1024]; + while !(decoder.is_finished() && decoder.can_collect() == 0) { + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + source = &source[read..]; + decoded.extend_from_slice(&chunk[..written]); + assert!(read > 0 || written > 0, "decode made no progress"); + } + assert_eq!(decoded, payload); + // The ring holds the 600 KiB content; the literal and block staging + // buffers take up to a block each on top. + let workspace = decoder.workspace_size(); + assert!( + workspace < 900 * 1024, + "ring rounded past the frame's content: workspace {workspace} bytes" + ); +} + +/// A streamed frame's one-shot buffer is reserved when its first block is in +/// hand, not on the header alone: a header declaring 64 MiB, followed by +/// nothing (a chunk boundary, or a truncated stream), costs no allocation. +#[test] +fn a_streamed_header_without_a_block_reserves_nothing() { + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x80, // window descriptor: 64 MiB + ]; + frame.extend_from_slice(&(64u32 << 20).to_le_bytes()); + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = alloc::vec![0u8; 1024]; + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("no block yet is not an error"); + assert_eq!((read, written), (0, 0)); + let workspace = decoder.workspace_size(); + assert!( + workspace < 1 << 20, + "a header with no block reserved {workspace} bytes" + ); +} + +/// A compressed block that expands past `MAX_BLOCK_SIZE` is malformed (RFC 8878 +/// 3.1.1.2.4, Block_Maximum_Size) however large the caller's slice. Hand-built: +/// a 1 MiB window, optionally a content size, and one block of 2 raw literals +/// plus two RLE-coded sequences (literal length 1, repeat offset 1, +/// match-length code 52 = 65,539) that together write 131,080 bytes. +const PAST_BLOCK_MAXIMUM_OUTPUT: u32 = 2 + 2 * 65_539; + +fn frame_with_a_block_past_the_block_maximum(content_size: Option) -> Vec { + let block: [u8; 13] = [ + 0x10, b'a', b'b', // raw literals section, 2 bytes + 0x02, // two sequences + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x34, // LL code 1, OF code 0, ML code 52 + 0x00, 0x00, 0x00, 0x00, // 16 zero extra bits per match length + 0x01, // stream start bit + ]; + let mut frame = alloc::vec![0x28, 0xB5, 0x2F, 0xFD]; // magic + match content_size { + // FHD: multi-segment, no checksum, no content size. + None => frame.extend_from_slice(&[0x00, 0x50]), + // FHD: multi-segment, 4-byte content size. + Some(size) => { + frame.extend_from_slice(&[0x80, 0x50]); + frame.extend_from_slice(&size.to_le_bytes()); + } + } + // (0x50 is the window descriptor: 1 MiB.) Last block, compressed, 13 bytes. + let header = (block.len() as u32) << 3 | 2 << 1 | 1; + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(&block); + frame +} + +#[test] +fn a_block_past_the_block_maximum_is_rejected_on_the_direct_path() { + let frame = frame_with_a_block_past_the_block_maximum(None); + let mut out = alloc::vec![0u8; 256 * 1024]; + let result = FrameDecoder::new().decode_all(&frame, &mut out); + // The slice had room: the block is malformed, not the target too small. + // (The block-body error variant carries its coordinates under `lsm`.) + match result { + Ok(_) | Err(super::FrameDecoderError::TargetTooSmall) => { + panic!("a block writing 131,080 bytes must be rejected as malformed, got {result:?}") + } + Err(_) => {} + } +} + +/// The same block in a frame whose declared size covers it: the content-size +/// bound does not stand in for the per-block one. +#[test] +fn a_block_past_the_block_maximum_is_rejected_under_a_declared_size() { + let frame = frame_with_a_block_past_the_block_maximum(Some(PAST_BLOCK_MAXIMUM_OUTPUT)); + let mut out = alloc::vec![0u8; PAST_BLOCK_MAXIMUM_OUTPUT as usize]; + let result = FrameDecoder::new().decode_all(&frame, &mut out); + assert!( + result.is_err(), + "a block writing 131,080 bytes must be rejected, got {result:?}" + ); +} + +#[test] +fn a_block_past_the_block_maximum_is_rejected_on_the_ring() { + let frame = frame_with_a_block_past_the_block_maximum(None); + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = alloc::vec![0u8; 256 * 1024]; + let result = decoder.decode_from_to(source, &mut chunk); + assert!( + result.is_err(), + "a block writing 131,080 bytes must be rejected, got {result:?}" + ); +} + #[test] fn dict_frame_decodes_through_direct_path() { // A dictionary frame decoded via `decode_all_with_dict_handle` diff --git a/zstd/src/decoding/user_slice_buf.rs b/zstd/src/decoding/user_slice_buf.rs index d644c9d75..5387b631d 100644 --- a/zstd/src/decoding/user_slice_buf.rs +++ b/zstd/src/decoding/user_slice_buf.rs @@ -3,21 +3,20 @@ //! //! Selected automatically by //! [`crate::decoding::FrameDecoder::decode_all`] (and -//! [`crate::decoding::FrameDecoder::decode_all_to_vec`]) when ALL of -//! the following hold: -//! - `frame_content_size > 0` — the header-derived content size -//! is non-zero. This is the actual eligibility condition (NOT -//! "FCS present"): an empty frame with an explicit FCS=0 -//! declaration on the wire stays on the fallback path because -//! there is no payload to write into the user slice. To -//! distinguish "FCS absent" from "FCS=0 explicit" elsewhere in -//! the decoder, use `FrameHeader::fcs_declared()` (e.g. the -//! fallback path's post-decode size check does). -//! - `output.len() >= frame_content_size` — the slice holds the -//! declared content. No `WILDCOPY_OVERLENGTH` slack is required: -//! when a sequence's literal+match bytes fit but the SIMD wildcopy -//! overshoot would not, the trailing sequence(s) take the bounded -//! (non-overshooting) copy in [`UserSliceBackend::exec_sequence_bounded`]. +//! [`crate::decoding::FrameDecoder::decode_all_to_vec`]) for a frame +//! that either: +//! - declares a non-zero content size the slice holds +//! (`output.len() >= frame_content_size`; an explicit FCS=0 stays on +//! the fallback path, there being no payload to write), or +//! - declares none, in which case the slice itself is the limit, as +//! upstream `ZSTD_decompressDCtx` decodes into `dst`. +//! +//! No `WILDCOPY_OVERLENGTH` slack is required: when a sequence's +//! literal+match bytes fit but the SIMD wildcopy overshoot would not, the +//! trailing sequence(s) take the bounded (non-overshooting) copy in +//! [`UserSliceBackend::exec_sequence_bounded`]. Each block is still held +//! to `MAX_BLOCK_SIZE` of output by the per-block ceiling +//! (`set_max_capacity`), whatever room the slice has. //! - No active dictionary (the persistent dict_content is not //! carried into the stack-local DecodeBuffer this backend //! builds; dict frames stay on the regular path). @@ -139,6 +138,13 @@ pub(crate) struct UserSliceBackend<'a> { /// for API parity with `FlatBuf` and `RingBuffer`. head: usize, tail: usize, + /// Per-block output ceiling on the live byte count, armed by + /// `set_block_output_ceiling` before each sequence section. A block may + /// write at most `MAX_BLOCK_SIZE` (RFC 8878 3.1.1.2.4) whatever room the + /// caller's slice has, which for a frame of unknown size is its only + /// other bound. Checked where `RingBuffer` checks it: the inline gate + /// and the match reservation. `usize::MAX` leaves the slice as the bound. + max_capacity: usize, } impl<'a> UserSliceBackend<'a> { @@ -154,9 +160,19 @@ impl<'a> UserSliceBackend<'a> { slice, head: 0, tail: 0, + max_capacity: usize::MAX, } } + /// Whether `n` more bytes keep the live byte count within the per-block + /// ceiling ([`Self::max_capacity`]). + #[inline(always)] + fn within_block_ceiling(&self, n: usize) -> bool { + (self.tail - self.head) + .checked_add(n) + .is_some_and(|live| live <= self.max_capacity) + } + /// Physical bytes `slice[from..tail]` — the output written since a /// previously-observed [`BufferBackend::tail`]. The direct decode /// path hashes each block's output through this right after the @@ -609,9 +625,26 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { slice: &mut [], head: 0, tail: 0, + max_capacity: usize::MAX, } } + /// The linear slice is always contiguous, so only the per-block output + /// ceiling can refuse the inline body; `sequence_output_fits` and the + /// tight-tail branch cover the slice's own bound. A refused sequence + /// takes the `push` / `repeat` path, whose `try_reserve` reports it. + #[inline(always)] + fn inline_exec_ok(&self, lit_length: usize, match_length: usize, _offset: usize) -> bool { + lit_length + .checked_add(match_length) + .is_some_and(|written| self.within_block_ceiling(written)) + } + + #[inline] + fn set_max_capacity(&mut self, max_capacity: usize) { + self.max_capacity = max_capacity; + } + #[inline] fn clear(&mut self) { self.head = 0; @@ -624,12 +657,16 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { // check. Lets safe public decode APIs catch a malformed-frame // overshoot here instead of via the `assert!` inside // `extend_from_within_unchecked` further down the call chain. + // The per-block ceiling bounds the match writes this reservation + // precedes, as `RingBuffer::try_reserve` bounds them. match self.tail.checked_add(n) { - Some(new_tail) if new_tail <= self.slice.len() => Ok(()), + Some(new_tail) if new_tail <= self.slice.len() && self.within_block_ceiling(n) => { + Ok(()) + } _ => Err(super::buffer_backend::BackendOverflow { tail: self.tail, requested: n, - capacity: self.slice.len(), + capacity: self.slice.len().min(self.max_capacity), }), } } From 09b20b6762bd7798154898e7105cb6660ad17e02 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 15 Sep 2026 23:28:40 +0300 Subject: [PATCH 05/29] perf(decoding): fold the block ceiling into the slice bound, keep a block of ring past the window - The direct path's per-block ceiling cost a check of its own on every sequence (inline gate plus reservation: +1.2..4.5% on the C ABI one-shot rows). It now narrows the one bound sequence writes already check, UserSliceBackend::cap, to the nearer of the slice's end and the ceiling, as upstream folds blockSizeMax into oend. - Capping a multi-segment frame's ring at its declared content left no block of room once the window filled whenever the content ran a few bytes past the window, and the block after doubled the ring (an 8 MiB level-19 frame through a new ZSTD_DStream: 6.95 -> 8.15 ms, page faults 6.5K -> 12.6K). The ring now reserves the content-capped window plus a block; a single-segment frame keeps exactly its content. Test: a frame 1000 bytes past its 1 MiB window, compressed blocks, streamed in 128 KiB steps, stays under 1.5 windows (it fails on the content-capped size). Part of #508 --- zstd/src/decoding/frame_decoder.rs | 29 +++++----- zstd/src/decoding/frame_decoder/tests.rs | 53 +++++++++++++++++ zstd/src/decoding/user_slice_buf.rs | 73 +++++++++++------------- 3 files changed, 102 insertions(+), 53 deletions(-) diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index c812836a9..730b9e24d 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -849,22 +849,25 @@ impl FrameDecoderState { } } - /// The up-front reservation for a frame that decodes through the buffer: - /// the window plus the block decoded into it, or the whole content when - /// that is smaller (upstream `ZSTD_decodingBufferSize_min`). Reserving the - /// window alone left the first block past a full window to grow the buffer - /// and copy the window across. + /// The up-front reservation for a frame that decodes through the buffer. + /// A single-segment frame's buffer holds its whole content, which is its + /// window. A multi-segment frame's ring holds the content-capped window + /// plus room for the next block, since each block reserves a whole block + /// before it decodes: reserving the window alone left the first block + /// past a full window to grow the ring and copy the window across, and + /// capping at the content left no block of room once the window filled. + /// Upstream sizes its stream buffer as window + block too + /// (`ZSTD_decodingBufferSize_min`); it can cap at the content because its + /// buffer is not a ring. fn decoding_buffer_size(&self) -> usize { - let window_size = self.frame_header.window_size().unwrap_or(0); - let block = window_size.min(u64::from(crate::common::MAX_BLOCK_SIZE)); + let useful_window = self.useful_window_size(); + if self.frame_header.descriptor.single_segment_flag() { + return useful_window; + } + let window_size = self.frame_header.window_size().unwrap_or(0) as usize; // No overflow: the window was checked against // `MAXIMUM_ALLOWED_WINDOW_SIZE` when the header was taken. - let needed = window_size + block; - if self.frame_header.fcs_declared() { - needed.min(self.frame_header.frame_content_size()) as usize - } else { - needed as usize - } + useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize) } /// Construct a new frame decoder state, reading the frame header diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 2e8b32371..9009c3d00 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1022,6 +1022,59 @@ fn a_frame_of_unknown_size_into_a_short_slice_is_target_too_small() { ); } +/// A frame a little longer than its window, with compressed blocks: each block +/// reserves a whole block of room before it decodes, so the ring needs the +/// window plus a block even though the content ends a few bytes past the +/// window. Capping the ring at the content made the block after a full window +/// double it. +#[test] +fn a_streamed_frame_just_past_its_window_keeps_room_for_a_block() { + use crate::encoding::CompressionParameters; + let window_log = 20u32; + let window = 1usize << window_log; + let mut state = 0x6A09_E667_F3BC_C908u64; + let payload: Vec = (0..window + 1000) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + b"abcdefgh"[(state >> 61) as usize] + }) + .collect(); + let params = CompressionParameters::builder(CompressionLevel::Level(1)) + .window_log(window_log) + .build() + .expect("window_log within bounds"); + let mut compressor = FrameCompressor::new(CompressionLevel::Level(1)); + compressor.set_parameters(¶ms); + compressor.set_source(payload.as_slice()); + let mut compressed = Vec::new(); + compressor.set_drain(&mut compressed); + compressor.compress(); + let header = crate::decoding::read_frame_header_info(&compressed, false).expect("header"); + assert_eq!(header.window_size, window as u64); + + let mut decoder = FrameDecoder::new(); + let mut source = compressed.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut decoded = Vec::with_capacity(payload.len()); + let mut chunk = alloc::vec![0u8; 128 * 1024]; + while !(decoder.is_finished() && decoder.can_collect() == 0) { + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + source = &source[read..]; + decoded.extend_from_slice(&chunk[..written]); + assert!(read > 0 || written > 0, "decode made no progress"); + } + assert_eq!(decoded, payload); + let workspace = decoder.workspace_size(); + assert!( + workspace < window + window / 2, + "ring grew past one window plus a block: workspace {workspace} bytes" + ); +} + /// A multi-segment frame whose declared content is smaller than its window /// gets a ring of its content, not of the window rounded up: the header /// carries both, and the ring's amortized growth would otherwise round the diff --git a/zstd/src/decoding/user_slice_buf.rs b/zstd/src/decoding/user_slice_buf.rs index 5387b631d..5454853bd 100644 --- a/zstd/src/decoding/user_slice_buf.rs +++ b/zstd/src/decoding/user_slice_buf.rs @@ -138,13 +138,17 @@ pub(crate) struct UserSliceBackend<'a> { /// for API parity with `FlatBuf` and `RingBuffer`. head: usize, tail: usize, - /// Per-block output ceiling on the live byte count, armed by - /// `set_block_output_ceiling` before each sequence section. A block may - /// write at most `MAX_BLOCK_SIZE` (RFC 8878 3.1.1.2.4) whatever room the - /// caller's slice has, which for a frame of unknown size is its only - /// other bound. Checked where `RingBuffer` checks it: the inline gate - /// and the match reservation. `usize::MAX` leaves the slice as the bound. - max_capacity: usize, + /// Where sequence writes must stop: the slice's end, or sooner under the + /// per-block output ceiling armed by `set_block_output_ceiling` before + /// each sequence section. A block may write at most `MAX_BLOCK_SIZE` + /// (RFC 8878 3.1.1.2.4) whatever room the caller's slice has, which for a + /// frame of unknown size is its only other bound. Folded into the one + /// bound every sequence write already checks ([`BufferBackend::cap`]), as + /// upstream folds `blockSizeMax` into `oend`, so the ceiling costs no + /// check of its own. Raw and RLE blocks write through `try_extend*`, + /// which keep the slice's end: the ceiling bounds sequences, as on + /// `RingBuffer`. + sequence_cap: usize, } impl<'a> UserSliceBackend<'a> { @@ -156,23 +160,15 @@ impl<'a> UserSliceBackend<'a> { /// back to [`Self::exec_sequence_bounded`] (exact, non-overshooting /// copies) for that trailing sequence. pub(crate) fn from_slice(slice: &'a mut [u8]) -> Self { + let sequence_cap = slice.len(); Self { slice, head: 0, tail: 0, - max_capacity: usize::MAX, + sequence_cap, } } - /// Whether `n` more bytes keep the live byte count within the per-block - /// ceiling ([`Self::max_capacity`]). - #[inline(always)] - fn within_block_ceiling(&self, n: usize) -> bool { - (self.tail - self.head) - .checked_add(n) - .is_some_and(|live| live <= self.max_capacity) - } - /// Physical bytes `slice[from..tail]` — the output written since a /// previously-observed [`BufferBackend::tail`]. The direct decode /// path hashes each block's output through this right after the @@ -282,7 +278,7 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { // wrapping past the slice length and letting the subsequent // unsafe pointer math go out of bounds. const MAX_WILDCOPY_OVERSHOOT: usize = 15; - let cap = self.slice.len(); + let cap = self.sequence_cap; // `self.tail <= cap` holds on entry (`from_slice` starts at 0 and // every prior sequence advanced `tail` only after this same check), // satisfying the `tail <= cap` precondition; see `sequence_output_fits`. @@ -393,7 +389,7 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { copy16, overlap_copy8, wildcopy_no_overlap, wildcopy_overlap_8byte_stride, }; const MAX_WILDCOPY_OVERSHOOT: usize = 15; - let cap = self.slice.len(); + let cap = self.sequence_cap; // `self.tail <= cap` precondition holds as in the SSE2 arm; see // `sequence_output_fits`. Hard guard with `overshoot = 0`; the // <=15-byte wildcopy slack is handled by the tight-tail branch @@ -500,7 +496,7 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { // this overshoot is handled by the tight-tail bounded branch below // rather than absorbed by slice capacity. const MAX_WILDCOPY_OVERSHOOT: usize = 31; - let cap = self.slice.len(); + let cap = self.sequence_cap; // `self.tail <= cap` holds on entry (`from_slice` starts at 0 and every // prior sequence advanced `tail` only after this same check), satisfying // the `tail <= cap` precondition; see `sequence_output_fits`. Hard guard @@ -625,24 +621,20 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { slice: &mut [], head: 0, tail: 0, - max_capacity: usize::MAX, + sequence_cap: 0, } } - /// The linear slice is always contiguous, so only the per-block output - /// ceiling can refuse the inline body; `sequence_output_fits` and the - /// tight-tail branch cover the slice's own bound. A refused sequence - /// takes the `push` / `repeat` path, whose `try_reserve` reports it. - #[inline(always)] - fn inline_exec_ok(&self, lit_length: usize, match_length: usize, _offset: usize) -> bool { - lit_length - .checked_add(match_length) - .is_some_and(|written| self.within_block_ceiling(written)) - } - + /// `max_capacity` bounds the live byte count, so the bound on the write + /// cursor is `head + max_capacity`, never past the slice. Saturating on + /// purpose: `usize::MAX` is "no ceiling", which lands on the slice's end. #[inline] fn set_max_capacity(&mut self, max_capacity: usize) { - self.max_capacity = max_capacity; + self.sequence_cap = self.slice.len().min(self.head.saturating_add(max_capacity)); + // The ceiling is armed as the live length plus a block, so it never + // lands behind the cursor; `sequence_output_fits` relies on + // `tail <= cap`. + debug_assert!(self.sequence_cap >= self.tail); } #[inline] @@ -657,16 +649,15 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { // check. Lets safe public decode APIs catch a malformed-frame // overshoot here instead of via the `assert!` inside // `extend_from_within_unchecked` further down the call chain. - // The per-block ceiling bounds the match writes this reservation - // precedes, as `RingBuffer::try_reserve` bounds them. + // Bounded by `sequence_cap`: the per-block ceiling bounds the match + // writes this reservation precedes, as `RingBuffer::try_reserve` + // bounds them. match self.tail.checked_add(n) { - Some(new_tail) if new_tail <= self.slice.len() && self.within_block_ceiling(n) => { - Ok(()) - } + Some(new_tail) if new_tail <= self.sequence_cap => Ok(()), _ => Err(super::buffer_backend::BackendOverflow { tail: self.tail, requested: n, - capacity: self.slice.len().min(self.max_capacity), + capacity: self.sequence_cap, }), } } @@ -691,9 +682,11 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { self.tail - self.head } + /// The bound sequence writes check ([`Self::sequence_cap`]): the slice's + /// end, or the per-block ceiling when it is nearer. #[inline] fn cap(&self) -> usize { - self.slice.len() + self.sequence_cap } #[inline] From 026dc019632a145186693354239832511d074266 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 00:03:14 +0300 Subject: [PATCH 06/29] fix(decoding): bound block literals, size small rings - Reject a block whose literals section regenerates more than the block maximum, and a block whose total output (sequences plus the literals left after the last one) runs past it, as upstream's ZSTD_decodeLiteralsBlock and the shared oend bound do. The ring path had no such check, so an oversized block decoded instead of failing. New error DecompressBlockError::ExpandsPastBlockMaximum. Regression tests: literals_past_the_block_maximum_are_rejected, trailing_literals_past_the_block_maximum_are_rejected. - A declared-size stream frame no larger than its window reserves just its content, not window plus a block: nothing ever drains out of the window, so the extra block was never written. The growth limit stays at window plus a block for frames that run past it. Test: a_streamed_frame_that_fits_its_window_reserves_just_its_content. - decode_from_to takes the pending drainable length from the read it already makes instead of querying the buffer a second time per block. --- zstd/src/decoding/block_decoder.rs | 19 ++++ zstd/src/decoding/decode_buffer.rs | 14 ++- zstd/src/decoding/errors.rs | 11 ++ zstd/src/decoding/frame_decoder.rs | 100 ++++++++++++------ zstd/src/decoding/frame_decoder/tests.rs | 127 ++++++++++++++++++++++- 5 files changed, 232 insertions(+), 39 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index e03bd2be8..5eda58dbe 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -313,8 +313,21 @@ impl BlockDecoder { raw: &[u8], dict: Option<&'d crate::decoding::dictionary::Dictionary>, ) -> Result<(), DecompressBlockError> { + // A block produces at most `MAX_BLOCK_SIZE` bytes, its literals and + // its matches together (RFC 8878 3.1.1.2.4). Upstream bounds both by + // the same `oend`: the literals up front (`litSize > blockSizeMax` is + // corruption in `ZSTD_decodeLiteralsBlock`) and every write after. + // Sequence writes stop at the per-block ceiling; the literals are + // checked here and the whole block after it decodes, which catches + // literals left over after the last sequence. + let len_before = buffer.len(); let mut section = LiteralsSection::new(); let bytes_in_literals_header = section.parse_from_header(raw)?; + if section.regenerated_size > MAX_BLOCK_SIZE { + return Err(DecompressBlockError::ExpandsPastBlockMaximum { + size: section.regenerated_size as usize, + }); + } let raw = &raw[bytes_in_literals_header as usize..]; vprintln!( "Found {} literalssection with regenerated size: {}, and compressed size: {:?}", @@ -421,6 +434,12 @@ impl BlockDecoder { buffer.push(literals_view); } + // Nothing drains the buffer inside a block, so the growth of its live + // length is this block's output. + let produced = buffer.len() - len_before; + if produced > MAX_BLOCK_SIZE as usize { + return Err(DecompressBlockError::ExpandsPastBlockMaximum { size: produced }); + } Ok(()) } diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index 8ddf42847..c3df3ec7c 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -74,6 +74,18 @@ pub(crate) struct DecodeBufferCheckpoint { impl Read for DecodeBuffer { fn read(&mut self, target: &mut [u8]) -> Result { + self.read_reporting_pending(target).map(|(read, _)| read) + } +} + +impl DecodeBuffer { + /// [`Read::read`], also reporting how many drainable bytes (past the + /// window) `target` had no room for, from the one length query the read + /// makes anyway. + pub(crate) fn read_reporting_pending( + &mut self, + target: &mut [u8], + ) -> Result<(usize, usize), Error> { let max_amount = self.can_drain_to_window_size().unwrap_or(0); let amount = max_amount.min(target.len()); @@ -83,7 +95,7 @@ impl Read for DecodeBuffer { written += buf.len(); (buf.len(), Ok(())) })?; - Ok(amount) + Ok((amount, max_amount - amount)) } } diff --git a/zstd/src/decoding/errors.rs b/zstd/src/decoding/errors.rs index 13f8f9a99..1d687f317 100644 --- a/zstd/src/decoding/errors.rs +++ b/zstd/src/decoding/errors.rs @@ -265,6 +265,12 @@ pub enum DecompressBlockError { SequencesHeaderParseError(SequencesHeaderParseError), DecodeSequenceError(DecodeSequenceError), ExecuteSequencesError(ExecuteSequencesError), + /// The block's literals, or its whole output, run past the block maximum + /// (RFC 8878 3.1.1.2.4, `Block_Maximum_Size`): `size` bytes where + /// `MAX_BLOCK_SIZE` is the most a block may produce. + ExpandsPastBlockMaximum { + size: usize, + }, } #[cfg(feature = "std")] @@ -302,6 +308,11 @@ impl core::fmt::Display for DecompressBlockError { DecompressBlockError::SequencesHeaderParseError(e) => write!(f, "{e:?}"), DecompressBlockError::DecodeSequenceError(e) => write!(f, "{e:?}"), DecompressBlockError::ExecuteSequencesError(e) => write!(f, "{e:?}"), + DecompressBlockError::ExpandsPastBlockMaximum { size } => write!( + f, + "Block expands to {size} bytes, past the maximum of {}", + crate::common::MAX_BLOCK_SIZE, + ), } } } diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 730b9e24d..7f918909d 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -513,17 +513,18 @@ impl DecoderScratchKind { /// frame writes only through `UserSliceBackend` and leaves this /// buffer empty. /// - /// `window_size` is the TARGET buffer capacity: callers pass the - /// frame's window plus a block (`decoding_buffer_size`), and the - /// method itself computes the shortfall past - /// the bytes already buffered before calling the backend's + /// `target` is the TARGET buffer capacity (`decoding_buffer_size`) and + /// `growth_limit` the most later growth may reach + /// (`decoding_buffer_limit`); the method itself computes the shortfall + /// past the bytes already buffered before calling the backend's /// ADDITIONAL-semantics `reserve_exact`. That keeps re-entries (the /// decode_all fallback loop runs `decode_blocks` once per strategy /// chunk, and streaming callers invoke it per call) from growing a /// window-full buffer toward 2x window, while per-block growth keeps /// the amortized `reserve`. #[inline] - fn reserve_buffer(&mut self, window_size: usize) { + fn reserve_buffer(&mut self, target: usize, growth_limit: usize) { + let window_size = target; // Exact growth: this is the one-shot pre-reservation, and a request // landing one slack past the retained capacity (e.g. a dictionary // prefix already loaded into the buffer) must not DOUBLE a @@ -537,12 +538,15 @@ impl DecoderScratchKind { // window-sized buffer toward 2x window. match self { Self::Ring(s) => { - // The target is the most the frame holds, so growth stops - // there: a content-capped size below the window would - // otherwise round up to the next power of two. + // The ring's growth rounds to the next power of two, capped at + // its limit; with the limit at the target for this one-shot + // reservation it lands exactly there, and later growth (a + // compressed block's own reservation) may then reach the + // frame's real limit, never double past it. s.buffer.set_growth_limit(window_size); let additional = window_size.saturating_sub(s.buffer.len()); s.buffer.reserve_exact(additional); + s.buffer.set_growth_limit(growth_limit); } Self::Flat(s) => { let additional = window_size.saturating_sub(s.buffer.len()); @@ -610,6 +614,18 @@ impl DecoderScratchKind { } } + /// [`Self::buffer_read`], also reporting the drainable bytes `target` had + /// no room for. + fn buffer_read_reporting_pending( + &mut self, + target: &mut [u8], + ) -> Result<(usize, usize), Error> { + match self { + Self::Ring(s) => s.buffer.read_reporting_pending(target), + Self::Flat(s) => s.buffer.read_reporting_pending(target), + } + } + fn buffer_read_all(&mut self, target: &mut [u8]) -> Result { match self { Self::Ring(s) => s.buffer.read_all(target), @@ -849,17 +865,17 @@ impl FrameDecoderState { } } - /// The up-front reservation for a frame that decodes through the buffer. - /// A single-segment frame's buffer holds its whole content, which is its - /// window. A multi-segment frame's ring holds the content-capped window - /// plus room for the next block, since each block reserves a whole block - /// before it decodes: reserving the window alone left the first block - /// past a full window to grow the ring and copy the window across, and - /// capping at the content left no block of room once the window filled. - /// Upstream sizes its stream buffer as window + block too + /// The most a frame decoding through the buffer holds, which its growth + /// stops at. A single-segment frame's buffer holds its whole content, + /// which is its window. A multi-segment frame's ring holds the + /// content-capped window plus room for the next block, since each + /// compressed block reserves a whole block before it decodes: without it + /// the first block past a full window grew the ring and copied the window + /// across, and a limit at the content left no block of room once the + /// window filled. Upstream sizes its stream buffer as window + block too /// (`ZSTD_decodingBufferSize_min`); it can cap at the content because its /// buffer is not a ring. - fn decoding_buffer_size(&self) -> usize { + fn decoding_buffer_limit(&self) -> usize { let useful_window = self.useful_window_size(); if self.frame_header.descriptor.single_segment_flag() { return useful_window; @@ -870,6 +886,30 @@ impl FrameDecoderState { useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize) } + /// What to reserve up front: the limit, except for a multi-segment frame + /// whose declared content fits its window, which reserves just its + /// content. Only compressed blocks need the block of room past it, and + /// they reserve it themselves; the limit then caps that growth at one + /// block. Such frames are rare (encoders mark a frame that fits its window + /// single-segment), and a small Raw or RLE one should not pay a block. + fn decoding_buffer_size(&self) -> usize { + let window_size = self.frame_header.window_size().unwrap_or(0); + if self.frame_header.fcs_declared() && self.frame_header.frame_content_size() <= window_size + { + self.useful_window_size() + } else { + self.decoding_buffer_limit() + } + } + + /// Reserve this frame's decode buffer ([`Self::decoding_buffer_size`]) + /// and cap its later growth ([`Self::decoding_buffer_limit`]). + fn reserve_decoding_buffer(&mut self) { + let target = self.decoding_buffer_size(); + let growth_limit = self.decoding_buffer_limit(); + self.decoder_scratch.reserve_buffer(target, growth_limit); + } + /// Construct a new frame decoder state, reading the frame header /// from `source`. When `magicless` is `true`, the 4-byte magic /// number prefix is NOT consumed (upstream zstd `ZSTD_f_zstd1_magicless`). @@ -1784,8 +1824,7 @@ impl FrameDecoder { // this mirrors it for callers driving `decode_blocks` directly. // Idempotent — the backend's `reserve` early-returns when capacity // is already sufficient. - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); + state.reserve_decoding_buffer(); let mut block_dec = decoding::block_decoder::new(); @@ -2024,8 +2063,7 @@ impl FrameDecoder { // the resume logic below bounds match reach by the frame's window // semantics, not by the (possibly smaller) reservation cap. let window_size = state.frame_header.window_size().unwrap_or(0) as usize; - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); + state.reserve_decoding_buffer(); // Cold resume: prime the match window + restore entropy/repcode state + // advance the block cursor BEFORE the loop, so the first in-range block @@ -2419,15 +2457,12 @@ impl FrameDecoder { // take it: the buffer then holds one window plus the block // being decoded, whatever the caller supplies, as upstream // `ZSTD_decompressStream` flushes each block before the next. - written += state + let (read, pending) = state .decoder_scratch - .buffer_read(&mut target[written..]) + .buffer_read_reporting_pending(&mut target[written..]) .map_err(err::FailedToDrainDecodebuffer)?; - if state - .decoder_scratch - .buffer_can_drain_to_window_size() - .is_some_and(|pending| pending > 0) - { + written += read; + if pending > 0 { break; } //check if there are enough bytes for the next header @@ -2458,8 +2493,7 @@ impl FrameDecoder { // small frame gets a small buffer; a frame of unknown size // keeps growing lazily rather than paying for its window. if state.block_counter == 0 && state.frame_header.fcs_declared() { - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); + state.reserve_decoding_buffer(); } // Only expose the held dictionary while THIS frame is dict-backed @@ -2939,8 +2973,7 @@ impl FrameDecoder { // `decode_blocks` applies, so its per-iteration reserve in // the loop below cannot grow the buffer back to the raw // frame window. - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); + state.reserve_decoding_buffer(); } let frame_start_total = total_bytes_written; loop { @@ -3082,8 +3115,7 @@ impl FrameDecoder { // `decode_blocks` applies, so its per-iteration reserve in // the loop below cannot grow the buffer back to the raw // frame window. - let buffer_size = state.decoding_buffer_size(); - state.decoder_scratch.reserve_buffer(buffer_size); + state.reserve_decoding_buffer(); } let frame_start_total = total_bytes_written; loop { diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 9009c3d00..4741a3180 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -847,13 +847,13 @@ fn reserve_buffer_reserves_the_shortfall_not_the_full_window_again() { use super::DecoderScratchKind; let window = 1usize << 20; let mut scratch = DecoderScratchKind::new_flat(window); - scratch.reserve_buffer(window); + scratch.reserve_buffer(window, window); let data = alloc::vec![0u8; window]; match &mut scratch { super::DecoderScratchKind::Flat(s) => s.buffer.push(&data), super::DecoderScratchKind::Ring(_) => unreachable!("new_flat builds Flat"), } - scratch.reserve_buffer(window); + scratch.reserve_buffer(window, window); let workspace = scratch.workspace_bytes(); assert!( workspace < window * 3 / 2, @@ -1124,6 +1124,46 @@ fn a_streamed_frame_smaller_than_its_window_gets_a_ring_of_its_content() { ); } +/// Capacity of the ring a multi-segment frame decoded into. +fn ring_capacity(decoder: &FrameDecoder) -> usize { + match &decoder + .state + .as_ref() + .expect("a frame was reset") + .decoder_scratch + { + super::DecoderScratchKind::Ring(s) => s.buffer.capacity(), + super::DecoderScratchKind::Flat(_) => panic!("a multi-segment frame decodes into the ring"), + } +} + +/// A multi-segment frame whose declared content fits its window reserves its +/// content, not a block past it: a one-byte Raw frame with a 1 MiB window +/// needs a byte of ring. +#[test] +fn a_streamed_frame_that_fits_its_window_reserves_just_its_content() { + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + frame.extend_from_slice(&1u32.to_le_bytes()); + frame.extend_from_slice(&[0x09, 0x00, 0x00, b'q']); // last Raw block, 1 byte + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = [0u8; 16]; + let (_, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + assert_eq!(&chunk[..written], b"q"); + let capacity = ring_capacity(&decoder); + assert!( + capacity < 1024, + "a one-byte frame reserved {capacity} bytes of ring" + ); +} + /// A streamed frame's one-shot buffer is reserved when its first block is in /// hand, not on the header alone: a header declaring 64 MiB, followed by /// nothing (a chunk boundary, or a truncated stream), costs no allocation. @@ -1166,6 +1206,12 @@ fn frame_with_a_block_past_the_block_maximum(content_size: Option) -> Vec) -> Vec { let mut frame = alloc::vec![0x28, 0xB5, 0x2F, 0xFD]; // magic match content_size { // FHD: multi-segment, no checksum, no content size. @@ -1176,13 +1222,86 @@ fn frame_with_a_block_past_the_block_maximum(content_size: Option) -> Vec [u8; 3] { + [ + ((regenerated & 0xF) << 4) as u8 | 0b11 << 2 | literals_type, + (regenerated >> 4) as u8, + (regenerated >> 12) as u8, + ] +} + +/// A block with no sequences whose RLE literals regenerate 200,000 bytes: past +/// the block maximum on literals alone. +fn block_of_literals_past_the_block_maximum() -> Vec { + let mut block = literals_header_20_bit(1, 200_000).to_vec(); + block.push(b'z'); // the repeated byte + block.push(0x00); // no sequences + block +} + +/// One sequence (literal length 1, repeat offset 1, match length 65,539) and +/// 65,539 literals left over after it: 131,079 bytes, where the sequence alone +/// stays within the block maximum and the trailing literals take it past. +fn block_with_trailing_literals_past_the_block_maximum() -> Vec { + let literals = 1 + 65_539; + let mut block = literals_header_20_bit(0, literals).to_vec(); + block.extend((0..literals).map(|i| i as u8)); + block.extend_from_slice(&[ + 0x01, // one sequence + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x34, // LL code 1, OF code 0, ML code 52 + 0x00, 0x00, // 16 zero extra bits for the match length + 0x01, // stream start bit + ]); + block +} + +/// `frame` fails to decode as malformed through the caller's slice (with +/// room to spare) and through the streaming ring alike. +fn assert_rejected_as_malformed(frame: &[u8], what: &str) { + let mut out = alloc::vec![0u8; 512 * 1024]; + match FrameDecoder::new().decode_all(frame, &mut out) { + Ok(_) | Err(super::FrameDecoderError::TargetTooSmall) => { + panic!("{what}: the direct path must reject the block as malformed") + } + Err(_) => {} + } + let mut decoder = FrameDecoder::new(); + let mut source = frame; + decoder.reset(&mut source).expect("header parses"); + let mut chunk = alloc::vec![0u8; 512 * 1024]; + assert!( + decoder.decode_from_to(source, &mut chunk).is_err(), + "{what}: the ring must reject the block" + ); +} + +#[test] +fn literals_past_the_block_maximum_are_rejected() { + let block = block_of_literals_past_the_block_maximum(); + assert_rejected_as_malformed(&frame_around_block(&block, None), "unsized"); + assert_rejected_as_malformed(&frame_around_block(&block, Some(200_000)), "declared size"); +} + +#[test] +fn trailing_literals_past_the_block_maximum_are_rejected() { + let block = block_with_trailing_literals_past_the_block_maximum(); + assert_rejected_as_malformed(&frame_around_block(&block, None), "unsized"); + assert_rejected_as_malformed( + &frame_around_block(&block, Some(1 + 65_539 + 65_539)), + "declared size", + ); +} + #[test] fn a_block_past_the_block_maximum_is_rejected_on_the_direct_path() { let frame = frame_with_a_block_past_the_block_maximum(None); From ac4953f358749bd6250b9afd4ccf1c1d386749d0 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 08:26:39 +0300 Subject: [PATCH 07/29] fix(decoding): hold a block to its frame's block maximum A frame's block maximum is the smaller of its window and 128 KiB (RFC 8878 3.1.1.2.4), which upstream derives once per frame as `blockSizeMax = MIN(windowSize, ZSTD_BLOCKSIZE_MAX)` and checks every block against. Ours bounded blocks at 128 KiB alone, so a frame with a 1 KiB window decoded blocks of 2 KiB: past what the format allows, and past what the window-sized buffer is meant to hold. The limit now comes from the window and bounds a compressed block's literals, its whole output and its sequence writes, plus a Raw or RLE block's size from its header before it writes, as upstream checks `rSize`. Regression test a_block_past_a_small_window_is_rejected covers all three block types and failed on each before this. Also pins that the per-block sequence ceiling does not narrow a following Raw block: a compressed block of 13 bytes followed by a Raw block of a whole block maximum fills the caller's slice (a_raw_block_after_a_compressed_one_fills_the_slice). --- zstd/src/decoding/block_decoder.rs | 69 +++++++++++---- zstd/src/decoding/frame_decoder/tests.rs | 88 +++++++++++++++++++ zstd/src/decoding/sequence_section_decoder.rs | 10 ++- zstd/src/decoding/user_slice_buf.rs | 14 +-- 4 files changed, 155 insertions(+), 26 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 5eda58dbe..7617bbfc3 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -33,6 +33,33 @@ enum DecoderState { Failed, //TODO put "self.internal_state = DecoderState::Failed;" everywhere an unresolvable error occurs } +/// The most a block of a frame with `window_size` may produce: the smaller of +/// the window and 128 KiB (RFC 8878 3.1.1.2.4), which upstream keeps per frame +/// as `blockSizeMax`. A single-segment frame's window is its content size, so +/// such a frame's blocks are bounded by the content as well. +#[inline] +pub(crate) fn block_maximum(window_size: usize) -> usize { + window_size.min(MAX_BLOCK_SIZE as usize) +} + +/// A Raw or RLE block states its output in its header, so it is held to the +/// frame's block maximum before anything is written, as upstream checks `rSize` +/// against `blockSizeMax` in `ZSTD_decompressContinue`. A compressed block's +/// output is only known as it decodes, and is checked there. +#[inline] +fn block_fits_the_maximum( + header: &BlockHeader, + window_size: usize, +) -> Result<(), DecodeBlockContentError> { + let size = header.decompressed_size as usize; + if size > block_maximum(window_size) { + return Err(DecodeBlockContentError::DecompressBlockError( + DecompressBlockError::ExpandsPastBlockMaximum { size }, + )); + } + Ok(()) +} + /// Create a new [BlockDecoder]. pub fn new() -> BlockDecoder { BlockDecoder { @@ -94,8 +121,9 @@ impl BlockDecoder { // path. Advance ONLY after the write succeeds, matching // the Raw arm's split_at-then-try_push-then-advance shape. let fill = source[0]; - workspace - .split() + let parts = workspace.split(); + block_fits_the_maximum(header, parts.buffer.window_size)?; + parts .buffer .try_extend_and_fill(fill, header.decompressed_size as usize) .map_err(|_| DecodeBlockContentError::BackendOverflow { step: block_type })?; @@ -120,8 +148,9 @@ impl BlockDecoder { // `UserSliceBackend` when the Raw payload would push // past the caller's output slice. Growable backends // grow on demand and always succeed. - workspace - .split() + let parts = workspace.split(); + block_fits_the_maximum(header, parts.buffer.window_size)?; + parts .buffer .try_push(payload) .map_err(|_| DecodeBlockContentError::BackendOverflow { step: block_type })?; @@ -174,8 +203,9 @@ impl BlockDecoder { source: err, } })?; - workspace - .split() + let parts = workspace.split(); + block_fits_the_maximum(header, parts.buffer.window_size)?; + parts .buffer .extend_and_fill(buf[0], header.decompressed_size as usize); @@ -189,8 +219,9 @@ impl BlockDecoder { // borrow-by-reference indirection. (Both io shims provide a // blanket `Read for &mut T`, so `&mut source` would also // compile; the by-value form is just cleaner here.) - workspace - .split() + let parts = workspace.split(); + block_fits_the_maximum(header, parts.buffer.window_size)?; + parts .buffer .extend_from_reader(source, header.decompressed_size as usize) .map_err(|err| DecodeBlockContentError::ReadError { @@ -313,17 +344,21 @@ impl BlockDecoder { raw: &[u8], dict: Option<&'d crate::decoding::dictionary::Dictionary>, ) -> Result<(), DecompressBlockError> { - // A block produces at most `MAX_BLOCK_SIZE` bytes, its literals and - // its matches together (RFC 8878 3.1.1.2.4). Upstream bounds both by - // the same `oend`: the literals up front (`litSize > blockSizeMax` is - // corruption in `ZSTD_decodeLiteralsBlock`) and every write after. - // Sequence writes stop at the per-block ceiling; the literals are - // checked here and the whole block after it decodes, which catches - // literals left over after the last sequence. + // A block produces at most its frame's block maximum, its literals and + // its matches together: the smaller of the window and 128 KiB (RFC 8878 + // 3.1.1.2.4), as upstream derives it once per frame + // (`zstd_decompress.c`: `blockSizeMax = MIN(windowSize, + // ZSTD_BLOCKSIZE_MAX)`). Upstream bounds both halves by the same + // `oend`: the literals up front (`litSize > blockSizeMax` is corruption + // in `ZSTD_decodeLiteralsBlock`) and every write after. Sequence writes + // stop at the per-block ceiling; the literals are checked here and the + // whole block after it decodes, which catches literals left over after + // the last sequence. + let block_maximum = block_maximum(buffer.window_size); let len_before = buffer.len(); let mut section = LiteralsSection::new(); let bytes_in_literals_header = section.parse_from_header(raw)?; - if section.regenerated_size > MAX_BLOCK_SIZE { + if section.regenerated_size as usize > block_maximum { return Err(DecompressBlockError::ExpandsPastBlockMaximum { size: section.regenerated_size as usize, }); @@ -437,7 +472,7 @@ impl BlockDecoder { // Nothing drains the buffer inside a block, so the growth of its live // length is this block's output. let produced = buffer.len() - len_before; - if produced > MAX_BLOCK_SIZE as usize { + if produced > block_maximum { return Err(DecompressBlockError::ExpandsPastBlockMaximum { size: produced }); } Ok(()) diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 4741a3180..0c96d6ae9 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1239,6 +1239,21 @@ fn literals_header_20_bit(literals_type: u8, regenerated: u32) -> [u8; 3] { ] } +/// A frame with a 1 KiB window (the smallest a frame may declare, so its block +/// maximum is 1 KiB rather than 128 KiB) around one last block of `block_type` +/// whose header carries `size_field`. No content size. +fn frame_with_a_tiny_window(payload: &[u8], block_type: u32, size_field: u32) -> Vec { + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x00, // FHD: multi-segment, no content size + 0x00, // window descriptor: 1 KiB + ]; + let header = size_field << 3 | block_type << 1 | 1; + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(payload); + frame +} + /// A block with no sequences whose RLE literals regenerate 200,000 bytes: past /// the block maximum on literals alone. fn block_of_literals_past_the_block_maximum() -> Vec { @@ -1302,6 +1317,79 @@ fn trailing_literals_past_the_block_maximum_are_rejected() { ); } +/// A compressed block (one sequence, literals left over) followed by a Raw +/// block of a whole block maximum, in a frame declaring both: the per-block +/// ceiling the compressed block armed bounds sequence writes only, so the Raw +/// block that follows is bounded by the caller's slice and decodes. +#[test] +fn a_raw_block_after_a_compressed_one_fills_the_slice() { + const LITERALS: usize = 10; + const RAW: usize = 128 * 1024; + // Literals, then one sequence: literal length 1, repeat offset 1, match + // length 3, leaving 9 literals after it. 13 bytes out. + let mut block = literals_header_20_bit(0, LITERALS as u32).to_vec(); + block.extend((0..LITERALS).map(|i| b'a' + i as u8)); + block.extend_from_slice(&[ + 0x01, // one sequence + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x00, // LL code 1, OF code 0, ML code 0 + 0x01, // stream start bit + ]); + let compressed_output = 1 + 3 + (LITERALS - 1); + + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + frame.extend_from_slice(&((compressed_output + RAW) as u32).to_le_bytes()); + // Compressed block, not last. + let header = (block.len() as u32) << 3 | 2 << 1; + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(&block); + // Last block, Raw, a whole block maximum of it. + let header = (RAW as u32) << 3 | 1; + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend((0..RAW).map(|i| (i * 31) as u8)); + + let mut out = alloc::vec![0u8; compressed_output + RAW]; + let written = FrameDecoder::new() + .decode_all(&frame, &mut out) + .expect("a Raw block after a compressed one decodes"); + assert_eq!(written, compressed_output + RAW); + assert_eq!(&out[..4], b"aaaa"); // one literal, then the match of three + assert_eq!( + &out[compressed_output..compressed_output + 4], + &[0u8, 31, 62, 93] + ); +} + +/// A frame's block maximum is the smaller of its window and 128 KiB (RFC 8878 +/// 3.1.1.2.4), so a 1 KiB window bounds every block at 1 KiB: literals, a +/// block's whole output, and a Raw or RLE block's size alike. +#[test] +fn a_block_past_a_small_window_is_rejected() { + let mut literals = literals_header_20_bit(1, 2048).to_vec(); + literals.push(b'z'); // the repeated byte + literals.push(0x00); // no sequences + let size_field = literals.len() as u32; + assert_rejected_as_malformed( + &frame_with_a_tiny_window(&literals, 2, size_field), + "literals of 2 KiB in a 1 KiB window", + ); + + let raw: Vec = (0..2048u32).map(|i| i as u8).collect(); + assert_rejected_as_malformed( + &frame_with_a_tiny_window(&raw, 0, 2048), + "a Raw block of 2 KiB in a 1 KiB window", + ); + + assert_rejected_as_malformed( + &frame_with_a_tiny_window(b"z", 1, 2048), + "an RLE block of 2 KiB in a 1 KiB window", + ); +} + #[test] fn a_block_past_the_block_maximum_is_rejected_on_the_direct_path() { let frame = frame_with_a_block_past_the_block_maximum(None); diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index c8ca3cc2e..af8501ade 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -168,9 +168,13 @@ where buffer.reserve_exact(MAX_BLOCK_SIZE as usize); // Arm the per-block output ceiling so a malformed / adversarial block // whose sequences over-produce cannot grow the buffer past - // `len + MAX_BLOCK_SIZE` (a decompression-bomb OOM on the growable - // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. - buffer.set_block_output_ceiling(MAX_BLOCK_SIZE as usize); + // `len + block_maximum` (a decompression-bomb OOM on the growable + // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. The + // ceiling is the frame's block maximum, which a narrow window lowers + // below 128 KiB. + buffer.set_block_output_ceiling(crate::decoding::block_decoder::block_maximum( + buffer.window_size, + )); let old_buffer_size = buffer.len(); let num_sequences = section.num_sequences as usize; diff --git a/zstd/src/decoding/user_slice_buf.rs b/zstd/src/decoding/user_slice_buf.rs index 5454853bd..754644d4a 100644 --- a/zstd/src/decoding/user_slice_buf.rs +++ b/zstd/src/decoding/user_slice_buf.rs @@ -140,14 +140,16 @@ pub(crate) struct UserSliceBackend<'a> { tail: usize, /// Where sequence writes must stop: the slice's end, or sooner under the /// per-block output ceiling armed by `set_block_output_ceiling` before - /// each sequence section. A block may write at most `MAX_BLOCK_SIZE` - /// (RFC 8878 3.1.1.2.4) whatever room the caller's slice has, which for a - /// frame of unknown size is its only other bound. Folded into the one + /// each sequence section. A block may write at most its frame's block + /// maximum (RFC 8878 3.1.1.2.4) whatever room the caller's slice has, which + /// for a frame of unknown size is its only other bound. Folded into the one /// bound every sequence write already checks ([`BufferBackend::cap`]), as /// upstream folds `blockSizeMax` into `oend`, so the ceiling costs no - /// check of its own. Raw and RLE blocks write through `try_extend*`, - /// which keep the slice's end: the ceiling bounds sequences, as on - /// `RingBuffer`. + /// check of its own. Raw and RLE blocks never read it: they write through + /// `try_extend*`, which bound at the slice's end, and their own size is + /// held to the block maximum from their header before they write. So the + /// ceiling of the block that armed it cannot narrow a later Raw or RLE + /// block, and there is nothing to re-arm at their boundaries. sequence_cap: usize, } From c347e578b6a18aea44d57164b827255fc596ad82 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 08:54:20 +0300 Subject: [PATCH 08/29] fix(decoding): report short targets on literal-only blocks - A compressed block with no sequences wrote its literals through the infallible path, which asserts on a fixed-capacity backend. Literals within the block maximum can still be longer than the caller's slice, so a valid frame decoded into a short target aborted where it must return TargetTooSmall. The write is now fallible (new DecompressBlockError::LiteralsOutputOverflow), and the direct path maps it the way it maps a sequence overshoot: every entry to that path holds output.len() >= limit, so a write past the slice is a write past the limit. Regression test literals_longer_than_the_slice_are_target_too_small. DecodeBuffer::push is now test-only: nothing in the decoder writes output through a path that cannot report a short target. - The pre-block reservation asks for the frame's block maximum rather than a flat 128 KiB. A frame with a 1 KiB window got a 131,073-byte ring where its peak is 2 KiB, since the ring's growth limit only clamps a need that fits under it. Test a_compressed_block_in_a_small_window_reserves_one_block_of_it measured that ring before the change. - decode_from_to's return contract said read == 0 means the same input cannot advance. Since output is drained before each block, read == 0 with written > 0 means the target filled and the same input decodes further once there is room. The contract now says so; both counters zero is the no-progress signal. --- zstd/src/decoding/block_decoder.rs | 12 ++++- zstd/src/decoding/decode_buffer.rs | 16 +++--- zstd/src/decoding/errors.rs | 19 +++++++ zstd/src/decoding/frame_decoder.rs | 26 +++++++++- zstd/src/decoding/frame_decoder/tests.rs | 51 +++++++++++++++++++ zstd/src/decoding/sequence_section_decoder.rs | 27 +++++----- 6 files changed, 129 insertions(+), 22 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 7617bbfc3..eb4cb6cd3 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -466,7 +466,17 @@ impl BlockDecoder { }, )); } - buffer.push(literals_view); + // Fallible: on a fixed-capacity backend literals within the block + // maximum can still be longer than the caller's slice, which is a + // short target rather than a corrupt frame. The infallible `push` + // asserts there instead of reporting it. + buffer.try_push(literals_view).map_err(|overflow| { + DecompressBlockError::LiteralsOutputOverflow { + tail: overflow.tail, + requested: overflow.requested, + capacity: overflow.capacity, + } + })?; } // Nothing drains the buffer inside a block, so the growth of its live diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index c3df3ec7c..276474c3b 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -155,6 +155,16 @@ impl DecodeBuffer { } } + /// Infallible append, for tests over a growable backend. The decoder + /// writes through [`Self::try_push`]: on a fixed-capacity backend the + /// infallible path asserts where a short target must be reported. + #[cfg(test)] + #[inline] + pub fn push(&mut self, data: &[u8]) { + self.buffer.extend(data); + self.total_output_counter += data.len() as u64; + } + /// Enable or disable the drain-time XXH64 pass. Set by the frame layer /// from the decoder's [`ContentChecksum`](crate::decoding::ContentChecksum) /// mode before each decode (`false` for `None`). @@ -397,12 +407,6 @@ impl DecodeBuffer { Ok(()) } - #[inline] - pub fn push(&mut self, data: &[u8]) { - self.buffer.extend(data); - self.total_output_counter += data.len() as u64; - } - /// Add `n` to the cumulative produced-byte counter for output produced /// outside `push` / `repeat` — namely the inline `exec_sequence_inline` /// path, which writes through the backend directly and so bypasses the diff --git a/zstd/src/decoding/errors.rs b/zstd/src/decoding/errors.rs index 1d687f317..ec6592b23 100644 --- a/zstd/src/decoding/errors.rs +++ b/zstd/src/decoding/errors.rs @@ -271,6 +271,17 @@ pub enum DecompressBlockError { ExpandsPastBlockMaximum { size: usize, }, + /// A block with no sequences whose literals do not fit a fixed-capacity + /// backend: `requested` bytes at `tail` against `capacity`. The block is + /// within the block maximum, so this says the caller's slice is short, and + /// the frame decoder turns it into `TargetTooSmall` (or a content-size + /// mismatch for a frame that declared one). Growable backends grow instead + /// and never produce it. + LiteralsOutputOverflow { + tail: usize, + requested: usize, + capacity: usize, + }, } #[cfg(feature = "std")] @@ -313,6 +324,14 @@ impl core::fmt::Display for DecompressBlockError { "Block expands to {size} bytes, past the maximum of {}", crate::common::MAX_BLOCK_SIZE, ), + DecompressBlockError::LiteralsOutputOverflow { + tail, + requested, + capacity, + } => write!( + f, + "Literals would write past the output buffer: tail={tail}, requested={requested}, capacity={capacity}" + ), } } } diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 7f918909d..0b3d312d8 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -2380,8 +2380,11 @@ impl FrameDecoder { /// By all means use decode_blocks if you have a io.Reader available. This is just for compatibility with other decompressors /// which try to serve an old-style c api /// - /// Returns (read, written), if read == 0 then the source did not contain a full block and further calls with the same - /// input will not make any progress! + /// Returns (read, written). Both zero means the call made no progress: the + /// source holds no full block and the buffer no drainable output, so the + /// same input cannot advance. `read == 0` with `written > 0` is progress of + /// the other kind: `target` filled from output already buffered, and the + /// same input decodes further once the caller offers more room. /// /// Note that no kind of block can be bigger than 128kb. /// So to be safe use at least 128*1024 (max block content size) + 3 (block_header size) + 18 (max frame_header size) bytes as your source buffer @@ -3514,6 +3517,25 @@ impl FrameDecoder { let tail = direct.buffer.buffer_ref().tail() as u64; return Err(overflow(tail.saturating_add(requested))); } + // A no-sequence block's literals did not fit the slice. Every + // direct-path entry holds `output.len() >= limit` (a declared + // size is checked against the slice before the path is chosen, + // and an undeclared frame's limit IS the slice), so a write + // past the slice is a write past `limit`: the frame outgrew its + // declared size, or the caller's target is short. + Err(crate::decoding::errors::DecodeBlockContentError::DecompressBlockError( + crate::decoding::errors::DecompressBlockError::LiteralsOutputOverflow { + tail, + requested, + capacity, + }, + )) => { + debug_assert!( + capacity as u64 >= limit, + "direct path entered with a short slice" + ); + return Err(overflow((tail as u64).saturating_add(requested as u64))); + } Err(e) => { return Err(block_body_decode_error( e, diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 0c96d6ae9..c7b16eb7a 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1364,6 +1364,57 @@ fn a_raw_block_after_a_compressed_one_fills_the_slice() { ); } +/// A block can produce at most its frame's block maximum, so that is what the +/// pre-block reservation asks for. A 1 KiB window asking for a full 128 KiB +/// grew the ring to 256 KiB: the growth limit clamps a need that fits under it, +/// and this one did not. +#[test] +fn a_compressed_block_in_a_small_window_reserves_one_block_of_it() { + // Literals, then one sequence: literal length 1, repeat offset 1, match + // length 3, leaving 9 literals after it. + let mut block = literals_header_20_bit(0, 10).to_vec(); + block.extend((0..10u32).map(|i| b'a' + i as u8)); + block.extend_from_slice(&[ + 0x01, // one sequence + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x00, // LL code 1, OF code 0, ML code 0 + 0x01, // stream start bit + ]); + let size_field = block.len() as u32; + let frame = frame_with_a_tiny_window(&block, 2, size_field); + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = [0u8; 64]; + let (_, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + assert_eq!(&chunk[..written], b"aaaabcdefghij"); + let capacity = ring_capacity(&decoder); + assert!( + capacity <= 4 * 1024, + "a 1 KiB-window frame reserved {capacity} bytes of ring" + ); +} + +/// A compressed block with no sequences writes its literals straight to the +/// buffer. Into a slice shorter than they are, that must be `TargetTooSmall` +/// like any other overshoot, not the infallible write's capacity assert. +#[test] +fn literals_longer_than_the_slice_are_target_too_small() { + let mut block = literals_header_20_bit(1, 2000).to_vec(); + block.push(b'z'); // the repeated byte + block.push(0x00); // no sequences + let frame = frame_around_block(&block, None); + let mut out = alloc::vec![0u8; 100]; + let result = FrameDecoder::new().decode_all(&frame, &mut out); + assert!( + matches!(result, Err(super::FrameDecoderError::TargetTooSmall)), + "2000 literals into 100 bytes must be TargetTooSmall, got {result:?}" + ); +} + /// A frame's block maximum is the smaller of its window and 128 KiB (RFC 8878 /// 3.1.1.2.4), so a 1 KiB window bounds every block at 1 KiB: literals, a /// block's whole output, and a Raw or RLE block's size alike. diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index af8501ade..ff3162665 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -6,7 +6,6 @@ use crate::bit_io::BitReaderReversed; use crate::blocks::sequence_section::{ MAX_LITERAL_LENGTH_CODE, MAX_MATCH_LENGTH_CODE, MAX_OFFSET_CODE, }; -use crate::common::MAX_BLOCK_SIZE; use crate::cpu_kernel::CpuKernelTag; use crate::decoding::errors::{DecodeSequenceError, DecompressBlockError, ExecuteSequencesError}; use crate::decoding::sequence_execution::do_offset_history; @@ -85,7 +84,7 @@ pub(crate) struct SeqStreamSetup<'src, 'fse, K: crate::cpu_kernel::CpuKernel> { /// if the block's mode bytes call for it, skips the start-of-stream /// padding, initialises the LL/OF/ML decoder states, reserves the /// block's output capacity AND arms the per-block output ceiling (the -/// decompression-bomb guard that bounds growth at `len + MAX_BLOCK_SIZE`), +/// decompression-bomb guard that bounds growth at `len + block_maximum`), /// and computes the long-pipeline gate. /// /// Centralising this is what keeps the ceiling (and every other @@ -165,16 +164,19 @@ where // amortized policy would DOUBLE the window-sized buffer for a tail // worth a fraction of a block. The ring backend keeps its own // amortized growth via the trait default. - buffer.reserve_exact(MAX_BLOCK_SIZE as usize); + // Both the reservation and the ceiling are the frame's block maximum, + // which a narrow window lowers below 128 KiB: reserving a full 128 KiB for + // a frame whose window is 1 KiB gave it a 131,073-byte ring where its peak + // is 2 KiB, the growth limit only clamping a need that fits under it. + // Derived here rather than carried on frame state: one `min` against a + // block decode, and a cached copy would have to be reset with the window. + let block_maximum = crate::decoding::block_decoder::block_maximum(buffer.window_size); + buffer.reserve_exact(block_maximum); // Arm the per-block output ceiling so a malformed / adversarial block // whose sequences over-produce cannot grow the buffer past // `len + block_maximum` (a decompression-bomb OOM on the growable - // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. The - // ceiling is the frame's block maximum, which a narrow window lowers - // below 128 KiB. - buffer.set_block_output_ceiling(crate::decoding::block_decoder::block_maximum( - buffer.window_size, - )); + // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. + buffer.set_block_output_ceiling(block_maximum); let old_buffer_size = buffer.len(); let num_sequences = section.num_sequences as usize; @@ -611,7 +613,7 @@ pub(crate) fn decode_and_execute_sequences_impl< if remaining != 0 { // try_restore_checkpoint succeeds when no reallocation happened // between the checkpoint and now (the common case: upfront - // reserve(MAX_BLOCK_SIZE) covers a well-formed block). When a + // reserve of the block maximum covers a well-formed block). When a // malformed block decodes past that bound, reserve_amortized // fires and compacts the ring buffer — the captured tail is no // longer meaningful and the rollback is skipped. Either way the @@ -1007,8 +1009,7 @@ pub(crate) unsafe fn execute_one_sequence_pipelined_resolved_avx2< /// pipeline already issued a PREFETCH_L1 ADVANCE iterations earlier). /// The per-call `buffer.reserve(match_length)` is preserved by that /// variant — required for memory safety against malformed inputs whose -/// `match_length` exceeds the upfront `reserve(MAX_BLOCK_SIZE)` -/// headroom. +/// `match_length` exceeds the upfront block-maximum headroom. #[inline(always)] #[allow(dead_code)] // live on aarch64 + tests only; see decode_and_execute_sequences_impl pub(crate) fn execute_one_sequence_pipelined( @@ -1132,7 +1133,7 @@ pub(crate) fn execute_one_sequence_pipelined= 3` for any valid sequence). // The wildcopy helpers assert this in debug builds. - // - Caller's upfront `reserve(MAX_BLOCK_SIZE)` plus the + // - Caller's upfront block-maximum reserve plus the // `WILDCOPY_OVERLENGTH = 32` slack on the user slice // guarantees the writable tail has room for // `lit_length + match_length + 15` (max wildcopy From 7c336b2d959aa8f8da600611b0c2194108eabb58 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 09:05:04 +0300 Subject: [PATCH 09/29] perf(decoding): keep the literal-only write off the per-block body --- zstd/src/decoding/block_decoder.rs | 33 ++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index eb4cb6cd3..78a190f0a 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -60,6 +60,27 @@ fn block_fits_the_maximum( Ok(()) } +/// A block with no sequences: its literals ARE its output, written fallibly so +/// a fixed-capacity backend reports a short target instead of asserting. +/// +/// Out of line on purpose. Inlined into the per-block body, the error shape cost +/// 9.9% of cycles on a 1 MiB level-19 stream decode while issuing 0.6% FEWER +/// instructions: the grown body displaced the sequence path's layout. Behind a +/// call, the same bound is free. +#[inline(never)] +fn write_literals_only( + buffer: &mut crate::decoding::decode_buffer::DecodeBuffer, + literals: &[u8], +) -> Result<(), DecompressBlockError> { + buffer + .try_push(literals) + .map_err(|overflow| DecompressBlockError::LiteralsOutputOverflow { + tail: overflow.tail, + requested: overflow.requested, + capacity: overflow.capacity, + }) +} + /// Create a new [BlockDecoder]. pub fn new() -> BlockDecoder { BlockDecoder { @@ -466,17 +487,7 @@ impl BlockDecoder { }, )); } - // Fallible: on a fixed-capacity backend literals within the block - // maximum can still be longer than the caller's slice, which is a - // short target rather than a corrupt frame. The infallible `push` - // asserts there instead of reporting it. - buffer.try_push(literals_view).map_err(|overflow| { - DecompressBlockError::LiteralsOutputOverflow { - tail: overflow.tail, - requested: overflow.requested, - capacity: overflow.capacity, - } - })?; + write_literals_only(buffer, literals_view)?; } // Nothing drains the buffer inside a block, so the growth of its live From 58d840ef6c1d39877df6ce9a2319ce77472c2748 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 09:08:33 +0300 Subject: [PATCH 10/29] perf(decoding): keep the infallible literal write for growable backends --- zstd/src/decoding/block_decoder.rs | 11 ++++++++++- zstd/src/decoding/buffer_backend.rs | 8 ++++++++ zstd/src/decoding/decode_buffer.rs | 8 ++++---- zstd/src/decoding/user_slice_buf.rs | 4 ++++ 4 files changed, 26 insertions(+), 5 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 78a190f0a..b72396242 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -487,7 +487,16 @@ impl BlockDecoder { }, )); } - write_literals_only(buffer, literals_view)?; + // A growable backend allocates rather than refusing, so its write + // cannot fail and takes the infallible path; the compile-time const + // folds the other arm away, leaving the block body as the optimiser + // saw it before (the fallible form here cost 9.9% of cycles on a + // 1 MiB level-19 stream while issuing 0.6% fewer instructions). + if B::FIXED_CAPACITY { + write_literals_only(buffer, literals_view)?; + } else { + buffer.push(literals_view); + } } // Nothing drains the buffer inside a block, so the growth of its live diff --git a/zstd/src/decoding/buffer_backend.rs b/zstd/src/decoding/buffer_backend.rs index 63dc764e2..b38bd45cd 100644 --- a/zstd/src/decoding/buffer_backend.rs +++ b/zstd/src/decoding/buffer_backend.rs @@ -113,6 +113,14 @@ pub(crate) trait BufferBackend: Sized { /// const: the dispatch-site branch folds away per backend. const INLINE_EXEC_MAINTAINS_OUTPUT_COUNTER: bool = true; + /// Whether a write can fail for want of room. `false` for the growable + /// backends, which allocate instead: their write sites take the infallible + /// path and the fallible arm is dead-eliminated, keeping the code the + /// optimiser sees over a hot block body unchanged. `UserSliceBackend` + /// overrides it, because the caller's slice cannot grow and a write past it + /// is an error to report rather than an assert to trip. + const FIXED_CAPACITY: bool = false; + /// Upstream zstd's `ZSTD_execSequence` body /// (zstd_decompress_block.c:1008-1105). Writes `lit_length` bytes /// from `lit_src` at the current tail, then writes `match_length` diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index 276474c3b..c93ff7eef 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -155,10 +155,10 @@ impl DecodeBuffer { } } - /// Infallible append, for tests over a growable backend. The decoder - /// writes through [`Self::try_push`]: on a fixed-capacity backend the - /// infallible path asserts where a short target must be reported. - #[cfg(test)] + /// Infallible append, for a backend that grows rather than refusing + /// ([`BufferBackend::FIXED_CAPACITY`] `== false`). On a fixed-capacity + /// backend the write asserts where a short target must be reported, so + /// those paths take [`Self::try_push`]. #[inline] pub fn push(&mut self, data: &[u8]) { self.buffer.extend(data); diff --git a/zstd/src/decoding/user_slice_buf.rs b/zstd/src/decoding/user_slice_buf.rs index 754644d4a..dbc19491f 100644 --- a/zstd/src/decoding/user_slice_buf.rs +++ b/zstd/src/decoding/user_slice_buf.rs @@ -250,6 +250,10 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { /// path). See the trait const's doc. const INLINE_EXEC_MAINTAINS_OUTPUT_COUNTER: bool = false; + /// The caller's slice cannot grow: a write past it is reported, never + /// asserted. See the trait const's doc. + const FIXED_CAPACITY: bool = true; + /// Upstream zstd `ZSTD_execSequence` body — see trait doc for /// preconditions / contract. #[cfg(target_arch = "x86_64")] From bf45e99b11f10b7fd163f3b624618454bc57f2bf Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 09:15:22 +0300 Subject: [PATCH 11/29] perf(decoding): hand the literal-only write over as a tail call --- zstd/src/decoding/block_decoder.rs | 34 ++++++++++++++++-------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index b72396242..27a62db9f 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -60,18 +60,22 @@ fn block_fits_the_maximum( Ok(()) } -/// A block with no sequences: its literals ARE its output, written fallibly so -/// a fixed-capacity backend reports a short target instead of asserting. +/// A block with no sequences: its literals ARE its output. Out of line, so the +/// per-block body keeps the shape the sequence executor is laid out around. /// -/// Out of line on purpose. Inlined into the per-block body, the error shape cost -/// 9.9% of cycles on a 1 MiB level-19 stream decode while issuing 0.6% FEWER -/// instructions: the grown body displaced the sequence path's layout. Behind a -/// call, the same bound is free. +/// A growable backend allocates rather than refusing, so its write cannot fail +/// and takes the infallible path; the compile-time const folds the other arm +/// away. A fixed-capacity backend reports a short target, where the infallible +/// write would assert. #[inline(never)] fn write_literals_only( buffer: &mut crate::decoding::decode_buffer::DecodeBuffer, literals: &[u8], ) -> Result<(), DecompressBlockError> { + if !B::FIXED_CAPACITY { + buffer.push(literals); + return Ok(()); + } buffer .try_push(literals) .map_err(|overflow| DecompressBlockError::LiteralsOutputOverflow { @@ -487,16 +491,14 @@ impl BlockDecoder { }, )); } - // A growable backend allocates rather than refusing, so its write - // cannot fail and takes the infallible path; the compile-time const - // folds the other arm away, leaving the block body as the optimiser - // saw it before (the fallible form here cost 9.9% of cycles on a - // 1 MiB level-19 stream while issuing 0.6% fewer instructions). - if B::FIXED_CAPACITY { - write_literals_only(buffer, literals_view)?; - } else { - buffer.push(literals_view); - } + // The literals ARE this block's output, and their length was held + // to the block maximum above, so the post-block check below has + // nothing left to say: hand the write over and return its result. + // A tail call rather than `?` on purpose. Carrying the fallible + // write's error path through this body cost 9.9% of cycles on a + // 1 MiB level-19 stream while issuing 0.6% FEWER instructions: the + // sequence executor it calls is laid out around this body. + return write_literals_only(buffer, literals_view); } // Nothing drains the buffer inside a block, so the growth of its live From 193e2f9820f369301218da0aada8a93cbe815a09 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 09:20:29 +0300 Subject: [PATCH 12/29] perf(decoding): arm the block ceiling where the block maximum is known --- zstd/src/decoding/block_decoder.rs | 12 ++++++++++ zstd/src/decoding/sequence_section_decoder.rs | 23 ++++--------------- 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 27a62db9f..9c80e8a7e 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -473,6 +473,18 @@ impl BlockDecoder { // (immutable view into block_content_buffer) can coexist // with the mutable borrows on the FSE / decode-buffer / // offset-hist fields. + // Room for this block's output, and the ceiling that bounds it. + // Exact growth: the reservation is a no-op while the frame-entry + // window reservation covers it, and on the frame's last block (a + // tail worth a fraction of a block) the amortized policy would + // DOUBLE a window-sized buffer. The ceiling is what stops a + // malformed block's sequences from growing the buffer past + // `len + block_maximum` (a decompression-bomb OOM on the growable + // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. + // Both belong here, where the block maximum is already in hand: the + // arithmetic then stays out of the per-kernel sequence monomorphs. + buffer.reserve_exact(block_maximum); + buffer.set_block_output_ceiling(block_maximum); decode_and_execute_sequences( &seq_section, raw, diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index ff3162665..a06959b22 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -158,25 +158,10 @@ where "sequence section update bits exceed 56-bit budget" ); - // Exact growth: this worst-case pre-block reservation is a no-op while - // the frame-entry window reservation covers it, and on the frame's LAST - // block (where the remaining content is smaller than a full block) the - // amortized policy would DOUBLE the window-sized buffer for a tail - // worth a fraction of a block. The ring backend keeps its own - // amortized growth via the trait default. - // Both the reservation and the ceiling are the frame's block maximum, - // which a narrow window lowers below 128 KiB: reserving a full 128 KiB for - // a frame whose window is 1 KiB gave it a 131,073-byte ring where its peak - // is 2 KiB, the growth limit only clamping a need that fits under it. - // Derived here rather than carried on frame state: one `min` against a - // block decode, and a cached copy would have to be reset with the window. - let block_maximum = crate::decoding::block_decoder::block_maximum(buffer.window_size); - buffer.reserve_exact(block_maximum); - // Arm the per-block output ceiling so a malformed / adversarial block - // whose sequences over-produce cannot grow the buffer past - // `len + block_maximum` (a decompression-bomb OOM on the growable - // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. - buffer.set_block_output_ceiling(block_maximum); + // The block's output room is reserved and its ceiling armed by the block + // decoder before it calls in: that is where the frame's block maximum is + // known, and keeping the arithmetic out of this body keeps it out of the + // per-kernel monomorphs this function is inlined into. let old_buffer_size = buffer.len(); let num_sequences = section.num_sequences as usize; From 7df1aa4044153b538320d4470ebfd089048de606 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 10:10:23 +0300 Subject: [PATCH 13/29] fix(decoding): stop a full target from pulling in another block - A drain that fills the caller's target exactly leaves nothing pending, which the block loop read as room to decode another block: its output had nowhere to go and its input was consumed for a caller that asked for no more. The loop now stops on a full target too. Regression test a_filled_target_stops_before_the_next_block read 3081 bytes of a three-block frame before the change and reads 2054 after. - The block-maximum error carried the global 128 KiB constant while the check compares against the frame's own maximum, so a 2 KiB block in a 1 KiB-window frame was reported as expanding past 131072. The variant carries the frame's maximum now and prints that. - A frame that declares no size takes the single-raw-block shortcut as well. The probe parsed the first block header and then failed a condition that could never hold for such a frame, leaving the general loop to parse the same header again. The shortcut checks the block maximum, which the general path checks and the shortcut previously did not. - CPU kernel detection moved to the decoder's entry: the block decoder takes a resolved kernel, so a chunked decode no longer reads the detection cache once per call. The detecting constructor is now test-only. --- zstd/src/decoding/block_decoder.rs | 29 +++++++++++++++---- zstd/src/decoding/errors.rs | 9 +++--- zstd/src/decoding/frame_decoder.rs | 36 +++++++++++++++++++----- zstd/src/decoding/frame_decoder/tests.rs | 31 ++++++++++++++++++++ 4 files changed, 88 insertions(+), 17 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 9c80e8a7e..8eefdb6e0 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -6,7 +6,9 @@ use super::super::blocks::sequence_section::SequencesHeader; use super::literals_section_decoder::{LiteralsView, decode_literals_zerocopy}; use super::sequence_section_decoder::decode_and_execute_sequences; use crate::common::MAX_BLOCK_SIZE; -use crate::cpu_kernel::{CpuKernelTag, detect_cpu_kernel}; +use crate::cpu_kernel::CpuKernelTag; +#[cfg(any(test, feature = "bench-internals"))] +use crate::cpu_kernel::detect_cpu_kernel; use crate::decoding::errors::DecodeSequenceError; use crate::decoding::errors::{ BlockHeaderReadError, BlockSizeError, BlockTypeError, DecodeBlockContentError, @@ -52,9 +54,10 @@ fn block_fits_the_maximum( window_size: usize, ) -> Result<(), DecodeBlockContentError> { let size = header.decompressed_size as usize; - if size > block_maximum(window_size) { + let maximum = block_maximum(window_size); + if size > maximum { return Err(DecodeBlockContentError::DecompressBlockError( - DecompressBlockError::ExpandsPastBlockMaximum { size }, + DecompressBlockError::ExpandsPastBlockMaximum { size, maximum }, )); } Ok(()) @@ -85,12 +88,22 @@ fn write_literals_only( }) } -/// Create a new [BlockDecoder]. +/// Create a new [BlockDecoder], detecting the CPU kernel. Detection belongs at +/// the decoder's entry, so the decode paths take [`with_kernel`] instead; this +/// is for callers that decode a block in isolation. +#[cfg(any(test, feature = "bench-internals"))] pub fn new() -> BlockDecoder { + with_kernel(detect_cpu_kernel()) +} + +/// Create a new [BlockDecoder] over a kernel the caller already resolved. A +/// decoder that builds one per call (a chunked decode does) detects once and +/// passes the tag here, rather than reading the detection cache every time. +pub(crate) fn with_kernel(kernel: CpuKernelTag) -> BlockDecoder { BlockDecoder { internal_state: DecoderState::ReadyToDecodeNextHeader, header_buffer: [0u8; 3], - kernel: detect_cpu_kernel(), + kernel, } } @@ -386,6 +399,7 @@ impl BlockDecoder { if section.regenerated_size as usize > block_maximum { return Err(DecompressBlockError::ExpandsPastBlockMaximum { size: section.regenerated_size as usize, + maximum: block_maximum, }); } let raw = &raw[bytes_in_literals_header as usize..]; @@ -517,7 +531,10 @@ impl BlockDecoder { // length is this block's output. let produced = buffer.len() - len_before; if produced > block_maximum { - return Err(DecompressBlockError::ExpandsPastBlockMaximum { size: produced }); + return Err(DecompressBlockError::ExpandsPastBlockMaximum { + size: produced, + maximum: block_maximum, + }); } Ok(()) } diff --git a/zstd/src/decoding/errors.rs b/zstd/src/decoding/errors.rs index ec6592b23..5aeb0a368 100644 --- a/zstd/src/decoding/errors.rs +++ b/zstd/src/decoding/errors.rs @@ -267,9 +267,11 @@ pub enum DecompressBlockError { ExecuteSequencesError(ExecuteSequencesError), /// The block's literals, or its whole output, run past the block maximum /// (RFC 8878 3.1.1.2.4, `Block_Maximum_Size`): `size` bytes where - /// `MAX_BLOCK_SIZE` is the most a block may produce. + /// `maximum` is the most a block of this frame may produce, the smaller of + /// its window and 128 KiB. ExpandsPastBlockMaximum { size: usize, + maximum: usize, }, /// A block with no sequences whose literals do not fit a fixed-capacity /// backend: `requested` bytes at `tail` against `capacity`. The block is @@ -319,10 +321,9 @@ impl core::fmt::Display for DecompressBlockError { DecompressBlockError::SequencesHeaderParseError(e) => write!(f, "{e:?}"), DecompressBlockError::DecodeSequenceError(e) => write!(f, "{e:?}"), DecompressBlockError::ExecuteSequencesError(e) => write!(f, "{e:?}"), - DecompressBlockError::ExpandsPastBlockMaximum { size } => write!( + DecompressBlockError::ExpandsPastBlockMaximum { size, maximum } => write!( f, - "Block expands to {size} bytes, past the maximum of {}", - crate::common::MAX_BLOCK_SIZE, + "Block expands to {size} bytes, past this frame's maximum of {maximum}" ), DecompressBlockError::LiteralsOutputOverflow { tail, diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 0b3d312d8..4ebda2509 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -203,6 +203,11 @@ pub struct FrameDecoder { /// `all(lsm, hash)` (see `per_block_checksums_enabled`). #[cfg(all(feature = "lsm", feature = "hash"))] computed_block_checksums: alloc::vec::Vec, + /// Best kernel this CPU offers, resolved once per decoder rather than per + /// block decoder built. A chunked decode enters `decode_from_to` once per + /// caller-sized target, so detecting there put an atomic read and a branch + /// on every call; feature detection belongs before the work, not inside it. + kernel: crate::cpu_kernel::CpuKernelTag, } /// How the decoder treats a frame's optional XXH64 content checksum @@ -1066,6 +1071,7 @@ impl FrameDecoder { per_block_checksums_enabled: false, #[cfg(all(feature = "lsm", feature = "hash"))] computed_block_checksums: alloc::vec::Vec::new(), + kernel: crate::cpu_kernel::detect_cpu_kernel(), } } @@ -1826,7 +1832,7 @@ impl FrameDecoder { // is already sufficient. state.reserve_decoding_buffer(); - let mut block_dec = decoding::block_decoder::new(); + let mut block_dec = decoding::block_decoder::with_kernel(self.kernel); let buffer_size_before = state.decoder_scratch.buffer_len(); let block_counter_before = state.block_counter; @@ -2138,7 +2144,7 @@ impl FrameDecoder { start_block }; - let mut block_dec = decoding::block_decoder::new(); + let mut block_dec = decoding::block_decoder::with_kernel(self.kernel); // Bytes of prefix-window output that physically precede the first // in-range block in the buffer. Captured at the prefix → in-range @@ -2412,11 +2418,12 @@ impl FrameDecoder { //pseudo block to scope "state" so we can borrow self again after the block { + let kernel = self.kernel; let state = match &mut self.state { Some(s) => s, None => panic!("Bug in library"), }; - let mut block_dec = decoding::block_decoder::new(); + let mut block_dec = decoding::block_decoder::with_kernel(kernel); // Honour the content-checksum mode on this hand-rolled decode // loop (it does not go through `decode_blocks`): hash only when @@ -2465,7 +2472,11 @@ impl FrameDecoder { .buffer_read_reporting_pending(&mut target[written..]) .map_err(err::FailedToDrainDecodebuffer)?; written += read; - if pending > 0 { + // Stop on a full target as well as on output left behind: + // a drain that empties the buffer into the last of `target` + // leaves nothing pending, and decoding another block then + // consumes input the caller cannot be handed the output of. + if pending > 0 || written == target.len() { break; } //check if there are enough bytes for the next header @@ -3269,6 +3280,7 @@ impl FrameDecoder { None => err::TargetTooSmall, }; + let kernel = self.kernel; let state = self .state .as_mut() @@ -3285,12 +3297,22 @@ impl FrameDecoder { // the 1-block copy, dominates. { let mut probe = *input; - let mut header_dec = block_decoder::new(); + let mut header_dec = block_decoder::with_kernel(kernel); if let Ok((bh, hsize)) = header_dec.read_block_header(&mut probe) { let n = bh.decompressed_size as usize; + // A frame that declares no size takes the shortcut too: the + // slice is its limit, and holding the block means holding the + // frame. Without this the probe parsed the header that the + // general loop below parses again, on the very path (a small + // frame from a streamed producer) this decode is for. The + // block maximum is checked here as the general path checks it: + // a block past it is malformed, and the shortcut must not be + // the way around that. + let window = state.frame_header.window_size().unwrap_or(0) as usize; if bh.last_block && matches!(bh.block_type, crate::blocks::block::BlockType::Raw) - && declared_size == Some(n as u64) + && declared_size.is_none_or(|declared| declared == n as u64) + && n <= block_decoder::block_maximum(window) && probe.len() >= n && output.len() >= n { @@ -3397,7 +3419,7 @@ impl FrameDecoder { // sync with `decode_blocks` so post-call accessors // (`bytes_read_from_source`, `blocks_decoded`) return // accurate values. - let mut block_dec = block_decoder::new(); + let mut block_dec = block_decoder::with_kernel(kernel); // Track total output bytes against the declared // `frame_content_size` via the buffer's actual write // counter — `BlockHeader.decompressed_size` is 0 for diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index c7b16eb7a..b21e61725 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1364,6 +1364,37 @@ fn a_raw_block_after_a_compressed_one_fills_the_slice() { ); } +/// A drain that fills `target` exactly leaves nothing pending, which is not a +/// reason to decode another block: its output would have nowhere to go, and its +/// input would be consumed for a caller that asked for no more. +#[test] +fn a_filled_target_stops_before_the_next_block() { + const BLOCK: u32 = 1024; + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x00, // FHD: multi-segment, no content size + 0x00, // window descriptor: 1 KiB + ]; + for i in 0..3u32 { + // Raw block header: last flag on the third, type 0, size. + let header = BLOCK << 3 | u32::from(i == 2); + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend((0..BLOCK).map(|b| (b + i) as u8)); + } + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = alloc::vec![0u8; BLOCK as usize]; + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + assert_eq!(written, BLOCK as usize); + // Two blocks fill the window and hand one block over; the third is left + // for the next call, with its header and body unread. + assert_eq!(read, 2 * (3 + BLOCK as usize)); +} + /// A block can produce at most its frame's block maximum, so that is what the /// pre-block reservation asks for. A 1 KiB window asking for a full 128 KiB /// grew the ring to 256 KiB: the growth limit clamps a need that fits under it, From e5e57b66a13b47da264021c1063cfaea7bc9fad1 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 10:19:02 +0300 Subject: [PATCH 14/29] refactor(decoding): resolve the kernel once, drop the side detects --- zstd/src/bit_io/bit_reader_reverse.rs | 97 +----------- zstd/src/cpu_kernel.rs | 19 +++ zstd/src/huff0/huff0_decoder.rs | 206 ++------------------------ zstd/src/huff0/huff0_decoder/tests.rs | 45 +++--- 4 files changed, 59 insertions(+), 308 deletions(-) diff --git a/zstd/src/bit_io/bit_reader_reverse.rs b/zstd/src/bit_io/bit_reader_reverse.rs index befcab966..e5a05ec6f 100644 --- a/zstd/src/bit_io/bit_reader_reverse.rs +++ b/zstd/src/bit_io/bit_reader_reverse.rs @@ -24,72 +24,6 @@ const BIT_MASK: [u64; 65] = { table }; -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -#[derive(Copy, Clone)] -struct TripleExtractDispatch { - use_pext: bool, -} - -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -static TRIPLE_EXTRACT_DISPATCH: OnceLock = OnceLock::new(); - -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -#[inline(always)] -fn should_use_pext(vendor: [u8; 12], family: u32) -> bool { - vendor != *b"AuthenticAMD" || family != 0x17 -} - -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -#[inline(always)] -fn triple_extract_dispatch() -> &'static TripleExtractDispatch { - TRIPLE_EXTRACT_DISPATCH.get_or_init(detect_triple_extract_dispatch) -} - -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -fn detect_triple_extract_dispatch() -> TripleExtractDispatch { - use core::arch::x86_64::__cpuid; - use std::arch::is_x86_feature_detected; - - if !is_x86_feature_detected!("bmi2") { - return TripleExtractDispatch { use_pext: false }; - } - - // AMD Zen1/Zen2 execute PEXT/PDEP through a slow microcode path. - // Keep scalar extraction there and enable PEXT on Intel and newer AMD. - let leaf0 = __cpuid(0); - let mut vendor = [0u8; 12]; - vendor[0..4].copy_from_slice(&leaf0.ebx.to_le_bytes()); - vendor[4..8].copy_from_slice(&leaf0.edx.to_le_bytes()); - vendor[8..12].copy_from_slice(&leaf0.ecx.to_le_bytes()); - let eax = __cpuid(1).eax; - let base_family = (eax >> 8) & 0xF; - let ext_family = (eax >> 20) & 0xFF; - let family = if base_family == 0xF { - base_family + ext_family - } else { - base_family - }; - - TripleExtractDispatch { - use_pext: should_use_pext(vendor, family), - } -} - -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] -#[target_feature(enable = "bmi2")] -unsafe fn extract_triple_pext(all_three: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { - use core::arch::x86_64::_pext_u64; - - let mask3 = BIT_MASK[n3 as usize]; - let mask2 = BIT_MASK[n2 as usize].wrapping_shl(u32::from(n3)); - let mask1 = BIT_MASK[n1 as usize].wrapping_shl(u32::from(n2) + u32::from(n3)); - - let val1 = _pext_u64(all_three, mask1); - let val2 = _pext_u64(all_three, mask2); - let val3 = _pext_u64(all_three, mask3); - (val1, val2, val3) -} - /// Zstandard encodes some types of data in a way that the data must be read /// back to front to decode it properly. `BitReaderReversed` provides a /// convenient interface to do that. @@ -144,16 +78,6 @@ pub struct BitReaderReversed<'s, K: CpuKernel = ScalarKernel> { /// drives monomorphisation of methods that route through `K::mask_lower_bits` /// without forcing the struct itself to carry runtime kernel state. _kernel: PhantomData, - - /// Cached `triple_extract_dispatch().use_pext` snapshot, populated - /// once in `new()`. `peek_bits_triple` reads this field instead of - /// re-checking the global `OnceLock` on every sequence — the - /// per-call atomic load + dispatch-branch was paying ~3 cycles on - /// every sequence decode (thousands per block × many blocks per - /// frame). One bool per `BitReaderReversed` lifetime, amortised - /// across every `peek_bits_triple` in the same decode pass. - #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] - pub(crate) use_pext_triple: bool, } impl<'s, K: CpuKernel> BitReaderReversed<'s, K> { @@ -196,8 +120,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> { bit_container: 0, extra_bits: 0, _kernel: PhantomData, - #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] - use_pext_triple: triple_extract_dispatch().use_pext, } } @@ -366,21 +288,10 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> { let shift_by = (64u8 - self.bits_consumed).wrapping_sub(sum); let all_three = self.bit_container.wrapping_shr(shift_by as u32); - #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] - if self.use_pext_triple { - // SAFETY: `use_pext_triple` was set in `new()` from - // `triple_extract_dispatch().use_pext`, which only returns - // `true` when BMI2 is runtime-detected; the unsafe call is - // gated on the same runtime check that the inline-form - // `try_extract_triple_with_pext` used to perform per-call. - return unsafe { extract_triple_pext(all_three, n1, n2, n3) }; - } - - let val1 = K::mask_lower_bits(all_three.wrapping_shr(u32::from(n3) + u32::from(n2)), n1); - let val2 = K::mask_lower_bits(all_three.wrapping_shr(u32::from(n3)), n2); - let val3 = K::mask_lower_bits(all_three, n3); - - (val1, val2, val3) + // The kernel was chosen where this decode was dispatched, so the split + // is the monomorph's own instruction sequence. The reader used to carry + // the choice as a flag and branch on it here, once per sequence. + K::extract_triple(all_three, n1, n2, n3) } /// BMI2-scoped variant of [`peek_bits`]. The whole body executes diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index 9a6040014..bcbd820bf 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -61,6 +61,25 @@ pub trait CpuKernel: Copy + 'static { /// per-stream table builders pin to `n <= MAX_*_BITS` at /// construction time; no per-call wrapper assert runs. fn mask_lower_bits(value: u64, n: u8) -> u64; + + /// Split the low `n1 + n2 + n3` bits of `packed` into three fields, the + /// highest first. The FSE sequence decoder reads its three state updates + /// this way, once per sequence. + /// + /// The default is three [`Self::mask_lower_bits`]; a kernel whose hardware + /// extracts them in one instruction overrides it. Every implementation + /// returns the same three values, so which one ran is invisible to the + /// stream being decoded. + /// + /// Precondition: `n1 + n2 + n3 <= 64`, as for `mask_lower_bits`. + #[inline(always)] + fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { + ( + Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3) + u32::from(n2)), n1), + Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3)), n2), + Self::mask_lower_bits(packed, n3), + ) + } } /// Scalar fallback — portable, no SIMD or BMI2 intrinsics. Selected diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index 3c2ef0338..c1f4becd5 100644 --- a/zstd/src/huff0/huff0_decoder.rs +++ b/zstd/src/huff0/huff0_decoder.rs @@ -4,144 +4,22 @@ use crate::bit_io::BitReaderReversed; use crate::decoding::errors::HuffmanTableError; use crate::fse::{FSEDecoder, FSETable}; use alloc::vec::Vec; -#[cfg(target_arch = "x86")] -use core::arch::x86::_bzhi_u32; -#[cfg(target_arch = "x86_64")] -use core::arch::x86_64::_bzhi_u64; -#[cfg(all(feature = "std", target_arch = "aarch64"))] -use std::arch::is_aarch64_feature_detected; -#[cfg(all(feature = "std", any(target_arch = "x86", target_arch = "x86_64")))] -use std::arch::is_x86_feature_detected; -#[cfg(feature = "std")] -use std::sync::OnceLock; /// The Zstandard specification limits the maximum length of a code to 11 bits. pub(crate) const MAX_MAX_NUM_BITS: u8 = 11; -#[derive(Copy, Clone, Debug, Eq, PartialEq)] -pub(crate) enum HuffmanDecodeKernel { - Scalar, - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - X86Bmi2, - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - X86Avx2, - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - X86Vbmi2, - #[cfg(target_arch = "aarch64")] - Aarch64Neon, - #[cfg(target_arch = "aarch64")] - Aarch64Sve, -} - -#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] -#[inline(always)] -const fn select_x86_huffman_decode_kernel( - has_avx512vbmi2: bool, - has_avx512f: bool, - has_avx512vl: bool, - has_avx512bw: bool, - has_bmi2: bool, - has_avx2: bool, -) -> HuffmanDecodeKernel { - if has_avx512vbmi2 && has_avx512f && has_avx512vl && has_avx512bw && has_bmi2 { - return HuffmanDecodeKernel::X86Vbmi2; - } - if has_avx2 && has_bmi2 { - return HuffmanDecodeKernel::X86Avx2; - } - if has_bmi2 { - return HuffmanDecodeKernel::X86Bmi2; - } - HuffmanDecodeKernel::Scalar -} - -#[cfg(feature = "std")] -#[inline(always)] -pub(crate) fn detect_huffman_decode_kernel() -> HuffmanDecodeKernel { - static KERNEL: OnceLock = OnceLock::new(); - *KERNEL.get_or_init(|| { - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - { - let kernel = select_x86_huffman_decode_kernel( - is_x86_feature_detected!("avx512vbmi2"), - is_x86_feature_detected!("avx512f"), - is_x86_feature_detected!("avx512vl"), - is_x86_feature_detected!("avx512bw"), - is_x86_feature_detected!("bmi2"), - is_x86_feature_detected!("avx2"), - ); - if kernel != HuffmanDecodeKernel::Scalar { - return kernel; - } - } - #[cfg(target_arch = "aarch64")] - { - if is_aarch64_feature_detected!("sve") { - return HuffmanDecodeKernel::Aarch64Sve; - } - if is_aarch64_feature_detected!("neon") { - return HuffmanDecodeKernel::Aarch64Neon; - } - } - HuffmanDecodeKernel::Scalar - }) -} - -#[cfg(not(feature = "std"))] -#[inline(always)] -pub(crate) fn detect_huffman_decode_kernel() -> HuffmanDecodeKernel { - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - { - let kernel = select_x86_huffman_decode_kernel( - cfg!(target_feature = "avx512vbmi2"), - cfg!(target_feature = "avx512f"), - cfg!(target_feature = "avx512vl"), - cfg!(target_feature = "avx512bw"), - cfg!(target_feature = "bmi2"), - cfg!(target_feature = "avx2"), - ); - if kernel != HuffmanDecodeKernel::Scalar { - return kernel; - } - } - #[cfg(target_arch = "aarch64")] - { - if cfg!(target_feature = "sve") { - return HuffmanDecodeKernel::Aarch64Sve; - } - if cfg!(target_feature = "neon") { - return HuffmanDecodeKernel::Aarch64Neon; - } - } - HuffmanDecodeKernel::Scalar -} - pub struct HuffmanDecoder<'table> { table: &'table HuffmanTable, - /// Read by `decode_symbol_and_advance` on x86 to pick between the - /// scalar and BMI2 single-symbol decode bodies (single-stream tail - /// loop after the 4-stream burst). On aarch64 and portable targets - /// the BMI2 arm doesn't exist and the field is unread — the - /// 4-stream SIMD-fallback path that previously consumed this - /// field now dispatches via the [`HufKernel`] trait at - /// `decompress_literals` entry instead. - #[cfg_attr( - not(any(target_arch = "x86", target_arch = "x86_64")), - allow(dead_code) - )] - kernel: HuffmanDecodeKernel, /// State is used to index into the table. pub state: u64, } impl<'t> HuffmanDecoder<'t> { - /// Create a new decoder with the provided table + /// Create a new decoder with the provided table. It holds no kernel of its + /// own: the decode methods take the `K` their caller was monomorphised for, + /// which is resolved once where the decode is dispatched. pub fn new(table: &'t HuffmanTable) -> HuffmanDecoder<'t> { - HuffmanDecoder { - table, - kernel: detect_huffman_decode_kernel(), - state: 0, - } + HuffmanDecoder { table, state: 0 } } /// Decode the symbol the internal state (cursor) is pointed at and return the @@ -205,57 +83,20 @@ impl<'t> HuffmanDecoder<'t> { } /// Decode symbol and advance state in one table lookup. + /// + /// The kernel is `K`, chosen once where the decode was dispatched, so the + /// state advance is the monomorph's own instruction: `bzhi` on the BMI2 + /// tiers, a mask elsewhere. `state_mask` is `(1 << max_num_bits) - 1`, the + /// same value `bzhi` produces, so the two agree bit for bit. #[inline(always)] pub fn decode_symbol_and_advance( &mut self, br: &mut BitReaderReversed<'_, K>, - ) -> u8 { - // On x86 the BMI2 kernel uses `_bzhi_u64` and is a real - // perf win over the scalar `((state << n) & mask) | new_bits` - // sequence, so the runtime match is load-bearing. On aarch64 - // both NEON and SVE arms previously aliased the scalar body - // verbatim — the match was paying a 3-arm dispatch cost for - // zero benefit. Collapsed to a direct scalar call there. - // The enum's Aarch64Neon / Aarch64Sve variants are themselves - // cfg-gated to target_arch = "aarch64", so under the outer - // x86 cfg below they don't exist — the match here is - // exhaustive on Scalar + X86Bmi2/Avx2/Vbmi2 alone, and an - // inner `cfg(target_arch = "aarch64")` arm would be dead - // (outer x86 cfg already false on aarch64). - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - { - match self.kernel { - HuffmanDecodeKernel::Scalar => self.decode_symbol_and_advance_scalar(br), - HuffmanDecodeKernel::X86Bmi2 - | HuffmanDecodeKernel::X86Avx2 - | HuffmanDecodeKernel::X86Vbmi2 => { - // SAFETY: This path is selected only after runtime/static feature checks. - unsafe { self.decode_symbol_and_advance_x86_bmi2(br) } - } - } - } - #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] - { - // aarch64 and portable targets: the X86* arms compile out - // entirely, so the match would collapse to a single arm. - // Bypass the match and call scalar directly — both - // Aarch64Neon and Aarch64Sve specialisations were - // verbatim clones of the scalar body (they were dropped - // in an earlier commit), and no NEON/SVE intrinsics - // exist for the single-symbol decode shape. - self.decode_symbol_and_advance_scalar(br) - } - } - - #[inline(always)] - fn decode_symbol_and_advance_scalar( - &mut self, - br: &mut BitReaderReversed<'_, K>, ) -> u8 { let packed = self.table.packed_decode[self.state as usize]; let num_bits = (packed >> 8) as u8; let new_bits = br.get_bits(num_bits); - self.state = ((self.state << num_bits) & self.table.state_mask) | new_bits; + self.state = K::mask_lower_bits(self.state << num_bits, self.table.max_num_bits) | new_bits; packed as u8 } @@ -279,33 +120,6 @@ impl<'t> HuffmanDecoder<'t> { packed as u8 } - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - #[target_feature(enable = "bmi2")] - unsafe fn decode_symbol_and_advance_x86_bmi2( - &mut self, - br: &mut BitReaderReversed<'_, K>, - ) -> u8 { - let packed = self.table.packed_decode[self.state as usize]; - let num_bits = (packed >> 8) as u8; - let new_bits = br.get_bits(num_bits); - self.state = unsafe { self.advance_state_x86_bmi2(num_bits, new_bits) }; - packed as u8 - } - - #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] - #[target_feature(enable = "bmi2")] - unsafe fn advance_state_x86_bmi2(&self, num_bits: u8, new_bits: u64) -> u64 { - #[cfg(target_arch = "x86_64")] - { - _bzhi_u64(self.state << num_bits, u32::from(self.table.max_num_bits)) | new_bits - } - #[cfg(target_arch = "x86")] - { - let shifted = ((self.state << num_bits) & u64::from(u32::MAX)) as u32; - u64::from(_bzhi_u32(shifted, u32::from(self.table.max_num_bits))) | new_bits - } - } - // aarch64 NEON / SVE kernels for `decode_symbol_and_advance` were // identical clones of the scalar body — no NEON/SVE intrinsics // were ever in use here (the SIMD kernels live in `decode4_*` diff --git a/zstd/src/huff0/huff0_decoder/tests.rs b/zstd/src/huff0/huff0_decoder/tests.rs index 2dbf563f7..121952722 100644 --- a/zstd/src/huff0/huff0_decoder/tests.rs +++ b/zstd/src/huff0/huff0_decoder/tests.rs @@ -84,7 +84,6 @@ fn decode_symbol_and_advance_scalar_matches_manual_transition() { let mut decoder = HuffmanDecoder { table: &table, - kernel: HuffmanDecodeKernel::Scalar, state: initial_state, }; let mut br = @@ -95,23 +94,31 @@ fn decode_symbol_and_advance_scalar_matches_manual_transition() { assert_eq!(decoder.state, expected_state); } -#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +/// The state advance is the kernel's own instruction (`bzhi` where the tier +/// has it, a mask elsewhere) and the two must agree bit for bit, since the +/// stream they decode does not know which one ran. #[test] -fn select_x86_kernel_ordering_is_stable() { - assert_eq!( - select_x86_huffman_decode_kernel(true, true, true, true, true, true), - HuffmanDecodeKernel::X86Vbmi2 - ); - assert_eq!( - select_x86_huffman_decode_kernel(false, false, false, false, true, true), - HuffmanDecodeKernel::X86Avx2 - ); - assert_eq!( - select_x86_huffman_decode_kernel(false, false, false, false, true, false), - HuffmanDecodeKernel::X86Bmi2 - ); - assert_eq!( - select_x86_huffman_decode_kernel(false, false, false, false, false, true), - HuffmanDecodeKernel::Scalar - ); +fn every_kernel_advances_the_state_alike() { + let table = test_table(); + let source = [0b10101010, 0b01010101]; + + let mut scalar = HuffmanDecoder::new(&table); + let mut scalar_br = BitReaderReversed::::new(&source); + let scalar_symbol = scalar.decode_symbol_and_advance(&mut scalar_br); + + #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + if std::arch::is_x86_feature_detected!("bmi2") { + let mut bmi2 = HuffmanDecoder::new(&table); + let mut bmi2_br = BitReaderReversed::::new(&source); + assert_eq!(bmi2.decode_symbol_and_advance(&mut bmi2_br), scalar_symbol); + assert_eq!(bmi2.state, scalar.state); + } + + #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))] + { + let mut neon = HuffmanDecoder::new(&table); + let mut neon_br = BitReaderReversed::::new(&source); + assert_eq!(neon.decode_symbol_and_advance(&mut neon_br), scalar_symbol); + assert_eq!(neon.state, scalar.state); + } } From 7673672975b28f0effdaa034eac68dc07be849ce Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 10:21:22 +0300 Subject: [PATCH 15/29] refactor(decoding): drop the pext side-branch from the sequence readers --- zstd/src/bit_io/bit_reader_reverse.rs | 86 ++------------------------ zstd/src/decoding/seq_decoder_bmi2.rs | 7 +-- zstd/src/decoding/seq_decoder_vbmi2.rs | 7 +-- 3 files changed, 6 insertions(+), 94 deletions(-) diff --git a/zstd/src/bit_io/bit_reader_reverse.rs b/zstd/src/bit_io/bit_reader_reverse.rs index e5a05ec6f..cd837235c 100644 --- a/zstd/src/bit_io/bit_reader_reverse.rs +++ b/zstd/src/bit_io/bit_reader_reverse.rs @@ -1,18 +1,12 @@ use crate::cpu_kernel::{CpuKernel, ScalarKernel}; use core::convert::TryInto; use core::marker::PhantomData; -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -use std::sync::OnceLock; /// Pre-computed mask table: `BIT_MASK[n]` equals the lower `n` bits set, -/// i.e. `(1u64 << n) - 1` for `n` in `0..=64`. -/// -/// `mask_lower_bits` no longer reads this table — it computes the mask -/// via `u64::MAX >> (64 - n)` to save a load. The table is still used -/// by the BMI2 PEXT triple-extract path on x86-64 (where the mask is -/// constructed once per call and then fed to `_pext_u64`), and by the -/// tests that verify mask values directly. -#[cfg(any(test, all(target_arch = "x86_64", feature = "kernel-bmi2")))] +/// i.e. `(1u64 << n) - 1` for `n` in `0..=64`. Kept for the tests that verify +/// mask values directly; `mask_lower_bits` computes the mask instead of +/// loading it. +#[cfg(test)] const BIT_MASK: [u64; 65] = { let mut table = [0u64; 65]; let mut i: u32 = 1; @@ -86,32 +80,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> { self.index as isize * 8 + (64 - self.bits_consumed as isize) - self.extra_bits as isize } - /// Returns `true` when the cached vendor policy says PEXT is fast - /// on the running CPU (Intel + AMD Zen3+) and the bmi2-direct - /// triple-extract path should be used. AMD Zen1/Zen2 microcode - /// PEXT is slower than the scalar 3× shift+mask path, so - /// [`should_use_pext`] caches `false` for those vendors. - /// - /// `no_std` x86_64 builds lack the runtime detection (`use_pext_triple` - /// is std-gated), so this falls back to `true`: callers on - /// `no_std` rely on compile-time `target_feature = "bmi2"` and - /// implicitly trust that the chosen target CPU advertises fast - /// PEXT. Vendor-specific microcode regression remains a - /// build-time concern there — pin a known-good target with - /// `RUSTFLAGS="-C target-cpu=..."`. - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - #[inline(always)] - pub(crate) fn use_pext_triple_fast(&self) -> bool { - #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] - { - self.use_pext_triple - } - #[cfg(not(all(feature = "std", target_arch = "x86_64")))] - { - true - } - } - pub fn new(source: &'s [u8]) -> BitReaderReversed<'s, K> { BitReaderReversed { index: source.len(), @@ -322,52 +290,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> { core::arch::x86_64::_bzhi_u64(self.bit_container.wrapping_shr(shift_by as u32), n as u32) } - /// BMI2-scoped variant of [`peek_bits_triple`]. Mirrors the - /// scalar/K-trait variant but inlines `_pext_u64` directly instead - /// of crossing the `extract_triple_pext` CALL boundary. - /// - /// On AMD Zen1/Zen2 (vendor=AuthenticAMD family=0x17) `_pext_u64` - /// goes through slow microcode; callers should still consult - /// `self.use_pext_triple` (populated at construction from the - /// global dispatch cache) and route to the scalar variant on - /// those CPUs. This method assumes the caller already gated on - /// `use_pext_triple == true`. - /// - /// # Safety - /// Caller MUST ensure BMI2 is available AND the running CPU - /// benefits from `_pext_u64` (i.e. not Zen1/Zen2). - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - #[target_feature(enable = "bmi2")] - #[inline] - pub(crate) unsafe fn peek_bits_triple_bmi2( - &mut self, - sum: u8, - n1: u8, - n2: u8, - n3: u8, - ) -> (u64, u64, u64) { - debug_assert_eq!( - u16::from(sum), - u16::from(n1) + u16::from(n2) + u16::from(n3), - "peek_bits_triple_bmi2: sum ({}) must equal n1+n2+n3 ({}+{}+{})", - sum, - n1, - n2, - n3 - ); - debug_assert!( - sum == 0 || self.bits_consumed + sum <= 64, - "peek_bits_triple_bmi2: not enough bits (consumed={}, requested={})", - self.bits_consumed, - sum - ); - let shift_by = (64u8 - self.bits_consumed).wrapping_sub(sum); - let all_three = self.bit_container.wrapping_shr(shift_by as u32); - // SAFETY: caller's target_feature includes BMI2 per `# Safety` - // contract; same scope as the enclosing fn. - unsafe { extract_triple_pext(all_three, n1, n2, n3) } - } - /// Consume `n` bits from the source. #[inline(always)] pub fn consume(&mut self, n: u8) { diff --git a/zstd/src/decoding/seq_decoder_bmi2.rs b/zstd/src/decoding/seq_decoder_bmi2.rs index ff01005ea..0d6c85b5b 100644 --- a/zstd/src/decoding/seq_decoder_bmi2.rs +++ b/zstd/src/decoding/seq_decoder_bmi2.rs @@ -40,12 +40,7 @@ macro_rules! decode_one_body { let (obits, ml_add, ll_add) = if sum_wide <= 56 { let sum = sum_wide as u8; $br.ensure_bits(sum); - // SAFETY: enclosing fn is target_feature(bmi2). - let triple = if $br.use_pext_triple_fast() { - unsafe { $br.peek_bits_triple_bmi2(sum, of_num_bits, ml_num_bits, ll_num_bits) } - } else { - $br.peek_bits_triple(sum, of_num_bits, ml_num_bits, ll_num_bits) - }; + let triple = $br.peek_bits_triple(sum, of_num_bits, ml_num_bits, ll_num_bits); $br.consume(sum); triple } else { diff --git a/zstd/src/decoding/seq_decoder_vbmi2.rs b/zstd/src/decoding/seq_decoder_vbmi2.rs index e595ad2a9..92e10d701 100644 --- a/zstd/src/decoding/seq_decoder_vbmi2.rs +++ b/zstd/src/decoding/seq_decoder_vbmi2.rs @@ -38,12 +38,7 @@ macro_rules! decode_one_body { let (obits, ml_add, ll_add) = if sum_wide <= 56 { let sum = sum_wide as u8; $br.ensure_bits(sum); - // SAFETY: enclosing fn carries full VBMI2+AVX2+BMI2 scope. - let triple = if $br.use_pext_triple_fast() { - unsafe { $br.peek_bits_triple_bmi2(sum, of_num_bits, ml_num_bits, ll_num_bits) } - } else { - $br.peek_bits_triple(sum, of_num_bits, ml_num_bits, ll_num_bits) - }; + let triple = $br.peek_bits_triple(sum, of_num_bits, ml_num_bits, ll_num_bits); $br.consume(sum); triple } else { From 97de827f9704ce4f3b939a68b6891dcb2830f8f2 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 10:34:13 +0300 Subject: [PATCH 16/29] test(decoding): compare the kernels instead of the removed pext policy --- zstd/src/bit_io/bit_reader_reverse/tests.rs | 73 +++++++-------------- 1 file changed, 25 insertions(+), 48 deletions(-) diff --git a/zstd/src/bit_io/bit_reader_reverse/tests.rs b/zstd/src/bit_io/bit_reader_reverse/tests.rs index 07c3739e9..9588c5c52 100644 --- a/zstd/src/bit_io/bit_reader_reverse/tests.rs +++ b/zstd/src/bit_io/bit_reader_reverse/tests.rs @@ -59,21 +59,6 @@ fn mask_lower_bits(value: u64, n: u8) -> u64 { value & mask } } -// Used only by the in-file extract_triple correctness tests after -// `peek_bits_triple` switched to the per-reader `use_pext_triple` -// cached flag (commit 8805122f) — production now calls -// `extract_triple_pext` directly via that path. Gating with -// `#[cfg(test)]` keeps the helper available for the tests while -// avoiding a `dead_code` warning under `-D warnings`. -#[cfg(all(test, feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -#[inline(always)] -fn try_extract_triple_with_pext(all_three: u64, n1: u8, n2: u8, n3: u8) -> Option<(u64, u64, u64)> { - if !triple_extract_dispatch().use_pext { - return None; - } - - Some(unsafe { extract_triple_pext(all_three, n1, n2, n3) }) -} #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] use std::arch::is_x86_feature_detected; @@ -301,12 +286,11 @@ fn peek_bits_bmi2_matches_scalar() { } } -/// `peek_bits_triple_bmi2` MUST produce the same triple as the -/// scalar variant for every width combination the FSE/HUF decoders -/// can reach. +/// Every kernel reads the same triple out of the same bits: the stream being +/// decoded cannot tell which monomorph ran. #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] #[test] -fn peek_bits_triple_bmi2_matches_scalar() { +fn peek_bits_triple_agrees_across_kernels() { if !is_x86_feature_detected!("bmi2") { return; } @@ -326,33 +310,22 @@ fn peek_bits_triple_bmi2_matches_scalar() { for &(n1, n2, n3) in &widths { let sum = n1 + n2 + n3; let mut scalar = super::BitReaderReversed::::new(&data); - let mut bmi2 = super::BitReaderReversed::::new(&data); + let mut bmi2 = super::BitReaderReversed::::new(&data); scalar.ensure_bits(sum); bmi2.ensure_bits(sum); let s = scalar.peek_bits_triple(sum, n1, n2, n3); - // SAFETY: gated on `is_x86_feature_detected!("bmi2")` above. - let b = unsafe { bmi2.peek_bits_triple_bmi2(sum, n1, n2, n3) }; + let b = bmi2.peek_bits_triple(sum, n1, n2, n3); assert_eq!(s, b, "mismatch at widths=({},{},{})", n1, n2, n3); } } +/// The kernel's own `extract_triple` against a plain masking reference, over +/// the widths the FSE and HUF decoders can reach. #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] #[test] -fn should_use_pext_policy_table() { - let cases = [ - (*b"AuthenticAMD", 0x17, false), - (*b"AuthenticAMD", 0x19, true), - (*b"GenuineIntel", 0x06, true), - ]; +fn extract_triple_matches_the_reference_under_every_kernel() { + use crate::cpu_kernel::CpuKernel; - for (vendor, family, expected) in cases { - assert_eq!(super::should_use_pext(vendor, family), expected); - } -} - -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -#[test] -fn bmi2_triple_extract_matches_scalar_reference() { if !is_x86_feature_detected!("bmi2") { return; } @@ -384,12 +357,14 @@ fn bmi2_triple_extract_matches_scalar_reference() { for &(n1, n2, n3) in &widths { for &all_three in &fixed_values { let expected = scalar_extract_triple(all_three, n1, n2, n3); - let pext = unsafe { super::extract_triple_pext(all_three, n1, n2, n3) }; - assert_eq!(pext, expected); - - if let Some(dispatched) = try_extract_triple_with_pext(all_three, n1, n2, n3) { - assert_eq!(dispatched, expected); - } + assert_eq!( + crate::cpu_kernel::ScalarKernel::extract_triple(all_three, n1, n2, n3), + expected + ); + assert_eq!( + crate::cpu_kernel::Bmi2Kernel::extract_triple(all_three, n1, n2, n3), + expected + ); } } @@ -398,12 +373,14 @@ fn bmi2_triple_extract_matches_scalar_reference() { for _ in 0..64 { let all_three = next_test_value(&mut state); let expected = scalar_extract_triple(all_three, n1, n2, n3); - let pext = unsafe { super::extract_triple_pext(all_three, n1, n2, n3) }; - assert_eq!(pext, expected); - - if let Some(dispatched) = try_extract_triple_with_pext(all_three, n1, n2, n3) { - assert_eq!(dispatched, expected); - } + assert_eq!( + crate::cpu_kernel::ScalarKernel::extract_triple(all_three, n1, n2, n3), + expected + ); + assert_eq!( + crate::cpu_kernel::Bmi2Kernel::extract_triple(all_three, n1, n2, n3), + expected + ); } } } From e5d093c47ecd05fccde8be116dd1687fa6a131bd Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 10:59:21 +0300 Subject: [PATCH 17/29] fix(decoding): give 32-bit x86 its BMI2 tier back - The kernel tiers were all declared for x86_64, so a 32-bit x86 build resolved to the scalar bodies whatever the CPU offered, losing the bit-extract the Huffman state advance had before the dispatch was unified. The BMI2 tier now covers both widths (32-bit `bzhi` applies to the halves), the detection picks it there, and the literals dispatch has its arm. The sequence monolith stays portable on 32-bit: its bodies are x86_64-only. - The Huffman state advance takes the table's precomputed mask again through a kernel operation: a tier with a bit-extract instruction ignores the mask and takes the width, the others take the mask instead of rebuilding it per symbol. - The kernel-parity tests sweep every tier the dispatcher can select on the running build rather than stopping at the first, and the Huffman one starts from a nonzero state whose masked result is nonzero, so a tier that masked wrongly would fail it. Two review findings are answered in comments rather than code, both about the per-block ceiling over literal writes. The ceiling bounds SEQUENCE writes, which is why it is armed beside the sequence reserve; literal writes go through `try_extend`, bounded by the caller's slice, and the literals section was held to the block maximum where it was parsed. Arming the ceiling over those writes would reject valid frames: a small block leaves the ceiling near its own output, and a following literal-only block of a whole block maximum would exceed it while the slice still had room. Test a_literal_only_block_after_a_compressed_one_fills_the_slice decodes exactly that pair. --- zstd/src/bit_io/bit_reader_reverse/tests.rs | 90 +++++++++++++++-- zstd/src/cpu_kernel.rs | 96 ++++++++++++++++--- zstd/src/decoding/block_decoder.rs | 10 ++ zstd/src/decoding/frame_decoder/tests.rs | 45 +++++++++ zstd/src/decoding/literals_section_decoder.rs | 15 ++- zstd/src/decoding/sequence_section_decoder.rs | 23 +++++ zstd/src/huff0/huff0_decoder.rs | 11 ++- zstd/src/huff0/huff0_decoder/tests.rs | 66 ++++++++++--- 8 files changed, 318 insertions(+), 38 deletions(-) diff --git a/zstd/src/bit_io/bit_reader_reverse/tests.rs b/zstd/src/bit_io/bit_reader_reverse/tests.rs index 9588c5c52..8005f9604 100644 --- a/zstd/src/bit_io/bit_reader_reverse/tests.rs +++ b/zstd/src/bit_io/bit_reader_reverse/tests.rs @@ -307,15 +307,91 @@ fn peek_bits_triple_agrees_across_kernels() { (15, 16, 17), (5, 0, 4), ]; + /// Read the same widths from the same bits under one kernel. + macro_rules! triple_under { + ($kernel:ty, $sum:expr, $n1:expr, $n2:expr, $n3:expr) => {{ + let mut reader = super::BitReaderReversed::<$kernel>::new(&data); + reader.ensure_bits($sum); + reader.peek_bits_triple($sum, $n1, $n2, $n3) + }}; + } + for &(n1, n2, n3) in &widths { let sum = n1 + n2 + n3; - let mut scalar = super::BitReaderReversed::::new(&data); - let mut bmi2 = super::BitReaderReversed::::new(&data); - scalar.ensure_bits(sum); - bmi2.ensure_bits(sum); - let s = scalar.peek_bits_triple(sum, n1, n2, n3); - let b = bmi2.peek_bits_triple(sum, n1, n2, n3); - assert_eq!(s, b, "mismatch at widths=({},{},{})", n1, n2, n3); + let expected = triple_under!(crate::cpu_kernel::ScalarKernel, sum, n1, n2, n3); + assert_eq!( + triple_under!(crate::cpu_kernel::Bmi2Kernel, sum, n1, n2, n3), + expected, + "Bmi2Kernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + #[cfg(feature = "kernel-avx2")] + if is_x86_feature_detected!("avx2") { + assert_eq!( + triple_under!(crate::cpu_kernel::Avx2Kernel, sum, n1, n2, n3), + expected, + "Avx2Kernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + } + #[cfg(feature = "kernel-vbmi2")] + if is_x86_feature_detected!("avx512vbmi2") { + assert_eq!( + triple_under!(crate::cpu_kernel::Vbmi2Kernel, sum, n1, n2, n3), + expected, + "Vbmi2Kernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + } + } +} + +/// The aarch64 tiers read the same triple as the scalar bodies they share. +#[cfg(all(feature = "std", target_arch = "aarch64", feature = "kernel-neon"))] +#[test] +fn peek_bits_triple_agrees_across_kernels() { + let data: [u8; 16] = [ + 0xDE, 0xAD, 0xBE, 0xEF, 0x42, 0x13, 0x37, 0xCA, 0xFE, 0x01, 0x99, 0x88, 0x77, 0x66, 0x55, + 0x44, + ]; + let widths = [(0, 0, 0), (1, 1, 1), (3, 5, 7), (8, 8, 8), (15, 16, 17)]; + + macro_rules! triple_under { + ($kernel:ty, $sum:expr, $n1:expr, $n2:expr, $n3:expr) => {{ + let mut reader = super::BitReaderReversed::<$kernel>::new(&data); + reader.ensure_bits($sum); + reader.peek_bits_triple($sum, $n1, $n2, $n3) + }}; + } + + for &(n1, n2, n3) in &widths { + let sum = n1 + n2 + n3; + let expected = triple_under!(crate::cpu_kernel::ScalarKernel, sum, n1, n2, n3); + assert_eq!( + triple_under!(crate::cpu_kernel::NeonKernel, sum, n1, n2, n3), + expected, + "NeonKernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + #[cfg(feature = "kernel-sve")] + if std::arch::is_aarch64_feature_detected!("sve") { + assert_eq!( + triple_under!(crate::cpu_kernel::SveKernel, sum, n1, n2, n3), + expected, + "SveKernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + } } } diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index bcbd820bf..112247ea7 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -62,6 +62,15 @@ pub trait CpuKernel: Copy + 'static { /// construction time; no per-call wrapper assert runs. fn mask_lower_bits(value: u64, n: u8) -> u64; + /// [`Self::mask_lower_bits`] for a caller that already holds the mask, + /// `mask == (1 << n) - 1`: the HUF table keeps one per decoder. A kernel + /// with a bit-extract instruction ignores the mask and takes `n`; the + /// others take the mask and skip building it per call. + #[inline(always)] + fn mask_lower_bits_precomputed(value: u64, mask: u64, _n: u8) -> u64 { + value & mask + } + /// Split the low `n1 + n2 + n3` bits of `packed` into three fields, the /// highest first. The FSE sequence decoder reads its three state updates /// this way, once per sequence. @@ -109,16 +118,24 @@ impl CpuKernel for ScalarKernel { // FSE/HUF paths. A dedicated `Sse2Kernel` lands when `copy_chunk` moves onto // the trait. -/// x86_64 BMI2-only kernel: `_bzhi_u64` for mask_lower_bits. Selected -/// when the CPU has BMI2 but not the AVX2 SIMD width to upgrade to -/// the Avx2 kernel. Treated as a stepping stone between Sse2 and -/// Avx2 on hardware that has BMI2 but not AVX2 (rare in practice but -/// matches upstream zstd's gating). -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +/// BMI2-only kernel: `bzhi` for mask_lower_bits. Selected when the CPU has +/// BMI2 but not the AVX2 SIMD width to upgrade to the Avx2 kernel. Treated as +/// a stepping stone between Sse2 and Avx2 on hardware that has BMI2 but not +/// AVX2 (rare in practice but matches upstream zstd's gating). Present on +/// 32-bit x86 as well as x86_64: the instruction is there, only its width +/// differs, and without this tier a 32-bit build would decode on the scalar +/// bodies whatever the CPU offers. +#[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" +))] #[derive(Copy, Clone, Default)] pub(crate) struct Bmi2Kernel; -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" +))] impl CpuKernel for Bmi2Kernel { #[inline(always)] fn mask_lower_bits(value: u64, n: u8) -> u64 { @@ -129,6 +146,13 @@ impl CpuKernel for Bmi2Kernel { // running CPU. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + /// `bzhi` takes the width, so the caller's mask is not needed. + #[inline(always)] + fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { + // SAFETY: as for `mask_lower_bits`. + unsafe { mask_lower_bits_bmi2_impl(value, n) } + } } /// x86_64 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern @@ -216,7 +240,10 @@ impl CpuKernel for SveKernel { /// same shared body. With `#[inline]` LLVM inlines the call into /// any caller that itself has BMI2 in scope; outside that scope the /// target_feature boundary is preserved. -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" +))] #[target_feature(enable = "bmi2")] #[inline] unsafe fn mask_lower_bits_bmi2_impl(value: u64, n: u8) -> u64 { @@ -226,7 +253,25 @@ unsafe fn mask_lower_bits_bmi2_impl(value: u64, n: u8) -> u64 { // already covers it). SAFETY: caller selected a kernel whose // CpuKernelTag was resolved after `is_x86_feature_detected!("bmi2")` // returned true, so the BMI2 instruction set is available. - core::arch::x86_64::_bzhi_u64(value, n as u32) + #[cfg(target_arch = "x86_64")] + { + core::arch::x86_64::_bzhi_u64(value, n as u32) + } + // 32-bit x86 has `bzhi` on 32-bit registers only. Widths up to 32 take one + // instruction on the low half; wider ones keep the low 32 bits whole and + // apply it to the high half, which is what a 64-bit `bzhi` does in one go. + #[cfg(target_arch = "x86")] + { + use core::arch::x86::_bzhi_u32; + if n >= 64 { + return value; + } + if n <= 32 { + return u64::from(_bzhi_u32(value as u32, u32::from(n))); + } + let high = _bzhi_u32((value >> 32) as u32, u32::from(n) - 32); + (value & u64::from(u32::MAX)) | (u64::from(high) << 32) + } } /// Pure boolean-input variant of the x86 kernel-tag selection. Both the @@ -286,7 +331,12 @@ pub(crate) enum CpuKernelTag { Scalar, #[cfg(all(target_arch = "x86_64", feature = "kernel-sse"))] Sse2, - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + /// Reachable on 32-bit x86 as well: `bzhi` is there, and without the tier + /// such a build would decode on the scalar bodies whatever the CPU offers. + #[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" + ))] Bmi2, #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] Avx2, @@ -336,6 +386,20 @@ fn detect_cpu_kernel_uncached() -> CpuKernelTag { cfg!(feature = "kernel-sse") && is_x86_feature_detected!("sse2"), ); } + // 32-bit x86 carries only the BMI2 tier: the wider tiers' kernels and + // their `target_feature` bodies are x86_64-only, so there is nothing + // above `bzhi` to select here. + #[cfg(target_arch = "x86")] + { + #[cfg(feature = "kernel-bmi2")] + { + use std::arch::is_x86_feature_detected; + if is_x86_feature_detected!("bmi2") { + return CpuKernelTag::Bmi2; + } + } + return CpuKernelTag::Scalar; + } #[cfg(target_arch = "aarch64")] { #[cfg(any(feature = "kernel-sve", feature = "kernel-neon"))] @@ -376,6 +440,13 @@ pub(crate) fn detect_cpu_kernel() -> CpuKernelTag { cfg!(target_feature = "sse2"), ); } + #[cfg(target_arch = "x86")] + { + #[cfg(all(feature = "kernel-bmi2", target_feature = "bmi2"))] + { + return CpuKernelTag::Bmi2; + } + } #[cfg(target_arch = "aarch64")] { #[cfg(all(feature = "kernel-sve", target_feature = "sve"))] @@ -401,7 +472,10 @@ impl CpuKernelTag { CpuKernelTag::Scalar => "scalar", #[cfg(all(target_arch = "x86_64", feature = "kernel-sse"))] CpuKernelTag::Sse2 => "sse2", - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + #[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" + ))] CpuKernelTag::Bmi2 => "bmi2", #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] CpuKernelTag::Avx2 => "avx2", diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 8eefdb6e0..28d92d7d7 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -70,6 +70,16 @@ fn block_fits_the_maximum( /// and takes the infallible path; the compile-time const folds the other arm /// away. A fixed-capacity backend reports a short target, where the infallible /// write would assert. +/// +/// No per-block ceiling is armed for this write, and none is needed. The +/// ceiling bounds SEQUENCE writes, which is why it is armed beside the +/// sequence reserve; the write here goes through `try_extend`, whose bound is +/// the caller's slice. Arming the previous block's ceiling over it would +/// REJECT valid frames: a small block leaves the ceiling near its own output, +/// and a following literal-only block of a whole block maximum would exceed it +/// while the slice still had room. The literals were already held to the block +/// maximum where their section was parsed, so this write cannot exceed it +/// either. #[inline(never)] fn write_literals_only( buffer: &mut crate::decoding::decode_buffer::DecodeBuffer, diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index b21e61725..1e17bc5ce 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1429,6 +1429,51 @@ fn a_compressed_block_in_a_small_window_reserves_one_block_of_it() { ); } +/// A compressed block with sequences, then one with none whose literals fill a +/// whole block maximum. The ceiling the first block armed bounds sequence +/// writes only, so the literal write that follows is bounded by the caller's +/// slice and both blocks decode. +#[test] +fn a_literal_only_block_after_a_compressed_one_fills_the_slice() { + const RLE_LITERALS: u32 = 128 * 1024; + // Literals, then one sequence: literal length 1, repeat offset 1, match + // length 3, leaving 9 literals after it. 13 bytes out. + let mut first = literals_header_20_bit(0, 10).to_vec(); + first.extend((0..10u32).map(|i| b'a' + i as u8)); + first.extend_from_slice(&[ + 0x01, // one sequence + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x00, // LL code 1, OF code 0, ML code 0 + 0x01, // stream start bit + ]); + // A block of RLE literals and no sequences: a whole block maximum of them. + let mut second = literals_header_20_bit(1, RLE_LITERALS).to_vec(); + second.push(b'z'); // the repeated byte + second.push(0x00); // no sequences + + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + let content = 13 + RLE_LITERALS; + frame.extend_from_slice(&content.to_le_bytes()); + let header = (first.len() as u32) << 3 | 2 << 1; // compressed, not last + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(&first); + let header = (second.len() as u32) << 3 | 2 << 1 | 1; // compressed, last + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(&second); + + let mut out = alloc::vec![0u8; content as usize]; + let written = FrameDecoder::new() + .decode_all(&frame, &mut out) + .expect("a literal-only block after a compressed one decodes"); + assert_eq!(written, content as usize); + assert_eq!(&out[..4], b"aaaa"); // one literal, then the match of three + assert!(out[13..].iter().all(|&b| b == b'z')); +} + /// A compressed block with no sequences writes its literals straight to the /// buffer. Into a slice shorter than they are, that must be `TargetTooSmall` /// like any other overshoot, not the infallible write's capacity assert. diff --git a/zstd/src/decoding/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index 192fbd355..acd87114b 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -6,7 +6,10 @@ use super::scratch::HuffmanScratch; use crate::bit_io::BitReaderReversed; #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] use crate::cpu_kernel::Avx2Kernel; -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" +))] use crate::cpu_kernel::Bmi2Kernel; #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] use crate::cpu_kernel::Vbmi2Kernel; @@ -163,7 +166,10 @@ fn decompress_literals( CpuKernelTag::Avx2 => unsafe { decompress_literals_avx2(section, scratch, dict, source, target) }, - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + #[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" + ))] CpuKernelTag::Bmi2 => unsafe { decompress_literals_bmi2(section, scratch, dict, source, target) }, @@ -183,7 +189,10 @@ unsafe fn decompress_literals_avx2( decompress_literals_impl::(section, scratch, dict, source, target) } -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" +))] #[target_feature(enable = "bmi2")] unsafe fn decompress_literals_bmi2( section: &LiteralsSection, diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index a06959b22..453a81fd6 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -288,6 +288,20 @@ pub fn decode_and_execute_sequences<'fse, B: super::buffer_backend::BufferBacken dict, ) } + // 32-bit x86 reaches the BMI2 tier for the entropy tables (the HUF + // state advance takes `bzhi` through `K`), but the sequence monolith + // has no 32-bit body: its `target_feature` modules are x86_64-only. + // The portable walk is what runs here until one exists. + #[cfg(all(target_arch = "x86", feature = "kernel-bmi2"))] + CpuKernelTag::Bmi2 => super::seq_decoder_scalar::decode_and_execute_sequences_scalar::( + section, + source, + fse, + buffer, + offset_hist, + literals_buffer, + dict, + ), #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] CpuKernelTag::Bmi2 => { // SAFETY: `detect_cpu_kernel()` only returns Bmi2 when @@ -630,6 +644,15 @@ pub(crate) fn decode_and_execute_sequences_impl< // as `OutputBufferOverflow` instead of panicking via the per-call // `assert!` inside `BufferBackend::extend`. Growable backends // (FlatBuf, RingBuffer) accept the write infallibly. + // + // The per-block ceiling is NOT re-checked here on purpose. It bounds the + // match writes, whose length a malformed block controls; these bytes are + // literals, and the whole literals section was held to the block maximum + // where it was parsed, so the ceiling would find nothing the parse did not + // already reject. Reserving against it would be worse than redundant: on + // the direct path the ceiling is relative to the block that armed it, so a + // valid frame whose blocks differ in size would start failing. The block's + // total output is checked once it has decoded. if lit_cur < literals_buffer_len { let rest = &literals_buffer[lit_cur..]; buffer.try_push(rest).map_err(ExecuteSequencesError::from)?; diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index c1f4becd5..2e4acf07c 100644 --- a/zstd/src/huff0/huff0_decoder.rs +++ b/zstd/src/huff0/huff0_decoder.rs @@ -86,8 +86,9 @@ impl<'t> HuffmanDecoder<'t> { /// /// The kernel is `K`, chosen once where the decode was dispatched, so the /// state advance is the monomorph's own instruction: `bzhi` on the BMI2 - /// tiers, a mask elsewhere. `state_mask` is `(1 << max_num_bits) - 1`, the - /// same value `bzhi` produces, so the two agree bit for bit. + /// tiers, the table's `state_mask` elsewhere, which is built once per table + /// rather than per symbol. `state_mask == (1 << max_num_bits) - 1` is the + /// value `bzhi` produces, so the two agree bit for bit. #[inline(always)] pub fn decode_symbol_and_advance( &mut self, @@ -96,7 +97,11 @@ impl<'t> HuffmanDecoder<'t> { let packed = self.table.packed_decode[self.state as usize]; let num_bits = (packed >> 8) as u8; let new_bits = br.get_bits(num_bits); - self.state = K::mask_lower_bits(self.state << num_bits, self.table.max_num_bits) | new_bits; + self.state = K::mask_lower_bits_precomputed( + self.state << num_bits, + self.table.state_mask, + self.table.max_num_bits, + ) | new_bits; packed as u8 } diff --git a/zstd/src/huff0/huff0_decoder/tests.rs b/zstd/src/huff0/huff0_decoder/tests.rs index 121952722..d769fd56e 100644 --- a/zstd/src/huff0/huff0_decoder/tests.rs +++ b/zstd/src/huff0/huff0_decoder/tests.rs @@ -95,30 +95,68 @@ fn decode_symbol_and_advance_scalar_matches_manual_transition() { } /// The state advance is the kernel's own instruction (`bzhi` where the tier -/// has it, a mask elsewhere) and the two must agree bit for bit, since the -/// stream they decode does not know which one ran. +/// has it, the table's mask elsewhere) and every tier must agree bit for bit, +/// since the stream they decode does not know which one ran. Every kernel the +/// dispatcher can select on this build runs here, not just the first one. +/// +/// The state starts nonzero and the entry decodes fewer bits than the table's +/// width, so the masked value is nonzero too: a kernel that masked wrongly +/// would show it. #[test] fn every_kernel_advances_the_state_alike() { - let table = test_table(); + let mut table = test_table(); + // State 3 decoding one bit leaves `(3 << 1) & 0b11 == 0b10` behind, so the + // mask has something to keep and a kernel that masked wrongly would show. + table.packed_decode[3] = u16::from(b'D') | (1u16 << 8); let source = [0b10101010, 0b01010101]; + const START: u64 = 3; let mut scalar = HuffmanDecoder::new(&table); + scalar.state = START; let mut scalar_br = BitReaderReversed::::new(&source); let scalar_symbol = scalar.decode_symbol_and_advance(&mut scalar_br); + assert_ne!(scalar.state, 0, "the masked state must be nonzero"); - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - if std::arch::is_x86_feature_detected!("bmi2") { - let mut bmi2 = HuffmanDecoder::new(&table); - let mut bmi2_br = BitReaderReversed::::new(&source); - assert_eq!(bmi2.decode_symbol_and_advance(&mut bmi2_br), scalar_symbol); - assert_eq!(bmi2.state, scalar.state); + /// Run one kernel over the same bits from the same state and compare. + macro_rules! same_as_scalar { + ($kernel:ty) => {{ + let mut decoder = HuffmanDecoder::new(&table); + decoder.state = START; + let mut reader = BitReaderReversed::<$kernel>::new(&source); + assert_eq!( + decoder.decode_symbol_and_advance(&mut reader), + scalar_symbol, + "{} decoded another symbol", + stringify!($kernel) + ); + assert_eq!( + decoder.state, + scalar.state, + "{} advanced the state differently", + stringify!($kernel) + ); + }}; } + #[cfg(all( + any(target_arch = "x86", target_arch = "x86_64"), + feature = "kernel-bmi2" + ))] + if std::arch::is_x86_feature_detected!("bmi2") { + same_as_scalar!(crate::cpu_kernel::Bmi2Kernel); + } + #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] + if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("bmi2") { + same_as_scalar!(crate::cpu_kernel::Avx2Kernel); + } + #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] + if std::arch::is_x86_feature_detected!("avx512vbmi2") { + same_as_scalar!(crate::cpu_kernel::Vbmi2Kernel); + } #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))] - { - let mut neon = HuffmanDecoder::new(&table); - let mut neon_br = BitReaderReversed::::new(&source); - assert_eq!(neon.decode_symbol_and_advance(&mut neon_br), scalar_symbol); - assert_eq!(neon.state, scalar.state); + same_as_scalar!(crate::cpu_kernel::NeonKernel); + #[cfg(all(target_arch = "aarch64", feature = "kernel-sve", feature = "std"))] + if std::arch::is_aarch64_feature_detected!("sve") { + same_as_scalar!(crate::cpu_kernel::SveKernel); } } From da99e65c25090ab6ebfd02cf562630f87870122f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:02:04 +0300 Subject: [PATCH 18/29] fix(decoding): let the aarch64 tiers reach the literals monomorph --- zstd/src/decoding/literals_section_decoder.rs | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/zstd/src/decoding/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index acd87114b..53d89086e 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -11,6 +11,14 @@ use crate::cpu_kernel::Avx2Kernel; feature = "kernel-bmi2" ))] use crate::cpu_kernel::Bmi2Kernel; +#[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))] +use crate::cpu_kernel::NeonKernel; +#[cfg(all( + target_arch = "aarch64", + feature = "kernel-sve", + any(feature = "std", target_feature = "sve"), +))] +use crate::cpu_kernel::SveKernel; #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] use crate::cpu_kernel::Vbmi2Kernel; #[cfg(test)] @@ -173,6 +181,23 @@ fn decompress_literals( CpuKernelTag::Bmi2 => unsafe { decompress_literals_bmi2(section, scratch, dict, source, target) }, + // The aarch64 tiers need no `target_feature` wrapper: NEON is part of + // the baseline ABI there, and SVE reaches nothing in this pipeline yet. + // They take their own monomorph all the same, so a body that diverges + // from the scalar one later arrives here instead of being dispatched + // past. + #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))] + CpuKernelTag::Neon => { + decompress_literals_impl::(section, scratch, dict, source, target) + } + #[cfg(all( + target_arch = "aarch64", + feature = "kernel-sve", + any(feature = "std", target_feature = "sve"), + ))] + CpuKernelTag::Sve => { + decompress_literals_impl::(section, scratch, dict, source, target) + } _ => decompress_literals_impl::(section, scratch, dict, source, target), } } From a7ecf1a24e3d15b03accfc12cf74b64810a165a0 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:19:05 +0300 Subject: [PATCH 19/29] fix(decoding): reserve no more than the frame has left to give The per-block reservation asked for a whole block maximum on every compressed block, so a frame that declared less than a block still grew its ring to one. The initial reservation already capped the ring at the declared content; this grew it straight back. The reservation now asks for the smaller of a block and what the frame has left to produce, which is what a declared size means. The ceiling stays the block maximum: it decides whether a block is malformed, and a frame that outruns its declared size is caught by the size check, which says so rather than blaming the block. Regression test a_compressed_block_reserves_no_more_than_the_frame_declares decodes a 13-byte frame with a 1 MiB window: the ring held 131,086 bytes before this and holds under 4 KiB after. --- zstd/src/decoding/block_decoder.rs | 13 +++++++- zstd/src/decoding/decode_buffer.rs | 28 ++++++++++++++++ zstd/src/decoding/frame_decoder.rs | 17 ++++++++++ zstd/src/decoding/frame_decoder/tests.rs | 41 ++++++++++++++++++++++++ 4 files changed, 98 insertions(+), 1 deletion(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 28d92d7d7..735abb401 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -507,7 +507,18 @@ impl BlockDecoder { // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. // Both belong here, where the block maximum is already in hand: the // arithmetic then stays out of the per-kernel sequence monomorphs. - buffer.reserve_exact(block_maximum); + // Reserve what the block can actually produce: its maximum, or + // what the frame has left to give when it declared a size. A frame + // declaring 13 bytes cannot produce 128 KiB, and reserving that + // for it left the ring mostly unused for the frame's lifetime. The + // ceiling stays the block maximum: it decides whether a block is + // malformed, and a frame that outruns its declared size is caught + // by the size check instead, which says so. + let room = match buffer.remaining_declared() { + Some(left) => block_maximum.min(left), + None => block_maximum, + }; + buffer.reserve_exact(room); buffer.set_block_output_ceiling(block_maximum); decode_and_execute_sequences( &seq_section, diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index c93ff7eef..4dca88892 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -32,6 +32,8 @@ pub struct DecodeBuffer { // refcount bump). The borrow checker guarantees the dictionary outlives // every read; `DecodeBuffer` itself stays `Send`/`Sync` by auto-derive. pub window_size: usize, + /// See [`DecodeBuffer::set_declared_content`]. + declared_content: Option, total_output_counter: u64, #[cfg(feature = "hash")] pub(crate) hash: twox_hash::XxHash64, @@ -117,6 +119,7 @@ impl DecodeBuffer { DecodeBuffer { buffer, window_size, + declared_content: None, total_output_counter: 0, #[cfg(feature = "hash")] hash: twox_hash::XxHash64::with_seed(0), @@ -145,6 +148,7 @@ impl DecodeBuffer { DecodeBuffer { buffer, window_size, + declared_content: None, total_output_counter: 0, #[cfg(feature = "hash")] hash: twox_hash::XxHash64::with_seed(0), @@ -207,8 +211,32 @@ impl DecodeBuffer { self.buffer.set_max_capacity(ceiling); } + /// What the frame says it will produce in total, when it says so. The + /// per-block reservation asks for no more than what is left of it: a frame + /// declaring less than a block cannot produce one, and reserving a whole + /// block for it leaves the ring mostly unused for the frame's lifetime. + /// `None` for a frame of unknown size, where a block is all we know. + #[inline] + pub(crate) fn set_declared_content(&mut self, content_size: Option) { + self.declared_content = content_size; + } + + /// Bytes the frame may still produce, for a frame that declared a size. + #[inline] + pub(crate) fn remaining_declared(&self) -> Option { + self.declared_content.map(|declared| { + // Saturating on purpose: a frame that has produced more than it + // declared is malformed, and the size check that rejects it runs + // where the block finishes. The answer here is just "nothing left + // worth reserving for". + let left = declared.saturating_sub(self.total_output_counter); + usize::try_from(left).unwrap_or(usize::MAX) + }) + } + pub fn reset(&mut self, window_size: usize) { self.window_size = window_size; + self.declared_content = None; self.buffer.clear(); self.buffer.set_growth_limit(peak_buffered_len(window_size)); // No reserve here: capacity decisions are pushed up to the frame diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index 4ebda2509..f72636cde 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -527,6 +527,16 @@ impl DecoderScratchKind { /// chunk, and streaming callers invoke it per call) from growing a /// window-full buffer toward 2x window, while per-block growth keeps /// the amortized `reserve`. + /// Hand the buffer what the frame declared it would produce, so the + /// per-block reservation can stop at the frame's remainder. + #[inline] + fn set_declared_content(&mut self, content_size: Option) { + match self { + Self::Ring(s) => s.buffer.set_declared_content(content_size), + Self::Flat(s) => s.buffer.set_declared_content(content_size), + } + } + #[inline] fn reserve_buffer(&mut self, target: usize, growth_limit: usize) { let window_size = target; @@ -912,6 +922,13 @@ impl FrameDecoderState { fn reserve_decoding_buffer(&mut self) { let target = self.decoding_buffer_size(); let growth_limit = self.decoding_buffer_limit(); + // What the frame promised to produce, so the per-block reservation can + // ask for the smaller of a block and what is left of that promise. + let declared = self + .frame_header + .fcs_declared() + .then(|| self.frame_header.frame_content_size()); + self.decoder_scratch.set_declared_content(declared); self.decoder_scratch.reserve_buffer(target, growth_limit); } diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 1e17bc5ce..7af47ffa5 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1429,6 +1429,47 @@ fn a_compressed_block_in_a_small_window_reserves_one_block_of_it() { ); } +/// A frame that declares 13 bytes cannot produce a block of 128 KiB, so the +/// per-block reservation asks for what is left of the frame instead. The ring +/// is otherwise grown to a block maximum by the first compressed block, whatever +/// the frame said it would produce. +#[test] +fn a_compressed_block_reserves_no_more_than_the_frame_declares() { + // Literals, then one sequence: literal length 1, repeat offset 1, match + // length 3, leaving 9 literals after it. 13 bytes out. + let mut block = literals_header_20_bit(0, 10).to_vec(); + block.extend((0..10u32).map(|i| b'a' + i as u8)); + block.extend_from_slice(&[ + 0x01, // one sequence + 0x54, // LL, OF and ML all RLE + 0x01, 0x00, 0x00, // LL code 1, OF code 0, ML code 0 + 0x01, // stream start bit + ]); + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + frame.extend_from_slice(&13u32.to_le_bytes()); + let header = (block.len() as u32) << 3 | 2 << 1 | 1; // compressed, last + frame.extend_from_slice(&header.to_le_bytes()[..3]); + frame.extend_from_slice(&block); + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut chunk = [0u8; 64]; + let (_, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + assert_eq!(&chunk[..written], b"aaaabcdefghij"); + let capacity = ring_capacity(&decoder); + assert!( + capacity < 4 * 1024, + "a 13-byte frame reserved {capacity} bytes of ring" + ); +} + /// A compressed block with sequences, then one with none whose literals fill a /// whole block maximum. The ceiling the first block armed bounds sequence /// writes only, so the literal write that follows is bounded by the caller's From 00997e9c2799a7003316630924b2e837a99170ca Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:22:00 +0300 Subject: [PATCH 20/29] perf(decoding): keep the block-reservation arithmetic off the block body --- zstd/src/decoding/block_decoder.rs | 15 ++++----------- zstd/src/decoding/decode_buffer.rs | 13 +++++++++++++ 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/zstd/src/decoding/block_decoder.rs b/zstd/src/decoding/block_decoder.rs index 735abb401..544398722 100644 --- a/zstd/src/decoding/block_decoder.rs +++ b/zstd/src/decoding/block_decoder.rs @@ -507,18 +507,11 @@ impl BlockDecoder { // RingBuffer); `DecodeBuffer::repeat` rejects the crossing match. // Both belong here, where the block maximum is already in hand: the // arithmetic then stays out of the per-kernel sequence monomorphs. - // Reserve what the block can actually produce: its maximum, or - // what the frame has left to give when it declared a size. A frame - // declaring 13 bytes cannot produce 128 KiB, and reserving that - // for it left the ring mostly unused for the frame's lifetime. The - // ceiling stays the block maximum: it decides whether a block is - // malformed, and a frame that outruns its declared size is caught + // The reservation stops at what the frame has left to produce; the + // ceiling stays the block maximum, since it decides whether a block + // is malformed and a frame that outruns its declared size is caught // by the size check instead, which says so. - let room = match buffer.remaining_declared() { - Some(left) => block_maximum.min(left), - None => block_maximum, - }; - buffer.reserve_exact(room); + buffer.reserve_for_block(block_maximum); buffer.set_block_output_ceiling(block_maximum); decode_and_execute_sequences( &seq_section, diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index 4dca88892..7cafdecb0 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -221,6 +221,19 @@ impl DecodeBuffer { self.declared_content = content_size; } + /// Room for one block's output: its maximum, or what the frame has left to + /// produce when it declared a size. A frame declaring less than a block + /// cannot produce one, and reserving a whole block for it leaves the buffer + /// mostly unused for the frame's lifetime. + #[inline] + pub(crate) fn reserve_for_block(&mut self, block_maximum: usize) { + let room = match self.remaining_declared() { + Some(left) => block_maximum.min(left), + None => block_maximum, + }; + self.reserve_exact(room); + } + /// Bytes the frame may still produce, for a frame that declared a size. #[inline] pub(crate) fn remaining_declared(&self) -> Option { From a98fd0b7b871b35268c9744a01b4dbc94b770e4d Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:24:45 +0300 Subject: [PATCH 21/29] perf(decoding): keep the declared-size field off the hot field layout --- zstd/src/decoding/decode_buffer.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zstd/src/decoding/decode_buffer.rs b/zstd/src/decoding/decode_buffer.rs index 7cafdecb0..40c6316f4 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -32,8 +32,6 @@ pub struct DecodeBuffer { // refcount bump). The borrow checker guarantees the dictionary outlives // every read; `DecodeBuffer` itself stays `Send`/`Sync` by auto-derive. pub window_size: usize, - /// See [`DecodeBuffer::set_declared_content`]. - declared_content: Option, total_output_counter: u64, #[cfg(feature = "hash")] pub(crate) hash: twox_hash::XxHash64, @@ -50,6 +48,10 @@ pub struct DecodeBuffer { /// without that flag ever being set. #[cfg(feature = "hash")] hash_dirty: bool, + /// See [`DecodeBuffer::set_declared_content`]. Last on purpose: the fields + /// above are read per sequence, and moving them would shift the layout the + /// decode monolith is built around. + declared_content: Option, } /// Rollback token produced by [`DecodeBuffer::checkpoint`]. From 70ba71b425f559773e7edf7ed018e46b04b4382f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:30:21 +0300 Subject: [PATCH 22/29] perf(decoding): take the low-bit mask from a table, not a guarded shift --- zstd/src/cpu_kernel.rs | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index 112247ea7..c79002008 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -96,18 +96,28 @@ pub trait CpuKernel: Copy + 'static { #[derive(Copy, Clone, Default)] pub struct ScalarKernel; +/// `BIT_MASK[n]` is the low `n` bits set for `n` in `0..=64`, and all bits for +/// anything past that (a width the formats cannot ask for). +/// +/// A table rather than `u64::MAX >> (64 - n)`: the shift form needs a guard for +/// `n == 0`, since a 64-bit shift is undefined, and that guard is a branch or a +/// `cmov` on every field of every sequence. Indexed by a `u8`, and sized for +/// every `u8`, so the load carries no bounds check either. The widths in use +/// are small, so the hot part is the first few cache lines of it. +pub(crate) const BIT_MASK: [u64; 256] = { + let mut table = [u64::MAX; 256]; + let mut i: usize = 0; + while i < 64 { + table[i] = (1u64 << i) - 1; + i += 1; + } + table +}; + impl CpuKernel for ScalarKernel { #[inline(always)] fn mask_lower_bits(value: u64, n: u8) -> u64 { - // `checked_shr` returns `None` for shift counts >= 64, which - // happens exactly when `n == 0` (`64 - 0 = 64`). Mapping - // both that case and the invalid `n > 64` underflow to 0 - // gives the mathematically-correct empty mask for n=0 and - // a safe-ish fallback for the invalid range. - let mask = u64::MAX - .checked_shr(64u32.wrapping_sub(n as u32)) - .unwrap_or(0); - value & mask + value & BIT_MASK[n as usize] } } From f3d68e0127bfbe8bfce50b749af0c6de922c0eae Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:37:16 +0300 Subject: [PATCH 23/29] perf(decoding): take pext back as a kernel operation --- zstd/src/cpu_kernel.rs | 120 +++++++++++++++++- zstd/src/decoding/literals_section_decoder.rs | 13 +- zstd/src/decoding/seq_decoder_bmi2.rs | 9 +- zstd/src/decoding/sequence_section_decoder.rs | 27 +++- .../sequence_section_decoder/tests.rs | 5 +- zstd/src/decoding/simd_copy.rs | 4 +- 6 files changed, 168 insertions(+), 10 deletions(-) diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index c79002008..8bb197d9d 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -121,6 +121,28 @@ impl CpuKernel for ScalarKernel { } } +/// One `pext` per field, the widths turned into masks from [`BIT_MASK`]. +/// +/// # Safety +/// The caller's kernel was selected after BMI2 was detected, and on hardware +/// where `pext` is not microcoded: see [`Bmi2SlowPextKernel`]. +#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[target_feature(enable = "bmi2")] +#[inline] +unsafe fn extract_triple_pext_impl(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { + use core::arch::x86_64::_pext_u64; + + let mask3 = BIT_MASK[n3 as usize]; + let mask2 = BIT_MASK[n2 as usize].wrapping_shl(u32::from(n3)); + let mask1 = BIT_MASK[n1 as usize].wrapping_shl(u32::from(n2) + u32::from(n3)); + + ( + _pext_u64(packed, mask1), + _pext_u64(packed, mask2), + _pext_u64(packed, mask3), + ) +} + // The SSE2 tier exists in `CpuKernelTag` (it carries the 128-bit copy-chunk // choice for the unified copy dispatch) but needs no `CpuKernel` ZST yet: the // only trait method, `mask_lower_bits`, has no SSE2-specific form (SSE2 has no @@ -163,6 +185,38 @@ impl CpuKernel for Bmi2Kernel { // SAFETY: as for `mask_lower_bits`. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + #[cfg(target_arch = "x86_64")] + #[inline(always)] + fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { + // SAFETY: as for `mask_lower_bits`, and the tier is only selected on + // hardware whose `pext` is not microcoded. + unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } + } +} + +/// x86_64 BMI2 kernel for hardware whose `pext` runs through microcode (AMD +/// Zen 1 and Zen 2, where it takes around 18 cycles against one for `bzhi`). +/// Identical to [`Bmi2Kernel`] except that the three-field extract keeps the +/// trait's mask form. The vendor question is answered once, where the CPU is +/// detected; below that there is a kernel, not a flag to branch on. +#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +#[derive(Copy, Clone, Default)] +pub(crate) struct Bmi2SlowPextKernel; + +#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] +impl CpuKernel for Bmi2SlowPextKernel { + #[inline(always)] + fn mask_lower_bits(value: u64, n: u8) -> u64 { + // SAFETY: as for `Bmi2Kernel::mask_lower_bits`. + unsafe { mask_lower_bits_bmi2_impl(value, n) } + } + + #[inline(always)] + fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { + // SAFETY: as for `mask_lower_bits`. + unsafe { mask_lower_bits_bmi2_impl(value, n) } + } } /// x86_64 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern @@ -181,6 +235,14 @@ impl CpuKernel for Avx2Kernel { // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + #[cfg(feature = "kernel-bmi2")] + #[inline(always)] + fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { + // SAFETY: as for `mask_lower_bits`; the tier is not selected on + // hardware whose `pext` is microcoded. + unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } + } } /// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU @@ -199,6 +261,14 @@ impl CpuKernel for Vbmi2Kernel { // at runtime before this kernel is instantiated. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + #[cfg(feature = "kernel-bmi2")] + #[inline(always)] + fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { + // SAFETY: as for `mask_lower_bits`; the tier is not selected on + // hardware whose `pext` is microcoded. + unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } + } } /// aarch64 NEON baseline kernel. Used on all aarch64 hardware that @@ -348,6 +418,10 @@ pub(crate) enum CpuKernelTag { feature = "kernel-bmi2" ))] Bmi2, + /// BMI2 hardware whose `pext` is microcoded (AMD Zen 1 and Zen 2). Takes + /// the mask form of the three-field extract; see [`Bmi2SlowPextKernel`]. + #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + Bmi2SlowPext, #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] Avx2, #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] @@ -377,6 +451,29 @@ pub(crate) fn detect_cpu_kernel() -> CpuKernelTag { *CACHED.get_or_init(detect_cpu_kernel_uncached) } +/// Whether this CPU runs `pext` in hardware. AMD Zen 1 and Zen 2 (family 0x17) +/// microcode it at around 18 cycles, where three masked shifts are faster; +/// every other vendor and family executes it in one. +#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] +fn fast_pext() -> bool { + use core::arch::x86_64::__cpuid; + + // Leaves 0 and 1 are architectural on every x86_64, so these are safe. + let leaf0 = __cpuid(0); + let mut vendor = [0u8; 12]; + vendor[0..4].copy_from_slice(&leaf0.ebx.to_le_bytes()); + vendor[4..8].copy_from_slice(&leaf0.edx.to_le_bytes()); + vendor[8..12].copy_from_slice(&leaf0.ecx.to_le_bytes()); + let eax = __cpuid(1).eax; + let base_family = (eax >> 8) & 0xF; + let family = if base_family == 0xF { + base_family + ((eax >> 20) & 0xFF) + } else { + base_family + }; + vendor != *b"AuthenticAMD" || family != 0x17 +} + #[cfg(feature = "std")] fn detect_cpu_kernel_uncached() -> CpuKernelTag { #[cfg(target_arch = "x86_64")] @@ -386,7 +483,7 @@ fn detect_cpu_kernel_uncached() -> CpuKernelTag { // `&&` short-circuits away the runtime `is_x86_feature_detected!` call // (and its CPUID/cache traffic) for tiers the build disabled — the // matching `select_x86_kernel` rung is `#[cfg]`-ed out anyway. - return select_x86_kernel( + let tier = select_x86_kernel( cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vbmi2"), cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512f"), cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vl"), @@ -395,6 +492,25 @@ fn detect_cpu_kernel_uncached() -> CpuKernelTag { cfg!(feature = "kernel-avx2") && is_x86_feature_detected!("avx2"), cfg!(feature = "kernel-sse") && is_x86_feature_detected!("sse2"), ); + // The one place a CPU quirk is allowed to be asked about: AMD Zen 1 and + // Zen 2 run `pext` through microcode, where the mask form of the + // three-field extract wins. Answering it here turns the quirk into a + // tier, so the decode paths below carry a kernel rather than a flag. + #[cfg(feature = "kernel-bmi2")] + if !fast_pext() { + let takes_pext = match tier { + CpuKernelTag::Bmi2 => true, + #[cfg(feature = "kernel-avx2")] + CpuKernelTag::Avx2 => true, + #[cfg(feature = "kernel-vbmi2")] + CpuKernelTag::Vbmi2 => true, + _ => false, + }; + if takes_pext { + return CpuKernelTag::Bmi2SlowPext; + } + } + return tier; } // 32-bit x86 carries only the BMI2 tier: the wider tiers' kernels and // their `target_feature` bodies are x86_64-only, so there is nothing @@ -487,6 +603,8 @@ impl CpuKernelTag { feature = "kernel-bmi2" ))] CpuKernelTag::Bmi2 => "bmi2", + #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + CpuKernelTag::Bmi2SlowPext => "bmi2-slow-pext", #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] CpuKernelTag::Avx2 => "avx2", #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] diff --git a/zstd/src/decoding/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index 53d89086e..363450cf9 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -179,7 +179,14 @@ fn decompress_literals( feature = "kernel-bmi2" ))] CpuKernelTag::Bmi2 => unsafe { - decompress_literals_bmi2(section, scratch, dict, source, target) + decompress_literals_bmi2::(section, scratch, dict, source, target) + }, + // Same body, the kernel whose three-field extract takes the mask form. + #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + CpuKernelTag::Bmi2SlowPext => unsafe { + decompress_literals_bmi2::( + section, scratch, dict, source, target, + ) }, // The aarch64 tiers need no `target_feature` wrapper: NEON is part of // the baseline ABI there, and SVE reaches nothing in this pipeline yet. @@ -219,14 +226,14 @@ unsafe fn decompress_literals_avx2( feature = "kernel-bmi2" ))] #[target_feature(enable = "bmi2")] -unsafe fn decompress_literals_bmi2( +unsafe fn decompress_literals_bmi2( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], target: &mut Vec, ) -> Result { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] diff --git a/zstd/src/decoding/seq_decoder_bmi2.rs b/zstd/src/decoding/seq_decoder_bmi2.rs index 0d6c85b5b..8286eec81 100644 --- a/zstd/src/decoding/seq_decoder_bmi2.rs +++ b/zstd/src/decoding/seq_decoder_bmi2.rs @@ -15,7 +15,6 @@ use super::sequence_section_decoder::{ ADVANCE, ADVANCE_MASK, ExecSeq, SeqStreamSetup, init_sequence_stream, }; use crate::blocks::sequence_section::{MAX_OFFSET_CODE, Sequence, SequencesHeader}; -use crate::cpu_kernel::Bmi2Kernel; use crate::decoding::errors::{DecodeSequenceError, DecompressBlockError, ExecuteSequencesError}; use crate::decoding::sequence_execution::do_offset_history; @@ -166,7 +165,11 @@ macro_rules! execute_one_body { /// Caller must have verified BMI2 availability. #[target_feature(enable = "bmi2")] #[allow(clippy::too_many_lines)] -pub(crate) unsafe fn decode_and_execute_sequences_bmi2<'fse, B: BufferBackend>( +pub(crate) unsafe fn decode_and_execute_sequences_bmi2< + 'fse, + B: BufferBackend, + K: crate::cpu_kernel::CpuKernel, +>( section: &SequencesHeader, source: &[u8], fse: &'fse mut FSEScratch, @@ -184,7 +187,7 @@ pub(crate) unsafe fn decode_and_execute_sequences_bmi2<'fse, B: BufferBackend>( old_buffer_size, num_sequences, use_long_pipeline, - } = init_sequence_stream::(section, source, fse, buffer, dict)?; + } = init_sequence_stream::(section, source, fse, buffer, dict)?; let literals_buffer_len = literals_buffer.len(); let mut lit_cur: usize = 0; let mut seq_sum: u32 = 0; diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index 453a81fd6..6cfea5810 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -311,7 +311,32 @@ pub fn decode_and_execute_sequences<'fse, B: super::buffer_backend::BufferBacken // divergence can be applied without touching the other // kernels — see #279 round 3. unsafe { - super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::( + super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< + B, + crate::cpu_kernel::Bmi2Kernel, + >( + section, + source, + fse, + buffer, + offset_hist, + literals_buffer, + dict, + ) + } + } + // Same monolith, a kernel whose three-field extract takes the mask + // form: this hardware microcodes `pext`. The wider monoliths are given + // up along with it, which costs little on the parts in question (Zen 1 + // splits a 256-bit operation in two anyway). + #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] + CpuKernelTag::Bmi2SlowPext => { + // SAFETY: detect confirmed BMI2, as for the arm above. + unsafe { + super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< + B, + crate::cpu_kernel::Bmi2SlowPextKernel, + >( section, source, fse, diff --git a/zstd/src/decoding/sequence_section_decoder/tests.rs b/zstd/src/decoding/sequence_section_decoder/tests.rs index 48724160b..38489db1d 100644 --- a/zstd/src/decoding/sequence_section_decoder/tests.rs +++ b/zstd/src/decoding/sequence_section_decoder/tests.rs @@ -425,7 +425,10 @@ mod init_sequence_stream_tests { let mut offset_hist = [1u32, 4, 8]; // SAFETY: BMI2 confirmed available by the runtime check above. let _ = unsafe { - crate::decoding::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::( + crate::decoding::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< + RingBuffer, + crate::cpu_kernel::Bmi2Kernel, + >( &header, &source, &mut fse, diff --git a/zstd/src/decoding/simd_copy.rs b/zstd/src/decoding/simd_copy.rs index 497960b59..a7933fb28 100644 --- a/zstd/src/decoding/simd_copy.rs +++ b/zstd/src/decoding/simd_copy.rs @@ -761,8 +761,10 @@ fn detect_x86_caps() -> X86Caps { avx2: true, sse2: true, }, + // Both BMI2 tiers copy the same way: they differ only in how + // the sequence decoder splits three bit fields. #[cfg(feature = "kernel-bmi2")] - CpuKernelTag::Bmi2 => X86Caps { + CpuKernelTag::Bmi2 | CpuKernelTag::Bmi2SlowPext => X86Caps { avx512f: false, avx2: false, sse2: true, From 6a03bea7b5aec580d5e5b5745b40e4a4b1d5c846 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 11:39:21 +0300 Subject: [PATCH 24/29] perf(decoding): keep the mask form of the three-field extract Measured `pext` against it on the i9, same frame, arms interleaved: the mask form issues FEWER instructions (6.7277e9 against 6.7320e9 on a 1 MiB level-19 stream decode) and the cycles overlap across repeats (2.640..2.643e9 against 2.628..2.639e9). `pext` needs three mask loads and two shifts to set up, which is what the three shifts and three `bzhi` of the mask form cost outright, and its latency is three cycles against one. So the instruction it saves is not saved, and paying for it means a second BMI2 tier plus a `__cpuid` vendor probe, because AMD Zen 1 and Zen 2 microcode `pext` at around 18 cycles. Reverted: the extract keeps the form every kernel shares, and the vendor question disappears with it. --- zstd/src/cpu_kernel.rs | 120 +----------------- zstd/src/decoding/literals_section_decoder.rs | 13 +- zstd/src/decoding/seq_decoder_bmi2.rs | 9 +- zstd/src/decoding/sequence_section_decoder.rs | 27 +--- .../sequence_section_decoder/tests.rs | 5 +- zstd/src/decoding/simd_copy.rs | 4 +- 6 files changed, 10 insertions(+), 168 deletions(-) diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index 8bb197d9d..c79002008 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -121,28 +121,6 @@ impl CpuKernel for ScalarKernel { } } -/// One `pext` per field, the widths turned into masks from [`BIT_MASK`]. -/// -/// # Safety -/// The caller's kernel was selected after BMI2 was detected, and on hardware -/// where `pext` is not microcoded: see [`Bmi2SlowPextKernel`]. -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] -#[target_feature(enable = "bmi2")] -#[inline] -unsafe fn extract_triple_pext_impl(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { - use core::arch::x86_64::_pext_u64; - - let mask3 = BIT_MASK[n3 as usize]; - let mask2 = BIT_MASK[n2 as usize].wrapping_shl(u32::from(n3)); - let mask1 = BIT_MASK[n1 as usize].wrapping_shl(u32::from(n2) + u32::from(n3)); - - ( - _pext_u64(packed, mask1), - _pext_u64(packed, mask2), - _pext_u64(packed, mask3), - ) -} - // The SSE2 tier exists in `CpuKernelTag` (it carries the 128-bit copy-chunk // choice for the unified copy dispatch) but needs no `CpuKernel` ZST yet: the // only trait method, `mask_lower_bits`, has no SSE2-specific form (SSE2 has no @@ -185,38 +163,6 @@ impl CpuKernel for Bmi2Kernel { // SAFETY: as for `mask_lower_bits`. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - #[cfg(target_arch = "x86_64")] - #[inline(always)] - fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { - // SAFETY: as for `mask_lower_bits`, and the tier is only selected on - // hardware whose `pext` is not microcoded. - unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } - } -} - -/// x86_64 BMI2 kernel for hardware whose `pext` runs through microcode (AMD -/// Zen 1 and Zen 2, where it takes around 18 cycles against one for `bzhi`). -/// Identical to [`Bmi2Kernel`] except that the three-field extract keeps the -/// trait's mask form. The vendor question is answered once, where the CPU is -/// detected; below that there is a kernel, not a flag to branch on. -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] -#[derive(Copy, Clone, Default)] -pub(crate) struct Bmi2SlowPextKernel; - -#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] -impl CpuKernel for Bmi2SlowPextKernel { - #[inline(always)] - fn mask_lower_bits(value: u64, n: u8) -> u64 { - // SAFETY: as for `Bmi2Kernel::mask_lower_bits`. - unsafe { mask_lower_bits_bmi2_impl(value, n) } - } - - #[inline(always)] - fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { - // SAFETY: as for `mask_lower_bits`. - unsafe { mask_lower_bits_bmi2_impl(value, n) } - } } /// x86_64 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern @@ -235,14 +181,6 @@ impl CpuKernel for Avx2Kernel { // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - #[cfg(feature = "kernel-bmi2")] - #[inline(always)] - fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { - // SAFETY: as for `mask_lower_bits`; the tier is not selected on - // hardware whose `pext` is microcoded. - unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } - } } /// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU @@ -261,14 +199,6 @@ impl CpuKernel for Vbmi2Kernel { // at runtime before this kernel is instantiated. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - #[cfg(feature = "kernel-bmi2")] - #[inline(always)] - fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) { - // SAFETY: as for `mask_lower_bits`; the tier is not selected on - // hardware whose `pext` is microcoded. - unsafe { extract_triple_pext_impl(packed, n1, n2, n3) } - } } /// aarch64 NEON baseline kernel. Used on all aarch64 hardware that @@ -418,10 +348,6 @@ pub(crate) enum CpuKernelTag { feature = "kernel-bmi2" ))] Bmi2, - /// BMI2 hardware whose `pext` is microcoded (AMD Zen 1 and Zen 2). Takes - /// the mask form of the three-field extract; see [`Bmi2SlowPextKernel`]. - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - Bmi2SlowPext, #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] Avx2, #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] @@ -451,29 +377,6 @@ pub(crate) fn detect_cpu_kernel() -> CpuKernelTag { *CACHED.get_or_init(detect_cpu_kernel_uncached) } -/// Whether this CPU runs `pext` in hardware. AMD Zen 1 and Zen 2 (family 0x17) -/// microcode it at around 18 cycles, where three masked shifts are faster; -/// every other vendor and family executes it in one. -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] -fn fast_pext() -> bool { - use core::arch::x86_64::__cpuid; - - // Leaves 0 and 1 are architectural on every x86_64, so these are safe. - let leaf0 = __cpuid(0); - let mut vendor = [0u8; 12]; - vendor[0..4].copy_from_slice(&leaf0.ebx.to_le_bytes()); - vendor[4..8].copy_from_slice(&leaf0.edx.to_le_bytes()); - vendor[8..12].copy_from_slice(&leaf0.ecx.to_le_bytes()); - let eax = __cpuid(1).eax; - let base_family = (eax >> 8) & 0xF; - let family = if base_family == 0xF { - base_family + ((eax >> 20) & 0xFF) - } else { - base_family - }; - vendor != *b"AuthenticAMD" || family != 0x17 -} - #[cfg(feature = "std")] fn detect_cpu_kernel_uncached() -> CpuKernelTag { #[cfg(target_arch = "x86_64")] @@ -483,7 +386,7 @@ fn detect_cpu_kernel_uncached() -> CpuKernelTag { // `&&` short-circuits away the runtime `is_x86_feature_detected!` call // (and its CPUID/cache traffic) for tiers the build disabled — the // matching `select_x86_kernel` rung is `#[cfg]`-ed out anyway. - let tier = select_x86_kernel( + return select_x86_kernel( cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vbmi2"), cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512f"), cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vl"), @@ -492,25 +395,6 @@ fn detect_cpu_kernel_uncached() -> CpuKernelTag { cfg!(feature = "kernel-avx2") && is_x86_feature_detected!("avx2"), cfg!(feature = "kernel-sse") && is_x86_feature_detected!("sse2"), ); - // The one place a CPU quirk is allowed to be asked about: AMD Zen 1 and - // Zen 2 run `pext` through microcode, where the mask form of the - // three-field extract wins. Answering it here turns the quirk into a - // tier, so the decode paths below carry a kernel rather than a flag. - #[cfg(feature = "kernel-bmi2")] - if !fast_pext() { - let takes_pext = match tier { - CpuKernelTag::Bmi2 => true, - #[cfg(feature = "kernel-avx2")] - CpuKernelTag::Avx2 => true, - #[cfg(feature = "kernel-vbmi2")] - CpuKernelTag::Vbmi2 => true, - _ => false, - }; - if takes_pext { - return CpuKernelTag::Bmi2SlowPext; - } - } - return tier; } // 32-bit x86 carries only the BMI2 tier: the wider tiers' kernels and // their `target_feature` bodies are x86_64-only, so there is nothing @@ -603,8 +487,6 @@ impl CpuKernelTag { feature = "kernel-bmi2" ))] CpuKernelTag::Bmi2 => "bmi2", - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - CpuKernelTag::Bmi2SlowPext => "bmi2-slow-pext", #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))] CpuKernelTag::Avx2 => "avx2", #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] diff --git a/zstd/src/decoding/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index 363450cf9..53d89086e 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -179,14 +179,7 @@ fn decompress_literals( feature = "kernel-bmi2" ))] CpuKernelTag::Bmi2 => unsafe { - decompress_literals_bmi2::(section, scratch, dict, source, target) - }, - // Same body, the kernel whose three-field extract takes the mask form. - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - CpuKernelTag::Bmi2SlowPext => unsafe { - decompress_literals_bmi2::( - section, scratch, dict, source, target, - ) + decompress_literals_bmi2(section, scratch, dict, source, target) }, // The aarch64 tiers need no `target_feature` wrapper: NEON is part of // the baseline ABI there, and SVE reaches nothing in this pipeline yet. @@ -226,14 +219,14 @@ unsafe fn decompress_literals_avx2( feature = "kernel-bmi2" ))] #[target_feature(enable = "bmi2")] -unsafe fn decompress_literals_bmi2( +unsafe fn decompress_literals_bmi2( section: &LiteralsSection, scratch: &mut HuffmanScratch, dict: Option<&Dictionary>, source: &[u8], target: &mut Vec, ) -> Result { - decompress_literals_impl::(section, scratch, dict, source, target) + decompress_literals_impl::(section, scratch, dict, source, target) } #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] diff --git a/zstd/src/decoding/seq_decoder_bmi2.rs b/zstd/src/decoding/seq_decoder_bmi2.rs index 8286eec81..0d6c85b5b 100644 --- a/zstd/src/decoding/seq_decoder_bmi2.rs +++ b/zstd/src/decoding/seq_decoder_bmi2.rs @@ -15,6 +15,7 @@ use super::sequence_section_decoder::{ ADVANCE, ADVANCE_MASK, ExecSeq, SeqStreamSetup, init_sequence_stream, }; use crate::blocks::sequence_section::{MAX_OFFSET_CODE, Sequence, SequencesHeader}; +use crate::cpu_kernel::Bmi2Kernel; use crate::decoding::errors::{DecodeSequenceError, DecompressBlockError, ExecuteSequencesError}; use crate::decoding::sequence_execution::do_offset_history; @@ -165,11 +166,7 @@ macro_rules! execute_one_body { /// Caller must have verified BMI2 availability. #[target_feature(enable = "bmi2")] #[allow(clippy::too_many_lines)] -pub(crate) unsafe fn decode_and_execute_sequences_bmi2< - 'fse, - B: BufferBackend, - K: crate::cpu_kernel::CpuKernel, ->( +pub(crate) unsafe fn decode_and_execute_sequences_bmi2<'fse, B: BufferBackend>( section: &SequencesHeader, source: &[u8], fse: &'fse mut FSEScratch, @@ -187,7 +184,7 @@ pub(crate) unsafe fn decode_and_execute_sequences_bmi2< old_buffer_size, num_sequences, use_long_pipeline, - } = init_sequence_stream::(section, source, fse, buffer, dict)?; + } = init_sequence_stream::(section, source, fse, buffer, dict)?; let literals_buffer_len = literals_buffer.len(); let mut lit_cur: usize = 0; let mut seq_sum: u32 = 0; diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index 6cfea5810..453a81fd6 100644 --- a/zstd/src/decoding/sequence_section_decoder.rs +++ b/zstd/src/decoding/sequence_section_decoder.rs @@ -311,32 +311,7 @@ pub fn decode_and_execute_sequences<'fse, B: super::buffer_backend::BufferBacken // divergence can be applied without touching the other // kernels — see #279 round 3. unsafe { - super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< - B, - crate::cpu_kernel::Bmi2Kernel, - >( - section, - source, - fse, - buffer, - offset_hist, - literals_buffer, - dict, - ) - } - } - // Same monolith, a kernel whose three-field extract takes the mask - // form: this hardware microcodes `pext`. The wider monoliths are given - // up along with it, which costs little on the parts in question (Zen 1 - // splits a 256-bit operation in two anyway). - #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))] - CpuKernelTag::Bmi2SlowPext => { - // SAFETY: detect confirmed BMI2, as for the arm above. - unsafe { - super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< - B, - crate::cpu_kernel::Bmi2SlowPextKernel, - >( + super::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::( section, source, fse, diff --git a/zstd/src/decoding/sequence_section_decoder/tests.rs b/zstd/src/decoding/sequence_section_decoder/tests.rs index 38489db1d..48724160b 100644 --- a/zstd/src/decoding/sequence_section_decoder/tests.rs +++ b/zstd/src/decoding/sequence_section_decoder/tests.rs @@ -425,10 +425,7 @@ mod init_sequence_stream_tests { let mut offset_hist = [1u32, 4, 8]; // SAFETY: BMI2 confirmed available by the runtime check above. let _ = unsafe { - crate::decoding::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::< - RingBuffer, - crate::cpu_kernel::Bmi2Kernel, - >( + crate::decoding::seq_decoder_bmi2::decode_and_execute_sequences_bmi2::( &header, &source, &mut fse, diff --git a/zstd/src/decoding/simd_copy.rs b/zstd/src/decoding/simd_copy.rs index a7933fb28..497960b59 100644 --- a/zstd/src/decoding/simd_copy.rs +++ b/zstd/src/decoding/simd_copy.rs @@ -761,10 +761,8 @@ fn detect_x86_caps() -> X86Caps { avx2: true, sse2: true, }, - // Both BMI2 tiers copy the same way: they differ only in how - // the sequence decoder splits three bit fields. #[cfg(feature = "kernel-bmi2")] - CpuKernelTag::Bmi2 | CpuKernelTag::Bmi2SlowPext => X86Caps { + CpuKernelTag::Bmi2 => X86Caps { avx512f: false, avx2: false, sse2: true, From 3aa5eb22ad56f61253b481d60239ef13263165bb Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 12:37:58 +0300 Subject: [PATCH 25/29] fix(decoding): finish an empty frame, and cap the ring at what a frame declares Two defects the per-block reservation work left, each with the regression test that fails without it. A frame that produces nothing could not be decoded into a slice that holds nothing. The block loop stops once the target is full, and a target of no bytes is full at its own length before any block is read, so the empty last block that ends such a frame was never reached: every call returned no progress on input that was complete. The stop now excludes an empty target, which reaches that block; a frame that does produce bytes buffers its first block and stops on the pending arm of the next pass, as before. A multi-segment frame keeps a block of room past its window because each compressed block reserves a block before it decodes, but a frame cannot produce past the size it declared. A 1 MiB window declaring one byte more reserved 1,179,649 bytes of ring for a byte of it, the per-block reservation having already been capped at the same remainder. The limit takes the declaration too, so that frame holds 1,048,578. --- zstd/src/decoding/frame_decoder.rs | 21 ++++++- zstd/src/decoding/frame_decoder/tests.rs | 78 ++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index f72636cde..f6d5ea5f6 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -890,6 +890,12 @@ impl FrameDecoderState { /// window filled. Upstream sizes its stream buffer as window + block too /// (`ZSTD_decodingBufferSize_min`); it can cap at the content because its /// buffer is not a ring. + /// + /// A frame that declares its size caps the whole of that at the declaration: + /// the block of room past the window is there for what a block still has to + /// produce, and a frame cannot produce past what it promised. Without the + /// cap a 1 MiB window declaring one byte more reserved a whole block of + /// room for that byte. fn decoding_buffer_limit(&self) -> usize { let useful_window = self.useful_window_size(); if self.frame_header.descriptor.single_segment_flag() { @@ -898,7 +904,12 @@ impl FrameDecoderState { let window_size = self.frame_header.window_size().unwrap_or(0) as usize; // No overflow: the window was checked against // `MAXIMUM_ALLOWED_WINDOW_SIZE` when the header was taken. - useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize) + let limit = useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize); + if self.frame_header.fcs_declared() { + let declared = self.frame_header.frame_content_size(); + return limit.min(usize::try_from(declared).unwrap_or(usize::MAX)); + } + limit } /// What to reserve up front: the limit, except for a multi-segment frame @@ -2493,7 +2504,13 @@ impl FrameDecoder { // a drain that empties the buffer into the last of `target` // leaves nothing pending, and decoding another block then // consumes input the caller cannot be handed the output of. - if pending > 0 || written == target.len() { + // A target of no bytes is not full in that sense: it starts + // at its own length, so the test would fire before any block + // was read and a frame that produces nothing could never + // reach the empty block that ends it. Such a frame decodes + // here; one that does produce bytes buffers its first block + // and stops on the `pending` arm of the next pass. + if pending > 0 || (!target.is_empty() && written == target.len()) { break; } //check if there are enough bytes for the next header diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 7af47ffa5..0c5c08d0f 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1124,6 +1124,57 @@ fn a_streamed_frame_smaller_than_its_window_gets_a_ring_of_its_content() { ); } +/// A multi-segment frame whose declared content is just past its window still +/// cannot produce more than it declared, so the block of room the ring keeps +/// past the window is capped by what is left to produce. A frame of a 1 MiB +/// window declaring one byte more reserved a whole block of that room up front, +/// where one byte is all any of it can ever hold. +#[test] +fn a_streamed_frame_just_past_its_window_reserves_only_what_it_declares() { + let window = 1024 * 1024u32; + let content = window + 1; + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x80, // FHD: multi-segment, 4-byte content size + 0x50, // window descriptor: 1 MiB + ]; + frame.extend_from_slice(&content.to_le_bytes()); + let mut payload = Vec::with_capacity(content as usize); + let mut left = content; + while left > 0 { + let size = left.min(128 * 1024); + left -= size; + // Raw block header: last flag, type 0, size. + let header = size << 3 | u32::from(left == 0); + frame.extend_from_slice(&header.to_le_bytes()[..3]); + let body: Vec = (0..size).map(|i| (i * 11 + left) as u8).collect(); + payload.extend_from_slice(&body); + frame.extend_from_slice(&body); + } + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let mut decoded = Vec::with_capacity(payload.len()); + let mut chunk = alloc::vec![0u8; 128 * 1024]; + while !(decoder.is_finished() && decoder.can_collect() == 0) { + let (read, written) = decoder + .decode_from_to(source, &mut chunk) + .expect("frame decodes"); + source = &source[read..]; + decoded.extend_from_slice(&chunk[..written]); + assert!(read > 0 || written > 0, "decode made no progress"); + } + assert_eq!(decoded, payload); + // The frame's own content, plus the byte the ring keeps to tell a full + // buffer from an empty one. No block of room on top of that. + let capacity = ring_capacity(&decoder); + assert!( + capacity <= content as usize + 1, + "a frame declaring {content} bytes reserved {capacity} bytes of ring" + ); +} + /// Capacity of the ring a multi-segment frame decoded into. fn ring_capacity(decoder: &FrameDecoder) -> usize { match &decoder @@ -1395,6 +1446,33 @@ fn a_filled_target_stops_before_the_next_block() { assert_eq!(read, 2 * (3 + BLOCK as usize)); } +/// A frame that produces nothing is decoded with a slice that holds nothing, +/// and it has to finish: its last block is empty, so there is no output the +/// caller is short of. Stopping on a full target at the top of the loop made an +/// empty target full before any block was read, so the block that ends the +/// frame was never reached and every further call reported no progress on input +/// that was complete. +#[test] +fn an_empty_frame_finishes_through_an_empty_slice() { + let mut frame = alloc::vec![ + 0x28, 0xB5, 0x2F, 0xFD, // magic + 0x00, // FHD: multi-segment, no content size + 0x00, // window descriptor: 1 KiB + ]; + // One last Raw block of no bytes. + frame.extend_from_slice(&1u32.to_le_bytes()[..3]); + + let mut decoder = FrameDecoder::new(); + let mut source = frame.as_slice(); + decoder.reset(&mut source).expect("header parses"); + let (read, written) = decoder + .decode_from_to(source, &mut []) + .expect("frame decodes"); + assert_eq!(written, 0); + assert_eq!(read, 3, "the block that ends the frame must be read"); + assert!(decoder.is_finished(), "the frame must finish"); +} + /// A block can produce at most its frame's block maximum, so that is what the /// pre-block reservation asks for. A 1 KiB window asking for a full 128 KiB /// grew the ring to 256 KiB: the growth limit clamps a need that fits under it, From c70fc8101b75c97d8e475a468c3e10e2c6e1b46f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 12:38:22 +0300 Subject: [PATCH 26/29] perf(decoding): bzhi for the precomputed HUF mask on Avx2 and Vbmi2 --- zstd/src/cpu_kernel.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index c79002008..da9f593f0 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -181,6 +181,13 @@ impl CpuKernel for Avx2Kernel { // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + /// `bzhi` takes the width, so the caller's mask is not needed. + #[inline(always)] + fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { + // SAFETY: as for `mask_lower_bits`. + unsafe { mask_lower_bits_bmi2_impl(value, n) } + } } /// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU @@ -199,6 +206,13 @@ impl CpuKernel for Vbmi2Kernel { // at runtime before this kernel is instantiated. unsafe { mask_lower_bits_bmi2_impl(value, n) } } + + /// `bzhi` takes the width, so the caller's mask is not needed. + #[inline(always)] + fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { + // SAFETY: as for `mask_lower_bits`. + unsafe { mask_lower_bits_bmi2_impl(value, n) } + } } /// aarch64 NEON baseline kernel. Used on all aarch64 hardware that From b09746b06aa9983ac4437904b2475629cd9b79b8 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 12:42:07 +0300 Subject: [PATCH 27/29] perf(decoding): advance the HUF state by the table's mask on every kernel The precomputed-mask op was overridden by the Bmi2 tier and not by Avx2 or Vbmi2, so the ladder disagreed with itself: the wider tiers took the default while the narrower one took `bzhi`. Measured which of the two forms is right, on the i9, both arms interleaved across five rounds on the level-19 frames that spend the most of their time in HUF. `bzhi` issues MORE instructions (6.7335e9 against 6.7278e9 on a 1 MiB stream decode, 6.7067e9 against 6.7010e9 one-shot): the mask is already in hand, so taking the width instead is a second load. The cycles overlap on every shape (2.6262-2.6361 against 2.6320-2.6615 streamed), which is no difference at this scale. So the mask is not what a tier settles for, it is the better form, and the op stops being per-kernel: the advance masks with `state_mask` everywhere and the trait loses the method. The numbers are recorded at the advance so the question does not come back. --- zstd/src/cpu_kernel.rs | 30 ------------------------------ zstd/src/huff0/huff0_decoder.rs | 21 +++++++++++---------- 2 files changed, 11 insertions(+), 40 deletions(-) diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index da9f593f0..987c9d9ca 100644 --- a/zstd/src/cpu_kernel.rs +++ b/zstd/src/cpu_kernel.rs @@ -62,15 +62,6 @@ pub trait CpuKernel: Copy + 'static { /// construction time; no per-call wrapper assert runs. fn mask_lower_bits(value: u64, n: u8) -> u64; - /// [`Self::mask_lower_bits`] for a caller that already holds the mask, - /// `mask == (1 << n) - 1`: the HUF table keeps one per decoder. A kernel - /// with a bit-extract instruction ignores the mask and takes `n`; the - /// others take the mask and skip building it per call. - #[inline(always)] - fn mask_lower_bits_precomputed(value: u64, mask: u64, _n: u8) -> u64 { - value & mask - } - /// Split the low `n1 + n2 + n3` bits of `packed` into three fields, the /// highest first. The FSE sequence decoder reads its three state updates /// this way, once per sequence. @@ -156,13 +147,6 @@ impl CpuKernel for Bmi2Kernel { // running CPU. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - /// `bzhi` takes the width, so the caller's mask is not needed. - #[inline(always)] - fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { - // SAFETY: as for `mask_lower_bits`. - unsafe { mask_lower_bits_bmi2_impl(value, n) } - } } /// x86_64 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern @@ -181,13 +165,6 @@ impl CpuKernel for Avx2Kernel { // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - /// `bzhi` takes the width, so the caller's mask is not needed. - #[inline(always)] - fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { - // SAFETY: as for `mask_lower_bits`. - unsafe { mask_lower_bits_bmi2_impl(value, n) } - } } /// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU @@ -206,13 +183,6 @@ impl CpuKernel for Vbmi2Kernel { // at runtime before this kernel is instantiated. unsafe { mask_lower_bits_bmi2_impl(value, n) } } - - /// `bzhi` takes the width, so the caller's mask is not needed. - #[inline(always)] - fn mask_lower_bits_precomputed(value: u64, _mask: u64, n: u8) -> u64 { - // SAFETY: as for `mask_lower_bits`. - unsafe { mask_lower_bits_bmi2_impl(value, n) } - } } /// aarch64 NEON baseline kernel. Used on all aarch64 hardware that diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index 2e4acf07c..059cbbb88 100644 --- a/zstd/src/huff0/huff0_decoder.rs +++ b/zstd/src/huff0/huff0_decoder.rs @@ -84,11 +84,16 @@ impl<'t> HuffmanDecoder<'t> { /// Decode symbol and advance state in one table lookup. /// - /// The kernel is `K`, chosen once where the decode was dispatched, so the - /// state advance is the monomorph's own instruction: `bzhi` on the BMI2 - /// tiers, the table's `state_mask` elsewhere, which is built once per table - /// rather than per symbol. `state_mask == (1 << max_num_bits) - 1` is the - /// value `bzhi` produces, so the two agree bit for bit. + /// The advance masks with the table's `state_mask`, built once per table + /// rather than per symbol, on every kernel alike. + /// + /// `state_mask == (1 << max_num_bits) - 1`, so a BMI2 `bzhi` on the width + /// produces the same value and was measured against this on the i9: it + /// issues MORE instructions (6.7335e9 against 6.7278e9 on a 1 MiB level-19 + /// stream decode) because the width is a second load where the mask is + /// already in hand, and the cycles overlap across repeats on every decode + /// shape. So the mask is not a fallback that the accelerated tiers give up + /// something by taking: it is the better form, and no tier overrides it. #[inline(always)] pub fn decode_symbol_and_advance( &mut self, @@ -97,11 +102,7 @@ impl<'t> HuffmanDecoder<'t> { let packed = self.table.packed_decode[self.state as usize]; let num_bits = (packed >> 8) as u8; let new_bits = br.get_bits(num_bits); - self.state = K::mask_lower_bits_precomputed( - self.state << num_bits, - self.table.state_mask, - self.table.max_num_bits, - ) | new_bits; + self.state = ((self.state << num_bits) & self.table.state_mask) | new_bits; packed as u8 } From 88fba73e4c3f4313166a1f06bf8a981df1a7666e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 13:05:42 +0300 Subject: [PATCH 28/29] fix(decoding): cap the ring's reservation at the declared size, not its limit Capping the growth limit at the declaration aborted the decode on a frame that declares less than its blocks produce: the ring ran out of buffer under a write it cannot refuse, and asserts on that rather than reporting it. The fuzz decode target finds such a frame in well under a minute. The declaration belongs on the up-front reservation instead, which is where the waste was: a 1 MiB window declaring one byte more reserved 1,179,649 bytes for that byte and now reserves 1,048,578. The limit stays at the window plus a block, so a frame that exceeds what it promised still has somewhere to put the bytes and is judged once they exist. Carries the input that aborted, straight from the fuzzer. --- zstd/src/decoding/frame_decoder.rs | 38 ++++++++++++++---------- zstd/src/decoding/frame_decoder/tests.rs | 26 ++++++++++++++++ 2 files changed, 48 insertions(+), 16 deletions(-) diff --git a/zstd/src/decoding/frame_decoder.rs b/zstd/src/decoding/frame_decoder.rs index f6d5ea5f6..f88b02ccd 100644 --- a/zstd/src/decoding/frame_decoder.rs +++ b/zstd/src/decoding/frame_decoder.rs @@ -891,11 +891,12 @@ impl FrameDecoderState { /// (`ZSTD_decodingBufferSize_min`); it can cap at the content because its /// buffer is not a ring. /// - /// A frame that declares its size caps the whole of that at the declaration: - /// the block of room past the window is there for what a block still has to - /// produce, and a frame cannot produce past what it promised. Without the - /// cap a 1 MiB window declaring one byte more reserved a whole block of - /// room for that byte. + /// A declared content size does NOT cap this. A frame can declare less than + /// its blocks go on to produce, and that is caught by the check against the + /// declaration once the bytes exist; a limit that stopped the ring short of + /// them would instead have the write run out of buffer, which the ring + /// asserts on rather than reports. Only the up-front reservation takes the + /// declaration ([`Self::decoding_buffer_size`]). fn decoding_buffer_limit(&self) -> usize { let useful_window = self.useful_window_size(); if self.frame_header.descriptor.single_segment_flag() { @@ -904,12 +905,7 @@ impl FrameDecoderState { let window_size = self.frame_header.window_size().unwrap_or(0) as usize; // No overflow: the window was checked against // `MAXIMUM_ALLOWED_WINDOW_SIZE` when the header was taken. - let limit = useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize); - if self.frame_header.fcs_declared() { - let declared = self.frame_header.frame_content_size(); - return limit.min(usize::try_from(declared).unwrap_or(usize::MAX)); - } - limit + useful_window + window_size.min(crate::common::MAX_BLOCK_SIZE as usize) } /// What to reserve up front: the limit, except for a multi-segment frame @@ -918,14 +914,24 @@ impl FrameDecoderState { /// they reserve it themselves; the limit then caps that growth at one /// block. Such frames are rare (encoders mark a frame that fits its window /// single-segment), and a small Raw or RLE one should not pay a block. + /// + /// A frame that declares its size never reserves past the declaration: the + /// block of room is there for what a block still has to produce, and an + /// honest frame produces exactly what it promised. A 1 MiB window declaring + /// one byte more reserved a whole block of room for that byte. A frame that + /// goes on to exceed its declaration grows into the limit and is then caught + /// by the check against it. fn decoding_buffer_size(&self) -> usize { let window_size = self.frame_header.window_size().unwrap_or(0); - if self.frame_header.fcs_declared() && self.frame_header.frame_content_size() <= window_size - { - self.useful_window_size() - } else { - self.decoding_buffer_limit() + if !self.frame_header.fcs_declared() { + return self.decoding_buffer_limit(); + } + let declared = self.frame_header.frame_content_size(); + if declared <= window_size { + return self.useful_window_size(); } + self.decoding_buffer_limit() + .min(usize::try_from(declared).unwrap_or(usize::MAX)) } /// Reserve this frame's decode buffer ([`Self::decoding_buffer_size`]) diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index 0c5c08d0f..41c1d3884 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -1175,6 +1175,32 @@ fn a_streamed_frame_just_past_its_window_reserves_only_what_it_declares() { ); } +/// A frame is free to declare less than its blocks go on to produce. The ring +/// has to be able to hold what arrives: a growth limit cut to the declaration +/// left the write short of buffer, and the ring aborts on that rather than +/// reporting it. Only the up-front reservation takes the declaration; the limit +/// stays at the window plus a block, so the bytes land and the decode carries +/// on to the checks that judge them. +#[test] +fn a_frame_that_produces_past_its_declared_size_does_not_abort_the_ring() { + // A frame declaring no content whose compressed block regenerates 512 + // bytes, which the sequence executor writes into the ring through the + // infallible path. + let frame: &[u8] = &[ + 0x28, 0xB5, 0x2F, 0xFD, 0x80, 0x14, 0x00, 0x00, 0x00, 0x00, 0x14, 0x02, 0x00, 0xA1, 0xA1, + 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0x9A, 0xA1, 0x81, 0xA1, 0xA1, 0xA1, 0xA1, 0x81, + 0x7A, 0x00, 0x30, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, + 0xCA, 0xCA, 0xCA, 0xCE, 0xA1, 0xA1, 0xA1, 0xA1, 0x81, 0xCA, 0x00, 0x7A, 0xA1, 0xA1, 0x5B, + 0xA1, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xAA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, 0xCA, + 0xCA, 0xCA, 0xCA, 0x2F, 0xCE, 0x0E, 0x39, + ]; + + // Whatever the verdict on the frame, reaching one is the point: the ring + // must not run out of buffer under the write. + let mut out = Vec::new(); + let _ = FrameDecoder::new().decode_all_to_vec(frame, &mut out); +} + /// Capacity of the ring a multi-segment frame decoded into. fn ring_capacity(decoder: &FrameDecoder) -> usize { match &decoder From ca1963b1fd8a9444532a0693bae069e3ff8c12d0 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Wed, 16 Sep 2026 13:08:46 +0300 Subject: [PATCH 29/29] test(decoding): gate the VBMI2 kernel tests on the selector's full predicate The tier mixes VBMI2 with AVX2 widths and BMI2 masking, and the sweeps asked only for VBMI2. A CPU that offers it while masking any of the rest would have reached the monomorph and decoded through instructions it does not have. Both sweeps now ask what the kernel selection asks. --- zstd/src/bit_io/bit_reader_reverse/tests.rs | 9 ++++++++- zstd/src/huff0/huff0_decoder/tests.rs | 12 +++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/zstd/src/bit_io/bit_reader_reverse/tests.rs b/zstd/src/bit_io/bit_reader_reverse/tests.rs index 8005f9604..15d928e87 100644 --- a/zstd/src/bit_io/bit_reader_reverse/tests.rs +++ b/zstd/src/bit_io/bit_reader_reverse/tests.rs @@ -338,8 +338,15 @@ fn peek_bits_triple_agrees_across_kernels() { n3 ); } + // The full predicate the kernel selection uses: the tier mixes VBMI2 + // with AVX2 widths, so a CPU offering VBMI2 alone must not reach it. #[cfg(feature = "kernel-vbmi2")] - if is_x86_feature_detected!("avx512vbmi2") { + if is_x86_feature_detected!("avx512vbmi2") + && is_x86_feature_detected!("avx512f") + && is_x86_feature_detected!("avx512vl") + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx2") + { assert_eq!( triple_under!(crate::cpu_kernel::Vbmi2Kernel, sum, n1, n2, n3), expected, diff --git a/zstd/src/huff0/huff0_decoder/tests.rs b/zstd/src/huff0/huff0_decoder/tests.rs index d769fd56e..85619d770 100644 --- a/zstd/src/huff0/huff0_decoder/tests.rs +++ b/zstd/src/huff0/huff0_decoder/tests.rs @@ -149,8 +149,18 @@ fn every_kernel_advances_the_state_alike() { if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("bmi2") { same_as_scalar!(crate::cpu_kernel::Avx2Kernel); } + // The same predicate the kernel selection uses, in full: the tier mixes + // VBMI2 with AVX2 widths and BMI2 masking, so a CPU that offers VBMI2 while + // masking any of the rest must not reach this monomorph. It would decode + // through instructions it does not have. #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))] - if std::arch::is_x86_feature_detected!("avx512vbmi2") { + if std::arch::is_x86_feature_detected!("avx512vbmi2") + && std::arch::is_x86_feature_detected!("avx512f") + && std::arch::is_x86_feature_detected!("avx512vl") + && std::arch::is_x86_feature_detected!("avx512bw") + && std::arch::is_x86_feature_detected!("bmi2") + && std::arch::is_x86_feature_detected!("avx2") + { same_as_scalar!(crate::cpu_kernel::Vbmi2Kernel); } #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]