BaseTools/Source/C/BrotliCompress/tools/bro.c

   1 /* Copyright 2014 Google Inc. All Rights Reserved.
   2
   3    Distributed under MIT license.
   4    See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
   5 */
   6
   7 /* Example main() function for Brotli library. */
   8
   9 #include <fcntl.h>
  10 #include <stdio.h>
  11 #include <stdlib.h>
  12 #include <string.h>
  13 #include <sys/stat.h>
  14 #include <sys/types.h>
  15 #include <time.h>
  16
  17 #include "../dec/decode.h"
  18 #include "../enc/encode.h"
  19
  20 #if !defined(_WIN32)
  21 #include <unistd.h>
  22 #else
  23 #include <io.h>
  24 #include <share.h>
  25
  26 #define MAKE_BINARY(FILENO) (_setmode((FILENO), _O_BINARY), (FILENO))
  27
  28 #if !defined(__MINGW32__)
  29 #define STDIN_FILENO MAKE_BINARY(_fileno(stdin))
  30 #define STDOUT_FILENO MAKE_BINARY(_fileno(stdout))
  31 #define S_IRUSR S_IREAD
  32 #define S_IWUSR S_IWRITE
  33 #endif
  34 #define fdopen _fdopen
  35 #define unlink _unlink
  36
  37 #define fopen ms_fopen
  38 #define open ms_open
  39
  40 #if defined(_MSC_VER) && (_MSC_VER >= 1400)
  41 #define fseek _fseeki64
  42 #define ftell _ftelli64
  43 #endif
  44
  45 static FILE* ms_fopen(const char *filename, const char *mode) {
  46   FILE* result = 0;
  47   fopen_s(&result, filename, mode);
  48   return result;
  49 }
  50
  51 static int ms_open(const char *filename, int oflag, int pmode) {
  52   int result = -1;
  53   _sopen_s(&result, filename, oflag | O_BINARY, _SH_DENYNO, pmode);
  54   return result;
  55 }
  56 #endif  /* WIN32 */
  57
  58 static int ParseQuality(const char* s, int* quality) {
  59   if (s[0] >= '0' && s[0] <= '9') {
  60     *quality = s[0] - '0';
  61     if (s[1] >= '0' && s[1] <= '9') {
  62       *quality = *quality * 10 + s[1] - '0';
  63       return (s[2] == 0) ? 1 : 0;
  64     }
  65     return (s[1] == 0) ? 1 : 0;
  66   }
  67   return 0;
  68 }
  69
  70 static void ParseArgv(int argc, char **argv,
  71                       char **input_path,
  72                       char **output_path,
  73                       char **dictionary_path,
  74                       int *force,
  75                       int *quality,
  76                       int *decompress,
  77                       int *repeat,
  78                       int *verbose,
  79                       int *lgwin) {
  80   int k;
  81   *force = 0;
  82   *input_path = 0;
  83   *output_path = 0;
  84   *repeat = 1;
  85   *verbose = 0;
  86   *lgwin = 22;
  87   {
  88     size_t argv0_len = strlen(argv[0]);
  89     *decompress =
  90         argv0_len >= 5 && strcmp(&argv[0][argv0_len - 5], "unbro") == 0;
  91   }
  92   for (k = 1; k < argc; ++k) {
  93     if (!strcmp("--force", argv[k]) ||
  94         !strcmp("-f", argv[k])) {
  95       if (*force != 0) {
  96         goto error;
  97       }
  98       *force = 1;
  99       continue;
 100     } else if (!strcmp("--decompress", argv[k]) ||
 101                !strcmp("--uncompress", argv[k]) ||
 102                !strcmp("-d", argv[k])) {
 103       *decompress = 1;
 104       continue;
 105     } else if (!strcmp("--verbose", argv[k]) ||
 106                !strcmp("-v", argv[k])) {
 107       if (*verbose != 0) {
 108         goto error;
 109       }
 110       *verbose = 1;
 111       continue;
 112     }
 113     if (k < argc - 1) {
 114       if (!strcmp("--input", argv[k]) ||
 115           !strcmp("--in", argv[k]) ||
 116           !strcmp("-i", argv[k])) {
 117         if (*input_path != 0) {
 118           goto error;
 119         }
 120         *input_path = argv[k + 1];
 121         ++k;
 122         continue;
 123       } else if (!strcmp("--output", argv[k]) ||
 124                  !strcmp("--out", argv[k]) ||
 125                  !strcmp("-o", argv[k])) {
 126         if (*output_path != 0) {
 127           goto error;
 128         }
 129         *output_path = argv[k + 1];
 130         ++k;
 131         continue;
 132       } else if (!strcmp("--custom-dictionary", argv[k])) {
 133         if (*dictionary_path != 0) {
 134           goto error;
 135         }
 136         *dictionary_path = argv[k + 1];
 137         ++k;
 138         continue;
 139       } else if (!strcmp("--quality", argv[k]) ||
 140                  !strcmp("-q", argv[k])) {
 141         if (!ParseQuality(argv[k + 1], quality)) {
 142           goto error;
 143         }
 144         ++k;
 145         continue;
 146       } else if (!strcmp("--repeat", argv[k]) ||
 147                  !strcmp("-r", argv[k])) {
 148         if (!ParseQuality(argv[k + 1], repeat)) {
 149           goto error;
 150         }
 151         ++k;
 152         continue;
 153       }  else if (!strcmp("--window", argv[k]) ||
 154                   !strcmp("-w", argv[k])) {
 155         if (!ParseQuality(argv[k + 1], lgwin)) {
 156           goto error;
 157         }
 158         if (*lgwin < 10 || *lgwin >= 25) {
 159           goto error;
 160         }
 161         ++k;
 162         continue;
 163       }
 164     }
 165     goto error;
 166   }
 167   return;
 168 error:
 169   fprintf(stderr,
 170           "Usage: %s [--force] [--quality n] [--decompress]"
 171           " [--input filename] [--output filename] [--repeat iters]"
 172           " [--verbose] [--window n] [--custom-dictionary filename]\n",
 173           argv[0]);
 174   exit(1);
 175 }
 176
 177 static FILE* OpenInputFile(const char* input_path) {
 178   FILE* f;
 179   if (input_path == 0) {
 180     return fdopen(STDIN_FILENO, "rb");
 181   }
 182   f = fopen(input_path, "rb");
 183   if (f == 0) {
 184     perror("fopen");
 185     exit(1);
 186   }
 187   return f;
 188 }
 189
 190 static FILE *OpenOutputFile(const char *output_path, const int force) {
 191   int fd;
 192   if (output_path == 0) {
 193     return fdopen(STDOUT_FILENO, "wb");
 194   }
 195   fd = open(output_path, O_CREAT | (force ? 0 : O_EXCL) | O_WRONLY | O_TRUNC,
 196             S_IRUSR | S_IWUSR);
 197   if (fd < 0) {
 198     if (!force) {
 199       struct stat statbuf;
 200       if (stat(output_path, &statbuf) == 0) {
 201         fprintf(stderr, "output file exists\n");
 202         exit(1);
 203       }
 204     }
 205     perror("open");
 206     exit(1);
 207   }
 208   return fdopen(fd, "wb");
 209 }
 210
 211 static int64_t FileSize(const char *path) {
 212   FILE *f = fopen(path, "rb");
 213   int64_t retval;
 214   if (f == NULL) {
 215     return -1;
 216   }
 217   if (fseek(f, 0L, SEEK_END) != 0) {
 218     fclose(f);
 219     return -1;
 220   }
 221   retval = ftell(f);
 222   if (fclose(f) != 0) {
 223     return -1;
 224   }
 225   return retval;
 226 }
 227
 228 /* Result ownersip is passed to caller.
 229    |*dictionary_size| is set to resulting buffer size. */
 230 static uint8_t* ReadDictionary(const char* path, size_t* dictionary_size) {
 231   static const int kMaxDictionarySize = (1 << 24) - 16;
 232   FILE *f = fopen(path, "rb");
 233   int64_t file_size_64;
 234   uint8_t* buffer;
 235   size_t bytes_read;
 236
 237   if (f == NULL) {
 238     perror("fopen");
 239     exit(1);
 240   }
 241
 242   file_size_64 = FileSize(path);
 243   if (file_size_64 == -1) {
 244     fprintf(stderr, "could not get size of dictionary file");
 245     exit(1);
 246   }
 247
 248   if (file_size_64 > kMaxDictionarySize) {
 249     fprintf(stderr, "dictionary is larger than maximum allowed: %d\n",
 250             kMaxDictionarySize);
 251     exit(1);
 252   }
 253   *dictionary_size = (size_t)file_size_64;
 254
 255   buffer = (uint8_t*)malloc(*dictionary_size);
 256   if (!buffer) {
 257     fprintf(stderr, "could not read dictionary: out of memory\n");
 258     exit(1);
 259   }
 260   bytes_read = fread(buffer, sizeof(uint8_t), *dictionary_size, f);
 261   if (bytes_read != *dictionary_size) {
 262     fprintf(stderr, "could not read dictionary\n");
 263     exit(1);
 264   }
 265   fclose(f);
 266   return buffer;
 267 }
 268
 269 static const size_t kFileBufferSize = 65536;
 270
 271 static int Decompress(FILE* fin, FILE* fout, const char* dictionary_path) {
 272   /* Dictionary should be kept during first rounds of decompression. */
 273   uint8_t* dictionary = NULL;
 274   uint8_t* input;
 275   uint8_t* output;
 276   size_t total_out;
 277   size_t available_in;
 278   const uint8_t* next_in;
 279   size_t available_out = kFileBufferSize;
 280   uint8_t* next_out;
 281   BrotliResult result = BROTLI_RESULT_ERROR;
 282   BrotliState* s = BrotliCreateState(NULL, NULL, NULL);
 283   if (!s) {
 284     fprintf(stderr, "out of memory\n");
 285     return 0;
 286   }
 287   if (dictionary_path != NULL) {
 288     size_t dictionary_size = 0;
 289     dictionary = ReadDictionary(dictionary_path, &dictionary_size);
 290     BrotliSetCustomDictionary(dictionary_size, dictionary, s);
 291   }
 292   input = (uint8_t*)malloc(kFileBufferSize);
 293   output = (uint8_t*)malloc(kFileBufferSize);
 294   if (!input || !output) {
 295     fprintf(stderr, "out of memory\n");
 296     goto end;
 297   }
 298   next_out = output;
 299   result = BROTLI_RESULT_NEEDS_MORE_INPUT;
 300   while (1) {
 301     if (result == BROTLI_RESULT_NEEDS_MORE_INPUT) {
 302       if (feof(fin)) {
 303         break;
 304       }
 305       available_in = fread(input, 1, kFileBufferSize, fin);
 306       next_in = input;
 307       if (ferror(fin)) {
 308         break;
 309       }
 310     } else if (result == BROTLI_RESULT_NEEDS_MORE_OUTPUT) {
 311       fwrite(output, 1, kFileBufferSize, fout);
 312       if (ferror(fout)) {
 313         break;
 314       }
 315       available_out = kFileBufferSize;
 316       next_out = output;
 317     } else {
 318       break; /* Error or success. */
 319     }
 320     result = BrotliDecompressStream(&available_in, &next_in,
 321         &available_out, &next_out, &total_out, s);
 322   }
 323   if (next_out != output) {
 324     fwrite(output, 1, (size_t)(next_out - output), fout);
 325   }
 326
 327   if ((result == BROTLI_RESULT_NEEDS_MORE_OUTPUT) || ferror(fout)) {
 328     fprintf(stderr, "failed to write output\n");
 329   } else if (result != BROTLI_RESULT_SUCCESS) { /* Error or needs more input. */
 330     fprintf(stderr, "corrupt input\n");
 331   }
 332
 333 end:
 334   free(dictionary);
 335   free(input);
 336   free(output);
 337   BrotliDestroyState(s);
 338   return (result == BROTLI_RESULT_SUCCESS) ? 1 : 0;
 339 }
 340
 341 static int Compress(int quality, int lgwin, FILE* fin, FILE* fout,
 342     const char *dictionary_path) {
 343   BrotliEncoderState* s = BrotliEncoderCreateInstance(0, 0, 0);
 344   uint8_t* buffer = (uint8_t*)malloc(kFileBufferSize << 1);
 345   uint8_t* input = buffer;
 346   uint8_t* output = buffer + kFileBufferSize;
 347   size_t available_in = 0;
 348   const uint8_t* next_in = NULL;
 349   size_t available_out = kFileBufferSize;
 350   uint8_t* next_out = output;
 351   int is_eof = 0;
 352   int is_ok = 1;
 353
 354   if (!s || !buffer) {
 355     is_ok = 0;
 356     goto finish;
 357   }
 358
 359   BrotliEncoderSetParameter(s, BROTLI_PARAM_QUALITY, (uint32_t)quality);
 360   BrotliEncoderSetParameter(s, BROTLI_PARAM_LGWIN, (uint32_t)lgwin);
 361   if (dictionary_path != NULL) {
 362     size_t dictionary_size = 0;
 363     uint8_t* dictionary = ReadDictionary(dictionary_path, &dictionary_size);
 364     BrotliEncoderSetCustomDictionary(s, dictionary_size, dictionary);
 365     free(dictionary);
 366   }
 367
 368   while (1) {
 369     if (available_in == 0 && !is_eof) {
 370       available_in = fread(input, 1, kFileBufferSize, fin);
 371       next_in = input;
 372       if (ferror(fin)) break;
 373       is_eof = feof(fin);
 374     }
 375
 376     if (!BrotliEncoderCompressStream(s,
 377         is_eof ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS,
 378         &available_in, &next_in, &available_out, &next_out, NULL)) {
 379       is_ok = 0;
 380       break;
 381     }
 382
 383     if (available_out != kFileBufferSize) {
 384       size_t out_size = kFileBufferSize - available_out;
 385       fwrite(output, 1, out_size, fout);
 386       if (ferror(fout)) break;
 387       available_out = kFileBufferSize;
 388       next_out = output;
 389     }
 390
 391     if (BrotliEncoderIsFinished(s)) break;
 392   }
 393
 394 finish:
 395   free(buffer);
 396   BrotliEncoderDestroyInstance(s);
 397
 398   if (!is_ok) {
 399     /* Should detect OOM? */
 400     fprintf(stderr, "failed to compress data\n");
 401     return 0;
 402   } else if (ferror(fout)) {
 403     fprintf(stderr, "failed to write output\n");
 404     return 0;
 405   } else if (ferror(fin)) {
 406     fprintf(stderr, "failed to read input\n");
 407     return 0;
 408   }
 409   return 1;
 410 }
 411
 412 int main(int argc, char** argv) {
 413   char *input_path = 0;
 414   char *output_path = 0;
 415   char *dictionary_path = 0;
 416   int force = 0;
 417   int quality = 11;
 418   int decompress = 0;
 419   int repeat = 1;
 420   int verbose = 0;
 421   int lgwin = 0;
 422   clock_t clock_start;
 423   int i;
 424   ParseArgv(argc, argv, &input_path, &output_path, &dictionary_path, &force,
 425             &quality, &decompress, &repeat, &verbose, &lgwin);
 426   clock_start = clock();
 427   for (i = 0; i < repeat; ++i) {
 428     FILE* fin = OpenInputFile(input_path);
 429     FILE* fout = OpenOutputFile(output_path, force || repeat);
 430     int is_ok = 0;
 431     if (decompress) {
 432       is_ok = Decompress(fin, fout, dictionary_path);
 433     } else {
 434       is_ok = Compress(quality, lgwin, fin, fout, dictionary_path);
 435     }
 436     if (!is_ok) {
 437       unlink(output_path);
 438       exit(1);
 439     }
 440     if (fclose(fin) != 0) {
 441       perror("fclose");
 442       exit(1);
 443     }
 444     if (fclose(fout) != 0) {
 445       perror("fclose");
 446       exit(1);
 447     }
 448   }
 449   if (verbose) {
 450     clock_t clock_end = clock();
 451     double duration = (double)(clock_end - clock_start) / CLOCKS_PER_SEC;
 452     int64_t uncompressed_size;
 453     double uncompressed_bytes_in_MB;
 454     if (duration < 1e-9) {
 455       duration = 1e-9;
 456     }
 457     uncompressed_size = FileSize(decompress ? output_path : input_path);
 458     if (uncompressed_size == -1) {
 459       fprintf(stderr, "failed to determine uncompressed file size\n");
 460       exit(1);
 461     }
 462     uncompressed_bytes_in_MB =
 463         (double)(repeat * uncompressed_size) / (1024.0 * 1024.0);
 464     if (decompress) {
 465       printf("Brotli decompression speed: ");
 466     } else {
 467       printf("Brotli compression speed: ");
 468     }
 469     printf("%g MB/s\n", uncompressed_bytes_in_MB / duration);
 470   }
 471   return 0;
 472 }