diff --git a/zstd/src/bit_io/bit_reader_reverse.rs b/zstd/src/bit_io/bit_reader_reverse.rs index befcab966..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; @@ -24,72 +18,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 +72,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> { @@ -162,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(), @@ -196,8 +88,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 +256,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 @@ -411,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/bit_io/bit_reader_reverse/tests.rs b/zstd/src/bit_io/bit_reader_reverse/tests.rs index 07c3739e9..15d928e87 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; } @@ -323,36 +307,108 @@ fn peek_bits_triple_bmi2_matches_scalar() { (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); - // SAFETY: gated on `is_x86_feature_detected!("bmi2")` above. - let b = unsafe { bmi2.peek_bits_triple_bmi2(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 + ); + } + // 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") + && 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, + "Vbmi2Kernel differs at widths=({},{},{})", + n1, + n2, + n3 + ); + } } } -#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))] +/// 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 should_use_pext_policy_table() { - let cases = [ - (*b"AuthenticAMD", 0x17, false), - (*b"AuthenticAMD", 0x19, true), - (*b"GenuineIntel", 0x06, true), +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 (vendor, family, expected) in cases { - assert_eq!(super::should_use_pext(vendor, family), expected); + 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 + ); + } } } +/// 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 bmi2_triple_extract_matches_scalar_reference() { +fn extract_triple_matches_the_reference_under_every_kernel() { + use crate::cpu_kernel::CpuKernel; + if !is_x86_feature_detected!("bmi2") { return; } @@ -384,12 +440,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 +456,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 + ); } } } diff --git a/zstd/src/cpu_kernel.rs b/zstd/src/cpu_kernel.rs index 9a6040014..987c9d9ca 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 @@ -68,18 +87,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] } } @@ -90,16 +119,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 { @@ -197,7 +234,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 { @@ -207,7 +247,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 @@ -267,7 +325,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, @@ -317,6 +380,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"))] @@ -357,6 +434,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"))] @@ -382,7 +466,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 e03bd2be8..544398722 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, @@ -33,12 +35,85 @@ enum DecoderState { Failed, //TODO put "self.internal_state = DecoderState::Failed;" everywhere an unresolvable error occurs } -/// Create a new [BlockDecoder]. +/// 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; + let maximum = block_maximum(window_size); + if size > maximum { + return Err(DecodeBlockContentError::DecompressBlockError( + DecompressBlockError::ExpandsPastBlockMaximum { size, maximum }, + )); + } + Ok(()) +} + +/// 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. +/// +/// 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. +/// +/// 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, + literals: &[u8], +) -> Result<(), DecompressBlockError> { + if !B::FIXED_CAPACITY { + buffer.push(literals); + return Ok(()); + } + buffer + .try_push(literals) + .map_err(|overflow| DecompressBlockError::LiteralsOutputOverflow { + tail: overflow.tail, + requested: overflow.requested, + capacity: overflow.capacity, + }) +} + +/// 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, } } @@ -94,8 +169,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 +196,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 +251,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 +267,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,8 +392,26 @@ impl BlockDecoder { raw: &[u8], dict: Option<&'d crate::decoding::dictionary::Dictionary>, ) -> Result<(), DecompressBlockError> { + // 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 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..]; vprintln!( "Found {} literalssection with regenerated size: {}, and compressed size: {:?}", @@ -400,6 +497,22 @@ 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. + // 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. + buffer.reserve_for_block(block_maximum); + buffer.set_block_output_ceiling(block_maximum); decode_and_execute_sequences( &seq_section, raw, @@ -418,9 +531,25 @@ impl BlockDecoder { }, )); } - 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 + // length is this block's output. + let produced = buffer.len() - len_before; + if produced > block_maximum { + return Err(DecompressBlockError::ExpandsPastBlockMaximum { + size: produced, + maximum: block_maximum, + }); + } Ok(()) } diff --git a/zstd/src/decoding/buffer_backend.rs b/zstd/src/decoding/buffer_backend.rs index d083a2c12..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` @@ -357,6 +365,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..40c6316f4 100644 --- a/zstd/src/decoding/decode_buffer.rs +++ b/zstd/src/decoding/decode_buffer.rs @@ -48,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`]. @@ -74,6 +78,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,15 +99,29 @@ impl Read for DecodeBuffer { written += buf.len(); (buf.len(), Ok(())) })?; - Ok(amount) + Ok((amount, max_amount - amount)) } } +/// 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, + declared_content: None, total_output_counter: 0, #[cfg(feature = "hash")] hash: twox_hash::XxHash64::with_seed(0), @@ -116,9 +146,11 @@ 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, + declared_content: None, total_output_counter: 0, #[cfg(feature = "hash")] hash: twox_hash::XxHash64::with_seed(0), @@ -129,6 +161,16 @@ impl DecodeBuffer { } } + /// 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); + 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`). @@ -171,9 +213,47 @@ 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; + } + + /// 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 { + 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 // layer. Direct-decode frames (`run_direct_decode`) write through // `UserSliceBackend` and never touch this buffer, so a long-lived @@ -308,6 +388,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 @@ -362,12 +450,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 13f8f9a99..5aeb0a368 100644 --- a/zstd/src/decoding/errors.rs +++ b/zstd/src/decoding/errors.rs @@ -265,6 +265,25 @@ 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 + /// `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 + /// 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")] @@ -302,6 +321,18 @@ 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, maximum } => write!( + f, + "Block expands to {size} bytes, past this frame's maximum of {maximum}" + ), + 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 b12268f47..f88b02ccd 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 @@ -513,16 +518,28 @@ 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 - /// 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`. + /// Hand the buffer what the frame declared it would produce, so the + /// per-block reservation can stop at the frame's remainder. #[inline] - fn reserve_buffer(&mut self, window_size: usize) { + 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; // 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 @@ -536,8 +553,15 @@ impl DecoderScratchKind { // window-sized buffer toward 2x window. match self { Self::Ring(s) => { + // 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()); @@ -605,6 +629,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), @@ -844,6 +880,75 @@ impl FrameDecoderState { } } + /// 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. + /// + /// 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() { + 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. + 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. + /// + /// 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() { + 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`]) + /// 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(); + // 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); + } + /// 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`). @@ -1000,6 +1105,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(), } } @@ -1751,17 +1857,16 @@ 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); - - let mut block_dec = decoding::block_decoder::new(); + // 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. + state.reserve_decoding_buffer(); + + 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; @@ -1993,13 +2098,12 @@ 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); + 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 @@ -2074,7 +2178,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 @@ -2307,14 +2411,20 @@ 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 /// 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 @@ -2330,6 +2440,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; @@ -2340,11 +2452,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 @@ -2383,6 +2496,29 @@ 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. + let (read, pending) = state + .decoder_scratch + .buffer_read_reporting_pending(&mut target[written..]) + .map_err(err::FailedToDrainDecodebuffer)?; + written += read; + // 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. + // 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 if mt_source.len() < 3 { break; @@ -2401,6 +2537,18 @@ 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() { + state.reserve_decoding_buffer(); + } // Only expose the held dictionary while THIS frame is dict-backed // (`using_dict` is set per dict-apply, cleared on reset). A reused @@ -2450,7 +2598,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. @@ -2717,15 +2868,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()?; @@ -2842,9 +2995,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 @@ -2861,12 +3023,11 @@ 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); + state.reserve_decoding_buffer(); } let frame_start_total = total_bytes_written; loop { @@ -2980,9 +3141,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 @@ -2995,12 +3165,11 @@ 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); + state.reserve_decoding_buffer(); } let frame_start_total = total_bytes_written; loop { @@ -3106,12 +3275,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 @@ -3127,7 +3297,7 @@ impl FrameDecoder { &mut self, input: &mut &[u8], output: &mut [u8], - content_size: u64, + declared_size: Option, ) -> Result { #[cfg(test)] { @@ -3140,6 +3310,17 @@ 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 kernel = self.kernel; let state = self .state .as_mut() @@ -3156,12 +3337,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) - && n as u64 == content_size + && declared_size.is_none_or(|declared| declared == n as u64) + && n <= block_decoder::block_maximum(window) && probe.len() >= n && output.len() >= n { @@ -3268,7 +3459,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 @@ -3327,13 +3518,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 @@ -3358,11 +3546,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 @@ -3372,23 +3558,45 @@ 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. + // 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; 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))); + } + // 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( @@ -3411,11 +3619,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; @@ -3439,15 +3644,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 ab80907cc..41c1d3884 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, @@ -862,6 +862,848 @@ 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); +} + +/// 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:?}" + ); +} + +/// 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 +/// 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 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" + ); +} + +/// 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 + .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. +#[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 + ]; + frame_around_block(&block, content_size) +} + +/// One frame holding `block` as its only (last, compressed) block, with a +/// 1 MiB window and optionally a declared content size. +fn frame_around_block(block: &[u8], content_size: Option) -> Vec { + 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. + 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 +} + +/// A 3-byte literals section header (size format 3, 20-bit regenerated size) +/// of the given type: 0 raw, 1 RLE. +fn literals_header_20_bit(literals_type: u8, regenerated: u32) -> [u8; 3] { + [ + ((regenerated & 0xF) << 4) as u8 | 0b11 << 2 | literals_type, + (regenerated >> 4) as u8, + (regenerated >> 12) as u8, + ] +} + +/// 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 { + 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", + ); +} + +/// 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 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 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, +/// 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 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 +/// 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. +#[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. +#[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); + 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/literals_section_decoder.rs b/zstd/src/decoding/literals_section_decoder.rs index 192fbd355..53d89086e 100644 --- a/zstd/src/decoding/literals_section_decoder.rs +++ b/zstd/src/decoding/literals_section_decoder.rs @@ -6,8 +6,19 @@ 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 = "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)] @@ -163,10 +174,30 @@ 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) }, + // 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), } } @@ -183,7 +214,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/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 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 { diff --git a/zstd/src/decoding/sequence_section_decoder.rs b/zstd/src/decoding/sequence_section_decoder.rs index c8ca3cc2e..453a81fd6 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 @@ -159,18 +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. - 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); + // 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; @@ -297,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 @@ -607,7 +612,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 @@ -639,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)?; @@ -1003,8 +1017,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( @@ -1128,7 +1141,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 diff --git a/zstd/src/decoding/user_slice_buf.rs b/zstd/src/decoding/user_slice_buf.rs index d644c9d75..dbc19491f 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,19 @@ pub(crate) struct UserSliceBackend<'a> { /// for API parity with `FlatBuf` and `RingBuffer`. head: usize, 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 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 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, } impl<'a> UserSliceBackend<'a> { @@ -150,10 +162,12 @@ 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, + sequence_cap, } } @@ -236,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")] @@ -266,7 +284,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`. @@ -377,7 +395,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 @@ -484,7 +502,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 @@ -609,9 +627,22 @@ impl<'a> BufferBackend for UserSliceBackend<'a> { slice: &mut [], head: 0, tail: 0, + sequence_cap: 0, } } + /// `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.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] fn clear(&mut self) { self.head = 0; @@ -624,12 +655,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. + // 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() => 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(), + capacity: self.sequence_cap, }), } } @@ -654,9 +688,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] diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index 3c2ef0338..059cbbb88 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,52 +83,21 @@ impl<'t> HuffmanDecoder<'t> { } /// Decode symbol and advance state in one table lookup. + /// + /// 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, 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; @@ -279,33 +126,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..85619d770 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,79 @@ 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, 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 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 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"); + + /// 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); + } + // 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") + && 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"))] + 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); + } }