src/api/internal/BgzfStream_p.cpp

   1 // ***************************************************************************
   2 // BgzfStream_p.cpp (c) 2011 Derek Barnett
   3 // Marth Lab, Department of Biology, Boston College
   4 // All rights reserved.
   5 // ---------------------------------------------------------------------------
   6 // Last modified: 21 March 2011(DB)
   7 // ---------------------------------------------------------------------------
   8 // Based on BGZF routines developed at the Broad Institute.
   9 // Provides the basic functionality for reading & writing BGZF files
  10 // Replaces the old BGZF.* files to avoid clashing with other toolkits
  11 // ***************************************************************************
  12
  13 #include <api/internal/BgzfStream_p.h>
  14 using namespace BamTools;
  15 using namespace BamTools::Internal;
  16
  17 #include <cstring>
  18 #include <algorithm>
  19 using namespace std;
  20
  21 // constructor
  22 BgzfStream::BgzfStream(void)
  23     : UncompressedBlockSize(Constants::BGZF_DEFAULT_BLOCK_SIZE)
  24     , CompressedBlockSize(Constants::BGZF_MAX_BLOCK_SIZE)
  25     , BlockLength(0)
  26     , BlockOffset(0)
  27     , BlockAddress(0)
  28     , IsOpen(false)
  29     , IsWriteOnly(false)
  30     , IsWriteCompressed(true)
  31     , Stream(NULL)
  32     , UncompressedBlock(NULL)
  33     , CompressedBlock(NULL)
  34 {
  35     try {
  36         CompressedBlock   = new char[CompressedBlockSize];
  37         UncompressedBlock = new char[UncompressedBlockSize];
  38     } catch( std::bad_alloc& ba ) {
  39         fprintf(stderr, "BgzfStream ERROR: unable to allocate memory\n");
  40         exit(1);
  41     }
  42 }
  43
  44 // destructor
  45 BgzfStream::~BgzfStream(void) {
  46     if( CompressedBlock   ) delete[] CompressedBlock;
  47     if( UncompressedBlock ) delete[] UncompressedBlock;
  48 }
  49
  50 // closes BGZF file
  51 void BgzfStream::Close(void) {
  52
  53     // skip if file not open
  54     if ( !IsOpen ) return;
  55
  56     // if writing to file, flush the current BGZF block,
  57     // then write an empty block (as EOF marker)
  58     if ( IsWriteOnly ) {
  59         FlushBlock();
  60         int blockLength = DeflateBlock();
  61         fwrite(CompressedBlock, 1, blockLength, Stream);
  62     }
  63
  64     // flush and close stream
  65     fflush(Stream);
  66     fclose(Stream);
  67
  68     // reset flags
  69     IsWriteCompressed = true;
  70     IsOpen = false;
  71 }
  72
  73 // compresses the current block
  74 int BgzfStream::DeflateBlock(void) {
  75
  76     // initialize the gzip header
  77     char* buffer = CompressedBlock;
  78     memset(buffer, 0, 18);
  79     buffer[0]  = Constants::GZIP_ID1;
  80     buffer[1]  = (char)Constants::GZIP_ID2;
  81     buffer[2]  = Constants::CM_DEFLATE;
  82     buffer[3]  = Constants::FLG_FEXTRA;
  83     buffer[9]  = (char)Constants::OS_UNKNOWN;
  84     buffer[10] = Constants::BGZF_XLEN;
  85     buffer[12] = Constants::BGZF_ID1;
  86     buffer[13] = Constants::BGZF_ID2;
  87     buffer[14] = Constants::BGZF_LEN;
  88
  89     // set compression level
  90     const int compressionLevel = ( IsWriteCompressed ? Z_DEFAULT_COMPRESSION : 0 );
  91
  92     // loop to retry for blocks that do not compress enough
  93     int inputLength = BlockOffset;
  94     int compressedLength = 0;
  95     unsigned int bufferSize = CompressedBlockSize;
  96
  97     while ( true ) {
  98
  99         // initialize zstream values
 100         z_stream zs;
 101         zs.zalloc    = NULL;
 102         zs.zfree     = NULL;
 103         zs.next_in   = (Bytef*)UncompressedBlock;
 104         zs.avail_in  = inputLength;
 105         zs.next_out  = (Bytef*)&buffer[Constants::BGZF_BLOCK_HEADER_LENGTH];
 106         zs.avail_out = bufferSize - Constants::BGZF_BLOCK_HEADER_LENGTH - Constants::BGZF_BLOCK_FOOTER_LENGTH;
 107
 108         // initialize the zlib compression algorithm
 109         if ( deflateInit2(&zs,
 110                           compressionLevel,
 111                           Z_DEFLATED,
 112                           Constants::GZIP_WINDOW_BITS,
 113                           Constants::Z_DEFAULT_MEM_LEVEL,
 114                           Z_DEFAULT_STRATEGY) != Z_OK )
 115         {
 116             fprintf(stderr, "BgzfStream ERROR: zlib deflate initialization failed\n");
 117             exit(1);
 118         }
 119
 120         // compress the data
 121         int status = deflate(&zs, Z_FINISH);
 122         if ( status != Z_STREAM_END ) {
 123
 124             deflateEnd(&zs);
 125
 126             // reduce the input length and try again
 127             if ( status == Z_OK ) {
 128                 inputLength -= 1024;
 129                 if ( inputLength < 0 ) {
 130                     fprintf(stderr, "BgzfStream ERROR: input reduction failed\n");
 131                     exit(1);
 132                 }
 133                 continue;
 134             }
 135
 136             fprintf(stderr, "BgzfStream ERROR: zlib::deflateEnd() failed\n");
 137             exit(1);
 138         }
 139
 140         // finalize the compression routine
 141         if ( deflateEnd(&zs) != Z_OK ) {
 142             fprintf(stderr, "BgzfStream ERROR: zlib::deflateEnd() failed\n");
 143             exit(1);
 144         }
 145
 146         compressedLength = zs.total_out;
 147         compressedLength += Constants::BGZF_BLOCK_HEADER_LENGTH + Constants::BGZF_BLOCK_FOOTER_LENGTH;
 148         if ( compressedLength > Constants::BGZF_MAX_BLOCK_SIZE ) {
 149             fprintf(stderr, "BgzfStream ERROR: deflate overflow\n");
 150             exit(1);
 151         }
 152
 153         break;
 154     }
 155
 156     // store the compressed length
 157     BamTools::PackUnsignedShort(&buffer[16], (unsigned short)(compressedLength - 1));
 158
 159     // store the CRC32 checksum
 160     unsigned int crc = crc32(0, NULL, 0);
 161     crc = crc32(crc, (Bytef*)UncompressedBlock, inputLength);
 162     BamTools::PackUnsignedInt(&buffer[compressedLength - 8], crc);
 163     BamTools::PackUnsignedInt(&buffer[compressedLength - 4], inputLength);
 164
 165     // ensure that we have less than a block of data left
 166     int remaining = BlockOffset - inputLength;
 167     if ( remaining > 0 ) {
 168         if ( remaining > inputLength ) {
 169             fprintf(stderr, "BgzfStream ERROR: after deflate, remainder too large\n");
 170             exit(1);
 171         }
 172         memcpy(UncompressedBlock, UncompressedBlock + inputLength, remaining);
 173     }
 174
 175     // update block data
 176     BlockOffset = remaining;
 177
 178     // return result
 179     return compressedLength;
 180 }
 181
 182 // flushes the data in the BGZF block
 183 void BgzfStream::FlushBlock(void) {
 184
 185     // flush all of the remaining blocks
 186     while ( BlockOffset > 0 ) {
 187
 188         // compress the data block
 189         int blockLength = DeflateBlock();
 190
 191         // flush the data to our output stream
 192         int numBytesWritten = fwrite(CompressedBlock, 1, blockLength, Stream);
 193         if ( numBytesWritten != blockLength ) {
 194             fprintf(stderr, "BgzfStream ERROR: expected to write %u bytes during flushing, but wrote %u bytes\n",
 195                     blockLength, numBytesWritten);
 196             exit(1);
 197         }
 198
 199         // update block data
 200         BlockAddress += blockLength;
 201     }
 202 }
 203
 204 // decompresses the current block
 205 int BgzfStream::InflateBlock(const int& blockLength) {
 206
 207     // inflate the data from compressed buffer into uncompressed buffer
 208     z_stream zs;
 209     zs.zalloc    = NULL;
 210     zs.zfree     = NULL;
 211     zs.next_in   = (Bytef*)CompressedBlock + 18;
 212     zs.avail_in  = blockLength - 16;
 213     zs.next_out  = (Bytef*)UncompressedBlock;
 214     zs.avail_out = UncompressedBlockSize;
 215
 216     int status = inflateInit2(&zs, Constants::GZIP_WINDOW_BITS);
 217     if ( status != Z_OK ) {
 218         fprintf(stderr, "BgzfStream ERROR: could not decompress block - zlib::inflateInit() failed\n");
 219         return -1;
 220     }
 221
 222     status = inflate(&zs, Z_FINISH);
 223     if ( status != Z_STREAM_END ) {
 224         inflateEnd(&zs);
 225         fprintf(stderr, "BgzfStream ERROR: could not decompress block - zlib::inflate() failed\n");
 226         return -1;
 227     }
 228
 229     status = inflateEnd(&zs);
 230     if ( status != Z_OK ) {
 231         fprintf(stderr, "BgzfStream ERROR: could not decompress block - zlib::inflateEnd() failed\n");
 232         return -1;
 233     }
 234
 235     // return result
 236     return zs.total_out;
 237 }
 238
 239 // opens the BGZF file for reading (mode is either "rb" for reading, or "wb" for writing)
 240 bool BgzfStream::Open(const string& filename, const char* mode) {
 241
 242     // close current stream, if necessary, before opening next
 243     if ( IsOpen ) Close();
 244
 245     // determine open mode
 246     if ( strcmp(mode, "rb") == 0 )
 247         IsWriteOnly = false;
 248     else if ( strcmp(mode, "wb") == 0)
 249         IsWriteOnly = true;
 250     else {
 251         fprintf(stderr, "BgzfStream ERROR: unknown file mode: %s\n", mode);
 252         return false;
 253     }
 254
 255     // open BGZF stream on a file
 256     if ( (filename != "stdin") && (filename != "stdout") )
 257         Stream = fopen(filename.c_str(), mode);
 258
 259     // open BGZF stream on stdin
 260     else if ( (filename == "stdin") && (strcmp(mode, "rb") == 0 ) )
 261         Stream = freopen(NULL, mode, stdin);
 262
 263     // open BGZF stream on stdout
 264     else if ( (filename == "stdout") && (strcmp(mode, "wb") == 0) )
 265         Stream = freopen(NULL, mode, stdout);
 266
 267     if ( !Stream ) {
 268         fprintf(stderr, "BgzfStream ERROR: unable to open file %s\n", filename.c_str() );
 269         return false;
 270     }
 271
 272     // set flag & return success
 273     IsOpen = true;
 274     return true;
 275 }
 276
 277 // reads BGZF data into a byte buffer
 278 int BgzfStream::Read(char* data, const unsigned int dataLength) {
 279
 280     // if stream not open for reading (or empty request)
 281     if ( !IsOpen || IsWriteOnly || dataLength == 0 )
 282         return 0;
 283
 284     // read blocks as needed until desired data length is retrieved
 285     char* output = data;
 286     unsigned int numBytesRead = 0;
 287     while ( numBytesRead < dataLength ) {
 288
 289         // determine bytes available in current block
 290         int bytesAvailable = BlockLength - BlockOffset;
 291
 292         // read (and decompress) next block if needed
 293         if ( bytesAvailable <= 0 ) {
 294             if ( !ReadBlock() ) return -1;
 295             bytesAvailable = BlockLength - BlockOffset;
 296             if ( bytesAvailable <= 0 ) break;
 297         }
 298
 299         // copy data from uncompressed source buffer into data destination buffer
 300         char* buffer   = UncompressedBlock;
 301         int copyLength = min( (int)(dataLength-numBytesRead), bytesAvailable );
 302         memcpy(output, buffer + BlockOffset, copyLength);
 303
 304         // update counters
 305         BlockOffset  += copyLength;
 306         output       += copyLength;
 307         numBytesRead += copyLength;
 308     }
 309
 310     // update block data
 311     if ( BlockOffset == BlockLength ) {
 312         BlockAddress = ftell64(Stream);
 313         BlockOffset  = 0;
 314         BlockLength  = 0;
 315     }
 316
 317     return numBytesRead;
 318 }
 319
 320 // reads a BGZF block
 321 bool BgzfStream::ReadBlock(void) {
 322
 323     char header[Constants::BGZF_BLOCK_HEADER_LENGTH];
 324     int64_t blockAddress = ftell64(Stream);
 325
 326     // read block header from file
 327     int count = fread(header, 1, sizeof(header), Stream);
 328
 329     // if block header empty
 330     if ( count == 0 ) {
 331         BlockLength = 0;
 332         return true;
 333     }
 334
 335     // if block header invalid size
 336     if ( count != sizeof(header) ) {
 337         fprintf(stderr, "BgzfStream ERROR: read block failed - could not read block header\n");
 338         return false;
 339     }
 340
 341     // validate block header contents
 342     if ( !BgzfStream::CheckBlockHeader(header) ) {
 343         fprintf(stderr, "BgzfStream ERROR: read block failed - invalid block header\n");
 344         return false;
 345     }
 346
 347     // copy header contents to compressed buffer
 348     int blockLength = BamTools::UnpackUnsignedShort(&header[16]) + 1;
 349     char* compressedBlock = CompressedBlock;
 350     memcpy(compressedBlock, header, Constants::BGZF_BLOCK_HEADER_LENGTH);
 351     int remaining = blockLength - Constants::BGZF_BLOCK_HEADER_LENGTH;
 352
 353     // read remainder of block
 354     count = fread(&compressedBlock[Constants::BGZF_BLOCK_HEADER_LENGTH], 1, remaining, Stream);
 355     if ( count != remaining ) {
 356         fprintf(stderr, "BgzfStream ERROR: read block failed - could not read data from block\n");
 357         return false;
 358     }
 359
 360     // decompress block data
 361     count = InflateBlock(blockLength);
 362     if ( count < 0 ) {
 363         fprintf(stderr, "BgzfStream ERROR: read block failed - could not decompress block data\n");
 364         return false;
 365     }
 366
 367     // update block data
 368     if ( BlockLength != 0 )
 369         BlockOffset = 0;
 370     BlockAddress = blockAddress;
 371     BlockLength  = count;
 372
 373     // return success
 374     return true;
 375 }
 376
 377 // seek to position in BGZF file
 378 bool BgzfStream::Seek(int64_t position) {
 379
 380     // skip if not open
 381     if ( !IsOpen ) return false;
 382
 383     // determine adjusted offset & address
 384     int     blockOffset  = (position & 0xFFFF);
 385     int64_t blockAddress = (position >> 16) & 0xFFFFFFFFFFFFLL;
 386
 387     // attempt seek in file
 388     if ( fseek64(Stream, blockAddress, SEEK_SET) != 0 ) {
 389         fprintf(stderr, "BgzfStream ERROR: unable to seek in file\n");
 390         return false;
 391     }
 392
 393     // update block data
 394     BlockLength  = 0;
 395     BlockAddress = blockAddress;
 396     BlockOffset  = blockOffset;
 397
 398     // return success
 399     return true;
 400 }
 401
 402 void BgzfStream::SetWriteCompressed(bool ok) {
 403     IsWriteCompressed = ok;
 404 }
 405
 406 // get file position in BGZF file
 407 int64_t BgzfStream::Tell(void) {
 408
 409     // skip if file not open
 410     if ( !IsOpen ) return false;
 411
 412     // otherwise return file pointer position
 413     return ( (BlockAddress << 16) | (BlockOffset & 0xFFFF) );
 414 }
 415
 416 // writes the supplied data into the BGZF buffer
 417 unsigned int BgzfStream::Write(const char* data, const unsigned int dataLen) {
 418
 419     // skip if file not open for writing
 420     if ( !IsOpen || !IsWriteOnly ) return false;
 421
 422     // write blocks as needed til all data is written
 423     unsigned int numBytesWritten = 0;
 424     const char* input = data;
 425     unsigned int blockLength = UncompressedBlockSize;
 426     while ( numBytesWritten < dataLen ) {
 427
 428         // copy data contents to uncompressed output buffer
 429         unsigned int copyLength = min(blockLength - BlockOffset, dataLen - numBytesWritten);
 430         char* buffer = UncompressedBlock;
 431         memcpy(buffer + BlockOffset, input, copyLength);
 432
 433         // update counter
 434         BlockOffset     += copyLength;
 435         input           += copyLength;
 436         numBytesWritten += copyLength;
 437
 438         // flush (& compress) output buffer when full
 439         if ( BlockOffset == blockLength ) FlushBlock();
 440     }
 441
 442     // return result
 443     return numBytesWritten;
 444 }