diff --git a/zstd/src/decoding/dictionary.rs b/zstd/src/decoding/dictionary.rs index 1321fcfb1..1714a3001 100644 --- a/zstd/src/decoding/dictionary.rs +++ b/zstd/src/decoding/dictionary.rs @@ -85,6 +85,42 @@ impl Dictionary { /// and returns a fully constructed [`Dictionary`] whose `id` can be /// checked against the frame's `dict_id`. pub fn decode_dict(raw: &[u8]) -> Result { + Self::decode_dict_inner(raw, true) + } + + /// Parse a dictionary for ENCODER use: builds the entropy + /// probabilities/weights needed by `to_encoder_table` but skips the + /// decode-only work the encoder never reads — the FSE *decoding* + /// tables + their `enrich_*` post-passes, and the HUF decode lookup + /// table (`packed_decode`). Produces a [`Dictionary`] whose FSE + /// `symbol_probabilities` / `accuracy_log` and HUF `bits` / + /// `max_num_bits` match `decode_dict` exactly, so the encoder entropy + /// tables — and thus the emitted frame — are byte-identical; only the + /// wasted decode-table builds are dropped. Offset history + content + /// are parsed the same way. + /// Crate-internal: the returned [`Dictionary`] deliberately has no + /// decode lookup tables (`packed_decode` / FSE `decode`), so it is + /// NOT safe to feed into a [`FrameDecoder`](crate::decoding::FrameDecoder) + /// — Huffman decode would index an empty `packed_decode`. The only caller + /// is `EncoderDictionary::from_bytes`, which wraps the result in the + /// encoder-only `EncoderDictionary` type (no decode path), so this + /// incomplete dictionary can never escape to the decode side. Keeping + /// this `pub(crate)` keeps it off the public `Dictionary` API entirely. + pub(crate) fn decode_dict_for_encoding( + raw: &[u8], + ) -> Result { + Self::decode_dict_inner(raw, false) + } + + /// Shared dictionary parser. `build_decode_tables` selects whether the + /// FSE/HUF tables get their full decoding tables (FSE decode table + + /// `enrich_*`, HUF `packed_decode`; decoder path) or only the + /// probability/weight parse (encoder path — see + /// [`Self::decode_dict_for_encoding`]). + fn decode_dict_inner( + raw: &[u8], + build_decode_tables: bool, + ) -> Result { const MIN_MAGIC_AND_ID_LEN: usize = 8; const OFFSET_HISTORY_LEN: usize = 12; @@ -117,34 +153,62 @@ impl Dictionary { let raw_tables = &raw[8..]; - let huf_size = new_dict.huf.table.build_decoder(raw_tables)?; + let huf_size = if build_decode_tables { + new_dict.huf.table.build_decoder(raw_tables)? + } else { + new_dict.huf.table.build_weights_only(raw_tables)? + }; let raw_tables = &raw_tables[huf_size as usize..]; - let of_size = new_dict.fse.offsets.build_decoder( - raw_tables, - crate::decoding::sequence_section_decoder::OF_MAX_LOG, - )?; - new_dict.fse.offsets.enrich_for_offsets(); + let of_size = if build_decode_tables { + let n = new_dict.fse.offsets.build_decoder( + raw_tables, + crate::decoding::sequence_section_decoder::OF_MAX_LOG, + )?; + new_dict.fse.offsets.enrich_for_offsets(); + n + } else { + new_dict.fse.offsets.read_table_probabilities( + raw_tables, + crate::decoding::sequence_section_decoder::OF_MAX_LOG, + )? + }; let raw_tables = &raw_tables[of_size..]; - let ml_size = new_dict.fse.match_lengths.build_decoder( - raw_tables, - crate::decoding::sequence_section_decoder::ML_MAX_LOG, - )?; - new_dict - .fse - .match_lengths - .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::ML_META); + let ml_size = if build_decode_tables { + let n = new_dict.fse.match_lengths.build_decoder( + raw_tables, + crate::decoding::sequence_section_decoder::ML_MAX_LOG, + )?; + new_dict + .fse + .match_lengths + .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::ML_META); + n + } else { + new_dict.fse.match_lengths.read_table_probabilities( + raw_tables, + crate::decoding::sequence_section_decoder::ML_MAX_LOG, + )? + }; let raw_tables = &raw_tables[ml_size..]; - let ll_size = new_dict.fse.literal_lengths.build_decoder( - raw_tables, - crate::decoding::sequence_section_decoder::LL_MAX_LOG, - )?; - new_dict - .fse - .literal_lengths - .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::LL_META); + let ll_size = if build_decode_tables { + let n = new_dict.fse.literal_lengths.build_decoder( + raw_tables, + crate::decoding::sequence_section_decoder::LL_MAX_LOG, + )?; + new_dict + .fse + .literal_lengths + .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::LL_META); + n + } else { + new_dict.fse.literal_lengths.read_table_probabilities( + raw_tables, + crate::decoding::sequence_section_decoder::LL_MAX_LOG, + )? + }; let raw_tables = &raw_tables[ll_size..]; if raw_tables.len() < OFFSET_HISTORY_LEN { diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index aaab29afb..d2aac224f 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -17,6 +17,54 @@ use crate::fse::fse_encoder::{FSETable, default_ll_table, default_ml_table, defa use crate::io::{Read, Write}; +/// A dictionary prepared for the ENCODER side, analogous to zstd's `CDict` +/// (vs the decoder's [`Dictionary`](crate::decoding::Dictionary) / `DDict`). +/// +/// It carries the entropy tables, content, and repeat-offset history the +/// compressor needs, but is a distinct type with **no decode path**: there is +/// no way to turn it into a [`DictionaryHandle`](crate::decoding::DictionaryHandle) +/// or feed it to a [`FrameDecoder`](crate::decoding::FrameDecoder). That keeps +/// the compress-only state (which may have been parsed without building the +/// decode lookup tables, see +/// [`set_dictionary_from_bytes`](FrameCompressor::set_dictionary_from_bytes)) +/// from ever reaching the decode side — the encoder/decoder dictionary split +/// mirrors C zstd's `CDict` / `DDict`. +pub struct EncoderDictionary { + pub(crate) inner: crate::decoding::Dictionary, +} + +impl EncoderDictionary { + /// Wrap an already-parsed [`Dictionary`](crate::decoding::Dictionary) for + /// encoder use. A fully-decoded dictionary is valid here; only the encoder + /// entropy tables, content, and offset history are read. + pub fn from_dictionary(dictionary: crate::decoding::Dictionary) -> Self { + Self { inner: dictionary } + } + + /// Parse a serialized dictionary blob for encoder use, skipping the decode + /// lookup-table build the encoder never reads (see + /// `Dictionary::decode_dict_for_encoding`). The encoder entropy tables — and + /// thus the emitted frame — are identical to a full parse. + pub fn from_bytes( + raw_dictionary: &[u8], + ) -> Result { + Ok(Self { + inner: crate::decoding::Dictionary::decode_dict_for_encoding(raw_dictionary)?, + }) + } + + /// The dictionary id. + /// + /// A dictionary attached for encoding always has a non-zero id (the + /// `set_dictionary*` / `set_encoder_dictionary` attach path rejects a + /// zero id). This getter, however, reflects the wrapped dictionary as-is: + /// an `EncoderDictionary` built via [`Self::from_dictionary`] from a raw + /// `Dictionary` with `id == 0` reports `0` here until it is attached. + pub fn id(&self) -> u32 { + self.inner.id + } +} + /// An interface for compressing arbitrary data with the ZStandard compression algorithm. /// /// `FrameCompressor` will generally be used by: @@ -40,7 +88,7 @@ pub struct FrameCompressor { uncompressed_data: Option, compressed_data: Option, compression_level: CompressionLevel, - dictionary: Option, + dictionary: Option, dictionary_entropy_cache: Option, source_size_hint: Option, state: CompressState, @@ -630,10 +678,10 @@ impl FrameCompressor { if use_dictionary_state && let Some(dict) = self.dictionary.as_ref() { // This state drives sequence encoding, while matcher priming below updates // the match generator's internal repeat-offset history for match finding. - self.state.offset_hist = dict.offset_hist; + self.state.offset_hist = dict.inner.offset_hist; self.state .matcher - .prime_with_dictionary(dict.dict_content.as_slice(), dict.offset_hist); + .prime_with_dictionary(dict.inner.dict_content.as_slice(), dict.inner.offset_hist); } if let Some(cache) = cached_entropy { self.state.last_huff_table.clone_from(&cache.huff); @@ -865,7 +913,7 @@ impl FrameCompressor { single_segment, content_checksum: cfg!(feature = "hash"), dictionary_id: if use_dictionary_state { - self.dictionary.as_ref().map(|dict| dict.id as u64) + self.dictionary.as_ref().map(|dict| dict.inner.id as u64) } else { None }, @@ -1114,8 +1162,53 @@ impl FrameCompressor { pub fn set_dictionary( &mut self, dictionary: crate::decoding::Dictionary, - ) -> Result, crate::decoding::errors::DictionaryDecodeError> - { + ) -> Result, crate::decoding::errors::DictionaryDecodeError> { + self.attach_dictionary(EncoderDictionary::from_dictionary(dictionary)) + } + + /// Parse and attach a serialized dictionary blob. + /// + /// Parses with the encoder-only path (skips the FSE/HUF decode lookup-table + /// build the encoder never reads); the entropy ENCODER tables — and thus + /// the emitted frame — are identical to a full parse. + pub fn set_dictionary_from_bytes( + &mut self, + raw_dictionary: &[u8], + ) -> Result, crate::decoding::errors::DictionaryDecodeError> { + self.attach_dictionary(EncoderDictionary::from_bytes(raw_dictionary)?) + } + + /// Attach an already-parsed [`EncoderDictionary`] without reparsing a raw + /// blob. + /// + /// Accepts an `EncoderDictionary` produced once via + /// [`EncoderDictionary::from_bytes`] / [`EncoderDictionary::from_dictionary`] + /// or handed back by [`Self::clear_dictionary`] / the `set_dictionary*` + /// return value, so callers can reattach or reuse a prepared dictionary + /// across compressions without re-running the dictionary parse each time. + /// Returns the previously-attached dictionary, if any. + pub fn set_encoder_dictionary( + &mut self, + dictionary: EncoderDictionary, + ) -> Result, crate::decoding::errors::DictionaryDecodeError> { + self.attach_dictionary(dictionary) + } + + /// Remove the attached dictionary, returning it as an [`EncoderDictionary`]. + pub fn clear_dictionary(&mut self) -> Option { + self.dictionary_entropy_cache = None; + self.dictionary.take() + } + + /// Validate `enc`, build the encoder entropy cache from it, store it, and + /// return the previously-attached dictionary. Shared by every public + /// attach entry point: `set_dictionary`, `set_dictionary_from_bytes`, and + /// `set_encoder_dictionary`. + fn attach_dictionary( + &mut self, + enc: EncoderDictionary, + ) -> Result, crate::decoding::errors::DictionaryDecodeError> { + let dictionary = &enc.inner; if dictionary.id == 0 { return Err(crate::decoding::errors::DictionaryDecodeError::ZeroDictionaryId); } @@ -1144,23 +1237,7 @@ impl FrameCompressor { .to_encoder_table() .map(|table| PreviousFseTable::Custom(Box::new(table))), }); - Ok(self.dictionary.replace(dictionary)) - } - - /// Parse and attach a serialized dictionary blob. - pub fn set_dictionary_from_bytes( - &mut self, - raw_dictionary: &[u8], - ) -> Result, crate::decoding::errors::DictionaryDecodeError> - { - let dictionary = crate::decoding::Dictionary::decode_dict(raw_dictionary)?; - self.set_dictionary(dictionary) - } - - /// Remove the attached dictionary. - pub fn clear_dictionary(&mut self) -> Option { - self.dictionary_entropy_cache = None; - self.dictionary.take() + Ok(self.dictionary.replace(enc)) } } @@ -1797,7 +1874,7 @@ mod tests { .set_dictionary(dict_for_encoder) .expect("valid dictionary should attach") .expect("set_dictionary_from_bytes inserted previous dictionary") - .id, + .id(), dict_for_decoder.id ); compressor.set_source(data.as_slice()); @@ -1939,6 +2016,110 @@ mod tests { ); } + #[test] + fn set_encoder_dictionary_reattaches_prepared_dict_without_reparse() { + let dict_raw = include_bytes!("../../dict_tests/dictionary"); + let payload = b"tenant=demo table=orders op=put key=1 value=aaaaabbbbbcccccdddddeeeee\n\ + tenant=demo table=orders op=put key=2 value=aaaaabbbbbcccccdddddeeeee\n"; + + // Prepare the EncoderDictionary once, then attach it via the prepared- + // dictionary API (no raw-blob reparse at attach time). + let prepared = + super::EncoderDictionary::from_bytes(dict_raw).expect("dict bytes should parse"); + let dict_id = prepared.id(); + + let mut with_dict = Vec::new(); + let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest); + let previous = compressor + .set_encoder_dictionary(prepared) + .expect("prepared dictionary should attach"); + assert!(previous.is_none()); + compressor.set_source(payload.as_slice()); + compressor.set_drain(&mut with_dict); + compressor.compress(); + // clear_dictionary hands the prepared dictionary back (last use of + // `compressor`, so its `&mut with_dict` drain borrow ends here). + let returned = compressor + .clear_dictionary() + .expect("dictionary was attached"); + assert_eq!(returned.id(), dict_id); + + // The reattached dictionary drives the frame: its id is advertised and + // the stream round-trips through a decoder primed with the same dict. + let (frame_header, _) = crate::decoding::frame::read_frame_header(with_dict.as_slice()) + .expect("encoded stream should have a frame header"); + assert_eq!(frame_header.dictionary_id(), Some(dict_id)); + let decoder_dict = crate::decoding::Dictionary::decode_dict(dict_raw).unwrap(); + let mut decoder = FrameDecoder::new(); + decoder.add_dict(decoder_dict).unwrap(); + let mut decoded = Vec::with_capacity(payload.len()); + decoder.decode_all_to_vec(&with_dict, &mut decoded).unwrap(); + assert_eq!(decoded.as_slice(), payload.as_slice()); + + // The dictionary handed back by clear_dictionary reattaches to another + // compressor without touching the raw bytes again, producing an + // identical frame. + let mut with_dict2 = Vec::new(); + let mut compressor2 = FrameCompressor::new(super::CompressionLevel::Fastest); + compressor2 + .set_encoder_dictionary(returned) + .expect("returned dictionary should reattach"); + compressor2.set_source(payload.as_slice()); + compressor2.set_drain(&mut with_dict2); + compressor2.compress(); + assert_eq!( + with_dict2, with_dict, + "reattached prepared dict must produce an identical frame" + ); + } + + #[test] + fn set_dictionary_from_bytes_matches_full_decode_byte_for_byte() { + // The encoder-only dict parse (`decode_dict_for_encoding`, used by + // `set_dictionary_from_bytes`) skips the FSE/HUF decoder-table build and + // the enrich passes. The encoder entropy tables are derived purely from + // the symbol probabilities / Huffman weights, so the compressed output + // MUST be byte-identical to the full-decode path. This pins the + // load-bearing equivalence so a future FSE/HUF parsing refactor that + // still round-trips but silently diverges on the probabilities/weights + // fails loudly here instead of producing a different (but valid) frame. + let dict_raw = include_bytes!("../../dict_tests/dictionary"); + let payload = b"tenant=demo table=orders op=put key=1 value=aaaaabbbbbcccccdddddeeeee\n\ + tenant=demo table=orders op=put key=2 value=aaaaabbbbbcccccdddddeeeee\n"; + + // Path A: encoder-only parse straight from the raw blob. + let mut from_bytes_out = Vec::new(); + { + let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest); + compressor + .set_dictionary_from_bytes(dict_raw) + .expect("dictionary bytes should parse"); + compressor.set_source(payload.as_slice()); + compressor.set_drain(&mut from_bytes_out); + compressor.compress(); + } + + // Path B: full decode (builds the decoder tables too), then attach for + // encoding via the `Dictionary` setter. + let full_decode = crate::decoding::Dictionary::decode_dict(dict_raw) + .expect("dictionary bytes should fully decode"); + let mut full_decode_out = Vec::new(); + { + let mut compressor = FrameCompressor::new(super::CompressionLevel::Fastest); + compressor + .set_dictionary(full_decode) + .expect("full-decode dictionary should attach"); + compressor.set_source(payload.as_slice()); + compressor.set_drain(&mut full_decode_out); + compressor.compress(); + } + + assert_eq!( + from_bytes_out, full_decode_out, + "encoder-only dict parse must produce byte-identical output to the full decode" + ); + } + #[test] fn set_dictionary_rejects_zero_dictionary_id() { let invalid = crate::decoding::Dictionary { diff --git a/zstd/src/encoding/mod.rs b/zstd/src/encoding/mod.rs index fc59c1180..8ecb9d5b6 100644 --- a/zstd/src/encoding/mod.rs +++ b/zstd/src/encoding/mod.rs @@ -87,7 +87,7 @@ mod levels; #[cfg(feature = "bench_internals")] pub mod sequence_capture; mod streaming_encoder; -pub use frame_compressor::FrameCompressor; +pub use frame_compressor::{EncoderDictionary, FrameCompressor}; #[cfg(feature = "lsm")] pub use frame_emit_info::{BlockType, FrameBlock, FrameEmitInfo}; pub use match_generator::MatchGeneratorDriver; diff --git a/zstd/src/fse/fse_decoder.rs b/zstd/src/fse/fse_decoder.rs index 48089bdb2..842d96744 100644 --- a/zstd/src/fse/fse_decoder.rs +++ b/zstd/src/fse/fse_decoder.rs @@ -321,6 +321,38 @@ impl FSETableImpl { Ok(bytes_read) } + /// Parse the table description into `symbol_probabilities` + + /// `accuracy_log` WITHOUT building the decoding table. + /// + /// Returns the same byte count as [`Self::build_decoder`] (the table + /// description length), so a caller stepping a cursor over a packed + /// stream of tables advances identically. Used by the encoder + /// dictionary load: [`Self::to_encoder_table`] reads only the + /// probabilities + accuracy log, so building the decode table (and + /// the `enrich_*` post-passes, which touch only decode entries) is + /// pure waste there. + /// + /// The existing `decode` table is cleared so a reused `FSETableImpl` + /// can't silently keep decoding against a stale table that no longer + /// matches the just-parsed probabilities (`init_state` would + /// otherwise pass whenever the old `decode.len()` still equalled + /// `1 << accuracy_log`). After this call the table is intentionally + /// non-decodable until `build_decoding_table` runs. + // `pub(crate)`: this leaves the table in an intentionally non-decodable + // partial-init state, so it must not be reachable from the public API + // (the module is re-exported with `pub use`). Only the crate-internal + // encoder-dictionary parse calls it. + pub(crate) fn read_table_probabilities( + &mut self, + source: &[u8], + max_log: u8, + ) -> Result { + let max_log = max_log.min(ENTRY_MAX_ACCURACY_LOG); + self.accuracy_log = 0; + self.decode.clear(); + self.read_probabilities(source, max_log) + } + /// Given the provided accuracy log, build a decoding table from that log. pub fn build_from_probabilities( &mut self, diff --git a/zstd/src/huff0/huff0_decoder.rs b/zstd/src/huff0/huff0_decoder.rs index 65c2d5bcd..c3bfa0261 100644 --- a/zstd/src/huff0/huff0_decoder.rs +++ b/zstd/src/huff0/huff0_decoder.rs @@ -410,6 +410,24 @@ impl HuffmanTable { Ok(bytes_used) } + /// Parse the weight header into `bits` + `max_num_bits` WITHOUT building + /// the decode lookup table. Returns the same byte count as + /// [`Self::build_decoder`] (the weight-header length), so a caller + /// stepping a cursor over packed tables advances identically. Used by + /// the encoder dictionary load: [`Self::to_encoder_table`] reads only + /// `bits` + `max_num_bits`, so filling `packed_decode` is pure waste. + /// + /// Crate-internal: the table it produces is intentionally non-decodable + /// (empty `packed_decode`); only `decode_dict_for_encoding` calls it, and + /// its result is wrapped in an `EncoderDictionary` that has no decode path. + pub(crate) fn build_weights_only(&mut self, source: &[u8]) -> Result { + self.packed_decode.clear(); + + let bytes_used = self.read_weights(source)?; + self.compute_huffman_bits()?; + Ok(bytes_used) + } + /// Read weights from the provided source. /// /// The huffman table is represented in the input data as a list of weights. @@ -566,11 +584,13 @@ impl HuffmanTable { Ok(bytes_read as u32) } - /// Once the weights have been read from the data, you can decode the weights - /// into a table, and use that table to decode the actual compressed data. - /// - /// This function populates the rest of the table from the series of weights. - fn build_table_from_weights(&mut self) -> Result<(), HuffmanTableError> { + /// Compute per-symbol code lengths (`bits`) + `max_num_bits` from the + /// parsed `weights`, returning `max_bits`. This is the slice of the + /// table build the ENCODER side needs: [`Self::to_encoder_table`] reads + /// only `bits` + `max_num_bits`. The decode lookup-table fill + /// (`bit_ranks` / `packed_decode` / `rank_indexes`) is decoder-only and + /// lives in [`Self::build_table_from_weights`]. + fn compute_huffman_bits(&mut self) -> Result { use HuffmanTableError as err; self.bits.clear(); @@ -615,6 +635,12 @@ impl HuffmanTable { return Err(err::MaxBitsTooHigh { got: max_bits }); } + Ok(max_bits) + } + + fn build_table_from_weights(&mut self) -> Result<(), HuffmanTableError> { + let max_bits = self.compute_huffman_bits()?; + self.bit_ranks.clear(); self.bit_ranks.resize((max_bits + 1) as usize, 0); for num_bits in &self.bits {