feat(rust): port the COVER dictionary trainer
Move COVER training, frequency mapping, segment selection, optimization, dictionary shrinking, and best-candidate synchronization into Rust. The remaining C translation unit is an ABI shim, while the public ZDICT and COVER symbols stay available to the existing C dictionary-builder callers. Test Plan: - cargo test --manifest-path rust/Cargo.toml dict_builder_cover - cargo check --manifest-path rust/Cargo.toml --no-default-features --features compression,decompression,dict-builder - cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings - make -B -C lib libzstd.a V=1 - git diff --check Depends-on: Rust compression, entropy, and dictionary-finalization leaves
This commit is contained in:
@@ -0,0 +1,1307 @@
|
||||
#![allow(non_camel_case_types)]
|
||||
#![allow(non_snake_case)]
|
||||
#![allow(clippy::missing_safety_doc)]
|
||||
#![allow(clippy::too_many_arguments)]
|
||||
|
||||
//! COVER dictionary training.
|
||||
//!
|
||||
//! The public dictionary-builder ABI is declared in `lib/dictBuilder/cover.h`.
|
||||
//! This module mirrors the implementation in `cover.c`; the C translation
|
||||
//! unit is now only an ABI anchor so that the remaining C dictionary builders
|
||||
//! can continue to include that header.
|
||||
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::c_void;
|
||||
use std::mem::size_of;
|
||||
use std::os::raw::{c_int, c_uint};
|
||||
use std::ptr;
|
||||
use std::slice;
|
||||
use std::sync::{Arc, Condvar, Mutex, OnceLock};
|
||||
|
||||
const ZDICT_DICTSIZE_MIN: usize = 256;
|
||||
const COVER_DEFAULT_SPLITPOINT: f64 = 1.0;
|
||||
const MAP_EMPTY_VALUE: u32 = u32::MAX;
|
||||
const COVER_PRIME4BYTES: u32 = 2_654_435_761;
|
||||
const MAX_ERROR_CODE: usize = 120;
|
||||
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);
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub struct ZDICT_params_t {
|
||||
pub compressionLevel: c_int,
|
||||
pub notificationLevel: c_uint,
|
||||
pub dictID: c_uint,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default)]
|
||||
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,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)]
|
||||
pub struct COVER_segment_t {
|
||||
pub begin: u32,
|
||||
pub end: u32,
|
||||
pub score: u32,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)]
|
||||
pub struct COVER_epoch_info_t {
|
||||
pub num: u32,
|
||||
pub size: u32,
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub struct COVER_dictSelection_t {
|
||||
pub dictContent: *mut u8,
|
||||
pub dictSize: usize,
|
||||
pub totalCompressedSize: usize,
|
||||
}
|
||||
|
||||
unsafe impl Send for COVER_dictSelection_t {}
|
||||
unsafe impl Sync for COVER_dictSelection_t {}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy)]
|
||||
struct CoverMapPair {
|
||||
key: u32,
|
||||
value: u32,
|
||||
}
|
||||
|
||||
struct CoverMap {
|
||||
data: Vec<CoverMapPair>,
|
||||
size_log: u32,
|
||||
size_mask: u32,
|
||||
}
|
||||
|
||||
struct CoverContext {
|
||||
samples: *const u8,
|
||||
samples_sizes: *const usize,
|
||||
offsets: Vec<usize>,
|
||||
nb_samples: usize,
|
||||
nb_train_samples: usize,
|
||||
suffix_size: usize,
|
||||
freqs: Vec<u32>,
|
||||
dmer_at: Vec<u32>,
|
||||
d: usize,
|
||||
}
|
||||
|
||||
/* Samples are immutable and the scratch arrays are never changed after the
|
||||
* context has been initialized. Each worker owns its frequency copy. */
|
||||
unsafe impl Send for CoverContext {}
|
||||
unsafe impl Sync for CoverContext {}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct BestData {
|
||||
dict: *mut u8,
|
||||
dict_capacity: usize,
|
||||
dict_size: usize,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
compressed_size: usize,
|
||||
live_jobs: usize,
|
||||
}
|
||||
|
||||
unsafe impl Send for BestData {}
|
||||
|
||||
struct BestSlot {
|
||||
data: Mutex<BestData>,
|
||||
cond: Condvar,
|
||||
}
|
||||
|
||||
impl BestSlot {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
data: Mutex::new(BestData {
|
||||
dict: ptr::null_mut(),
|
||||
dict_capacity: 0,
|
||||
dict_size: 0,
|
||||
parameters: ZDICT_cover_params_t::default(),
|
||||
compressed_size: usize::MAX,
|
||||
live_jobs: 0,
|
||||
}),
|
||||
cond: Condvar::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn best_slots() -> &'static Mutex<HashMap<usize, Arc<BestSlot>>> {
|
||||
static SLOTS: OnceLock<Mutex<HashMap<usize, Arc<BestSlot>>>> = OnceLock::new();
|
||||
SLOTS.get_or_init(|| Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn is_error(value: usize) -> bool {
|
||||
value > 0usize.wrapping_sub(MAX_ERROR_CODE)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn error(code: usize) -> usize {
|
||||
0usize.wrapping_sub(code)
|
||||
}
|
||||
|
||||
fn try_vec<T>(len: usize) -> Result<Vec<T>, ()> {
|
||||
let mut result = Vec::new();
|
||||
result.try_reserve_exact(len).map_err(|_| ())?;
|
||||
unsafe { result.set_len(len) };
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn malloc_bytes(size: usize) -> *mut u8 {
|
||||
unsafe { libc::malloc(size) }.cast::<u8>()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn free_bytes(ptr: *mut u8) {
|
||||
unsafe { libc::free(ptr.cast::<c_void>()) };
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn copy_bytes(dst: *mut u8, src: *const u8, size: usize) {
|
||||
unsafe { ptr::copy_nonoverlapping(src, dst, size) };
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn samples_slice<'a>(samples: *const u8, size: usize) -> &'a [u8] {
|
||||
unsafe { slice::from_raw_parts(samples, size) }
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn compare_dmers(ctx: &CoverContext, lhs: u32, rhs: u32) -> Ordering {
|
||||
let lhs = unsafe { samples_slice(ctx.samples.add(lhs as usize), ctx.d) };
|
||||
let rhs = unsafe { samples_slice(ctx.samples.add(rhs as usize), ctx.d) };
|
||||
lhs.cmp(rhs)
|
||||
}
|
||||
|
||||
fn map_init(size: u32) -> Result<CoverMap, ()> {
|
||||
let size_log = 32 - size.leading_zeros() + 1;
|
||||
let table_size = 1usize.checked_shl(size_log).ok_or(())?;
|
||||
let mut data = try_vec::<CoverMapPair>(table_size)?;
|
||||
for pair in &mut data {
|
||||
pair.key = 0;
|
||||
pair.value = MAP_EMPTY_VALUE;
|
||||
}
|
||||
Ok(CoverMap {
|
||||
data,
|
||||
size_log,
|
||||
size_mask: table_size as u32 - 1,
|
||||
})
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn map_clear(map: &mut CoverMap) {
|
||||
for pair in &mut map.data {
|
||||
pair.value = MAP_EMPTY_VALUE;
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn map_hash(map: &CoverMap, key: u32) -> u32 {
|
||||
(key.wrapping_mul(COVER_PRIME4BYTES)) >> (32 - map.size_log)
|
||||
}
|
||||
|
||||
fn map_index(map: &CoverMap, key: u32) -> usize {
|
||||
let hash = map_hash(map, key);
|
||||
let mut index = hash;
|
||||
loop {
|
||||
let pair = &map.data[index as usize];
|
||||
if pair.value == MAP_EMPTY_VALUE || pair.key == key {
|
||||
return index as usize;
|
||||
}
|
||||
index = (index + 1) & map.size_mask;
|
||||
}
|
||||
}
|
||||
|
||||
fn map_at(map: &mut CoverMap, key: u32) -> &mut u32 {
|
||||
let index = map_index(map, key);
|
||||
let pair = &mut map.data[index];
|
||||
if pair.value == MAP_EMPTY_VALUE {
|
||||
pair.key = key;
|
||||
pair.value = 0;
|
||||
}
|
||||
&mut pair.value
|
||||
}
|
||||
|
||||
fn map_remove(map: &mut CoverMap, key: u32) {
|
||||
let mut index = map_index(map, key);
|
||||
if map.data[index].value == MAP_EMPTY_VALUE {
|
||||
return;
|
||||
}
|
||||
let mut del = index;
|
||||
let mut shift = 1u32;
|
||||
loop {
|
||||
index = (index + 1) & map.size_mask as usize;
|
||||
if map.data[index].value == MAP_EMPTY_VALUE {
|
||||
map.data[del].value = MAP_EMPTY_VALUE;
|
||||
return;
|
||||
}
|
||||
let position = &map.data[index];
|
||||
let distance = (index as u32).wrapping_sub(map_hash(map, position.key)) & map.size_mask;
|
||||
if distance >= shift {
|
||||
map.data[del] = *position;
|
||||
del = index;
|
||||
shift = 1;
|
||||
} else {
|
||||
shift += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn lower_bound(values: &[usize], value: usize) -> usize {
|
||||
let mut first = 0usize;
|
||||
let mut count = values.len();
|
||||
while count != 0 {
|
||||
let step = count / 2;
|
||||
let index = first + step;
|
||||
if values[index] < value {
|
||||
first = index + 1;
|
||||
count -= step + 1;
|
||||
} else {
|
||||
count = step;
|
||||
}
|
||||
}
|
||||
first
|
||||
}
|
||||
|
||||
fn context_init(
|
||||
samples: *const u8,
|
||||
samples_sizes: *const usize,
|
||||
nb_samples: usize,
|
||||
d: usize,
|
||||
split_point: f64,
|
||||
) -> Result<CoverContext, usize> {
|
||||
if nb_samples == 0 || samples_sizes.is_null() {
|
||||
return Err(ERROR_SRC_SIZE_WRONG);
|
||||
}
|
||||
let sizes = unsafe { slice::from_raw_parts(samples_sizes, nb_samples) };
|
||||
let total_samples_size = sizes
|
||||
.iter()
|
||||
.fold(0usize, |sum, size| sum.wrapping_add(*size));
|
||||
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.saturating_sub(nb_train_samples)
|
||||
} else {
|
||||
nb_samples
|
||||
};
|
||||
let training_samples_size = if split_point < 1.0 {
|
||||
sizes[..nb_train_samples]
|
||||
.iter()
|
||||
.fold(0usize, |sum, size| sum.wrapping_add(*size))
|
||||
} else {
|
||||
total_samples_size
|
||||
};
|
||||
let _test_samples_size = if split_point < 1.0 {
|
||||
sizes[nb_train_samples..]
|
||||
.iter()
|
||||
.fold(0usize, |sum, size| sum.wrapping_add(*size))
|
||||
} else {
|
||||
total_samples_size
|
||||
};
|
||||
let max_dmer = d.max(size_of::<u64>());
|
||||
let max_samples_size = if size_of::<usize>() == 8 {
|
||||
usize::MAX
|
||||
} else {
|
||||
1usize << 30
|
||||
};
|
||||
if total_samples_size < max_dmer
|
||||
|| total_samples_size >= max_samples_size
|
||||
|| nb_train_samples < 5
|
||||
|| nb_test_samples < 1
|
||||
|| training_samples_size < max_dmer
|
||||
{
|
||||
return Err(ERROR_SRC_SIZE_WRONG);
|
||||
}
|
||||
let suffix_size = training_samples_size - max_dmer + 1;
|
||||
let mut suffix = try_vec::<u32>(suffix_size).map_err(|_| ERROR_MEMORY_ALLOCATION)?;
|
||||
let dmer_at = try_vec::<u32>(suffix_size).map_err(|_| ERROR_MEMORY_ALLOCATION)?;
|
||||
let mut offsets = try_vec::<usize>(nb_samples + 1).map_err(|_| ERROR_MEMORY_ALLOCATION)?;
|
||||
offsets[0] = 0;
|
||||
for (index, size) in sizes.iter().enumerate() {
|
||||
offsets[index + 1] = offsets[index].wrapping_add(*size);
|
||||
}
|
||||
for (index, value) in suffix.iter_mut().enumerate() {
|
||||
*value = index as u32;
|
||||
}
|
||||
let mut context = CoverContext {
|
||||
samples,
|
||||
samples_sizes,
|
||||
offsets,
|
||||
nb_samples,
|
||||
nb_train_samples,
|
||||
suffix_size,
|
||||
freqs: Vec::new(),
|
||||
dmer_at,
|
||||
d,
|
||||
};
|
||||
suffix.sort_by(|lhs, rhs| compare_dmers(&context, *lhs, *rhs));
|
||||
|
||||
let mut group_begin = 0usize;
|
||||
while group_begin < suffix_size {
|
||||
let mut group_end = group_begin + 1;
|
||||
while group_end < suffix_size
|
||||
&& compare_dmers(&context, suffix[group_begin], suffix[group_end]) == Ordering::Equal
|
||||
{
|
||||
group_end += 1;
|
||||
}
|
||||
let dmer_id = group_begin as u32;
|
||||
let mut frequency = 0u32;
|
||||
let mut current_offset_index = 0usize;
|
||||
let mut current_sample_end = context.offsets[0];
|
||||
let group = &suffix[group_begin..group_end];
|
||||
for (relative_index, &position_raw) in group.iter().enumerate() {
|
||||
let position = position_raw as usize;
|
||||
context.dmer_at[position] = dmer_id;
|
||||
if position < current_sample_end {
|
||||
continue;
|
||||
}
|
||||
frequency = frequency.wrapping_add(1);
|
||||
if relative_index + 1 != group.len() {
|
||||
let search_end = context.nb_samples;
|
||||
let relative =
|
||||
lower_bound(&context.offsets[current_offset_index..search_end], position);
|
||||
let sample_end_index = current_offset_index + relative;
|
||||
current_sample_end = context.offsets[sample_end_index.min(context.nb_samples)];
|
||||
current_offset_index = (sample_end_index + 1).min(context.nb_samples);
|
||||
}
|
||||
}
|
||||
suffix[group_begin] = frequency;
|
||||
group_begin = group_end;
|
||||
}
|
||||
context.freqs = suffix;
|
||||
Ok(context)
|
||||
}
|
||||
|
||||
fn select_segment(
|
||||
ctx: &CoverContext,
|
||||
freqs: &mut [u32],
|
||||
active_dmers: &mut CoverMap,
|
||||
begin: u32,
|
||||
end: u32,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
) -> 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,
|
||||
};
|
||||
map_clear(active_dmers);
|
||||
while active.end < end {
|
||||
let new_dmer = ctx.dmer_at[active.end as usize];
|
||||
let occurrence = map_at(active_dmers, new_dmer);
|
||||
if *occurrence == 0 {
|
||||
active.score = active.score.wrapping_add(freqs[new_dmer as usize]);
|
||||
}
|
||||
active.end += 1;
|
||||
*occurrence += 1;
|
||||
if active.end - active.begin == dmers_in_k + 1 {
|
||||
let deleted_dmer = ctx.dmer_at[active.begin as usize];
|
||||
let deleted_occurrence = map_at(active_dmers, deleted_dmer);
|
||||
active.begin += 1;
|
||||
*deleted_occurrence -= 1;
|
||||
if *deleted_occurrence == 0 {
|
||||
map_remove(active_dmers, deleted_dmer);
|
||||
active.score = active.score.wrapping_sub(freqs[deleted_dmer as usize]);
|
||||
}
|
||||
}
|
||||
if active.score > best.score {
|
||||
best = active;
|
||||
}
|
||||
}
|
||||
let mut new_begin = best.end;
|
||||
let mut new_end = best.begin;
|
||||
for position in best.begin..best.end {
|
||||
let frequency = freqs[ctx.dmer_at[position as usize] as usize];
|
||||
if frequency != 0 {
|
||||
new_begin = new_begin.min(position);
|
||||
new_end = position + 1;
|
||||
}
|
||||
}
|
||||
best.begin = new_begin;
|
||||
best.end = new_end;
|
||||
for position in best.begin..best.end {
|
||||
freqs[ctx.dmer_at[position as usize] as usize] = 0;
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
fn build_dictionary(
|
||||
ctx: &CoverContext,
|
||||
freqs: &mut [u32],
|
||||
active_dmers: &mut CoverMap,
|
||||
dict: *mut u8,
|
||||
dict_capacity: usize,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
) -> usize {
|
||||
let epochs = COVER_computeEpochs(
|
||||
dict_capacity.min(u32::MAX as usize) as u32,
|
||||
ctx.suffix_size.min(u32::MAX as usize) as u32,
|
||||
parameters.k,
|
||||
4,
|
||||
);
|
||||
if epochs.num == 0 || epochs.size == 0 {
|
||||
return dict_capacity;
|
||||
}
|
||||
let max_zero_score_run = 10usize.max((epochs.num as usize >> 3).min(100));
|
||||
let mut zero_score_run = 0usize;
|
||||
let mut tail = dict_capacity;
|
||||
let mut epoch = 0usize;
|
||||
while tail > 0 {
|
||||
let epoch_begin = (epoch * epochs.size as usize).min(ctx.suffix_size) as u32;
|
||||
let epoch_end = (epoch_begin as usize + epochs.size as usize).min(ctx.suffix_size) as u32;
|
||||
let segment = select_segment(ctx, freqs, active_dmers, epoch_begin, epoch_end, parameters);
|
||||
if segment.score == 0 {
|
||||
zero_score_run += 1;
|
||||
if zero_score_run >= max_zero_score_run {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
zero_score_run = 0;
|
||||
let segment_size = ((segment.end - segment.begin) as usize)
|
||||
.saturating_add(parameters.d as usize)
|
||||
.saturating_sub(1)
|
||||
.min(tail);
|
||||
if segment_size < parameters.d as usize {
|
||||
break;
|
||||
}
|
||||
tail -= segment_size;
|
||||
unsafe {
|
||||
copy_bytes(
|
||||
dict.add(tail),
|
||||
ctx.samples.add(segment.begin as usize),
|
||||
segment_size,
|
||||
);
|
||||
}
|
||||
}
|
||||
epoch = (epoch + 1) % epochs.num as usize;
|
||||
}
|
||||
tail
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn slot_for(best: *mut c_void) -> Option<Arc<BestSlot>> {
|
||||
if best.is_null() {
|
||||
return None;
|
||||
}
|
||||
let slots = best_slots()
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
slots.get(&(best as usize)).cloned()
|
||||
}
|
||||
|
||||
fn best_start_slot(slot: &BestSlot) {
|
||||
let mut data = slot
|
||||
.data
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
data.live_jobs = data.live_jobs.wrapping_add(1);
|
||||
}
|
||||
|
||||
fn best_wait_slot(slot: &BestSlot) {
|
||||
let mut data = slot
|
||||
.data
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
while data.live_jobs != 0 {
|
||||
data = slot
|
||||
.cond
|
||||
.wait(data)
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
}
|
||||
}
|
||||
|
||||
fn best_finish_slot(
|
||||
slot: &BestSlot,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
selection: COVER_dictSelection_t,
|
||||
) {
|
||||
let mut data = slot
|
||||
.data
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
data.live_jobs = data.live_jobs.wrapping_sub(1);
|
||||
if selection.totalCompressedSize < data.compressed_size {
|
||||
if data.dict.is_null() || data.dict_capacity < selection.dictSize {
|
||||
let replacement = unsafe { malloc_bytes(selection.dictSize) };
|
||||
if replacement.is_null() && selection.dictSize != 0 {
|
||||
if !data.dict.is_null() {
|
||||
unsafe { free_bytes(data.dict) };
|
||||
}
|
||||
data.dict = ptr::null_mut();
|
||||
data.dict_capacity = 0;
|
||||
data.dict_size = 0;
|
||||
data.compressed_size = ERROR_GENERIC;
|
||||
slot.cond.notify_one();
|
||||
return;
|
||||
}
|
||||
if !data.dict.is_null() {
|
||||
unsafe { free_bytes(data.dict) };
|
||||
}
|
||||
data.dict = replacement;
|
||||
data.dict_capacity = selection.dictSize;
|
||||
}
|
||||
if !selection.dictContent.is_null() {
|
||||
unsafe { copy_bytes(data.dict, selection.dictContent, selection.dictSize) };
|
||||
data.dict_size = selection.dictSize;
|
||||
data.parameters = parameters;
|
||||
data.compressed_size = selection.totalCompressedSize;
|
||||
}
|
||||
}
|
||||
if data.live_jobs == 0 {
|
||||
slot.cond.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct BestSnapshot {
|
||||
dict: *mut u8,
|
||||
dict_size: usize,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
compressed_size: usize,
|
||||
}
|
||||
|
||||
unsafe impl Send for BestSnapshot {}
|
||||
|
||||
fn best_snapshot(slot: &BestSlot) -> BestSnapshot {
|
||||
let data = slot
|
||||
.data
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
BestSnapshot {
|
||||
dict: data.dict,
|
||||
dict_size: data.dict_size,
|
||||
parameters: data.parameters,
|
||||
compressed_size: data.compressed_size,
|
||||
}
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
struct ZSTD_CCtx {
|
||||
_private: [u8; 0],
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
struct ZSTD_CDict {
|
||||
_private: [u8; 0],
|
||||
}
|
||||
|
||||
unsafe extern "C" {
|
||||
fn ZDICT_finalizeDictionary(
|
||||
dst_dict_buffer: *mut c_void,
|
||||
max_dict_size: usize,
|
||||
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;
|
||||
fn ZSTD_compressBound(src_size: usize) -> usize;
|
||||
fn ZSTD_createCCtx() -> *mut ZSTD_CCtx;
|
||||
fn ZSTD_freeCCtx(cctx: *mut ZSTD_CCtx) -> usize;
|
||||
fn ZSTD_createCDict(
|
||||
dict: *const c_void,
|
||||
dict_size: usize,
|
||||
compression_level: c_int,
|
||||
) -> *mut ZSTD_CDict;
|
||||
fn ZSTD_freeCDict(cdict: *mut ZSTD_CDict) -> usize;
|
||||
fn ZSTD_compress_usingCDict(
|
||||
cctx: *mut ZSTD_CCtx,
|
||||
dst: *mut c_void,
|
||||
dst_capacity: usize,
|
||||
src: *const c_void,
|
||||
src_size: usize,
|
||||
cdict: *const ZSTD_CDict,
|
||||
) -> usize;
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_sum(samples_sizes: *const usize, nb_samples: c_uint) -> usize {
|
||||
if nb_samples == 0 {
|
||||
return 0;
|
||||
}
|
||||
if samples_sizes.is_null() {
|
||||
return 0;
|
||||
}
|
||||
unsafe {
|
||||
slice::from_raw_parts(samples_sizes, nb_samples as usize)
|
||||
.iter()
|
||||
.fold(0usize, |sum, size| sum.wrapping_add(*size))
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub extern "C" fn COVER_computeEpochs(
|
||||
max_dict_size: u32,
|
||||
nb_dmers: u32,
|
||||
k: u32,
|
||||
passes: u32,
|
||||
) -> COVER_epoch_info_t {
|
||||
let min_epoch_size = k.saturating_mul(10);
|
||||
let k = k.max(1);
|
||||
let passes = passes.max(1);
|
||||
let mut epochs = COVER_epoch_info_t {
|
||||
num: (max_dict_size / k / passes).max(1),
|
||||
size: 0,
|
||||
};
|
||||
epochs.size = nb_dmers / epochs.num;
|
||||
if epochs.size < min_epoch_size {
|
||||
epochs.size = min_epoch_size.min(nb_dmers);
|
||||
epochs.num = if epochs.size == 0 {
|
||||
1
|
||||
} else {
|
||||
nb_dmers.checked_div(epochs.size).unwrap_or(1).max(1)
|
||||
};
|
||||
}
|
||||
epochs
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub extern "C" fn COVER_warnOnSmallCorpus(
|
||||
max_dict_size: usize,
|
||||
nb_dmers: usize,
|
||||
display_level: c_int,
|
||||
) {
|
||||
let ratio = nb_dmers as f64 / max_dict_size as f64;
|
||||
if ratio < 10.0 && display_level >= 1 {
|
||||
eprintln!(
|
||||
"WARNING: The maximum dictionary size {} is too large compared to the source size {}! size(source)/size(dictionary) = {}, but it should be >= 10! This may lead to a subpar dictionary! We recommend training on sources at least 10x, and preferably 100x the size of the dictionary! ",
|
||||
max_dict_size, nb_dmers, ratio
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" 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,
|
||||
mut parameters: ZDICT_cover_params_t,
|
||||
) -> usize {
|
||||
parameters.splitPoint = 1.0;
|
||||
if parameters.k == 0
|
||||
|| parameters.d == 0
|
||||
|| parameters.k as usize > dict_buffer_capacity
|
||||
|| parameters.d > parameters.k
|
||||
|| parameters.splitPoint <= 0.0
|
||||
|| parameters.splitPoint > 1.0
|
||||
{
|
||||
return error(42);
|
||||
}
|
||||
if nb_samples == 0 {
|
||||
return error(72);
|
||||
}
|
||||
if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
|
||||
return error(70);
|
||||
}
|
||||
let context = match context_init(
|
||||
samples_buffer.cast(),
|
||||
samples_sizes,
|
||||
nb_samples as usize,
|
||||
parameters.d as usize,
|
||||
parameters.splitPoint,
|
||||
) {
|
||||
Ok(context) => context,
|
||||
Err(error_code) => return error_code,
|
||||
};
|
||||
let mut active_dmers = match map_init(parameters.k - parameters.d + 1) {
|
||||
Ok(map) => map,
|
||||
Err(()) => return ERROR_MEMORY_ALLOCATION,
|
||||
};
|
||||
let mut freqs = match context.freqs.clone().try_reserve(0) {
|
||||
Ok(()) => context.freqs.clone(),
|
||||
Err(_) => return ERROR_MEMORY_ALLOCATION,
|
||||
};
|
||||
let tail = build_dictionary(
|
||||
&context,
|
||||
&mut freqs,
|
||||
&mut active_dmers,
|
||||
dict_buffer.cast(),
|
||||
dict_buffer_capacity,
|
||||
parameters,
|
||||
);
|
||||
unsafe {
|
||||
ZDICT_finalizeDictionary(
|
||||
dict_buffer,
|
||||
dict_buffer_capacity,
|
||||
dict_buffer.cast::<u8>().add(tail).cast(),
|
||||
dict_buffer_capacity - tail,
|
||||
samples_buffer,
|
||||
samples_sizes,
|
||||
nb_samples,
|
||||
parameters.zParams,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_checkTotalCompressedSize(
|
||||
parameters: ZDICT_cover_params_t,
|
||||
samples_sizes: *const usize,
|
||||
samples: *const u8,
|
||||
offsets: *mut usize,
|
||||
nb_train_samples: usize,
|
||||
nb_samples: usize,
|
||||
dict: *mut u8,
|
||||
dict_buffer_capacity: usize,
|
||||
) -> usize {
|
||||
if samples_sizes.is_null() || samples.is_null() || offsets.is_null() {
|
||||
return ERROR_GENERIC;
|
||||
}
|
||||
let sizes = unsafe { slice::from_raw_parts(samples_sizes, nb_samples) };
|
||||
let offset_values = unsafe { slice::from_raw_parts(offsets, nb_samples + 1) };
|
||||
let start = if parameters.splitPoint < 1.0 {
|
||||
nb_train_samples
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let mut max_sample_size = 0usize;
|
||||
for size in sizes.iter().skip(start) {
|
||||
max_sample_size = max_sample_size.max(*size);
|
||||
}
|
||||
let dst_capacity = unsafe { ZSTD_compressBound(max_sample_size) };
|
||||
if is_error(dst_capacity) {
|
||||
return dst_capacity;
|
||||
}
|
||||
let dst = unsafe { malloc_bytes(dst_capacity) };
|
||||
let cctx = unsafe { ZSTD_createCCtx() };
|
||||
let cdict = unsafe {
|
||||
ZSTD_createCDict(
|
||||
dict.cast(),
|
||||
dict_buffer_capacity,
|
||||
parameters.zParams.compressionLevel,
|
||||
)
|
||||
};
|
||||
if dst.is_null() || cctx.is_null() || cdict.is_null() {
|
||||
if !cctx.is_null() {
|
||||
unsafe { ZSTD_freeCCtx(cctx) };
|
||||
}
|
||||
if !cdict.is_null() {
|
||||
unsafe { ZSTD_freeCDict(cdict) };
|
||||
}
|
||||
if !dst.is_null() {
|
||||
unsafe { free_bytes(dst) };
|
||||
}
|
||||
return ERROR_GENERIC;
|
||||
}
|
||||
let mut total = dict_buffer_capacity;
|
||||
for index in start..nb_samples {
|
||||
let result = unsafe {
|
||||
ZSTD_compress_usingCDict(
|
||||
cctx,
|
||||
dst.cast(),
|
||||
dst_capacity,
|
||||
samples.add(offset_values[index]).cast(),
|
||||
sizes[index],
|
||||
cdict,
|
||||
)
|
||||
};
|
||||
if is_error(result) {
|
||||
total = result;
|
||||
break;
|
||||
}
|
||||
total = total.wrapping_add(result);
|
||||
}
|
||||
unsafe {
|
||||
ZSTD_freeCCtx(cctx);
|
||||
ZSTD_freeCDict(cdict);
|
||||
free_bytes(dst);
|
||||
}
|
||||
total
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_best_init(best: *mut c_void) {
|
||||
if best.is_null() {
|
||||
return;
|
||||
}
|
||||
let slot = Arc::new(BestSlot::new());
|
||||
let mut slots = best_slots()
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
slots.insert(best as usize, slot);
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_best_wait(best: *mut c_void) {
|
||||
if let Some(slot) = slot_for(best) {
|
||||
best_wait_slot(&slot);
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_best_destroy(best: *mut c_void) {
|
||||
if best.is_null() {
|
||||
return;
|
||||
}
|
||||
let slot = {
|
||||
let slots = best_slots()
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
slots.get(&(best as usize)).cloned()
|
||||
};
|
||||
if let Some(slot) = slot {
|
||||
best_wait_slot(&slot);
|
||||
let mut slots = best_slots()
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
slots.remove(&(best as usize));
|
||||
let snapshot = best_snapshot(&slot);
|
||||
if !snapshot.dict.is_null() {
|
||||
unsafe { free_bytes(snapshot.dict) };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_best_start(best: *mut c_void) {
|
||||
if let Some(slot) = slot_for(best) {
|
||||
best_start_slot(&slot);
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_best_finish(
|
||||
best: *mut c_void,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
selection: COVER_dictSelection_t,
|
||||
) {
|
||||
if let Some(slot) = slot_for(best) {
|
||||
best_finish_slot(&slot, parameters, selection);
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub extern "C" fn COVER_dictSelectionError(error_code: usize) -> COVER_dictSelection_t {
|
||||
COVER_dictSelection_t {
|
||||
dictContent: ptr::null_mut(),
|
||||
dictSize: 0,
|
||||
totalCompressedSize: error_code,
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub extern "C" fn COVER_dictSelectionIsError(selection: COVER_dictSelection_t) -> c_uint {
|
||||
(is_error(selection.totalCompressedSize) || selection.dictContent.is_null()) as c_uint
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_dictSelectionFree(selection: COVER_dictSelection_t) {
|
||||
if !selection.dictContent.is_null() {
|
||||
unsafe { free_bytes(selection.dictContent) };
|
||||
}
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" fn COVER_selectDict(
|
||||
custom_dict_content: *mut u8,
|
||||
dict_buffer_capacity: usize,
|
||||
mut dict_content_size: usize,
|
||||
samples_buffer: *const u8,
|
||||
samples_sizes: *const usize,
|
||||
nb_finalize_samples: c_uint,
|
||||
nb_check_samples: usize,
|
||||
nb_samples: usize,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
offsets: *mut usize,
|
||||
_total_compressed_size: usize,
|
||||
) -> COVER_dictSelection_t {
|
||||
let mut total_compressed_size: usize;
|
||||
let custom_dict_content_size = dict_content_size;
|
||||
let custom_dict_end = custom_dict_content.wrapping_add(custom_dict_content_size);
|
||||
let largest_dict_buffer = unsafe { malloc_bytes(dict_buffer_capacity) };
|
||||
let candidate_dict_buffer = unsafe { malloc_bytes(dict_buffer_capacity) };
|
||||
if (largest_dict_buffer.is_null() || candidate_dict_buffer.is_null())
|
||||
&& dict_buffer_capacity != 0
|
||||
{
|
||||
if !largest_dict_buffer.is_null() {
|
||||
unsafe { free_bytes(largest_dict_buffer) };
|
||||
}
|
||||
if !candidate_dict_buffer.is_null() {
|
||||
unsafe { free_bytes(candidate_dict_buffer) };
|
||||
}
|
||||
return COVER_dictSelectionError(dict_content_size);
|
||||
}
|
||||
unsafe {
|
||||
copy_bytes(largest_dict_buffer, custom_dict_content, dict_content_size);
|
||||
}
|
||||
dict_content_size = unsafe {
|
||||
ZDICT_finalizeDictionary(
|
||||
largest_dict_buffer.cast(),
|
||||
dict_buffer_capacity,
|
||||
custom_dict_content.cast(),
|
||||
dict_content_size,
|
||||
samples_buffer.cast(),
|
||||
samples_sizes,
|
||||
nb_finalize_samples,
|
||||
parameters.zParams,
|
||||
)
|
||||
};
|
||||
if is_error(dict_content_size) {
|
||||
unsafe {
|
||||
free_bytes(largest_dict_buffer);
|
||||
free_bytes(candidate_dict_buffer);
|
||||
}
|
||||
return COVER_dictSelectionError(dict_content_size);
|
||||
}
|
||||
total_compressed_size = unsafe {
|
||||
COVER_checkTotalCompressedSize(
|
||||
parameters,
|
||||
samples_sizes,
|
||||
samples_buffer,
|
||||
offsets,
|
||||
nb_check_samples,
|
||||
nb_samples,
|
||||
largest_dict_buffer,
|
||||
dict_content_size,
|
||||
)
|
||||
};
|
||||
if is_error(total_compressed_size) {
|
||||
unsafe {
|
||||
free_bytes(largest_dict_buffer);
|
||||
free_bytes(candidate_dict_buffer);
|
||||
}
|
||||
return COVER_dictSelectionError(total_compressed_size);
|
||||
}
|
||||
if parameters.shrinkDict == 0 {
|
||||
unsafe { free_bytes(candidate_dict_buffer) };
|
||||
return COVER_dictSelection(
|
||||
largest_dict_buffer,
|
||||
dict_content_size,
|
||||
total_compressed_size,
|
||||
);
|
||||
}
|
||||
|
||||
let largest_dict = dict_content_size;
|
||||
let largest_compressed = total_compressed_size;
|
||||
let regression_tolerance = parameters.shrinkDictMaxRegression as f64 / 100.0 + 1.0;
|
||||
dict_content_size = ZDICT_DICTSIZE_MIN;
|
||||
while dict_content_size < largest_dict {
|
||||
unsafe {
|
||||
copy_bytes(candidate_dict_buffer, largest_dict_buffer, largest_dict);
|
||||
let content = custom_dict_end.sub(dict_content_size);
|
||||
dict_content_size = ZDICT_finalizeDictionary(
|
||||
candidate_dict_buffer.cast(),
|
||||
dict_buffer_capacity,
|
||||
content.cast(),
|
||||
dict_content_size,
|
||||
samples_buffer.cast(),
|
||||
samples_sizes,
|
||||
nb_finalize_samples,
|
||||
parameters.zParams,
|
||||
);
|
||||
}
|
||||
if is_error(dict_content_size) {
|
||||
unsafe {
|
||||
free_bytes(largest_dict_buffer);
|
||||
free_bytes(candidate_dict_buffer);
|
||||
}
|
||||
return COVER_dictSelectionError(dict_content_size);
|
||||
}
|
||||
total_compressed_size = unsafe {
|
||||
COVER_checkTotalCompressedSize(
|
||||
parameters,
|
||||
samples_sizes,
|
||||
samples_buffer,
|
||||
offsets,
|
||||
nb_check_samples,
|
||||
nb_samples,
|
||||
candidate_dict_buffer,
|
||||
dict_content_size,
|
||||
)
|
||||
};
|
||||
if is_error(total_compressed_size) {
|
||||
unsafe {
|
||||
free_bytes(largest_dict_buffer);
|
||||
free_bytes(candidate_dict_buffer);
|
||||
}
|
||||
return COVER_dictSelectionError(total_compressed_size);
|
||||
}
|
||||
if (total_compressed_size as f64) <= (largest_compressed as f64) * regression_tolerance {
|
||||
unsafe { free_bytes(largest_dict_buffer) };
|
||||
return COVER_dictSelection(
|
||||
candidate_dict_buffer,
|
||||
dict_content_size,
|
||||
total_compressed_size,
|
||||
);
|
||||
}
|
||||
dict_content_size = dict_content_size.saturating_mul(2);
|
||||
}
|
||||
unsafe { free_bytes(candidate_dict_buffer) };
|
||||
COVER_dictSelection(largest_dict_buffer, largest_dict, largest_compressed)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn COVER_dictSelection(
|
||||
dict_content: *mut u8,
|
||||
dict_size: usize,
|
||||
total_compressed_size: usize,
|
||||
) -> COVER_dictSelection_t {
|
||||
COVER_dictSelection_t {
|
||||
dictContent: dict_content,
|
||||
dictSize: dict_size,
|
||||
totalCompressedSize: total_compressed_size,
|
||||
}
|
||||
}
|
||||
|
||||
fn try_parameters(
|
||||
ctx: &CoverContext,
|
||||
slot: &BestSlot,
|
||||
dict_buffer_capacity: usize,
|
||||
parameters: ZDICT_cover_params_t,
|
||||
) {
|
||||
let mut active_dmers = match map_init(parameters.k - parameters.d + 1) {
|
||||
Ok(map) => map,
|
||||
Err(()) => {
|
||||
best_finish_slot(slot, parameters, COVER_dictSelectionError(ERROR_GENERIC));
|
||||
return;
|
||||
}
|
||||
};
|
||||
let dict = unsafe { malloc_bytes(dict_buffer_capacity) };
|
||||
let mut freqs = match ctx.freqs.clone().try_reserve(0) {
|
||||
Ok(()) => ctx.freqs.clone(),
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
if dict.is_null() || freqs.len() != ctx.freqs.len() {
|
||||
if !dict.is_null() {
|
||||
unsafe { free_bytes(dict) };
|
||||
}
|
||||
best_finish_slot(slot, parameters, COVER_dictSelectionError(ERROR_GENERIC));
|
||||
return;
|
||||
}
|
||||
let tail = build_dictionary(
|
||||
ctx,
|
||||
&mut freqs,
|
||||
&mut active_dmers,
|
||||
dict,
|
||||
dict_buffer_capacity,
|
||||
parameters,
|
||||
);
|
||||
let selection = unsafe {
|
||||
COVER_selectDict(
|
||||
dict.add(tail),
|
||||
dict_buffer_capacity,
|
||||
dict_buffer_capacity - tail,
|
||||
ctx.samples,
|
||||
ctx.samples_sizes,
|
||||
ctx.nb_train_samples as c_uint,
|
||||
ctx.nb_train_samples,
|
||||
ctx.nb_samples,
|
||||
parameters,
|
||||
ctx.offsets.as_ptr().cast_mut(),
|
||||
ERROR_GENERIC,
|
||||
)
|
||||
};
|
||||
unsafe { free_bytes(dict) };
|
||||
best_finish_slot(slot, parameters, selection);
|
||||
unsafe { COVER_dictSelectionFree(selection) };
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
pub unsafe extern "C" 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 {
|
||||
let parameters_ref = unsafe { &mut *parameters };
|
||||
let nb_threads = parameters_ref.nbThreads as usize;
|
||||
let split_point = if parameters_ref.splitPoint <= 0.0 {
|
||||
COVER_DEFAULT_SPLITPOINT
|
||||
} else {
|
||||
parameters_ref.splitPoint
|
||||
};
|
||||
let k_min_d = if parameters_ref.d == 0 {
|
||||
6
|
||||
} else {
|
||||
parameters_ref.d
|
||||
};
|
||||
let k_max_d = if parameters_ref.d == 0 {
|
||||
8
|
||||
} else {
|
||||
parameters_ref.d
|
||||
};
|
||||
let k_min_k = if parameters_ref.k == 0 {
|
||||
50
|
||||
} else {
|
||||
parameters_ref.k
|
||||
};
|
||||
let k_max_k = if parameters_ref.k == 0 {
|
||||
2000
|
||||
} else {
|
||||
parameters_ref.k
|
||||
};
|
||||
let k_steps = if parameters_ref.steps == 0 {
|
||||
40
|
||||
} else {
|
||||
parameters_ref.steps
|
||||
};
|
||||
if split_point <= 0.0
|
||||
|| split_point > 1.0
|
||||
|| k_min_k < k_max_d
|
||||
|| k_max_k < k_min_k
|
||||
|| nb_samples == 0
|
||||
|| dict_buffer_capacity < ZDICT_DICTSIZE_MIN
|
||||
{
|
||||
return if nb_samples == 0 {
|
||||
ERROR_SRC_SIZE_WRONG
|
||||
} else if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
|
||||
ERROR_DST_SIZE_TOO_SMALL
|
||||
} else {
|
||||
ERROR_PARAMETER_OUT_OF_BOUND
|
||||
};
|
||||
}
|
||||
let k_step_size = ((k_max_k - k_min_k) / k_steps).max(1);
|
||||
let _k_iterations = (1 + (k_max_d - k_min_d) / 2) * (1 + (k_max_k - k_min_k) / k_step_size);
|
||||
let mut token = 0u8;
|
||||
let token_ptr = (&mut token as *mut u8).cast::<c_void>();
|
||||
unsafe { COVER_best_init(token_ptr) };
|
||||
let slot = match slot_for(token_ptr) {
|
||||
Some(slot) => slot,
|
||||
None => return ERROR_MEMORY_ALLOCATION,
|
||||
};
|
||||
let mut warned = false;
|
||||
for d in (k_min_d..=k_max_d).step_by(2) {
|
||||
let context = match context_init(
|
||||
samples_buffer.cast(),
|
||||
samples_sizes,
|
||||
nb_samples as usize,
|
||||
d as usize,
|
||||
split_point,
|
||||
) {
|
||||
Ok(context) => context,
|
||||
Err(error_code) => {
|
||||
unsafe { COVER_best_destroy(token_ptr) };
|
||||
return error_code;
|
||||
}
|
||||
};
|
||||
if !warned {
|
||||
COVER_warnOnSmallCorpus(
|
||||
dict_buffer_capacity,
|
||||
context.suffix_size,
|
||||
parameters_ref.zParams.notificationLevel as c_int,
|
||||
);
|
||||
warned = true;
|
||||
}
|
||||
let mut jobs = Vec::new();
|
||||
let mut k = k_min_k;
|
||||
while k <= k_max_k {
|
||||
let mut job = *parameters_ref;
|
||||
job.k = k;
|
||||
job.d = d;
|
||||
job.splitPoint = split_point;
|
||||
job.steps = k_steps;
|
||||
job.shrinkDict = 0;
|
||||
best_start_slot(&slot);
|
||||
jobs.push(job);
|
||||
k = k.saturating_add(k_step_size);
|
||||
if k == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if nb_threads > 1 && jobs.len() > 1 {
|
||||
let worker_count = nb_threads.min(jobs.len());
|
||||
std::thread::scope(|scope| {
|
||||
let next_job = Arc::new(Mutex::new(jobs.into_iter()));
|
||||
let context_ref = &context;
|
||||
for _ in 0..worker_count {
|
||||
let next_job = Arc::clone(&next_job);
|
||||
let slot = Arc::clone(&slot);
|
||||
scope.spawn(move || loop {
|
||||
let job = {
|
||||
let mut jobs =
|
||||
next_job.lock().unwrap_or_else(|poison| poison.into_inner());
|
||||
jobs.next()
|
||||
};
|
||||
match job {
|
||||
Some(job) => {
|
||||
try_parameters(context_ref, &slot, dict_buffer_capacity, job)
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
} else {
|
||||
for job in jobs {
|
||||
try_parameters(&context, &slot, dict_buffer_capacity, job);
|
||||
}
|
||||
}
|
||||
best_wait_slot(&slot);
|
||||
}
|
||||
let snapshot = best_snapshot(&slot);
|
||||
if is_error(snapshot.compressed_size) {
|
||||
let result = snapshot.compressed_size;
|
||||
unsafe { COVER_best_destroy(token_ptr) };
|
||||
return result;
|
||||
}
|
||||
unsafe {
|
||||
ptr::copy_nonoverlapping(snapshot.dict, dict_buffer.cast(), snapshot.dict_size);
|
||||
*parameters = snapshot.parameters;
|
||||
COVER_best_destroy(token_ptr);
|
||||
}
|
||||
snapshot.dict_size
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn epoch_calculation_matches_cover_rules() {
|
||||
assert_eq!(
|
||||
COVER_computeEpochs(1024, 10_000, 100, 4),
|
||||
COVER_epoch_info_t { num: 2, size: 5000 }
|
||||
);
|
||||
assert_eq!(
|
||||
COVER_computeEpochs(1024, 100, 100, 4),
|
||||
COVER_epoch_info_t { num: 1, size: 100 }
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sum_handles_empty_and_multiple_samples() {
|
||||
let sizes = [3usize, 5, 8];
|
||||
assert_eq!(unsafe { COVER_sum(sizes.as_ptr(), 0) }, 0);
|
||||
assert_eq!(
|
||||
unsafe { COVER_sum(sizes.as_ptr(), sizes.len() as c_uint) },
|
||||
16
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn context_groups_repeated_dmers_by_sample() {
|
||||
let samples = [b'a'; 45];
|
||||
let sizes = [9usize, 9, 9, 9, 9];
|
||||
let context = context_init(samples.as_ptr(), sizes.as_ptr(), sizes.len(), 4, 1.0)
|
||||
.expect("valid COVER context");
|
||||
assert_eq!(context.suffix_size, samples.len() - 8 + 1);
|
||||
assert_eq!(context.freqs.len(), context.suffix_size);
|
||||
assert!(context.freqs.iter().any(|frequency| *frequency >= 5));
|
||||
}
|
||||
}
|
||||
@@ -18,6 +18,8 @@ pub mod hist;
|
||||
pub mod huf_compress;
|
||||
#[cfg(feature = "decompression")]
|
||||
pub mod huf_decompress;
|
||||
#[cfg(feature = "dict-builder")]
|
||||
pub mod dict_builder_cover;
|
||||
pub mod legacy;
|
||||
pub mod mem;
|
||||
pub mod pool;
|
||||
|
||||
Reference in New Issue
Block a user