mzstd 0.1.0

//! Segment jobs and frame assembly of parallel compression.
//!
//! A frame is cut into fixed segments. Each segment is compressed by its own
//! engine into non-final blocks: the engine first runs over up to an
//! overlap's worth of the preceding input, whose blocks are discarded, so the
//! segment's matches may reach back into it; then its entropy tables and
//! repeat offsets are reset, so the first kept block refers to nothing the
//! decoder has not seen. Assembly writes the frame header, the segments'
//! blocks in order, an empty last block and the checksum.
//!
//! A segment's bytes depend only on the parameters, the overlap and the
//! segment itself, so the frame is the same however many tasks produced it
//! and in whatever order they ran.

use alloc::vec;
use alloc::vec::Vec;

use crate::compressor::core::cctx::{AppliedParams, Engine, RequestedParams, compress_bound};
use crate::compressor::core::error::CoreError;
use crate::compressor::core::params::{CONTENT_SIZE_UNKNOWN, Strategy};
use crate::compressor::core::window::{Addr, Segment, Segments};
use crate::shared::format::{MAGIC_NUMBER, WINDOW_LOG_ABSOLUTE_MIN};
use crate::shared::xxhash::Xxh64;

/// Parameters every segment of one frame shares.
#[derive(Copy, Clone, Debug)]
pub(crate) struct FramePlan {
    pub(crate) applied: AppliedParams,
    /// Declared content size, when known.
    pub(crate) content_size: Option<u64>,
}

impl FramePlan {
    /// Resolves the frame's parameters for `content_size` (unknown when
    /// `None`).
    pub(crate) fn new(requested: &RequestedParams, content_size: Option<u64>) -> Self {
        let mut applied = requested.resolve(content_size.unwrap_or(CONTENT_SIZE_UNKNOWN));
        if content_size.is_none() {
            applied.fparams.content_size = false;
        }
        Self {
            applied,
            content_size,
        }
    }

    /// History a segment may reach into: an eighth of the window, or the
    /// whole window for the ultra strategies.
    pub(crate) fn overlap(&self) -> usize {
        let window = 1usize << self.applied.cparams.window_log;
        if self.applied.cparams.strategy >= Strategy::BtUltra {
            window
        } else {
            window >> 3
        }
    }

    /// The frame header (`ZSTD_writeFrameHeader` without dictionaries).
    pub(crate) fn header(&self) -> Vec<u8> {
        let fparams = self.applied.fparams;
        let window_log = self.applied.cparams.window_log;
        let window_size = 1u64 << window_log;
        let pledged = self.content_size.unwrap_or(0);
        let content_size = fparams.content_size && self.content_size.is_some();
        let single_segment = content_size && window_size >= pledged;
        let fcs_code: u8 = if content_size {
            u8::from(pledged >= 256)
                + u8::from(pledged >= 65536 + 256)
                + u8::from(pledged >= 0xFFFF_FFFF)
        } else {
            0
        };
        let descriptor =
            (u8::from(fparams.checksum) << 2) + (u8::from(single_segment) << 5) + (fcs_code << 6);
        let mut out = Vec::with_capacity(18);
        out.extend_from_slice(&MAGIC_NUMBER.to_le_bytes());
        out.push(descriptor);
        if !single_segment {
            out.push(((window_log - WINDOW_LOG_ABSOLUTE_MIN) << 3) as u8);
        }
        match fcs_code {
            0 if single_segment => out.push(pledged as u8),
            0 => {}
            1 => out.extend_from_slice(&((pledged - 256) as u16).to_le_bytes()),
            2 => out.extend_from_slice(&(pledged as u32).to_le_bytes()),
            _ => out.extend_from_slice(&pledged.to_le_bytes()),
        }
        out
    }

    /// The frame's end: an empty last block and, when enabled, the
    /// checksum of `hasher`'s content.
    pub(crate) fn epilogue(&self, hasher: &Xxh64) -> Vec<u8> {
        let mut out = vec![1, 0, 0];
        if self.applied.fparams.checksum {
            out.extend_from_slice(&(hasher.digest() as u32).to_le_bytes());
        }
        out
    }
}

/// Compresses one segment.
///
/// `data` holds the overlap followed by the segment; `overlap` is the
/// overlap's length. Returns the segment's blocks, none of them last.
///
/// # Errors
///
/// Whatever the engine reports; with a destination of `compressBound` bytes
/// it reports nothing for valid input.
#[cfg_attr(all(feature = "hotpath", not(feature = "no_std")), hotpath::measure)]
pub(crate) fn compress_segment(
    engine: &mut Engine,
    plan: &FramePlan,
    data: &[u8],
    overlap: usize,
) -> Result<Vec<u8>, CoreError> {
    let mut applied = plan.applied;
    applied.fparams.checksum = false;
    engine.begin(applied, CONTENT_SIZE_UNKNOWN);
    let segments = Segments {
        buffer: &[],
        input: (data, 1),
    };
    let segment_len = data.len() - overlap;
    let mut scratch = vec![0u8; compress_bound(overlap.max(1)).unwrap_or(0) + 32];
    // Writes the header and compresses the overlap; all of it is discarded.
    engine.compress_continue(
        &segments,
        &mut scratch,
        Addr::new(Segment::Input(1), 0),
        overlap,
        false,
    )?;
    engine.prev.entropy.reset();
    engine.prev.rep = [0; 3];
    let mut out = vec![0u8; compress_bound(segment_len).unwrap_or(0) + 32];
    let size = engine.compress_continue(
        &segments,
        &mut out,
        Addr::new(Segment::Input(1), overlap as i64),
        segment_len,
        false,
    )?;
    out.truncate(size);
    Ok(out)
}

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

    fn frame(plan: &FramePlan, input: &[u8], segment: usize) -> Vec<u8> {
        let mut out = plan.header();
        let mut engine = Engine::new();
        let mut start = 0;
        while start < input.len() {
            let end = (start + segment).min(input.len());
            let from = start.saturating_sub(plan.overlap());
            out.extend(
                compress_segment(&mut engine, plan, &input[from..end], start - from).unwrap(),
            );
            start = end;
        }
        let mut hasher = Xxh64::new(0);
        hasher.update(input);
        out.extend(plan.epilogue(&hasher));
        out
    }

    #[test]
    fn assembled_frames_decode_with_the_reference() {
        let input = corpus::mixed(1_200_000, 9);
        for level in [-5, 1, 3, 7, 13, 17, 19] {
            for checksum in [false, true] {
                for size in [None, Some(input.len() as u64)] {
                    let mut requested = RequestedParams {
                        compression_level: level,
                        ..RequestedParams::default()
                    };
                    requested.fparams.checksum = checksum;
                    let plan = FramePlan::new(&requested, size);
                    let compressed = frame(&plan, &input, 300_000);
                    let decoded =
                        zstd_ffi::oracle::decompress_stream(&compressed, 1 << 16, Some(31))
                            .expect("decodes");
                    assert_eq!(decoded, input, "level {level}");
                }
            }
        }
    }
}