zstd-rs/rust/src/zstd_opt_tree.rs

#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(clippy::missing_safety_doc)]

//! Binary-tree maintenance for optimal block compression.
//!
//! `ZSTD_MatchState_t` remains a C-owned, configuration-sensitive structure.
//! The companion C translation unit projects only the matching-table and window
//! leaves used to populate its binary tree.  The parser itself remains separate
//! while this module moves the shared tree-maintenance path into Rust.

use crate::mem::{
    MEM_64bits, MEM_isLittleEndian, MEM_read16, MEM_read32, MEM_readLE32, MEM_readLE64, MEM_readST,
};
use std::cmp::{max, min};
use std::ffi::c_void;
use std::mem::size_of;
use std::os::raw::c_int;

const HASH_READ_SIZE: usize = 8;

/// Leaf view built by `zstd_opt.c` from its C-owned match state.
#[repr(C)]
pub struct ZSTD_RustOptTreeState {
    hash_table: *mut u32,
    chain_table: *mut u32,
    base: *const u8,
    dict_base: *const u8,
    dict_limit: u32,
    low_limit: u32,
    loaded_dict_end: u32,
    next_to_update: *mut u32,
    hash_log: u32,
    chain_log: u32,
    search_log: u32,
    window_log: u32,
    use_c_predict: c_int,
}

#[inline]
fn ptr_lt(left: *const u8, right: *const u8) -> bool {
    (left as usize) < (right as usize)
}

#[inline]
unsafe fn hash_shift32(value: u32, hbits: u32) -> usize {
    if hbits == 0 {
        0
    } else {
        (value >> (32 - hbits)) as usize
    }
}

#[inline]
unsafe fn hash_shift64(value: u64, hbits: u32) -> usize {
    if hbits == 0 {
        0
    } else {
        (value >> (64 - hbits)) as usize
    }
}

#[inline]
unsafe fn hash_ptr(ptr: *const u8, hbits: u32, mls: u32) -> usize {
    match mls {
        5 => {
            let value = unsafe { MEM_readLE64(ptr.cast::<c_void>()) };
            unsafe { hash_shift64(value.wrapping_shl(24).wrapping_mul(889_523_592_379), hbits) }
        }
        6 => {
            let value = unsafe { MEM_readLE64(ptr.cast::<c_void>()) };
            unsafe {
                hash_shift64(
                    value.wrapping_shl(16).wrapping_mul(227_718_039_650_203),
                    hbits,
                )
            }
        }
        7 => {
            let value = unsafe { MEM_readLE64(ptr.cast::<c_void>()) };
            unsafe {
                hash_shift64(
                    value.wrapping_shl(8).wrapping_mul(58_295_818_150_454_627),
                    hbits,
                )
            }
        }
        8 => {
            let value = unsafe { MEM_readLE64(ptr.cast::<c_void>()) };
            unsafe { hash_shift64(value.wrapping_mul(0xCF1B_BCDC_B7A5_6463), hbits) }
        }
        _ => {
            let value = unsafe { MEM_readLE32(ptr.cast::<c_void>()) };
            unsafe { hash_shift32(value.wrapping_mul(2_654_435_761), hbits) }
        }
    }
}

#[inline]
fn common_bytes(word: usize) -> usize {
    let zeros = if MEM_isLittleEndian() {
        word.trailing_zeros()
    } else {
        word.leading_zeros()
    };
    (zeros / 8) as usize
}

/// Equivalent to C's `ZSTD_count()` including its word-at-a-time tail.
unsafe fn count(mut input: *const u8, mut matched: *const u8, input_limit: *const u8) -> usize {
    let input_start = input;
    let word_size = size_of::<usize>();
    while unsafe { input_limit.offset_from(input) as usize } >= word_size {
        let difference =
            unsafe { MEM_readST(matched.cast::<c_void>()) ^ MEM_readST(input.cast::<c_void>()) };
        if difference != 0 {
            return unsafe { input.offset_from(input_start) as usize } + common_bytes(difference);
        }
        input = input.wrapping_add(word_size);
        matched = matched.wrapping_add(word_size);
    }
    if MEM_64bits()
        && unsafe { input_limit.offset_from(input) as usize } >= 4
        && unsafe { MEM_read32(matched.cast::<c_void>()) == MEM_read32(input.cast::<c_void>()) }
    {
        input = input.wrapping_add(4);
        matched = matched.wrapping_add(4);
    }
    if unsafe { input_limit.offset_from(input) as usize } >= 2
        && unsafe { MEM_read16(matched.cast::<c_void>()) == MEM_read16(input.cast::<c_void>()) }
    {
        input = input.wrapping_add(2);
        matched = matched.wrapping_add(2);
    }
    if ptr_lt(input, input_limit) && unsafe { *matched == *input } {
        input = input.wrapping_add(1);
    }
    unsafe { input.offset_from(input_start) as usize }
}

unsafe fn count_2segments(
    input: *const u8,
    matched: *const u8,
    input_end: *const u8,
    match_end: *const u8,
    input_start: *const u8,
) -> usize {
    let match_remaining = unsafe { match_end.offset_from(matched) as usize };
    let input_remaining = unsafe { input_end.offset_from(input) as usize };
    let first_end = input.wrapping_add(min(match_remaining, input_remaining));
    let first_count = unsafe { count(input, matched, first_end) };
    if matched.wrapping_add(first_count) != match_end {
        return first_count;
    }
    first_count + unsafe { count(input.wrapping_add(first_count), input_start, input_end) }
}

#[inline]
fn lowest_match_index(state: &ZSTD_RustOptTreeState, current: u32) -> u32 {
    let max_distance = 1u32.wrapping_shl(state.window_log);
    let within_window = if current.wrapping_sub(state.low_limit) > max_distance {
        current.wrapping_sub(max_distance)
    } else {
        state.low_limit
    };
    if state.loaded_dict_end != 0 {
        state.low_limit
    } else {
        within_window
    }
}

unsafe fn insert_bt1(
    state: &ZSTD_RustOptTreeState,
    ip: *const u8,
    iend: *const u8,
    target: u32,
    mls: u32,
    ext_dict: bool,
) -> u32 {
    let hash = unsafe { hash_ptr(ip, state.hash_log, mls) };
    let bt_log = state.chain_log - 1;
    let bt_mask = (1u32 << bt_log) - 1;
    let mut match_index = unsafe { *state.hash_table.add(hash) };
    let mut common_smaller = 0usize;
    let mut common_larger = 0usize;
    let dict_end = state.dict_base.wrapping_add(state.dict_limit as usize);
    let prefix_start = state.base.wrapping_add(state.dict_limit as usize);
    let current = unsafe { ip.offset_from(state.base) as u32 };
    let bt_low = current.saturating_sub(bt_mask);
    let mut smaller = unsafe { state.chain_table.add(2 * (current & bt_mask) as usize) };
    let mut larger = unsafe { smaller.add(1) };
    let mut dummy = 0u32;
    let window_low = lowest_match_index(state, target);
    let mut match_end_index = current.wrapping_add(HASH_READ_SIZE as u32 + 1);
    let mut best_length = HASH_READ_SIZE;
    let mut comparisons = 1u32 << state.search_log;

    unsafe { *state.hash_table.add(hash) = current };

    let mut predicted_small = 0u32;
    let mut predicted_large = 0u32;
    if state.use_c_predict != 0 {
        let prior = unsafe {
            state
                .chain_table
                .add(2 * (current.wrapping_sub(1) & bt_mask) as usize)
        };
        predicted_small = unsafe { *prior }.wrapping_add(u32::from(unsafe { *prior } > 0));
        predicted_large =
            unsafe { *prior.add(1) }.wrapping_add(u32::from(unsafe { *prior.add(1) } > 0));
    }

    while comparisons != 0 && match_index >= window_low {
        comparisons -= 1;
        let next = unsafe { state.chain_table.add(2 * (match_index & bt_mask) as usize) };
        let mut match_length = min(common_smaller, common_larger);

        if state.use_c_predict != 0 && match_index == predicted_small {
            let prediction = unsafe {
                state
                    .chain_table
                    .add(2 * (match_index.wrapping_sub(1) & bt_mask) as usize)
            };
            unsafe { *smaller = match_index };
            if match_index <= bt_low {
                smaller = &mut dummy;
                break;
            }
            smaller = unsafe { next.add(1) };
            match_index = unsafe { *next.add(1) };
            predicted_small = unsafe { *prediction.add(1) }
                .wrapping_add(u32::from(unsafe { *prediction.add(1) } > 0));
            continue;
        }
        if state.use_c_predict != 0 && match_index == predicted_large {
            let prediction = unsafe {
                state
                    .chain_table
                    .add(2 * (match_index.wrapping_sub(1) & bt_mask) as usize)
            };
            unsafe { *larger = match_index };
            if match_index <= bt_low {
                larger = &mut dummy;
                break;
            }
            larger = next;
            match_index = unsafe { *next };
            predicted_large =
                unsafe { *prediction }.wrapping_add(u32::from(unsafe { *prediction } > 0));
            continue;
        }

        let mut matched;
        if !ext_dict || match_index.wrapping_add(match_length as u32) >= state.dict_limit {
            matched = state.base.wrapping_add(match_index as usize);
            match_length += unsafe {
                count(
                    ip.wrapping_add(match_length),
                    matched.wrapping_add(match_length),
                    iend,
                )
            };
        } else {
            matched = state.dict_base.wrapping_add(match_index as usize);
            match_length += unsafe {
                count_2segments(
                    ip.wrapping_add(match_length),
                    matched.wrapping_add(match_length),
                    iend,
                    dict_end,
                    prefix_start,
                )
            };
            if match_index.wrapping_add(match_length as u32) >= state.dict_limit {
                matched = state.base.wrapping_add(match_index as usize);
            }
        }

        if match_length > best_length {
            best_length = match_length;
            if match_length > match_end_index.wrapping_sub(match_index) as usize {
                match_end_index = match_index.wrapping_add(match_length as u32);
            }
        }

        if ip.wrapping_add(match_length) == iend {
            break;
        }

        if unsafe { *matched.add(match_length) < *ip.add(match_length) } {
            unsafe { *smaller = match_index };
            common_smaller = match_length;
            if match_index <= bt_low {
                smaller = &mut dummy;
                break;
            }
            smaller = unsafe { next.add(1) };
            match_index = unsafe { *next.add(1) };
        } else {
            unsafe { *larger = match_index };
            common_larger = match_length;
            if match_index <= bt_low {
                larger = &mut dummy;
                break;
            }
            larger = next;
            match_index = unsafe { *next };
        }
    }

    unsafe {
        *smaller = 0;
        *larger = 0;
    }
    let positions = if best_length > 384 {
        min(192, (best_length - 384) as u32)
    } else {
        0
    };
    max(
        positions,
        match_end_index.wrapping_sub(current.wrapping_add(HASH_READ_SIZE as u32)),
    )
}

#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_opt_updateTreeInternal(
    state: *mut ZSTD_RustOptTreeState,
    ip: *const c_void,
    iend: *const c_void,
    mls: u32,
    ext_dict: c_int,
) {
    let state = unsafe { &*state };
    let input = ip.cast::<u8>();
    let input_end = iend.cast::<u8>();
    let target = unsafe { input.offset_from(state.base) as u32 };
    let mut index = unsafe { *state.next_to_update };
    while index < target {
        let forward = unsafe {
            insert_bt1(
                state,
                state.base.wrapping_add(index as usize),
                input_end,
                target,
                mls,
                ext_dict != 0,
            )
        };
        debug_assert!(index < index.wrapping_add(forward));
        index = index.wrapping_add(forward);
    }
    unsafe { *state.next_to_update = target };
}

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

    #[test]
    fn tree_update_records_prefix_positions() {
        let input = b"abcabcabcabcabcabcabcabc";
        let mut hashes = [0u32; 1 << 8];
        let mut chain = [0u32; 1 << 9];
        let mut next = 2u32;
        let mut state = ZSTD_RustOptTreeState {
            hash_table: hashes.as_mut_ptr(),
            chain_table: chain.as_mut_ptr(),
            base: input.as_ptr(),
            dict_base: input.as_ptr(),
            dict_limit: 0,
            low_limit: 1,
            loaded_dict_end: 0,
            next_to_update: &mut next,
            hash_log: 8,
            chain_log: 9,
            search_log: 4,
            window_log: 10,
            use_c_predict: 0,
        };
        unsafe {
            ZSTD_rust_opt_updateTreeInternal(
                &mut state,
                input.as_ptr().add(input.len() - HASH_READ_SIZE).cast(),
                input.as_ptr().add(input.len()).cast(),
                4,
                0,
            );
        }
        assert_eq!(next, (input.len() - HASH_READ_SIZE) as u32);
        assert!(hashes.iter().any(|&index| index != 0));
        assert!(chain.iter().all(|&index| index < input.len() as u32));
    }
}