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 *gapmem,
  77                       int *decompress,
  78                       int *repeat,
  79                       int *verbose,
  80                       int *lgwin) {
  81   int k;
  82   *force = 0;
  83   *input_path = 0;
  84   *output_path = 0;
  85   *repeat = 1;
  86   *verbose = 0;
  87   *lgwin = 22;
  88   {
  89     size_t argv0_len = strlen(argv[0]);
  90     *decompress =
  91         argv0_len >= 5 && strcmp(&argv[0][argv0_len - 5], "unbro") == 0;
  92   }
  93   for (k = 1; k < argc; ++k) {
  94     if (!strcmp("--force", argv[k]) ||
  95         !strcmp("-f", argv[k])) {
  96       if (*force != 0) {
  97         goto error;
  98       }
  99       *force = 1;
 100       continue;
 101     } else if (!strcmp("--decompress", argv[k]) ||
 102                !strcmp("--uncompress", argv[k]) ||
 103                !strcmp("-d", argv[k])) {
 104       *decompress = 1;
 105       continue;
 106     } else if (!strcmp("--verbose", argv[k]) ||
 107                !strcmp("-v", argv[k])) {
 108       if (*verbose != 0) {
 109         goto error;
 110       }
 111       *verbose = 1;
 112       continue;
 113     }
 114     if (k < argc - 1) {
 115       if (!strcmp("--input", argv[k]) ||
 116           !strcmp("--in", argv[k]) ||
 117           !strcmp("-i", argv[k])) {
 118         if (*input_path != 0) {
 119           goto error;
 120         }
 121         *input_path = argv[k + 1];
 122         ++k;
 123         continue;
 124       } else if (!strcmp("--output", argv[k]) ||
 125                  !strcmp("--out", argv[k]) ||
 126                  !strcmp("-o", argv[k])) {
 127         if (*output_path != 0) {
 128           goto error;
 129         }
 130         *output_path = argv[k + 1];
 131         ++k;
 132         continue;
 133       } else if (!strcmp("--custom-dictionary", argv[k])) {
 134         if (*dictionary_path != 0) {
 135           goto error;
 136         }
 137         *dictionary_path = argv[k + 1];
 138         ++k;
 139         continue;
 140       } else if (!strcmp("--quality", argv[k]) ||
 141                  !strcmp("-q", argv[k])) {
 142         if (!ParseQuality(argv[k + 1], quality)) {
 143           goto error;
 144         }
 145         ++k;
 146         continue;
 147       } else if (!strcmp("--repeat", argv[k]) ||
 148                  !strcmp("-r", argv[k])) {
 149         if (!ParseQuality(argv[k + 1], repeat)) {
 150           goto error;
 151         }
 152         ++k;
 153         continue;
 154       }  else if (!strcmp("--window", argv[k]) ||
 155                   !strcmp("-w", argv[k])) {
 156         if (!ParseQuality(argv[k + 1], lgwin)) {
 157           goto error;
 158         }
 159         if (*lgwin < 10 || *lgwin >= 25) {
 160           goto error;
 161         }
 162         ++k;
 163         continue;
 164       } else if (!strcmp("--gap", argv[k]) ||
 165                   !strcmp("-g", argv[k])) {
 166         if (!ParseQuality(argv[k + 1], gapmem)) {
 167           goto error;
 168         }
 169         ++k;
 170         continue;
 171       }
 172     }
 173     goto error;
 174   }
 175   return;
 176 error:
 177   fprintf(stderr,
 178           "Usage: %s [--force] [--quality n] [--gap n] [--decompress]"
 179           " [--input filename] [--output filename] [--repeat iters]"
 180           " [--verbose] [--window n] [--custom-dictionary filename]\n",
 181           argv[0]);
 182   exit(1);
 183 }
 184
 185 static FILE* OpenInputFile(const char* input_path) {
 186   FILE* f;
 187   if (input_path == 0) {
 188     return fdopen(STDIN_FILENO, "rb");
 189   }
 190   f = fopen(input_path, "rb");
 191   if (f == 0) {
 192     perror("fopen");
 193     exit(1);
 194   }
 195   return f;
 196 }
 197
 198 static FILE *OpenOutputFile(const char *output_path, const int force) {
 199   int fd;
 200   if (output_path == 0) {
 201     return fdopen(STDOUT_FILENO, "wb");
 202   }
 203   fd = open(output_path, O_CREAT | (force ? 0 : O_EXCL) | O_WRONLY | O_TRUNC,
 204             S_IRUSR | S_IWUSR);
 205   if (fd < 0) {
 206     if (!force) {
 207       struct stat statbuf;
 208       if (stat(output_path, &statbuf) == 0) {
 209         fprintf(stderr, "output file exists\n");
 210         exit(1);
 211       }
 212     }
 213     perror("open");
 214     exit(1);
 215   }
 216   return fdopen(fd, "wb");
 217 }
 218
 219 static int64_t FileSize(const char *path) {
 220   FILE *f = fopen(path, "rb");
 221   int64_t retval;
 222   if (f == NULL) {
 223     return -1;
 224   }
 225   if (fseek(f, 0L, SEEK_END) != 0) {
 226     fclose(f);
 227     return -1;
 228   }
 229   retval = ftell(f);
 230   if (fclose(f) != 0) {
 231     return -1;
 232   }
 233   return retval;
 234 }
 235
 236 /* Brotli specified memory allocate function */
 237 static void *BrAlloc (void *memsize, size_t size) {
 238   *(int64_t *)memsize += size;
 239   return malloc(size);
 240 }
 241
 242 /* Brotli specified memory free function */
 243 static void BrFree (void *memsize, void *ptr) {
 244   free(ptr);
 245 }
 246
 247 /* Result ownersip is passed to caller.
 248    |*dictionary_size| is set to resulting buffer size. */
 249 static uint8_t* ReadDictionary(const char* path, size_t* dictionary_size, void *memsize) {
 250   static const int kMaxDictionarySize = (1 << 24) - 16;
 251   FILE *f = fopen(path, "rb");
 252   int64_t file_size_64;
 253   uint8_t* buffer;
 254   size_t bytes_read;
 255
 256   if (f == NULL) {
 257     perror("fopen");
 258     exit(1);
 259   }
 260
 261   file_size_64 = FileSize(path);
 262   if (file_size_64 == -1) {
 263     fprintf(stderr, "could not get size of dictionary file");
 264     exit(1);
 265   }
 266
 267   if (file_size_64 > kMaxDictionarySize) {
 268     fprintf(stderr, "dictionary is larger than maximum allowed: %d\n",
 269             kMaxDictionarySize);
 270     exit(1);
 271   }
 272   *dictionary_size = (size_t)file_size_64;
 273
 274   buffer = (uint8_t *)BrAlloc(memsize, *dictionary_size);
 275   if (!buffer) {
 276     fprintf(stderr, "could not read dictionary: out of memory\n");
 277     exit(1);
 278   }
 279   bytes_read = fread(buffer, sizeof(uint8_t), *dictionary_size, f);
 280   if (bytes_read != *dictionary_size) {
 281     fprintf(stderr, "could not read dictionary\n");
 282     exit(1);
 283   }
 284   fclose(f);
 285   return buffer;
 286 }
 287
 288 static const size_t kFileBufferSize = 65536;
 289
 290 static int Decompress(FILE* fin, FILE* fout, const char* dictionary_path, void *memsize) {
 291   /* Dictionary should be kept during first rounds of decompression. */
 292   uint8_t* dictionary = NULL;
 293   uint8_t* input;
 294   uint8_t* output;
 295   size_t total_out;
 296   size_t available_in;
 297   const uint8_t* next_in;
 298   size_t available_out = kFileBufferSize;
 299   uint8_t* next_out;
 300   BrotliResult result = BROTLI_RESULT_ERROR;
 301   BrotliState* s = BrotliCreateState(BrAlloc, BrFree, memsize);
 302   if (!s) {
 303     fprintf(stderr, "out of memory\n");
 304     return 0;
 305   }
 306   if (dictionary_path != NULL) {
 307     size_t dictionary_size = 0;
 308     dictionary = ReadDictionary(dictionary_path, &dictionary_size, memsize);
 309     BrotliSetCustomDictionary(dictionary_size, dictionary, s);
 310   }
 311   input = (uint8_t*)BrAlloc(memsize, kFileBufferSize);
 312   output = (uint8_t*)BrAlloc(memsize, kFileBufferSize);
 313   if (!input || !output) {
 314     fprintf(stderr, "out of memory\n");
 315     goto end;
 316   }
 317   next_out = output;
 318   result = BROTLI_RESULT_NEEDS_MORE_INPUT;
 319   while (1) {
 320     if (result == BROTLI_RESULT_NEEDS_MORE_INPUT) {
 321       if (feof(fin)) {
 322         break;
 323       }
 324       available_in = fread(input, 1, kFileBufferSize, fin);
 325       next_in = input;
 326       if (ferror(fin)) {
 327         break;
 328       }
 329     } else if (result == BROTLI_RESULT_NEEDS_MORE_OUTPUT) {
 330       fwrite(output, 1, kFileBufferSize, fout);
 331       if (ferror(fout)) {
 332         break;
 333       }
 334       available_out = kFileBufferSize;
 335       next_out = output;
 336     } else {
 337       break; /* Error or success. */
 338     }
 339     result = BrotliDecompressStream(&available_in, &next_in,
 340         &available_out, &next_out, &total_out, s);
 341   }
 342   if (next_out != output) {
 343     fwrite(output, 1, (size_t)(next_out - output), fout);
 344   }
 345
 346   if ((result == BROTLI_RESULT_NEEDS_MORE_OUTPUT) || ferror(fout)) {
 347     fprintf(stderr, "failed to write output\n");
 348   } else if (result != BROTLI_RESULT_SUCCESS) { /* Error or needs more input. */
 349     fprintf(stderr, "corrupt input\n");
 350   }
 351
 352 end:
 353   BrFree(memsize, dictionary);
 354   BrFree(memsize, input);
 355   BrFree(memsize, output);
 356   BrotliDestroyState(s);
 357   return (result == BROTLI_RESULT_SUCCESS) ? 1 : 0;
 358 }
 359
 360 static int Compress(int quality, int lgwin, FILE* fin, FILE* fout,
 361     const char *dictionary_path, void *memsize) {
 362   BrotliEncoderState* s = BrotliEncoderCreateInstance(BrAlloc, BrFree, memsize);
 363   uint8_t* buffer = (uint8_t*)BrAlloc(memsize, kFileBufferSize << 1);
 364   uint8_t* input = buffer;
 365   uint8_t* output = buffer + kFileBufferSize;
 366   size_t available_in = 0;
 367   const uint8_t* next_in = NULL;
 368   size_t available_out = kFileBufferSize;
 369   uint8_t* next_out = output;
 370   int is_eof = 0;
 371   int is_ok = 1;
 372
 373   if (!s || !buffer) {
 374     is_ok = 0;
 375     goto finish;
 376   }
 377
 378   BrotliEncoderSetParameter(s, BROTLI_PARAM_QUALITY, (uint32_t)quality);
 379   BrotliEncoderSetParameter(s, BROTLI_PARAM_LGWIN, (uint32_t)lgwin);
 380   if (dictionary_path != NULL) {
 381     size_t dictionary_size = 0;
 382     uint8_t* dictionary = ReadDictionary(dictionary_path, &dictionary_size, memsize);
 383     BrotliEncoderSetCustomDictionary(s, dictionary_size, dictionary);
 384     BrFree(memsize, dictionary);
 385   }
 386
 387   while (1) {
 388     if (available_in == 0 && !is_eof) {
 389       available_in = fread(input, 1, kFileBufferSize, fin);
 390       next_in = input;
 391       if (ferror(fin)) break;
 392       is_eof = feof(fin);
 393     }
 394
 395     if (!BrotliEncoderCompressStream(s,
 396         is_eof ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS,
 397         &available_in, &next_in, &available_out, &next_out, NULL)) {
 398       is_ok = 0;
 399       break;
 400     }
 401
 402     if (available_out != kFileBufferSize) {
 403       size_t out_size = kFileBufferSize - available_out;
 404       fwrite(output, 1, out_size, fout);
 405       if (ferror(fout)) break;
 406       available_out = kFileBufferSize;
 407       next_out = output;
 408     }
 409
 410     if (BrotliEncoderIsFinished(s)) break;
 411   }
 412
 413 finish:
 414   BrFree(memsize, buffer);
 415   BrotliEncoderDestroyInstance(s);
 416
 417   if (!is_ok) {
 418     /* Should detect OOM? */
 419     fprintf(stderr, "failed to compress data\n");
 420     return 0;
 421   } else if (ferror(fout)) {
 422     fprintf(stderr, "failed to write output\n");
 423     return 0;
 424   } else if (ferror(fin)) {
 425     fprintf(stderr, "failed to read input\n");
 426     return 0;
 427   }
 428   return 1;
 429 }
 430
 431 #define GAP_MEM_BLOCK  4096
 432
 433 int main(int argc, char** argv) {
 434   char *input_path = 0;
 435   char *output_path = 0;
 436   char *dictionary_path = 0;
 437   int force = 0;
 438   int quality = 11;
 439   int gmem = 1;
 440   int decompress = 0;
 441   int repeat = 1;
 442   int verbose = 0;
 443   int lgwin = 0;
 444   clock_t clock_start;
 445   int i;
 446   int64_t originsize = 0;
 447   int64_t msize = 0;
 448   ParseArgv(argc, argv, &input_path, &output_path, &dictionary_path, &force,
 449             &quality, &gmem, &decompress, &repeat, &verbose, &lgwin);
 450   clock_start = clock();
 451   for (i = 0; i < repeat; ++i) {
 452     FILE* fin = OpenInputFile(input_path);
 453     FILE* fout = OpenOutputFile(output_path, force || repeat);
 454     int is_ok = 0;
 455     if (decompress) {
 456       if (fseek(fin, 16, SEEK_SET) != 0) {
 457         fclose(fin);
 458         return -1;
 459       }
 460       is_ok = Decompress(fin, fout, dictionary_path, (void *)&msize);
 461     } else {
 462       originsize = FileSize(input_path);  /* get original file size */
 463       fwrite(&originsize, 1, sizeof(int64_t), fout); /* add in original binary file size */
 464       fwrite(&msize, 1, sizeof(int64_t), fout); /* add in dummy decompression required memory size */
 465       is_ok = Compress(quality, lgwin, fin, fout, dictionary_path, (void *)&msize);
 466     }
 467     if (!is_ok) {
 468       unlink(output_path);
 469       exit(1);
 470     }
 471     if (fclose(fin) != 0) {
 472       perror("fclose");
 473       exit(1);
 474     }
 475     if (fclose(fout) != 0) {
 476       perror("fclose");
 477       exit(1);
 478     }
 479     /* after compression operation then execute decompression operation
 480        to get decompression required memory size. */
 481     if (decompress == 0) {
 482       fin = OpenInputFile(output_path);
 483       fout = tmpfile ();
 484       msize = 0;
 485       if (fseek(fin, 16, SEEK_SET) != 0) {
 486         fclose(fin);
 487         return -1;
 488       }
 489       is_ok = Decompress(fin, fout, dictionary_path, (void *)&msize);
 490       if (!is_ok) {
 491         exit(1);
 492       }
 493       if (fclose(fin) != 0) {
 494         perror("fclose");
 495         exit(1);
 496       }
 497       if (fclose(fout) != 0) {
 498         perror("fclose");
 499         exit(1);
 500       }
 501       fout = fopen(output_path, "rb+");  /* open output_path file and add in head info */
 502       /* seek to the offset of decompression required memory size */
 503       if (fseek(fout, 8, SEEK_SET) != 0) {
 504         fclose(fout);
 505         return -1;
 506       }
 507       msize += gmem * GAP_MEM_BLOCK;  /* there is a memory gap between IA32 and X64 environment*/
 508       fwrite(&msize, 1, sizeof(int64_t), fout); /* update final decompression required memory size */
 509       if (fclose(fout) != 0) {
 510         perror("fclose");
 511         exit(1);
 512       }
 513     }
 514   }
 515   if (verbose) {
 516     clock_t clock_end = clock();
 517     double duration = (double)(clock_end - clock_start) / CLOCKS_PER_SEC;
 518     int64_t uncompressed_size;
 519     double uncompressed_bytes_in_MB;
 520     if (duration < 1e-9) {
 521       duration = 1e-9;
 522     }
 523     uncompressed_size = FileSize(decompress ? output_path : input_path);
 524     if (uncompressed_size == -1) {
 525       fprintf(stderr, "failed to determine uncompressed file size\n");
 526       exit(1);
 527     }
 528     uncompressed_bytes_in_MB =
 529         (double)(repeat * uncompressed_size) / (1024.0 * 1024.0);
 530     if (decompress) {
 531       printf("Brotli decompression speed: ");
 532     } else {
 533       printf("Brotli compression speed: ");
 534     }
 535     printf("%g MB/s\n", uncompressed_bytes_in_MB / duration);
 536   }
 537   return 0;
 538 }