From 468ccfa38828513622e3ff7654dc4cc22f78ee29 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sun, 12 Jul 2026 10:08:03 +0200 Subject: [PATCH] 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 --- lib/dictBuilder/cover.c | 1296 +------------------------------ rust/src/dict_builder_cover.rs | 1307 ++++++++++++++++++++++++++++++++ rust/src/lib.rs | 2 + 3 files changed, 1314 insertions(+), 1291 deletions(-) create mode 100644 rust/src/dict_builder_cover.rs diff --git a/lib/dictBuilder/cover.c b/lib/dictBuilder/cover.c index 2ef33c73e..45dd23704 100644 --- a/lib/dictBuilder/cover.c +++ b/lib/dictBuilder/cover.c @@ -8,1295 +8,9 @@ * You may select, at your option, one of the above-listed licenses. */ -/* ***************************************************************************** - * Constructs a dictionary using a heuristic based on the following paper: - * - * Liao, Petri, Moffat, Wirth - * 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 /* fprintf */ -#include /* malloc, free, qsort_r */ - -#include /* memset */ -#include /* 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" +/* + * COVER is implemented by rust/src/dict_builder_cover.rs. Keep this + * translation unit so existing C build descriptions continue to include the + * public COVER ABI header while the Rust static library supplies the symbols. + */ #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; - } -} diff --git a/rust/src/dict_builder_cover.rs b/rust/src/dict_builder_cover.rs new file mode 100644 index 000000000..d52ca62be --- /dev/null +++ b/rust/src/dict_builder_cover.rs @@ -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, + size_log: u32, + size_mask: u32, +} + +struct CoverContext { + samples: *const u8, + samples_sizes: *const usize, + offsets: Vec, + nb_samples: usize, + nb_train_samples: usize, + suffix_size: usize, + freqs: Vec, + dmer_at: Vec, + 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, + 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>> { + static SLOTS: OnceLock>>> = 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(len: usize) -> Result, ()> { + 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::() +} + +#[inline] +unsafe fn free_bytes(ptr: *mut u8) { + unsafe { libc::free(ptr.cast::()) }; +} + +#[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 { + let size_log = 32 - size.leading_zeros() + 1; + let table_size = 1usize.checked_shl(size_log).ok_or(())?; + let mut data = try_vec::(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 { + 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::()); + let max_samples_size = if size_of::() == 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::(suffix_size).map_err(|_| ERROR_MEMORY_ALLOCATION)?; + let dmer_at = try_vec::(suffix_size).map_err(|_| ERROR_MEMORY_ALLOCATION)?; + let mut offsets = try_vec::(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> { + 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::().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::(); + 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)); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 7b30d4504..9661f848a 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -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;