ZDICTLIB_API size_t ZDICT_optimizeTrainFromBuffer_cover()

in lib/dictBuilder/cover.c [1119:1254]


ZDICTLIB_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;
  }
}