zstd-rs/rust/src/zstd_compress_api.rs

//! Public, context-free compression API helpers.
//!
//! The stateful compression context remains in the C translation unit for
//! now.  These helpers have no context layout dependency, so they can expose
//! the public ABI directly from Rust.

use crate::errors::{ZstdErrorCode, ERROR};

#[cfg(target_pointer_width = "64")]
const ZSTD_MAX_INPUT_SIZE: usize = 0xFF00_FF00_FF00_FF00;
#[cfg(target_pointer_width = "32")]
const ZSTD_MAX_INPUT_SIZE: usize = 0xFF00_FF00;

const SMALL_INPUT_THRESHOLD: usize = 128 << 10;
const ZSTD_MINMATCH_MIN: usize = 3;
const ZSTD_BLOCKSIZE_MAX_MIN: usize = 1 << 10;

/// ABI-compatible `ZSTD_Sequence` from `zstd.h`.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZstdSequence {
    offset: u32,
    lit_length: u32,
    match_length: u32,
    rep: u32,
}

/// Returns the maximum one-pass compressed size, matching
/// `ZSTD_COMPRESSBOUND()` and its public error contract.
#[no_mangle]
pub extern "C" fn ZSTD_compressBound(src_size: usize) -> usize {
    if src_size >= ZSTD_MAX_INPUT_SIZE {
        return ERROR(ZstdErrorCode::SrcSizeWrong);
    }

    let margin = if src_size < SMALL_INPUT_THRESHOLD {
        (SMALL_INPUT_THRESHOLD - src_size) >> 11
    } else {
        0
    };
    src_size.wrapping_add(src_size >> 8).wrapping_add(margin)
}

/// Returns the maximum number of public sequences generated for an input.
#[no_mangle]
pub extern "C" fn ZSTD_sequenceBound(src_size: usize) -> usize {
    (src_size / ZSTD_MINMATCH_MIN + 1) + (src_size / ZSTD_BLOCKSIZE_MAX_MIN + 1)
}

/// Removes explicit block delimiters, carrying their literals into the next
/// sequence exactly as the C public helper does.
///
/// # Safety
///
/// When `sequences_size` is nonzero, `sequences` must reference that many
/// initialized `ZSTD_Sequence` values.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_mergeBlockDelimiters(
    sequences: *mut ZstdSequence,
    sequences_size: usize,
) -> usize {
    let mut output = 0usize;
    for input in 0..sequences_size {
        let sequence = unsafe { sequences.add(input).read() };
        if sequence.offset == 0 && sequence.match_length == 0 {
            if input != sequences_size - 1 {
                let next = unsafe { &mut *sequences.add(input + 1) };
                next.lit_length = next.lit_length.wrapping_add(sequence.lit_length);
            }
        } else {
            unsafe { sequences.add(output).write(sequence) };
            output += 1;
        }
    }
    output
}

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

    #[test]
    fn compress_bound_matches_the_public_macro_edges() {
        assert_eq!(ZSTD_compressBound(0), 64);
        assert_eq!(ZSTD_compressBound(1), 64);
        assert_eq!(ZSTD_compressBound(SMALL_INPUT_THRESHOLD - 1), 131_582);
        assert_eq!(ZSTD_compressBound(SMALL_INPUT_THRESHOLD), 131_584);
        assert!(ERR_isError(ZSTD_compressBound(ZSTD_MAX_INPUT_SIZE)));
    }

    #[test]
    fn sequence_bound_matches_the_public_formula() {
        assert_eq!(ZSTD_sequenceBound(0), 2);
        assert_eq!(ZSTD_sequenceBound(3), 3);
        assert_eq!(ZSTD_sequenceBound(1024), 344);
    }

    #[test]
    fn merge_block_delimiters_carries_literals_forward() {
        let mut sequences = [
            ZstdSequence {
                offset: 0,
                lit_length: 3,
                match_length: 0,
                rep: 0,
            },
            ZstdSequence {
                offset: 7,
                lit_length: 1,
                match_length: 4,
                rep: 2,
            },
            ZstdSequence {
                offset: 0,
                lit_length: 9,
                match_length: 0,
                rep: 0,
            },
        ];

        assert_eq!(
            unsafe { ZSTD_mergeBlockDelimiters(sequences.as_mut_ptr(), sequences.len()) },
            1
        );
        assert_eq!(
            sequences[0],
            ZstdSequence {
                offset: 7,
                lit_length: 4,
                match_length: 4,
                rep: 2,
            }
        );
    }
}