zstd-rs/rust/src/zstd_presplit.rs

#![allow(non_snake_case)]

//! Compression block pre-splitting.
//!
//! This is the small distribution heuristic used by the compressor to decide
//! whether a full 128 KiB block should be split before match finding.  It uses
//! caller-provided workspace exactly like the C implementation so it does not
//! add an allocation on the compression hot path.

use crate::hist::HIST_add;
use crate::mem::MEM_read16;
use std::ffi::c_void;
use std::mem::size_of;
use std::os::raw::c_int;
use std::ptr;

const BLOCK_SIZE: usize = 128 << 10;
const BLOCK_SIZE_MIN: usize = 3_500;
const THRESHOLD_PENALTY_RATE: u64 = 16;
const THRESHOLD_BASE: u64 = THRESHOLD_PENALTY_RATE - 2;
const THRESHOLD_PENALTY: u64 = 3;

const HASH_LENGTH: usize = 2;
const HASH_LOG_MAX: u32 = 10;
const HASH_TABLE_SIZE: usize = 1 << HASH_LOG_MAX;
const KNUTH: u32 = 0x9e37_79b9;
const CHUNK_SIZE: usize = 8 << 10;
const SEGMENT_SIZE: usize = 512;
const WORKSPACE_SIZE: usize = 8_208;

#[repr(C)]
struct Fingerprint {
    events: [u32; HASH_TABLE_SIZE],
    nb_events: usize,
}

#[repr(C)]
struct FpStats {
    past_events: Fingerprint,
    new_events: Fingerprint,
}

#[inline]
unsafe fn fingerprint_events(fp: *mut Fingerprint) -> *mut u32 {
    unsafe { ptr::addr_of_mut!((*fp).events).cast() }
}

#[inline]
unsafe fn fingerprint_nb_events(fp: *const Fingerprint) -> usize {
    unsafe { ptr::addr_of!((*fp).nb_events).read() }
}

#[inline]
unsafe fn set_fingerprint_nb_events(fp: *mut Fingerprint, value: usize) {
    unsafe { ptr::addr_of_mut!((*fp).nb_events).write(value) }
}

#[inline]
unsafe fn hash2(source: *const u8, hash_log: u32) -> usize {
    debug_assert!((8..=HASH_LOG_MAX).contains(&hash_log));
    if hash_log == 8 {
        unsafe { *source as usize }
    } else {
        let value = unsafe { MEM_read16(source.cast()) as u32 };
        (value.wrapping_mul(KNUTH) >> (32 - hash_log)) as usize
    }
}

unsafe fn init_stats(stats: *mut FpStats) {
    unsafe { ptr::write_bytes(stats.cast::<u8>(), 0, size_of::<FpStats>()) }
}

unsafe fn add_events(
    fp: *mut Fingerprint,
    source: *const u8,
    source_size: usize,
    sampling_rate: usize,
    hash_log: u32,
) {
    debug_assert!(source_size >= HASH_LENGTH);
    let limit = source_size - HASH_LENGTH + 1;
    let events = unsafe { fingerprint_events(fp) };
    let mut index = 0;
    while index < limit {
        let event = unsafe { events.add(hash2(source.add(index), hash_log)) };
        unsafe { event.write(event.read().wrapping_add(1)) };
        index += sampling_rate;
    }
    let nb_events = unsafe { fingerprint_nb_events(fp) };
    unsafe { set_fingerprint_nb_events(fp, nb_events.wrapping_add(limit / sampling_rate)) };
}

unsafe fn record_fingerprint(
    fp: *mut Fingerprint,
    source: *const u8,
    source_size: usize,
    sampling_rate: usize,
    hash_log: u32,
) {
    unsafe {
        ptr::write_bytes(fingerprint_events(fp), 0, 1usize << hash_log);
        set_fingerprint_nb_events(fp, 0);
        add_events(fp, source, source_size, sampling_rate, hash_log);
    }
}

#[inline]
fn abs64(value: i64) -> u64 {
    if value < 0 {
        value.wrapping_neg() as u64
    } else {
        value as u64
    }
}

unsafe fn fp_distance(fp1: *const Fingerprint, fp2: *const Fingerprint, hash_log: u32) -> u64 {
    debug_assert!(hash_log <= HASH_LOG_MAX);
    let fp1_events = unsafe { fingerprint_events(fp1.cast_mut()) };
    let fp2_events = unsafe { fingerprint_events(fp2.cast_mut()) };
    let fp1_nb_events = unsafe { fingerprint_nb_events(fp1) } as i64;
    let fp2_nb_events = unsafe { fingerprint_nb_events(fp2) } as i64;
    let mut distance = 0u64;
    for index in 0..(1usize << hash_log) {
        let lhs = unsafe { fp1_events.add(index).read() as i64 } * fp2_nb_events;
        let rhs = unsafe { fp2_events.add(index).read() as i64 } * fp1_nb_events;
        distance = distance.wrapping_add(abs64(lhs.wrapping_sub(rhs)));
    }
    distance
}

unsafe fn compare_fingerprints(
    reference: *const Fingerprint,
    new_fingerprint: *const Fingerprint,
    penalty: u64,
    hash_log: u32,
) -> bool {
    let reference_events = unsafe { fingerprint_nb_events(reference) } as u64;
    let new_events = unsafe { fingerprint_nb_events(new_fingerprint) } as u64;
    debug_assert_ne!(reference_events, 0);
    debug_assert_ne!(new_events, 0);

    let p50 = reference_events * new_events;
    let deviation = unsafe { fp_distance(reference, new_fingerprint, hash_log) };
    let threshold = p50 * (THRESHOLD_BASE + penalty) / THRESHOLD_PENALTY_RATE;
    deviation >= threshold
}

unsafe fn merge_events(accumulator: *mut Fingerprint, new_fingerprint: *const Fingerprint) {
    let accumulator_events = unsafe { fingerprint_events(accumulator) };
    let new_events = unsafe { fingerprint_events(new_fingerprint.cast_mut()) };
    for index in 0..HASH_TABLE_SIZE {
        let value = unsafe {
            accumulator_events
                .add(index)
                .read()
                .wrapping_add(new_events.add(index).read())
        };
        unsafe { accumulator_events.add(index).write(value) };
    }
    let nb_events = unsafe { fingerprint_nb_events(accumulator) }
        .wrapping_add(unsafe { fingerprint_nb_events(new_fingerprint) });
    unsafe { set_fingerprint_nb_events(accumulator, nb_events) };
}

unsafe fn split_block_by_chunks(
    block_start: *const u8,
    block_size: usize,
    level: usize,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    const RECORD_PARAMS: [(usize, u32); 4] = [(43, 8), (11, 9), (5, 10), (1, 10)];

    debug_assert!(level <= 3);
    debug_assert_eq!(block_size, BLOCK_SIZE);
    debug_assert!(!workspace.is_null());
    debug_assert_eq!(workspace as usize % std::mem::align_of::<FpStats>(), 0);
    debug_assert!(workspace_size >= size_of::<FpStats>());
    debug_assert!(workspace_size >= WORKSPACE_SIZE);

    let (sampling_rate, hash_log) = RECORD_PARAMS[level];
    let stats = workspace.cast::<FpStats>();
    let past_events = unsafe { ptr::addr_of_mut!((*stats).past_events) };
    let new_events = unsafe { ptr::addr_of_mut!((*stats).new_events) };
    unsafe {
        init_stats(stats);
        record_fingerprint(
            past_events,
            block_start,
            CHUNK_SIZE,
            sampling_rate,
            hash_log,
        );
    }

    let mut penalty = THRESHOLD_PENALTY;
    let mut position = CHUNK_SIZE;
    while position <= block_size - CHUNK_SIZE {
        unsafe {
            record_fingerprint(
                new_events,
                block_start.add(position),
                CHUNK_SIZE,
                sampling_rate,
                hash_log,
            );
        }
        if unsafe { compare_fingerprints(past_events, new_events, penalty, hash_log) } {
            return position;
        }
        unsafe { merge_events(past_events, new_events) };
        penalty = penalty.saturating_sub(1);
        position += CHUNK_SIZE;
    }
    debug_assert_eq!(position, block_size);
    block_size
}

unsafe fn split_block_from_borders(
    block_start: *const u8,
    block_size: usize,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    debug_assert_eq!(block_size, BLOCK_SIZE);
    debug_assert!(!workspace.is_null());
    debug_assert_eq!(workspace as usize % std::mem::align_of::<FpStats>(), 0);
    debug_assert!(workspace_size >= size_of::<FpStats>());
    debug_assert!(workspace_size >= WORKSPACE_SIZE);

    let stats = workspace.cast::<FpStats>();
    let past_events = unsafe { ptr::addr_of_mut!((*stats).past_events) };
    let new_events = unsafe { ptr::addr_of_mut!((*stats).new_events) };
    let middle_events = unsafe { workspace.cast::<u8>().add(SEGMENT_SIZE * size_of::<u32>()) }
        .cast::<Fingerprint>();

    unsafe {
        init_stats(stats);
        HIST_add(
            fingerprint_events(past_events),
            block_start.cast(),
            SEGMENT_SIZE,
        );
        HIST_add(
            fingerprint_events(new_events),
            block_start.add(block_size - SEGMENT_SIZE).cast(),
            SEGMENT_SIZE,
        );
        set_fingerprint_nb_events(past_events, SEGMENT_SIZE);
        set_fingerprint_nb_events(new_events, SEGMENT_SIZE);
    }
    if !unsafe { compare_fingerprints(past_events, new_events, 0, 8) } {
        return block_size;
    }

    unsafe {
        HIST_add(
            fingerprint_events(middle_events),
            block_start.add(block_size / 2 - SEGMENT_SIZE / 2).cast(),
            SEGMENT_SIZE,
        );
        set_fingerprint_nb_events(middle_events, SEGMENT_SIZE);
    }
    let distance_from_begin = unsafe { fp_distance(past_events, middle_events, 8) };
    let distance_from_end = unsafe { fp_distance(new_events, middle_events, 8) };
    let minimum_distance = (SEGMENT_SIZE * SEGMENT_SIZE / 3) as u64;
    if abs64((distance_from_begin as i64).wrapping_sub(distance_from_end as i64)) < minimum_distance
    {
        64 << 10
    } else if distance_from_begin > distance_from_end {
        32 << 10
    } else {
        96 << 10
    }
}

/// Chooses a split point for a full-size compression block.
///
/// # Safety
/// `block_start` must reference exactly one readable 128 KiB block. `workspace`
/// must be aligned for `FpStats` and provide at least 8,208 bytes.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_splitBlock(
    block_start: *const c_void,
    block_size: usize,
    level: c_int,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    debug_assert!((0..=4).contains(&level));
    debug_assert!(block_size >= BLOCK_SIZE_MIN);
    if level == 0 {
        unsafe {
            split_block_from_borders(block_start.cast(), block_size, workspace, workspace_size)
        }
    } else {
        unsafe {
            split_block_by_chunks(
                block_start.cast(),
                block_size,
                (level - 1) as usize,
                workspace,
                workspace_size,
            )
        }
    }
}

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

    fn split(input: &[u8], level: c_int) -> usize {
        let mut workspace = vec![0usize; WORKSPACE_SIZE.div_ceil(size_of::<usize>())];
        unsafe {
            ZSTD_splitBlock(
                input.as_ptr().cast(),
                input.len(),
                level,
                workspace.as_mut_ptr().cast(),
                workspace.len() * size_of::<usize>(),
            )
        }
    }

    fn assert_all_levels(input: &[u8], expected: [usize; 5]) {
        for (level, expected) in expected.into_iter().enumerate() {
            assert_eq!(split(input, level as c_int), expected, "level {level}");
        }
    }

    #[test]
    fn workspace_layout_fits_the_public_contract() {
        assert!(size_of::<FpStats>() <= WORKSPACE_SIZE);
        assert_eq!(
            size_of::<Fingerprint>(),
            HASH_TABLE_SIZE * size_of::<u32>() + size_of::<usize>()
        );
    }

    #[test]
    fn split_points_match_c_reference_vectors() {
        let mut input = vec![0u8; BLOCK_SIZE];
        assert_all_levels(&input, [BLOCK_SIZE; 5]);

        input[..BLOCK_SIZE / 2].fill(0);
        input[BLOCK_SIZE / 2..].fill(1);
        assert_all_levels(&input, [64 << 10; 5]);

        input[..BLOCK_SIZE / 4].fill(0);
        input[BLOCK_SIZE / 4..].fill(1);
        assert_all_levels(&input, [32 << 10; 5]);
    }

    #[test]
    fn sampled_levels_match_c_reference_vectors() {
        let mut input = vec![0u8; BLOCK_SIZE];
        let mut seed = 0x1234_5678u32;
        for byte in &mut input {
            seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            *byte = (seed >> 24) as u8;
        }
        assert_all_levels(
            &input,
            [BLOCK_SIZE, CHUNK_SIZE, BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE],
        );

        for (index, byte) in input.iter_mut().enumerate() {
            *byte = ((index / CHUNK_SIZE) & 1) as u8;
        }
        assert_all_levels(
            &input,
            [64 << 10, CHUNK_SIZE, CHUNK_SIZE, CHUNK_SIZE, CHUNK_SIZE],
        );
    }
}