lz4rip_encode/

dict.rs

//! COVER dictionary trainer.
#![forbid(unsafe_code)]

use alloc::collections::VecDeque;
use alloc::vec;
use alloc::vec::Vec;

use lz4rip_core::MAX_DISTANCE;
use lz4rip_core::MINMATCH;

/// Trains an LZ4 dictionary from sample messages using the COVER
/// algorithm.
///
/// # Example
/// ```
/// use lz4rip_encode::DictTrainer;
///
/// let mut trainer = DictTrainer::new(2048);
/// for msg in &[b"hello world" as &[u8], b"hello rust", b"hello lz4"] {
///     trainer.add_sample(msg);
/// }
/// let dict = trainer.train();
/// let compressor = lz4rip_encode::Compressor::with_dict(&dict);
/// ```
pub struct DictTrainer {
    max_dict_size: usize,
    samples: VecDeque<Vec<u8>>,
    total_bytes: usize,
}

impl DictTrainer {
    /// Create a trainer targeting `max_dict_size` bytes of output.
    ///
    /// Typical values: 2048 for small messages, 4096 for larger ones.
    /// The dict is capped at 65535 bytes (LZ4 max match distance).
    pub fn new(max_dict_size: usize) -> Self {
        let max_dict_size = max_dict_size.min(MAX_DISTANCE);
        DictTrainer {
            max_dict_size,
            samples: VecDeque::new(),
            total_bytes: 0,
        }
    }

    /// Add a training sample.
    ///
    /// Samples shorter than 4 bytes or longer than `max_dict_size` are
    /// silently skipped.
    pub fn add_sample(&mut self, data: &[u8]) {
        if data.len() < MINMATCH || data.len() > self.max_dict_size {
            return;
        }

        let budget = self.max_dict_size * 8;

        while self.total_bytes + data.len() > budget && !self.samples.is_empty() {
            self.total_bytes -= self.samples.pop_front().unwrap().len();
        }

        self.total_bytes += data.len();
        self.samples.push_back(data.to_vec());
    }

    /// Number of samples added so far.
    pub fn sample_count(&self) -> usize {
        self.samples.len()
    }

    /// Total bytes of sample data added so far.
    pub fn total_bytes(&self) -> usize {
        self.total_bytes
    }

    /// Train a dictionary from the collected samples. Consumes the
    /// trainer, freeing all sample data.
    ///
    /// Returns a raw byte buffer. If fewer than 2 samples were added,
    /// returns an empty vec.
    pub fn train(self) -> Vec<u8> {
        if self.samples.len() < 2 {
            return Vec::new();
        }
        let sample_refs: Vec<&[u8]> = self.samples.iter().map(|s| s.as_slice()).collect();
        cover_select(&sample_refs, self.max_dict_size)
    }
}

const D: usize = 8;
const FREQ_BITS: usize = 16;
const FREQ_SIZE: usize = 1 << FREQ_BITS;
const FREQ_MASK: usize = FREQ_SIZE - 1;

#[inline(always)]
fn hash_dmer(data: &[u8]) -> usize {
    let v = u64::from_le_bytes(data[..8].try_into().unwrap());
    v.wrapping_mul(0x9E3779B97F4A7C15) as usize
}

fn cover_select(samples: &[&[u8]], dict_size: usize) -> Vec<u8> {
    let mut concat = Vec::new();
    let mut offsets = Vec::new();
    for &sample in samples {
        offsets.push(concat.len());
        concat.extend_from_slice(sample);
    }
    offsets.push(concat.len());

    if concat.len() < D {
        return concat[..dict_size.min(concat.len())].to_vec();
    }

    let num_dmers = concat.len() - D + 1;

    let mut hashes = vec![0u32; num_dmers];
    for i in 0..num_dmers {
        hashes[i] = (hash_dmer(&concat[i..i + D]) & FREQ_MASK) as u32;
    }

    let mut freqs = vec![0u32; FREQ_SIZE];
    for s in 0..samples.len() {
        let start = offsets[s];
        let end = offsets[s + 1];
        if end - start < D {
            continue;
        }
        for i in start..end - D + 1 {
            freqs[hashes[i] as usize] += 1;
        }
    }

    let k = dict_size / 4;
    if concat.len() < k {
        return concat;
    }

    let seg_dmers = k - D + 1;
    let mut used = vec![false; concat.len()];
    let mut segments: Vec<(usize, u64)> = Vec::new();
    let mut collected = 0usize;

    while collected < dict_size {
        let mut prefix = vec![0u64; num_dmers + 1];
        prefix[0] = 0;
        for i in 0..num_dmers {
            prefix[i + 1] = prefix[i] + freqs[hashes[i] as usize] as u64;
        }

        let mut best_pos = 0;
        let mut best_score = 0u64;
        for pos in 0..=concat.len() - k {
            if !used[pos] {
                let score = prefix[pos + seg_dmers] - prefix[pos];
                if score > best_score {
                    best_score = score;
                    best_pos = pos;
                }
            }
        }

        if best_score == 0 {
            break;
        }

        segments.push((best_pos, best_score));

        for i in best_pos..best_pos + seg_dmers {
            freqs[hashes[i] as usize] = 0;
        }
        used[best_pos..best_pos + k].fill(true);

        collected += k;
    }

    let mut dict = Vec::with_capacity(dict_size);
    for &(pos, _) in segments.iter().rev() {
        let end = (pos + k).min(concat.len());
        dict.extend_from_slice(&concat[pos..end]);
        if dict.len() >= dict_size {
            break;
        }
    }
    dict.truncate(dict_size);
    dict
}

#[cfg(test)]
mod tests {
    use super::*;

    fn json_msg(i: u32) -> Vec<u8> {
        alloc::format!(
            r#"{{"ts":"2026-04-27T12:00:00.{i:04}Z","level":"INFO","service":"api-gw","trace":"{i:08x}","method":"GET","path":"/v1/users/{i:04}","status":200,"latency_ms":{lat},"region":"us-east-1"}}"#,
            i = i,
            lat = 10 + i % 490,
        )
        .into_bytes()
    }

    #[test]
    fn train_produces_nonempty_dict() {
        let mut trainer = DictTrainer::new(2048);
        for i in 0..100 {
            trainer.add_sample(&json_msg(i));
        }
        let dict = trainer.train();
        assert!(!dict.is_empty(), "dict should not be empty");
        assert!(dict.len() <= 2048, "dict should respect max size");
    }

    #[test]
    fn skips_too_short() {
        let mut trainer = DictTrainer::new(2048);
        trainer.add_sample(b"hi");
        assert_eq!(trainer.sample_count(), 0);
    }

    #[test]
    fn skips_too_long() {
        let mut trainer = DictTrainer::new(64);
        trainer.add_sample(&[0u8; 100]);
        assert_eq!(trainer.sample_count(), 0);
    }

    #[test]
    fn evicts_old_samples() {
        let mut trainer = DictTrainer::new(2048);
        for i in 0..200 {
            trainer.add_sample(&json_msg(i));
        }
        assert!(trainer.total_bytes() <= 2048 * 8);
        assert!(trainer.sample_count() < 200);
    }

    #[test]
    fn too_few_samples_returns_empty() {
        let mut trainer = DictTrainer::new(2048);
        trainer.add_sample(b"hello world");
        assert!(trainer.train().is_empty());
    }
}