zstd-rs/rust/src/zstd_compress_block_split.rs

//! Post-block-split partition search.
//!
//! The C compressor owns `ZSTD_CCtx`, block-state ownership, entropy-table
//! selection, and block emission. This module owns only the recursive search
//! over shallow `SeqStore_t` views. The estimator is supplied with projected
//! entropy state and parameter scalars rather than the private C context.

use crate::errors::ERR_isError;
use crate::zstd_compress_stats::{
    SeqStore_t, ZSTD_entropyCTablesMetadata_t, ZSTD_entropyCTables_t,
    ZSTD_rust_buildBlockEntropyStats, ZSTD_rust_deriveSeqStoreChunk, ZSTD_rust_estimateBlockSize,
};
use std::ffi::c_void;
use std::os::raw::c_int;

const MIN_SEQUENCES_BLOCK_SPLITTING: usize = 300;
const MAX_NB_BLOCK_SPLITS: usize = 196;
const SET_COMPRESSED: c_int = 2;

/// Recursively searches one sequence range for useful midpoint partitions.
///
/// The callback receives each derived view in the same order as the C
/// implementation: full range, left half, right half. It returns either an
/// estimated size or a zstd error code. Errors stop the current search path
/// and therefore cannot create a split.
#[allow(clippy::too_many_arguments)]
unsafe fn derive_block_splits_helper<F>(
    partitions: *mut u32,
    split_index: &mut usize,
    start_index: usize,
    end_index: usize,
    original_seq_store: *const SeqStore_t,
    full_seq_store_chunk: *mut SeqStore_t,
    first_half_seq_store: *mut SeqStore_t,
    second_half_seq_store: *mut SeqStore_t,
    estimate: &mut F,
) where
    F: FnMut(*mut SeqStore_t) -> usize,
{
    debug_assert!(end_index >= start_index);
    let sequence_count = end_index - start_index;
    let midpoint = start_index.wrapping_add(end_index) / 2;

    if sequence_count < MIN_SEQUENCES_BLOCK_SPLITTING || *split_index >= MAX_NB_BLOCK_SPLITS {
        return;
    }

    unsafe {
        ZSTD_rust_deriveSeqStoreChunk(
            full_seq_store_chunk,
            original_seq_store,
            start_index,
            end_index,
        );
        ZSTD_rust_deriveSeqStoreChunk(
            first_half_seq_store,
            original_seq_store,
            start_index,
            midpoint,
        );
        ZSTD_rust_deriveSeqStoreChunk(
            second_half_seq_store,
            original_seq_store,
            midpoint,
            end_index,
        );
    }

    let estimated_original_size = estimate(full_seq_store_chunk);
    let estimated_first_half_size = estimate(first_half_seq_store);
    let estimated_second_half_size = estimate(second_half_seq_store);
    if ERR_isError(estimated_original_size)
        || ERR_isError(estimated_first_half_size)
        || ERR_isError(estimated_second_half_size)
    {
        return;
    }

    if estimated_first_half_size.wrapping_add(estimated_second_half_size) < estimated_original_size
    {
        unsafe {
            derive_block_splits_helper(
                partitions,
                split_index,
                start_index,
                midpoint,
                original_seq_store,
                full_seq_store_chunk,
                first_half_seq_store,
                second_half_seq_store,
                estimate,
            );
        }

        /* Keep the configured split limit a true upper bound even if a left
         * subtree reaches it before its parent gets to midpoint insertion. */
        if *split_index >= MAX_NB_BLOCK_SPLITS {
            return;
        }
        unsafe {
            *partitions.add(*split_index) = midpoint as u32;
        }
        *split_index += 1;

        unsafe {
            derive_block_splits_helper(
                partitions,
                split_index,
                midpoint,
                end_index,
                original_seq_store,
                full_seq_store_chunk,
                first_half_seq_store,
                second_half_seq_store,
                estimate,
            );
        }
    }
}

/// Runs the partition search with a caller-provided estimator.
unsafe fn derive_block_splits_with_estimator<F>(
    partitions: *mut u32,
    nb_sequences: u32,
    original_seq_store: *const SeqStore_t,
    full_seq_store_chunk: *mut SeqStore_t,
    first_half_seq_store: *mut SeqStore_t,
    second_half_seq_store: *mut SeqStore_t,
    mut estimate: F,
) -> usize
where
    F: FnMut(*mut SeqStore_t) -> usize,
{
    if nb_sequences <= 4 {
        return 0;
    }

    let mut split_index = 0usize;
    unsafe {
        derive_block_splits_helper(
            partitions,
            &mut split_index,
            0,
            nb_sequences as usize,
            original_seq_store,
            full_seq_store_chunk,
            first_half_seq_store,
            second_half_seq_store,
            &mut estimate,
        );
        *partitions.add(split_index) = nb_sequences;
    }
    split_index
}

/// Rust leaf for C's `ZSTD_deriveBlockSplits()`.
///
/// The entropy and parameter arguments are the state projected from the C
/// context by the narrow shim. Each estimate rebuilds `next_entropy` and the
/// shared metadata exactly as the original C helper did before comparing the
/// three candidate sizes.
#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_deriveBlockSplits(
    partitions: *mut u32,
    nb_sequences: u32,
    original_seq_store: *const SeqStore_t,
    full_seq_store_chunk: *mut SeqStore_t,
    first_half_seq_store: *mut SeqStore_t,
    second_half_seq_store: *mut SeqStore_t,
    prev_entropy: *const ZSTD_entropyCTables_t,
    next_entropy: *mut ZSTD_entropyCTables_t,
    strategy: c_int,
    disable_literal_compression: c_int,
    entropy_metadata: *mut ZSTD_entropyCTablesMetadata_t,
    workspace: *mut c_void,
    workspace_size: usize,
) -> usize {
    let mut estimate = |seq_store: *mut SeqStore_t| -> usize {
        let entropy_stats = unsafe {
            ZSTD_rust_buildBlockEntropyStats(
                seq_store,
                prev_entropy,
                next_entropy,
                strategy,
                disable_literal_compression,
                entropy_metadata,
                workspace,
                workspace_size,
            )
        };
        if ERR_isError(entropy_stats) {
            return entropy_stats;
        }

        let seq_store = unsafe { &*seq_store };
        let lit_size = unsafe { seq_store.lit.offset_from(seq_store.litStart) } as usize;
        let nb_seq = unsafe { seq_store.sequences.offset_from(seq_store.sequencesStart) } as usize;
        let write_lit_entropy =
            unsafe { c_int::from((*entropy_metadata).hufMetadata.hType == SET_COMPRESSED) };
        unsafe {
            ZSTD_rust_estimateBlockSize(
                seq_store.litStart,
                lit_size,
                seq_store.ofCode,
                seq_store.llCode,
                seq_store.mlCode,
                nb_seq,
                next_entropy,
                entropy_metadata,
                workspace,
                workspace_size,
                write_lit_entropy,
                1,
            )
        }
    };

    unsafe {
        derive_block_splits_with_estimator(
            partitions,
            nb_sequences,
            original_seq_store,
            full_seq_store_chunk,
            first_half_seq_store,
            second_half_seq_store,
            &mut estimate,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::errors::{ZstdErrorCode, ERROR};
    use crate::zstd_compress_sequences::SeqDef;

    fn empty_seq_store() -> SeqStore_t {
        SeqStore_t {
            sequencesStart: std::ptr::null_mut(),
            sequences: std::ptr::null_mut(),
            litStart: std::ptr::null_mut(),
            lit: std::ptr::null_mut(),
            llCode: std::ptr::null_mut(),
            mlCode: std::ptr::null_mut(),
            ofCode: std::ptr::null_mut(),
            maxNbSeq: 0,
            maxNbLit: 0,
            longLengthType: 0,
            longLengthPos: 0,
        }
    }

    fn make_seq_store(
        sequences: &mut [SeqDef],
        codes: &mut [u8],
        literals: &mut [u8],
        long_length_type: c_int,
        long_length_pos: u32,
    ) -> SeqStore_t {
        for sequence in sequences.iter_mut() {
            *sequence = SeqDef {
                offBase: 4,
                litLength: 0,
                mlBase: 0,
            };
        }
        let sequences_start = sequences.as_mut_ptr();
        let literals_start = literals.as_mut_ptr();
        let codes_start = codes.as_mut_ptr();
        SeqStore_t {
            sequencesStart: sequences_start,
            sequences: unsafe { sequences_start.add(sequences.len()) },
            litStart: literals_start,
            lit: literals_start,
            llCode: codes_start,
            mlCode: codes_start,
            ofCode: codes_start,
            maxNbSeq: sequences.len(),
            maxNbLit: literals.len(),
            longLengthType: long_length_type,
            longLengthPos: long_length_pos,
        }
    }

    unsafe fn sequence_range(
        seq_store: *mut SeqStore_t,
        original_sequences: *mut SeqDef,
    ) -> (usize, usize) {
        let seq_store = unsafe { &*seq_store };
        let start = unsafe { seq_store.sequencesStart.offset_from(original_sequences) } as usize;
        let end = unsafe { seq_store.sequences.offset_from(original_sequences) } as usize;
        (start, end)
    }

    unsafe fn run_search<F>(
        sequences: &mut [SeqDef],
        long_length_type: c_int,
        long_length_pos: u32,
        partitions: &mut [u32],
        estimate: F,
    ) -> usize
    where
        F: FnMut(*mut SeqStore_t) -> usize,
    {
        let mut codes = vec![0u8; sequences.len()];
        let mut literals = vec![0u8; sequences.len() + 1];
        let original = make_seq_store(
            sequences,
            &mut codes,
            &mut literals,
            long_length_type,
            long_length_pos,
        );
        let mut full = empty_seq_store();
        let mut first = empty_seq_store();
        let mut second = empty_seq_store();
        derive_block_splits_with_estimator(
            partitions.as_mut_ptr(),
            sequences.len() as u32,
            &original,
            &mut full,
            &mut first,
            &mut second,
            estimate,
        )
    }

    #[test]
    fn recursive_search_is_left_first_with_midpoint_insertion() {
        let mut sequences = vec![SeqDef::default(); 1200];
        let original_sequences = sequences.as_mut_ptr();
        let mut partitions = vec![u32::MAX; 8];
        let mut calls = Vec::new();
        let mut call_index = 0usize;
        let splits = unsafe {
            run_search(&mut sequences, 0, 0, &mut partitions, |chunk| {
                let range = sequence_range(chunk, original_sequences);
                calls.push(range);
                let estimate = if call_index.is_multiple_of(3) {
                    1000
                } else {
                    0
                };
                call_index += 1;
                estimate
            })
        };

        assert_eq!(splits, 7);
        assert_eq!(&partitions[..splits], &[150, 300, 450, 600, 750, 900, 1050]);
        assert_eq!(
            &calls[..6],
            &[
                (0, 1200),
                (0, 600),
                (600, 1200),
                (0, 600),
                (0, 300),
                (300, 600)
            ]
        );
        assert_eq!(partitions[splits], 1200);
    }

    #[test]
    fn minimum_sequence_count_skips_estimation() {
        let mut sequences = vec![SeqDef::default(); 299];
        let mut partitions = [u32::MAX; 2];
        let mut calls = 0usize;
        let splits = unsafe {
            run_search(&mut sequences, 0, 0, &mut partitions, |_| {
                calls += 1;
                0
            })
        };

        assert_eq!(splits, 0);
        assert_eq!(calls, 0);
        assert_eq!(partitions[0], 299);
    }

    #[test]
    fn split_count_is_capped_at_configured_limit() {
        let mut sequences = vec![SeqDef::default(); 60_000];
        let mut partitions = vec![u32::MAX; MAX_NB_BLOCK_SPLITS + 1];
        let mut call_index = 0usize;
        let splits = unsafe {
            run_search(&mut sequences, 0, 0, &mut partitions, |_| {
                let estimate = if call_index.is_multiple_of(3) {
                    1000
                } else {
                    0
                };
                call_index += 1;
                estimate
            })
        };

        assert_eq!(splits, MAX_NB_BLOCK_SPLITS);
        assert!(partitions[..splits]
            .windows(2)
            .all(|window| window[0] < window[1]));
        assert_eq!(partitions[splits], 60_000);
    }

    #[test]
    fn estimator_errors_leave_the_current_range_unsplit() {
        let mut sequences = vec![SeqDef::default(); 600];
        let mut partitions = [u32::MAX; 2];
        let mut calls = 0usize;
        let splits = unsafe {
            run_search(&mut sequences, 0, 0, &mut partitions, |_| {
                calls += 1;
                ERROR(ZstdErrorCode::Generic)
            })
        };

        assert_eq!(splits, 0);
        assert_eq!(calls, 3);
        assert_eq!(partitions[0], 600);
    }

    #[test]
    fn an_estimation_error_in_any_candidate_prevents_the_split() {
        for error_index in 0..3 {
            let mut sequences = vec![SeqDef::default(); 600];
            let mut partitions = [u32::MAX; 2];
            let mut calls = 0;
            let splits = unsafe {
                run_search(&mut sequences, 0, 0, &mut partitions, |_| {
                    let current_call = calls;
                    calls += 1;
                    if current_call == error_index {
                        ERROR(ZstdErrorCode::Generic)
                    } else {
                        0
                    }
                })
            };

            assert_eq!(splits, 0, "error in candidate {error_index}");
            assert_eq!(calls, 3, "error in candidate {error_index}");
            assert_eq!(partitions[0], 600, "error in candidate {error_index}");
        }
    }

    #[test]
    fn equal_estimated_cost_does_not_split() {
        let mut sequences = vec![SeqDef::default(); 600];
        let mut partitions = [u32::MAX; 2];
        let mut calls = 0;
        let splits = unsafe {
            run_search(&mut sequences, 0, 0, &mut partitions, |_| {
                calls += 1;
                0
            })
        };

        assert_eq!(splits, 0);
        assert_eq!(calls, 3);
        assert_eq!(partitions[0], 600);
    }

    #[test]
    fn derived_views_keep_long_length_metadata_in_their_ranges() {
        let mut sequences = vec![SeqDef::default(); 600];
        let original_sequences = sequences.as_mut_ptr();
        let mut partitions = [u32::MAX; 2];
        let mut metadata = Vec::new();
        let splits = unsafe {
            run_search(&mut sequences, 1, 450, &mut partitions, |chunk| {
                let range = sequence_range(chunk, original_sequences);
                let view = &*chunk;
                metadata.push((range, view.longLengthType, view.longLengthPos));
                0
            })
        };

        assert_eq!(splits, 0);
        assert_eq!(
            metadata,
            vec![((0, 600), 1, 450), ((0, 300), 0, 450), ((300, 600), 1, 150)]
        );
    }
}