From 89bdaa8e35f4b7242ca54d1ca83c8c1e1615e012 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Sat, 17 Oct 2026 10:27:38 +0300 Subject: [PATCH] fix(encoding): single-segment frames whenever the window covers the content The streaming encoder wrote a single-segment frame only for pledged sizes of 512 bytes to 16 KiB and never with a dictionary, and the frame compressor only when a size was hinted before the frame. Upstream's rule (`ZSTD_writeFrameHeader`) is one condition: the content size is written and the window covers it. Elsewhere our frames declared the power of two above the content as their window, so a decoder reserved more memory than the data needs: a 140 MiB file at `--long=28` asked for 256 MiB where upstream's frame asks for 140 MiB. - One rule in `frame_header::single_segment`, read by both encoders, dictionary frames included (RFC 8878 3.1.1.1.2: the window is the content size, and the dictionary stays referenceable before it). - The CLI's frames change in the header alone: everything after it is byte-identical at levels 1-19 on four fixture sizes, with and without a dictionary, and the header is up to four bytes shorter. Tests check both encoders against libzstd's header layout from 100 bytes to 1 MiB with and without a dictionary, with a C round-trip; two tests that read the window from the header now use content past it. Closes #570 --- ffi-bench/tests/frame_compressor_ffi.rs | 62 ++++++++++++++++++++ zstd/src/decoding/frame_decoder/tests.rs | 7 ++- zstd/src/encoding/frame_compressor.rs | 34 +++-------- zstd/src/encoding/frame_compressor/tests.rs | 28 +++++++++ zstd/src/encoding/frame_header.rs | 11 ++++ zstd/src/encoding/streaming_encoder.rs | 31 +++------- zstd/src/encoding/streaming_encoder/tests.rs | 32 ++++++++++ zstd/src/tests/roundtrip_integrity.rs | 24 ++++---- 8 files changed, 166 insertions(+), 63 deletions(-) diff --git a/ffi-bench/tests/frame_compressor_ffi.rs b/ffi-bench/tests/frame_compressor_ffi.rs index afaec4bbd..83baf63d4 100644 --- a/ffi-bench/tests/frame_compressor_ffi.rs +++ b/ffi-bench/tests/frame_compressor_ffi.rs @@ -119,6 +119,68 @@ fn fcs_header_written_and_c_zstd_compatible() { } } +/// Frames of a known size take the layout the reference takes for them, with +/// a dictionary and without, from both encoders: single-segment whenever the +/// window covers the content (`ZSTD_writeFrameHeader`), so a decoder reserves +/// the content rather than the power of two above it. Each frame decodes +/// through the C library. +#[test] +fn known_size_frames_choose_the_reference_header_layout() { + use std::io::Write; + use structured_zstd::encoding::{EncoderDictionary, StreamingEncoder}; + + let text: Vec = (0..60_000u32) + .flat_map(|i| format!("row {} key {} val {}\n", i % 97, i % 13, i % 7).into_bytes()) + .collect(); + let dictionary = &text[..8_192]; + for size in [100usize, 600, 16 * 1024 + 1, 64 * 1024, 1 << 20] { + let data = &text[..size]; + for with_dictionary in [false, true] { + let reference = if with_dictionary { + zstd::bulk::Compressor::with_dictionary(3, dictionary) + .and_then(|mut c| c.compress(data)) + } else { + zstd::bulk::compress(data, 3) + } + .expect("C compress"); + let (reference_single, reference_fcs, _) = frame_header_info(&reference); + + let mut streaming = StreamingEncoder::new(Vec::new(), CompressionLevel::Level(3)); + let mut one_shot: FrameCompressor = FrameCompressor::new(CompressionLevel::Level(3)); + if with_dictionary { + let encoder_dictionary = + EncoderDictionary::from_serialized_or_raw_content(dictionary).unwrap(); + streaming + .set_encoder_dictionary(encoder_dictionary.clone()) + .unwrap(); + one_shot.set_encoder_dictionary(encoder_dictionary).unwrap(); + } + streaming.set_pledged_content_size(size as u64).unwrap(); + streaming.write_all(data).unwrap(); + let streamed = streaming.finish().unwrap(); + one_shot.set_source_size_hint(size as u64); + let one_shot_frame = one_shot.compress_independent_frame(data); + + for (path, frame) in [("stream", &streamed), ("one-shot", &one_shot_frame)] { + let (single, fcs, _) = frame_header_info(frame); + assert_eq!( + (single, fcs), + (reference_single, reference_fcs), + "{path}, {size} bytes, dictionary {with_dictionary}: header layout" + ); + let decoded = if with_dictionary { + zstd::bulk::Decompressor::with_dictionary(dictionary) + .and_then(|mut d| d.decompress(frame, size)) + } else { + zstd::bulk::decompress(frame, size) + } + .expect("C decode"); + assert_eq!(decoded, data, "{path}, {size} bytes: C roundtrip"); + } + } + } +} + #[test] fn source_size_hint_fastest_remains_ffi_compatible_small_input() { let data = vec![0xAB; 2047]; diff --git a/zstd/src/decoding/frame_decoder/tests.rs b/zstd/src/decoding/frame_decoder/tests.rs index ced345ea7..d9008a7e7 100644 --- a/zstd/src/decoding/frame_decoder/tests.rs +++ b/zstd/src/decoding/frame_decoder/tests.rs @@ -172,9 +172,10 @@ fn decode_all_multi_segment_frame_decodes_correctly() { // FlatBuf/Ring-backed path. // // Make the payload structured so multi-segment behavior - // actually kicks in: 2 MiB of repeating + random-ish bytes - // forces window_size lower than content_size at the encoder. - let mut payload: Vec = Vec::with_capacity(2 * 1024 * 1024); + // actually kicks in: 64 KiB past the default level's 2 MiB + // window keeps window_size below content_size, so the frame + // cannot be single-segment. + let mut payload: Vec = Vec::with_capacity(2 * 1024 * 1024 + 64 * 1024); for i in 0..payload.capacity() { payload.push((i.wrapping_mul(2_654_435_761) & 0xFF) as u8); } diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 5f5660f8d..aea0af9d5 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -1475,7 +1475,6 @@ impl CompressState { struct FramePrep { window_size: u64, use_dictionary_state: bool, - source_size_hint_known: bool, initial_size_hint: Option, } @@ -2487,7 +2486,6 @@ impl FrameCompressor { } } let initial_size_hint = self.source_size_hint; - let source_size_hint_known = initial_size_hint.is_some(); let use_dictionary_state = !matches!(self.compression_level, CompressionLevel::Uncompressed) && self.state.matcher.supports_dictionary_priming() @@ -2652,7 +2650,6 @@ impl FrameCompressor { FramePrep { window_size, use_dictionary_state, - source_size_hint_known, initial_size_hint, } } @@ -2901,29 +2898,16 @@ impl FrameCompressor { /// Appends rather than returns so the one-shot path serializes straight /// into the reused output buffer with no per-frame header `Vec`. fn append_frame_header(&self, total_uncompressed: u64, prep: &FramePrep, out: &mut Vec) { - // Match the upstream zstd framing policy (`ZSTD_writeFrameHeader`): - // single-segment whenever the content size is known and the whole - // source fits the active window (`contentSizeFlag && windowSize >= - // srcSize`). A single-segment frame REQUIRES an FCS field, so - // suppressing the content size (`content_size_flag` off) forces the - // windowed layout. There is no lower size bound: small payloads - // benefit most, since a windowed frame cannot encode a content size - // below 256 in fewer than 4 FCS bytes (the 1-byte FCS class is - // single-segment-only, see `find_fcs_field_size`), whereas a - // single-segment frame stores it in one byte and omits the window - // descriptor. The single-segment window equals the FCS, so a block - // must never reference past the content: the post-hoc raw fallback in - // the block emitters guarantees any non-shrinking block is stored raw, - // and genuine matches stay within the already-emitted output. - // Dictionary frames qualify too (the dictionary is decoder setup - // state, not part of the regenerated segment), keeping the decoder's - // single-allocation path (our decoder caps reservation to - // min(window, FCS) either way). - let single_segment = self.content_size_flag - && prep.source_size_hint_known - && total_uncompressed <= prep.window_size; + // The header goes out after the content, so its size is known whether + // or not it was hinted. A windowed frame cannot encode a content size + // below 256 in fewer than 4 FCS bytes (the 1-byte class is + // single-segment-only, see `find_fcs_field_size`), so small payloads + // gain the most. + let frame_content_size = self.content_size_flag.then_some(total_uncompressed); + let single_segment = + crate::encoding::frame_header::single_segment(frame_content_size, prep.window_size); let header = FrameHeader { - frame_content_size: self.content_size_flag.then_some(total_uncompressed), + frame_content_size, single_segment, content_checksum: cfg!(feature = "hash") && self.content_checksum, dictionary_id: if prep.use_dictionary_state && self.dict_id_flag { diff --git a/zstd/src/encoding/frame_compressor/tests.rs b/zstd/src/encoding/frame_compressor/tests.rs index 7aaae2f3c..88ee7100e 100644 --- a/zstd/src/encoding/frame_compressor/tests.rs +++ b/zstd/src/encoding/frame_compressor/tests.rs @@ -629,6 +629,34 @@ fn compressible_stream_output_capacity_stays_at_output_scale() { assert_eq!(decoded, input); } +// The header goes out after the content, so its size is written and the +// single-segment rule applies even when no size was hinted before the frame: +// a frame the window covers declares the content as its window. +#[test] +fn unhinted_frame_the_window_covers_is_single_segment() { + let payload = b"unhinted payload, written once the content is known ".repeat(200); + let mut compressor = super::FrameCompressor::new(super::CompressionLevel::Default); + compressor.set_source(payload.as_slice()); + compressor.set_drain(Vec::::new()); + compressor.compress(); + let frame = compressor.take_drain().expect("the drain was set"); + let parsed = crate::decoding::frame::read_frame_header(frame.as_slice()) + .expect("frame header must parse") + .0; + assert!( + parsed.descriptor.single_segment_flag(), + "a frame of {} bytes under a larger window is windowed", + payload.len() + ); + assert_eq!(parsed.frame_content_size(), payload.len() as u64); + let mut decoder = crate::decoding::FrameDecoder::new(); + let mut decoded = Vec::with_capacity(payload.len()); + decoder + .decode_all_to_vec(&frame, &mut decoded) + .expect("our frame decodes"); + assert_eq!(decoded, payload); +} + // A dictionary frame with a known content size that fits the window // must take the single-segment layout (reference parity): the // dictionary is decoder setup state, not part of the regenerated diff --git a/zstd/src/encoding/frame_header.rs b/zstd/src/encoding/frame_header.rs index ff91e0b58..ab068001d 100644 --- a/zstd/src/encoding/frame_header.rs +++ b/zstd/src/encoding/frame_header.rs @@ -140,6 +140,17 @@ impl FrameHeader { } } +/// Whether a frame whose header writes `content_size` (`None` when it leaves +/// the size out) under a window of `window_size` bytes is single-segment: +/// whenever the size is written and the window covers it, with a dictionary or +/// without, as upstream decides it (`ZSTD_writeFrameHeader`, +/// `zstd_compress.c:4728`). The decoder's window is then the content size +/// (RFC 8878 3.1.1.1.2) rather than the power of two above it, and a +/// dictionary stays referenceable before the content. +pub(crate) fn single_segment(content_size: Option, window_size: u64) -> bool { + content_size.is_some_and(|size| size <= window_size) +} + /// Serialize a `Frame_Content_Size` value into `field_size` bytes. /// /// For 1, 4, and 8-byte fields the value is stored directly. diff --git a/zstd/src/encoding/streaming_encoder.rs b/zstd/src/encoding/streaming_encoder.rs index 3ab885e4b..762b6e03a 100644 --- a/zstd/src/encoding/streaming_encoder.rs +++ b/zstd/src/encoding/streaming_encoder.rs @@ -859,10 +859,10 @@ impl CompressionContext { // A dictionary is only active when it can actually be primed: the level // compresses (not `Uncompressed`) AND the matcher supports priming AND a // dictionary is attached. Mirrors `FrameCompressor`'s `use_dictionary_state` - // so a streaming frame never advertises a `Dictionary_ID`, disables - // single-segment, or seeds dict entropy/offsets unless the dictionary is - // genuinely in play (otherwise it would emit frames that needlessly - // require a dictionary at decode time). + // so a streaming frame never advertises a `Dictionary_ID` or seeds dict + // entropy/offsets unless the dictionary is genuinely in play (otherwise + // it would emit frames that needlessly require a dictionary at decode + // time). let use_dictionary_state = !matches!(self.compression_level, CompressionLevel::Uncompressed) && self.state.matcher.supports_dictionary_priming() @@ -1001,27 +1001,12 @@ impl CompressionContext { )); } - // Single-segment is incompatible with a dictionary (the dictionary - // pushes referenceable history before the content, so the frame needs - // an explicit window descriptor); gate it off when a dict is attached, - // mirroring `FrameCompressor`'s `!use_dictionary_state` guard. - // Single-segment also requires the FCS field to be present - // (`content_size_flag`): the layout drops the window descriptor, - // so the header must carry the content size for decoders to size - // their window. - let single_segment = self.content_size_flag - && !use_dictionary_state - && self - .pledged_content_size - .map(|size| (512..=(1 << 14)).contains(&size) && size <= window_size) - .unwrap_or(false); + let frame_content_size = self.pledged_content_size.filter(|_| self.content_size_flag); + let single_segment = + crate::encoding::frame_header::single_segment(frame_content_size, window_size); let header = FrameHeader { - frame_content_size: if self.content_size_flag { - self.pledged_content_size - } else { - None - }, + frame_content_size, single_segment, content_checksum: cfg!(feature = "hash") && self.content_checksum, dictionary_id: if use_dictionary_state && self.dictionary_id_flag { diff --git a/zstd/src/encoding/streaming_encoder/tests.rs b/zstd/src/encoding/streaming_encoder/tests.rs index 43d633bee..e0be196fb 100644 --- a/zstd/src/encoding/streaming_encoder/tests.rs +++ b/zstd/src/encoding/streaming_encoder/tests.rs @@ -1177,6 +1177,38 @@ fn source_size_hint_directly_reduces_window_header() { assert_eq!(decoded, payload); } +/// A frame whose pledged size the window covers is single-segment at any size +/// and with a dictionary, as upstream writes it (`ZSTD_writeFrameHeader`): +/// the decoder then reserves the content, not the next power of two above it. +#[test] +fn a_pledged_frame_the_window_covers_is_single_segment() { + let text: Vec = (0..60_000u32) + .flat_map(|i| alloc::format!("row {} key {} val {}\n", i % 97, i % 13, i % 7).into_bytes()) + .collect(); + let dictionary = EncoderDictionary::from_serialized_or_raw_content(&text[..8_192]) + .expect("raw content is a dictionary"); + for size in [100usize, 600, 16 * 1024 + 1, 64 * 1024, 1 << 20] { + let payload = &text[..size]; + for with_dictionary in [false, true] { + let mut encoder = StreamingEncoder::new(Vec::new(), CompressionLevel::from_level(3)); + if with_dictionary { + encoder.set_encoder_dictionary(dictionary.clone()).unwrap(); + } + encoder.set_pledged_content_size(size as u64).unwrap(); + encoder.write_all(payload).unwrap(); + let frame = encoder.finish().unwrap(); + let header = crate::decoding::frame::read_frame_header(frame.as_slice()) + .unwrap() + .0; + assert!( + header.descriptor.single_segment_flag(), + "{size} bytes, dictionary {with_dictionary}: windowed frame" + ); + assert_eq!(header.frame_content_size(), size as u64); + } + } +} + #[test] fn single_segment_requires_pledged_to_fit_matcher_window() { let payload = b"streaming-window-gate-".repeat(60); // 1320 bytes diff --git a/zstd/src/tests/roundtrip_integrity.rs b/zstd/src/tests/roundtrip_integrity.rs index 12e00c308..b0e1c092a 100644 --- a/zstd/src/tests/roundtrip_integrity.rs +++ b/zstd/src/tests/roundtrip_integrity.rs @@ -857,7 +857,9 @@ fn streaming_pledged_size_uses_source_hint() { } /// One-shot `compress_to_vec` should propagate source-size hint for Default -/// so tiny payloads avoid oversized dfast window/table sizing. +/// so tiny payloads avoid oversized dfast window/table sizing: it writes the +/// frame an explicitly hinted compressor writes, and that frame advertises the +/// content as its window. #[test] fn compress_to_vec_default_small_input_uses_source_size_hint() { let data = generate_compressible(9604, 4 * 1024); // 4 KiB @@ -870,31 +872,29 @@ fn compress_to_vec_default_small_input_uses_source_size_hint() { .window_size() .unwrap(); - // Manual compressor path without hint (legacy behavior baseline). - let no_hint = { + let hinted = { let mut compressor = FrameCompressor::new(CompressionLevel::Default); + compressor.set_source_size_hint(data.len() as u64); compressor.set_source(data.as_slice()); let mut out = Vec::new(); compressor.set_drain(&mut out); compressor.compress(); out }; - let no_hint_window = crate::decoding::frame::read_frame_header(no_hint.as_slice()) - .unwrap() - .0 - .window_size() - .unwrap(); let mut decoder = StreamingDecoder::new(auto_hint.as_slice()); let mut decoded = Vec::new(); decoder.read_to_end(&mut decoded).unwrap(); assert_eq!(decoded, data); - assert!( - auto_hint_window < no_hint_window, - "compress_to_vec(default) should advertise a smaller window on tiny payloads: auto_hint={} no_hint={}", + assert_eq!( + auto_hint, hinted, + "compress_to_vec(default) should write the hinted frame" + ); + assert_eq!( auto_hint_window, - no_hint_window + data.len() as u64, + "a tiny payload should advertise its content as the window" ); }