lib/Support/Compressor.cpp

   1 //===- lib/Support/Compressor.cpp -------------------------------*- C++ -*-===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file was developed by Reid Spencer and is distributed under the
   6 // University of Illinois Open Source License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This file implements the llvm::Compressor class, an abstraction for memory
  11 // block compression.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "llvm/Config/config.h"
  16 #include "llvm/Support/Compressor.h"
  17 #include "llvm/ADT/StringExtras.h"
  18 #include <cassert>
  19 #include <string>
  20 #include <ostream>
  21 #include "bzip2/bzlib.h"
  22 using namespace llvm;
  23
  24 enum CompressionTypes {
  25   COMP_TYPE_NONE  = '0',
  26   COMP_TYPE_BZIP2 = '2',
  27 };
  28
  29 static int getdata(char*& buffer, size_t &size,
  30                    llvm::Compressor::OutputDataCallback* cb, void* context) {
  31   buffer = 0;
  32   size = 0;
  33   int result = (*cb)(buffer, size, context);
  34   assert(buffer != 0 && "Invalid result from Compressor callback");
  35   assert(size != 0 && "Invalid result from Compressor callback");
  36   return result;
  37 }
  38
  39 //===----------------------------------------------------------------------===//
  40 //=== NULLCOMP - a compression like set of routines that just copies data
  41 //===            without doing any compression. This is provided so that if the
  42 //===            configured environment doesn't have a compression library the
  43 //===            program can still work, albeit using more data/memory.
  44 //===----------------------------------------------------------------------===//
  45
  46 struct NULLCOMP_stream {
  47   // User provided fields
  48   char*  next_in;
  49   size_t avail_in;
  50   char*  next_out;
  51   size_t avail_out;
  52
  53   // Information fields
  54   size_t output_count; // Total count of output bytes
  55 };
  56
  57 static void NULLCOMP_init(NULLCOMP_stream* s) {
  58   s->output_count = 0;
  59 }
  60
  61 static bool NULLCOMP_compress(NULLCOMP_stream* s) {
  62   assert(s && "Invalid NULLCOMP_stream");
  63   assert(s->next_in != 0);
  64   assert(s->next_out != 0);
  65   assert(s->avail_in >= 1);
  66   assert(s->avail_out >= 1);
  67
  68   if (s->avail_out >= s->avail_in) {
  69     ::memcpy(s->next_out, s->next_in, s->avail_in);
  70     s->output_count += s->avail_in;
  71     s->avail_out -= s->avail_in;
  72     s->next_in += s->avail_in;
  73     s->avail_in = 0;
  74     return true;
  75   } else {
  76     ::memcpy(s->next_out, s->next_in, s->avail_out);
  77     s->output_count += s->avail_out;
  78     s->avail_in -= s->avail_out;
  79     s->next_in += s->avail_out;
  80     s->avail_out = 0;
  81     return false;
  82   }
  83 }
  84
  85 static bool NULLCOMP_decompress(NULLCOMP_stream* s) {
  86   assert(s && "Invalid NULLCOMP_stream");
  87   assert(s->next_in != 0);
  88   assert(s->next_out != 0);
  89   assert(s->avail_in >= 1);
  90   assert(s->avail_out >= 1);
  91
  92   if (s->avail_out >= s->avail_in) {
  93     ::memcpy(s->next_out, s->next_in, s->avail_in);
  94     s->output_count += s->avail_in;
  95     s->avail_out -= s->avail_in;
  96     s->next_in += s->avail_in;
  97     s->avail_in = 0;
  98     return true;
  99   } else {
 100     ::memcpy(s->next_out, s->next_in, s->avail_out);
 101     s->output_count += s->avail_out;
 102     s->avail_in -= s->avail_out;
 103     s->next_in += s->avail_out;
 104     s->avail_out = 0;
 105     return false;
 106   }
 107 }
 108
 109 static void NULLCOMP_end(NULLCOMP_stream* strm) {
 110 }
 111
 112 namespace {
 113
 114 /// This structure is only used when a bytecode file is compressed.
 115 /// As bytecode is being decompressed, the memory buffer might need
 116 /// to be reallocated. The buffer allocation is handled in a callback
 117 /// and this structure is needed to retain information across calls
 118 /// to the callback.
 119 /// @brief An internal buffer object used for handling decompression
 120 struct BufferContext {
 121   char* buff;
 122   size_t size;
 123   BufferContext(size_t compressedSize) {
 124     // Null to indicate malloc of a new block
 125     buff = 0;
 126
 127     // Compute the initial length of the uncompression buffer. Note that this
 128     // is twice the length of the compressed buffer and will be doubled again
 129     // in the callback for an initial allocation of 4x compressedSize.  This
 130     // calculation is based on the typical compression ratio of bzip2 on LLVM
 131     // bytecode files which typically ranges in the 50%-75% range.   Since we
 132     // typically get at least 50%, doubling is insufficient. By using a 4x
 133     // multiplier on the first allocation, we minimize the impact of having to
 134     // copy the buffer on reallocation.
 135     size = compressedSize*2;
 136   }
 137
 138   /// trimTo - Reduce the size of the buffer down to the specified amount.  This
 139   /// is useful after have read in the bytecode file to discard extra unused
 140   /// memory.
 141   ///
 142   void trimTo(size_t NewSize) {
 143     buff = (char*)::realloc(buff, NewSize);
 144     size = NewSize;
 145   }
 146
 147   /// This function handles allocation of the buffer used for decompression of
 148   /// compressed bytecode files. It is called by Compressor::decompress which is
 149   /// called by BytecodeReader::ParseBytecode.
 150   static size_t callback(char*&buff, size_t &sz, void* ctxt){
 151     // Case the context variable to our BufferContext
 152     BufferContext* bc = reinterpret_cast<BufferContext*>(ctxt);
 153
 154     // Compute the new, doubled, size of the block
 155     size_t new_size = bc->size * 2;
 156
 157     // Extend or allocate the block (realloc(0,n) == malloc(n))
 158     char* new_buff = (char*) ::realloc(bc->buff, new_size);
 159
 160     // Figure out what to return to the Compressor. If this is the first call,
 161     // then bc->buff will be null. In this case we want to return the entire
 162     // buffer because there was no previous allocation.  Otherwise, when the
 163     // buffer is reallocated, we save the new base pointer in the
 164     // BufferContext.buff field but return the address of only the extension,
 165     // mid-way through the buffer (since its size was doubled). Furthermore,
 166     // the sz result must be 1/2 the total size of the buffer.
 167     if (bc->buff == 0 ) {
 168       buff = bc->buff = new_buff;
 169       sz = new_size;
 170     } else {
 171       bc->buff = new_buff;
 172       buff = new_buff + bc->size;
 173       sz = bc->size;
 174     }
 175
 176     // Retain the size of the allocated block
 177     bc->size = new_size;
 178
 179     // Make sure we fail (return 1) if we didn't get any memory.
 180     return (bc->buff == 0 ? 1 : 0);
 181   }
 182 };
 183
 184 } // end anonymous namespace
 185
 186
 187 namespace {
 188
 189 // This structure retains the context when compressing the bytecode file. The
 190 // WriteCompressedData function below uses it to keep track of the previously
 191 // filled chunk of memory (which it writes) and how many bytes have been
 192 // written.
 193 struct WriterContext {
 194   // Initialize the context
 195   WriterContext(std::ostream*OS, size_t CS)
 196     : chunk(0), sz(0), written(0), compSize(CS), Out(OS) {}
 197
 198   // Make sure we clean up memory
 199   ~WriterContext() {
 200     if (chunk)
 201       delete [] chunk;
 202   }
 203
 204   // Write the chunk
 205   void write(size_t size = 0) {
 206     size_t write_size = (size == 0 ? sz : size);
 207     Out->write(chunk,write_size);
 208     written += write_size;
 209     delete [] chunk;
 210     chunk = 0;
 211     sz = 0;
 212   }
 213
 214   // This function is a callback used by the Compressor::compress function to
 215   // allocate memory for the compression buffer. This function fulfills that
 216   // responsibility but also writes the previous (now filled) buffer out to the
 217   // stream.
 218   static size_t callback(char*& buffer, size_t &size, void* context) {
 219     // Cast the context to the structure it must point to.
 220     WriterContext* ctxt = reinterpret_cast<WriterContext*>(context);
 221
 222     // If there's a previously allocated chunk, it must now be filled with
 223     // compressed data, so we write it out and deallocate it.
 224     if (ctxt->chunk != 0 && ctxt->sz > 0 ) {
 225       ctxt->write();
 226     }
 227
 228     // Compute the size of the next chunk to allocate. We attempt to allocate
 229     // enough memory to handle the compression in a single memory allocation. In
 230     // general, the worst we do on compression of bytecode is about 50% so we
 231     // conservatively estimate compSize / 2 as the size needed for the
 232     // compression buffer. compSize is the size of the compressed data, provided
 233     // by WriteBytecodeToFile.
 234     size = ctxt->sz = ctxt->compSize / 2;
 235
 236     // Allocate the chunks
 237     buffer = ctxt->chunk = new char [size];
 238
 239     // We must return 1 if the allocation failed so that the Compressor knows
 240     // not to use the buffer pointer.
 241     return (ctxt->chunk == 0 ? 1 : 0);
 242   }
 243
 244   char* chunk;       // pointer to the chunk of memory filled by compression
 245   size_t sz;         // size of chunk
 246   size_t written;    // aggregate total of bytes written in all chunks
 247   size_t compSize;   // size of the uncompressed buffer
 248   std::ostream* Out; // The stream we write the data to.
 249 };
 250
 251 }  // end anonymous namespace
 252
 253 // Compress in one of three ways
 254 size_t Compressor::compress(const char* in, size_t size,
 255                             OutputDataCallback* cb, void* context) {
 256   assert(in && "Can't compress null buffer");
 257   assert(size && "Can't compress empty buffer");
 258   assert(cb && "Can't compress without a callback function");
 259
 260   size_t result = 0;
 261
 262   // For small files, we just don't bother compressing. bzip2 isn't very good
 263   // with tiny files and can actually make the file larger, so we just avoid
 264   // it altogether.
 265   if (size > 64*1024) {
 266     // Set up the bz_stream
 267     bz_stream bzdata;
 268     bzdata.bzalloc = 0;
 269     bzdata.bzfree = 0;
 270     bzdata.opaque = 0;
 271     bzdata.next_in = (char*)in;
 272     bzdata.avail_in = size;
 273     bzdata.next_out = 0;
 274     bzdata.avail_out = 0;
 275     switch ( BZ2_bzCompressInit(&bzdata, 5, 0, 100) ) {
 276       case BZ_CONFIG_ERROR: throw std::string("bzip2 library mis-compiled");
 277       case BZ_PARAM_ERROR:  throw std::string("Compressor internal error");
 278       case BZ_MEM_ERROR:    throw std::string("Out of memory");
 279       case BZ_OK:
 280       default:
 281         break;
 282     }
 283
 284     // Get a block of memory
 285     if (0 != getdata(bzdata.next_out, bzdata.avail_out,cb,context)) {
 286       BZ2_bzCompressEnd(&bzdata);
 287       throw std::string("Can't allocate output buffer");
 288     }
 289
 290     // Put compression code in first byte
 291     (*bzdata.next_out++) = COMP_TYPE_BZIP2;
 292     bzdata.avail_out--;
 293
 294     // Compress it
 295     int bzerr = BZ_FINISH_OK;
 296     while (BZ_FINISH_OK == (bzerr = BZ2_bzCompress(&bzdata, BZ_FINISH))) {
 297       if (0 != getdata(bzdata.next_out, bzdata.avail_out,cb,context)) {
 298         BZ2_bzCompressEnd(&bzdata);
 299         throw std::string("Can't allocate output buffer");
 300       }
 301     }
 302     switch (bzerr) {
 303       case BZ_SEQUENCE_ERROR:
 304       case BZ_PARAM_ERROR: throw std::string("Param/Sequence error");
 305       case BZ_FINISH_OK:
 306       case BZ_STREAM_END: break;
 307       default: throw std::string("Oops: ") + utostr(unsigned(bzerr));
 308     }
 309
 310     // Finish
 311     result = bzdata.total_out_lo32 + 1;
 312     if (sizeof(size_t) == sizeof(uint64_t))
 313       result |= static_cast<uint64_t>(bzdata.total_out_hi32) << 32;
 314
 315     BZ2_bzCompressEnd(&bzdata);
 316   } else {
 317     // Do null compression, for small files
 318     NULLCOMP_stream sdata;
 319     sdata.next_in = (char*)in;
 320     sdata.avail_in = size;
 321     NULLCOMP_init(&sdata);
 322
 323     if (0 != getdata(sdata.next_out, sdata.avail_out,cb,context)) {
 324       throw std::string("Can't allocate output buffer");
 325     }
 326
 327     *(sdata.next_out++) = COMP_TYPE_NONE;
 328     sdata.avail_out--;
 329
 330     while (!NULLCOMP_compress(&sdata)) {
 331       if (0 != getdata(sdata.next_out, sdata.avail_out,cb,context)) {
 332         throw std::string("Can't allocate output buffer");
 333       }
 334     }
 335
 336     result = sdata.output_count + 1;
 337     NULLCOMP_end(&sdata);
 338   }
 339   return result;
 340 }
 341
 342 size_t Compressor::compressToNewBuffer(const char* in, size_t size, char*&out) {
 343   BufferContext bc(size);
 344   size_t result = compress(in,size,BufferContext::callback,(void*)&bc);
 345   bc.trimTo(result);
 346   out = bc.buff;
 347   return result;
 348 }
 349
 350 size_t
 351 Compressor::compressToStream(const char*in, size_t size, std::ostream& out) {
 352   // Set up the context and writer
 353   WriterContext ctxt(&out, size / 2);
 354
 355   // Compress everything after the magic number (which we'll alter).
 356   size_t zipSize = Compressor::compress(in,size,
 357     WriterContext::callback, (void*)&ctxt);
 358
 359   if (ctxt.chunk) {
 360     ctxt.write(zipSize - ctxt.written);
 361   }
 362   return zipSize;
 363 }
 364
 365 // Decompress in one of three ways
 366 size_t Compressor::decompress(const char *in, size_t size,
 367                               OutputDataCallback* cb, void* context) {
 368   assert(in && "Can't decompress null buffer");
 369   assert(size > 1 && "Can't decompress empty buffer");
 370   assert(cb && "Can't decompress without a callback function");
 371
 372   size_t result = 0;
 373
 374   switch (*in++) {
 375     case COMP_TYPE_BZIP2: {
 376       // Set up the bz_stream
 377       bz_stream bzdata;
 378       bzdata.bzalloc = 0;
 379       bzdata.bzfree = 0;
 380       bzdata.opaque = 0;
 381       bzdata.next_in = (char*)in;
 382       bzdata.avail_in = size - 1;
 383       bzdata.next_out = 0;
 384       bzdata.avail_out = 0;
 385       switch ( BZ2_bzDecompressInit(&bzdata, 0, 0) ) {
 386         case BZ_CONFIG_ERROR: throw std::string("bzip2 library mis-compiled");
 387         case BZ_PARAM_ERROR:  throw std::string("Compressor internal error");
 388         case BZ_MEM_ERROR:    throw std::string("Out of memory");
 389         case BZ_OK:
 390         default:
 391           break;
 392       }
 393
 394       // Get a block of memory
 395       if (0 != getdata(bzdata.next_out, bzdata.avail_out,cb,context)) {
 396         BZ2_bzDecompressEnd(&bzdata);
 397         throw std::string("Can't allocate output buffer");
 398       }
 399
 400       // Decompress it
 401       int bzerr = BZ_OK;
 402       while (BZ_OK == (bzerr = BZ2_bzDecompress(&bzdata))) {
 403         if (0 != getdata(bzdata.next_out, bzdata.avail_out,cb,context)) {
 404           BZ2_bzDecompressEnd(&bzdata);
 405           throw std::string("Can't allocate output buffer");
 406         }
 407       }
 408
 409       switch (bzerr) {
 410         case BZ_PARAM_ERROR:  throw std::string("Compressor internal error");
 411         case BZ_MEM_ERROR:    throw std::string("Out of memory");
 412         case BZ_DATA_ERROR:   throw std::string("Data integrity error");
 413         case BZ_DATA_ERROR_MAGIC:throw std::string("Data is not BZIP2");
 414         default: throw("Ooops");
 415         case BZ_STREAM_END:
 416           break;
 417       }
 418
 419       // Finish
 420       result = bzdata.total_out_lo32;
 421       if (sizeof(size_t) == sizeof(uint64_t))
 422         result |= (static_cast<uint64_t>(bzdata.total_out_hi32) << 32);
 423       BZ2_bzDecompressEnd(&bzdata);
 424       break;
 425     }
 426
 427     case COMP_TYPE_NONE: {
 428       NULLCOMP_stream sdata;
 429       sdata.next_in = (char*)in;
 430       sdata.avail_in = size - 1;
 431       NULLCOMP_init(&sdata);
 432
 433       if (0 != getdata(sdata.next_out, sdata.avail_out,cb,context)) {
 434         throw std::string("Can't allocate output buffer");
 435       }
 436
 437       while (!NULLCOMP_decompress(&sdata)) {
 438         if (0 != getdata(sdata.next_out, sdata.avail_out,cb,context)) {
 439           throw std::string("Can't allocate output buffer");
 440         }
 441       }
 442
 443       result = sdata.output_count;
 444       NULLCOMP_end(&sdata);
 445       break;
 446     }
 447
 448     default:
 449       throw std::string("Unknown type of compressed data");
 450   }
 451
 452   return result;
 453 }
 454
 455 size_t
 456 Compressor::decompressToNewBuffer(const char* in, size_t size, char*&out) {
 457   BufferContext bc(size);
 458   size_t result = decompress(in,size,BufferContext::callback,(void*)&bc);
 459   out = bc.buff;
 460   return result;
 461 }
 462
 463 size_t
 464 Compressor::decompressToStream(const char*in, size_t size, std::ostream& out){
 465   // Set up the context and writer
 466   WriterContext ctxt(&out,size / 2);
 467
 468   // Compress everything after the magic number (which we'll alter)
 469   size_t zipSize = Compressor::decompress(in,size,
 470     WriterContext::callback, (void*)&ctxt);
 471
 472   if (ctxt.chunk) {
 473     ctxt.write(zipSize - ctxt.written);
 474   }
 475   return zipSize;
 476 }
 477
 478 // vim: sw=2 ai