zstd-rs/rust/src/benchfn.rs

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

//! Benchmark loop for arbitrary functions over a set of blocks.
//!
//! Port of `programs/benchfn.c`.  `BMK_benchFunction` measures one batch of
//! runs; `BMK_benchTimedFn` repeats batches, growing the loop count until a
//! run lasts long enough to be reported reliably against `run_ms`, within a
//! `total_ms` budget tracked by `BMK_timedFnState_t`.
//!
//! ABI notes: `BMK_runOutcome_t` and `BMK_runTime_t` are returned by value
//! across the C boundary and `BMK_benchParams_t` is passed by value, so all
//! three are `repr(C)` mirrors of the `benchfn.h` layout, pinned by asserts
//! here and in the C shim.  `BMK_timedFnState_t` is opaque to C, but
//! `BMK_initStatic_timedFnState` guarantees it fits the 64-byte
//! `BMK_timedFnState_shell`, and `BMK_createTimedFnState` uses `malloc` so
//! creation and destruction stay interchangeable with C callers.

use std::os::raw::{c_int, c_uint, c_void};
use std::ptr;

use crate::timefn::{PTime, UTIL_clockSpanNano, UTIL_getTime, UTIL_time_t};

const TIMELOOP_NANOSEC: PTime = 1_000_000_000;

/// Valid benchmark result (`BMK_runTime_t` in benchfn.h).
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct BMK_runTime_t {
    /// Time per iteration, over all blocks.
    pub nanoSecPerRun: f64,
    /// Sum of the benchmarked function's return values, first loop only.
    pub sumOfReturn: usize,
}

/// Outcome variant of a benchmark run (`BMK_runOutcome_t` in benchfn.h):
/// either a valid `BMK_runTime_t` or an error result. C callers treat it as
/// opaque and use the accessor functions below.
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct BMK_runOutcome_t {
    pub internal_never_ever_use_directly: BMK_runTime_t,
    pub error_result_never_ever_use_directly: usize,
    pub error_tag_never_ever_use_directly: c_int,
}

// These mirror the static asserts in the programs/benchfn.c shim: the structs
// cross the ABI by value, so field offsets must match the C header exactly.
const _: () =
    assert!(std::mem::offset_of!(BMK_runTime_t, sumOfReturn) == std::mem::size_of::<f64>());
const _: () = assert!(
    std::mem::offset_of!(BMK_runOutcome_t, error_result_never_ever_use_directly)
        == std::mem::size_of::<BMK_runTime_t>()
);
const _: () = assert!(
    std::mem::offset_of!(BMK_runOutcome_t, error_tag_never_ever_use_directly)
        == std::mem::size_of::<BMK_runTime_t>() + std::mem::size_of::<usize>()
);

/// `size_t (*BMK_benchFn_t)(const void*, size_t, void*, size_t, void*)`
pub type BMK_benchFn_t = Option<
    unsafe extern "C" fn(
        src: *const c_void,
        srcSize: usize,
        dst: *mut c_void,
        dstCapacity: usize,
        customPayload: *mut c_void,
    ) -> usize,
>;
/// `size_t (*BMK_initFn_t)(void*)`
pub type BMK_initFn_t = Option<unsafe extern "C" fn(initPayload: *mut c_void) -> usize>;
/// `unsigned (*BMK_errorFn_t)(size_t)`
pub type BMK_errorFn_t = Option<unsafe extern "C" fn(result: usize) -> c_uint>;

/// Parameters of `BMK_benchFunction`, passed by value (`BMK_benchParams_t`).
#[repr(C)]
#[derive(Clone, Copy)]
pub struct BMK_benchParams_t {
    pub benchFn: BMK_benchFn_t,
    pub benchPayload: *mut c_void,
    pub initFn: BMK_initFn_t,
    pub initPayload: *mut c_void,
    pub errorFn: BMK_errorFn_t,
    pub blockCount: usize,
    pub srcBuffers: *const *const c_void,
    pub srcSizes: *const usize,
    pub dstBuffers: *const *mut c_void,
    pub dstCapacities: *const usize,
    pub blockResults: *mut usize,
}

/// Aborts, like benchfn.c's `CONTROL`, when an accessor is used on the wrong
/// outcome variant.
fn control(condition: bool) {
    if !condition {
        std::process::abort();
    }
}

fn error_outcome(errorResult: usize) -> BMK_runOutcome_t {
    BMK_runOutcome_t {
        internal_never_ever_use_directly: BMK_runTime_t {
            nanoSecPerRun: 0.0,
            sumOfReturn: 0,
        },
        error_result_never_ever_use_directly: errorResult,
        error_tag_never_ever_use_directly: 1,
    }
}

fn valid_outcome(runTime: BMK_runTime_t) -> BMK_runOutcome_t {
    BMK_runOutcome_t {
        internal_never_ever_use_directly: runTime,
        error_result_never_ever_use_directly: 0,
        error_tag_never_ever_use_directly: 0,
    }
}

/// Tells if the outcome carries a valid measurement.
#[no_mangle]
pub extern "C" fn BMK_isSuccessful_runOutcome(outcome: BMK_runOutcome_t) -> c_int {
    c_int::from(outcome.error_tag_never_ever_use_directly == 0)
}

/// Extracts the measurement; aborts if the outcome is an error, so validity
/// must be checked first with `BMK_isSuccessful_runOutcome`.
#[no_mangle]
pub extern "C" fn BMK_extract_runTime(outcome: BMK_runOutcome_t) -> BMK_runTime_t {
    control(outcome.error_tag_never_ever_use_directly == 0);
    outcome.internal_never_ever_use_directly
}

/// Extracts the faulty `benchFn` return value; aborts if the outcome is
/// valid, so failure must be checked first.
#[no_mangle]
pub extern "C" fn BMK_extract_errorResult(outcome: BMK_runOutcome_t) -> usize {
    control(outcome.error_tag_never_ever_use_directly != 0);
    outcome.error_result_never_ever_use_directly
}

/// Runs `initFn` once, then `benchFn` `nbLoops` times over every block, and
/// reports the mean time per loop.  On the first loop, per-block results are
/// stored into `blockResults` (when provided) and checked with `errorFn`
/// (when provided); the first failing block aborts the measurement and
/// produces an error outcome carrying the faulty return value.
#[no_mangle]
pub unsafe extern "C" fn BMK_benchFunction(
    p: BMK_benchParams_t,
    mut nbLoops: c_uint,
) -> BMK_runOutcome_t {
    // Minimum nbLoops is 1.
    nbLoops += c_uint::from(nbLoops == 0);

    // Warm up and erase the result buffers.
    for blockNb in 0..p.blockCount {
        unsafe {
            let dst = *p.dstBuffers.add(blockNb);
            ptr::write_bytes(dst.cast::<u8>(), 0xE5, *p.dstCapacities.add(blockNb));
        }
    }

    let benchFn = p.benchFn.expect("benchFn is mandatory");
    let mut dstSize = 0usize;
    let clockStart = UTIL_getTime();
    if let Some(initFn) = p.initFn {
        unsafe { initFn(p.initPayload) };
    }
    for loopNb in 0..nbLoops {
        for blockNb in 0..p.blockCount {
            let res = unsafe {
                benchFn(
                    *p.srcBuffers.add(blockNb),
                    *p.srcSizes.add(blockNb),
                    *p.dstBuffers.add(blockNb),
                    *p.dstCapacities.add(blockNb),
                    p.benchPayload,
                )
            };
            if loopNb == 0 {
                if !p.blockResults.is_null() {
                    unsafe { *p.blockResults.add(blockNb) = res };
                }
                if let Some(errorFn) = p.errorFn {
                    if unsafe { errorFn(res) } != 0 {
                        return error_outcome(res);
                    }
                }
                dstSize = dstSize.wrapping_add(res);
            }
        }
    }

    let totalTime = UTIL_clockSpanNano(clockStart);
    valid_outcome(BMK_runTime_t {
        nanoSecPerRun: totalTime as f64 / f64::from(nbLoops),
        sumOfReturn: dstSize,
    })
}

/// Benchmark session state (`struct BMK_timedFnState_s`), opaque to C.
#[repr(C)]
pub struct BMK_timedFnState_t {
    timeSpent_ns: PTime,
    timeBudget_ns: PTime,
    runBudget_ns: PTime,
    fastestRun: BMK_runTime_t,
    nbLoops: c_uint,
    coolTime: UTIL_time_t,
}

/// `BMK_TIMEDFNSTATE_SIZE` in benchfn.h: capacity of the caller-provided
/// `BMK_timedFnState_shell`, which the state must always fit.
const BMK_TIMEDFNSTATE_SIZE: usize = 64;
const _: () = assert!(std::mem::size_of::<BMK_timedFnState_t>() <= BMK_TIMEDFNSTATE_SIZE);
// The shell aligns via a `long long` member; the state must not need more.
const _: () = assert!(std::mem::align_of::<BMK_timedFnState_t>() <= std::mem::align_of::<u64>());

/// Allocates and initializes a benchmark session lasting a minimum of
/// `total_ms`, paced at intervals of approximately `run_ms`.  Uses `malloc`
/// so ownership stays interchangeable with the original C implementation.
#[no_mangle]
pub extern "C" fn BMK_createTimedFnState(
    total_ms: c_uint,
    run_ms: c_uint,
) -> *mut BMK_timedFnState_t {
    let state = unsafe { libc::malloc(std::mem::size_of::<BMK_timedFnState_t>()) }
        .cast::<BMK_timedFnState_t>();
    if state.is_null() {
        return ptr::null_mut();
    }
    unsafe { BMK_resetTimedFnState(state, total_ms, run_ms) };
    state
}

/// Releases a state obtained from `BMK_createTimedFnState`.
#[no_mangle]
pub unsafe extern "C" fn BMK_freeTimedFnState(state: *mut BMK_timedFnState_t) {
    unsafe { libc::free(state.cast()) };
}

/// Places the session state into a caller-provided buffer, typically a
/// `BMK_timedFnState_shell`.  Returns NULL when the buffer is missing, too
/// small, or misaligned.
#[no_mangle]
pub unsafe extern "C" fn BMK_initStatic_timedFnState(
    buffer: *mut c_void,
    size: usize,
    total_ms: c_uint,
    run_ms: c_uint,
) -> *mut BMK_timedFnState_t {
    if buffer.is_null() {
        return ptr::null_mut();
    }
    if size < std::mem::size_of::<BMK_timedFnState_t>() {
        return ptr::null_mut();
    }
    if !(buffer as usize).is_multiple_of(std::mem::align_of::<BMK_timedFnState_t>()) {
        return ptr::null_mut();
    }
    let state = buffer.cast::<BMK_timedFnState_t>();
    unsafe { BMK_resetTimedFnState(state, total_ms, run_ms) };
    state
}

/// Re-arms a session for a new benchmark of `total_ms`, paced at `run_ms`.
#[no_mangle]
pub unsafe extern "C" fn BMK_resetTimedFnState(
    timedFnState: *mut BMK_timedFnState_t,
    total_ms: c_uint,
    run_ms: c_uint,
) {
    let total_ms = if total_ms == 0 { 1 } else { total_ms };
    let mut run_ms = if run_ms == 0 { 1 } else { run_ms };
    if run_ms > total_ms {
        run_ms = total_ms;
    }
    let state = BMK_timedFnState_t {
        timeSpent_ns: 0,
        timeBudget_ns: PTime::from(total_ms) * TIMELOOP_NANOSEC / 1000,
        runBudget_ns: PTime::from(run_ms) * TIMELOOP_NANOSEC / 1000,
        fastestRun: BMK_runTime_t {
            // Must be larger than any potential measurement.
            nanoSecPerRun: TIMELOOP_NANOSEC as f64 * 2_000_000_000.0,
            sumOfReturn: usize::MAX,
        },
        nbLoops: 1,
        coolTime: UTIL_getTime(),
    };
    unsafe { timedFnState.write(state) };
}

/// Tells if the total time budget of the session is spent.  Also reports 1
/// after `BMK_benchTimedFn` returned an error.
#[no_mangle]
pub unsafe extern "C" fn BMK_isCompleted_TimedFn(timedFnState: *const BMK_timedFnState_t) -> c_int {
    let state = unsafe { &*timedFnState };
    c_int::from(state.timeSpent_ns >= state.timeBudget_ns)
}

/// Runs one measurement supposed to last about `run_ms`, automatically
/// scaling `nbLoops`.  Runs shorter than half the run budget are re-tried
/// with a larger workload instead of being reported, limiting rounding-error
/// risks; the best (fastest) qualifying run is returned.
#[no_mangle]
pub unsafe extern "C" fn BMK_benchTimedFn(
    cont: *mut BMK_timedFnState_t,
    p: BMK_benchParams_t,
) -> BMK_runOutcome_t {
    let cont = unsafe { &mut *cont };
    let runBudget_ns = cont.runBudget_ns;
    let runTimeMin_ns = runBudget_ns / 2;
    let mut bestRunTime = cont.fastestRun;

    loop {
        let runResult = unsafe { BMK_benchFunction(p, cont.nbLoops) };

        if BMK_isSuccessful_runOutcome(runResult) == 0 {
            // Error: move out.
            return runResult;
        }

        let newRunTime = BMK_extract_runTime(runResult);
        let loopDuration_ns = newRunTime.nanoSecPerRun * f64::from(cont.nbLoops);

        cont.timeSpent_ns = cont.timeSpent_ns.wrapping_add(loopDuration_ns as PTime);

        // Estimate nbLoops for the next run to last approximately run_ms.
        if loopDuration_ns > runBudget_ns as f64 / 50.0 {
            let fastestRun_ns = bestRunTime.nanoSecPerRun.min(newRunTime.nanoSecPerRun);
            cont.nbLoops = ((runBudget_ns as f64 / fastestRun_ns) as c_uint).wrapping_add(1);
        } else {
            // Previous run was too short: blindly increase workload by a
            // x10 multiplier.
            const MULTIPLIER: c_uint = 10;
            debug_assert!(cont.nbLoops < c_uint::MAX / MULTIPLIER); // avoid overflow
            cont.nbLoops = cont.nbLoops.wrapping_mul(MULTIPLIER);
        }

        if loopDuration_ns < runTimeMin_ns as f64 {
            // Don't report results when the run time was too small, which
            // increases the risk of rounding errors.
            continue;
        }
        if newRunTime.nanoSecPerRun < bestRunTime.nanoSecPerRun {
            bestRunTime = newRunTime;
        }
        return valid_outcome(bestRunTime);
    }
}

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

    /// Test payload observed through `benchPayload`/`initPayload` pointers.
    #[derive(Default)]
    struct CallLog {
        bench_calls: usize,
        init_calls: usize,
    }

    /// Counts invocations and reports `srcSize`, like a size-preserving codec.
    unsafe extern "C" fn counting_bench_fn(
        _src: *const c_void,
        srcSize: usize,
        _dst: *mut c_void,
        _dstCapacity: usize,
        payload: *mut c_void,
    ) -> usize {
        let log = unsafe { &mut *payload.cast::<CallLog>() };
        log.bench_calls += 1;
        srcSize
    }

    unsafe extern "C" fn counting_init_fn(payload: *mut c_void) -> usize {
        let log = unsafe { &mut *payload.cast::<CallLog>() };
        log.init_calls += 1;
        0
    }

    /// Flags results of 5 bytes and above as errors.
    unsafe extern "C" fn error_on_5(result: usize) -> c_uint {
        c_uint::from(result >= 5)
    }

    struct Fixture {
        srcs: Vec<Vec<u8>>,
        dsts: Vec<Vec<u8>>,
        src_ptrs: Vec<*const c_void>,
        src_sizes: Vec<usize>,
        dst_ptrs: Vec<*mut c_void>,
        dst_capacities: Vec<usize>,
        block_results: Vec<usize>,
        log: CallLog,
    }

    impl Fixture {
        fn new(block_sizes: &[usize]) -> Box<Self> {
            let srcs: Vec<Vec<u8>> = block_sizes.iter().map(|size| vec![0u8; *size]).collect();
            let mut dsts: Vec<Vec<u8>> = block_sizes.iter().map(|size| vec![0u8; *size]).collect();
            let src_ptrs = srcs.iter().map(|src| src.as_ptr().cast()).collect();
            let src_sizes = srcs.iter().map(Vec::len).collect();
            let dst_ptrs = dsts.iter_mut().map(|dst| dst.as_mut_ptr().cast()).collect();
            let dst_capacities = dsts.iter().map(Vec::len).collect();
            let block_results = vec![0usize; block_sizes.len()];
            Box::new(Self {
                srcs,
                dsts,
                src_ptrs,
                src_sizes,
                dst_ptrs,
                dst_capacities,
                block_results,
                log: CallLog::default(),
            })
        }

        fn params(&mut self, errorFn: BMK_errorFn_t) -> BMK_benchParams_t {
            let payload: *mut CallLog = &mut self.log;
            BMK_benchParams_t {
                benchFn: Some(counting_bench_fn),
                benchPayload: payload.cast(),
                initFn: Some(counting_init_fn),
                initPayload: payload.cast(),
                errorFn,
                blockCount: self.srcs.len(),
                srcBuffers: self.src_ptrs.as_ptr(),
                srcSizes: self.src_sizes.as_ptr(),
                dstBuffers: self.dst_ptrs.as_ptr(),
                dstCapacities: self.dst_capacities.as_ptr(),
                blockResults: self.block_results.as_mut_ptr(),
            }
        }
    }

    #[test]
    fn bench_function_accounts_loops_blocks_and_first_loop_results() {
        let mut fixture = Fixture::new(&[3, 8]);
        let params = fixture.params(None);

        let outcome = unsafe { BMK_benchFunction(params, 4) };

        assert_eq!(BMK_isSuccessful_runOutcome(outcome), 1);
        let run_time = BMK_extract_runTime(outcome);
        // benchFn ran nbLoops times over each block; initFn ran once.
        assert_eq!(fixture.log.bench_calls, 4 * 2);
        assert_eq!(fixture.log.init_calls, 1);
        // sumOfReturn and blockResults reflect the first loop only.
        assert_eq!(run_time.sumOfReturn, 3 + 8);
        assert_eq!(fixture.block_results, vec![3, 8]);
        assert!(run_time.nanoSecPerRun >= 0.0);
    }

    #[test]
    fn bench_function_treats_zero_loops_as_one_and_warms_up_buffers() {
        let mut fixture = Fixture::new(&[4]);
        let params = fixture.params(None);

        let outcome = unsafe { BMK_benchFunction(params, 0) };

        assert_eq!(BMK_isSuccessful_runOutcome(outcome), 1);
        assert_eq!(fixture.log.bench_calls, 1);
        // The result buffer was erased with the 0xE5 warm-up pattern.
        assert_eq!(fixture.dsts[0], vec![0xE5; 4]);
    }

    #[test]
    fn bench_function_reports_the_first_failing_block() {
        let mut fixture = Fixture::new(&[3, 5, 7]);
        let params = fixture.params(Some(error_on_5));

        let outcome = unsafe { BMK_benchFunction(params, 10) };

        assert_eq!(BMK_isSuccessful_runOutcome(outcome), 0);
        assert_eq!(BMK_extract_errorResult(outcome), 5);
        // Execution stopped at the failing block, before the third one.
        assert_eq!(fixture.log.bench_calls, 2);
        // blockResults were recorded up to and including the failure.
        assert_eq!(fixture.block_results[..2], [3, 5]);
    }

    #[test]
    fn reset_clamps_budgets_and_rearms_the_loop_counter() {
        let state = BMK_createTimedFnState(0, 7);
        assert!(!state.is_null());

        {
            let state = unsafe { &*state };
            // total_ms 0 becomes 1ms, and run_ms is clamped to total_ms.
            assert_eq!(state.timeBudget_ns, 1_000_000);
            assert_eq!(state.runBudget_ns, 1_000_000);
            assert_eq!(state.nbLoops, 1);
            assert_eq!(state.timeSpent_ns, 0);
            assert_eq!(state.fastestRun.sumOfReturn, usize::MAX);
        }
        assert_eq!(unsafe { BMK_isCompleted_TimedFn(state) }, 0);

        unsafe { BMK_resetTimedFnState(state, 2_000, 500) };
        {
            let state = unsafe { &*state };
            assert_eq!(state.timeBudget_ns, 2_000_000_000);
            assert_eq!(state.runBudget_ns, 500_000_000);
        }

        unsafe { BMK_freeTimedFnState(state) };
    }

    #[test]
    fn static_state_initialization_validates_its_buffer() {
        let mut shell = [0u64; BMK_TIMEDFNSTATE_SIZE / 8];
        let buffer: *mut c_void = shell.as_mut_ptr().cast();

        // A properly sized and aligned buffer is accepted.
        let state = unsafe { BMK_initStatic_timedFnState(buffer, 64, 1_000, 100) };
        assert!(!state.is_null());
        assert_eq!(unsafe { BMK_isCompleted_TimedFn(state) }, 0);

        // NULL, undersized, and misaligned buffers are rejected.
        let too_small = std::mem::size_of::<BMK_timedFnState_t>() - 1;
        unsafe {
            assert!(BMK_initStatic_timedFnState(ptr::null_mut(), 64, 1, 1).is_null());
            assert!(BMK_initStatic_timedFnState(buffer, too_small, 1, 1).is_null());
            assert!(
                BMK_initStatic_timedFnState(buffer.cast::<u8>().add(1).cast(), 63, 1, 1).is_null()
            );
        }
    }

    #[test]
    fn timed_runs_grow_the_workload_and_spend_the_budget() {
        let mut fixture = Fixture::new(&[16]);
        let params = fixture.params(None);
        let state = BMK_createTimedFnState(4, 2);
        assert!(!state.is_null());

        let mut rounds = 0usize;
        while unsafe { BMK_isCompleted_TimedFn(state) } == 0 {
            let outcome = unsafe { BMK_benchTimedFn(state, params) };
            assert_eq!(BMK_isSuccessful_runOutcome(outcome), 1);
            let run_time = BMK_extract_runTime(outcome);
            assert_eq!(run_time.sumOfReturn, 16);
            rounds += 1;
            assert!(rounds < 1_000, "the time budget must eventually be spent");
        }

        {
            let state = unsafe { &*state };
            // A reported run had to last at least runBudget/2, which is only
            // reachable for this trivial function with a grown loop counter.
            assert!(state.nbLoops > 1);
            assert!(state.timeSpent_ns >= state.timeBudget_ns);
        }
        // Every reported outcome came from a run of >= runBudget/2, and the
        // budget accounting matches BMK_isCompleted_TimedFn.
        assert!(rounds >= 1);
        assert!(fixture.log.bench_calls >= rounds);

        unsafe { BMK_freeTimedFnState(state) };
    }

    #[test]
    fn timed_runs_propagate_errors_without_aborting() {
        let mut fixture = Fixture::new(&[9]);
        let params = fixture.params(Some(error_on_5));
        let state = BMK_createTimedFnState(1_000, 100);

        let outcome = unsafe { BMK_benchTimedFn(state, params) };

        assert_eq!(BMK_isSuccessful_runOutcome(outcome), 0);
        assert_eq!(BMK_extract_errorResult(outcome), 9);

        unsafe { BMK_freeTimedFnState(state) };
    }
}