ceph/src/zstd/contrib/experimental_dict_builders/fastCover/fastCover.c

   1 /*-*************************************
   2 *  Dependencies
   3 ***************************************/
   4 #include <stdio.h>  /* fprintf */
   5 #include <stdlib.h> /* malloc, free, qsort */
   6 #include <string.h> /* memset */
   7 #include <time.h>   /* clock */
   8 #include "mem.h" /* read */
   9 #include "pool.h"
  10 #include "threading.h"
  11 #include "fastCover.h"
  12 #include "zstd_internal.h" /* includes zstd.h */
  13 #include "zdict.h"
  14
  15
  16 /*-*************************************
  17 *  Constants
  18 ***************************************/
  19 #define FASTCOVER_MAX_SAMPLES_SIZE (sizeof(size_t) == 8 ? ((U32)-1) : ((U32)1 GB))
  20 #define FASTCOVER_MAX_F 32
  21 #define DEFAULT_SPLITPOINT 1.0
  22
  23 /*-*************************************
  24 *  Console display
  25 ***************************************/
  26 static int g_displayLevel = 2;
  27 #define DISPLAY(...)                                                           \
  28   {                                                                            \
  29     fprintf(stderr, __VA_ARGS__);                                              \
  30     fflush(stderr);                                                            \
  31   }
  32 #define LOCALDISPLAYLEVEL(displayLevel, l, ...)                                \
  33   if (displayLevel >= l) {                                                     \
  34     DISPLAY(__VA_ARGS__);                                                      \
  35   } /* 0 : no display;   1: errors;   2: default;  3: details;  4: debug */
  36 #define DISPLAYLEVEL(l, ...) LOCALDISPLAYLEVEL(g_displayLevel, l, __VA_ARGS__)
  37
  38 #define LOCALDISPLAYUPDATE(displayLevel, l, ...)                               \
  39   if (displayLevel >= l) {                                                     \
  40     if ((clock() - g_time > refreshRate) || (displayLevel >= 4)) {             \
  41       g_time = clock();                                                        \
  42       DISPLAY(__VA_ARGS__);                                                    \
  43     }                                                                          \
  44   }
  45 #define DISPLAYUPDATE(l, ...) LOCALDISPLAYUPDATE(g_displayLevel, l, __VA_ARGS__)
  46 static const clock_t refreshRate = CLOCKS_PER_SEC * 15 / 100;
  47 static clock_t g_time = 0;
  48
  49
  50 /*-*************************************
  51 * Hash Functions
  52 ***************************************/
  53 static const U64 prime6bytes = 227718039650203ULL;
  54 static size_t ZSTD_hash6(U64 u, U32 h) { return (size_t)(((u  << (64-48)) * prime6bytes) >> (64-h)) ; }
  55 static size_t ZSTD_hash6Ptr(const void* p, U32 h) { return ZSTD_hash6(MEM_readLE64(p), h); }
  56
  57 static const U64 prime8bytes = 0xCF1BBCDCB7A56463ULL;
  58 static size_t ZSTD_hash8(U64 u, U32 h) { return (size_t)(((u) * prime8bytes) >> (64-h)) ; }
  59 static size_t ZSTD_hash8Ptr(const void* p, U32 h) { return ZSTD_hash8(MEM_readLE64(p), h); }
  60
  61
  62 /**
  63  * Hash the d-byte value pointed to by p and mod 2^f
  64  */
  65 static size_t FASTCOVER_hashPtrToIndex(const void* p, U32 h, unsigned d) {
  66   if (d == 6) {
  67     return ZSTD_hash6Ptr(p, h) & ((1 << h) - 1);
  68   }
  69   return ZSTD_hash8Ptr(p, h) & ((1 << h) - 1);
  70 }
  71
  72
  73 /*-*************************************
  74 * Context
  75 ***************************************/
  76 typedef struct {
  77   const BYTE *samples;
  78   size_t *offsets;
  79   const size_t *samplesSizes;
  80   size_t nbSamples;
  81   size_t nbTrainSamples;
  82   size_t nbTestSamples;
  83   size_t nbDmers;
  84   U32 *freqs;
  85   U16 *segmentFreqs;
  86   unsigned d;
  87 } FASTCOVER_ctx_t;
  88
  89
  90 /*-*************************************
  91 *  Helper functions
  92 ***************************************/
  93 /**
  94  * Returns the sum of the sample sizes.
  95  */
  96 static size_t FASTCOVER_sum(const size_t *samplesSizes, unsigned nbSamples) {
  97   size_t sum = 0;
  98   unsigned i;
  99   for (i = 0; i < nbSamples; ++i) {
 100     sum += samplesSizes[i];
 101   }
 102   return sum;
 103 }
 104
 105
 106 /*-*************************************
 107 *  fast functions
 108 ***************************************/
 109 /**
 110  * A segment is a range in the source as well as the score of the segment.
 111  */
 112 typedef struct {
 113   U32 begin;
 114   U32 end;
 115   U32 score;
 116 } FASTCOVER_segment_t;
 117
 118
 119 /**
 120  * Selects the best segment in an epoch.
 121  * Segments of are scored according to the function:
 122  *
 123  * Let F(d) be the frequency of all dmers with hash value d.
 124  * Let S_i be hash value of the dmer at position i of segment S which has length k.
 125  *
 126  *     Score(S) = F(S_1) + F(S_2) + ... + F(S_{k-d+1})
 127  *
 128  * Once the dmer with hash value d is in the dictionary we set F(d) = F(d)/2.
 129  */
 130 static FASTCOVER_segment_t FASTCOVER_selectSegment(const FASTCOVER_ctx_t *ctx,
 131                                                   U32 *freqs, U32 begin,U32 end,
 132                                                   ZDICT_fastCover_params_t parameters) {
 133   /* Constants */
 134   const U32 k = parameters.k;
 135   const U32 d = parameters.d;
 136   const U32 dmersInK = k - d + 1;
 137   /* Try each segment (activeSegment) and save the best (bestSegment) */
 138   FASTCOVER_segment_t bestSegment = {0, 0, 0};
 139   FASTCOVER_segment_t activeSegment;
 140   /* Reset the activeDmers in the segment */
 141   /* The activeSegment starts at the beginning of the epoch. */
 142   activeSegment.begin = begin;
 143   activeSegment.end = begin;
 144   activeSegment.score = 0;
 145   {
 146     /* Slide the activeSegment through the whole epoch.
 147      * Save the best segment in bestSegment.
 148      */
 149     while (activeSegment.end < end) {
 150       /* Get hash value of current dmer */
 151       const size_t index = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.end, parameters.f, ctx->d);
 152       /* Add frequency of this index to score if this is the first occurrence of index in active segment */
 153       if (ctx->segmentFreqs[index] == 0) {
 154         activeSegment.score += freqs[index];
 155       }
 156       ctx->segmentFreqs[index] += 1;
 157       /* Increment end of segment */
 158       activeSegment.end += 1;
 159       /* If the window is now too large, drop the first position */
 160       if (activeSegment.end - activeSegment.begin == dmersInK + 1) {
 161         /* Get hash value of the dmer to be eliminated from active segment */
 162         const size_t delIndex = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.begin, parameters.f, ctx->d);
 163         ctx->segmentFreqs[delIndex] -= 1;
 164         /* Subtract frequency of this index from score if this is the last occurrence of this index in active segment */
 165         if (ctx->segmentFreqs[delIndex] == 0) {
 166           activeSegment.score -= freqs[delIndex];
 167         }
 168         /* Increment start of segment */
 169         activeSegment.begin += 1;
 170       }
 171       /* If this segment is the best so far save it */
 172       if (activeSegment.score > bestSegment.score) {
 173         bestSegment = activeSegment;
 174       }
 175     }
 176     /* Zero out rest of segmentFreqs array */
 177     while (activeSegment.begin < end) {
 178       const size_t delIndex = FASTCOVER_hashPtrToIndex(ctx->samples + activeSegment.begin, parameters.f, ctx->d);
 179       ctx->segmentFreqs[delIndex] -= 1;
 180       activeSegment.begin += 1;
 181     }
 182   }
 183   {
 184     /* Trim off the zero frequency head and tail from the segment. */
 185     U32 newBegin = bestSegment.end;
 186     U32 newEnd = bestSegment.begin;
 187     U32 pos;
 188     for (pos = bestSegment.begin; pos != bestSegment.end; ++pos) {
 189       const size_t index = FASTCOVER_hashPtrToIndex(ctx->samples + pos, parameters.f, ctx->d);
 190       U32 freq = freqs[index];
 191       if (freq != 0) {
 192         newBegin = MIN(newBegin, pos);
 193         newEnd = pos + 1;
 194       }
 195     }
 196     bestSegment.begin = newBegin;
 197     bestSegment.end = newEnd;
 198   }
 199   {
 200     /*  Zero the frequency of hash value of each dmer covered by the chosen segment. */
 201     U32 pos;
 202     for (pos = bestSegment.begin; pos != bestSegment.end; ++pos) {
 203       const size_t i = FASTCOVER_hashPtrToIndex(ctx->samples + pos, parameters.f, ctx->d);
 204       freqs[i] = 0;
 205     }
 206   }
 207   return bestSegment;
 208 }
 209
 210 /**
 211  * Check the validity of the parameters.
 212  * Returns non-zero if the parameters are valid and 0 otherwise.
 213  */
 214 static int FASTCOVER_checkParameters(ZDICT_fastCover_params_t parameters,
 215                                  size_t maxDictSize) {
 216   /* k, d, and f are required parameters */
 217   if (parameters.d == 0 || parameters.k == 0 || parameters.f == 0) {
 218     return 0;
 219   }
 220   /* d has to be 6 or 8 */
 221   if (parameters.d != 6 && parameters.d != 8) {
 222     return 0;
 223   }
 224   /* 0 < f <= FASTCOVER_MAX_F */
 225   if (parameters.f > FASTCOVER_MAX_F) {
 226     return 0;
 227   }
 228   /* k <= maxDictSize */
 229   if (parameters.k > maxDictSize) {
 230     return 0;
 231   }
 232   /* d <= k */
 233   if (parameters.d > parameters.k) {
 234     return 0;
 235   }
 236   /* 0 < splitPoint <= 1 */
 237   if (parameters.splitPoint <= 0 || parameters.splitPoint > 1) {
 238     return 0;
 239   }
 240   return 1;
 241 }
 242
 243
 244 /**
 245  * Clean up a context initialized with `FASTCOVER_ctx_init()`.
 246  */
 247 static void FASTCOVER_ctx_destroy(FASTCOVER_ctx_t *ctx) {
 248   if (!ctx) {
 249     return;
 250   }
 251   if (ctx->segmentFreqs) {
 252     free(ctx->segmentFreqs);
 253     ctx->segmentFreqs = NULL;
 254   }
 255   if (ctx->freqs) {
 256     free(ctx->freqs);
 257     ctx->freqs = NULL;
 258   }
 259   if (ctx->offsets) {
 260     free(ctx->offsets);
 261     ctx->offsets = NULL;
 262   }
 263 }
 264
 265 /**
 266  * Calculate for frequency of hash value of each dmer in ctx->samples
 267  */
 268 static void FASTCOVER_computeFrequency(U32 *freqs, unsigned f, FASTCOVER_ctx_t *ctx){
 269   size_t start; /* start of current dmer */
 270   for (unsigned i = 0; i < ctx->nbTrainSamples; i++) {
 271     size_t currSampleStart = ctx->offsets[i];
 272     size_t currSampleEnd = ctx->offsets[i+1];
 273     start = currSampleStart;
 274     while (start + ctx->d <= currSampleEnd) {
 275       const size_t dmerIndex = FASTCOVER_hashPtrToIndex(ctx->samples + start, f, ctx->d);
 276       freqs[dmerIndex]++;
 277       start++;
 278     }
 279   }
 280 }
 281
 282 /**
 283  * Prepare a context for dictionary building.
 284  * The context is only dependent on the parameter `d` and can used multiple
 285  * times.
 286  * Returns 1 on success or zero on error.
 287  * The context must be destroyed with `FASTCOVER_ctx_destroy()`.
 288  */
 289 static int FASTCOVER_ctx_init(FASTCOVER_ctx_t *ctx, const void *samplesBuffer,
 290                           const size_t *samplesSizes, unsigned nbSamples,
 291                           unsigned d, double splitPoint, unsigned f) {
 292   const BYTE *const samples = (const BYTE *)samplesBuffer;
 293   const size_t totalSamplesSize = FASTCOVER_sum(samplesSizes, nbSamples);
 294   /* Split samples into testing and training sets */
 295   const unsigned nbTrainSamples = splitPoint < 1.0 ? (unsigned)((double)nbSamples * splitPoint) : nbSamples;
 296   const unsigned nbTestSamples = splitPoint < 1.0 ? nbSamples - nbTrainSamples : nbSamples;
 297   const size_t trainingSamplesSize = splitPoint < 1.0 ? FASTCOVER_sum(samplesSizes, nbTrainSamples) : totalSamplesSize;
 298   const size_t testSamplesSize = splitPoint < 1.0 ? FASTCOVER_sum(samplesSizes + nbTrainSamples, nbTestSamples) : totalSamplesSize;
 299   /* Checks */
 300   if (totalSamplesSize < MAX(d, sizeof(U64)) ||
 301       totalSamplesSize >= (size_t)FASTCOVER_MAX_SAMPLES_SIZE) {
 302     DISPLAYLEVEL(1, "Total samples size is too large (%u MB), maximum size is %u MB\n",
 303                  (U32)(totalSamplesSize >> 20), (FASTCOVER_MAX_SAMPLES_SIZE >> 20));
 304     return 0;
 305   }
 306   /* Check if there are at least 5 training samples */
 307   if (nbTrainSamples < 5) {
 308     DISPLAYLEVEL(1, "Total number of training samples is %u and is invalid.", nbTrainSamples);
 309     return 0;
 310   }
 311   /* Check if there's testing sample */
 312   if (nbTestSamples < 1) {
 313     DISPLAYLEVEL(1, "Total number of testing samples is %u and is invalid.", nbTestSamples);
 314     return 0;
 315   }
 316   /* Zero the context */
 317   memset(ctx, 0, sizeof(*ctx));
 318   DISPLAYLEVEL(2, "Training on %u samples of total size %u\n", nbTrainSamples,
 319                (U32)trainingSamplesSize);
 320   DISPLAYLEVEL(2, "Testing on %u samples of total size %u\n", nbTestSamples,
 321                (U32)testSamplesSize);
 322
 323   ctx->samples = samples;
 324   ctx->samplesSizes = samplesSizes;
 325   ctx->nbSamples = nbSamples;
 326   ctx->nbTrainSamples = nbTrainSamples;
 327   ctx->nbTestSamples = nbTestSamples;
 328   ctx->nbDmers = trainingSamplesSize - d + 1;
 329   ctx->d = d;
 330
 331   /* The offsets of each file */
 332   ctx->offsets = (size_t *)malloc((nbSamples + 1) * sizeof(size_t));
 333   if (!ctx->offsets) {
 334     DISPLAYLEVEL(1, "Failed to allocate scratch buffers\n");
 335     FASTCOVER_ctx_destroy(ctx);
 336     return 0;
 337   }
 338
 339   /* Fill offsets from the samplesSizes */
 340   {
 341     U32 i;
 342     ctx->offsets[0] = 0;
 343     for (i = 1; i <= nbSamples; ++i) {
 344       ctx->offsets[i] = ctx->offsets[i - 1] + samplesSizes[i - 1];
 345     }
 346   }
 347
 348   /* Initialize frequency array of size 2^f */
 349   ctx->freqs = (U32 *)calloc((1 << f), sizeof(U32));
 350   ctx->segmentFreqs = (U16 *)calloc((1 << f), sizeof(U16));
 351   DISPLAYLEVEL(2, "Computing frequencies\n");
 352   FASTCOVER_computeFrequency(ctx->freqs, f, ctx);
 353
 354   return 1;
 355 }
 356
 357
 358 /**
 359  * Given the prepared context build the dictionary.
 360  */
 361 static size_t FASTCOVER_buildDictionary(const FASTCOVER_ctx_t *ctx, U32 *freqs,
 362                                     void *dictBuffer,
 363                                     size_t dictBufferCapacity,
 364                                     ZDICT_fastCover_params_t parameters){
 365   BYTE *const dict = (BYTE *)dictBuffer;
 366   size_t tail = dictBufferCapacity;
 367   /* Divide the data up into epochs of equal size.
 368    * We will select at least one segment from each epoch.
 369    */
 370   const U32 epochs = MAX(1, (U32)(dictBufferCapacity / parameters.k));
 371   const U32 epochSize = (U32)(ctx->nbDmers / epochs);
 372   size_t epoch;
 373   DISPLAYLEVEL(2, "Breaking content into %u epochs of size %u\n", epochs,
 374                epochSize);
 375   /* Loop through the epochs until there are no more segments or the dictionary
 376    * is full.
 377    */
 378   for (epoch = 0; tail > 0; epoch = (epoch + 1) % epochs) {
 379     const U32 epochBegin = (U32)(epoch * epochSize);
 380     const U32 epochEnd = epochBegin + epochSize;
 381     size_t segmentSize;
 382     /* Select a segment */
 383     FASTCOVER_segment_t segment = FASTCOVER_selectSegment(
 384         ctx, freqs, epochBegin, epochEnd, parameters);
 385
 386     /* If the segment covers no dmers, then we are out of content */
 387     if (segment.score == 0) {
 388       break;
 389     }
 390
 391     /* Trim the segment if necessary and if it is too small then we are done */
 392     segmentSize = MIN(segment.end - segment.begin + parameters.d - 1, tail);
 393     if (segmentSize < parameters.d) {
 394       break;
 395     }
 396
 397     /* We fill the dictionary from the back to allow the best segments to be
 398      * referenced with the smallest offsets.
 399      */
 400     tail -= segmentSize;
 401     memcpy(dict + tail, ctx->samples + segment.begin, segmentSize);
 402     DISPLAYUPDATE(
 403         2, "\r%u%%       ",
 404         (U32)(((dictBufferCapacity - tail) * 100) / dictBufferCapacity));
 405   }
 406   DISPLAYLEVEL(2, "\r%79s\r", "");
 407   return tail;
 408 }
 409
 410
 411 /**
 412  * FASTCOVER_best_t is used for two purposes:
 413  * 1. Synchronizing threads.
 414  * 2. Saving the best parameters and dictionary.
 415  *
 416  * All of the methods except FASTCOVER_best_init() are thread safe if zstd is
 417  * compiled with multithreaded support.
 418  */
 419 typedef struct fast_best_s {
 420   ZSTD_pthread_mutex_t mutex;
 421   ZSTD_pthread_cond_t cond;
 422   size_t liveJobs;
 423   void *dict;
 424   size_t dictSize;
 425   ZDICT_fastCover_params_t parameters;
 426   size_t compressedSize;
 427 } FASTCOVER_best_t;
 428
 429 /**
 430  * Initialize the `FASTCOVER_best_t`.
 431  */
 432 static void FASTCOVER_best_init(FASTCOVER_best_t *best) {
 433   if (best==NULL) return; /* compatible with init on NULL */
 434   (void)ZSTD_pthread_mutex_init(&best->mutex, NULL);
 435   (void)ZSTD_pthread_cond_init(&best->cond, NULL);
 436   best->liveJobs = 0;
 437   best->dict = NULL;
 438   best->dictSize = 0;
 439   best->compressedSize = (size_t)-1;
 440   memset(&best->parameters, 0, sizeof(best->parameters));
 441 }
 442
 443 /**
 444  * Wait until liveJobs == 0.
 445  */
 446 static void FASTCOVER_best_wait(FASTCOVER_best_t *best) {
 447   if (!best) {
 448     return;
 449   }
 450   ZSTD_pthread_mutex_lock(&best->mutex);
 451   while (best->liveJobs != 0) {
 452     ZSTD_pthread_cond_wait(&best->cond, &best->mutex);
 453   }
 454   ZSTD_pthread_mutex_unlock(&best->mutex);
 455 }
 456
 457 /**
 458  * Call FASTCOVER_best_wait() and then destroy the FASTCOVER_best_t.
 459  */
 460 static void FASTCOVER_best_destroy(FASTCOVER_best_t *best) {
 461   if (!best) {
 462     return;
 463   }
 464   FASTCOVER_best_wait(best);
 465   if (best->dict) {
 466     free(best->dict);
 467   }
 468   ZSTD_pthread_mutex_destroy(&best->mutex);
 469   ZSTD_pthread_cond_destroy(&best->cond);
 470 }
 471
 472 /**
 473  * Called when a thread is about to be launched.
 474  * Increments liveJobs.
 475  */
 476 static void FASTCOVER_best_start(FASTCOVER_best_t *best) {
 477   if (!best) {
 478     return;
 479   }
 480   ZSTD_pthread_mutex_lock(&best->mutex);
 481   ++best->liveJobs;
 482   ZSTD_pthread_mutex_unlock(&best->mutex);
 483 }
 484
 485 /**
 486  * Called when a thread finishes executing, both on error or success.
 487  * Decrements liveJobs and signals any waiting threads if liveJobs == 0.
 488  * If this dictionary is the best so far save it and its parameters.
 489  */
 490 static void FASTCOVER_best_finish(FASTCOVER_best_t *best, size_t compressedSize,
 491                               ZDICT_fastCover_params_t parameters, void *dict,
 492                               size_t dictSize) {
 493   if (!best) {
 494     return;
 495   }
 496   {
 497     size_t liveJobs;
 498     ZSTD_pthread_mutex_lock(&best->mutex);
 499     --best->liveJobs;
 500     liveJobs = best->liveJobs;
 501     /* If the new dictionary is better */
 502     if (compressedSize < best->compressedSize) {
 503       /* Allocate space if necessary */
 504       if (!best->dict || best->dictSize < dictSize) {
 505         if (best->dict) {
 506           free(best->dict);
 507         }
 508         best->dict = malloc(dictSize);
 509         if (!best->dict) {
 510           best->compressedSize = ERROR(GENERIC);
 511           best->dictSize = 0;
 512           return;
 513         }
 514       }
 515       /* Save the dictionary, parameters, and size */
 516       memcpy(best->dict, dict, dictSize);
 517       best->dictSize = dictSize;
 518       best->parameters = parameters;
 519       best->compressedSize = compressedSize;
 520     }
 521     ZSTD_pthread_mutex_unlock(&best->mutex);
 522     if (liveJobs == 0) {
 523       ZSTD_pthread_cond_broadcast(&best->cond);
 524     }
 525   }
 526 }
 527
 528 /**
 529  * Parameters for FASTCOVER_tryParameters().
 530  */
 531 typedef struct FASTCOVER_tryParameters_data_s {
 532   const FASTCOVER_ctx_t *ctx;
 533   FASTCOVER_best_t *best;
 534   size_t dictBufferCapacity;
 535   ZDICT_fastCover_params_t parameters;
 536 } FASTCOVER_tryParameters_data_t;
 537
 538 /**
 539  * Tries a set of parameters and updates the FASTCOVER_best_t with the results.
 540  * This function is thread safe if zstd is compiled with multithreaded support.
 541  * It takes its parameters as an *OWNING* opaque pointer to support threading.
 542  */
 543 static void FASTCOVER_tryParameters(void *opaque) {
 544   /* Save parameters as local variables */
 545   FASTCOVER_tryParameters_data_t *const data = (FASTCOVER_tryParameters_data_t *)opaque;
 546   const FASTCOVER_ctx_t *const ctx = data->ctx;
 547   const ZDICT_fastCover_params_t parameters = data->parameters;
 548   size_t dictBufferCapacity = data->dictBufferCapacity;
 549   size_t totalCompressedSize = ERROR(GENERIC);
 550   /* Allocate space for hash table, dict, and freqs */
 551   BYTE *const dict = (BYTE * const)malloc(dictBufferCapacity);
 552   U32 *freqs = (U32*) malloc((1 << parameters.f) * sizeof(U32));
 553   if (!dict || !freqs) {
 554     DISPLAYLEVEL(1, "Failed to allocate buffers: out of memory\n");
 555     goto _cleanup;
 556   }
 557   /* Copy the frequencies because we need to modify them */
 558   memcpy(freqs, ctx->freqs, (1 << parameters.f) * sizeof(U32));
 559   /* Build the dictionary */
 560   {
 561     const size_t tail = FASTCOVER_buildDictionary(ctx, freqs, dict,
 562                                               dictBufferCapacity, parameters);
 563
 564     dictBufferCapacity = ZDICT_finalizeDictionary(
 565         dict, dictBufferCapacity, dict + tail, dictBufferCapacity - tail,
 566         ctx->samples, ctx->samplesSizes, (unsigned)ctx->nbTrainSamples,
 567         parameters.zParams);
 568     if (ZDICT_isError(dictBufferCapacity)) {
 569       DISPLAYLEVEL(1, "Failed to finalize dictionary\n");
 570       goto _cleanup;
 571     }
 572   }
 573   /* Check total compressed size */
 574   {
 575     /* Pointers */
 576     ZSTD_CCtx *cctx;
 577     ZSTD_CDict *cdict;
 578     void *dst;
 579     /* Local variables */
 580     size_t dstCapacity;
 581     size_t i;
 582     /* Allocate dst with enough space to compress the maximum sized sample */
 583     {
 584       size_t maxSampleSize = 0;
 585       i = parameters.splitPoint < 1.0 ? ctx->nbTrainSamples : 0;
 586       for (; i < ctx->nbSamples; ++i) {
 587         maxSampleSize = MAX(ctx->samplesSizes[i], maxSampleSize);
 588       }
 589       dstCapacity = ZSTD_compressBound(maxSampleSize);
 590       dst = malloc(dstCapacity);
 591     }
 592     /* Create the cctx and cdict */
 593     cctx = ZSTD_createCCtx();
 594     cdict = ZSTD_createCDict(dict, dictBufferCapacity,
 595                              parameters.zParams.compressionLevel);
 596     if (!dst || !cctx || !cdict) {
 597       goto _compressCleanup;
 598     }
 599     /* Compress each sample and sum their sizes (or error) */
 600     totalCompressedSize = dictBufferCapacity;
 601     i = parameters.splitPoint < 1.0 ? ctx->nbTrainSamples : 0;
 602     for (; i < ctx->nbSamples; ++i) {
 603       const size_t size = ZSTD_compress_usingCDict(
 604           cctx, dst, dstCapacity, ctx->samples + ctx->offsets[i],
 605           ctx->samplesSizes[i], cdict);
 606       if (ZSTD_isError(size)) {
 607         totalCompressedSize = ERROR(GENERIC);
 608         goto _compressCleanup;
 609       }
 610       totalCompressedSize += size;
 611     }
 612   _compressCleanup:
 613     ZSTD_freeCCtx(cctx);
 614     ZSTD_freeCDict(cdict);
 615     if (dst) {
 616       free(dst);
 617     }
 618   }
 619
 620 _cleanup:
 621   FASTCOVER_best_finish(data->best, totalCompressedSize, parameters, dict,
 622                     dictBufferCapacity);
 623   free(data);
 624   if (dict) {
 625     free(dict);
 626   }
 627   if (freqs) {
 628     free(freqs);
 629   }
 630 }
 631
 632 ZDICTLIB_API size_t ZDICT_trainFromBuffer_fastCover(
 633     void *dictBuffer, size_t dictBufferCapacity, const void *samplesBuffer,
 634     const size_t *samplesSizes, unsigned nbSamples, ZDICT_fastCover_params_t parameters) {
 635     BYTE* const dict = (BYTE*)dictBuffer;
 636     FASTCOVER_ctx_t ctx;
 637     parameters.splitPoint = 1.0;
 638     /* Initialize global data */
 639     g_displayLevel = parameters.zParams.notificationLevel;
 640     /* Checks */
 641     if (!FASTCOVER_checkParameters(parameters, dictBufferCapacity)) {
 642       DISPLAYLEVEL(1, "FASTCOVER parameters incorrect\n");
 643       return ERROR(GENERIC);
 644     }
 645     if (nbSamples == 0) {
 646       DISPLAYLEVEL(1, "FASTCOVER must have at least one input file\n");
 647       return ERROR(GENERIC);
 648     }
 649     if (dictBufferCapacity < ZDICT_DICTSIZE_MIN) {
 650       DISPLAYLEVEL(1, "dictBufferCapacity must be at least %u\n",
 651                    ZDICT_DICTSIZE_MIN);
 652       return ERROR(dstSize_tooSmall);
 653     }
 654     /* Initialize context */
 655     if (!FASTCOVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples,
 656                             parameters.d, parameters.splitPoint, parameters.f)) {
 657       DISPLAYLEVEL(1, "Failed to initialize context\n");
 658       return ERROR(GENERIC);
 659     }
 660     /* Build the dictionary */
 661     DISPLAYLEVEL(2, "Building dictionary\n");
 662     {
 663       const size_t tail = FASTCOVER_buildDictionary(&ctx, ctx.freqs, dictBuffer,
 664                                                 dictBufferCapacity, parameters);
 665
 666       const size_t dictionarySize = ZDICT_finalizeDictionary(
 667           dict, dictBufferCapacity, dict + tail, dictBufferCapacity - tail,
 668           samplesBuffer, samplesSizes, (unsigned)ctx.nbTrainSamples,
 669           parameters.zParams);
 670       if (!ZSTD_isError(dictionarySize)) {
 671           DISPLAYLEVEL(2, "Constructed dictionary of size %u\n",
 672                       (U32)dictionarySize);
 673       }
 674       FASTCOVER_ctx_destroy(&ctx);
 675       return dictionarySize;
 676     }
 677 }
 678
 679
 680
 681 ZDICTLIB_API size_t ZDICT_optimizeTrainFromBuffer_fastCover(
 682     void *dictBuffer, size_t dictBufferCapacity, const void *samplesBuffer,
 683     const size_t *samplesSizes, unsigned nbSamples,
 684     ZDICT_fastCover_params_t *parameters) {
 685     /* constants */
 686     const unsigned nbThreads = parameters->nbThreads;
 687     const double splitPoint =
 688         parameters->splitPoint <= 0.0 ? DEFAULT_SPLITPOINT : parameters->splitPoint;
 689     const unsigned kMinD = parameters->d == 0 ? 6 : parameters->d;
 690     const unsigned kMaxD = parameters->d == 0 ? 8 : parameters->d;
 691     const unsigned kMinK = parameters->k == 0 ? 50 : parameters->k;
 692     const unsigned kMaxK = parameters->k == 0 ? 2000 : parameters->k;
 693     const unsigned kSteps = parameters->steps == 0 ? 40 : parameters->steps;
 694     const unsigned kStepSize = MAX((kMaxK - kMinK) / kSteps, 1);
 695     const unsigned kIterations =
 696         (1 + (kMaxD - kMinD) / 2) * (1 + (kMaxK - kMinK) / kStepSize);
 697     const unsigned f = parameters->f == 0 ? 23 : parameters->f;
 698
 699     /* Local variables */
 700     const int displayLevel = parameters->zParams.notificationLevel;
 701     unsigned iteration = 1;
 702     unsigned d;
 703     unsigned k;
 704     FASTCOVER_best_t best;
 705     POOL_ctx *pool = NULL;
 706
 707     /* Checks */
 708     if (splitPoint <= 0 || splitPoint > 1) {
 709       LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect splitPoint\n");
 710       return ERROR(GENERIC);
 711     }
 712     if (kMinK < kMaxD || kMaxK < kMinK) {
 713       LOCALDISPLAYLEVEL(displayLevel, 1, "Incorrect k\n");
 714       return ERROR(GENERIC);
 715     }
 716     if (nbSamples == 0) {
 717       DISPLAYLEVEL(1, "FASTCOVER must have at least one input file\n");
 718       return ERROR(GENERIC);
 719     }
 720     if (dictBufferCapacity < ZDICT_DICTSIZE_MIN) {
 721       DISPLAYLEVEL(1, "dictBufferCapacity must be at least %u\n",
 722                    ZDICT_DICTSIZE_MIN);
 723       return ERROR(dstSize_tooSmall);
 724     }
 725     if (nbThreads > 1) {
 726       pool = POOL_create(nbThreads, 1);
 727       if (!pool) {
 728         return ERROR(memory_allocation);
 729       }
 730     }
 731     /* Initialization */
 732     FASTCOVER_best_init(&best);
 733     /* Turn down global display level to clean up display at level 2 and below */
 734     g_displayLevel = displayLevel == 0 ? 0 : displayLevel - 1;
 735     /* Loop through d first because each new value needs a new context */
 736     LOCALDISPLAYLEVEL(displayLevel, 2, "Trying %u different sets of parameters\n",
 737                       kIterations);
 738     for (d = kMinD; d <= kMaxD; d += 2) {
 739       /* Initialize the context for this value of d */
 740       FASTCOVER_ctx_t ctx;
 741       LOCALDISPLAYLEVEL(displayLevel, 3, "d=%u\n", d);
 742       if (!FASTCOVER_ctx_init(&ctx, samplesBuffer, samplesSizes, nbSamples, d, splitPoint, f)) {
 743         LOCALDISPLAYLEVEL(displayLevel, 1, "Failed to initialize context\n");
 744         FASTCOVER_best_destroy(&best);
 745         POOL_free(pool);
 746         return ERROR(GENERIC);
 747       }
 748       /* Loop through k reusing the same context */
 749       for (k = kMinK; k <= kMaxK; k += kStepSize) {
 750         /* Prepare the arguments */
 751         FASTCOVER_tryParameters_data_t *data = (FASTCOVER_tryParameters_data_t *)malloc(
 752             sizeof(FASTCOVER_tryParameters_data_t));
 753         LOCALDISPLAYLEVEL(displayLevel, 3, "k=%u\n", k);
 754         if (!data) {
 755           LOCALDISPLAYLEVEL(displayLevel, 1, "Failed to allocate parameters\n");
 756           FASTCOVER_best_destroy(&best);
 757           FASTCOVER_ctx_destroy(&ctx);
 758           POOL_free(pool);
 759           return ERROR(GENERIC);
 760         }
 761         data->ctx = &ctx;
 762         data->best = &best;
 763         data->dictBufferCapacity = dictBufferCapacity;
 764         data->parameters = *parameters;
 765         data->parameters.k = k;
 766         data->parameters.d = d;
 767         data->parameters.f = f;
 768         data->parameters.splitPoint = splitPoint;
 769         data->parameters.steps = kSteps;
 770         data->parameters.zParams.notificationLevel = g_displayLevel;
 771         /* Check the parameters */
 772         if (!FASTCOVER_checkParameters(data->parameters, dictBufferCapacity)) {
 773           DISPLAYLEVEL(1, "fastCover parameters incorrect\n");
 774           free(data);
 775           continue;
 776         }
 777         /* Call the function and pass ownership of data to it */
 778         FASTCOVER_best_start(&best);
 779         if (pool) {
 780           POOL_add(pool, &FASTCOVER_tryParameters, data);
 781         } else {
 782           FASTCOVER_tryParameters(data);
 783         }
 784         /* Print status */
 785         LOCALDISPLAYUPDATE(displayLevel, 2, "\r%u%%       ",
 786                            (U32)((iteration * 100) / kIterations));
 787         ++iteration;
 788       }
 789       FASTCOVER_best_wait(&best);
 790       FASTCOVER_ctx_destroy(&ctx);
 791     }
 792     LOCALDISPLAYLEVEL(displayLevel, 2, "\r%79s\r", "");
 793     /* Fill the output buffer and parameters with output of the best parameters */
 794     {
 795       const size_t dictSize = best.dictSize;
 796       if (ZSTD_isError(best.compressedSize)) {
 797         const size_t compressedSize = best.compressedSize;
 798         FASTCOVER_best_destroy(&best);
 799         POOL_free(pool);
 800         return compressedSize;
 801       }
 802       *parameters = best.parameters;
 803       memcpy(dictBuffer, best.dict, dictSize);
 804       FASTCOVER_best_destroy(&best);
 805       POOL_free(pool);
 806       return dictSize;
 807     }
 808
 809 }