deps/libchdr/deps/zstd-1.5.5/contrib/seekable_format/examples/parallel_compression.c

   1 /*
   2  * Copyright (c) Meta Platforms, Inc. and affiliates.
   3  * All rights reserved.
   4  *
   5  * This source code is licensed under both the BSD-style license (found in the
   6  * LICENSE file in the root directory of this source tree) and the GPLv2 (found
   7  * in the COPYING file in the root directory of this source tree).
   8  */
   9
  10 #include <stdlib.h>    // malloc, free, exit, atoi
  11 #include <stdio.h>     // fprintf, perror, feof, fopen, etc.
  12 #include <string.h>    // strlen, memset, strcat
  13 #define ZSTD_STATIC_LINKING_ONLY
  14 #include <zstd.h>      // presumes zstd library is installed
  15 #include <zstd_errors.h>
  16 #if defined(WIN32) || defined(_WIN32)
  17 #  include <windows.h>
  18 #  define SLEEP(x) Sleep(x)
  19 #else
  20 #  include <unistd.h>
  21 #  define SLEEP(x) usleep(x * 1000)
  22 #endif
  23
  24 #include "xxhash.h"
  25
  26 #include "pool.h"      // use zstd thread pool for demo
  27
  28 #include "../zstd_seekable.h"
  29
  30 static void* malloc_orDie(size_t size)
  31 {
  32     void* const buff = malloc(size);
  33     if (buff) return buff;
  34     /* error */
  35     perror("malloc:");
  36     exit(1);
  37 }
  38
  39 static FILE* fopen_orDie(const char *filename, const char *instruction)
  40 {
  41     FILE* const inFile = fopen(filename, instruction);
  42     if (inFile) return inFile;
  43     /* error */
  44     perror(filename);
  45     exit(3);
  46 }
  47
  48 static size_t fread_orDie(void* buffer, size_t sizeToRead, FILE* file)
  49 {
  50     size_t const readSize = fread(buffer, 1, sizeToRead, file);
  51     if (readSize == sizeToRead) return readSize;   /* good */
  52     if (feof(file)) return readSize;   /* good, reached end of file */
  53     /* error */
  54     perror("fread");
  55     exit(4);
  56 }
  57
  58 static size_t fwrite_orDie(const void* buffer, size_t sizeToWrite, FILE* file)
  59 {
  60     size_t const writtenSize = fwrite(buffer, 1, sizeToWrite, file);
  61     if (writtenSize == sizeToWrite) return sizeToWrite;   /* good */
  62     /* error */
  63     perror("fwrite");
  64     exit(5);
  65 }
  66
  67 static size_t fclose_orDie(FILE* file)
  68 {
  69     if (!fclose(file)) return 0;
  70     /* error */
  71     perror("fclose");
  72     exit(6);
  73 }
  74
  75 static void fseek_orDie(FILE* file, long int offset, int origin)
  76 {
  77     if (!fseek(file, offset, origin)) {
  78         if (!fflush(file)) return;
  79     }
  80     /* error */
  81     perror("fseek");
  82     exit(7);
  83 }
  84
  85 static long int ftell_orDie(FILE* file)
  86 {
  87     long int off = ftell(file);
  88     if (off != -1) return off;
  89     /* error */
  90     perror("ftell");
  91     exit(8);
  92 }
  93
  94 struct job {
  95     const void* src;
  96     size_t srcSize;
  97     void* dst;
  98     size_t dstSize;
  99
 100     unsigned checksum;
 101
 102     int compressionLevel;
 103     int done;
 104 };
 105
 106 static void compressFrame(void* opaque)
 107 {
 108     struct job* job = opaque;
 109
 110     job->checksum = XXH64(job->src, job->srcSize, 0);
 111
 112     size_t ret = ZSTD_compress(job->dst, job->dstSize, job->src, job->srcSize, job->compressionLevel);
 113     if (ZSTD_isError(ret)) {
 114         fprintf(stderr, "ZSTD_compress() error : %s \n", ZSTD_getErrorName(ret));
 115         exit(20);
 116     }
 117
 118     job->dstSize = ret;
 119     job->done = 1;
 120 }
 121
 122 static void compressFile_orDie(const char* fname, const char* outName, int cLevel, unsigned frameSize, int nbThreads)
 123 {
 124     POOL_ctx* pool = POOL_create(nbThreads, nbThreads);
 125     if (pool == NULL) { fprintf(stderr, "POOL_create() error \n"); exit(9); }
 126
 127     FILE* const fin  = fopen_orDie(fname, "rb");
 128     FILE* const fout = fopen_orDie(outName, "wb");
 129
 130     if (ZSTD_compressBound(frameSize) > 0xFFFFFFFFU) { fprintf(stderr, "Frame size too large \n"); exit(10); }
 131     unsigned dstSize = ZSTD_compressBound(frameSize);
 132
 133
 134     fseek_orDie(fin, 0, SEEK_END);
 135     long int length = ftell_orDie(fin);
 136     fseek_orDie(fin, 0, SEEK_SET);
 137
 138     size_t numFrames = (length + frameSize - 1) / frameSize;
 139
 140     struct job* jobs = malloc_orDie(sizeof(struct job) * numFrames);
 141
 142     size_t i;
 143     for(i = 0; i < numFrames; i++) {
 144         void* in = malloc_orDie(frameSize);
 145         void* out = malloc_orDie(dstSize);
 146
 147         size_t inSize = fread_orDie(in, frameSize, fin);
 148
 149         jobs[i].src = in;
 150         jobs[i].srcSize = inSize;
 151         jobs[i].dst = out;
 152         jobs[i].dstSize = dstSize;
 153         jobs[i].compressionLevel = cLevel;
 154         jobs[i].done = 0;
 155         POOL_add(pool, compressFrame, &jobs[i]);
 156     }
 157
 158     ZSTD_frameLog* fl = ZSTD_seekable_createFrameLog(1);
 159     if (fl == NULL) { fprintf(stderr, "ZSTD_seekable_createFrameLog() failed \n"); exit(11); }
 160     for (i = 0; i < numFrames; i++) {
 161         while (!jobs[i].done) SLEEP(5); /* wake up every 5 milliseconds to check */
 162         fwrite_orDie(jobs[i].dst, jobs[i].dstSize, fout);
 163         free((void*)jobs[i].src);
 164         free(jobs[i].dst);
 165
 166         size_t ret = ZSTD_seekable_logFrame(fl, jobs[i].dstSize, jobs[i].srcSize, jobs[i].checksum);
 167         if (ZSTD_isError(ret)) { fprintf(stderr, "ZSTD_seekable_logFrame() error : %s \n", ZSTD_getErrorName(ret)); }
 168     }
 169
 170     {   unsigned char seekTableBuff[1024];
 171         ZSTD_outBuffer out = {seekTableBuff, 1024, 0};
 172         while (ZSTD_seekable_writeSeekTable(fl, &out) != 0) {
 173             fwrite_orDie(seekTableBuff, out.pos, fout);
 174             out.pos = 0;
 175         }
 176         fwrite_orDie(seekTableBuff, out.pos, fout);
 177     }
 178
 179     ZSTD_seekable_freeFrameLog(fl);
 180     free(jobs);
 181     fclose_orDie(fout);
 182     fclose_orDie(fin);
 183 }
 184
 185 static const char* createOutFilename_orDie(const char* filename)
 186 {
 187     size_t const inL = strlen(filename);
 188     size_t const outL = inL + 5;
 189     void* outSpace = malloc_orDie(outL);
 190     memset(outSpace, 0, outL);
 191     strcat(outSpace, filename);
 192     strcat(outSpace, ".zst");
 193     return (const char*)outSpace;
 194 }
 195
 196 int main(int argc, const char** argv) {
 197     const char* const exeName = argv[0];
 198     if (argc!=4) {
 199         printf("wrong arguments\n");
 200         printf("usage:\n");
 201         printf("%s FILE FRAME_SIZE NB_THREADS\n", exeName);
 202         return 1;
 203     }
 204
 205     {   const char* const inFileName = argv[1];
 206         unsigned const frameSize = (unsigned)atoi(argv[2]);
 207         int const nbThreads = atoi(argv[3]);
 208
 209         const char* const outFileName = createOutFilename_orDie(inFileName);
 210         compressFile_orDie(inFileName, outFileName, 5, frameSize, nbThreads);
 211     }
 212
 213     return 0;
 214 }