1mod 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,
58 vec::Vec,
59};
60
61const MAX_TRAINING_PREALLOC_BYTES: usize = 8 * 1024 * 1024;
62const MAX_HUFFMAN_STATS_BYTES: usize = 64 * 1024;
63
64pub const MIN_TRAINED_DICT_SIZE: usize = DICT_MAGIC_NUM.len() + 4 + 12 + 8;
72
73#[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
108pub(super) struct DictParams {
113 pub segment_size: u32,
123}
124
125pub 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 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 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_raw_dict_from_source(chained_files, total_file_len as usize, output, dict_size)?;
180 Ok(())
181}
182
183pub 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 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 let mut pool: BinaryHeap<Reverse<Segment>> = BinaryHeap::new();
243 let (num_epochs, epoch_size_kmers) = compute_epoch_info(¶ms, dict_size, source_size / K);
244 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 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(¶ms, &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 ctx.frequencies.clear();
277 }
278 vprintln!(
279 "create_dict: {epoch_counter} epochs written, writing {} segments",
280 pool.len()
281 );
282 while let Some(segment) = pool.pop() {
285 output.write_all(&segment.0.raw)?;
286 }
287 Ok(())
288}
289
290fn 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 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 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
459pub 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 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
535fn 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
571pub 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
593pub 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
610pub 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
626pub 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#[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;