zstd/

dict.rs

//! Train a dictionary from various sources.
//!
//! A dictionary can help improve the compression of small files.
//! The dictionary must be present during decompression,
//! but can be shared across multiple "similar" files.
//!
//! Creating a dictionary using the `zstd` C library,
//! using the `zstd` command-line interface, using this library,
//! or using the `train` binary provided, should give the same result,
//! and are therefore completely compatible.
//!
//! To use, see [`Encoder::with_dictionary`] or [`Decoder::with_dictionary`].
//!
//! [`Encoder::with_dictionary`]: ../struct.Encoder.html#method.with_dictionary
//! [`Decoder::with_dictionary`]: ../struct.Decoder.html#method.with_dictionary

#[cfg(feature = "zdict_builder")]
use std::io::{self, Read};

pub use zstd_safe::{CDict, DDict};

/// Prepared dictionary for compression
///
/// A dictionary can include its own copy of the data (if it is `'static`), or it can merely point
/// to a separate buffer (if it has another lifetime).
pub struct EncoderDictionary<'a> {
    cdict: CDict<'a>,
}

/// The error for a dictionary zstd could not load.
fn invalid_dictionary() -> std::io::Error {
    std::io::Error::new(
        std::io::ErrorKind::InvalidData,
        "zstd could not load the dictionary",
    )
}

impl EncoderDictionary<'static> {
    /// Creates a prepared dictionary for compression.
    ///
    /// This will copy the dictionary internally.
    ///
    /// # Panics
    ///
    /// If zstd cannot load the dictionary - for instance, it starts with the
    /// dictionary magic number but its header is corrupt. Use
    /// [`EncoderDictionary::try_copy`] for a dictionary you did not make
    /// yourself.
    pub fn copy(dictionary: &[u8], level: i32) -> Self {
        Self::try_copy(dictionary, level)
            .expect("zstd returned null pointer when creating dict")
    }

    /// Creates a prepared dictionary for compression, or an error if zstd
    /// cannot load it.
    ///
    /// This will copy the dictionary internally.
    pub fn try_copy(dictionary: &[u8], level: i32) -> std::io::Result<Self> {
        let cdict = zstd_safe::CDict::try_create(dictionary, level)
            .ok_or_else(invalid_dictionary)?;
        Ok(Self { cdict })
    }
}

impl<'a> EncoderDictionary<'a> {
    #[cfg(feature = "experimental")]
    #[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "experimental")))]
    /// Create prepared dictionary for compression
    ///
    /// A level of `0` uses zstd's default (currently `3`).
    ///
    /// Only available with the `experimental` feature. Use `EncoderDictionary::copy` otherwise.
    ///
    /// # Panics
    ///
    /// If zstd cannot load the dictionary. Use [`EncoderDictionary::try_new`]
    /// for a dictionary you did not make yourself.
    pub fn new(dictionary: &'a [u8], level: i32) -> Self {
        Self::try_new(dictionary, level).expect("zstd returned null pointer")
    }

    #[cfg(feature = "experimental")]
    #[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "experimental")))]
    /// Create prepared dictionary for compression, or an error if zstd cannot
    /// load it.
    ///
    /// Only available with the `experimental` feature. Use
    /// `EncoderDictionary::try_copy` otherwise.
    pub fn try_new(dictionary: &'a [u8], level: i32) -> std::io::Result<Self> {
        let cdict =
            zstd_safe::CDict::try_create_by_reference(dictionary, level)
                .ok_or_else(invalid_dictionary)?;
        Ok(Self { cdict })
    }

    /// Returns reference to `CDict` inner object
    pub fn as_cdict(&self) -> &CDict<'a> {
        &self.cdict
    }
}

/// Prepared dictionary for decompression
pub struct DecoderDictionary<'a> {
    ddict: DDict<'a>,
}

impl DecoderDictionary<'static> {
    /// Create a prepared dictionary for decompression.
    ///
    /// This will copy the dictionary internally.
    ///
    /// # Panics
    ///
    /// If zstd cannot load the dictionary - for instance, it starts with the
    /// dictionary magic number but its header is corrupt. Use
    /// [`DecoderDictionary::try_copy`] for a dictionary you did not make
    /// yourself.
    pub fn copy(dictionary: &[u8]) -> Self {
        Self::try_copy(dictionary)
            .expect("zstd returned null pointer when creating dict")
    }

    /// Create a prepared dictionary for decompression, or an error if zstd
    /// cannot load it.
    ///
    /// This will copy the dictionary internally.
    pub fn try_copy(dictionary: &[u8]) -> std::io::Result<Self> {
        let ddict = zstd_safe::DDict::try_create(dictionary)
            .ok_or_else(invalid_dictionary)?;
        Ok(Self { ddict })
    }
}

impl<'a> DecoderDictionary<'a> {
    #[cfg(feature = "experimental")]
    #[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "experimental")))]
    /// Create prepared dictionary for decompression
    ///
    /// Only available with the `experimental` feature. Use `DecoderDictionary::copy` otherwise.
    ///
    /// # Panics
    ///
    /// If zstd cannot load the dictionary. Use [`DecoderDictionary::try_new`]
    /// for a dictionary you did not make yourself.
    pub fn new(dict: &'a [u8]) -> Self {
        Self::try_new(dict).expect("zstd returned null pointer")
    }

    #[cfg(feature = "experimental")]
    #[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "experimental")))]
    /// Create prepared dictionary for decompression, or an error if zstd
    /// cannot load it.
    ///
    /// Only available with the `experimental` feature. Use
    /// `DecoderDictionary::try_copy` otherwise.
    pub fn try_new(dict: &'a [u8]) -> std::io::Result<Self> {
        let ddict = zstd_safe::DDict::try_create_by_reference(dict)
            .ok_or_else(invalid_dictionary)?;
        Ok(Self { ddict })
    }

    /// Returns reference to `DDict` inner object
    pub fn as_ddict(&self) -> &DDict<'a> {
        &self.ddict
    }
}

/// Train a dictionary from a big continuous chunk of data, with all samples
/// contiguous in memory.
///
/// This is the most efficient way to train a dictionary,
/// since this is directly fed into `zstd`.
///
/// * `sample_data` is the concatenation of all sample data.
/// * `sample_sizes` is the size of each sample in `sample_data`.
///   The sum of all `sample_sizes` should equal the length of `sample_data`.
/// * `max_size` is the maximum size of the dictionary to generate.
///
/// The result is the dictionary data. You can, for example, feed it to [`CDict::create`].
#[cfg(feature = "zdict_builder")]
#[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "zdict_builder")))]
pub fn from_continuous(
    sample_data: &[u8],
    sample_sizes: &[usize],
    max_size: usize,
) -> io::Result<Vec<u8>> {
    use crate::map_error_code;

    // Complain if the lengths don't add up to the entire data.
    // Add them with overflow detection: a wrapping sum would accept sizes
    // reaching far outside of `sample_data`.
    let total_size = sample_sizes
        .iter()
        .try_fold(0usize, |total, &size| total.checked_add(size));
    if total_size != Some(sample_data.len()) {
        return Err(io::Error::new(
            io::ErrorKind::Other,
            "sample sizes don't add up".to_string(),
        ));
    }

    let mut result = Vec::with_capacity(max_size);
    zstd_safe::train_from_buffer(&mut result, sample_data, sample_sizes)
        .map_err(map_error_code)?;
    Ok(result)
}

/// Train a dictionary from multiple samples.
///
/// The samples will internally be copied to a single continuous buffer,
/// so make sure you have enough memory available.
///
/// If you need to stretch your system's limits,
/// [`from_continuous`] directly uses the given slice.
///
/// [`from_continuous`]: ./fn.from_continuous.html
///
/// * `samples` is a list of individual samples to train on.
/// * `max_size` is the maximum size of the dictionary to generate.
///
/// The result is the dictionary data. You can, for example, feed it to [`CDict::create`].
#[cfg(feature = "zdict_builder")]
#[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "zdict_builder")))]
pub fn from_samples<S: AsRef<[u8]>>(
    samples: &[S],
    max_size: usize,
) -> io::Result<Vec<u8>> {
    // Pre-allocate the entire required size.
    let total_length: usize =
        samples.iter().map(|sample| sample.as_ref().len()).sum();

    let mut data = Vec::with_capacity(total_length);

    // Copy every sample to a big chunk of memory
    data.extend(samples.iter().flat_map(|s| s.as_ref()).cloned());

    let sizes: Vec<_> = samples.iter().map(|s| s.as_ref().len()).collect();

    from_continuous(&data, &sizes, max_size)
}

/// Train a dictionary from multiple samples.
///
/// Unlike [`from_samples`], this does not require having a list of all samples.
/// It also allows running into an error when iterating through the samples.
///
/// They will still be copied to a continuous array and fed to [`from_continuous`].
///
/// * `samples` is an iterator of individual samples to train on.
/// * `max_size` is the maximum size of the dictionary to generate.
///
/// The result is the dictionary data. You can, for example, feed it to [`CDict::create`].
///
/// # Examples
///
/// ```rust,no_run
/// // Train from a couple of json files.
/// let dict_buffer = zstd::dict::from_sample_iterator(
///     ["file_a.json", "file_b.json"]
///         .into_iter()
///         .map(|filename| std::fs::File::open(filename)),
///     10_000,  // 10kB dictionary
/// ).unwrap();
/// ```
///
/// ```rust,no_run
/// use std::io::BufRead as _;
/// // Treat each line from stdin as a separate sample.
/// let dict_buffer = zstd::dict::from_sample_iterator(
///     std::io::stdin().lock().lines().map(|line: std::io::Result<String>| {
///         // Transform each line into a `Cursor<Vec<u8>>` so they implement Read.
///         line.map(String::into_bytes)
///             .map(std::io::Cursor::new)
///     }),
///     10_000,  // 10kB dictionary
/// ).unwrap();
/// ```
#[cfg(feature = "zdict_builder")]
#[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "zdict_builder")))]
pub fn from_sample_iterator<I, R>(
    samples: I,
    max_size: usize,
) -> io::Result<Vec<u8>>
where
    I: IntoIterator<Item = io::Result<R>>,
    R: Read,
{
    let mut data = Vec::new();
    let mut sizes = Vec::new();

    for sample in samples {
        let mut sample = sample?;
        let len = sample.read_to_end(&mut data)?;
        sizes.push(len);
    }

    from_continuous(&data, &sizes, max_size)
}

/// Train a dict from a list of files.
///
/// * `filenames` is an iterator of files to load. Each file will be treated as an individual
///   sample.
/// * `max_size` is the maximum size of the dictionary to generate.
///
/// The result is the dictionary data. You can, for example, feed it to [`CDict::create`].
#[cfg(feature = "zdict_builder")]
#[cfg_attr(feature = "doc-cfg", doc(cfg(feature = "zdict_builder")))]
pub fn from_files<I, P>(filenames: I, max_size: usize) -> io::Result<Vec<u8>>
where
    P: AsRef<std::path::Path>,
    I: IntoIterator<Item = P>,
{
    from_sample_iterator(
        filenames
            .into_iter()
            .map(|filename| std::fs::File::open(filename)),
        max_size,
    )
}

#[cfg(test)]
#[cfg(feature = "zdict_builder")]
mod tests {
    use std::fs;
    use std::io;
    use std::io::Read;

    use walkdir;

    #[test]
    fn test_dict_training() {
        // Train a dictionary
        let paths: Vec<_> = walkdir::WalkDir::new("src")
            .into_iter()
            .map(|entry| entry.unwrap())
            .map(|entry| entry.into_path())
            .filter(|path| path.to_str().unwrap().ends_with(".rs"))
            .collect();

        let dict = super::from_files(&paths, 4000).unwrap();

        for path in paths {
            let mut buffer = Vec::new();
            let mut file = fs::File::open(path).unwrap();
            let mut content = Vec::new();
            file.read_to_end(&mut content).unwrap();
            io::copy(
                &mut &content[..],
                &mut crate::stream::Encoder::with_dictionary(
                    &mut buffer,
                    1,
                    &dict,
                )
                .unwrap()
                .auto_finish(),
            )
            .unwrap();

            let mut result = Vec::new();
            io::copy(
                &mut crate::stream::Decoder::with_dictionary(
                    &buffer[..],
                    &dict[..],
                )
                .unwrap(),
                &mut result,
            )
            .unwrap();

            assert_eq!(&content, &result);
        }
    }

    #[test]
    fn test_from_continuous_overflowing_sizes() {
        // These sizes only add up to the data length by wrapping around.
        let data = [0u8; 100];
        let sizes = [usize::MAX - 99, 200];
        assert!(super::from_continuous(&data, &sizes, 4000).is_err());
    }
}