zstd-rs/rust/src/hist.rs

#![allow(non_snake_case)]

use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
use std::ffi::c_void;
use std::os::raw::c_uint;

pub const HIST_WKSP_SIZE_U32: usize = 1024;
pub const HIST_WKSP_SIZE: usize = HIST_WKSP_SIZE_U32 * std::mem::size_of::<c_uint>();

#[no_mangle]
pub extern "C" fn HIST_isError(code: usize) -> c_uint {
    ERR_isError(code) as c_uint
}

#[no_mangle]
pub unsafe extern "C" fn HIST_add(count: *mut c_uint, src: *const c_void, src_size: usize) {
    let mut ip = src.cast::<u8>();
    let iend = ip.wrapping_add(src_size);
    while ip < iend {
        let counter = count.add(*ip as usize);
        *counter = (*counter).wrapping_add(1);
        ip = ip.add(1);
    }
}

#[no_mangle]
pub unsafe extern "C" fn HIST_count_simple(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    src: *const c_void,
    src_size: usize,
) -> c_uint {
    let mut max_symbol_value = *max_symbol_value_ptr;
    std::ptr::write_bytes(count, 0, max_symbol_value as usize + 1);

    if src_size == 0 {
        *max_symbol_value_ptr = 0;
        return 0;
    }

    let mut ip = src.cast::<u8>();
    let iend = ip.add(src_size);
    while ip < iend {
        let counter = count.add(*ip as usize);
        *counter = (*counter).wrapping_add(1);
        ip = ip.add(1);
    }

    while *count.add(max_symbol_value as usize) == 0 {
        max_symbol_value -= 1;
    }
    *max_symbol_value_ptr = max_symbol_value;

    let mut largest_count = 0;
    for symbol in 0..=max_symbol_value as usize {
        largest_count = largest_count.max(*count.add(symbol));
    }
    largest_count
}

unsafe fn HIST_count_parallel_wksp(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    source: *const c_void,
    source_size: usize,
    check_max_symbol_value: bool,
    workspace: *mut c_uint,
) -> usize {
    let requested_max = *max_symbol_value_ptr as usize;
    let count_size = requested_max + 1;

    if source_size == 0 {
        std::ptr::write_bytes(count, 0, count_size);
        *max_symbol_value_ptr = 0;
        return 0;
    }

    std::ptr::write_bytes(workspace, 0, HIST_WKSP_SIZE_U32);
    let counting1 = workspace;
    let counting2 = workspace.add(256);
    let counting3 = workspace.add(512);
    let counting4 = workspace.add(768);

    let source = source.cast::<u8>();
    let mut offset = 0;

    // Keep at least four bytes for the scalar tail, matching the C loop's
    // pipelined 16-byte stripes without its intentional short-input over-read.
    while source_size - offset >= 20 {
        for word in 0..4 {
            let bytes = source.add(offset + word * 4);
            let counters = [counting1, counting2, counting3, counting4];
            for (byte, counting) in counters.into_iter().enumerate() {
                let counter = counting.add(*bytes.add(byte) as usize);
                *counter = (*counter).wrapping_add(1);
            }
        }
        offset += 16;
    }

    while offset < source_size {
        let counter = counting1.add(*source.add(offset) as usize);
        *counter = (*counter).wrapping_add(1);
        offset += 1;
    }

    let mut largest_count = 0;
    for symbol in 0..256 {
        let total = (*counting1.add(symbol))
            .wrapping_add(*counting2.add(symbol))
            .wrapping_add(*counting3.add(symbol))
            .wrapping_add(*counting4.add(symbol));
        *counting1.add(symbol) = total;
        largest_count = largest_count.max(total);
    }

    let mut max_symbol_value = 255usize;
    while *counting1.add(max_symbol_value) == 0 {
        max_symbol_value -= 1;
    }
    if check_max_symbol_value && max_symbol_value > requested_max {
        return ERROR(ZstdErrorCode::MaxSymbolValueTooSmall);
    }

    *max_symbol_value_ptr = max_symbol_value as c_uint;
    std::ptr::copy(counting1, count, count_size);
    largest_count as usize
}

#[no_mangle]
pub unsafe extern "C" fn HIST_countFast_wksp(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    source: *const c_void,
    source_size: usize,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    if source_size < 1500 {
        return HIST_count_simple(count, max_symbol_value_ptr, source, source_size) as usize;
    }
    if workspace as usize & 3 != 0 {
        return ERROR(ZstdErrorCode::Generic);
    }
    if workspace_size < HIST_WKSP_SIZE {
        return ERROR(ZstdErrorCode::WorkSpaceTooSmall);
    }
    HIST_count_parallel_wksp(
        count,
        max_symbol_value_ptr,
        source,
        source_size,
        false,
        workspace.cast(),
    )
}

#[no_mangle]
pub unsafe extern "C" fn HIST_count_wksp(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    source: *const c_void,
    source_size: usize,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    if workspace as usize & 3 != 0 {
        return ERROR(ZstdErrorCode::Generic);
    }
    if workspace_size < HIST_WKSP_SIZE {
        return ERROR(ZstdErrorCode::WorkSpaceTooSmall);
    }
    if *max_symbol_value_ptr < 255 {
        return HIST_count_parallel_wksp(
            count,
            max_symbol_value_ptr,
            source,
            source_size,
            true,
            workspace.cast(),
        );
    }

    *max_symbol_value_ptr = 255;
    HIST_countFast_wksp(
        count,
        max_symbol_value_ptr,
        source,
        source_size,
        workspace,
        workspace_size,
    )
}

#[no_mangle]
pub unsafe extern "C" fn HIST_countFast(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    source: *const c_void,
    source_size: usize,
) -> usize {
    let mut workspace = [0u32; HIST_WKSP_SIZE_U32];
    HIST_countFast_wksp(
        count,
        max_symbol_value_ptr,
        source,
        source_size,
        workspace.as_mut_ptr().cast(),
        HIST_WKSP_SIZE,
    )
}

#[no_mangle]
pub unsafe extern "C" fn HIST_count(
    count: *mut c_uint,
    max_symbol_value_ptr: *mut c_uint,
    source: *const c_void,
    source_size: usize,
) -> usize {
    let mut workspace = [0u32; HIST_WKSP_SIZE_U32];
    HIST_count_wksp(
        count,
        max_symbol_value_ptr,
        source,
        source_size,
        workspace.as_mut_ptr().cast(),
        HIST_WKSP_SIZE,
    )
}

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

    #[test]
    fn add_accumulates_counts() {
        let source = [0u8, 1, 1, 3, 3, 3];
        let mut counts = [0u32; 256];
        counts[1] = 4;

        unsafe { HIST_add(counts.as_mut_ptr(), source.as_ptr().cast(), source.len()) };

        assert_eq!(counts[0], 1);
        assert_eq!(counts[1], 6);
        assert_eq!(counts[2], 0);
        assert_eq!(counts[3], 3);
    }

    #[test]
    fn simple_count_reports_largest_count_and_symbol() {
        let source = [0u8, 3, 3, 2, 3, 2];
        let mut counts = [99u32; 256];
        let mut max_symbol = 255;

        let largest = unsafe {
            HIST_count_simple(
                counts.as_mut_ptr(),
                &mut max_symbol,
                source.as_ptr().cast(),
                source.len(),
            )
        };

        assert_eq!(largest, 3);
        assert_eq!(max_symbol, 3);
        assert_eq!(&counts[..4], &[1, 0, 2, 3]);
        assert!(counts[4..].iter().all(|&value| value == 0));
    }

    #[test]
    fn empty_input_clears_requested_alphabet() {
        let mut counts = [7u32; 8];
        let mut max_symbol = 7;
        let mut workspace = [0u32; HIST_WKSP_SIZE_U32];

        let largest = unsafe {
            HIST_count_wksp(
                counts.as_mut_ptr(),
                &mut max_symbol,
                std::ptr::null(),
                0,
                workspace.as_mut_ptr().cast(),
                HIST_WKSP_SIZE,
            )
        };

        assert_eq!(largest, 0);
        assert_eq!(max_symbol, 0);
        assert_eq!(counts, [0; 8]);
    }

    #[test]
    fn checked_count_rejects_symbols_outside_the_alphabet() {
        let source = [1u8, 2, 7];
        let mut counts = [55u32; 8];
        let mut max_symbol = 3;
        let mut workspace = [0u32; HIST_WKSP_SIZE_U32];

        let result = unsafe {
            HIST_count_wksp(
                counts.as_mut_ptr(),
                &mut max_symbol,
                source.as_ptr().cast(),
                source.len(),
                workspace.as_mut_ptr().cast(),
                HIST_WKSP_SIZE,
            )
        };

        assert_eq!(result, ERROR(ZstdErrorCode::MaxSymbolValueTooSmall));
        assert_eq!(max_symbol, 3);
        assert_eq!(counts, [55; 8]);
    }

    #[test]
    fn workspace_contract_matches_the_c_api() {
        let source = vec![42u8; 1500];
        let mut counts = [0u32; 256];
        let mut max_symbol = 255;
        let mut workspace = vec![0u8; HIST_WKSP_SIZE + 1];

        let misaligned = unsafe {
            HIST_countFast_wksp(
                counts.as_mut_ptr(),
                &mut max_symbol,
                source.as_ptr().cast(),
                source.len(),
                workspace.as_mut_ptr().add(1).cast(),
                HIST_WKSP_SIZE,
            )
        };
        assert_eq!(misaligned, ERROR(ZstdErrorCode::Generic));

        let mut aligned_workspace = [0u32; HIST_WKSP_SIZE_U32];
        let too_small = unsafe {
            HIST_countFast_wksp(
                counts.as_mut_ptr(),
                &mut max_symbol,
                source.as_ptr().cast(),
                source.len(),
                aligned_workspace.as_mut_ptr().cast(),
                HIST_WKSP_SIZE - 1,
            )
        };
        assert_eq!(too_small, ERROR(ZstdErrorCode::WorkSpaceTooSmall));
    }

    #[test]
    fn parallel_count_handles_stripes_and_tail() {
        let source: Vec<u8> = (0..1537).map(|value| (value % 17) as u8).collect();
        let mut counts = [0u32; 256];
        let mut max_symbol = 255;
        let mut workspace = [0u32; HIST_WKSP_SIZE_U32];

        let largest = unsafe {
            HIST_countFast_wksp(
                counts.as_mut_ptr(),
                &mut max_symbol,
                source.as_ptr().cast(),
                source.len(),
                workspace.as_mut_ptr().cast(),
                HIST_WKSP_SIZE,
            )
        };

        assert_eq!(largest, 91);
        assert_eq!(max_symbol, 16);
        assert_eq!(counts[..7], [91; 7]);
        assert_eq!(counts[7..17], [90; 10]);
        assert!(counts[17..].iter().all(|&value| value == 0));
    }
}