ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(dict): port the FastCover trainer to Rust

Browse Source 
Move FastCover parameter validation, corpus preparation, optimization, and
dictionary training into Rust. Retain the C source as a static-linking ABI
anchor and register the implementation with the compression dictionary-builder
feature set.

Test Plan:
- rustfmt +nightly --check --edition 2021 rust/src/dict_builder_fastcover.rs rust/src/lib.rs
- RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo clippy --manifest-path rust/Cargo.toml --all-targets --no-default-features --features compression,dict-builder -- -D warnings
- FastCover focused Rust tests and deterministic C-reference harness
- git diff --cached --check
This commit is contained in:
ddidderr
2026-07-12 10:40:51 +02:00
parent 60877175a4
commit b8ccfd7b87
3 changed files with 1080 additions and 751 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/dict_builder_fastcover.rs
+1072
View File
@@ -0,0 +1,1072 @@
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(clippy::missing_safety_doc)]
#![allow(clippy::too_many_arguments)]
//! FastCOVER dictionary training.
//!
//! The public FastCOVER declarations live in `lib/zdict.h`. This module is a
//! Rust translation of `lib/dictBuilder/fastcover.c`; the C translation unit
//! remains in the native source lists as a declaration-only ABI anchor.
use crate::dict_builder_cover::{
COVER_computeEpochs, COVER_dictSelectionError, COVER_dictSelectionFree,
COVER_dictSelectionIsError, COVER_dictSelection_t, COVER_segment_t, COVER_selectDict,
ZDICT_cover_params_t, ZDICT_params_t,
};
use crate::errors::ERR_isError;
use crate::mem::MEM_readLE64;
use std::ffi::c_void;
use std::mem::size_of;
use std::os::raw::{c_int, c_uint};
use std::ptr;
use std::sync::atomic::{AtomicI32, Ordering as AtomicOrdering};
use std::sync::{Arc, Condvar, Mutex};
const ZDICT_DICTSIZE_MIN: usize = 256;
const FASTCOVER_MAX_SAMPLES_SIZE_64: usize = u32::MAX as usize;
const FASTCOVER_MAX_SAMPLES_SIZE_32: usize = 1usize << 30;
const FASTCOVER_MAX_F: c_uint = 31;
const FASTCOVER_MAX_ACCEL: c_uint = 10;
const FASTCOVER_DEFAULT_SPLITPOINT: f64 = 0.75;
const DEFAULT_F: c_uint = 20;
const DEFAULT_ACCEL: c_uint = 1;
const ERROR_GENERIC: usize = 0usize.wrapping_sub(1);
const ERROR_MEMORY_ALLOCATION: usize = 0usize.wrapping_sub(64);
const ERROR_PARAMETER_OUT_OF_BOUND: usize = 0usize.wrapping_sub(42);
const ERROR_SRC_SIZE_WRONG: usize = 0usize.wrapping_sub(72);
const ERROR_DST_SIZE_TOO_SMALL: usize = 0usize.wrapping_sub(70);
const PRIME6BYTES: u64 = 227_718_039_650_203;
const PRIME8BYTES: u64 = 0xCF1B_BCDC_B7A5_6463;
static DISPLAY_LEVEL: AtomicI32 = AtomicI32::new(0);
#[repr(C)]
#[derive(Clone, Copy, Default)]
/// ABI-compatible `ZDICT_fastCover_params_t` from `zdict.h`.
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,
}
#[derive(Clone, Copy)]
struct FastCoverAccel {
finalize: c_uint,
skip: c_uint,
}
const DEFAULT_ACCEL_PARAMETERS: [FastCoverAccel; 11] = [
FastCoverAccel {
finalize: 100,
skip: 0,
},
FastCoverAccel {
finalize: 100,
skip: 0,
},
FastCoverAccel {
finalize: 50,
skip: 1,
},
FastCoverAccel {
finalize: 34,
skip: 2,
},
FastCoverAccel {
finalize: 25,
skip: 3,
},
FastCoverAccel {
finalize: 20,
skip: 4,
},
FastCoverAccel {
finalize: 17,
skip: 5,
},
FastCoverAccel {
finalize: 14,
skip: 6,
},
FastCoverAccel {
finalize: 13,
skip: 7,
},
FastCoverAccel {
finalize: 11,
skip: 8,
},
FastCoverAccel {
finalize: 10,
skip: 9,
},
];
struct FastCoverContext {
samples: *const u8,
offsets: *mut usize,
samples_sizes: *const usize,
nb_samples: usize,
nb_train_samples: usize,
_nb_test_samples: usize,
nb_dmers: usize,
freqs: *mut u32,
d: c_uint,
f: c_uint,
accel_params: FastCoverAccel,
}
// The input buffer and sample sizes are immutable. Each worker only mutates
// its own frequency copy and segment-frequency table.
unsafe impl Send for FastCoverContext {}
unsafe impl Sync for FastCoverContext {}
impl Drop for FastCoverContext {
fn drop(&mut self) {
unsafe {
free_bytes(self.freqs);
free_bytes(self.offsets);
}
self.freqs = ptr::null_mut();
self.offsets = ptr::null_mut();
}
}
#[inline]
fn display_level() -> c_int {
DISPLAY_LEVEL.load(AtomicOrdering::Relaxed)
}
#[inline]
fn set_display_level(level: c_uint) {
DISPLAY_LEVEL.store(level as c_int, AtomicOrdering::Relaxed);
}
#[inline]
fn display(level: c_int, message: &str) {
if display_level() >= level {
eprint!("{message}");
}
}
#[inline]
unsafe fn free_bytes<T>(allocation: *mut T) {
unsafe { libc::free(allocation.cast::<c_void>()) };
}
#[inline]
unsafe fn malloc_bytes(size: usize) -> *mut u8 {
unsafe { libc::malloc(size) }.cast::<u8>()
}
#[inline]
unsafe fn calloc_array<T>(count: usize) -> *mut T {
unsafe { libc::calloc(count, size_of::<T>()) }.cast::<T>()
}
#[inline]
fn frequency_count(f: c_uint) -> usize {
// f is validated before this helper is used. The wrapping shift keeps
// the C unsigned-shift behavior for malformed internal callers.
1usize.wrapping_shl(f)
}
#[inline]
fn frequency_bytes(f: c_uint, element_size: usize) -> usize {
((1u64 << f) * element_size as u64) as usize
}
#[inline]
unsafe fn sample_sum(samples_sizes: *const usize, nb_samples: usize) -> usize {
let mut total = 0usize;
for index in 0..nb_samples {
total = total.wrapping_add(unsafe { *samples_sizes.add(index) });
}
total
}
#[inline]
unsafe fn hash_ptr_to_index(pointer: *const u8, f: c_uint, d: c_uint) -> usize {
let value = unsafe { MEM_readLE64(pointer.cast::<c_void>()) };
if d == 6 {
((value << 16).wrapping_mul(PRIME6BYTES) >> (64 - f)) as usize
} else {
(value.wrapping_mul(PRIME8BYTES) >> (64 - f)) as usize
}
}
fn check_parameters(
parameters: ZDICT_cover_params_t,
max_dict_size: usize,
f: c_uint,
accel: c_uint,
) -> bool {
if parameters.d == 0 || parameters.k == 0 {
return false;
}
if parameters.d != 6 && parameters.d != 8 {
return false;
}
if parameters.k as usize > max_dict_size {
return false;
}
if parameters.d > parameters.k {
return false;
}
if f > FASTCOVER_MAX_F || f == 0 {
return false;
}
if parameters.splitPoint <= 0.0 || parameters.splitPoint > 1.0 {
return false;
}
if accel > FASTCOVER_MAX_ACCEL || accel == 0 {
return false;
}
true
}
unsafe fn compute_frequency(freqs: *mut u32, context: &FastCoverContext) {
let read_length = context.d.max(8) as usize;
for sample in 0..context.nb_train_samples {
let mut start = unsafe { *context.offsets.add(sample) };
let sample_end = unsafe { *context.offsets.add(sample + 1) };
while start.wrapping_add(read_length) <= sample_end {
let index =
unsafe { hash_ptr_to_index(context.samples.add(start), context.f, context.d) };
let frequency = unsafe { freqs.add(index) };
unsafe { *frequency = (*frequency).wrapping_add(1) };
start = start.wrapping_add(context.accel_params.skip as usize + 1);
}
}
}
unsafe fn context_init(
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: usize,
d: c_uint,
split_point: f64,
f: c_uint,
accel_params: FastCoverAccel,
) -> Result<FastCoverContext, usize> {
let samples = samples_buffer.cast::<u8>();
let total_samples_size = unsafe { sample_sum(samples_sizes, nb_samples) };
let nb_train_samples = if split_point < 1.0 {
((nb_samples as f64) * split_point) as usize
} else {
nb_samples
};
let nb_test_samples = if split_point < 1.0 {
nb_samples - nb_train_samples
} else {
nb_samples
};
let training_samples_size = if split_point < 1.0 {
unsafe { sample_sum(samples_sizes, nb_train_samples) }
} else {
total_samples_size
};
let _test_samples_size = if split_point < 1.0 {
unsafe { sample_sum(samples_sizes.add(nb_train_samples), nb_test_samples) }
} else {
total_samples_size
};
let max_dmer = (d as usize).max(size_of::<u64>());
let max_samples_size = if size_of::<usize>() == 8 {
FASTCOVER_MAX_SAMPLES_SIZE_64
} else {
FASTCOVER_MAX_SAMPLES_SIZE_32
};
if total_samples_size < max_dmer || total_samples_size >= max_samples_size {
display(1, "Total samples size is too large\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
if nb_train_samples < 5 {
display(1, "Total number of training samples is invalid\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
if nb_test_samples < 1 {
display(1, "Total number of testing samples is invalid\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
let offsets_count = nb_samples.wrapping_add(1);
let offsets = unsafe { calloc_array::<usize>(offsets_count) };
if offsets.is_null() {
display(1, "Failed to allocate scratch buffers\n");
return Err(ERROR_MEMORY_ALLOCATION);
}
unsafe {
*offsets = 0;
for index in 0..nb_samples {
*offsets.add(index + 1) = (*offsets.add(index)).wrapping_add(*samples_sizes.add(index));
}
}
let freqs = unsafe { calloc_array::<u32>(frequency_count(f)) };
if freqs.is_null() {
unsafe { free_bytes(offsets) };
display(1, "Failed to allocate frequency table\n");
return Err(ERROR_MEMORY_ALLOCATION);
}
let context = FastCoverContext {
samples,
offsets,
samples_sizes,
nb_samples,
nb_train_samples,
_nb_test_samples: nb_test_samples,
nb_dmers: training_samples_size.wrapping_sub(max_dmer).wrapping_add(1),
freqs,
d,
f,
accel_params,
};
unsafe { compute_frequency(context.freqs, &context) };
Ok(context)
}
unsafe fn select_segment(
context: &FastCoverContext,
freqs: *mut u32,
begin: u32,
end: u32,
parameters: ZDICT_cover_params_t,
segment_freqs: *mut u16,
) -> COVER_segment_t {
let dmers_in_k = parameters.k - parameters.d + 1;
let mut best = COVER_segment_t::default();
let mut active = COVER_segment_t {
begin,
end: begin,
score: 0,
};
while active.end < end {
let index = unsafe {
hash_ptr_to_index(
context.samples.add(active.end as usize),
context.f,
context.d,
)
};
let active_frequency = unsafe { segment_freqs.add(index) };
if unsafe { *active_frequency } == 0 {
active.score = active.score.wrapping_add(unsafe { *freqs.add(index) });
}
active.end = active.end.wrapping_add(1);
unsafe { *active_frequency = (*active_frequency).wrapping_add(1) };
if active.end - active.begin == dmers_in_k + 1 {
let deleted_index = unsafe {
hash_ptr_to_index(
context.samples.add(active.begin as usize),
context.f,
context.d,
)
};
let deleted_frequency = unsafe { segment_freqs.add(deleted_index) };
unsafe { *deleted_frequency = (*deleted_frequency).wrapping_sub(1) };
if unsafe { *deleted_frequency } == 0 {
active.score = active
.score
.wrapping_sub(unsafe { *freqs.add(deleted_index) });
}
active.begin = active.begin.wrapping_add(1);
}
if active.score > best.score {
best = active;
}
}
while active.begin < end {
let index = unsafe {
hash_ptr_to_index(
context.samples.add(active.begin as usize),
context.f,
context.d,
)
};
let frequency = unsafe { segment_freqs.add(index) };
unsafe { *frequency = (*frequency).wrapping_sub(1) };
active.begin = active.begin.wrapping_add(1);
}
let mut position = best.begin;
while position != best.end {
let index = unsafe {
hash_ptr_to_index(context.samples.add(position as usize), context.f, context.d)
};
unsafe { *freqs.add(index) = 0 };
position = position.wrapping_add(1);
}
best
}
unsafe fn build_dictionary(
context: &FastCoverContext,
freqs: *mut u32,
dict_buffer: *mut u8,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
segment_freqs: *mut u16,
) -> usize {
let epochs = COVER_computeEpochs(
dict_buffer_capacity as u32,
context.nb_dmers as u32,
parameters.k,
1,
);
if epochs.num == 0 {
return dict_buffer_capacity;
}
let mut tail = dict_buffer_capacity;
let mut zero_score_run = 0usize;
let mut epoch = 0usize;
while tail > 0 {
let epoch_begin = epoch.wrapping_mul(epochs.size as usize) as u32;
let epoch_end = epoch_begin.wrapping_add(epochs.size);
let segment = unsafe {
select_segment(
context,
freqs,
epoch_begin,
epoch_end,
parameters,
segment_freqs,
)
};
if segment.score == 0 {
zero_score_run += 1;
if zero_score_run >= 10 {
break;
}
} else {
zero_score_run = 0;
let segment_size = (segment
.end
.wrapping_sub(segment.begin)
.wrapping_add(parameters.d)
.wrapping_sub(1) as usize)
.min(tail);
if segment_size < parameters.d as usize {
break;
}
tail -= segment_size;
unsafe {
ptr::copy_nonoverlapping(
context.samples.add(segment.begin as usize),
dict_buffer.add(tail),
segment_size,
);
}
}
epoch = (epoch + 1) % epochs.num as usize;
}
tail
}
#[derive(Clone, Copy)]
struct FastBestState {
dict: *mut u8,
dict_size: usize,
parameters: ZDICT_cover_params_t,
compressed_size: usize,
live_jobs: usize,
}
unsafe impl Send for FastBestState {}
struct FastBest {
state: Mutex<FastBestState>,
cond: Condvar,
}
impl FastBest {
fn new() -> Self {
Self {
state: Mutex::new(FastBestState {
dict: ptr::null_mut(),
dict_size: 0,
parameters: ZDICT_cover_params_t::default(),
compressed_size: ERROR_GENERIC,
live_jobs: 0,
}),
cond: Condvar::new(),
}
}
fn start(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
state.live_jobs = state.live_jobs.wrapping_add(1);
}
fn finish(&self, parameters: ZDICT_cover_params_t, selection: COVER_dictSelection_t) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
state.live_jobs = state.live_jobs.wrapping_sub(1);
if selection.totalCompressedSize < state.compressed_size {
if state.dict.is_null() || state.dict_size < selection.dictSize {
unsafe { free_bytes(state.dict) };
let replacement = unsafe { malloc_bytes(selection.dictSize) };
if replacement.is_null() && selection.dictSize != 0 {
state.dict = ptr::null_mut();
state.dict_size = 0;
state.compressed_size = ERROR_GENERIC;
self.cond.notify_one();
return;
}
state.dict = replacement;
}
if !selection.dictContent.is_null() {
unsafe {
ptr::copy_nonoverlapping(selection.dictContent, state.dict, selection.dictSize);
}
state.dict_size = selection.dictSize;
state.parameters = parameters;
state.compressed_size = selection.totalCompressedSize;
}
}
if state.live_jobs == 0 {
self.cond.notify_all();
}
}
fn wait(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
while state.live_jobs != 0 {
state = self
.cond
.wait(state)
.unwrap_or_else(|poison| poison.into_inner());
}
}
fn snapshot(&self) -> FastBestState {
*self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn destroy(&self) {
self.wait();
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
unsafe { free_bytes(state.dict) };
state.dict = ptr::null_mut();
state.dict_size = 0;
}
}
unsafe fn try_parameters(
context: &FastCoverContext,
best: &FastBest,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
) {
let count = frequency_count(context.f);
let segment_freqs = unsafe { calloc_array::<u16>(count) };
let dict = unsafe { malloc_bytes(dict_buffer_capacity) };
let freqs = unsafe { malloc_bytes(frequency_bytes(context.f, size_of::<u32>())) }.cast::<u32>();
let mut selection = COVER_dictSelectionError(ERROR_GENERIC);
if segment_freqs.is_null() || dict.is_null() || freqs.is_null() {
display(1, "Failed to allocate buffers: out of memory\n");
} else {
unsafe {
ptr::copy_nonoverlapping(context.freqs, freqs, frequency_count(context.f));
let tail = build_dictionary(
context,
freqs,
dict,
dict_buffer_capacity,
parameters,
segment_freqs,
);
selection = select_dictionary(
context,
dict,
tail,
dict_buffer_capacity,
parameters,
context.accel_params.finalize,
);
}
if COVER_dictSelectionIsError(selection) != 0 {
display(1, "Failed to select dictionary\n");
}
}
best.finish(parameters, selection);
unsafe {
free_bytes(dict);
free_bytes(segment_freqs);
free_bytes(freqs);
COVER_dictSelectionFree(selection);
}
}
unsafe fn select_dictionary(
context: &FastCoverContext,
dict: *mut u8,
tail: usize,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
finalize_percentage: c_uint,
) -> COVER_dictSelection_t {
let nb_finalize_samples = (context
.nb_train_samples
.wrapping_mul(finalize_percentage as usize)
/ 100) as c_uint;
unsafe {
COVER_selectDict(
dict.add(tail),
dict_buffer_capacity,
dict_buffer_capacity - tail,
context.samples,
context.samples_sizes,
nb_finalize_samples,
context.nb_train_samples,
context.nb_samples,
parameters,
context.offsets,
ERROR_GENERIC,
)
}
}
fn convert_to_cover_params(parameters: ZDICT_fastCover_params_t) -> ZDICT_cover_params_t {
ZDICT_cover_params_t {
k: parameters.k,
d: parameters.d,
steps: parameters.steps,
nbThreads: parameters.nbThreads,
splitPoint: parameters.splitPoint,
shrinkDict: parameters.shrinkDict,
shrinkDictMaxRegression: 0,
zParams: parameters.zParams,
}
}
fn convert_to_fastcover_params(
parameters: ZDICT_cover_params_t,
output: &mut ZDICT_fastCover_params_t,
f: c_uint,
accel: c_uint,
) {
output.k = parameters.k;
output.d = parameters.d;
output.steps = parameters.steps;
output.nbThreads = parameters.nbThreads;
output.splitPoint = parameters.splitPoint;
output.f = f;
output.accel = accel;
output.zParams = parameters.zParams;
output.shrinkDict = parameters.shrinkDict;
}
unsafe fn finalize_dictionary(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
custom_dict_content: *const c_void,
dict_content_size: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
parameters: ZDICT_params_t,
) -> usize {
unsafe {
ZDICT_finalizeDictionary(
dict_buffer,
dict_buffer_capacity,
custom_dict_content,
dict_content_size,
samples_buffer,
samples_sizes,
nb_samples,
parameters,
)
}
}
unsafe extern "C" {
fn ZDICT_finalizeDictionary(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
custom_dict_content: *const c_void,
dict_content_size: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
parameters: ZDICT_params_t,
) -> usize;
}
#[no_mangle]
pub unsafe extern "C" 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,
mut parameters: ZDICT_fastCover_params_t,
) -> usize {
set_display_level(parameters.zParams.notificationLevel);
parameters.splitPoint = 1.0;
if parameters.f == 0 {
parameters.f = DEFAULT_F;
}
if parameters.accel == 0 {
parameters.accel = DEFAULT_ACCEL;
}
let cover_parameters = convert_to_cover_params(parameters);
if !check_parameters(
cover_parameters,
dict_buffer_capacity,
parameters.f,
parameters.accel,
) {
display(1, "FASTCOVER parameters incorrect\n");
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if nb_samples == 0 {
display(1, "FASTCOVER must have at least one input file\n");
return ERROR_SRC_SIZE_WRONG;
}
if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
display(1, "dictBufferCapacity must be at least 256\n");
return ERROR_DST_SIZE_TOO_SMALL;
}
let context = match context_init(
samples_buffer,
samples_sizes,
nb_samples as usize,
cover_parameters.d,
parameters.splitPoint,
parameters.f,
DEFAULT_ACCEL_PARAMETERS[parameters.accel as usize],
) {
Ok(context) => context,
Err(error_code) => {
display(1, "Failed to initialize context\n");
return error_code;
}
};
COVER_warn_on_small_corpus(dict_buffer_capacity, context.nb_dmers, display_level());
let segment_freqs = unsafe { calloc_array::<u16>(frequency_count(parameters.f)) };
if segment_freqs.is_null() {
unsafe { free_bytes(segment_freqs) };
return ERROR_MEMORY_ALLOCATION;
}
let tail = unsafe {
build_dictionary(
&context,
context.freqs,
dict_buffer.cast::<u8>(),
dict_buffer_capacity,
cover_parameters,
segment_freqs,
)
};
let nb_finalize_samples = (context
.nb_train_samples
.wrapping_mul(context.accel_params.finalize as usize)
/ 100) as c_uint;
let dictionary_size = unsafe {
finalize_dictionary(
dict_buffer,
dict_buffer_capacity,
dict_buffer.cast::<u8>().add(tail).cast(),
dict_buffer_capacity - tail,
samples_buffer,
samples_sizes,
nb_finalize_samples,
cover_parameters.zParams,
)
};
unsafe { free_bytes(segment_freqs) };
dictionary_size
}
#[no_mangle]
pub unsafe extern "C" 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 {
let parameters = unsafe { &mut *parameters };
let nb_threads = parameters.nbThreads as usize;
let split_point = if parameters.splitPoint <= 0.0 {
FASTCOVER_DEFAULT_SPLITPOINT
} else {
parameters.splitPoint
};
let k_min_d = if parameters.d == 0 { 6 } else { parameters.d };
let k_max_d = if parameters.d == 0 { 8 } else { parameters.d };
let k_min_k = if parameters.k == 0 { 50 } else { parameters.k };
let k_max_k = if parameters.k == 0 {
2000
} else {
parameters.k
};
let k_steps = if parameters.steps == 0 {
40
} else {
parameters.steps
};
let k_step_size = (k_max_k.wrapping_sub(k_min_k) / k_steps).max(1);
let _k_iterations = (1 + (k_max_d.wrapping_sub(k_min_d) / 2))
.wrapping_mul(1 + (k_max_k.wrapping_sub(k_min_k) / k_step_size));
let f = if parameters.f == 0 {
DEFAULT_F
} else {
parameters.f
};
let accel = if parameters.accel == 0 {
DEFAULT_ACCEL
} else {
parameters.accel
};
let display_level = parameters.zParams.notificationLevel as c_int;
if split_point <= 0.0 || split_point > 1.0 {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if accel == 0 || accel > FASTCOVER_MAX_ACCEL {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if k_min_k < k_max_d || k_max_k < k_min_k {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if nb_samples == 0 {
return ERROR_SRC_SIZE_WRONG;
}
if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
return ERROR_DST_SIZE_TOO_SMALL;
}
let cover_parameters = convert_to_cover_params(*parameters);
set_display_level(if display_level == 0 {
0
} else {
(display_level - 1) as c_uint
});
let best = FastBest::new();
let mut warned = false;
let mut d = k_min_d;
loop {
if d > k_max_d {
break;
}
let context = match unsafe {
context_init(
samples_buffer,
samples_sizes,
nb_samples as usize,
d,
split_point,
f,
DEFAULT_ACCEL_PARAMETERS[accel as usize],
)
} {
Ok(context) => context,
Err(error_code) => {
best.destroy();
return error_code;
}
};
if !warned {
COVER_warn_on_small_corpus(dict_buffer_capacity, context.nb_dmers, display_level);
warned = true;
}
let mut jobs = Vec::new();
let mut k = k_min_k;
while k <= k_max_k {
let mut job = cover_parameters;
job.k = k;
job.d = d;
job.splitPoint = split_point;
job.steps = k_steps;
job.shrinkDict = 0;
job.zParams.notificationLevel = display_level.saturating_sub(1) as c_uint;
if check_parameters(job, dict_buffer_capacity, f, accel) {
if jobs.try_reserve(1).is_err() {
best.destroy();
return ERROR_MEMORY_ALLOCATION;
}
jobs.push(job);
}
let next = k.wrapping_add(k_step_size);
if next <= k {
break;
}
k = next;
}
for _ in 0..jobs.len() {
best.start();
}
if nb_threads > 1 && jobs.len() > 1 {
let jobs = Arc::new(Mutex::new(jobs.into_iter()));
let worker_count = nb_threads.min(
jobs.lock()
.unwrap_or_else(|poison| poison.into_inner())
.size_hint()
.0,
);
std::thread::scope(|scope| {
for _ in 0..worker_count {
let jobs = Arc::clone(&jobs);
let context_ref = &context;
let best_ref = &best;
scope.spawn(move || loop {
let job = jobs
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.next();
match job {
Some(job) => unsafe {
try_parameters(context_ref, best_ref, dict_buffer_capacity, job)
},
None => break,
}
});
}
});
} else {
for job in jobs {
unsafe { try_parameters(&context, &best, dict_buffer_capacity, job) };
}
}
best.wait();
let next = d.wrapping_add(2);
if next <= d {
break;
}
d = next;
}
let snapshot = best.snapshot();
if ERR_isError(snapshot.compressed_size) {
let result = snapshot.compressed_size;
best.destroy();
return result;
}
unsafe {
ptr::copy_nonoverlapping(snapshot.dict, dict_buffer.cast::<u8>(), snapshot.dict_size);
}
convert_to_fastcover_params(snapshot.parameters, parameters, f, accel);
let result = snapshot.dict_size;
best.destroy();
result
}
#[inline]
fn COVER_warn_on_small_corpus(max_dict_size: usize, nb_dmers: usize, display: c_int) {
crate::dict_builder_cover::COVER_warnOnSmallCorpus(max_dict_size, nb_dmers, display);
}
#[cfg(test)]
mod tests {
use super::*;
use std::mem::size_of;
#[test]
fn fastcover_params_have_the_c_layout() {
assert_eq!(size_of::<ZDICT_fastCover_params_t>(), 56);
assert_eq!(ZDICT_fastCover_params_t::default().k, 0);
}
#[test]
fn hash_formulas_match_the_zstd_constants() {
let bytes = [1u8, 2, 3, 4, 5, 6, 7, 8];
let value = u64::from_le_bytes(bytes);
let expected6 = ((value << 16).wrapping_mul(PRIME6BYTES) >> (64 - 9)) as usize;
let expected8 = (value.wrapping_mul(PRIME8BYTES) >> (64 - 9)) as usize;
assert_eq!(
unsafe { hash_ptr_to_index(bytes.as_ptr(), 9, 6) },
expected6
);
assert_eq!(
unsafe { hash_ptr_to_index(bytes.as_ptr(), 9, 8) },
expected8
);
}
#[test]
fn acceleration_table_matches_fastcover_reference() {
let finalize = [100, 100, 50, 34, 25, 20, 17, 14, 13, 11, 10];
let skip = [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
for index in 0..=10 {
assert_eq!(DEFAULT_ACCEL_PARAMETERS[index].finalize, finalize[index]);
assert_eq!(DEFAULT_ACCEL_PARAMETERS[index].skip, skip[index]);
}
}
#[test]
fn direct_api_keeps_error_precedence() {
let parameters = ZDICT_fastCover_params_t {
k: 50,
d: 8,
..ZDICT_fastCover_params_t::default()
};
assert_eq!(
unsafe {
ZDICT_trainFromBuffer_fastCover(
ptr::null_mut(),
0,
ptr::null(),
ptr::null(),
0,
parameters,
)
},
ERROR_PARAMETER_OUT_OF_BOUND
);
assert_eq!(
unsafe {
ZDICT_trainFromBuffer_fastCover(
ptr::null_mut(),
ZDICT_DICTSIZE_MIN - 1,
ptr::null(),
ptr::null(),
1,
parameters,
)
},
ERROR_DST_SIZE_TOO_SMALL
);
}
}
rust/src/lib.rs
+2
View File
@@ -8,6 +8,8 @@ pub mod debug;
#[cfg(feature = "dict-builder")]
pub mod dict_builder_cover;
#[cfg(all(feature = "compression", feature = "dict-builder"))]
pub mod dict_builder_fastcover;
#[cfg(all(feature = "compression", feature = "dict-builder"))]
pub mod dict_builder_zdict;
#[cfg(feature = "dict-builder")]
pub mod divsufsort;