zstd-rs/rust/src/dibio.rs

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

//! Dictionary training I/O for the command-line interface.
//!
//! The public entry point keeps the `programs/dibio.h` ABI.  File discovery,
//! sample selection, buffering, and dictionary output are Rust-owned; the
//! actual legacy, COVER, and fastCOVER trainers are reached through their
//! existing C ABI symbols.

use std::cmp::min;
use std::ffi::{c_char, c_void, CStr};
use std::fmt;
use std::fs;
use std::io::{self, Write};
use std::mem::size_of;
use std::os::raw::{c_int, c_uint};
use std::path::PathBuf;
use std::time::{Duration, Instant};

#[cfg(unix)]
use std::ffi::OsStr;
#[cfg(unix)]
use std::os::unix::ffi::OsStrExt;

const KB: usize = 1 << 10;
const MB: usize = 1 << 20;
const GB: usize = 1 << 30;

const SAMPLESIZE_MAX: usize = 128 * KB;
const MEMMULT: u64 = 11;
const COVER_MEMMULT: u64 = 9;
const FASTCOVER_MEMMULT: u64 = 1;
const NOISELENGTH: usize = 32;
const MAX_SAMPLES_SIZE: usize = 2 * GB;
const REFRESH_RATE: Duration = Duration::from_micros(1_000_000 / 6);

// The C implementation initializes `dictSize` with the public enum value
// `ZSTD_error_GENERIC` (1), rather than an encoded size_t error.
const ERROR_GENERIC: usize = 1;

const MODE_READ: &[u8] = b"rb\0";
const MODE_WRITE: &[u8] = b"wb\0";

/// ABI-compatible `ZDICT_params_t` from `lib/zdict.h`.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZDICT_params_t {
    pub compressionLevel: c_int,
    pub notificationLevel: c_uint,
    pub dictID: c_uint,
}

/// ABI-compatible `ZDICT_legacy_params_t` from `lib/zdict.h`.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ZDICT_legacy_params_t {
    pub selectivityLevel: c_uint,
    pub zParams: ZDICT_params_t,
}

/// ABI-compatible `ZDICT_cover_params_t` from `lib/zdict.h`.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct ZDICT_cover_params_t {
    pub k: c_uint,
    pub d: c_uint,
    pub steps: c_uint,
    pub nbThreads: c_uint,
    pub splitPoint: f64,
    pub shrinkDict: c_uint,
    pub shrinkDictMaxRegression: c_uint,
    pub zParams: ZDICT_params_t,
}

/// ABI-compatible `ZDICT_fastCover_params_t` from `lib/zdict.h`.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct ZDICT_fastCover_params_t {
    pub k: c_uint,
    pub d: c_uint,
    pub f: c_uint,
    pub steps: c_uint,
    pub nbThreads: c_uint,
    pub splitPoint: f64,
    pub accel: c_uint,
    pub shrinkDict: c_uint,
    pub shrinkDictMaxRegression: c_uint,
    pub zParams: ZDICT_params_t,
}

unsafe extern "C" {
    fn malloc(size: usize) -> *mut c_void;
    fn free(pointer: *mut c_void);

    fn fopen(filename: *const c_char, mode: *const c_char) -> *mut c_void;
    fn fread(pointer: *mut c_void, size: usize, count: usize, stream: *mut c_void) -> usize;
    fn fwrite(pointer: *const c_void, size: usize, count: usize, stream: *mut c_void) -> usize;
    fn fclose(stream: *mut c_void) -> c_int;

    fn ZDICT_isError(error_code: usize) -> c_uint;
    fn ZDICT_getErrorName(error_code: usize) -> *const c_char;
    fn ZDICT_trainFromBuffer_legacy(
        dict_buffer: *mut c_void,
        dict_buffer_capacity: usize,
        samples_buffer: *const c_void,
        samples_sizes: *const usize,
        nb_samples: c_uint,
        parameters: ZDICT_legacy_params_t,
    ) -> usize;
    fn ZDICT_trainFromBuffer_cover(
        dict_buffer: *mut c_void,
        dict_buffer_capacity: usize,
        samples_buffer: *const c_void,
        samples_sizes: *const usize,
        nb_samples: c_uint,
        parameters: ZDICT_cover_params_t,
    ) -> usize;
    fn ZDICT_optimizeTrainFromBuffer_cover(
        dict_buffer: *mut c_void,
        dict_buffer_capacity: usize,
        samples_buffer: *const c_void,
        samples_sizes: *const usize,
        nb_samples: c_uint,
        parameters: *mut ZDICT_cover_params_t,
    ) -> usize;
    fn ZDICT_trainFromBuffer_fastCover(
        dict_buffer: *mut c_void,
        dict_buffer_capacity: usize,
        samples_buffer: *const c_void,
        samples_sizes: *const usize,
        nb_samples: c_uint,
        parameters: ZDICT_fastCover_params_t,
    ) -> usize;
    fn ZDICT_optimizeTrainFromBuffer_fastCover(
        dict_buffer: *mut c_void,
        dict_buffer_capacity: usize,
        samples_buffer: *const c_void,
        samples_sizes: *const usize,
        nb_samples: c_uint,
        parameters: *mut ZDICT_fastCover_params_t,
    ) -> usize;
}

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

    static ERROR_NAME: &[u8] = b"test dictionary trainer error\0";

    unsafe fn write_test_dictionary(buffer: *mut c_void, capacity: usize) -> usize {
        let size = capacity.min(64);
        if !buffer.is_null() {
            std::ptr::write_bytes(buffer.cast::<u8>(), 0xA5, size);
        }
        size
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_isError(_code: usize) -> c_uint {
        0
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_getErrorName(_code: usize) -> *const c_char {
        ERROR_NAME.as_ptr().cast()
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_trainFromBuffer_legacy(
        buffer: *mut c_void,
        capacity: usize,
        _samples: *const c_void,
        _sample_sizes: *const usize,
        _nb_samples: c_uint,
        _parameters: ZDICT_legacy_params_t,
    ) -> usize {
        write_test_dictionary(buffer, capacity)
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_trainFromBuffer_cover(
        buffer: *mut c_void,
        capacity: usize,
        _samples: *const c_void,
        _sample_sizes: *const usize,
        _nb_samples: c_uint,
        _parameters: ZDICT_cover_params_t,
    ) -> usize {
        write_test_dictionary(buffer, capacity)
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_optimizeTrainFromBuffer_cover(
        buffer: *mut c_void,
        capacity: usize,
        _samples: *const c_void,
        _sample_sizes: *const usize,
        _nb_samples: c_uint,
        _parameters: *mut ZDICT_cover_params_t,
    ) -> usize {
        write_test_dictionary(buffer, capacity)
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_trainFromBuffer_fastCover(
        buffer: *mut c_void,
        capacity: usize,
        _samples: *const c_void,
        _sample_sizes: *const usize,
        _nb_samples: c_uint,
        _parameters: ZDICT_fastCover_params_t,
    ) -> usize {
        write_test_dictionary(buffer, capacity)
    }

    #[no_mangle]
    pub unsafe extern "C" fn ZDICT_optimizeTrainFromBuffer_fastCover(
        buffer: *mut c_void,
        capacity: usize,
        _samples: *const c_void,
        _sample_sizes: *const usize,
        _nb_samples: c_uint,
        _parameters: *mut ZDICT_fastCover_params_t,
    ) -> usize {
        write_test_dictionary(buffer, capacity)
    }
}

#[derive(Clone, Copy, Debug, Default)]
struct FileStats {
    total_size_to_load: i64,
    nb_samples: c_int,
    one_sample_too_large: bool,
}

struct DisplayClock {
    last_update: Instant,
}

impl DisplayClock {
    fn new() -> Self {
        Self {
            // `g_displayClock` is zero-initialized in the C implementation,
            // so its first span always permits a progress update.
            last_update: Instant::now() - REFRESH_RATE,
        }
    }

    fn update(&mut self, display_level: c_int, level: c_int, message: fmt::Arguments<'_>) {
        if display_level < level {
            return;
        }
        if self.last_update.elapsed() > REFRESH_RATE || display_level >= 4 {
            self.last_update = Instant::now();
            eprint!("{message}");
            if display_level >= 4 {
                let _ = io::stderr().flush();
            }
        }
    }
}

fn display(display_level: c_int, level: c_int, message: fmt::Arguments<'_>) {
    if display_level >= level {
        eprint!("{message}");
    }
}

fn fatal(code: c_int, message: impl fmt::Display) -> ! {
    eprintln!("Error {code} : {message}");
    std::process::exit(code);
}

fn c_string_display(pointer: *const c_char) -> String {
    unsafe { CStr::from_ptr(pointer) }
        .to_string_lossy()
        .into_owned()
}

#[cfg(unix)]
fn path_from_c_string(value: &CStr) -> PathBuf {
    PathBuf::from(OsStr::from_bytes(value.to_bytes()))
}

#[cfg(not(unix))]
fn path_from_c_string(value: &CStr) -> PathBuf {
    PathBuf::from(value.to_string_lossy().into_owned())
}

/// Returns `-1` for a missing or non-regular file, matching
/// `DiB_getFileSize()`'s `UTIL_FILESIZE_UNKNOWN` conversion.
fn get_file_size(file_name: *const c_char) -> i64 {
    let path = path_from_c_string(unsafe { CStr::from_ptr(file_name) });
    match fs::metadata(path) {
        Ok(metadata) if metadata.is_file() => i64::try_from(metadata.len()).unwrap_or(-1),
        _ => -1,
    }
}

/// `MIN(fileSize, SAMPLESIZE_MAX)` in the original C source has different
/// usual-arithmetic-conversion behavior on 32-bit and 64-bit targets.  Keep
/// that detail here because unknown files are represented by `-1`.
fn c_min_file_size(file_size: i64) -> i64 {
    if size_of::<usize>() == 8 {
        if (file_size as u64) < SAMPLESIZE_MAX as u64 {
            file_size
        } else {
            SAMPLESIZE_MAX as i64
        }
    } else {
        min(file_size, SAMPLESIZE_MAX as i64)
    }
}

unsafe fn shuffle(file_names: *const *const c_char, nb_files: c_int) {
    if nb_files <= 1 {
        return;
    }

    // The C API takes `const char**` but deliberately shuffles the caller's
    // table in place.  The CLI supplies a writable pointer array.
    let table = unsafe { std::slice::from_raw_parts_mut(file_names.cast_mut(), nb_files as usize) };
    let mut seed = 0xFD2F_B528u32;
    for index in (1..table.len()).rev() {
        seed = seed.wrapping_mul(2_654_435_761);
        seed ^= 2_246_822_519;
        seed = seed.rotate_left(13);
        let random = (seed >> 5) as usize;
        table.swap(index, random % (index + 1));
    }
}

fn file_stats(
    file_names: *const *const c_char,
    nb_files: c_int,
    chunk_size: usize,
    display_level: c_int,
) -> FileStats {
    debug_assert!(chunk_size <= SAMPLESIZE_MAX);
    let mut stats = FileStats::default();

    for index in 0..nb_files.max(0) as usize {
        let file_name = unsafe { *file_names.add(index) };
        let file_size = get_file_size(file_name);
        if file_size == 0 {
            display(
                display_level,
                3,
                format_args!(
                    "Sample file '{}' has zero size, skipping...\n",
                    c_string_display(file_name)
                ),
            );
            continue;
        }

        if chunk_size > 0 {
            // This is the C expression `(fileSize + chunkSize - 1) /
            // chunkSize`, including its unsigned wrap for an unknown file on
            // the 64-bit targets where size_t and S64 have equal width.
            let chunks = (file_size as u64)
                .wrapping_add(chunk_size as u64)
                .wrapping_sub(1)
                / chunk_size as u64;
            stats.nb_samples = stats.nb_samples.wrapping_add(chunks as c_int);
            stats.total_size_to_load = stats.total_size_to_load.wrapping_add(file_size);
        } else {
            if file_size > SAMPLESIZE_MAX as i64 {
                if file_size > (2 * SAMPLESIZE_MAX) as i64 {
                    stats.one_sample_too_large = true;
                }
                display(
                    display_level,
                    3,
                    format_args!(
                        "Sample file '{}' is too large, limiting to {} KB\n",
                        c_string_display(file_name),
                        SAMPLESIZE_MAX / KB
                    ),
                );
            }
            stats.nb_samples = stats.nb_samples.wrapping_add(1);
            stats.total_size_to_load = stats
                .total_size_to_load
                .wrapping_add(c_min_file_size(file_size));
        }
    }

    display(
        display_level,
        4,
        format_args!(
            "Found training data {} files, {} KB, {} samples\n",
            nb_files,
            (stats.total_size_to_load / KB as i64) as c_int,
            stats.nb_samples
        ),
    );
    stats
}

fn maximum_memory() -> usize {
    if size_of::<usize>() == 4 {
        2 * GB - 64 * MB
    } else {
        (512 * MB) << size_of::<usize>()
    }
}

unsafe fn find_max_memory(required_memory: u64) -> usize {
    let step = (8 * MB) as u64;
    let mut required = (((required_memory >> 23) + 1) << 23).wrapping_add(step);
    required = min(required, maximum_memory() as u64);

    loop {
        let test_memory = unsafe { malloc(required as usize) };
        if !test_memory.is_null() {
            unsafe { free(test_memory) };
            return required as usize;
        }
        required = required.wrapping_sub(step);
    }
}

unsafe fn fill_noise(buffer: *mut u8, length: usize) {
    let mut accumulator = 2_654_435_761u32;
    for index in 0..length {
        accumulator = accumulator.wrapping_mul(2_246_822_519);
        unsafe { buffer.add(index).write((accumulator >> 21) as u8) };
    }
}

unsafe fn read_chunk(
    stream: *mut c_void,
    destination: *mut u8,
    length: usize,
    file_name: *const c_char,
) {
    let read = unsafe { fread(destination.cast(), 1, length, stream) };
    if read != length {
        fatal(
            11,
            format_args!("Pb reading {}", c_string_display(file_name)),
        );
    }
}

unsafe fn load_files(
    buffer: *mut u8,
    buffer_size: usize,
    sample_sizes: *mut usize,
    sample_count: usize,
    file_names: *const *const c_char,
    nb_files: c_int,
    target_chunk_size: usize,
    display_level: c_int,
    display_clock: &mut DisplayClock,
) -> (usize, usize) {
    debug_assert!(target_chunk_size <= SAMPLESIZE_MAX);
    let mut total_data_loaded = 0usize;
    let mut nb_samples_loaded = 0usize;
    let mut file_index = 0usize;
    let file_count = nb_files.max(0) as usize;

    while nb_samples_loaded < sample_count && file_index < file_count {
        let file_name = unsafe { *file_names.add(file_index) };
        let file_size = get_file_size(file_name);
        if file_size <= 0 {
            file_index += 1;
            continue;
        }

        let stream = unsafe { fopen(file_name, MODE_READ.as_ptr().cast()) };
        if stream.is_null() {
            let error = io::Error::last_os_error();
            fatal(
                10,
                format_args!(
                    "zstd: dictBuilder: {} {} ",
                    c_string_display(file_name),
                    error
                ),
            );
        }
        display_clock.update(
            display_level,
            2,
            format_args!("Loading {}...       \r", c_string_display(file_name)),
        );

        let first_chunk = if target_chunk_size > 0 {
            min(file_size as usize, target_chunk_size)
        } else {
            min(file_size as usize, SAMPLESIZE_MAX)
        };
        if total_data_loaded.wrapping_add(first_chunk) > buffer_size {
            unsafe { fclose(stream) };
            break;
        }

        unsafe {
            read_chunk(
                stream,
                buffer.add(total_data_loaded),
                first_chunk,
                file_name,
            );
            sample_sizes.add(nb_samples_loaded).write(first_chunk);
        }
        nb_samples_loaded += 1;
        total_data_loaded += first_chunk;

        if target_chunk_size > 0 {
            let mut file_data_loaded = first_chunk;
            while (file_data_loaded as i64) < file_size && nb_samples_loaded < sample_count {
                let remaining = (file_size - file_data_loaded as i64) as usize;
                let chunk_size = min(remaining, target_chunk_size);
                if total_data_loaded.wrapping_add(chunk_size) > buffer_size {
                    break;
                }
                unsafe {
                    read_chunk(stream, buffer.add(total_data_loaded), chunk_size, file_name);
                    sample_sizes.add(nb_samples_loaded).write(chunk_size);
                }
                nb_samples_loaded += 1;
                total_data_loaded += chunk_size;
                file_data_loaded += chunk_size;
            }
        }

        file_index += 1;
        unsafe { fclose(stream) };
    }

    display(display_level, 2, format_args!("\r{:79}\r", ""));
    display(
        display_level,
        4,
        format_args!(
            "Loaded {} KB total training data, {} nb samples \n",
            total_data_loaded / KB,
            nb_samples_loaded
        ),
    );
    (total_data_loaded, nb_samples_loaded)
}

fn loaded_size(total_size_to_load: i64, max_memory: usize) -> usize {
    let first = min(max_memory as i64, total_size_to_load);
    if size_of::<usize>() == 8 {
        min(first as u64, MAX_SAMPLES_SIZE as u64) as usize
    } else {
        min(first, MAX_SAMPLES_SIZE as i64) as usize
    }
}

unsafe fn save_dictionary(
    dictionary_file_name: *const c_char,
    dictionary: *const c_void,
    dictionary_size: usize,
) {
    let stream = unsafe { fopen(dictionary_file_name, MODE_WRITE.as_ptr().cast()) };
    if stream.is_null() {
        fatal(
            3,
            format_args!("cannot open {} ", c_string_display(dictionary_file_name)),
        );
    }

    let written = unsafe { fwrite(dictionary, 1, dictionary_size, stream) };
    if written != dictionary_size {
        fatal(
            4,
            format_args!("{} : write error", c_string_display(dictionary_file_name)),
        );
    }

    if unsafe { fclose(stream) } != 0 {
        fatal(
            5,
            format_args!("{} : flush error", c_string_display(dictionary_file_name)),
        );
    }
}

unsafe fn error_name(error_code: usize) -> String {
    let name = unsafe { ZDICT_getErrorName(error_code) };
    if name.is_null() {
        return "unknown error".to_owned();
    }
    unsafe { CStr::from_ptr(name) }
        .to_string_lossy()
        .into_owned()
}

unsafe fn free_buffers(
    source_buffer: *mut c_void,
    sample_sizes: *mut usize,
    dictionary_buffer: *mut c_void,
) {
    unsafe {
        free(source_buffer);
        free(sample_sizes.cast());
        free(dictionary_buffer);
    }
}

/// Train a dictionary from files and write it to `dictionary_file_name`.
///
/// This is the Rust implementation of the public `DiB_trainFromFiles()` ABI
/// declared in `programs/dibio.h`.  Its exit-on-I/O-error behavior is kept for
/// compatibility with the command-line program; trainer errors return `1`.
#[no_mangle]
pub unsafe extern "C" fn DiB_trainFromFiles(
    dictionary_file_name: *const c_char,
    max_dictionary_size: usize,
    file_names: *const *const c_char,
    nb_files: c_int,
    chunk_size: usize,
    legacy_params: *mut ZDICT_legacy_params_t,
    cover_params: *mut ZDICT_cover_params_t,
    fast_cover_params: *mut ZDICT_fastCover_params_t,
    optimize: c_int,
    memory_limit: c_uint,
) -> c_int {
    let display_level = if !legacy_params.is_null() {
        unsafe { (*legacy_params).zParams.notificationLevel as c_int }
    } else if !cover_params.is_null() {
        unsafe { (*cover_params).zParams.notificationLevel as c_int }
    } else if !fast_cover_params.is_null() {
        unsafe { (*fast_cover_params).zParams.notificationLevel as c_int }
    } else {
        0
    };
    let mut display_clock = DisplayClock::new();
    let dictionary_buffer = unsafe { malloc(max_dictionary_size) };

    display(display_level, 3, format_args!("Shuffling input files\n"));
    unsafe { shuffle(file_names, nb_files) };

    let stats = file_stats(file_names, nb_files, chunk_size, display_level);
    let memory_multiplier = if !legacy_params.is_null() {
        MEMMULT
    } else if !cover_params.is_null() {
        COVER_MEMMULT
    } else {
        FASTCOVER_MEMMULT
    };
    let required_memory = (stats.total_size_to_load as u64).wrapping_mul(memory_multiplier);
    let max_memory = unsafe { find_max_memory(required_memory) } / memory_multiplier as usize;
    let mut loaded_size = loaded_size(stats.total_size_to_load, max_memory);
    if memory_limit != 0 {
        display(
            display_level,
            2,
            format_args!(
                "!  Warning : setting manual memory limit for dictionary training data at {} MB \n",
                (memory_limit as usize) / MB
            ),
        );
        loaded_size = min(loaded_size, memory_limit as usize);
    }

    let source_buffer_size = loaded_size.wrapping_add(NOISELENGTH);
    let source_buffer = unsafe { malloc(source_buffer_size) };
    let sample_count = stats.nb_samples.max(0) as usize;
    let sample_sizes_bytes = sample_count.wrapping_mul(size_of::<usize>());
    let sample_sizes = unsafe { malloc(sample_sizes_bytes) }.cast::<usize>();

    if ((stats.nb_samples > 0) && sample_sizes.is_null())
        || source_buffer.is_null()
        || dictionary_buffer.is_null()
    {
        fatal(12, "not enough memory for DiB_trainFiles");
    }
    if stats.one_sample_too_large {
        display(
            display_level,
            2,
            format_args!("!  Warning : some sample(s) are very large \n"),
        );
        display(
            display_level,
            2,
            format_args!("!  Note that dictionary is only useful for small samples. \n"),
        );
        display(
            display_level,
            2,
            format_args!(
                "!  As a consequence, only the first {} bytes of each sample are loaded \n",
                SAMPLESIZE_MAX
            ),
        );
    }
    if stats.nb_samples < 5 {
        display(
            display_level,
            2,
            format_args!("!  Warning : nb of samples too low for proper processing ! \n"),
        );
        display(
            display_level,
            2,
            format_args!("!  Please provide _one file per sample_. \n"),
        );
        display(
            display_level,
            2,
            format_args!(
                "!  Alternatively, split files into fixed-size blocks representative of samples, with -B# \n"
            ),
        );
        fatal(14, "nb of samples too low");
    }
    let dictionary_size_threshold = (max_dictionary_size as i64).wrapping_mul(8);
    if stats.total_size_to_load < dictionary_size_threshold {
        display(
            display_level,
            2,
            format_args!(
                "!  Warning : data size of samples too small for target dictionary size \n"
            ),
        );
        display(
            display_level,
            2,
            format_args!("!  Samples should be about 100x larger than target dictionary size \n"),
        );
    }

    if (loaded_size as i64) < stats.total_size_to_load {
        display(
            display_level,
            1,
            format_args!(
                "Training samples set too large ({} MB); training on {} MB only...\n",
                (stats.total_size_to_load / MB as i64) as c_uint,
                loaded_size / MB
            ),
        );
    }

    let (loaded_size, nb_samples_loaded) = unsafe {
        load_files(
            source_buffer.cast(),
            loaded_size,
            sample_sizes,
            sample_count,
            file_names,
            nb_files,
            chunk_size,
            display_level,
            &mut display_clock,
        )
    };

    let mut dictionary_size = ERROR_GENERIC;
    if !legacy_params.is_null() {
        unsafe { fill_noise(source_buffer.cast::<u8>().add(loaded_size), NOISELENGTH) };
        dictionary_size = unsafe {
            ZDICT_trainFromBuffer_legacy(
                dictionary_buffer,
                max_dictionary_size,
                source_buffer,
                sample_sizes,
                nb_samples_loaded as c_uint,
                *legacy_params,
            )
        };
    } else if !cover_params.is_null() {
        if optimize != 0 {
            dictionary_size = unsafe {
                ZDICT_optimizeTrainFromBuffer_cover(
                    dictionary_buffer,
                    max_dictionary_size,
                    source_buffer,
                    sample_sizes,
                    nb_samples_loaded as c_uint,
                    cover_params,
                )
            };
            if unsafe { ZDICT_isError(dictionary_size) } == 0 {
                let parameters = unsafe { &*cover_params };
                display(
                    display_level,
                    2,
                    format_args!(
                        "k={}\nd={}\nsteps={}\nsplit={}\n",
                        parameters.k,
                        parameters.d,
                        parameters.steps,
                        (parameters.splitPoint * 100.0) as c_uint
                    ),
                );
            }
        } else {
            dictionary_size = unsafe {
                ZDICT_trainFromBuffer_cover(
                    dictionary_buffer,
                    max_dictionary_size,
                    source_buffer,
                    sample_sizes,
                    nb_samples_loaded as c_uint,
                    *cover_params,
                )
            };
        }
    } else if !fast_cover_params.is_null() {
        if optimize != 0 {
            dictionary_size = unsafe {
                ZDICT_optimizeTrainFromBuffer_fastCover(
                    dictionary_buffer,
                    max_dictionary_size,
                    source_buffer,
                    sample_sizes,
                    nb_samples_loaded as c_uint,
                    fast_cover_params,
                )
            };
            if unsafe { ZDICT_isError(dictionary_size) } == 0 {
                let parameters = unsafe { &*fast_cover_params };
                display(
                    display_level,
                    2,
                    format_args!(
                        "k={}\nd={}\nf={}\nsteps={}\nsplit={}\naccel={}\n",
                        parameters.k,
                        parameters.d,
                        parameters.f,
                        parameters.steps,
                        (parameters.splitPoint * 100.0) as c_uint,
                        parameters.accel
                    ),
                );
            }
        } else {
            dictionary_size = unsafe {
                ZDICT_trainFromBuffer_fastCover(
                    dictionary_buffer,
                    max_dictionary_size,
                    source_buffer,
                    sample_sizes,
                    nb_samples_loaded as c_uint,
                    *fast_cover_params,
                )
            };
        }
    } else {
        debug_assert!(false, "one dictionary trainer parameter set is required");
    }

    if unsafe { ZDICT_isError(dictionary_size) } != 0 {
        display(
            display_level,
            1,
            format_args!("dictionary training failed : {} \n", unsafe {
                error_name(dictionary_size)
            }),
        );
        unsafe { free_buffers(source_buffer, sample_sizes, dictionary_buffer) };
        return 1;
    }

    display(
        display_level,
        2,
        format_args!(
            "Save dictionary of size {} into file {} \n",
            dictionary_size as c_uint,
            c_string_display(dictionary_file_name)
        ),
    );
    unsafe {
        save_dictionary(dictionary_file_name, dictionary_buffer, dictionary_size);
        free_buffers(source_buffer, sample_sizes, dictionary_buffer);
    }
    0
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::ffi::CString;
    use std::fs;
    use std::mem::{offset_of, size_of};
    use std::path::{Path, PathBuf};
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[cfg(unix)]
    fn path_to_c_string(path: &Path) -> CString {
        CString::new(path.as_os_str().as_bytes()).expect("temporary path has no NUL")
    }

    #[cfg(not(unix))]
    fn path_to_c_string(path: &Path) -> CString {
        CString::new(path.to_string_lossy().as_bytes()).expect("temporary path has no NUL")
    }

    fn sample(index: usize) -> Vec<u8> {
        let line = format!(
            "record={index:02};common-key=common-value;payload=abcdefghijklmnopqrstuvwxyz\n"
        );
        let mut result = Vec::with_capacity(1024);
        while result.len() < 1024 {
            result.extend_from_slice(line.as_bytes());
        }
        result.truncate(1024);
        result
    }

    fn temporary_paths() -> (Vec<PathBuf>, PathBuf) {
        static NEXT: AtomicUsize = AtomicUsize::new(0);
        let prefix = format!(
            "zstd-dibio-{}-{}",
            std::process::id(),
            NEXT.fetch_add(1, Ordering::Relaxed)
        );
        let root = std::env::temp_dir();
        let samples = (0..5)
            .map(|index| root.join(format!("{prefix}-sample-{index}")))
            .collect::<Vec<_>>();
        let dictionary = root.join(format!("{prefix}-dictionary"));
        (samples, dictionary)
    }

    #[test]
    fn dictionary_parameter_layouts_match_zdict_h() {
        assert_eq!(size_of::<ZDICT_params_t>(), 12);
        assert_eq!(size_of::<ZDICT_legacy_params_t>(), 16);
        assert_eq!(size_of::<ZDICT_cover_params_t>(), 48);
        assert_eq!(size_of::<ZDICT_fastCover_params_t>(), 56);
        assert_eq!(offset_of!(ZDICT_cover_params_t, splitPoint), 16);
        assert_eq!(offset_of!(ZDICT_fastCover_params_t, splitPoint), 24);
    }

    #[test]
    fn trains_from_chunked_temporary_sample_files() {
        let (sample_paths, dictionary_path) = temporary_paths();
        let _cleanup = TemporaryPathCleanup {
            paths: sample_paths
                .iter()
                .cloned()
                .chain(std::iter::once(dictionary_path.clone()))
                .collect(),
        };
        for (index, path) in sample_paths.iter().enumerate() {
            fs::write(path, sample(index)).expect("write temporary training sample");
        }

        let sample_names = sample_paths
            .iter()
            .map(|path| path_to_c_string(path))
            .collect::<Vec<_>>();
        let file_names = sample_names
            .iter()
            .map(|name| name.as_ptr())
            .collect::<Vec<_>>();
        let dictionary_name = path_to_c_string(&dictionary_path);
        let mut parameters = ZDICT_legacy_params_t {
            zParams: ZDICT_params_t {
                compressionLevel: 3,
                ..ZDICT_params_t::default()
            },
            ..ZDICT_legacy_params_t::default()
        };

        let result = unsafe {
            DiB_trainFromFiles(
                dictionary_name.as_ptr(),
                2048,
                file_names.as_ptr(),
                file_names.len() as c_int,
                256,
                &mut parameters,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                0,
                0,
            )
        };
        assert_eq!(result, 0);
        let dictionary = fs::read(&dictionary_path).expect("read trained dictionary");
        assert!(!dictionary.is_empty());
    }

    #[test]
    fn legacy_trainer_abi_smoke() {
        let samples = (0..5).map(sample).collect::<Vec<_>>();
        let sample_sizes = samples.iter().map(Vec::len).collect::<Vec<_>>();
        let flat_samples = samples.concat();
        let mut dictionary = vec![0u8; 2048];
        let parameters = ZDICT_legacy_params_t::default();
        let result = unsafe {
            ZDICT_trainFromBuffer_legacy(
                dictionary.as_mut_ptr().cast(),
                dictionary.len(),
                flat_samples.as_ptr().cast(),
                sample_sizes.as_ptr(),
                sample_sizes.len() as c_uint,
                parameters,
            )
        };
        assert_eq!(unsafe { ZDICT_isError(result) }, 0);
        assert!(result > 0);
    }

    impl Drop for TemporaryPathCleanup {
        fn drop(&mut self) {
            for path in &self.paths {
                let _ = fs::remove_file(path);
            }
        }
    }

    struct TemporaryPathCleanup {
        paths: Vec<PathBuf>,
    }
}