zstd-rs/rust/src/zstdmt_compress.rs

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

//! Rust-owned resource pools used by the multithreaded compressor.
//!
//! The serial LDM state, job descriptor fields, worker callback, and streaming
//! state still use private C layouts.  `zstdmt_compress.c` therefore keeps
//! those operations and projects only allocation/lifecycle pieces and pure
//! sizing policy into this module.  The entry points below are narrow C ABIs:
//! buffers, `ZSTD_CCtx *` values, and job descriptors remain opaque to Rust,
//! while allocation, reuse, expansion, synchronization, and sizing policy are
//! Rust-owned.

use std::mem::{self, MaybeUninit};
use std::os::raw::{c_int, c_uint, c_void};
use std::ptr;
use std::sync::Mutex;

const ZSTDMT_JOBLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 29 } else { 30 };
const ZSTD_WINDOWLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 30 } else { 31 };

const ZSTD_FAST: c_int = 1;
const ZSTD_DFAST: c_int = 2;
const ZSTD_GREEDY: c_int = 3;
const ZSTD_LAZY: c_int = 4;
const ZSTD_LAZY2: c_int = 5;
const ZSTD_BTLAZY2: c_int = 6;
const ZSTD_BTOPT: c_int = 7;
const ZSTD_BTULTRA: c_int = 8;
const ZSTD_BTULTRA2: c_int = 9;
const ZSTD_PS_ENABLE: c_int = 1;
#[cfg(test)]
const ZSTD_PS_DISABLE: c_int = 2;

#[inline]
fn cycle_log(chain_log: c_uint, strategy: c_int) -> c_uint {
    chain_log.wrapping_sub((strategy >= ZSTD_BTLAZY2) as c_uint)
}

#[inline]
fn compute_target_job_log(
    window_log: c_uint,
    chain_log: c_uint,
    strategy: c_int,
    enable_ldm: c_int,
) -> c_uint {
    let job_log = if enable_ldm == ZSTD_PS_ENABLE {
        /* In Long Range Mode, the windowLog is typically oversized.
         * In which case, it's preferable to determine the jobSize
         * based on cycleLog instead. */
        21.max(cycle_log(chain_log, strategy).wrapping_add(3))
    } else {
        20.max(window_log.wrapping_add(2))
    };
    job_log.min(ZSTDMT_JOBLOG_MAX)
}

#[inline]
fn overlap_log_default(strategy: c_int) -> c_int {
    match strategy {
        ZSTD_BTULTRA2 => 9,
        ZSTD_BTULTRA | ZSTD_BTOPT => 8,
        ZSTD_BTLAZY2 | ZSTD_LAZY2 => 7,
        ZSTD_LAZY | ZSTD_GREEDY | ZSTD_DFAST | ZSTD_FAST => 6,
        _ => 6,
    }
}

#[inline]
fn overlap_log(overlap_log: c_int, strategy: c_int) -> c_int {
    debug_assert!((0..=9).contains(&overlap_log));
    if overlap_log == 0 {
        overlap_log_default(strategy)
    } else {
        overlap_log
    }
}

#[inline]
fn compute_overlap_size(
    window_log: c_uint,
    chain_log: c_uint,
    strategy: c_int,
    overlap_log_value: c_int,
    enable_ldm: c_int,
) -> usize {
    let overlap_r_log = 9 - overlap_log(overlap_log_value, strategy);
    let mut overlap_log_value = if overlap_r_log >= 8 {
        0
    } else {
        window_log as c_int - overlap_r_log
    };
    debug_assert!((0..=8).contains(&overlap_r_log));
    if enable_ldm == ZSTD_PS_ENABLE {
        /* In Long Range Mode, the windowLog is typically oversized.
         * In which case, it's preferable to determine the jobSize
         * based on chainLog instead.
         * Then, ovLog becomes a fraction of the jobSize, rather than windowSize */
        let target_job_log = compute_target_job_log(window_log, chain_log, strategy, enable_ldm);
        overlap_log_value = window_log.min(target_job_log.wrapping_sub(2)) as c_int - overlap_r_log;
    }
    debug_assert!(overlap_log_value >= 0);
    debug_assert!(overlap_log_value <= ZSTD_WINDOWLOG_MAX as c_int);
    if overlap_log_value == 0 {
        0
    } else {
        1usize << overlap_log_value as usize
    }
}

/// C ABI for the pure MT target job-log policy.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_computeTargetJobLog(
    windowLog: c_uint,
    chainLog: c_uint,
    strategy: c_int,
    enableLdm: c_int,
) -> c_uint {
    compute_target_job_log(windowLog, chainLog, strategy, enableLdm)
}

/// C ABI for the MT overlap-log default/selection policy.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_overlapLog(overlapLog: c_int, strategy: c_int) -> c_int {
    overlap_log(overlapLog, strategy)
}

/// C ABI for the pure MT overlap-size policy.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_computeOverlapSize(
    windowLog: c_uint,
    chainLog: c_uint,
    strategy: c_int,
    overlapLog: c_int,
    enableLdm: c_int,
) -> usize {
    compute_overlap_size(windowLog, chainLog, strategy, overlapLog, enableLdm)
}

type ZstdAllocFunction = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void;
type ZstdFreeFunction = unsafe extern "C" fn(*mut c_void, *mut c_void);

/// ABI-compatible representation of `ZSTD_customMem` from `zstd.h`.
#[repr(C)]
#[derive(Clone, Copy)]
pub struct ZstdCustomMem {
    pub customAlloc: Option<ZstdAllocFunction>,
    pub customFree: Option<ZstdFreeFunction>,
    pub opaque: *mut c_void,
}

/// ABI-compatible buffer returned to the C adapter.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct ZstdMtBuffer {
    pub start: *mut c_void,
    pub capacity: usize,
}

/// ABI-compatible representation of `rawSeq` from `zstd_compress_internal.h`.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct ZstdMtRawSeq {
    pub offset: u32,
    pub litLength: u32,
    pub matchLength: u32,
}

/// ABI-compatible representation of `RawSeqStore_t` from
/// `zstd_compress_internal.h`.
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct ZstdMtRawSeqStore {
    pub seq: *mut ZstdMtRawSeq,
    pub pos: usize,
    pub posInSequence: usize,
    pub size: usize,
    pub capacity: usize,
}

#[inline]
fn buffer_to_seq(buffer: ZstdMtBuffer) -> ZstdMtRawSeqStore {
    ZstdMtRawSeqStore {
        seq: buffer.start.cast::<ZstdMtRawSeq>(),
        pos: 0,
        posInSequence: 0,
        size: 0,
        capacity: buffer.capacity / mem::size_of::<ZstdMtRawSeq>(),
    }
}

#[inline]
fn seq_to_buffer(seq: ZstdMtRawSeqStore) -> ZstdMtBuffer {
    ZstdMtBuffer {
        start: seq.seq.cast::<c_void>(),
        capacity: seq.capacity.wrapping_mul(mem::size_of::<ZstdMtRawSeq>()),
    }
}

#[inline]
fn is_overlapped(
    buffer_start: *const c_void,
    buffer_capacity: usize,
    range_start: *const c_void,
    range_size: usize,
) -> c_int {
    if buffer_start.is_null() || range_start.is_null() {
        return 0;
    }

    let buffer_start = buffer_start.cast::<u8>();
    let range_start = range_start.cast::<u8>();
    let buffer_end = buffer_start.wrapping_add(buffer_capacity);
    let range_end = range_start.wrapping_add(range_size);

    /* Empty ranges cannot overlap. */
    if buffer_start == buffer_end || range_start == range_end {
        return 0;
    }

    (buffer_start < range_end && range_start < buffer_end) as c_int
}

/// Convert a byte buffer into the raw-sequence store view used by the MT
/// sequence pool. The capacity is expressed in whole `rawSeq` elements, just
/// like the original C conversion leaf.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_bufferToSeq(buffer: ZstdMtBuffer) -> ZstdMtRawSeqStore {
    buffer_to_seq(buffer)
}

/// Convert the raw-sequence store view back to a byte buffer for pool APIs.
/// The returned capacity is measured in bytes and follows C `size_t` wraparound
/// semantics for the multiplication.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_seqToBuffer(seq: ZstdMtRawSeqStore) -> ZstdMtBuffer {
    seq_to_buffer(seq)
}

/// Return non-zero when two non-empty byte ranges overlap.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_isOverlapped(
    bufferStart: *const c_void,
    bufferCapacity: usize,
    rangeStart: *const c_void,
    rangeSize: usize,
) -> c_int {
    is_overlapped(bufferStart, bufferCapacity, rangeStart, rangeSize)
}

/// Return non-zero when a buffer overlaps either the external dictionary or
/// the active prefix represented by a C `ZSTD_window_t`.
///
/// The C window stores the two ranges as pointer/index pairs.  Passing those
/// scalar fields separately keeps the private C struct out of the Rust ABI;
/// the byte-distance arithmetic mirrors the original pointer subtraction.
#[no_mangle]
pub extern "C" fn ZSTDMT_rust_doesOverlapWindow(
    bufferStart: *const c_void,
    bufferCapacity: usize,
    nextSrc: *const c_void,
    base: *const c_void,
    dictBase: *const c_void,
    dictLimit: u32,
    lowLimit: u32,
) -> c_int {
    let ext_dict_start = dictBase.cast::<u8>().wrapping_add(lowLimit as usize);
    let ext_dict_size = dictLimit.wrapping_sub(lowLimit) as usize;
    let prefix_start = base.cast::<u8>().wrapping_add(dictLimit as usize);
    let prefix_size = (nextSrc as usize)
        .wrapping_sub(base as usize)
        .wrapping_sub(dictLimit as usize);

    (is_overlapped(
        bufferStart,
        bufferCapacity,
        ext_dict_start.cast(),
        ext_dict_size,
    ) != 0
        || is_overlapped(
            bufferStart,
            bufferCapacity,
            prefix_start.cast(),
            prefix_size,
        ) != 0) as c_int
}

#[derive(Default)]
struct BufferPoolState {
    buffer_size: usize,
    nb_buffers: usize,
}

/// The object is allocated with the caller's `ZSTD_customMem`; its mutex
/// protects all mutable pool state and the raw reusable-buffer array.
pub struct RustBufferPool {
    custom_mem: ZstdCustomMem,
    total_buffers: usize,
    buffers: *mut MaybeUninit<ZstdMtBuffer>,
    state: Mutex<BufferPoolState>,
}

// The raw array is only accessed while `state` is held.  The pool itself is
// passed between C worker threads as an opaque pointer.
unsafe impl Send for RustBufferPool {}
unsafe impl Sync for RustBufferPool {}

/// The `ZSTD_CCtx *` values are created and destroyed by the existing C API;
/// Rust owns only the synchronized reusable-pointer pool.
pub struct RustCCtxPool {
    custom_mem: ZstdCustomMem,
    total_cctx: usize,
    cctxs: *mut MaybeUninit<*mut c_void>,
    state: Mutex<usize>,
}

unsafe impl Send for RustCCtxPool {}
unsafe impl Sync for RustCCtxPool {}

unsafe extern "C" {
    fn ZSTD_createCCtx_advanced(custom_mem: ZstdCustomMem) -> *mut c_void;
    fn ZSTD_freeCCtx(cctx: *mut c_void) -> usize;
    fn ZSTD_sizeof_CCtx(cctx: *const c_void) -> usize;
}

unsafe fn custom_calloc(size: usize, custom_mem: ZstdCustomMem) -> *mut c_void {
    if let Some(alloc) = custom_mem.customAlloc {
        let allocation = unsafe { alloc(custom_mem.opaque, size) };
        if !allocation.is_null() {
            unsafe { ptr::write_bytes(allocation, 0, size) };
        }
        allocation
    } else {
        unsafe { libc::calloc(1, size) }
    }
}

unsafe fn custom_malloc(size: usize, custom_mem: ZstdCustomMem) -> *mut c_void {
    if let Some(alloc) = custom_mem.customAlloc {
        unsafe { alloc(custom_mem.opaque, size) }
    } else {
        unsafe { libc::malloc(size) }
    }
}

unsafe fn custom_free(allocation: *mut c_void, custom_mem: ZstdCustomMem) {
    if allocation.is_null() {
        return;
    }
    if let Some(free) = custom_mem.customFree {
        unsafe { free(custom_mem.opaque, allocation) };
    } else {
        unsafe { libc::free(allocation) };
    }
}

fn checked_array_size<T>(len: usize) -> Option<usize> {
    mem::size_of::<T>().checked_mul(len)
}

fn rounded_job_count(requested: c_uint) -> Option<c_uint> {
    if requested == 0 {
        return None;
    }
    // The original C expression is `1 << (highbit(requested) + 1)`, which
    // deliberately chooses a strictly larger power of two when requested is
    // already a power of two.
    let shift = usize::BITS - (requested as usize).leading_zeros();
    let count = 1usize.checked_shl(shift)?;
    if count > c_uint::MAX as usize {
        return None;
    }
    Some(count as c_uint)
}

unsafe fn create_job_table(
    nb_jobs_ptr: *mut c_uint,
    job_size: usize,
    custom_mem: ZstdCustomMem,
) -> *mut c_void {
    if nb_jobs_ptr.is_null() || job_size == 0 {
        return ptr::null_mut();
    }
    let Some(nb_jobs) = (unsafe { rounded_job_count(*nb_jobs_ptr) }) else {
        return ptr::null_mut();
    };
    let Some(table_size) = job_size.checked_mul(nb_jobs as usize) else {
        return ptr::null_mut();
    };
    let table = unsafe { custom_calloc(table_size, custom_mem) };
    if table.is_null() {
        return ptr::null_mut();
    }

    unsafe { *nb_jobs_ptr = nb_jobs };
    table
}

unsafe fn free_job_table_storage(job_table: *mut c_void, custom_mem: ZstdCustomMem) {
    if job_table.is_null() {
        return;
    }
    unsafe { custom_free(job_table, custom_mem) };
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_job_table_create(
    nb_jobs_ptr: *mut c_uint,
    job_size: usize,
    custom_mem: ZstdCustomMem,
) -> *mut c_void {
    unsafe { create_job_table(nb_jobs_ptr, job_size, custom_mem) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_job_table_free(
    job_table: *mut c_void,
    _nb_jobs: c_uint,
    _job_size: usize,
    custom_mem: ZstdCustomMem,
) {
    // The MT C adapter destroys its platform mutexes and condition variables
    // before calling this storage-only release function.  Keeping this Rust
    // side free of MT-only C references also preserves single-threaded builds
    // where zstdmt_compress.c is intentionally omitted.
    unsafe { free_job_table_storage(job_table, custom_mem) }
}

unsafe fn create_buffer_pool(
    max_nb_buffers: usize,
    custom_mem: ZstdCustomMem,
) -> *mut RustBufferPool {
    if max_nb_buffers == 0 {
        return ptr::null_mut();
    }
    let Some(buffer_bytes) = checked_array_size::<MaybeUninit<ZstdMtBuffer>>(max_nb_buffers) else {
        return ptr::null_mut();
    };

    let pool = unsafe { custom_calloc(mem::size_of::<RustBufferPool>(), custom_mem) }
        .cast::<RustBufferPool>();
    if pool.is_null() {
        return ptr::null_mut();
    }
    let buffers =
        unsafe { custom_calloc(buffer_bytes, custom_mem) }.cast::<MaybeUninit<ZstdMtBuffer>>();
    if buffers.is_null() {
        unsafe { custom_free(pool.cast(), custom_mem) };
        return ptr::null_mut();
    }

    unsafe {
        pool.write(RustBufferPool {
            custom_mem,
            total_buffers: max_nb_buffers,
            buffers,
            state: Mutex::new(BufferPoolState {
                buffer_size: 64 << 10,
                nb_buffers: 0,
            }),
        });
    }
    pool
}

unsafe fn destroy_buffer_pool(pool: *mut RustBufferPool) {
    if pool.is_null() {
        return;
    }
    let custom_mem = unsafe { (*pool).custom_mem };
    let buffers = unsafe { (*pool).buffers };
    let total_buffers = unsafe { (*pool).total_buffers };
    for index in 0..total_buffers {
        let buffer = unsafe { buffers.add(index).read().assume_init() };
        unsafe { custom_free(buffer.start, custom_mem) };
    }
    unsafe { ptr::drop_in_place(pool) };
    unsafe {
        custom_free(buffers.cast(), custom_mem);
        custom_free(pool.cast(), custom_mem);
    }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_create(
    max_nb_buffers: c_uint,
    custom_mem: ZstdCustomMem,
) -> *mut RustBufferPool {
    unsafe { create_buffer_pool(max_nb_buffers as usize, custom_mem) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_free(pool: *mut RustBufferPool) {
    unsafe { destroy_buffer_pool(pool) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_sizeof(pool: *const RustBufferPool) -> usize {
    if pool.is_null() {
        return 0;
    }
    let pool_ref = unsafe { &*pool };
    let _state = pool_ref
        .state
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    let mut total_buffer_size = 0usize;
    for index in 0..pool_ref.total_buffers {
        let buffer = unsafe { pool_ref.buffers.add(index).read().assume_init() };
        total_buffer_size = total_buffer_size.saturating_add(buffer.capacity);
    }
    mem::size_of::<RustBufferPool>()
        .saturating_add(
            pool_ref
                .total_buffers
                .saturating_mul(mem::size_of::<MaybeUninit<ZstdMtBuffer>>()),
        )
        .saturating_add(total_buffer_size)
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_set_size(
    pool: *mut RustBufferPool,
    buffer_size: usize,
) {
    if pool.is_null() {
        return;
    }
    let pool_ref = unsafe { &*pool };
    let mut state = pool_ref
        .state
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    state.buffer_size = buffer_size;
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_expand(
    pool: *mut RustBufferPool,
    max_nb_buffers: c_uint,
) -> *mut RustBufferPool {
    if pool.is_null() {
        return ptr::null_mut();
    }
    let max_nb_buffers = max_nb_buffers as usize;
    let pool_ref = unsafe { &*pool };
    let (total_buffers, buffer_size, custom_mem) = {
        let state = pool_ref
            .state
            .lock()
            .unwrap_or_else(|error| error.into_inner());
        (
            pool_ref.total_buffers,
            state.buffer_size,
            pool_ref.custom_mem,
        )
    };
    if total_buffers >= max_nb_buffers {
        return pool;
    }

    // This matches the original resize contract: the old pool is consumed
    // before creating the larger replacement.
    unsafe { destroy_buffer_pool(pool) };
    let replacement = unsafe { create_buffer_pool(max_nb_buffers, custom_mem) };
    if !replacement.is_null() {
        unsafe { ZSTDMT_rust_buffer_pool_set_size(replacement, buffer_size) };
    }
    replacement
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_get(pool: *mut RustBufferPool) -> ZstdMtBuffer {
    if pool.is_null() {
        return ZstdMtBuffer::default();
    }
    let pool_ref = unsafe { &*pool };
    let (buffer_size, reusable) = {
        let mut state = pool_ref
            .state
            .lock()
            .unwrap_or_else(|error| error.into_inner());
        if state.buffer_size == 0 {
            return ZstdMtBuffer::default();
        }
        if state.nb_buffers == 0 {
            (state.buffer_size, ZstdMtBuffer::default())
        } else {
            state.nb_buffers -= 1;
            let index = state.nb_buffers;
            let buffer = unsafe { pool_ref.buffers.add(index).read().assume_init() };
            unsafe {
                pool_ref
                    .buffers
                    .add(index)
                    .write(MaybeUninit::new(ZstdMtBuffer::default()));
            }
            (state.buffer_size, buffer)
        }
    };

    if !reusable.start.is_null()
        && reusable.capacity >= buffer_size
        && (reusable.capacity >> 3) <= buffer_size
    {
        return reusable;
    }
    if !reusable.start.is_null() {
        unsafe { custom_free(reusable.start, pool_ref.custom_mem) };
    }

    let start = unsafe { custom_malloc(buffer_size, pool_ref.custom_mem) };
    ZstdMtBuffer {
        start,
        capacity: if start.is_null() { 0 } else { buffer_size },
    }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_release(
    pool: *mut RustBufferPool,
    buffer: ZstdMtBuffer,
) {
    if pool.is_null() || buffer.start.is_null() {
        return;
    }
    let pool_ref = unsafe { &*pool };
    let mut state = pool_ref
        .state
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    if state.nb_buffers < pool_ref.total_buffers {
        let index = state.nb_buffers;
        unsafe {
            pool_ref.buffers.add(index).write(MaybeUninit::new(buffer));
        }
        state.nb_buffers += 1;
        return;
    }
    drop(state);
    unsafe { custom_free(buffer.start, pool_ref.custom_mem) };
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_resize(
    pool: *mut RustBufferPool,
    buffer: ZstdMtBuffer,
) -> ZstdMtBuffer {
    if pool.is_null() || buffer.start.is_null() {
        return buffer;
    }
    let pool_ref = unsafe { &*pool };
    let buffer_size = {
        let state = pool_ref
            .state
            .lock()
            .unwrap_or_else(|error| error.into_inner());
        state.buffer_size
    };
    if buffer.capacity >= buffer_size {
        return buffer;
    }
    let start = unsafe { custom_malloc(buffer_size, pool_ref.custom_mem) };
    if start.is_null() {
        return buffer;
    }
    unsafe {
        ptr::copy_nonoverlapping(
            buffer.start.cast::<u8>(),
            start.cast::<u8>(),
            buffer.capacity,
        );
        custom_free(buffer.start, pool_ref.custom_mem);
    }
    ZstdMtBuffer {
        start,
        capacity: buffer_size,
    }
}

unsafe fn create_cctx_pool(nb_workers: usize, custom_mem: ZstdCustomMem) -> *mut RustCCtxPool {
    if nb_workers == 0 {
        return ptr::null_mut();
    }
    let Some(cctx_bytes) = checked_array_size::<MaybeUninit<*mut c_void>>(nb_workers) else {
        return ptr::null_mut();
    };
    let pool =
        unsafe { custom_calloc(mem::size_of::<RustCCtxPool>(), custom_mem) }.cast::<RustCCtxPool>();
    if pool.is_null() {
        return ptr::null_mut();
    }
    let cctxs = unsafe { custom_calloc(cctx_bytes, custom_mem) }.cast::<MaybeUninit<*mut c_void>>();
    if cctxs.is_null() {
        unsafe { custom_free(pool.cast(), custom_mem) };
        return ptr::null_mut();
    }
    let first = unsafe { ZSTD_createCCtx_advanced(custom_mem) };
    if first.is_null() {
        unsafe {
            custom_free(cctxs.cast(), custom_mem);
            custom_free(pool.cast(), custom_mem);
        }
        return ptr::null_mut();
    }
    unsafe { cctxs.write(MaybeUninit::new(first)) };
    unsafe {
        pool.write(RustCCtxPool {
            custom_mem,
            total_cctx: nb_workers,
            cctxs,
            state: Mutex::new(1),
        });
    }
    pool
}

unsafe fn destroy_cctx_pool(pool: *mut RustCCtxPool) {
    if pool.is_null() {
        return;
    }
    let custom_mem = unsafe { (*pool).custom_mem };
    let cctxs = unsafe { (*pool).cctxs };
    let total_cctx = unsafe { (*pool).total_cctx };
    for index in 0..total_cctx {
        let cctx = unsafe { cctxs.add(index).read().assume_init() };
        unsafe { ZSTD_freeCCtx(cctx) };
    }
    unsafe { ptr::drop_in_place(pool) };
    unsafe {
        custom_free(cctxs.cast(), custom_mem);
        custom_free(pool.cast(), custom_mem);
    }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_create(
    nb_workers: c_uint,
    custom_mem: ZstdCustomMem,
) -> *mut RustCCtxPool {
    unsafe { create_cctx_pool(nb_workers as usize, custom_mem) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_free(pool: *mut RustCCtxPool) {
    unsafe { destroy_cctx_pool(pool) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_sizeof(pool: *const RustCCtxPool) -> usize {
    if pool.is_null() {
        return 0;
    }
    let pool_ref = unsafe { &*pool };
    let _state = pool_ref
        .state
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    let mut total_cctx_size = 0usize;
    for index in 0..pool_ref.total_cctx {
        let cctx = unsafe { pool_ref.cctxs.add(index).read().assume_init() };
        total_cctx_size =
            total_cctx_size.saturating_add(unsafe { ZSTD_sizeof_CCtx(cctx.cast_const()) });
    }
    mem::size_of::<RustCCtxPool>()
        .saturating_add(
            pool_ref
                .total_cctx
                .saturating_mul(mem::size_of::<MaybeUninit<*mut c_void>>()),
        )
        .saturating_add(total_cctx_size)
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_expand(
    pool: *mut RustCCtxPool,
    nb_workers: c_uint,
) -> *mut RustCCtxPool {
    if pool.is_null() {
        return ptr::null_mut();
    }
    let nb_workers = nb_workers as usize;
    let pool_ref = unsafe { &*pool };
    if pool_ref.total_cctx >= nb_workers {
        return pool;
    }
    let custom_mem = pool_ref.custom_mem;
    unsafe { destroy_cctx_pool(pool) };
    unsafe { create_cctx_pool(nb_workers, custom_mem) }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_get(pool: *mut RustCCtxPool) -> *mut c_void {
    if pool.is_null() {
        return ptr::null_mut();
    }
    let pool_ref = unsafe { &*pool };
    let cctx = {
        let mut available = pool_ref
            .state
            .lock()
            .unwrap_or_else(|error| error.into_inner());
        if *available == 0 {
            ptr::null_mut()
        } else {
            *available -= 1;
            let index = *available;
            unsafe { pool_ref.cctxs.add(index).read().assume_init() }
        }
    };
    if !cctx.is_null() {
        cctx
    } else {
        unsafe { ZSTD_createCCtx_advanced(pool_ref.custom_mem) }
    }
}

#[no_mangle]
pub unsafe extern "C" fn ZSTDMT_rust_cctx_pool_release(pool: *mut RustCCtxPool, cctx: *mut c_void) {
    if pool.is_null() || cctx.is_null() {
        return;
    }
    let pool_ref = unsafe { &*pool };
    let mut available = pool_ref
        .state
        .lock()
        .unwrap_or_else(|error| error.into_inner());
    if *available < pool_ref.total_cctx {
        let index = *available;
        unsafe {
            pool_ref.cctxs.add(index).write(MaybeUninit::new(cctx));
        }
        *available += 1;
        return;
    }
    drop(available);
    unsafe { ZSTD_freeCCtx(cctx) };
}

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

    const DEFAULT_MEM: ZstdCustomMem = ZstdCustomMem {
        customAlloc: None,
        customFree: None,
        opaque: ptr::null_mut(),
    };

    #[test]
    fn raw_seq_buffer_conversion_uses_whole_element_capacity() {
        let mut sequences = [ZstdMtRawSeq::default(); 3];
        let element_size = mem::size_of::<ZstdMtRawSeq>();
        let buffer = ZstdMtBuffer {
            start: sequences.as_mut_ptr().cast(),
            capacity: element_size * sequences.len() + element_size - 1,
        };

        let seq = ZSTDMT_rust_bufferToSeq(buffer);
        assert_eq!(seq.seq, sequences.as_mut_ptr());
        assert_eq!(seq.pos, 0);
        assert_eq!(seq.posInSequence, 0);
        assert_eq!(seq.size, 0);
        assert_eq!(seq.capacity, sequences.len());

        let roundtrip = ZSTDMT_rust_seqToBuffer(seq);
        assert_eq!(roundtrip.start, buffer.start);
        assert_eq!(roundtrip.capacity, element_size * sequences.len());
    }

    #[test]
    fn raw_seq_to_buffer_preserves_pointer_and_size_t_multiplication() {
        let seq = ZstdMtRawSeqStore {
            seq: ptr::null_mut(),
            pos: 4,
            posInSequence: 5,
            size: 6,
            capacity: usize::MAX,
        };
        let buffer = ZSTDMT_rust_seqToBuffer(seq);
        assert!(buffer.start.is_null());
        assert_eq!(
            buffer.capacity,
            usize::MAX.wrapping_mul(mem::size_of::<ZstdMtRawSeq>())
        );

        let empty = ZSTDMT_rust_bufferToSeq(ZstdMtBuffer {
            start: ptr::null_mut(),
            capacity: mem::size_of::<ZstdMtRawSeq>() - 1,
        });
        assert!(empty.seq.is_null());
        assert_eq!(empty.capacity, 0);
        assert_eq!(empty.pos, 0);
        assert_eq!(empty.posInSequence, 0);
        assert_eq!(empty.size, 0);
    }

    #[test]
    fn overlapped_rejects_null_ranges() {
        let bytes = [0u8; 8];
        let start = bytes.as_ptr().cast::<c_void>();

        assert_eq!(ZSTDMT_rust_isOverlapped(ptr::null(), 4, start, 4), 0);
        assert_eq!(ZSTDMT_rust_isOverlapped(start, 4, ptr::null(), 4), 0);
    }

    #[test]
    fn overlapped_rejects_empty_ranges() {
        let bytes = [0u8; 8];
        let start = bytes.as_ptr().cast::<c_void>();

        assert_eq!(ZSTDMT_rust_isOverlapped(start, 0, start, 4), 0);
        assert_eq!(ZSTDMT_rust_isOverlapped(start, 4, start, 0), 0);
    }

    #[test]
    fn overlapped_uses_half_open_range_boundaries() {
        let bytes = [0u8; 16];
        let start = bytes.as_ptr();
        let middle = start.wrapping_add(4);

        assert_eq!(
            ZSTDMT_rust_isOverlapped(start.cast(), 4, middle.cast(), 4),
            0
        );
        assert_eq!(
            ZSTDMT_rust_isOverlapped(middle.cast(), 4, start.cast(), 4),
            0
        );
    }

    #[test]
    fn overlapped_detects_contained_ranges() {
        let bytes = [0u8; 16];
        let start = bytes.as_ptr();
        let inner = start.wrapping_add(4);

        assert_eq!(
            ZSTDMT_rust_isOverlapped(start.cast(), 12, inner.cast(), 4),
            1
        );
        assert_eq!(
            ZSTDMT_rust_isOverlapped(inner.cast(), 4, start.cast(), 12),
            1
        );
    }

    #[test]
    fn overlapped_rejects_disjoint_ranges() {
        let bytes = [0u8; 16];
        let start = bytes.as_ptr();
        let after = start.wrapping_add(8);

        assert_eq!(
            ZSTDMT_rust_isOverlapped(start.cast(), 4, after.cast(), 4),
            0
        );
        assert_eq!(
            ZSTDMT_rust_isOverlapped(after.cast(), 4, start.cast(), 4),
            0
        );
    }

    #[test]
    fn overlap_window_checks_external_dictionary_and_prefix() {
        let bytes = [0u8; 32];
        let base = bytes.as_ptr();
        let next_src = base.wrapping_add(16);
        let dict_base = base.wrapping_add(16);

        assert_eq!(
            ZSTDMT_rust_doesOverlapWindow(
                base.wrapping_add(20).cast(),
                4,
                next_src.cast(),
                base.cast(),
                dict_base.cast(),
                8,
                4,
            ),
            1
        );
        assert_eq!(
            ZSTDMT_rust_doesOverlapWindow(
                base.wrapping_add(12).cast(),
                4,
                next_src.cast(),
                base.cast(),
                dict_base.cast(),
                8,
                4,
            ),
            1
        );
        assert_eq!(
            ZSTDMT_rust_doesOverlapWindow(
                base.cast(),
                4,
                next_src.cast(),
                base.cast(),
                dict_base.cast(),
                8,
                4,
            ),
            0
        );
    }

    #[test]
    fn buffer_pool_reuses_and_resizes_buffers() {
        let pool = unsafe { ZSTDMT_rust_buffer_pool_create(2, DEFAULT_MEM) };
        assert!(!pool.is_null());

        let first = unsafe { ZSTDMT_rust_buffer_pool_get(pool) };
        assert!(!first.start.is_null());
        assert_eq!(first.capacity, 64 << 10);
        unsafe { ZSTDMT_rust_buffer_pool_release(pool, first) };

        let reused = unsafe { ZSTDMT_rust_buffer_pool_get(pool) };
        assert_eq!(reused.start, first.start);
        assert_eq!(reused.capacity, first.capacity);
        unsafe { ZSTDMT_rust_buffer_pool_release(pool, reused) };

        unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 128 << 10) };
        let resized = unsafe { ZSTDMT_rust_buffer_pool_get(pool) };
        assert_eq!(resized.capacity, 128 << 10);
        unsafe {
            ZSTDMT_rust_buffer_pool_release(pool, resized);
            ZSTDMT_rust_buffer_pool_free(pool);
        }
    }

    #[test]
    fn zero_size_pool_does_not_allocate_sequence_storage() {
        let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) };
        assert!(!pool.is_null());
        unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 0) };
        let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(pool) };
        assert!(buffer.start.is_null());
        assert_eq!(buffer.capacity, 0);
        unsafe { ZSTDMT_rust_buffer_pool_free(pool) };
    }

    #[test]
    fn expansion_preserves_requested_buffer_size() {
        let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) };
        assert!(!pool.is_null());
        unsafe { ZSTDMT_rust_buffer_pool_set_size(pool, 96 << 10) };
        let expanded = unsafe { ZSTDMT_rust_buffer_pool_expand(pool, 3) };
        assert!(!expanded.is_null());
        let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(expanded) };
        assert_eq!(buffer.capacity, 96 << 10);
        unsafe {
            ZSTDMT_rust_buffer_pool_release(expanded, buffer);
            ZSTDMT_rust_buffer_pool_free(expanded);
        }
    }

    #[test]
    fn resize_preserves_existing_buffer_contents() {
        let pool = unsafe { ZSTDMT_rust_buffer_pool_create(1, DEFAULT_MEM) };
        assert!(!pool.is_null());
        let buffer = unsafe { ZSTDMT_rust_buffer_pool_get(pool) };
        assert!(!buffer.start.is_null());
        let sample = b"multithreaded-buffer-pool";
        unsafe {
            ptr::copy_nonoverlapping(sample.as_ptr(), buffer.start.cast(), sample.len());
            ZSTDMT_rust_buffer_pool_set_size(pool, buffer.capacity + 1);
        }
        let resized = unsafe { ZSTDMT_rust_buffer_pool_resize(pool, buffer) };
        assert_eq!(resized.capacity, buffer.capacity + 1);
        let contents =
            unsafe { std::slice::from_raw_parts(resized.start.cast::<u8>(), sample.len()) };
        assert_eq!(contents, sample);
        unsafe {
            ZSTDMT_rust_buffer_pool_release(pool, resized);
            ZSTDMT_rust_buffer_pool_free(pool);
        }
    }

    #[test]
    fn job_table_count_matches_c_power_of_two_contract() {
        assert_eq!(rounded_job_count(0), None);
        assert_eq!(rounded_job_count(1), Some(2));
        assert_eq!(rounded_job_count(3), Some(4));
        assert_eq!(rounded_job_count(4), Some(8));
        assert_eq!(rounded_job_count(255), Some(256));
        assert_eq!(rounded_job_count(256), Some(512));
    }

    #[test]
    fn target_job_log_preserves_ldm_and_non_ldm_policy() {
        assert_eq!(
            compute_target_job_log(18, 25, ZSTD_FAST, ZSTD_PS_DISABLE),
            20
        );
        assert_eq!(
            compute_target_job_log(30, 25, ZSTD_FAST, ZSTD_PS_DISABLE),
            ZSTDMT_JOBLOG_MAX
        );
        assert_eq!(
            compute_target_job_log(10, 25, ZSTD_FAST, ZSTD_PS_ENABLE),
            28
        );
        assert_eq!(
            compute_target_job_log(30, 25, ZSTD_BTULTRA2, ZSTD_PS_ENABLE),
            27
        );
        assert_eq!(
            compute_target_job_log(10, 30, ZSTD_BTULTRA2, ZSTD_PS_ENABLE),
            ZSTDMT_JOBLOG_MAX
        );
    }

    #[test]
    fn overlap_log_defaults_follow_strategy_groups() {
        assert_eq!(overlap_log_default(ZSTD_FAST), 6);
        assert_eq!(overlap_log_default(ZSTD_DFAST), 6);
        assert_eq!(overlap_log_default(ZSTD_GREEDY), 6);
        assert_eq!(overlap_log_default(ZSTD_LAZY), 6);
        assert_eq!(overlap_log_default(ZSTD_LAZY2), 7);
        assert_eq!(overlap_log_default(ZSTD_BTLAZY2), 7);
        assert_eq!(overlap_log_default(ZSTD_BTOPT), 8);
        assert_eq!(overlap_log_default(ZSTD_BTULTRA), 8);
        assert_eq!(overlap_log_default(ZSTD_BTULTRA2), 9);
        assert_eq!(overlap_log_default(0), 6);

        assert_eq!(overlap_log(0, ZSTD_BTULTRA2), 9);
        assert_eq!(overlap_log(0, ZSTD_FAST), 6);
        assert_eq!(overlap_log(5, ZSTD_BTULTRA2), 5);
        assert_eq!(overlap_log(9, ZSTD_FAST), 9);
    }

    #[test]
    fn overlap_size_uses_window_or_target_job_log_as_expected() {
        assert_eq!(
            compute_overlap_size(20, 25, ZSTD_FAST, 0, ZSTD_PS_DISABLE),
            1usize << 17
        );
        assert_eq!(
            compute_overlap_size(20, 25, ZSTD_BTULTRA2, 0, ZSTD_PS_DISABLE),
            1usize << 20
        );
        assert_eq!(
            compute_overlap_size(20, 25, ZSTD_FAST, 1, ZSTD_PS_DISABLE),
            0
        );
        assert_eq!(
            compute_overlap_size(30, 25, ZSTD_FAST, 0, ZSTD_PS_ENABLE),
            1usize << 23
        );
        assert_eq!(
            compute_overlap_size(30, 25, ZSTD_BTULTRA2, 0, ZSTD_PS_ENABLE),
            1usize << 25
        );
    }
}