zstd_compress.cpp

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2024 Blender Foundation
 *
 * SPDX-License-Identifier: Apache-2.0 */
 
#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <thread>
#include <vector>
 
#include <zstd.h>
 
int main(const int argc, const char **argv)
{
  if (argc < 3) {
    return -1;
  }
 
  /* TODO: This might fail for non-ASCII paths on Windows... */
  std::ifstream in(argv[1], std::ios_base::binary);
  std::ofstream out(argv[2], std::ios_base::binary);
  if (!in || !out) {
    return -1;
  }
 
  in.seekg(0, std::ios_base::end);
  size_t in_size = in.tellg();
  in.seekg(0, std::ios_base::beg);
  if (!in) {
    return -1;
  }
 
  std::vector<char> in_data(in_size);
  in.read(in_data.data(), in_size);
  if (!in) {
    return -1;
  }
 
  size_t out_size = ZSTD_compressBound(in_size);
  if (ZSTD_isError(out_size)) {
    return -1;
  }
  std::vector<char> out_data(out_size);
 
  ZSTD_CCtx *cctx = ZSTD_createCCtx();
  if (!cctx) {
    return -1;
  }
  ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, 19);
  /* Enable multi-threaded compression. Output is a standard zstd frame,
   * fully compatible with regular single-threaded decompression.
   * If libzstd lacks MT support the parameter is ignored (still works).
   *
   * Cap the worker count so that several parallel ninja jobs do not
   * multiply into an out-of-memory situation on many-core build hosts.
   * Override with the CYCLES_ZSTD_WORKERS environment variable. */
  unsigned int workers = std::thread::hardware_concurrency();
  if (workers == 0) {
    workers = 4;
  }
  if (workers > 4) {
    workers = 4;
  }
  if (const char *env = std::getenv("CYCLES_ZSTD_WORKERS")) {
    const int env_workers = std::atoi(env);
    if (env_workers >= 0) {
      workers = static_cast<unsigned int>(env_workers);
    }
  }
  ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, static_cast<int>(workers));
  out_size = ZSTD_compress2(cctx, out_data.data(), out_data.size(), in_data.data(), in_data.size());
  ZSTD_freeCCtx(cctx);
  if (ZSTD_isError(out_size)) {
    return -1;
  }
 
  out.write(out_data.data(), out_size);
  if (!out) {
    return -1;
  }
 
  return 0;
}