feat(rust): port compression block pre-splitting

Move ZSTD_splitBlock into Rust while retaining its caller-owned workspace contract. The implementation preserves the C fingerprint sampling heuristic and calls the migrated histogram primitive without adding a hot-path allocation. Reference-vector and randomized differential tests compare every split level with the original C translation unit. The C source remains as a declaration shim so existing source lists resolve the Rust ABI during the migration. Test Plan: - cargo fmt, cargo clippy, cargo clippy --benches, cargo clippy --tests - cargo test --all-targets and cargo build --release - cargo test/build --target i686-unknown-linux-gnu - 264-case original-C/Rust differential harness - C fuzzer, zstreamtest, and fuzzer32 smoke runs Refs: rust/README.md
2026-07-10 21:28:56 +02:00
parent ca70ca8061
commit 24d01b92fb
4 changed files with 379 additions and 237 deletions
@@ -17,3 +17,4 @@ pub mod threading;
 pub mod xxhash;
 pub mod zstd_common;
 pub mod zstd_ddict;
+pub mod zstd_presplit;
@@ -0,0 +1,376 @@
+#![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],
+        );
+    }
+}