fix(encoding): single-segment frames whenever the window covers the content

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<u8> = (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];

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<u8> = 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<u8> = 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);
     }

zstd/src/encoding/frame_compressor.rs

@@ -1475,7 +1475,6 @@ impl<M: Matcher> CompressState<M> {
 struct FramePrep {
     window_size: u64,
     use_dictionary_state: bool,
-    source_size_hint_known: bool,
     initial_size_hint: Option<u64>,
 }
 
@@ -2487,7 +2486,6 @@ impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
             }
         }
         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<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
         FramePrep {
             window_size,
             use_dictionary_state,
-            source_size_hint_known,
             initial_size_hint,
         }
     }
@@ -2901,29 +2898,16 @@ impl<R: Read, W: Write, M: Matcher> FrameCompressor<R, W, M> {
     /// 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<u8>) {
-        // 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 {

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::<u8>::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

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<u64>, 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.

zstd/src/encoding/streaming_encoder.rs

@@ -859,10 +859,10 @@ impl<M: Matcher> CompressionContext<M> {
         // 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<M: Matcher> CompressionContext<M> {
             ));
         }
 
-        // 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 {

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<u8> = (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

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"
     );
 }