Skip to main content

structured_zstd/dictionary/
mod.rs

1//! Code for creating a separate content dictionary.
2//!
3//! Effective dictionaries are up to 1% the size of the complete training body,
4//! and are trained on many examples of the original data.
5//!
6//! Implemented following the paper "Effective construction of
7//! Relative Lempel-Ziv Dictionaries", by Kewen Liao, Matthias Petri,
8//! Alistair Moffat, and Anthony Wirth
9
10// The algorithm is summarized here
11// 1. The text is split into "epochs", or chunks from the original source
12// 2. From within each epoch, we select the "segment", or 1 KiB contiguous section
13//    that's predicted to be the best option to include in the dictionary. Concatenated,
14//    these segments form the dictionary.
15//
16// This segment scoring algorithm operates as follows:
17// For a given epoch:
18//  - Run a reservoir sampler over the entire epoch, creating a
19//    reservoir of n/t, where `t` is the desired number of occurrences
20//    we want the most common k-mers to have
21//  - Have the ability to estimate
22//    the frequency of a given k-mer: `f(w: k-mer)` calculates
23//    the frequency of w in the reservoir using a rolling karp-rabin hash
24//  - The score of a segment is the sum of `f(w)` called on every kmer within the segment
25mod cover;
26mod fastcover;
27mod frequency;
28mod reservoir;
29
30use crate::bit_io::BitWriter;
31use crate::blocks::sequence_section::{
32    MAX_LITERAL_LENGTH_CODE, MAX_MATCH_LENGTH_CODE, MAX_OFFSET_CODE,
33};
34use crate::decoding::dictionary::MAGIC_NUM as DICT_MAGIC_NUM;
35use crate::decoding::sequence_section_decoder::{LL_MAX_LOG, ML_MAX_LOG, OF_MAX_LOG};
36use crate::dictionary::reservoir::create_sample;
37use crate::fse::fse_encoder::{self, build_table_from_symbol_counts};
38use crate::huff0::HuffmanTable as HuffmanDecoderTable;
39use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable as HuffmanEncoderTable};
40use core::cmp::Reverse;
41use cover::*;
42pub use fastcover::{
43    DEFAULT_D_CANDIDATES, DEFAULT_F_CANDIDATES, DEFAULT_K_CANDIDATES, FastCoverParams,
44    FastCoverTuned,
45};
46use std::{
47    boxed::Box,
48    collections::{BinaryHeap, HashMap},
49    format,
50    fs::{self, File},
51    io::{self, Read},
52    path::{Path, PathBuf},
53    // `vec` import covers the `vec![..]` macro used below: this crate is
54    // no_std-with-std-feature, so the std prelude isn't pulled in implicitly
55    // for top-level items in this module. Removing this import fails the
56    // build with `cannot find macro 'vec' in this scope` — verified.
57    vec,
58    vec::Vec,
59};
60
61const MAX_TRAINING_PREALLOC_BYTES: usize = 8 * 1024 * 1024;
62const MAX_HUFFMAN_STATS_BYTES: usize = 64 * 1024;
63
64/// Smallest size a trained dictionary can occupy, whatever it was trained on.
65///
66/// The magic number, the dictionary ID, the three repeat offsets and the
67/// shortest content the writers emit are unconditional; a real dictionary is
68/// larger still, since the entropy tables between them are never empty. Use it
69/// to reject an impossible `dict_size` before spending the corpus: the training
70/// entry points can only discover the true bound once those tables are built.
71pub const MIN_TRAINED_DICT_SIZE: usize = DICT_MAGIC_NUM.len() + 4 + 12 + 8;
72
73/// Tuning knobs for pure-Rust FastCOVER training.
74#[derive(Debug, Clone)]
75pub struct FastCoverOptions {
76    pub optimize: bool,
77    pub split_point: f64,
78    pub accel: usize,
79    pub k: usize,
80    pub d: usize,
81    pub f: u32,
82    pub k_candidates: Vec<usize>,
83    pub d_candidates: Vec<usize>,
84    pub f_candidates: Vec<u32>,
85}
86
87impl Default for FastCoverOptions {
88    fn default() -> Self {
89        Self {
90            optimize: true,
91            split_point: 0.75,
92            accel: 1,
93            k: 256,
94            d: 8,
95            f: 20,
96            k_candidates: DEFAULT_K_CANDIDATES.to_vec(),
97            d_candidates: DEFAULT_D_CANDIDATES.to_vec(),
98            f_candidates: DEFAULT_F_CANDIDATES.to_vec(),
99        }
100    }
101}
102
103#[derive(Debug, Clone, Copy, Default)]
104pub struct FinalizeOptions {
105    pub dict_id: Option<u32>,
106}
107
108/// A set of values that are used during dictionary construction.
109///
110/// Changing these values can improve the resulting dictionary size for certain datasets.
111// TODO: move `k` here.
112pub(super) struct DictParams {
113    /// Segment size.
114    ///
115    /// As found under "4. Experiments - Varying Segment Size" in the original paper, a
116    /// segment size of 2 kiB was effective.
117    ///
118    /// "We explored a range of \[`segment_size`\] values and found the performance of LMC is insensitive
119    /// to \[`segment_size`\]. We fix \[`segment_size`\] to 2kiB
120    ///
121    /// Reasonable range: [16, 2048+]
122    pub segment_size: u32,
123}
124
125/// Creates a "raw content" dictionary, training off of every file in this directory and all
126/// sub-directories.
127///
128/// The resulting dictionary will be approximately `dict_size` or less, and written to `output`.
129///
130/// # Errors
131/// This function returns `Ok(())` if the dictionary was created successfully, and an
132/// `Err(io::Error)` if an error was encountered reading the input directory or
133/// writing dictionary bytes to `output`.
134///
135/// # Examples
136/// ```no_run
137/// use std::fs::File;
138/// // Create a roughly 1mb dictionary, training off of file in `sample_files`
139/// let input_folder = "sample_files/";
140/// let mut output = File::create("output.dict").unwrap();
141/// structured_zstd::dictionary::create_raw_dict_from_dir(input_folder, &mut output, 1_000_000)
142///     .expect("dictionary training from sample_files should succeed");
143/// ```
144pub fn create_raw_dict_from_dir<P: AsRef<Path>, W: io::Write>(
145    path: P,
146    output: &mut W,
147    dict_size: usize,
148) -> Result<(), io::Error> {
149    // Collect a list of a path to every file in the directory into `file_paths`
150    let mut file_paths: Vec<PathBuf> = Vec::new();
151    let dir: fs::ReadDir = fs::read_dir(path)?;
152    fn recurse_read(dir: fs::ReadDir, file_paths: &mut Vec<PathBuf>) -> Result<(), io::Error> {
153        for entry in dir {
154            let entry = entry?;
155            if entry.file_type()?.is_dir() {
156                recurse_read(fs::read_dir(entry.path())?, file_paths)?;
157            } else {
158                file_paths.push(entry.path());
159            }
160        }
161        Ok(())
162    }
163    recurse_read(dir, &mut file_paths)?;
164
165    // Open each file and chain the readers together
166    let mut total_file_len: u64 = 0;
167    let mut file_handles: Vec<fs::File> = Vec::new();
168    for path in file_paths {
169        let handle = File::open(path)?;
170        total_file_len += handle.metadata()?.len();
171        file_handles.push(handle);
172    }
173    let empty_reader: Box<dyn Read> = Box::new(io::empty());
174    let chained_files = file_handles
175        .iter()
176        .fold(empty_reader, |acc, reader| Box::new(acc.chain(reader)));
177
178    // Create a dict using the new reader
179    create_raw_dict_from_source(chained_files, total_file_len as usize, output, dict_size)?;
180    Ok(())
181}
182
183/// Read from `source` to create a "raw content" dictionary of `dict_size`.
184/// The completed dictionary is written to `output`.
185///
186/// - `source` will be used as training data for the entire dictionary.
187/// - `source_size` is used only as a preallocation hint before reading `source` and
188///   does not affect sampling once all data has been buffered.
189/// - `output` is where the completed dictionary will be written.
190/// - `dict_size` determines how large the complete dictionary should be. The completed
191///   dictionary will be this size or smaller.
192///
193/// This function reads the entire `source` into an in-memory `Vec<u8>` before building
194/// the dictionary. The provided reader need not be buffered, but callers should avoid
195/// sources too large to fit comfortably in memory.
196///
197/// # API note
198/// This public API returns `io::Result<()>` and propagates source/output I/O failures.
199pub fn create_raw_dict_from_source<R: io::Read, W: io::Write>(
200    mut source: R,
201    source_size: usize,
202    output: &mut W,
203    dict_size: usize,
204) -> io::Result<()> {
205    if dict_size == 0 {
206        return Ok(());
207    }
208    let prealloc = source_size.min(MAX_TRAINING_PREALLOC_BYTES);
209    let mut all = Vec::with_capacity(prealloc);
210    source.read_to_end(&mut all)?;
211    if all.is_empty() {
212        return Ok(());
213    }
214
215    if all.len() < K {
216        let keep = usize::min(all.len(), dict_size);
217        output.write_all(&all[all.len() - keep..])?;
218        return Ok(());
219    }
220
221    let source_size = all.len();
222    vprintln!("create_dict: creating {dict_size} byte dict from {source_size} byte source");
223
224    let params = DictParams { segment_size: 2048 };
225    let num_segments = usize::max(1, source_size / params.segment_size as usize);
226    // According to 4. Experiments - Varying Reservoir Sampler Thresholds,
227    // setting reservoir size to collection size / min{collection size / (2 * number of segments),
228    // 256} was effective
229    let denom = usize::max(1, source_size / (2 * num_segments));
230    let sample_scale = usize::max(1, usize::min(denom, 256));
231    let mut sample_size = source_size / sample_scale;
232    sample_size = usize::max(sample_size, usize::min(source_size, 16));
233    vprintln!("create_dict: creating {sample_size} byte sample of collection");
234    let mut sample_reader = all.as_slice();
235    let collection_sample = create_sample(&mut sample_reader, sample_size);
236
237    // A collection of segments to be used in the final dictionary.
238    //
239    // Contains the best segment from every epoch.
240    // Reverse is used because we want a min heap, where
241    // the lowest scoring items come first
242    let mut pool: BinaryHeap<Reverse<Segment>> = BinaryHeap::new();
243    let (num_epochs, epoch_size_kmers) = compute_epoch_info(&params, dict_size, source_size / K);
244    // Plain `*`/`+` throughout the epoch walk below: epochs partition the
245    // training source, so `epoch_size_kmers * K`, `epoch_idx * epoch_size`, and
246    // `start + epoch_size` are all bounded by the source length (<= isize::MAX)
247    // and cannot overflow usize.
248    let epoch_size = usize::max(K, epoch_size_kmers * K);
249    vprintln!("create_dict: computed epoch info, using {num_epochs} epochs of {epoch_size} bytes");
250    let mut epoch_counter = 0;
251    let mut ctx = Context {
252        frequencies: HashMap::with_capacity(epoch_size / K),
253    };
254    // Score each segment in each planned epoch and select the highest-scoring
255    // segment for the pool. Keep exactly `num_epochs` windows to avoid
256    // emitting more segments than the requested dictionary budget allows.
257    for epoch_idx in 0..num_epochs {
258        let start = epoch_idx * epoch_size;
259        if start >= all.len() {
260            break;
261        }
262        let end = if epoch_idx + 1 == num_epochs {
263            all.len()
264        } else {
265            usize::min(start + epoch_size, all.len())
266        };
267        let epoch = &all[start..end];
268        epoch_counter += 1;
269        let best_segment = pick_best_segment(&params, &mut ctx, epoch, &collection_sample);
270        vprintln!(
271            "\tcreate_dict: epoch {epoch_counter}/{num_epochs} has best segment score {}",
272            best_segment.score
273        );
274        pool.push(Reverse(best_segment));
275        // Wipe frequency list for next epoch
276        ctx.frequencies.clear();
277    }
278    vprintln!(
279        "create_dict: {epoch_counter} epochs written, writing {} segments",
280        pool.len()
281    );
282    // Write the dictionary with the highest scoring segment last because
283    // closer items can be represented with a smaller offset
284    while let Some(segment) = pool.pop() {
285        output.write_all(&segment.0.raw)?;
286    }
287    Ok(())
288}
289
290/// The `i`th of [`MAX_HUFFMAN_STATS_BYTES`] samples spread evenly over `len`
291/// bytes.
292///
293/// Computed in 64 bits: `i * len` reaches 2^48 for an addressable corpus, which
294/// a 32-bit `usize` cannot hold — the product overflows for any corpus past
295/// 64 KiB there, and the multiply panics rather than sampling. The quotient is
296/// always below `len`, so the narrowing back is exact.
297fn strided_index(i: usize, len: usize) -> usize {
298    ((i as u64 * len as u64) / MAX_HUFFMAN_STATS_BYTES as u64) as usize
299}
300
301fn serialize_huffman_table(sample_data: &[u8], raw_content: &[u8]) -> io::Result<Vec<u8>> {
302    fn bounded_huffman_stats(data: &[u8]) -> Vec<u8> {
303        if data.len() <= MAX_HUFFMAN_STATS_BYTES {
304            return data.to_vec();
305        }
306
307        let mut stats = Vec::with_capacity(MAX_HUFFMAN_STATS_BYTES);
308        for i in 0..MAX_HUFFMAN_STATS_BYTES {
309            stats.push(data[strided_index(i, data.len())]);
310        }
311        stats
312    }
313
314    let source = if sample_data.len() >= 2 {
315        sample_data
316    } else {
317        raw_content
318    };
319    let mut stats = bounded_huffman_stats(source);
320    if stats.len() < 2 || stats.iter().all(|b| *b == stats[0]) {
321        // A corpus with no distribution to measure gets a synthetic one. It
322        // stops at 128 symbols because a perfectly flat alphabet gives every
323        // symbol the same weight: FSE cannot encode that (an RLE weight
324        // stream), and the direct nibble form addresses at most 128 symbols, so
325        // a full 0..=255 alphabet would have no description at all.
326        stats = (0u8..128).collect();
327    }
328
329    let mut table = HuffmanEncoderTable::build_from_data(stats.as_slice());
330    if table.writeable_table_description_size().is_none() {
331        // Sampled real data can land on the same shape: a flat alphabet wider
332        // than 128 symbols. Fall back to the synthetic narrow one, which always
333        // has a description.
334        stats = (0u8..128).collect();
335        table = HuffmanEncoderTable::build_from_data(stats.as_slice());
336    }
337    let mut writer = BitWriter::new();
338    let mut encoder = HuffmanEncoder::new(&table, &mut writer);
339    encoder.encode(&[stats[0]], true);
340    let encoded = writer.dump();
341
342    let mut decoder = HuffmanDecoderTable::new();
343    let table_size = decoder
344        .build_decoder(encoded.as_slice())
345        .map_err(|e| io::Error::other(format!("failed to decode generated huffman table: {e}")))?;
346    Ok(encoded[..table_size as usize].to_vec())
347}
348
349fn serialize_fse_table(table: &fse_encoder::FSETable) -> Vec<u8> {
350    let mut writer = BitWriter::new();
351    table.write_table(&mut writer);
352    writer.dump()
353}
354
355fn bounded_fse_symbols(data: &[u8], max_symbol: u8) -> Vec<u8> {
356    let modulo = u16::from(max_symbol) + 1;
357    if data.is_empty() {
358        return Vec::from([0u8]);
359    }
360    if data.len() <= MAX_HUFFMAN_STATS_BYTES {
361        return data
362            .iter()
363            .map(|b| (u16::from(*b) % modulo) as u8)
364            .collect();
365    }
366
367    let mut out = Vec::with_capacity(MAX_HUFFMAN_STATS_BYTES);
368    for i in 0..MAX_HUFFMAN_STATS_BYTES {
369        let idx = strided_index(i, data.len());
370        out.push((u16::from(data[idx]) % modulo) as u8);
371    }
372    out
373}
374
375fn serialize_fse_table_from_corpus(
376    sample_data: &[u8],
377    raw_content: &[u8],
378    max_symbol: u8,
379    max_log: u8,
380) -> io::Result<Vec<u8>> {
381    fn counts_total_for_source(source: &[u8], max_symbol: u8, counts: &mut [usize]) -> usize {
382        counts.fill(0);
383        for symbol in bounded_fse_symbols(source, max_symbol) {
384            counts[usize::from(symbol)] += 1;
385        }
386        counts.iter().sum::<usize>()
387    }
388
389    let mut counts = vec![0usize; usize::from(max_symbol) + 1];
390    let using_sample = !sample_data.is_empty();
391    let mut total = counts_total_for_source(
392        if using_sample {
393            sample_data
394        } else {
395            raw_content
396        },
397        max_symbol,
398        &mut counts,
399    );
400    if total <= 1 && using_sample && !raw_content.is_empty() {
401        total = counts_total_for_source(raw_content, max_symbol, &mut counts);
402    }
403    if total <= 1 {
404        return Err(io::Error::new(
405            io::ErrorKind::InvalidInput,
406            "insufficient symbol statistics for FSE table",
407        ));
408    }
409    let table = build_table_from_symbol_counts(&counts, max_log, false);
410    Ok(serialize_fse_table(&table))
411}
412
413fn finalized_content_budget(
414    sample_data: &[u8],
415    raw_fallback: &[u8],
416    dict_size: usize,
417) -> io::Result<usize> {
418    let min_content_size = 8usize;
419    let huf_len = serialize_huffman_table(sample_data, raw_fallback)?.len();
420    let of_len =
421        serialize_fse_table_from_corpus(sample_data, raw_fallback, MAX_OFFSET_CODE, OF_MAX_LOG)?
422            .len();
423    let ml_len = serialize_fse_table_from_corpus(
424        sample_data,
425        raw_fallback,
426        MAX_MATCH_LENGTH_CODE,
427        ML_MAX_LOG,
428    )?
429    .len();
430    let ll_len = serialize_fse_table_from_corpus(
431        sample_data,
432        raw_fallback,
433        MAX_LITERAL_LENGTH_CODE,
434        LL_MAX_LOG,
435    )?
436    .len();
437
438    let header_len = DICT_MAGIC_NUM.len() + 4 + huf_len + of_len + ml_len + ll_len + 12;
439    let max_content_budget = dict_size.saturating_sub(header_len);
440    if max_content_budget < min_content_size {
441        return Err(io::Error::new(
442            io::ErrorKind::InvalidInput,
443            "dictionary size too small to fit header and offset history",
444        ));
445    }
446    Ok(max_content_budget)
447}
448
449fn derive_dict_id(raw_content: &[u8]) -> u32 {
450    let mut h = 0xcbf29ce484222325u64;
451    for &b in raw_content {
452        h ^= u64::from(b);
453        h = h.wrapping_mul(0x100000001b3);
454    }
455    let compliant = (h % ((1u64 << 31) - 32768)) + 32768;
456    compliant as u32
457}
458
459/// Finalize raw dictionary content into a full zstd dictionary binary
460/// (`magic + dict_id + entropy tables + offset history + content`).
461pub fn finalize_raw_dict(
462    raw_content: &[u8],
463    sample_data: &[u8],
464    dict_size: usize,
465    options: FinalizeOptions,
466) -> io::Result<Vec<u8>> {
467    if raw_content.is_empty() {
468        return Err(io::Error::new(
469            io::ErrorKind::InvalidInput,
470            "raw dictionary content must not be empty",
471        ));
472    }
473    let mut out = Vec::with_capacity(dict_size.max(256));
474    out.extend_from_slice(&DICT_MAGIC_NUM);
475    let dict_id = options
476        .dict_id
477        .unwrap_or_else(|| derive_dict_id(raw_content));
478    if dict_id == 0 {
479        return Err(io::Error::new(
480            io::ErrorKind::InvalidInput,
481            "dictionary id must be non-zero",
482        ));
483    }
484    out.extend_from_slice(&dict_id.to_le_bytes());
485    out.extend_from_slice(serialize_huffman_table(sample_data, raw_content)?.as_slice());
486    out.extend_from_slice(
487        serialize_fse_table_from_corpus(sample_data, raw_content, MAX_OFFSET_CODE, OF_MAX_LOG)?
488            .as_slice(),
489    );
490    out.extend_from_slice(
491        serialize_fse_table_from_corpus(
492            sample_data,
493            raw_content,
494            MAX_MATCH_LENGTH_CODE,
495            ML_MAX_LOG,
496        )?
497        .as_slice(),
498    );
499    out.extend_from_slice(
500        serialize_fse_table_from_corpus(
501            sample_data,
502            raw_content,
503            MAX_LITERAL_LENGTH_CODE,
504            LL_MAX_LOG,
505        )?
506        .as_slice(),
507    );
508
509    // Repeat offsets: keep default bootstrap history.
510    out.extend_from_slice(&1u32.to_le_bytes());
511    out.extend_from_slice(&4u32.to_le_bytes());
512    out.extend_from_slice(&8u32.to_le_bytes());
513
514    let min_content_size = 8usize;
515    let max_content_budget = dict_size.saturating_sub(out.len());
516    if max_content_budget < min_content_size {
517        return Err(io::Error::new(
518            io::ErrorKind::InvalidInput,
519            "dictionary size too small to fit header and offset history",
520        ));
521    }
522
523    let content = if raw_content.len() > max_content_budget {
524        &raw_content[raw_content.len() - max_content_budget..]
525    } else {
526        raw_content
527    };
528    if content.len() < min_content_size {
529        out.resize(out.len() + (min_content_size - content.len()), 0);
530    }
531    out.extend_from_slice(content);
532    Ok(out)
533}
534
535/// Train a raw FastCOVER dictionary from a source stream.
536fn train_fastcover_internal(
537    sample: &[u8],
538    dict_size: usize,
539    options: &FastCoverOptions,
540) -> (Vec<u8>, FastCoverTuned) {
541    if options.optimize {
542        fastcover::optimize_fastcover_raw(
543            sample,
544            dict_size,
545            options.split_point,
546            options.accel,
547            options.d_candidates.as_slice(),
548            options.f_candidates.as_slice(),
549            options.k_candidates.as_slice(),
550        )
551    } else {
552        let params = fastcover::normalize_fastcover_params(FastCoverParams {
553            k: options.k,
554            d: options.d,
555            f: options.f,
556            accel: options.accel,
557        });
558        (
559            fastcover::train_fastcover_raw(sample, dict_size, params),
560            FastCoverTuned {
561                k: params.k,
562                d: params.d,
563                f: params.f,
564                accel: params.accel,
565                score: 0,
566            },
567        )
568    }
569}
570
571/// Train a raw FastCOVER dictionary directly from an in-memory sample.
572pub fn train_fastcover_raw_from_slice(
573    sample: &[u8],
574    dict_size: usize,
575    options: &FastCoverOptions,
576) -> io::Result<(Vec<u8>, FastCoverTuned)> {
577    if sample.is_empty() {
578        return Err(io::Error::new(
579            io::ErrorKind::InvalidInput,
580            "source stream is empty",
581        ));
582    }
583    let (dict, tuned) = train_fastcover_internal(sample, dict_size, options);
584    if dict.is_empty() && dict_size > 0 {
585        return Err(io::Error::new(
586            io::ErrorKind::InvalidInput,
587            "training sample is too small for FastCOVER",
588        ));
589    }
590    Ok((dict, tuned))
591}
592
593/// Train a raw FastCOVER dictionary from a source stream.
594///
595/// This function fully buffers the entire training corpus into memory via
596/// `read_to_end`, which can consume significant RAM for large inputs.
597pub fn create_fastcover_raw_dict_from_source<R: io::Read, W: io::Write>(
598    mut source: R,
599    output: &mut W,
600    dict_size: usize,
601    options: &FastCoverOptions,
602) -> io::Result<FastCoverTuned> {
603    let mut sample = Vec::new();
604    source.read_to_end(&mut sample)?;
605    let (dict, tuned) = train_fastcover_raw_from_slice(sample.as_slice(), dict_size, options)?;
606    output.write_all(dict.as_slice())?;
607    Ok(tuned)
608}
609
610/// Train and finalize a FastCOVER dictionary in pure Rust.
611///
612/// This function fully buffers the entire training corpus into memory via
613/// `read_to_end`, which can consume significant RAM for large inputs.
614pub fn create_fastcover_dict_from_source<R: io::Read, W: io::Write>(
615    mut source: R,
616    output: &mut W,
617    dict_size: usize,
618    fastcover: &FastCoverOptions,
619    finalize: FinalizeOptions,
620) -> io::Result<FastCoverTuned> {
621    let mut sample = Vec::new();
622    source.read_to_end(&mut sample)?;
623    create_fastcover_dict_from_slice(sample.as_slice(), output, dict_size, fastcover, finalize)
624}
625
626/// Train and finalize a FastCOVER dictionary from a corpus already in memory.
627///
628/// The same work as [`create_fastcover_dict_from_source`] for a caller that
629/// holds the bytes: the corpus is the largest allocation training makes, and
630/// handing it over as a slice keeps it to one copy rather than buffering it a
631/// second time inside.
632pub fn create_fastcover_dict_from_slice<W: io::Write>(
633    sample: &[u8],
634    output: &mut W,
635    dict_size: usize,
636    fastcover: &FastCoverOptions,
637    finalize: FinalizeOptions,
638) -> io::Result<FastCoverTuned> {
639    if sample.is_empty() {
640        return Err(io::Error::new(
641            io::ErrorKind::InvalidInput,
642            "source stream is empty",
643        ));
644    }
645    let content_budget = finalized_content_budget(sample, sample, dict_size)?;
646    let (raw_dict, tuned) = train_fastcover_raw_from_slice(sample, content_budget, fastcover)?;
647
648    let finalized = finalize_raw_dict(raw_dict.as_slice(), sample, dict_size, finalize)?;
649    output.write_all(finalized.as_slice())?;
650    Ok(tuned)
651}
652
653/// Build a finalized FastCOVER dictionary, attach it to a fastest-level
654/// frame compressor, and compress a fresh payload. Returns
655/// `(finalized_dictionary, compressed_frame, original_payload)` so a
656/// roundtrip check can decode `compressed_frame` against
657/// `finalized_dictionary` and compare to `original_payload`. The
658/// C-decoder roundtrip that consumes this lives in the `ffi-bench` crate;
659/// this side stays pure Rust.
660#[cfg(feature = "bench-internals")]
661pub(crate) fn dict_roundtrip_fixture() -> (
662    alloc::vec::Vec<u8>,
663    alloc::vec::Vec<u8>,
664    alloc::vec::Vec<u8>,
665) {
666    use crate::decoding::Dictionary;
667    use crate::encoding::{CompressionLevel, FrameCompressor};
668
669    let mut sample = alloc::vec::Vec::new();
670    for i in 0..512u32 {
671        sample.extend_from_slice(
672            alloc::format!(
673                "tenant=demo table=orders key={i} region=eu payload=aaaaabbbbbcccccdddddeeeee\n"
674            )
675            .as_bytes(),
676        );
677    }
678
679    let dict_size = 4096usize;
680    let content_budget = finalized_content_budget(sample.as_slice(), sample.as_slice(), dict_size)
681        .expect("content budget should be computable");
682    let raw = fastcover::train_fastcover_raw(
683        sample.as_slice(),
684        content_budget,
685        fastcover::FastCoverParams {
686            k: 256,
687            d: 8,
688            f: 20,
689            accel: 1,
690        },
691    );
692    let finalized = finalize_raw_dict(
693        raw.as_slice(),
694        sample.as_slice(),
695        dict_size,
696        FinalizeOptions::default(),
697    )
698    .expect("finalization should succeed");
699    let parsed =
700        Dictionary::decode_dict(finalized.as_slice()).expect("finalized dictionary should parse");
701    assert!(!parsed.dict_content.is_empty());
702
703    let mut payload = alloc::vec::Vec::new();
704    for idx in 0..96u32 {
705        payload.extend_from_slice(
706            alloc::format!("tenant=demo op=put key={idx} value=aaaaabbbbbcccccdddddeeeee\n")
707                .as_bytes(),
708        );
709    }
710
711    let mut compressed = alloc::vec::Vec::new();
712    let mut compressor = FrameCompressor::new(CompressionLevel::Fastest);
713    compressor
714        .set_dictionary(parsed)
715        .expect("dictionary should attach");
716    compressor.set_source(payload.as_slice());
717    compressor.set_drain(&mut compressed);
718    compressor.compress();
719
720    (finalized, compressed, payload)
721}
722
723#[cfg(test)]
724mod tests;