1*01826a49SYabin Cui /*
2*01826a49SYabin Cui * Copyright (c) Meta Platforms, Inc. and affiliates.
3*01826a49SYabin Cui * All rights reserved.
4*01826a49SYabin Cui *
5*01826a49SYabin Cui * This source code is licensed under both the BSD-style license (found in the
6*01826a49SYabin Cui * LICENSE file in the root directory of this source tree) and the GPLv2 (found
7*01826a49SYabin Cui * in the COPYING file in the root directory of this source tree).
8*01826a49SYabin Cui * You may select, at your option, one of the above-listed licenses.
9*01826a49SYabin Cui */
10*01826a49SYabin Cui
11*01826a49SYabin Cui
12*01826a49SYabin Cui #include <stdio.h> // printf
13*01826a49SYabin Cui #include <stdlib.h> // free
14*01826a49SYabin Cui #include <string.h> // memset, strcat, strlen
15*01826a49SYabin Cui #include <zstd.h> // presumes zstd library is installed
16*01826a49SYabin Cui #include "common.h" // Helper functions, CHECK(), and CHECK_ZSTD()
17*01826a49SYabin Cui
compressFile_orDie(const char * fname,const char * outName,int cLevel,int nbThreads)18*01826a49SYabin Cui static void compressFile_orDie(const char* fname, const char* outName, int cLevel,
19*01826a49SYabin Cui int nbThreads)
20*01826a49SYabin Cui {
21*01826a49SYabin Cui fprintf (stderr, "Starting compression of %s with level %d, using %d threads\n",
22*01826a49SYabin Cui fname, cLevel, nbThreads);
23*01826a49SYabin Cui
24*01826a49SYabin Cui /* Open the input and output files. */
25*01826a49SYabin Cui FILE* const fin = fopen_orDie(fname, "rb");
26*01826a49SYabin Cui FILE* const fout = fopen_orDie(outName, "wb");
27*01826a49SYabin Cui /* Create the input and output buffers.
28*01826a49SYabin Cui * They may be any size, but we recommend using these functions to size them.
29*01826a49SYabin Cui * Performance will only suffer significantly for very tiny buffers.
30*01826a49SYabin Cui */
31*01826a49SYabin Cui size_t const buffInSize = ZSTD_CStreamInSize();
32*01826a49SYabin Cui void* const buffIn = malloc_orDie(buffInSize);
33*01826a49SYabin Cui size_t const buffOutSize = ZSTD_CStreamOutSize();
34*01826a49SYabin Cui void* const buffOut = malloc_orDie(buffOutSize);
35*01826a49SYabin Cui
36*01826a49SYabin Cui /* Create the context. */
37*01826a49SYabin Cui ZSTD_CCtx* const cctx = ZSTD_createCCtx();
38*01826a49SYabin Cui CHECK(cctx != NULL, "ZSTD_createCCtx() failed!");
39*01826a49SYabin Cui
40*01826a49SYabin Cui /* Set any parameters you want.
41*01826a49SYabin Cui * Here we set the compression level, and enable the checksum.
42*01826a49SYabin Cui */
43*01826a49SYabin Cui CHECK_ZSTD( ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, cLevel) );
44*01826a49SYabin Cui CHECK_ZSTD( ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 1) );
45*01826a49SYabin Cui if (nbThreads > 1) {
46*01826a49SYabin Cui size_t const r = ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, nbThreads);
47*01826a49SYabin Cui if (ZSTD_isError(r)) {
48*01826a49SYabin Cui fprintf (stderr, "Note: the linked libzstd library doesn't support multithreading. "
49*01826a49SYabin Cui "Reverting to single-thread mode. \n");
50*01826a49SYabin Cui }
51*01826a49SYabin Cui }
52*01826a49SYabin Cui
53*01826a49SYabin Cui /* This loop read from the input file, compresses that entire chunk,
54*01826a49SYabin Cui * and writes all output produced to the output file.
55*01826a49SYabin Cui */
56*01826a49SYabin Cui size_t const toRead = buffInSize;
57*01826a49SYabin Cui for (;;) {
58*01826a49SYabin Cui size_t read = fread_orDie(buffIn, toRead, fin);
59*01826a49SYabin Cui /* Select the flush mode.
60*01826a49SYabin Cui * If the read may not be finished (read == toRead) we use
61*01826a49SYabin Cui * ZSTD_e_continue. If this is the last chunk, we use ZSTD_e_end.
62*01826a49SYabin Cui * Zstd optimizes the case where the first flush mode is ZSTD_e_end,
63*01826a49SYabin Cui * since it knows it is compressing the entire source in one pass.
64*01826a49SYabin Cui */
65*01826a49SYabin Cui int const lastChunk = (read < toRead);
66*01826a49SYabin Cui ZSTD_EndDirective const mode = lastChunk ? ZSTD_e_end : ZSTD_e_continue;
67*01826a49SYabin Cui /* Set the input buffer to what we just read.
68*01826a49SYabin Cui * We compress until the input buffer is empty, each time flushing the
69*01826a49SYabin Cui * output.
70*01826a49SYabin Cui */
71*01826a49SYabin Cui ZSTD_inBuffer input = { buffIn, read, 0 };
72*01826a49SYabin Cui int finished;
73*01826a49SYabin Cui do {
74*01826a49SYabin Cui /* Compress into the output buffer and write all of the output to
75*01826a49SYabin Cui * the file so we can reuse the buffer next iteration.
76*01826a49SYabin Cui */
77*01826a49SYabin Cui ZSTD_outBuffer output = { buffOut, buffOutSize, 0 };
78*01826a49SYabin Cui size_t const remaining = ZSTD_compressStream2(cctx, &output , &input, mode);
79*01826a49SYabin Cui CHECK_ZSTD(remaining);
80*01826a49SYabin Cui fwrite_orDie(buffOut, output.pos, fout);
81*01826a49SYabin Cui /* If we're on the last chunk we're finished when zstd returns 0,
82*01826a49SYabin Cui * which means its consumed all the input AND finished the frame.
83*01826a49SYabin Cui * Otherwise, we're finished when we've consumed all the input.
84*01826a49SYabin Cui */
85*01826a49SYabin Cui finished = lastChunk ? (remaining == 0) : (input.pos == input.size);
86*01826a49SYabin Cui } while (!finished);
87*01826a49SYabin Cui CHECK(input.pos == input.size,
88*01826a49SYabin Cui "Impossible: zstd only returns 0 when the input is completely consumed!");
89*01826a49SYabin Cui
90*01826a49SYabin Cui if (lastChunk) {
91*01826a49SYabin Cui break;
92*01826a49SYabin Cui }
93*01826a49SYabin Cui }
94*01826a49SYabin Cui
95*01826a49SYabin Cui ZSTD_freeCCtx(cctx);
96*01826a49SYabin Cui fclose_orDie(fout);
97*01826a49SYabin Cui fclose_orDie(fin);
98*01826a49SYabin Cui free(buffIn);
99*01826a49SYabin Cui free(buffOut);
100*01826a49SYabin Cui }
101*01826a49SYabin Cui
102*01826a49SYabin Cui
createOutFilename_orDie(const char * filename)103*01826a49SYabin Cui static char* createOutFilename_orDie(const char* filename)
104*01826a49SYabin Cui {
105*01826a49SYabin Cui size_t const inL = strlen(filename);
106*01826a49SYabin Cui size_t const outL = inL + 5;
107*01826a49SYabin Cui void* const outSpace = malloc_orDie(outL);
108*01826a49SYabin Cui memset(outSpace, 0, outL);
109*01826a49SYabin Cui strcat(outSpace, filename);
110*01826a49SYabin Cui strcat(outSpace, ".zst");
111*01826a49SYabin Cui return (char*)outSpace;
112*01826a49SYabin Cui }
113*01826a49SYabin Cui
main(int argc,const char ** argv)114*01826a49SYabin Cui int main(int argc, const char** argv)
115*01826a49SYabin Cui {
116*01826a49SYabin Cui const char* const exeName = argv[0];
117*01826a49SYabin Cui
118*01826a49SYabin Cui if (argc < 2) {
119*01826a49SYabin Cui printf("wrong arguments\n");
120*01826a49SYabin Cui printf("usage:\n");
121*01826a49SYabin Cui printf("%s FILE [LEVEL] [THREADS]\n", exeName);
122*01826a49SYabin Cui return 1;
123*01826a49SYabin Cui }
124*01826a49SYabin Cui
125*01826a49SYabin Cui int cLevel = 1;
126*01826a49SYabin Cui int nbThreads = 1;
127*01826a49SYabin Cui
128*01826a49SYabin Cui if (argc >= 3) {
129*01826a49SYabin Cui cLevel = atoi (argv[2]);
130*01826a49SYabin Cui CHECK(cLevel != 0, "can't parse LEVEL!");
131*01826a49SYabin Cui }
132*01826a49SYabin Cui
133*01826a49SYabin Cui if (argc >= 4) {
134*01826a49SYabin Cui nbThreads = atoi (argv[3]);
135*01826a49SYabin Cui CHECK(nbThreads != 0, "can't parse THREADS!");
136*01826a49SYabin Cui }
137*01826a49SYabin Cui
138*01826a49SYabin Cui const char* const inFilename = argv[1];
139*01826a49SYabin Cui
140*01826a49SYabin Cui char* const outFilename = createOutFilename_orDie(inFilename);
141*01826a49SYabin Cui compressFile_orDie(inFilename, outFilename, cLevel, nbThreads);
142*01826a49SYabin Cui
143*01826a49SYabin Cui free(outFilename); /* not strictly required, since program execution stops there,
144*01826a49SYabin Cui * but some static analyzer may complain otherwise */
145*01826a49SYabin Cui return 0;
146*01826a49SYabin Cui }
147