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

feat(rust): port the COVER dictionary trainer

Browse Source 
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:
ddidderr
2026-07-12 10:08:03 +02:00
parent 70260b308b
commit 468ccfa388
3 changed files with 1314 additions and 1291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/dictBuilder/cover.c
+5 -1291
View File
@@ -8,1295 +8,9 @@
* You may select, at your option, one of the above-listed licenses. * You may select, at your option, one of the above-listed licenses.
*/ */
/* ***************************************************************************** /*
* Constructs a dictionary using a heuristic based on the following paper: * COVER is implemented by rust/src/dict_builder_cover.rs. Keep this
* * translation unit so existing C build descriptions continue to include the
* Liao, Petri, Moffat, Wirth * public COVER ABI header while the Rust static library supplies the symbols.
* Effective Construction of Relative Lempel-Ziv Dictionaries */
* Published in WWW 2016.
*
* Adapted from code originally written by @ot (Giuseppe Ottaviano).
******************************************************************************/
/*-*************************************
* Dependencies
***************************************/
/* qsort_r is an extension. */
#if defined(__linux) || defined(__linux__) || defined(linux) || defined(__gnu_linux__) || \
defined(__CYGWIN__) || defined(__MSYS__)
#if !defined(_GNU_SOURCE) && !defined(__ANDROID__) /* NDK doesn't ship qsort_r(). */
#define _GNU_SOURCE
#endif
#endif
#include <stdio.h> /* fprintf */
#include <stdlib.h> /* malloc, free, qsort_r */
#include <string.h> /* memset */
#include <time.h> /* clock */
#ifndef ZDICT_STATIC_LINKING_ONLY
# define ZDICT_STATIC_LINKING_ONLY
#endif
#include "../common/mem.h" /* read */
#include "../common/pool.h" /* POOL_ctx */
#include "../common/threading.h" /* ZSTD_pthread_mutex_t */
#include "../common/zstd_internal.h" /* includes zstd.h */
#include "../common/bits.h" /* ZSTD_highbit32 */
#include "../zdict.h"
#include "cover.h" #include "cover.h"
/*-*************************************
* Constants
***************************************/
/**
* There are 32bit indexes used to ref samples, so limit samples size to 4GB
* on 64bit builds.
* For 32bit builds we choose 1 GB.
* Most 32bit platforms have 2GB user-mode addressable space and we allocate a large
* contiguous buffer, so 1GB is already a high limit.
*/
#define COVER_MAX_SAMPLES_SIZE (sizeof(size_t) == 8 ? ((unsigned)-1) : ((unsigned)1 GB))
#define COVER_DEFAULT_SPLITPOINT 1.0
/*-*************************************
* Console display
***************************************/
#ifndef LOCALDISPLAYLEVEL
static int g_displayLevel = 0;
#endif
#undef DISPLAY
#define DISPLAY(...) \
{ \
fprintf(stderr, __VA_ARGS__); \
fflush(stderr); \
}
#undef LOCALDISPLAYLEVEL
#define LOCALDISPLAYLEVEL(displayLevel, l, ...) \
if (displayLevel >= l) { \
DISPLAY(__VA_ARGS__); \
} /* 0 : no display; 1: errors; 2: default; 3: details; 4: debug */
#undef DISPLAYLEVEL
#define DISPLAYLEVEL(l, ...) LOCALDISPLAYLEVEL(g_displayLevel, l, __VA_ARGS__)
#ifndef LOCALDISPLAYUPDATE
static const clock_t g_refreshRate = CLOCKS_PER_SEC * 15 / 100;
static clock_t g_time = 0;
#endif
#undef LOCALDISPLAYUPDATE
#define LOCALDISPLAYUPDATE(displayLevel, l, ...) \
if (displayLevel >= l) { \
if ((clock() - g_time > g_refreshRate) || (displayLevel >= 4)) { \
g_time = clock(); \
DISPLAY(__VA_ARGS__); \
} \
}
#undef DISPLAYUPDATE
#define DISPLAYUPDATE(l, ...) LOCALDISPLAYUPDATE(g_displayLevel, l, __VA_ARGS__)
/*-*************************************
* Hash table
***************************************
* A small specialized hash map for storing activeDmers.
* The map does not resize, so if it becomes full it will loop forever.
* Thus, the map must be large enough to store every value.
* The map implements linear probing and keeps its load less than 0.5.
*/
#define MAP_EMPTY_VALUE ((U32)-1)
typedef struct COVER_map_pair_t_s {
U32 key;
U32 value;
} COVER_map_pair_t;
typedef struct COVER_map_s {
COVER_map_pair_t *data;
U32 sizeLog;
U32 size;
U32 sizeMask;
} COVER_map_t;
/**
* Clear the map.
*/
static void COVER_map_clear(COVER_map_t *map) {
memset(map->data, MAP_EMPTY_VALUE, map->size * sizeof(COVER_map_pair_t));
}
/**
* Initializes a map of the given size.
* Returns 1 on success and 0 on failure.
* The map must be destroyed with COVER_map_destroy().
* The map is only guaranteed to be large enough to hold size elements.
*/
static int COVER_map_init(COVER_map_t *map, U32 size) {
map->sizeLog = ZSTD_highbit32(size) + 2;
map->size = (U32)1 << map->sizeLog;
map->sizeMask = map->size - 1;
map->data = (COVER_map_pair_t *)malloc(map->size * sizeof(COVER_map_pair_t));
if (!map->data) {
map->sizeLog = 0;
map->size = 0;
return 0;
}
COVER_map_clear(map);
return 1;
}
/**
* Internal hash function
*/
static const U32 COVER_prime4bytes = 2654435761U;
static U32 COVER_map_hash(COVER_map_t *map, U32 key) {
return (key * COVER_prime4bytes) >> (32 - map->sizeLog);
}
/**
* Helper function that returns the index that a key should be placed into.
*/
static U32 COVER_map_index(COVER_map_t *map, U32 key) {
const U32 hash = COVER_map_hash(map, key);
U32 i;
for (i = hash;; i = (i + 1) & map->sizeMask) {
COVER_map_pair_t *pos = &map->data[i];
if (pos->value == MAP_EMPTY_VALUE) {
return i;
}
if (pos->key == key) {
return i;
}
}
}
/**
* Returns the pointer to the value for key.
* If key is not in the map, it is inserted and the value is set to 0.
* The map must not be full.
*/
static U32 *COVER_map_at(COVER_map_t *map, U32 key) {
COVER_map_pair_t *pos = &map->data[COVER_map_index(map, key)];
if (pos->value == MAP_EMPTY_VALUE) {
pos->key = key;
pos->value = 0;
}
return &pos->value;
}
/**
* Deletes key from the map if present.
*/
static void COVER_map_remove(COVER_map_t *map, U32 key) {
U32 i = COVER_map_index(map, key);
COVER_map_pair_t *del = &map->data[i];
U32 shift = 1;
if (del->value == MAP_EMPTY_VALUE) {
return;
}
for (i = (i + 1) & map->sizeMask;; i = (i + 1) & map->sizeMask) {
COVER_map_pair_t *const pos = &map->data[i];
/* If the position is empty we are done */
if (pos->value == MAP_EMPTY_VALUE) {
del->value = MAP_EMPTY_VALUE;
return;
}
/* If pos can be moved to del do so */
if (((i - COVER_map_hash(map, pos->key)) & map->sizeMask) >= shift) {
del->key = pos->key;
del->value = pos->value;
del = pos;
shift = 1;
} else {
++shift;
}
}
}
/**
* Destroys a map that is inited with COVER_map_init().
*/
static void COVER_map_destroy(COVER_map_t *map) {
if (map->data) {
free(map->data);
}
map->data = NULL;
map->size = 0;
}
/*-*************************************
* Context
***************************************/
typedef struct {
const BYTE *samples;
size_t *offsets;
const size_t *samplesSizes;
size_t nbSamples;
size_t nbTrainSamples;
size_t nbTestSamples;
U32 *suffix;
size_t suffixSize;
U32 *freqs;
U32 *dmerAt;
unsigned d;
} COVER_ctx_t;
#if !defined(_GNU_SOURCE) && !defined(__APPLE__) && !defined(_MSC_VER)
/* C90 only offers qsort() that needs a global context. */
static COVER_ctx_t *g_coverCtx = NULL;
#endif
/*-*************************************
* Helper functions
***************************************/
/**
* Returns the sum of the sample sizes.
*/
size_t COVER_sum(const size_t *samplesSizes, unsigned nbSamples) {
size_t sum = 0;
unsigned i;
for (i = 0; i < nbSamples; ++i) {
sum += samplesSizes[i];
}
return sum;
}
/**
* Returns -1 if the dmer at lp is less than the dmer at rp.
* Return 0 if the dmers at lp and rp are equal.
* Returns 1 if the dmer at lp is greater than the dmer at rp.
*/
static int COVER_cmp(COVER_ctx_t *ctx, const void *lp, const void *rp) {
U32 const lhs = *(U32 const *)lp;
U32 const rhs = *(U32 const *)rp;
return memcmp(ctx->samples + lhs, ctx->samples + rhs, ctx->d);
}
/**
* Faster version for d <= 8.
*/
static int COVER_cmp8(COVER_ctx_t *ctx, const void *lp, const void *rp) {
U64 const mask = (ctx->d == 8) ? (U64)-1 : (((U64)1 << (8 * ctx->d)) - 1);
U64 const lhs = MEM_readLE64(ctx->samples + *(U32 const *)lp) & mask;
U64 const rhs = MEM_readLE64(ctx->samples + *(U32 const *)rp) & mask;
if (lhs < rhs) {
return -1;
}
return (lhs > rhs);
}
/**
* Same as COVER_cmp() except ties are broken by pointer value
*/
#if (defined(_WIN32) && defined(_MSC_VER)) || defined(__APPLE__)
static int WIN_CDECL COVER_strict_cmp(void* g_coverCtx, const void* lp, const void* rp) {
#elif defined(_GNU_SOURCE)
static int COVER_strict_cmp(const void *lp, const void *rp, void *g_coverCtx) {
#else /* C90 fallback.*/
static int COVER_strict_cmp(const void *lp, const void *rp) {
#endif
int result = COVER_cmp((COVER_ctx_t*)g_coverCtx, lp, rp);
if (result == 0) {
result = lp < rp ? -1 : 1;
}
return result;
}
/**
* Faster version for d <= 8.
*/
#if (defined(_WIN32) && defined(_MSC_VER)) || defined(__APPLE__)
static int WIN_CDECL COVER_strict_cmp8(void* g_coverCtx, const void* lp, const void* rp) {
#elif defined(_GNU_SOURCE)
static int COVER_strict_cmp8(const void *lp, const void *rp, void *g_coverCtx) {
#else /* C90 fallback.*/
static int COVER_strict_cmp8(const void *lp, const void *rp) {
#endif
int result = COVER_cmp8((COVER_ctx_t*)g_coverCtx, lp, rp);
if (result == 0) {
result = lp < rp ? -1 : 1;
}
return result;
}
/**
* Abstract away divergence of qsort_r() parameters.
* Hopefully when C11 become the norm, we will be able
* to clean it up.
*/
static void stableSort(COVER_ctx_t *ctx) {
#if defined(__APPLE__)
qsort_r(ctx->suffix, ctx->suffixSize, sizeof(U32),
ctx,
(ctx->d <= 8 ? &COVER_strict_cmp8 : &COVER_strict_cmp));
#elif defined(_GNU_SOURCE)
qsort_r(ctx->suffix, ctx->suffixSize, sizeof(U32),
(ctx->d <= 8 ? &COVER_strict_cmp8 : &COVER_strict_cmp),
ctx);
#elif defined(_WIN32) && defined(_MSC_VER)
qsort_s(ctx->suffix, ctx->suffixSize, sizeof(U32),
(ctx->d <= 8 ? &COVER_strict_cmp8 : &COVER_strict_cmp),
ctx);
#elif defined(__OpenBSD__)
g_coverCtx = ctx;
mergesort(ctx->suffix, ctx->suffixSize, sizeof(U32),
(ctx->d <= 8 ? &COVER_strict_cmp8 : &COVER_strict_cmp));
#else /* C90 fallback.*/
g_coverCtx = ctx;
/* TODO(cavalcanti): implement a reentrant qsort() when is not available. */
qsort(ctx->suffix, ctx->suffixSize, sizeof(U32),
(ctx->d <= 8 ? &COVER_strict_cmp8 : &COVER_strict_cmp));
#endif
}
/**
* Returns the first pointer in [first, last) whose element does not compare
* less than value. If no such element exists it returns last.
*/
static const size_t *COVER_lower_bound(const size_t* first, const size_t* last,
size_t value) {
size_t count = (size_t)(last - first);
assert(last >= first);
while (count != 0) {
size_t step = count / 2;
const size_t *ptr = first;
ptr += step;
if (*ptr < value) {
first = ++ptr;
count -= step + 1;
} else {
count = step;
}
}
return first;
}
/**
* Generic groupBy function.
* Groups an array sorted by cmp into groups with equivalent values.
* Calls grp for each group.
*/
static void
COVER_groupBy(const void *data, size_t count, size_t size, COVER_ctx_t *ctx,
int (*cmp)(COVER_ctx_t *, const void *, const void *),
void (*grp)(COVER_ctx_t *, const void *, const void *)) {
const BYTE *ptr = (const BYTE *)data;
size_t num = 0;
while (num < count) {
const BYTE *grpEnd = ptr + size;
++num;
while (num < count && cmp(ctx, ptr, grpEnd) == 0) {
grpEnd += size;
++num;
}
grp(ctx, ptr, grpEnd);
ptr = grpEnd;
}
}
/*-*************************************
* Cover functions
***************************************/
/**
* Called on each group of positions with the same dmer.
* Counts the frequency of each dmer and saves it in the suffix array.
* Fills `ctx->dmerAt`.
*/
static void COVER_group(COVER_ctx_t *ctx, const void *group,
const void *groupEnd) {
/* The group consists of all the positions with the same first d bytes. */
const U32 *grpPtr = (const U32 *)group;
const U32 *grpEnd = (const U32 *)groupEnd;
/* The dmerId is how we will reference this dmer.
* This allows us to map the whole dmer space to a much smaller space, the
* size of the suffix array.
*/
const U32 dmerId = (U32)(grpPtr - ctx->suffix);
/* Count the number of samples this dmer shows up in */
U32 freq = 0;
/* Details */
const size_t *curOffsetPtr = ctx->offsets;
const size_t *offsetsEnd = ctx->offsets + ctx->nbSamples;
/* Once *grpPtr >= curSampleEnd this occurrence of the dmer is in a
* different sample than the last.
*/
size_t curSampleEnd = ctx->offsets[0];
for (; grpPtr != grpEnd; ++grpPtr) {
/* Save the dmerId for this position so we can get back to it. */
ctx->dmerAt[*grpPtr] = dmerId;
/* Dictionaries only help for the first reference to the dmer.
* After that zstd can reference the match from the previous reference.
* So only count each dmer once for each sample it is in.
*/
if (*grpPtr < curSampleEnd) {
continue;
}
freq += 1;
/* Binary search to find the end of the sample *grpPtr is in.
* In the common case that grpPtr + 1 == grpEnd we can skip the binary
* search because the loop is over.
*/
if (grpPtr + 1 != grpEnd) {
const size_t *sampleEndPtr =
COVER_lower_bound(curOffsetPtr, offsetsEnd, *grpPtr);
curSampleEnd = *sampleEndPtr;
curOffsetPtr = sampleEndPtr + 1;
}
}
/* At this point we are never going to look at this segment of the suffix
* array again. We take advantage of this fact to save memory.
* We store the frequency of the dmer in the first position of the group,
* which is dmerId.
*/
ctx->suffix[dmerId] = freq;
}
/**
* Selects the best segment in an epoch.
* Segments of are scored according to the function:
*
* Let F(d) be the frequency of dmer d.
* Let S_i be the dmer at position i of segment S which has length k.
*
* Score(S) = F(S_1) + F(S_2) + ... + F(S_{k-d+1})
*
* Once the dmer d is in the dictionary we set F(d) = 0.
*/
static COVER_segment_t COVER_selectSegment(const COVER_ctx_t *ctx, U32 *freqs,
COVER_map_t *activeDmers, U32 begin,
U32 end,
ZDICT_cover_params_t parameters) {
/* Constants */
const U32 k = parameters.k;
const U32 d = parameters.d;
const U32 dmersInK = k - d + 1;
/* Try each segment (activeSegment) and save the best (bestSegment) */
COVER_segment_t bestSegment = {0, 0, 0};
COVER_segment_t activeSegment;
/* Reset the activeDmers in the segment */
COVER_map_clear(activeDmers);
/* The activeSegment starts at the beginning of the epoch. */
activeSegment.begin = begin;
activeSegment.end = begin;
activeSegment.score = 0;
/* Slide the activeSegment through the whole epoch.
* Save the best segment in bestSegment.
*/
while (activeSegment.end < end) {
/* The dmerId for the dmer at the next position */
U32 newDmer = ctx->dmerAt[activeSegment.end];
/* The entry in activeDmers for this dmerId */
U32 *newDmerOcc = COVER_map_at(activeDmers, newDmer);
/* If the dmer isn't already present in the segment add its score. */
if (*newDmerOcc == 0) {
/* The paper suggest using the L-0.5 norm, but experiments show that it
* doesn't help.
*/
activeSegment.score += freqs[newDmer];
}
/* Add the dmer to the segment */
activeSegment.end += 1;
*newDmerOcc += 1;
/* If the window is now too large, drop the first position */
if (activeSegment.end - activeSegment.begin == dmersInK + 1) {
U32 delDmer = ctx->dmerAt[activeSegment.begin];
U32 *delDmerOcc = COVER_map_at(activeDmers, delDmer);
activeSegment.begin += 1;
*delDmerOcc -= 1;
/* If this is the last occurrence of the dmer, subtract its score */
if (*delDmerOcc == 0) {
COVER_map_remove(activeDmers, delDmer);
activeSegment.score -= freqs[delDmer];
}
}
/* If this segment is the best so far save it */
if (activeSegment.score > bestSegment.score) {
bestSegment = activeSegment;
}
}
{
/* Trim off the zero frequency head and tail from the segment. */
U32 newBegin = bestSegment.end;
U32 newEnd = bestSegment.begin;
U32 pos;
for (pos = bestSegment.begin; pos != bestSegment.end; ++pos) {
U32 freq = freqs[ctx->dmerAt[pos]];
if (freq != 0) {
newBegin = MIN(newBegin, pos);
newEnd = pos + 1;
}
}
bestSegment.begin = newBegin;
bestSegment.end = newEnd;
}
{
/* Zero out the frequency of each dmer covered by the chosen segment. */
U32 pos;
for (pos = bestSegment.begin; pos != bestSegment.end; ++pos) {
freqs[ctx->dmerAt[pos]] = 0;
}
}
return bestSegment;
}
/**
* Check the validity of the parameters.
* Returns non-zero if the parameters are valid and 0 otherwise.
*/
static int COVER_checkParameters(ZDICT_cover_params_t parameters,
size_t maxDictSize) {
/* k and d are required parameters */
if (parameters.d == 0 || parameters.k == 0) {
return 0;
}
/* k <= maxDictSize */
if (parameters.k > maxDictSize) {
return 0;
}
/* d <= k */
if (parameters.d > parameters.k) {
return 0;
}
/* 0 < splitPoint <= 1 */
if (parameters.splitPoint <= 0 || parameters.splitPoint > 1){
return 0;
}
return 1;
}
/**
* Clean up a context initialized with `COVER_ctx_init()`.
*/
static void COVER_ctx_destroy(COVER_ctx_t *ctx) {
if (!ctx) {
return;
}
if (ctx->suffix) {
free(ctx->suffix);
ctx->suffix = NULL;
}
if (ctx->freqs) {
free(ctx->freqs);
ctx->freqs = NULL;
}
if (ctx->dmerAt) {
free(ctx->dmerAt);
ctx->dmerAt = NULL;
}
if (ctx->offsets) {
free(ctx->offsets);
ctx->offsets = NULL;
}
}
/**
* Prepare a context for dictionary building.
* The context is only dependent on the parameter `d` and can be used multiple
* times.
* Returns 0 on success or error code on error.
* The context must be destroyed with `COVER_ctx_destroy()`.
*/
static size_t COVER_ctx_init(COVER_ctx_t *ctx, const void *samplesBuffer,
const size_t *samplesSizes, unsigned nbSamples,
unsigned d, double splitPoint)
{
const BYTE *const samples = (const BYTE *)samplesBuffer;
const size_t totalSamplesSize = COVER_sum(samplesSizes, nbSamples);
/* Split samples into testing and training sets */
const unsigned nbTrainSamples = splitPoint < 1.0 ? (unsigned)((double)nbSamples * splitPoint) : nbSamples;
const unsigned nbTestSamples = splitPoint < 1.0 ? nbSamples - nbTrainSamples : nbSamples;
const size_t trainingSamplesSize = splitPoint < 1.0 ? COVER_sum(samplesSizes, nbTrainSamples) : totalSamplesSize;
const size_t testSamplesSize = splitPoint < 1.0 ? COVER_sum(samplesSizes + nbTrainSamples, nbTestSamples) : totalSamplesSize;
/* Checks */
if (totalSamplesSize < MAX(d, sizeof(U64)) ||
totalSamplesSize >= (size_t)COVER_MAX_SAMPLES_SIZE) {
DISPLAYLEVEL(1, "Total samples size is too large (%u MB), maximum size is %u MB\n",
(unsigned)(totalSamplesSize>>20), (COVER_MAX_SAMPLES_SIZE >> 20));
return ERROR(srcSize_wrong);
}
/* Check if there are at least 5 training samples */
if (nbTrainSamples < 5) {
DISPLAYLEVEL(1, "Total number of training samples is %u and is invalid.", nbTrainSamples);
return ERROR(srcSize_wrong);
}
/* Check if there's testing sample */
if (nbTestSamples < 1) {
DISPLAYLEVEL(1, "Total number of testing samples is %u and is invalid.", nbTestSamples);
return ERROR(srcSize_wrong);
}
/* Zero the context */
memset(ctx, 0, sizeof(*ctx));
DISPLAYLEVEL(2, "Training on %u samples of total size %u\n", nbTrainSamples,
(unsigned)trainingSamplesSize);
DISPLAYLEVEL(2, "Testing on %u samples of total size %u\n", nbTestSamples,
(unsigned)testSamplesSize);
ctx->samples = samples;
ctx->samplesSizes = samplesSizes;
ctx->nbSamples = nbSamples;
ctx->nbTrainSamples = nbTrainSamples;
ctx->nbTestSamples = nbTestSamples;
/* Partial suffix array */
ctx->suffixSize = trainingSamplesSize - MAX(d, sizeof(U64)) + 1;
ctx->suffix = (U32 *)malloc(ctx->suffixSize * sizeof(U32));
/* Maps index to the dmerID */
ctx->dmerAt = (U32 *)malloc(ctx->suffixSize * sizeof(U32));
/* The offsets of each file */
ctx->offsets = (size_t *)malloc((nbSamples + 1) * sizeof(size_t));
if (!ctx->suffix || !ctx->dmerAt || !ctx->offsets) {
DISPLAYLEVEL(1, "Failed to allocate scratch buffers\n");
COVER_ctx_destroy(ctx);
return ERROR(memory_allocation);
}
ctx->freqs = NULL;
ctx->d = d;
/* Fill offsets from the samplesSizes */
{
U32 i;
ctx->offsets[0] = 0;
for (i = 1; i <= nbSamples; ++i) {
ctx->offsets[i] = ctx->offsets[i - 1] + samplesSizes[i - 1];
}
}
DISPLAYLEVEL(2, "Constructing partial suffix array\n");
{
/* suffix is a partial suffix array.
* It only sorts suffixes by their first parameters.d bytes.
* The sort is stable, so each dmer group is sorted by position in input.
*/
U32 i;
for (i = 0; i < ctx->suffixSize; ++i) {
ctx->suffix[i] = i;
}
stableSort(ctx);
}
DISPLAYLEVEL(2, "Computing frequencies\n");
/* For each dmer group (group of positions with the same first d bytes):
* 1. For each position we set dmerAt[position] = dmerID. The dmerID is
* (groupBeginPtr - suffix). This allows us to go from position to
* dmerID so we can look up values in freq.
* 2. We calculate how many samples the dmer occurs in and save it in
* freqs[dmerId].
*/
COVER_groupBy(ctx->suffix, ctx->suffixSize, sizeof(U32), ctx,
(ctx->d <= 8 ? &COVER_cmp8 : &COVER_cmp), &COVER_group);
ctx->freqs = ctx->suffix;
ctx->suffix = NULL;
return 0;
}
void COVER_warnOnSmallCorpus(size_t maxDictSize, size_t nbDmers, int displayLevel)
{
const double ratio = (double)nbDmers / (double)maxDictSize;
if (ratio >= 10) {
return;
}
LOCALDISPLAYLEVEL(displayLevel, 1,
"WARNING: The maximum dictionary size %u is too large "
"compared to the source size %u! "
"size(source)/size(dictionary) = %f, 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! \n", (U32)maxDictSize,
(U32)nbDmers, ratio);
}
COVER_epoch_info_t COVER_computeEpochs(U32 maxDictSize,
U32 nbDmers, U32 k, U32 passes)
{
const U32 minEpochSize = k * 10;
COVER_epoch_info_t epochs;
epochs.num = MAX(1, maxDictSize / k / passes);
epochs.size = nbDmers / epochs.num;
if (epochs.size >= minEpochSize) {
assert(epochs.size * epochs.num <= nbDmers);
return epochs;
}
epochs.size = MIN(minEpochSize, nbDmers);
epochs.num = nbDmers / epochs.size;
assert(epochs.size * epochs.num <= nbDmers);
return epochs;
}
/**
* Given the prepared context build the dictionary.
*/
static size_t COVER_buildDictionary(const COVER_ctx_t *ctx, U32 *freqs,
COVER_map_t *activeDmers, void *dictBuffer,
size_t dictBufferCapacity,
ZDICT_cover_params_t parameters) {
BYTE *const dict = (BYTE *)dictBuffer;
size_t tail = dictBufferCapacity;
/* Divide the data into epochs. We will select one segment from each epoch. */
const COVER_epoch_info_t epochs = COVER_computeEpochs(
(U32)dictBufferCapacity, (U32)ctx->suffixSize, parameters.k, 4);
const size_t maxZeroScoreRun = MAX(10, MIN(100, epochs.num >> 3));
size_t zeroScoreRun = 0;
size_t epoch;
DISPLAYLEVEL(2, "Breaking content into %u epochs of size %u\n",
(U32)epochs.num, (U32)epochs.size);
/* Loop through the epochs until there are no more segments or the dictionary
* is full.
*/
for (epoch = 0; tail > 0; epoch = (epoch + 1) % epochs.num) {
const U32 epochBegin = (U32)(epoch * epochs.size);
const U32 epochEnd = epochBegin + epochs.size;
size_t segmentSize;
/* Select a segment */
COVER_segment_t segment = COVER_selectSegment(
ctx, freqs, activeDmers, epochBegin, epochEnd, parameters);
/* If the segment covers no dmers, then we are out of content.
* There may be new content in other epochs, for continue for some time.
*/
if (segment.score == 0) {
if (++zeroScoreRun >= maxZeroScoreRun) {
break;
}
continue;
}
zeroScoreRun = 0;
/* Trim the segment if necessary and if it is too small then we are done */
segmentSize = MIN(segment.end - segment.begin + parameters.d - 1, tail);
if (segmentSize < parameters.d) {
break;
}
/* We fill the dictionary from the back to allow the best segments to be
* referenced with the smallest offsets.
*/
tail -= segmentSize;
memcpy(dict + tail, ctx->samples + segment.begin, segmentSize);
DISPLAYUPDATE(
2, "\r%u%% ",
(unsigned)(((dictBufferCapacity - tail) * 100) / dictBufferCapacity));
}
DISPLAYLEVEL(2, "\r%79s\r", "");
return tail;
}
ZDICTLIB_STATIC_API size_t ZDICT_trainFromBuffer_cover(
void *dictBuffer, size_t dictBufferCapacity,
const void *samplesBuffer, const size_t *samplesSizes, unsigned nbSamples,
ZDICT_cover_params_t parameters)
{
BYTE* const dict = (BYTE*)dictBuffer;
COVER_ctx_t ctx;
COVER_map_t activeDmers;
parameters.splitPoint = 1.0;
/* Initialize global data */
g_displayLevel = (int)parameters.zParams.notificationLevel;
/* Checks */
if (!COVER_checkParameters(parameters, dictBufferCapacity)) {
DISPLAYLEVEL(1, "Cover parameters incorrect\n");
return ERROR(parameter_outOfBound);
}
if (nbSamples == 0) {
DISPLAYLEVEL(1, "Cover must have at least one input file\n");
return ERROR(srcSize_wrong);
}
if (dictBufferCapacity < ZDICT_DICTSIZE_MIN) {
DISPLAYLEVEL(1, "dictBufferCapacity must be at least %u\n",
ZDICT_DICTSIZE_MIN);
return ERROR(dstSize_tooSmall);
}
/* Initialize context and activeDmers */
{
size_t const initVal = COVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples,
parameters.d, parameters.splitPoint);
if (ZSTD_isError(initVal)) {
return initVal;
}
}
COVER_warnOnSmallCorpus(dictBufferCapacity, ctx.suffixSize, g_displayLevel);
if (!COVER_map_init(&activeDmers, parameters.k - parameters.d + 1)) {
DISPLAYLEVEL(1, "Failed to allocate dmer map: out of memory\n");
COVER_ctx_destroy(&ctx);
return ERROR(memory_allocation);
}
DISPLAYLEVEL(2, "Building dictionary\n");
{
const size_t tail =
COVER_buildDictionary(&ctx, ctx.freqs, &activeDmers, dictBuffer,
dictBufferCapacity, parameters);
const size_t dictionarySize = ZDICT_finalizeDictionary(
dict, dictBufferCapacity, dict + tail, dictBufferCapacity - tail,
samplesBuffer, samplesSizes, nbSamples, parameters.zParams);
if (!ZSTD_isError(dictionarySize)) {
DISPLAYLEVEL(2, "Constructed dictionary of size %u\n",
(unsigned)dictionarySize);
}
COVER_ctx_destroy(&ctx);
COVER_map_destroy(&activeDmers);
return dictionarySize;
}
}
size_t COVER_checkTotalCompressedSize(const ZDICT_cover_params_t parameters,
const size_t *samplesSizes, const BYTE *samples,
size_t *offsets,
size_t nbTrainSamples, size_t nbSamples,
BYTE *const dict, size_t dictBufferCapacity) {
size_t totalCompressedSize = ERROR(GENERIC);
/* Pointers */
ZSTD_CCtx *cctx;
ZSTD_CDict *cdict;
void *dst;
/* Local variables */
size_t dstCapacity;
size_t i;
/* Allocate dst with enough space to compress the maximum sized sample */
{
size_t maxSampleSize = 0;
i = parameters.splitPoint < 1.0 ? nbTrainSamples : 0;
for (; i < nbSamples; ++i) {
maxSampleSize = MAX(samplesSizes[i], maxSampleSize);
}
dstCapacity = ZSTD_compressBound(maxSampleSize);
dst = malloc(dstCapacity);
}
/* Create the cctx and cdict */
cctx = ZSTD_createCCtx();
cdict = ZSTD_createCDict(dict, dictBufferCapacity,
parameters.zParams.compressionLevel);
if (!dst || !cctx || !cdict) {
goto _compressCleanup;
}
/* Compress each sample and sum their sizes (or error) */
totalCompressedSize = dictBufferCapacity;
i = parameters.splitPoint < 1.0 ? nbTrainSamples : 0;
for (; i < nbSamples; ++i) {
const size_t size = ZSTD_compress_usingCDict(
cctx, dst, dstCapacity, samples + offsets[i],
samplesSizes[i], cdict);
if (ZSTD_isError(size)) {
totalCompressedSize = size;
goto _compressCleanup;
}
totalCompressedSize += size;
}
_compressCleanup:
ZSTD_freeCCtx(cctx);
ZSTD_freeCDict(cdict);
if (dst) {
free(dst);
}
return totalCompressedSize;
}
/**
* Initialize the `COVER_best_t`.
*/
void COVER_best_init(COVER_best_t *best) {
if (best==NULL) return; /* compatible with init on NULL */
(void)ZSTD_pthread_mutex_init(&best->mutex, NULL);
(void)ZSTD_pthread_cond_init(&best->cond, NULL);
best->liveJobs = 0;
best->dict = NULL;
best->dictSize = 0;
best->compressedSize = (size_t)-1;
memset(&best->parameters, 0, sizeof(best->parameters));
}
/**
* Wait until liveJobs == 0.
*/
void COVER_best_wait(COVER_best_t *best) {
if (!best) {
return;
}
ZSTD_pthread_mutex_lock(&best->mutex);
while (best->liveJobs != 0) {
ZSTD_pthread_cond_wait(&best->cond, &best->mutex);
}
ZSTD_pthread_mutex_unlock(&best->mutex);
}
/**
* Call COVER_best_wait() and then destroy the COVER_best_t.
*/
void COVER_best_destroy(COVER_best_t *best) {
if (!best) {
return;
}
COVER_best_wait(best);
if (best->dict) {
free(best->dict);
}
ZSTD_pthread_mutex_destroy(&best->mutex);
ZSTD_pthread_cond_destroy(&best->cond);
}
/**
* Called when a thread is about to be launched.
* Increments liveJobs.
*/
void COVER_best_start(COVER_best_t *best) {
if (!best) {
return;
}
ZSTD_pthread_mutex_lock(&best->mutex);
++best->liveJobs;
ZSTD_pthread_mutex_unlock(&best->mutex);
}
/**
* Called when a thread finishes executing, both on error or success.
* Decrements liveJobs and signals any waiting threads if liveJobs == 0.
* If this dictionary is the best so far save it and its parameters.
*/
void COVER_best_finish(COVER_best_t* best,
ZDICT_cover_params_t parameters,
COVER_dictSelection_t selection)
{
void* dict = selection.dictContent;
size_t compressedSize = selection.totalCompressedSize;
size_t dictSize = selection.dictSize;
if (!best) {
return;
}
{
size_t liveJobs;
ZSTD_pthread_mutex_lock(&best->mutex);
--best->liveJobs;
liveJobs = best->liveJobs;
/* If the new dictionary is better */
if (compressedSize < best->compressedSize) {
/* Allocate space if necessary */
if (!best->dict || best->dictSize < dictSize) {
if (best->dict) {
free(best->dict);
}
best->dict = malloc(dictSize);
if (!best->dict) {
best->compressedSize = ERROR(GENERIC);
best->dictSize = 0;
ZSTD_pthread_cond_signal(&best->cond);
ZSTD_pthread_mutex_unlock(&best->mutex);
return;
}
}
/* Save the dictionary, parameters, and size */
if (dict) {
memcpy(best->dict, dict, dictSize);
best->dictSize = dictSize;
best->parameters = parameters;
best->compressedSize = compressedSize;
}
}
if (liveJobs == 0) {
ZSTD_pthread_cond_broadcast(&best->cond);
}
ZSTD_pthread_mutex_unlock(&best->mutex);
}
}
static COVER_dictSelection_t setDictSelection(BYTE* buf, size_t s, size_t csz)
{
COVER_dictSelection_t ds;
ds.dictContent = buf;
ds.dictSize = s;
ds.totalCompressedSize = csz;
return ds;
}
COVER_dictSelection_t COVER_dictSelectionError(size_t error) {
return setDictSelection(NULL, 0, error);
}
unsigned COVER_dictSelectionIsError(COVER_dictSelection_t selection) {
return (ZSTD_isError(selection.totalCompressedSize) || !selection.dictContent);
}
void COVER_dictSelectionFree(COVER_dictSelection_t selection){
free(selection.dictContent);
}
COVER_dictSelection_t COVER_selectDict(BYTE* customDictContent, size_t dictBufferCapacity,
size_t dictContentSize, const BYTE* samplesBuffer, const size_t* samplesSizes, unsigned nbFinalizeSamples,
size_t nbCheckSamples, size_t nbSamples, ZDICT_cover_params_t params, size_t* offsets, size_t totalCompressedSize) {
size_t largestDict = 0;
size_t largestCompressed = 0;
BYTE* customDictContentEnd = customDictContent + dictContentSize;
BYTE* largestDictbuffer = (BYTE*)malloc(dictBufferCapacity);
BYTE* candidateDictBuffer = (BYTE*)malloc(dictBufferCapacity);
double regressionTolerance = ((double)params.shrinkDictMaxRegression / 100.0) + 1.00;
if (!largestDictbuffer || !candidateDictBuffer) {
free(largestDictbuffer);
free(candidateDictBuffer);
return COVER_dictSelectionError(dictContentSize);
}
/* Initial dictionary size and compressed size */
memcpy(largestDictbuffer, customDictContent, dictContentSize);
dictContentSize = ZDICT_finalizeDictionary(
largestDictbuffer, dictBufferCapacity, customDictContent, dictContentSize,
samplesBuffer, samplesSizes, nbFinalizeSamples, params.zParams);
if (ZDICT_isError(dictContentSize)) {
free(largestDictbuffer);
free(candidateDictBuffer);
return COVER_dictSelectionError(dictContentSize);
}
totalCompressedSize = COVER_checkTotalCompressedSize(params, samplesSizes,
samplesBuffer, offsets,
nbCheckSamples, nbSamples,
largestDictbuffer, dictContentSize);
if (ZSTD_isError(totalCompressedSize)) {
free(largestDictbuffer);
free(candidateDictBuffer);
return COVER_dictSelectionError(totalCompressedSize);
}
if (params.shrinkDict == 0) {
free(candidateDictBuffer);
return setDictSelection(largestDictbuffer, dictContentSize, totalCompressedSize);
}
largestDict = dictContentSize;
largestCompressed = totalCompressedSize;
dictContentSize = ZDICT_DICTSIZE_MIN;
/* Largest dict is initially at least ZDICT_DICTSIZE_MIN */
while (dictContentSize < largestDict) {
memcpy(candidateDictBuffer, largestDictbuffer, largestDict);
dictContentSize = ZDICT_finalizeDictionary(
candidateDictBuffer, dictBufferCapacity, customDictContentEnd - dictContentSize, dictContentSize,
samplesBuffer, samplesSizes, nbFinalizeSamples, params.zParams);
if (ZDICT_isError(dictContentSize)) {
free(largestDictbuffer);
free(candidateDictBuffer);
return COVER_dictSelectionError(dictContentSize);
}
totalCompressedSize = COVER_checkTotalCompressedSize(params, samplesSizes,
samplesBuffer, offsets,
nbCheckSamples, nbSamples,
candidateDictBuffer, dictContentSize);
if (ZSTD_isError(totalCompressedSize)) {
free(largestDictbuffer);
free(candidateDictBuffer);
return COVER_dictSelectionError(totalCompressedSize);
}
if ((double)totalCompressedSize <= (double)largestCompressed * regressionTolerance) {
free(largestDictbuffer);
return setDictSelection( candidateDictBuffer, dictContentSize, totalCompressedSize );
}
dictContentSize *= 2;
}
dictContentSize = largestDict;
totalCompressedSize = largestCompressed;
free(candidateDictBuffer);
return setDictSelection( largestDictbuffer, dictContentSize, totalCompressedSize );
}
/**
* Parameters for COVER_tryParameters().
*/
typedef struct COVER_tryParameters_data_s {
const COVER_ctx_t *ctx;
COVER_best_t *best;
size_t dictBufferCapacity;
ZDICT_cover_params_t parameters;
} COVER_tryParameters_data_t;
/**
* Tries a set of parameters and updates the COVER_best_t with the results.
* This function is thread safe if zstd is compiled with multithreaded support.
* It takes its parameters as an *OWNING* opaque pointer to support threading.
*/
static void COVER_tryParameters(void *opaque)
{
/* Save parameters as local variables */
COVER_tryParameters_data_t *const data = (COVER_tryParameters_data_t*)opaque;
const COVER_ctx_t *const ctx = data->ctx;
const ZDICT_cover_params_t parameters = data->parameters;
size_t dictBufferCapacity = data->dictBufferCapacity;
size_t totalCompressedSize = ERROR(GENERIC);
/* Allocate space for hash table, dict, and freqs */
COVER_map_t activeDmers;
BYTE* const dict = (BYTE*)malloc(dictBufferCapacity);
COVER_dictSelection_t selection = COVER_dictSelectionError(ERROR(GENERIC));
U32* const freqs = (U32*)malloc(ctx->suffixSize * sizeof(U32));
if (!COVER_map_init(&activeDmers, parameters.k - parameters.d + 1)) {
DISPLAYLEVEL(1, "Failed to allocate dmer map: out of memory\n");
goto _cleanup;
}
if (!dict || !freqs) {
DISPLAYLEVEL(1, "Failed to allocate buffers: out of memory\n");
goto _cleanup;
}
/* Copy the frequencies because we need to modify them */
memcpy(freqs, ctx->freqs, ctx->suffixSize * sizeof(U32));
/* Build the dictionary */
{
const size_t tail = COVER_buildDictionary(ctx, freqs, &activeDmers, dict,
dictBufferCapacity, parameters);
selection = COVER_selectDict(dict + tail, dictBufferCapacity, dictBufferCapacity - tail,
ctx->samples, ctx->samplesSizes, (unsigned)ctx->nbTrainSamples, ctx->nbTrainSamples, ctx->nbSamples, parameters, ctx->offsets,
totalCompressedSize);
if (COVER_dictSelectionIsError(selection)) {
DISPLAYLEVEL(1, "Failed to select dictionary\n");
goto _cleanup;
}
}
_cleanup:
free(dict);
COVER_best_finish(data->best, parameters, selection);
free(data);
COVER_map_destroy(&activeDmers);
COVER_dictSelectionFree(selection);
free(freqs);
}
ZDICTLIB_STATIC_API size_t ZDICT_optimizeTrainFromBuffer_cover(
void* dictBuffer, size_t dictBufferCapacity, const void* samplesBuffer,
const size_t* samplesSizes, unsigned nbSamples,
ZDICT_cover_params_t* parameters)
{
/* constants */
const unsigned nbThreads = parameters->nbThreads;
const double splitPoint =
parameters->splitPoint <= 0.0 ? COVER_DEFAULT_SPLITPOINT : parameters->splitPoint;
const unsigned kMinD = parameters->d == 0 ? 6 : parameters->d;
const unsigned kMaxD = parameters->d == 0 ? 8 : parameters->d;
const unsigned kMinK = parameters->k == 0 ? 50 : parameters->k;
const unsigned kMaxK = parameters->k == 0 ? 2000 : parameters->k;
const unsigned kSteps = parameters->steps == 0 ? 40 : parameters->steps;
const unsigned kStepSize = MAX((kMaxK - kMinK) / kSteps, 1);
const unsigned kIterations =
(1 + (kMaxD - kMinD) / 2) * (1 + (kMaxK - kMinK) / kStepSize);
const unsigned shrinkDict = 0;
/* Local variables */
const int displayLevel = parameters->zParams.notificationLevel;
unsigned iteration = 1;
unsigned d;
unsigned k;
COVER_best_t best;
POOL_ctx *pool = NULL;
int warned = 0;
/* Checks */
if (splitPoint <= 0 || splitPoint > 1) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect parameters\n");
return ERROR(parameter_outOfBound);
}
if (kMinK < kMaxD || kMaxK < kMinK) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect parameters\n");
return ERROR(parameter_outOfBound);
}
if (nbSamples == 0) {
DISPLAYLEVEL(1, "Cover must have at least one input file\n");
return ERROR(srcSize_wrong);
}
if (dictBufferCapacity < ZDICT_DICTSIZE_MIN) {
DISPLAYLEVEL(1, "dictBufferCapacity must be at least %u\n",
ZDICT_DICTSIZE_MIN);
return ERROR(dstSize_tooSmall);
}
if (nbThreads > 1) {
pool = POOL_create(nbThreads, 1);
if (!pool) {
return ERROR(memory_allocation);
}
}
/* Initialization */
COVER_best_init(&best);
/* Turn down global display level to clean up display at level 2 and below */
g_displayLevel = displayLevel == 0 ? 0 : displayLevel - 1;
/* Loop through d first because each new value needs a new context */
LOCALDISPLAYLEVEL(displayLevel, 2, "Trying %u different sets of parameters\n",
kIterations);
for (d = kMinD; d <= kMaxD; d += 2) {
/* Initialize the context for this value of d */
COVER_ctx_t ctx;
LOCALDISPLAYLEVEL(displayLevel, 3, "d=%u\n", d);
{
const size_t initVal = COVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples, d, splitPoint);
if (ZSTD_isError(initVal)) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Failed to initialize context\n");
COVER_best_destroy(&best);
POOL_free(pool);
return initVal;
}
}
if (!warned) {
COVER_warnOnSmallCorpus(dictBufferCapacity, ctx.suffixSize, displayLevel);
warned = 1;
}
/* Loop through k reusing the same context */
for (k = kMinK; k <= kMaxK; k += kStepSize) {
/* Prepare the arguments */
COVER_tryParameters_data_t *data = (COVER_tryParameters_data_t *)malloc(
sizeof(COVER_tryParameters_data_t));
LOCALDISPLAYLEVEL(displayLevel, 3, "k=%u\n", k);
if (!data) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Failed to allocate parameters\n");
COVER_best_destroy(&best);
COVER_ctx_destroy(&ctx);
POOL_free(pool);
return ERROR(memory_allocation);
}
data->ctx = &ctx;
data->best = &best;
data->dictBufferCapacity = dictBufferCapacity;
data->parameters = *parameters;
data->parameters.k = k;
data->parameters.d = d;
data->parameters.splitPoint = splitPoint;
data->parameters.steps = kSteps;
data->parameters.shrinkDict = shrinkDict;
data->parameters.zParams.notificationLevel = g_displayLevel;
/* Check the parameters */
if (!COVER_checkParameters(data->parameters, dictBufferCapacity)) {
DISPLAYLEVEL(1, "Cover parameters incorrect\n");
free(data);
continue;
}
/* Call the function and pass ownership of data to it */
COVER_best_start(&best);
if (pool) {
POOL_add(pool, &COVER_tryParameters, data);
} else {
COVER_tryParameters(data);
}
/* Print status */
LOCALDISPLAYUPDATE(displayLevel, 2, "\r%u%% ",
(unsigned)((iteration * 100) / kIterations));
++iteration;
}
COVER_best_wait(&best);
COVER_ctx_destroy(&ctx);
}
LOCALDISPLAYLEVEL(displayLevel, 2, "\r%79s\r", "");
/* Fill the output buffer and parameters with output of the best parameters */
{
const size_t dictSize = best.dictSize;
if (ZSTD_isError(best.compressedSize)) {
const size_t compressedSize = best.compressedSize;
COVER_best_destroy(&best);
POOL_free(pool);
return compressedSize;
}
*parameters = best.parameters;
memcpy(dictBuffer, best.dict, dictSize);
COVER_best_destroy(&best);
POOL_free(pool);
return dictSize;
}
}
rust/src/dict_builder_cover.rs
+1307
View File
@@ -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));
}
}
rust/src/lib.rs
+2
View File
@@ -18,6 +18,8 @@ pub mod hist;
pub mod huf_compress; pub mod huf_compress;
#[cfg(feature = "decompression")] #[cfg(feature = "decompression")]
pub mod huf_decompress; pub mod huf_decompress;
#[cfg(feature = "dict-builder")]
pub mod dict_builder_cover;
pub mod legacy; pub mod legacy;
pub mod mem; pub mod mem;
pub mod pool; pub mod pool;