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

feat(dict): port the FastCover trainer to Rust

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

Test Plan:
- rustfmt +nightly --check --edition 2021 rust/src/dict_builder_fastcover.rs rust/src/lib.rs
- RUSTC_WRAPPER= CARGO_BUILD_RUSTC_WRAPPER= cargo clippy --manifest-path rust/Cargo.toml --all-targets --no-default-features --features compression,dict-builder -- -D warnings
- FastCover focused Rust tests and deterministic C-reference harness
- git diff --cached --check
This commit is contained in:
ddidderr
2026-07-12 10:40:51 +02:00
parent 60877175a4
commit b8ccfd7b87
3 changed files with 1080 additions and 751 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/dictBuilder/fastcover.c
+6 -751
View File
@@ -8,759 +8,14 @@
* You may select, at your option, one of the above-listed licenses.
*/
/*-*************************************
* Dependencies
***************************************/
#include <stdio.h> /* fprintf */
#include <stdlib.h> /* malloc, free, qsort */
#include <string.h> /* memset */
#include <time.h> /* clock */
/*
* FastCOVER is implemented by rust/src/dict_builder_fastcover.rs. Keep this
* translation unit so existing C build descriptions continue to include the
* public static-only declarations while the Rust static library supplies the
* symbols.
*/
#ifndef ZDICT_STATIC_LINKING_ONLY
# define ZDICT_STATIC_LINKING_ONLY
#endif
#include "../common/mem.h" /* read */
#include "../common/pool.h"
#include "../common/threading.h"
#include "../common/zstd_internal.h" /* includes zstd.h */
#include "../compress/zstd_compress_internal.h" /* ZSTD_hash*() */
#include "../zdict.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 FASTCOVER_MAX_SAMPLES_SIZE (sizeof(size_t) == 8 ? ((unsigned)-1) : ((unsigned)1 GB))
#define FASTCOVER_MAX_F 31
#define FASTCOVER_MAX_ACCEL 10
#define FASTCOVER_DEFAULT_SPLITPOINT 0.75
#define DEFAULT_F 20
#define DEFAULT_ACCEL 1
/*-*************************************
* 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 Functions
***************************************/
/**
* Hash the d-byte value pointed to by p and mod 2^f into the frequency vector
*/
static size_t FASTCOVER_hashPtrToIndex(const void* p, U32 f, unsigned d) {
if (d == 6) {
return ZSTD_hash6Ptr(p, f);
}
return ZSTD_hash8Ptr(p, f);
}
/*-*************************************
* Acceleration
***************************************/
typedef struct {
unsigned finalize; /* Percentage of training samples used for ZDICT_finalizeDictionary */
unsigned skip; /* Number of dmer skipped between each dmer counted in computeFrequency */
} FASTCOVER_accel_t;
static const FASTCOVER_accel_t FASTCOVER_defaultAccelParameters[FASTCOVER_MAX_ACCEL+1] = {
{ 100, 0 }, /* accel = 0, should not happen because accel = 0 defaults to accel = 1 */
{ 100, 0 }, /* accel = 1 */
{ 50, 1 }, /* accel = 2 */
{ 34, 2 }, /* accel = 3 */
{ 25, 3 }, /* accel = 4 */
{ 20, 4 }, /* accel = 5 */
{ 17, 5 }, /* accel = 6 */
{ 14, 6 }, /* accel = 7 */
{ 13, 7 }, /* accel = 8 */
{ 11, 8 }, /* accel = 9 */
{ 10, 9 }, /* accel = 10 */
};
/*-*************************************
* Context
***************************************/
typedef struct {
const BYTE *samples;
size_t *offsets;
const size_t *samplesSizes;
size_t nbSamples;
size_t nbTrainSamples;
size_t nbTestSamples;
size_t nbDmers;
U32 *freqs;
unsigned d;
unsigned f;
FASTCOVER_accel_t accelParams;
} FASTCOVER_ctx_t;
/*-*************************************
* Helper functions
***************************************/
/**
* Selects the best segment in an epoch.
* Segments of are scored according to the function:
*
* Let F(d) be the frequency of all dmers with hash value d.
* Let S_i be hash value of 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 with hash value d is in the dictionary we set F(d) = 0.
*/
static COVER_segment_t FASTCOVER_selectSegment(const FASTCOVER_ctx_t *ctx,
U32 *freqs, U32 begin, U32 end,
ZDICT_cover_params_t parameters,
U16* segmentFreqs) {
/* Constants */
const U32 k = parameters.k;
const U32 d = parameters.d;
const U32 f = ctx->f;
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 */
/* 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) {
/* Get hash value of current dmer */
const size_t idx = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.end, f, d);
/* Add frequency of this index to score if this is the first occurrence of index in active segment */
if (segmentFreqs[idx] == 0) {
activeSegment.score += freqs[idx];
}
/* Increment end of segment and segmentFreqs*/
activeSegment.end += 1;
segmentFreqs[idx] += 1;
/* If the window is now too large, drop the first position */
if (activeSegment.end - activeSegment.begin == dmersInK + 1) {
/* Get hash value of the dmer to be eliminated from active segment */
const size_t delIndex = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.begin, f, d);
segmentFreqs[delIndex] -= 1;
/* Subtract frequency of this index from score if this is the last occurrence of this index in active segment */
if (segmentFreqs[delIndex] == 0) {
activeSegment.score -= freqs[delIndex];
}
/* Increment start of segment */
activeSegment.begin += 1;
}
/* If this segment is the best so far save it */
if (activeSegment.score > bestSegment.score) {
bestSegment = activeSegment;
}
}
/* Zero out rest of segmentFreqs array */
while (activeSegment.begin < end) {
const size_t delIndex = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.begin, f, d);
segmentFreqs[delIndex] -= 1;
activeSegment.begin += 1;
}
{
/* Zero the frequency of hash value of each dmer covered by the chosen segment. */
U32 pos;
for (pos = bestSegment.begin; pos != bestSegment.end; ++pos) {
const size_t i = FASTCOVER_hashPtrToIndex(ctx->samples + pos, f, d);
freqs[i] = 0;
}
}
return bestSegment;
}
static int FASTCOVER_checkParameters(ZDICT_cover_params_t parameters,
size_t maxDictSize, unsigned f,
unsigned accel) {
/* k, d, and f are required parameters */
if (parameters.d == 0 || parameters.k == 0) {
return 0;
}
/* d has to be 6 or 8 */
if (parameters.d != 6 && parameters.d != 8) {
return 0;
}
/* k <= maxDictSize */
if (parameters.k > maxDictSize) {
return 0;
}
/* d <= k */
if (parameters.d > parameters.k) {
return 0;
}
/* 0 < f <= FASTCOVER_MAX_F*/
if (f > FASTCOVER_MAX_F || f == 0) {
return 0;
}
/* 0 < splitPoint <= 1 */
if (parameters.splitPoint <= 0 || parameters.splitPoint > 1) {
return 0;
}
/* 0 < accel <= 10 */
if (accel > 10 || accel == 0) {
return 0;
}
return 1;
}
/**
* Clean up a context initialized with `FASTCOVER_ctx_init()`.
*/
static void
FASTCOVER_ctx_destroy(FASTCOVER_ctx_t* ctx)
{
if (!ctx) return;
free(ctx->freqs);
ctx->freqs = NULL;
free(ctx->offsets);
ctx->offsets = NULL;
}
/**
* Calculate for frequency of hash value of each dmer in ctx->samples
*/
static void
FASTCOVER_computeFrequency(U32* freqs, const FASTCOVER_ctx_t* ctx)
{
const unsigned f = ctx->f;
const unsigned d = ctx->d;
const unsigned skip = ctx->accelParams.skip;
const unsigned readLength = MAX(d, 8);
size_t i;
assert(ctx->nbTrainSamples >= 5);
assert(ctx->nbTrainSamples <= ctx->nbSamples);
for (i = 0; i < ctx->nbTrainSamples; i++) {
size_t start = ctx->offsets[i]; /* start of current dmer */
size_t const currSampleEnd = ctx->offsets[i+1];
while (start + readLength <= currSampleEnd) {
const size_t dmerIndex = FASTCOVER_hashPtrToIndex(ctx->samples + start, f, d);
freqs[dmerIndex]++;
start = start + skip + 1;
}
}
}
/**
* 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 `FASTCOVER_ctx_destroy()`.
*/
static size_t
FASTCOVER_ctx_init(FASTCOVER_ctx_t* ctx,
const void* samplesBuffer,
const size_t* samplesSizes, unsigned nbSamples,
unsigned d, double splitPoint, unsigned f,
FASTCOVER_accel_t accelParams)
{
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)FASTCOVER_MAX_SAMPLES_SIZE) {
DISPLAYLEVEL(1, "Total samples size is too large (%u MB), maximum size is %u MB\n",
(unsigned)(totalSamplesSize >> 20), (FASTCOVER_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\n", 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.\n", 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;
ctx->nbDmers = trainingSamplesSize - MAX(d, sizeof(U64)) + 1;
ctx->d = d;
ctx->f = f;
ctx->accelParams = accelParams;
/* The offsets of each file */
ctx->offsets = (size_t*)calloc((nbSamples + 1), sizeof(size_t));
if (ctx->offsets == NULL) {
DISPLAYLEVEL(1, "Failed to allocate scratch buffers \n");
FASTCOVER_ctx_destroy(ctx);
return ERROR(memory_allocation);
}
/* Fill offsets from the samplesSizes */
{ U32 i;
ctx->offsets[0] = 0;
assert(nbSamples >= 5);
for (i = 1; i <= nbSamples; ++i) {
ctx->offsets[i] = ctx->offsets[i - 1] + samplesSizes[i - 1];
}
}
/* Initialize frequency array of size 2^f */
ctx->freqs = (U32*)calloc(((U64)1 << f), sizeof(U32));
if (ctx->freqs == NULL) {
DISPLAYLEVEL(1, "Failed to allocate frequency table \n");
FASTCOVER_ctx_destroy(ctx);
return ERROR(memory_allocation);
}
DISPLAYLEVEL(2, "Computing frequencies\n");
FASTCOVER_computeFrequency(ctx->freqs, ctx);
return 0;
}
/**
* Given the prepared context build the dictionary.
*/
static size_t
FASTCOVER_buildDictionary(const FASTCOVER_ctx_t* ctx,
U32* freqs,
void* dictBuffer, size_t dictBufferCapacity,
ZDICT_cover_params_t parameters,
U16* segmentFreqs)
{
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->nbDmers, parameters.k, 1);
const size_t maxZeroScoreRun = 10;
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 = FASTCOVER_selectSegment(
ctx, freqs, epochBegin, epochEnd, parameters, segmentFreqs);
/* 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;
}
/**
* Parameters for FASTCOVER_tryParameters().
*/
typedef struct FASTCOVER_tryParameters_data_s {
const FASTCOVER_ctx_t* ctx;
COVER_best_t* best;
size_t dictBufferCapacity;
ZDICT_cover_params_t parameters;
} FASTCOVER_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 FASTCOVER_tryParameters(void* opaque)
{
/* Save parameters as local variables */
FASTCOVER_tryParameters_data_t *const data = (FASTCOVER_tryParameters_data_t*)opaque;
const FASTCOVER_ctx_t *const ctx = data->ctx;
const ZDICT_cover_params_t parameters = data->parameters;
size_t dictBufferCapacity = data->dictBufferCapacity;
size_t totalCompressedSize = ERROR(GENERIC);
/* Initialize array to keep track of frequency of dmer within activeSegment */
U16* segmentFreqs = (U16*)calloc(((U64)1 << ctx->f), sizeof(U16));
/* Allocate space for hash table, dict, and freqs */
BYTE *const dict = (BYTE*)malloc(dictBufferCapacity);
COVER_dictSelection_t selection = COVER_dictSelectionError(ERROR(GENERIC));
U32* freqs = (U32*) malloc(((U64)1 << ctx->f) * sizeof(U32));
if (!segmentFreqs || !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, ((U64)1 << ctx->f) * sizeof(U32));
/* Build the dictionary */
{ const size_t tail = FASTCOVER_buildDictionary(ctx, freqs, dict, dictBufferCapacity,
parameters, segmentFreqs);
const unsigned nbFinalizeSamples = (unsigned)(ctx->nbTrainSamples * ctx->accelParams.finalize / 100);
selection = COVER_selectDict(dict + tail, dictBufferCapacity, dictBufferCapacity - tail,
ctx->samples, ctx->samplesSizes, nbFinalizeSamples, 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);
free(segmentFreqs);
COVER_dictSelectionFree(selection);
free(freqs);
}
static void
FASTCOVER_convertToCoverParams(ZDICT_fastCover_params_t fastCoverParams,
ZDICT_cover_params_t* coverParams)
{
coverParams->k = fastCoverParams.k;
coverParams->d = fastCoverParams.d;
coverParams->steps = fastCoverParams.steps;
coverParams->nbThreads = fastCoverParams.nbThreads;
coverParams->splitPoint = fastCoverParams.splitPoint;
coverParams->zParams = fastCoverParams.zParams;
coverParams->shrinkDict = fastCoverParams.shrinkDict;
}
static void
FASTCOVER_convertToFastCoverParams(ZDICT_cover_params_t coverParams,
ZDICT_fastCover_params_t* fastCoverParams,
unsigned f, unsigned accel)
{
fastCoverParams->k = coverParams.k;
fastCoverParams->d = coverParams.d;
fastCoverParams->steps = coverParams.steps;
fastCoverParams->nbThreads = coverParams.nbThreads;
fastCoverParams->splitPoint = coverParams.splitPoint;
fastCoverParams->f = f;
fastCoverParams->accel = accel;
fastCoverParams->zParams = coverParams.zParams;
fastCoverParams->shrinkDict = coverParams.shrinkDict;
}
ZDICTLIB_STATIC_API size_t
ZDICT_trainFromBuffer_fastCover(void* dictBuffer, size_t dictBufferCapacity,
const void* samplesBuffer,
const size_t* samplesSizes, unsigned nbSamples,
ZDICT_fastCover_params_t parameters)
{
BYTE* const dict = (BYTE*)dictBuffer;
FASTCOVER_ctx_t ctx;
ZDICT_cover_params_t coverParams;
FASTCOVER_accel_t accelParams;
/* Initialize global data */
g_displayLevel = (int)parameters.zParams.notificationLevel;
/* Assign splitPoint and f if not provided */
parameters.splitPoint = 1.0;
parameters.f = parameters.f == 0 ? DEFAULT_F : parameters.f;
parameters.accel = parameters.accel == 0 ? DEFAULT_ACCEL : parameters.accel;
/* Convert to cover parameter */
memset(&coverParams, 0 , sizeof(coverParams));
FASTCOVER_convertToCoverParams(parameters, &coverParams);
/* Checks */
if (!FASTCOVER_checkParameters(coverParams, dictBufferCapacity, parameters.f,
parameters.accel)) {
DISPLAYLEVEL(1, "FASTCOVER parameters incorrect\n");
return ERROR(parameter_outOfBound);
}
if (nbSamples == 0) {
DISPLAYLEVEL(1, "FASTCOVER 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);
}
/* Assign corresponding FASTCOVER_accel_t to accelParams*/
accelParams = FASTCOVER_defaultAccelParameters[parameters.accel];
/* Initialize context */
{
size_t const initVal = FASTCOVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples,
coverParams.d, parameters.splitPoint, parameters.f,
accelParams);
if (ZSTD_isError(initVal)) {
DISPLAYLEVEL(1, "Failed to initialize context\n");
return initVal;
}
}
COVER_warnOnSmallCorpus(dictBufferCapacity, ctx.nbDmers, g_displayLevel);
/* Build the dictionary */
DISPLAYLEVEL(2, "Building dictionary\n");
{
/* Initialize array to keep track of frequency of dmer within activeSegment */
U16* segmentFreqs = (U16 *)calloc(((U64)1 << parameters.f), sizeof(U16));
const size_t tail = FASTCOVER_buildDictionary(&ctx, ctx.freqs, dictBuffer,
dictBufferCapacity, coverParams, segmentFreqs);
const unsigned nbFinalizeSamples = (unsigned)(ctx.nbTrainSamples * ctx.accelParams.finalize / 100);
const size_t dictionarySize = ZDICT_finalizeDictionary(
dict, dictBufferCapacity, dict + tail, dictBufferCapacity - tail,
samplesBuffer, samplesSizes, nbFinalizeSamples, coverParams.zParams);
if (!ZSTD_isError(dictionarySize)) {
DISPLAYLEVEL(2, "Constructed dictionary of size %u\n",
(unsigned)dictionarySize);
}
FASTCOVER_ctx_destroy(&ctx);
free(segmentFreqs);
return dictionarySize;
}
}
ZDICTLIB_STATIC_API size_t
ZDICT_optimizeTrainFromBuffer_fastCover(
void* dictBuffer, size_t dictBufferCapacity,
const void* samplesBuffer,
const size_t* samplesSizes, unsigned nbSamples,
ZDICT_fastCover_params_t* parameters)
{
ZDICT_cover_params_t coverParams;
FASTCOVER_accel_t accelParams;
/* constants */
const unsigned nbThreads = parameters->nbThreads;
const double splitPoint =
parameters->splitPoint <= 0.0 ? FASTCOVER_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 f = parameters->f == 0 ? DEFAULT_F : parameters->f;
const unsigned accel = parameters->accel == 0 ? DEFAULT_ACCEL : parameters->accel;
const unsigned shrinkDict = 0;
/* Local variables */
const int displayLevel = (int)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 splitPoint\n");
return ERROR(parameter_outOfBound);
}
if (accel == 0 || accel > FASTCOVER_MAX_ACCEL) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect accel\n");
return ERROR(parameter_outOfBound);
}
if (kMinK < kMaxD || kMaxK < kMinK) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect k\n");
return ERROR(parameter_outOfBound);
}
if (nbSamples == 0) {
LOCALDISPLAYLEVEL(displayLevel, 1, "FASTCOVER must have at least one input file\n");
return ERROR(srcSize_wrong);
}
if (dictBufferCapacity < ZDICT_DICTSIZE_MIN) {
LOCALDISPLAYLEVEL(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);
memset(&coverParams, 0 , sizeof(coverParams));
FASTCOVER_convertToCoverParams(*parameters, &coverParams);
accelParams = FASTCOVER_defaultAccelParameters[accel];
/* 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 */
FASTCOVER_ctx_t ctx;
LOCALDISPLAYLEVEL(displayLevel, 3, "d=%u\n", d);
{
size_t const initVal = FASTCOVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples, d, splitPoint, f, accelParams);
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.nbDmers, displayLevel);
warned = 1;
}
/* Loop through k reusing the same context */
for (k = kMinK; k <= kMaxK; k += kStepSize) {
/* Prepare the arguments */
FASTCOVER_tryParameters_data_t *data = (FASTCOVER_tryParameters_data_t *)malloc(
sizeof(FASTCOVER_tryParameters_data_t));
LOCALDISPLAYLEVEL(displayLevel, 3, "k=%u\n", k);
if (!data) {
LOCALDISPLAYLEVEL(displayLevel, 1, "Failed to allocate parameters\n");
COVER_best_destroy(&best);
FASTCOVER_ctx_destroy(&ctx);
POOL_free(pool);
return ERROR(memory_allocation);
}
data->ctx = &ctx;
data->best = &best;
data->dictBufferCapacity = dictBufferCapacity;
data->parameters = coverParams;
data->parameters.k = k;
data->parameters.d = d;
data->parameters.splitPoint = splitPoint;
data->parameters.steps = kSteps;
data->parameters.shrinkDict = shrinkDict;
data->parameters.zParams.notificationLevel = (unsigned)g_displayLevel;
/* Check the parameters */
if (!FASTCOVER_checkParameters(data->parameters, dictBufferCapacity,
data->ctx->f, accel)) {
DISPLAYLEVEL(1, "FASTCOVER 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, &FASTCOVER_tryParameters, data);
} else {
FASTCOVER_tryParameters(data);
}
/* Print status */
LOCALDISPLAYUPDATE(displayLevel, 2, "\r%u%% ",
(unsigned)((iteration * 100) / kIterations));
++iteration;
}
COVER_best_wait(&best);
FASTCOVER_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;
}
FASTCOVER_convertToFastCoverParams(best.parameters, parameters, f, accel);
memcpy(dictBuffer, best.dict, dictSize);
COVER_best_destroy(&best);
POOL_free(pool);
return dictSize;
}
}
rust/src/dict_builder_fastcover.rs
+1072
View File
@@ -0,0 +1,1072 @@
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(clippy::missing_safety_doc)]
#![allow(clippy::too_many_arguments)]
//! FastCOVER dictionary training.
//!
//! The public FastCOVER declarations live in `lib/zdict.h`. This module is a
//! Rust translation of `lib/dictBuilder/fastcover.c`; the C translation unit
//! remains in the native source lists as a declaration-only ABI anchor.
use crate::dict_builder_cover::{
COVER_computeEpochs, COVER_dictSelectionError, COVER_dictSelectionFree,
COVER_dictSelectionIsError, COVER_dictSelection_t, COVER_segment_t, COVER_selectDict,
ZDICT_cover_params_t, ZDICT_params_t,
};
use crate::errors::ERR_isError;
use crate::mem::MEM_readLE64;
use std::ffi::c_void;
use std::mem::size_of;
use std::os::raw::{c_int, c_uint};
use std::ptr;
use std::sync::atomic::{AtomicI32, Ordering as AtomicOrdering};
use std::sync::{Arc, Condvar, Mutex};
const ZDICT_DICTSIZE_MIN: usize = 256;
const FASTCOVER_MAX_SAMPLES_SIZE_64: usize = u32::MAX as usize;
const FASTCOVER_MAX_SAMPLES_SIZE_32: usize = 1usize << 30;
const FASTCOVER_MAX_F: c_uint = 31;
const FASTCOVER_MAX_ACCEL: c_uint = 10;
const FASTCOVER_DEFAULT_SPLITPOINT: f64 = 0.75;
const DEFAULT_F: c_uint = 20;
const DEFAULT_ACCEL: c_uint = 1;
const ERROR_GENERIC: usize = 0usize.wrapping_sub(1);
const ERROR_MEMORY_ALLOCATION: usize = 0usize.wrapping_sub(64);
const ERROR_PARAMETER_OUT_OF_BOUND: usize = 0usize.wrapping_sub(42);
const ERROR_SRC_SIZE_WRONG: usize = 0usize.wrapping_sub(72);
const ERROR_DST_SIZE_TOO_SMALL: usize = 0usize.wrapping_sub(70);
const PRIME6BYTES: u64 = 227_718_039_650_203;
const PRIME8BYTES: u64 = 0xCF1B_BCDC_B7A5_6463;
static DISPLAY_LEVEL: AtomicI32 = AtomicI32::new(0);
#[repr(C)]
#[derive(Clone, Copy, Default)]
/// ABI-compatible `ZDICT_fastCover_params_t` from `zdict.h`.
pub struct ZDICT_fastCover_params_t {
pub k: c_uint,
pub d: c_uint,
pub f: c_uint,
pub steps: c_uint,
pub nbThreads: c_uint,
pub splitPoint: f64,
pub accel: c_uint,
pub shrinkDict: c_uint,
pub shrinkDictMaxRegression: c_uint,
pub zParams: ZDICT_params_t,
}
#[derive(Clone, Copy)]
struct FastCoverAccel {
finalize: c_uint,
skip: c_uint,
}
const DEFAULT_ACCEL_PARAMETERS: [FastCoverAccel; 11] = [
FastCoverAccel {
finalize: 100,
skip: 0,
},
FastCoverAccel {
finalize: 100,
skip: 0,
},
FastCoverAccel {
finalize: 50,
skip: 1,
},
FastCoverAccel {
finalize: 34,
skip: 2,
},
FastCoverAccel {
finalize: 25,
skip: 3,
},
FastCoverAccel {
finalize: 20,
skip: 4,
},
FastCoverAccel {
finalize: 17,
skip: 5,
},
FastCoverAccel {
finalize: 14,
skip: 6,
},
FastCoverAccel {
finalize: 13,
skip: 7,
},
FastCoverAccel {
finalize: 11,
skip: 8,
},
FastCoverAccel {
finalize: 10,
skip: 9,
},
];
struct FastCoverContext {
samples: *const u8,
offsets: *mut usize,
samples_sizes: *const usize,
nb_samples: usize,
nb_train_samples: usize,
_nb_test_samples: usize,
nb_dmers: usize,
freqs: *mut u32,
d: c_uint,
f: c_uint,
accel_params: FastCoverAccel,
}
// The input buffer and sample sizes are immutable. Each worker only mutates
// its own frequency copy and segment-frequency table.
unsafe impl Send for FastCoverContext {}
unsafe impl Sync for FastCoverContext {}
impl Drop for FastCoverContext {
fn drop(&mut self) {
unsafe {
free_bytes(self.freqs);
free_bytes(self.offsets);
}
self.freqs = ptr::null_mut();
self.offsets = ptr::null_mut();
}
}
#[inline]
fn display_level() -> c_int {
DISPLAY_LEVEL.load(AtomicOrdering::Relaxed)
}
#[inline]
fn set_display_level(level: c_uint) {
DISPLAY_LEVEL.store(level as c_int, AtomicOrdering::Relaxed);
}
#[inline]
fn display(level: c_int, message: &str) {
if display_level() >= level {
eprint!("{message}");
}
}
#[inline]
unsafe fn free_bytes<T>(allocation: *mut T) {
unsafe { libc::free(allocation.cast::<c_void>()) };
}
#[inline]
unsafe fn malloc_bytes(size: usize) -> *mut u8 {
unsafe { libc::malloc(size) }.cast::<u8>()
}
#[inline]
unsafe fn calloc_array<T>(count: usize) -> *mut T {
unsafe { libc::calloc(count, size_of::<T>()) }.cast::<T>()
}
#[inline]
fn frequency_count(f: c_uint) -> usize {
// f is validated before this helper is used. The wrapping shift keeps
// the C unsigned-shift behavior for malformed internal callers.
1usize.wrapping_shl(f)
}
#[inline]
fn frequency_bytes(f: c_uint, element_size: usize) -> usize {
((1u64 << f) * element_size as u64) as usize
}
#[inline]
unsafe fn sample_sum(samples_sizes: *const usize, nb_samples: usize) -> usize {
let mut total = 0usize;
for index in 0..nb_samples {
total = total.wrapping_add(unsafe { *samples_sizes.add(index) });
}
total
}
#[inline]
unsafe fn hash_ptr_to_index(pointer: *const u8, f: c_uint, d: c_uint) -> usize {
let value = unsafe { MEM_readLE64(pointer.cast::<c_void>()) };
if d == 6 {
((value << 16).wrapping_mul(PRIME6BYTES) >> (64 - f)) as usize
} else {
(value.wrapping_mul(PRIME8BYTES) >> (64 - f)) as usize
}
}
fn check_parameters(
parameters: ZDICT_cover_params_t,
max_dict_size: usize,
f: c_uint,
accel: c_uint,
) -> bool {
if parameters.d == 0 || parameters.k == 0 {
return false;
}
if parameters.d != 6 && parameters.d != 8 {
return false;
}
if parameters.k as usize > max_dict_size {
return false;
}
if parameters.d > parameters.k {
return false;
}
if f > FASTCOVER_MAX_F || f == 0 {
return false;
}
if parameters.splitPoint <= 0.0 || parameters.splitPoint > 1.0 {
return false;
}
if accel > FASTCOVER_MAX_ACCEL || accel == 0 {
return false;
}
true
}
unsafe fn compute_frequency(freqs: *mut u32, context: &FastCoverContext) {
let read_length = context.d.max(8) as usize;
for sample in 0..context.nb_train_samples {
let mut start = unsafe { *context.offsets.add(sample) };
let sample_end = unsafe { *context.offsets.add(sample + 1) };
while start.wrapping_add(read_length) <= sample_end {
let index =
unsafe { hash_ptr_to_index(context.samples.add(start), context.f, context.d) };
let frequency = unsafe { freqs.add(index) };
unsafe { *frequency = (*frequency).wrapping_add(1) };
start = start.wrapping_add(context.accel_params.skip as usize + 1);
}
}
}
unsafe fn context_init(
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: usize,
d: c_uint,
split_point: f64,
f: c_uint,
accel_params: FastCoverAccel,
) -> Result<FastCoverContext, usize> {
let samples = samples_buffer.cast::<u8>();
let total_samples_size = unsafe { sample_sum(samples_sizes, nb_samples) };
let nb_train_samples = if split_point < 1.0 {
((nb_samples as f64) * split_point) as usize
} else {
nb_samples
};
let nb_test_samples = if split_point < 1.0 {
nb_samples - nb_train_samples
} else {
nb_samples
};
let training_samples_size = if split_point < 1.0 {
unsafe { sample_sum(samples_sizes, nb_train_samples) }
} else {
total_samples_size
};
let _test_samples_size = if split_point < 1.0 {
unsafe { sample_sum(samples_sizes.add(nb_train_samples), nb_test_samples) }
} else {
total_samples_size
};
let max_dmer = (d as usize).max(size_of::<u64>());
let max_samples_size = if size_of::<usize>() == 8 {
FASTCOVER_MAX_SAMPLES_SIZE_64
} else {
FASTCOVER_MAX_SAMPLES_SIZE_32
};
if total_samples_size < max_dmer || total_samples_size >= max_samples_size {
display(1, "Total samples size is too large\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
if nb_train_samples < 5 {
display(1, "Total number of training samples is invalid\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
if nb_test_samples < 1 {
display(1, "Total number of testing samples is invalid\n");
return Err(ERROR_SRC_SIZE_WRONG);
}
let offsets_count = nb_samples.wrapping_add(1);
let offsets = unsafe { calloc_array::<usize>(offsets_count) };
if offsets.is_null() {
display(1, "Failed to allocate scratch buffers\n");
return Err(ERROR_MEMORY_ALLOCATION);
}
unsafe {
*offsets = 0;
for index in 0..nb_samples {
*offsets.add(index + 1) = (*offsets.add(index)).wrapping_add(*samples_sizes.add(index));
}
}
let freqs = unsafe { calloc_array::<u32>(frequency_count(f)) };
if freqs.is_null() {
unsafe { free_bytes(offsets) };
display(1, "Failed to allocate frequency table\n");
return Err(ERROR_MEMORY_ALLOCATION);
}
let context = FastCoverContext {
samples,
offsets,
samples_sizes,
nb_samples,
nb_train_samples,
_nb_test_samples: nb_test_samples,
nb_dmers: training_samples_size.wrapping_sub(max_dmer).wrapping_add(1),
freqs,
d,
f,
accel_params,
};
unsafe { compute_frequency(context.freqs, &context) };
Ok(context)
}
unsafe fn select_segment(
context: &FastCoverContext,
freqs: *mut u32,
begin: u32,
end: u32,
parameters: ZDICT_cover_params_t,
segment_freqs: *mut u16,
) -> COVER_segment_t {
let dmers_in_k = parameters.k - parameters.d + 1;
let mut best = COVER_segment_t::default();
let mut active = COVER_segment_t {
begin,
end: begin,
score: 0,
};
while active.end < end {
let index = unsafe {
hash_ptr_to_index(
context.samples.add(active.end as usize),
context.f,
context.d,
)
};
let active_frequency = unsafe { segment_freqs.add(index) };
if unsafe { *active_frequency } == 0 {
active.score = active.score.wrapping_add(unsafe { *freqs.add(index) });
}
active.end = active.end.wrapping_add(1);
unsafe { *active_frequency = (*active_frequency).wrapping_add(1) };
if active.end - active.begin == dmers_in_k + 1 {
let deleted_index = unsafe {
hash_ptr_to_index(
context.samples.add(active.begin as usize),
context.f,
context.d,
)
};
let deleted_frequency = unsafe { segment_freqs.add(deleted_index) };
unsafe { *deleted_frequency = (*deleted_frequency).wrapping_sub(1) };
if unsafe { *deleted_frequency } == 0 {
active.score = active
.score
.wrapping_sub(unsafe { *freqs.add(deleted_index) });
}
active.begin = active.begin.wrapping_add(1);
}
if active.score > best.score {
best = active;
}
}
while active.begin < end {
let index = unsafe {
hash_ptr_to_index(
context.samples.add(active.begin as usize),
context.f,
context.d,
)
};
let frequency = unsafe { segment_freqs.add(index) };
unsafe { *frequency = (*frequency).wrapping_sub(1) };
active.begin = active.begin.wrapping_add(1);
}
let mut position = best.begin;
while position != best.end {
let index = unsafe {
hash_ptr_to_index(context.samples.add(position as usize), context.f, context.d)
};
unsafe { *freqs.add(index) = 0 };
position = position.wrapping_add(1);
}
best
}
unsafe fn build_dictionary(
context: &FastCoverContext,
freqs: *mut u32,
dict_buffer: *mut u8,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
segment_freqs: *mut u16,
) -> usize {
let epochs = COVER_computeEpochs(
dict_buffer_capacity as u32,
context.nb_dmers as u32,
parameters.k,
1,
);
if epochs.num == 0 {
return dict_buffer_capacity;
}
let mut tail = dict_buffer_capacity;
let mut zero_score_run = 0usize;
let mut epoch = 0usize;
while tail > 0 {
let epoch_begin = epoch.wrapping_mul(epochs.size as usize) as u32;
let epoch_end = epoch_begin.wrapping_add(epochs.size);
let segment = unsafe {
select_segment(
context,
freqs,
epoch_begin,
epoch_end,
parameters,
segment_freqs,
)
};
if segment.score == 0 {
zero_score_run += 1;
if zero_score_run >= 10 {
break;
}
} else {
zero_score_run = 0;
let segment_size = (segment
.end
.wrapping_sub(segment.begin)
.wrapping_add(parameters.d)
.wrapping_sub(1) as usize)
.min(tail);
if segment_size < parameters.d as usize {
break;
}
tail -= segment_size;
unsafe {
ptr::copy_nonoverlapping(
context.samples.add(segment.begin as usize),
dict_buffer.add(tail),
segment_size,
);
}
}
epoch = (epoch + 1) % epochs.num as usize;
}
tail
}
#[derive(Clone, Copy)]
struct FastBestState {
dict: *mut u8,
dict_size: usize,
parameters: ZDICT_cover_params_t,
compressed_size: usize,
live_jobs: usize,
}
unsafe impl Send for FastBestState {}
struct FastBest {
state: Mutex<FastBestState>,
cond: Condvar,
}
impl FastBest {
fn new() -> Self {
Self {
state: Mutex::new(FastBestState {
dict: ptr::null_mut(),
dict_size: 0,
parameters: ZDICT_cover_params_t::default(),
compressed_size: ERROR_GENERIC,
live_jobs: 0,
}),
cond: Condvar::new(),
}
}
fn start(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
state.live_jobs = state.live_jobs.wrapping_add(1);
}
fn finish(&self, parameters: ZDICT_cover_params_t, selection: COVER_dictSelection_t) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
state.live_jobs = state.live_jobs.wrapping_sub(1);
if selection.totalCompressedSize < state.compressed_size {
if state.dict.is_null() || state.dict_size < selection.dictSize {
unsafe { free_bytes(state.dict) };
let replacement = unsafe { malloc_bytes(selection.dictSize) };
if replacement.is_null() && selection.dictSize != 0 {
state.dict = ptr::null_mut();
state.dict_size = 0;
state.compressed_size = ERROR_GENERIC;
self.cond.notify_one();
return;
}
state.dict = replacement;
}
if !selection.dictContent.is_null() {
unsafe {
ptr::copy_nonoverlapping(selection.dictContent, state.dict, selection.dictSize);
}
state.dict_size = selection.dictSize;
state.parameters = parameters;
state.compressed_size = selection.totalCompressedSize;
}
}
if state.live_jobs == 0 {
self.cond.notify_all();
}
}
fn wait(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
while state.live_jobs != 0 {
state = self
.cond
.wait(state)
.unwrap_or_else(|poison| poison.into_inner());
}
}
fn snapshot(&self) -> FastBestState {
*self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn destroy(&self) {
self.wait();
let mut state = self
.state
.lock()
.unwrap_or_else(|poison| poison.into_inner());
unsafe { free_bytes(state.dict) };
state.dict = ptr::null_mut();
state.dict_size = 0;
}
}
unsafe fn try_parameters(
context: &FastCoverContext,
best: &FastBest,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
) {
let count = frequency_count(context.f);
let segment_freqs = unsafe { calloc_array::<u16>(count) };
let dict = unsafe { malloc_bytes(dict_buffer_capacity) };
let freqs = unsafe { malloc_bytes(frequency_bytes(context.f, size_of::<u32>())) }.cast::<u32>();
let mut selection = COVER_dictSelectionError(ERROR_GENERIC);
if segment_freqs.is_null() || dict.is_null() || freqs.is_null() {
display(1, "Failed to allocate buffers: out of memory\n");
} else {
unsafe {
ptr::copy_nonoverlapping(context.freqs, freqs, frequency_count(context.f));
let tail = build_dictionary(
context,
freqs,
dict,
dict_buffer_capacity,
parameters,
segment_freqs,
);
selection = select_dictionary(
context,
dict,
tail,
dict_buffer_capacity,
parameters,
context.accel_params.finalize,
);
}
if COVER_dictSelectionIsError(selection) != 0 {
display(1, "Failed to select dictionary\n");
}
}
best.finish(parameters, selection);
unsafe {
free_bytes(dict);
free_bytes(segment_freqs);
free_bytes(freqs);
COVER_dictSelectionFree(selection);
}
}
unsafe fn select_dictionary(
context: &FastCoverContext,
dict: *mut u8,
tail: usize,
dict_buffer_capacity: usize,
parameters: ZDICT_cover_params_t,
finalize_percentage: c_uint,
) -> COVER_dictSelection_t {
let nb_finalize_samples = (context
.nb_train_samples
.wrapping_mul(finalize_percentage as usize)
/ 100) as c_uint;
unsafe {
COVER_selectDict(
dict.add(tail),
dict_buffer_capacity,
dict_buffer_capacity - tail,
context.samples,
context.samples_sizes,
nb_finalize_samples,
context.nb_train_samples,
context.nb_samples,
parameters,
context.offsets,
ERROR_GENERIC,
)
}
}
fn convert_to_cover_params(parameters: ZDICT_fastCover_params_t) -> ZDICT_cover_params_t {
ZDICT_cover_params_t {
k: parameters.k,
d: parameters.d,
steps: parameters.steps,
nbThreads: parameters.nbThreads,
splitPoint: parameters.splitPoint,
shrinkDict: parameters.shrinkDict,
shrinkDictMaxRegression: 0,
zParams: parameters.zParams,
}
}
fn convert_to_fastcover_params(
parameters: ZDICT_cover_params_t,
output: &mut ZDICT_fastCover_params_t,
f: c_uint,
accel: c_uint,
) {
output.k = parameters.k;
output.d = parameters.d;
output.steps = parameters.steps;
output.nbThreads = parameters.nbThreads;
output.splitPoint = parameters.splitPoint;
output.f = f;
output.accel = accel;
output.zParams = parameters.zParams;
output.shrinkDict = parameters.shrinkDict;
}
unsafe fn finalize_dictionary(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
custom_dict_content: *const c_void,
dict_content_size: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
parameters: ZDICT_params_t,
) -> usize {
unsafe {
ZDICT_finalizeDictionary(
dict_buffer,
dict_buffer_capacity,
custom_dict_content,
dict_content_size,
samples_buffer,
samples_sizes,
nb_samples,
parameters,
)
}
}
unsafe extern "C" {
fn ZDICT_finalizeDictionary(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
custom_dict_content: *const c_void,
dict_content_size: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
parameters: ZDICT_params_t,
) -> usize;
}
#[no_mangle]
pub unsafe extern "C" fn ZDICT_trainFromBuffer_fastCover(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
mut parameters: ZDICT_fastCover_params_t,
) -> usize {
set_display_level(parameters.zParams.notificationLevel);
parameters.splitPoint = 1.0;
if parameters.f == 0 {
parameters.f = DEFAULT_F;
}
if parameters.accel == 0 {
parameters.accel = DEFAULT_ACCEL;
}
let cover_parameters = convert_to_cover_params(parameters);
if !check_parameters(
cover_parameters,
dict_buffer_capacity,
parameters.f,
parameters.accel,
) {
display(1, "FASTCOVER parameters incorrect\n");
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if nb_samples == 0 {
display(1, "FASTCOVER must have at least one input file\n");
return ERROR_SRC_SIZE_WRONG;
}
if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
display(1, "dictBufferCapacity must be at least 256\n");
return ERROR_DST_SIZE_TOO_SMALL;
}
let context = match context_init(
samples_buffer,
samples_sizes,
nb_samples as usize,
cover_parameters.d,
parameters.splitPoint,
parameters.f,
DEFAULT_ACCEL_PARAMETERS[parameters.accel as usize],
) {
Ok(context) => context,
Err(error_code) => {
display(1, "Failed to initialize context\n");
return error_code;
}
};
COVER_warn_on_small_corpus(dict_buffer_capacity, context.nb_dmers, display_level());
let segment_freqs = unsafe { calloc_array::<u16>(frequency_count(parameters.f)) };
if segment_freqs.is_null() {
unsafe { free_bytes(segment_freqs) };
return ERROR_MEMORY_ALLOCATION;
}
let tail = unsafe {
build_dictionary(
&context,
context.freqs,
dict_buffer.cast::<u8>(),
dict_buffer_capacity,
cover_parameters,
segment_freqs,
)
};
let nb_finalize_samples = (context
.nb_train_samples
.wrapping_mul(context.accel_params.finalize as usize)
/ 100) as c_uint;
let dictionary_size = unsafe {
finalize_dictionary(
dict_buffer,
dict_buffer_capacity,
dict_buffer.cast::<u8>().add(tail).cast(),
dict_buffer_capacity - tail,
samples_buffer,
samples_sizes,
nb_finalize_samples,
cover_parameters.zParams,
)
};
unsafe { free_bytes(segment_freqs) };
dictionary_size
}
#[no_mangle]
pub unsafe extern "C" fn ZDICT_optimizeTrainFromBuffer_fastCover(
dict_buffer: *mut c_void,
dict_buffer_capacity: usize,
samples_buffer: *const c_void,
samples_sizes: *const usize,
nb_samples: c_uint,
parameters: *mut ZDICT_fastCover_params_t,
) -> usize {
let parameters = unsafe { &mut *parameters };
let nb_threads = parameters.nbThreads as usize;
let split_point = if parameters.splitPoint <= 0.0 {
FASTCOVER_DEFAULT_SPLITPOINT
} else {
parameters.splitPoint
};
let k_min_d = if parameters.d == 0 { 6 } else { parameters.d };
let k_max_d = if parameters.d == 0 { 8 } else { parameters.d };
let k_min_k = if parameters.k == 0 { 50 } else { parameters.k };
let k_max_k = if parameters.k == 0 {
2000
} else {
parameters.k
};
let k_steps = if parameters.steps == 0 {
40
} else {
parameters.steps
};
let k_step_size = (k_max_k.wrapping_sub(k_min_k) / k_steps).max(1);
let _k_iterations = (1 + (k_max_d.wrapping_sub(k_min_d) / 2))
.wrapping_mul(1 + (k_max_k.wrapping_sub(k_min_k) / k_step_size));
let f = if parameters.f == 0 {
DEFAULT_F
} else {
parameters.f
};
let accel = if parameters.accel == 0 {
DEFAULT_ACCEL
} else {
parameters.accel
};
let display_level = parameters.zParams.notificationLevel as c_int;
if split_point <= 0.0 || split_point > 1.0 {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if accel == 0 || accel > FASTCOVER_MAX_ACCEL {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if k_min_k < k_max_d || k_max_k < k_min_k {
return ERROR_PARAMETER_OUT_OF_BOUND;
}
if nb_samples == 0 {
return ERROR_SRC_SIZE_WRONG;
}
if dict_buffer_capacity < ZDICT_DICTSIZE_MIN {
return ERROR_DST_SIZE_TOO_SMALL;
}
let cover_parameters = convert_to_cover_params(*parameters);
set_display_level(if display_level == 0 {
0
} else {
(display_level - 1) as c_uint
});
let best = FastBest::new();
let mut warned = false;
let mut d = k_min_d;
loop {
if d > k_max_d {
break;
}
let context = match unsafe {
context_init(
samples_buffer,
samples_sizes,
nb_samples as usize,
d,
split_point,
f,
DEFAULT_ACCEL_PARAMETERS[accel as usize],
)
} {
Ok(context) => context,
Err(error_code) => {
best.destroy();
return error_code;
}
};
if !warned {
COVER_warn_on_small_corpus(dict_buffer_capacity, context.nb_dmers, display_level);
warned = true;
}
let mut jobs = Vec::new();
let mut k = k_min_k;
while k <= k_max_k {
let mut job = cover_parameters;
job.k = k;
job.d = d;
job.splitPoint = split_point;
job.steps = k_steps;
job.shrinkDict = 0;
job.zParams.notificationLevel = display_level.saturating_sub(1) as c_uint;
if check_parameters(job, dict_buffer_capacity, f, accel) {
if jobs.try_reserve(1).is_err() {
best.destroy();
return ERROR_MEMORY_ALLOCATION;
}
jobs.push(job);
}
let next = k.wrapping_add(k_step_size);
if next <= k {
break;
}
k = next;
}
for _ in 0..jobs.len() {
best.start();
}
if nb_threads > 1 && jobs.len() > 1 {
let jobs = Arc::new(Mutex::new(jobs.into_iter()));
let worker_count = nb_threads.min(
jobs.lock()
.unwrap_or_else(|poison| poison.into_inner())
.size_hint()
.0,
);
std::thread::scope(|scope| {
for _ in 0..worker_count {
let jobs = Arc::clone(&jobs);
let context_ref = &context;
let best_ref = &best;
scope.spawn(move || loop {
let job = jobs
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.next();
match job {
Some(job) => unsafe {
try_parameters(context_ref, best_ref, dict_buffer_capacity, job)
},
None => break,
}
});
}
});
} else {
for job in jobs {
unsafe { try_parameters(&context, &best, dict_buffer_capacity, job) };
}
}
best.wait();
let next = d.wrapping_add(2);
if next <= d {
break;
}
d = next;
}
let snapshot = best.snapshot();
if ERR_isError(snapshot.compressed_size) {
let result = snapshot.compressed_size;
best.destroy();
return result;
}
unsafe {
ptr::copy_nonoverlapping(snapshot.dict, dict_buffer.cast::<u8>(), snapshot.dict_size);
}
convert_to_fastcover_params(snapshot.parameters, parameters, f, accel);
let result = snapshot.dict_size;
best.destroy();
result
}
#[inline]
fn COVER_warn_on_small_corpus(max_dict_size: usize, nb_dmers: usize, display: c_int) {
crate::dict_builder_cover::COVER_warnOnSmallCorpus(max_dict_size, nb_dmers, display);
}
#[cfg(test)]
mod tests {
use super::*;
use std::mem::size_of;
#[test]
fn fastcover_params_have_the_c_layout() {
assert_eq!(size_of::<ZDICT_fastCover_params_t>(), 56);
assert_eq!(ZDICT_fastCover_params_t::default().k, 0);
}
#[test]
fn hash_formulas_match_the_zstd_constants() {
let bytes = [1u8, 2, 3, 4, 5, 6, 7, 8];
let value = u64::from_le_bytes(bytes);
let expected6 = ((value << 16).wrapping_mul(PRIME6BYTES) >> (64 - 9)) as usize;
let expected8 = (value.wrapping_mul(PRIME8BYTES) >> (64 - 9)) as usize;
assert_eq!(
unsafe { hash_ptr_to_index(bytes.as_ptr(), 9, 6) },
expected6
);
assert_eq!(
unsafe { hash_ptr_to_index(bytes.as_ptr(), 9, 8) },
expected8
);
}
#[test]
fn acceleration_table_matches_fastcover_reference() {
let finalize = [100, 100, 50, 34, 25, 20, 17, 14, 13, 11, 10];
let skip = [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
for index in 0..=10 {
assert_eq!(DEFAULT_ACCEL_PARAMETERS[index].finalize, finalize[index]);
assert_eq!(DEFAULT_ACCEL_PARAMETERS[index].skip, skip[index]);
}
}
#[test]
fn direct_api_keeps_error_precedence() {
let parameters = ZDICT_fastCover_params_t {
k: 50,
d: 8,
..ZDICT_fastCover_params_t::default()
};
assert_eq!(
unsafe {
ZDICT_trainFromBuffer_fastCover(
ptr::null_mut(),
0,
ptr::null(),
ptr::null(),
0,
parameters,
)
},
ERROR_PARAMETER_OUT_OF_BOUND
);
assert_eq!(
unsafe {
ZDICT_trainFromBuffer_fastCover(
ptr::null_mut(),
ZDICT_DICTSIZE_MIN - 1,
ptr::null(),
ptr::null(),
1,
parameters,
)
},
ERROR_DST_SIZE_TOO_SMALL
);
}
}
rust/src/lib.rs
+2
View File
@@ -8,6 +8,8 @@ pub mod debug;
#[cfg(feature = "dict-builder")]
pub mod dict_builder_cover;
#[cfg(all(feature = "compression", feature = "dict-builder"))]
pub mod dict_builder_fastcover;
#[cfg(all(feature = "compression", feature = "dict-builder"))]
pub mod dict_builder_zdict;
#[cfg(feature = "dict-builder")]
pub mod divsufsort;