bgzf.c

   1 /*
   2  * The Broad Institute
   3  * SOFTWARE COPYRIGHT NOTICE AGREEMENT
   4  * This software and its documentation are copyright 2008 by the
   5  * Broad Institute/Massachusetts Institute of Technology. All rights are reserved.
   6  *
   7  * This software is supplied without any warranty or guaranteed support whatsoever.
   8  * Neither the Broad Institute nor MIT can be responsible for its use, misuse,
   9  * or functionality.
  10  */
  11
  12 /*
  13   2009-06-29 by lh3: cache recent uncompressed blocks.
  14   2009-06-25 by lh3: optionally use my knetfile library to access file on a FTP.
  15   2009-06-12 by lh3: support a mode string like "wu" where 'u' for uncompressed output */
  16
  17 #include <stdio.h>
  18 #include <stdlib.h>
  19 #include <string.h>
  20 #include <unistd.h>
  21 #include <fcntl.h>
  22 #include <sys/types.h>
  23 #include <sys/stat.h>
  24 #include "bgzf.h"
  25
  26 #include "khash.h"
  27 typedef struct {
  28         int size;
  29         uint8_t *block;
  30         int64_t end_offset;
  31 } cache_t;
  32 KHASH_MAP_INIT_INT64(cache, cache_t)
  33
  34 #ifdef _NO_LFS
  35 #define ftello(fp) ftell(fp)
  36 #define fseeko(fp, offset, whence) fseek(fp, offset, whence)
  37 #else
  38 extern off_t ftello(FILE *stream);
  39 extern int fseeko(FILE *stream, off_t offset, int whence);
  40 #endif
  41
  42 typedef int8_t byte;
  43
  44 static const int DEFAULT_BLOCK_SIZE = 64 * 1024;
  45 static const int MAX_BLOCK_SIZE = 64 * 1024;
  46
  47 static const int BLOCK_HEADER_LENGTH = 18;
  48 static const int BLOCK_FOOTER_LENGTH = 8;
  49
  50 static const int GZIP_ID1 = 31;
  51 static const int GZIP_ID2 = 139;
  52 static const int CM_DEFLATE = 8;
  53 static const int FLG_FEXTRA = 4;
  54 static const int OS_UNKNOWN = 255;
  55 static const int BGZF_ID1 = 66; // 'B'
  56 static const int BGZF_ID2 = 67; // 'C'
  57 static const int BGZF_LEN = 2;
  58 static const int BGZF_XLEN = 6; // BGZF_LEN+4
  59
  60 static const int GZIP_WINDOW_BITS = -15; // no zlib header
  61 static const int Z_DEFAULT_MEM_LEVEL = 8;
  62
  63
  64 inline
  65 void
  66 packInt16(uint8_t* buffer, uint16_t value)
  67 {
  68     buffer[0] = value;
  69     buffer[1] = value >> 8;
  70 }
  71
  72 inline
  73 int
  74 unpackInt16(const uint8_t* buffer)
  75 {
  76     return (buffer[0] | (buffer[1] << 8));
  77 }
  78
  79 inline
  80 void
  81 packInt32(uint8_t* buffer, uint32_t value)
  82 {
  83     buffer[0] = value;
  84     buffer[1] = value >> 8;
  85     buffer[2] = value >> 16;
  86     buffer[3] = value >> 24;
  87 }
  88
  89 inline
  90 int
  91 min(int x, int y)
  92 {
  93     return (x < y) ? x : y;
  94 }
  95
  96 static
  97 void
  98 report_error(BGZF* fp, const char* message) {
  99     fp->error = message;
 100 }
 101
 102 static BGZF *bgzf_read_init()
 103 {
 104         BGZF *fp;
 105         fp = calloc(1, sizeof(BGZF));
 106     fp->uncompressed_block_size = MAX_BLOCK_SIZE;
 107     fp->uncompressed_block = malloc(MAX_BLOCK_SIZE);
 108     fp->compressed_block_size = MAX_BLOCK_SIZE;
 109     fp->compressed_block = malloc(MAX_BLOCK_SIZE);
 110         fp->cache_size = 0;
 111         fp->cache = kh_init(cache);
 112         return fp;
 113 }
 114
 115 static
 116 BGZF*
 117 open_read(int fd)
 118 {
 119 #ifdef _USE_KNETFILE
 120     knetFile *file = knet_dopen(fd, "r");
 121 #else
 122     FILE* file = fdopen(fd, "r");
 123 #endif
 124     BGZF* fp;
 125         if (file == 0) return 0;
 126         fp = bgzf_read_init();
 127     fp->file_descriptor = fd;
 128     fp->open_mode = 'r';
 129 #ifdef _USE_KNETFILE
 130     fp->x.fpr = file;
 131 #else
 132     fp->file = file;
 133 #endif
 134     return fp;
 135 }
 136
 137 static
 138 BGZF*
 139 open_write(int fd, bool is_uncompressed)
 140 {
 141     FILE* file = fdopen(fd, "w");
 142     BGZF* fp;
 143         if (file == 0) return 0;
 144         fp = malloc(sizeof(BGZF));
 145     fp->file_descriptor = fd;
 146     fp->open_mode = 'w';
 147     fp->owned_file = 0; fp->is_uncompressed = is_uncompressed;
 148 #ifdef _USE_KNETFILE
 149     fp->x.fpw = file;
 150 #else
 151     fp->file = file;
 152 #endif
 153     fp->uncompressed_block_size = DEFAULT_BLOCK_SIZE;
 154     fp->uncompressed_block = NULL;
 155     fp->compressed_block_size = MAX_BLOCK_SIZE;
 156     fp->compressed_block = malloc(MAX_BLOCK_SIZE);
 157     fp->block_address = 0;
 158     fp->block_offset = 0;
 159     fp->block_length = 0;
 160     fp->error = NULL;
 161     return fp;
 162 }
 163
 164 BGZF*
 165 bgzf_open(const char* __restrict path, const char* __restrict mode)
 166 {
 167     BGZF* fp = NULL;
 168     if (mode[0] == 'r' || mode[0] == 'R') { /* The reading mode is preferred. */
 169 #ifdef _USE_KNETFILE
 170                 knetFile *file = knet_open(path, mode);
 171                 if (file == 0) return 0;
 172                 fp = bgzf_read_init();
 173                 fp->file_descriptor = -1;
 174                 fp->open_mode = 'r';
 175                 fp->x.fpr = file;
 176 #else
 177                 int oflag = O_RDONLY;
 178                 int fd = open(path, oflag);
 179                 if (fd == -1) return 0;
 180         fp = open_read(fd);
 181 #endif
 182     } else if (mode[0] == 'w' || mode[0] == 'W') {
 183                 int oflag = O_WRONLY | O_CREAT | O_TRUNC;
 184                 int fd = open(path, oflag, 0644);
 185                 if (fd == -1) return 0;
 186         fp = open_write(fd, strstr(mode, "u")? 1 : 0);
 187     }
 188     if (fp != NULL) {
 189         fp->owned_file = 1;
 190     }
 191     return fp;
 192 }
 193
 194 BGZF*
 195 bgzf_fdopen(int fd, const char * __restrict mode)
 196 {
 197         if (fd == -1) return 0;
 198     if (mode[0] == 'r' || mode[0] == 'R') {
 199         return open_read(fd);
 200     } else if (mode[0] == 'w' || mode[0] == 'W') {
 201         return open_write(fd, strstr(mode, "u")? 1 : 0);
 202     } else {
 203         return NULL;
 204     }
 205 }
 206
 207 static
 208 int
 209 deflate_block(BGZF* fp, int block_length)
 210 {
 211     // Deflate the block in fp->uncompressed_block into fp->compressed_block.
 212     // Also adds an extra field that stores the compressed block length.
 213
 214     byte* buffer = fp->compressed_block;
 215     int buffer_size = fp->compressed_block_size;
 216
 217     // Init gzip header
 218     buffer[0] = GZIP_ID1;
 219     buffer[1] = GZIP_ID2;
 220     buffer[2] = CM_DEFLATE;
 221     buffer[3] = FLG_FEXTRA;
 222     buffer[4] = 0; // mtime
 223     buffer[5] = 0;
 224     buffer[6] = 0;
 225     buffer[7] = 0;
 226     buffer[8] = 0;
 227     buffer[9] = OS_UNKNOWN;
 228     buffer[10] = BGZF_XLEN;
 229     buffer[11] = 0;
 230     buffer[12] = BGZF_ID1;
 231     buffer[13] = BGZF_ID2;
 232     buffer[14] = BGZF_LEN;
 233     buffer[15] = 0;
 234     buffer[16] = 0; // placeholder for block length
 235     buffer[17] = 0;
 236
 237     // loop to retry for blocks that do not compress enough
 238     int input_length = block_length;
 239     int compressed_length = 0;
 240     while (1) {
 241                 int compress_level = fp->is_uncompressed? 0 : Z_DEFAULT_COMPRESSION;
 242         z_stream zs;
 243         zs.zalloc = NULL;
 244         zs.zfree = NULL;
 245         zs.next_in = fp->uncompressed_block;
 246         zs.avail_in = input_length;
 247         zs.next_out = (void*)&buffer[BLOCK_HEADER_LENGTH];
 248         zs.avail_out = buffer_size - BLOCK_HEADER_LENGTH - BLOCK_FOOTER_LENGTH;
 249
 250         int status = deflateInit2(&zs, compress_level, Z_DEFLATED,
 251                                   GZIP_WINDOW_BITS, Z_DEFAULT_MEM_LEVEL, Z_DEFAULT_STRATEGY);
 252         if (status != Z_OK) {
 253             report_error(fp, "deflate init failed");
 254             return -1;
 255         }
 256         status = deflate(&zs, Z_FINISH);
 257         if (status != Z_STREAM_END) {
 258             deflateEnd(&zs);
 259             if (status == Z_OK) {
 260                 // Not enough space in buffer.
 261                 // Can happen in the rare case the input doesn't compress enough.
 262                 // Reduce the amount of input until it fits.
 263                 input_length -= 1024;
 264                 if (input_length <= 0) {
 265                     // should never happen
 266                     report_error(fp, "input reduction failed");
 267                     return -1;
 268                 }
 269                 continue;
 270             }
 271             report_error(fp, "deflate failed");
 272             return -1;
 273         }
 274         status = deflateEnd(&zs);
 275         if (status != Z_OK) {
 276             report_error(fp, "deflate end failed");
 277             return -1;
 278         }
 279         compressed_length = zs.total_out;
 280         compressed_length += BLOCK_HEADER_LENGTH + BLOCK_FOOTER_LENGTH;
 281         if (compressed_length > MAX_BLOCK_SIZE) {
 282             // should never happen
 283             report_error(fp, "deflate overflow");
 284             return -1;
 285         }
 286         break;
 287     }
 288
 289     packInt16((uint8_t*)&buffer[16], compressed_length-1);
 290     uint32_t crc = crc32(0L, NULL, 0L);
 291     crc = crc32(crc, fp->uncompressed_block, input_length);
 292     packInt32((uint8_t*)&buffer[compressed_length-8], crc);
 293     packInt32((uint8_t*)&buffer[compressed_length-4], input_length);
 294
 295     int remaining = block_length - input_length;
 296     if (remaining > 0) {
 297         if (remaining > input_length) {
 298             // should never happen (check so we can use memcpy)
 299             report_error(fp, "remainder too large");
 300             return -1;
 301         }
 302         memcpy(fp->uncompressed_block,
 303                fp->uncompressed_block + input_length,
 304                remaining);
 305     }
 306     fp->block_offset = remaining;
 307     return compressed_length;
 308 }
 309
 310 static
 311 int
 312 inflate_block(BGZF* fp, int block_length)
 313 {
 314     // Inflate the block in fp->compressed_block into fp->uncompressed_block
 315
 316     z_stream zs;
 317     zs.zalloc = NULL;
 318     zs.zfree = NULL;
 319     zs.next_in = fp->compressed_block + 18;
 320     zs.avail_in = block_length - 16;
 321     zs.next_out = fp->uncompressed_block;
 322     zs.avail_out = fp->uncompressed_block_size;
 323
 324     int status = inflateInit2(&zs, GZIP_WINDOW_BITS);
 325     if (status != Z_OK) {
 326         report_error(fp, "inflate init failed");
 327         return -1;
 328     }
 329     status = inflate(&zs, Z_FINISH);
 330     if (status != Z_STREAM_END) {
 331         inflateEnd(&zs);
 332         report_error(fp, "inflate failed");
 333         return -1;
 334     }
 335     status = inflateEnd(&zs);
 336     if (status != Z_OK) {
 337         report_error(fp, "inflate failed");
 338         return -1;
 339     }
 340     return zs.total_out;
 341 }
 342
 343 static
 344 int
 345 check_header(const byte* header)
 346 {
 347     return (header[0] == GZIP_ID1 &&
 348             header[1] == (byte) GZIP_ID2 &&
 349             header[2] == Z_DEFLATED &&
 350             (header[3] & FLG_FEXTRA) != 0 &&
 351             unpackInt16((uint8_t*)&header[10]) == BGZF_XLEN &&
 352             header[12] == BGZF_ID1 &&
 353             header[13] == BGZF_ID2 &&
 354             unpackInt16((uint8_t*)&header[14]) == BGZF_LEN);
 355 }
 356
 357 static void free_cache(BGZF *fp)
 358 {
 359         khint_t k;
 360         khash_t(cache) *h = (khash_t(cache)*)fp->cache;
 361         if (fp->open_mode != 'r') return;
 362         for (k = kh_begin(h); k < kh_end(h); ++k)
 363                 if (kh_exist(h, k)) free(kh_val(h, k).block);
 364         kh_destroy(cache, h);
 365 }
 366
 367 static int load_block_from_cache(BGZF *fp, int64_t block_address)
 368 {
 369         khint_t k;
 370         cache_t *p;
 371         khash_t(cache) *h = (khash_t(cache)*)fp->cache;
 372         k = kh_get(cache, h, block_address);
 373         if (k == kh_end(h)) return 0;
 374         p = &kh_val(h, k);
 375         if (fp->block_length != 0) fp->block_offset = 0;
 376         fp->block_address = block_address;
 377         fp->block_length = p->size;
 378         memcpy(fp->uncompressed_block, p->block, MAX_BLOCK_SIZE);
 379 #ifdef _USE_KNETFILE
 380         knet_seek(fp->x.fpr, p->end_offset, SEEK_SET);
 381 #else
 382         fseeko(fp->file, p->end_offset, SEEK_SET);
 383 #endif
 384         return p->size;
 385 }
 386
 387 static void cache_block(BGZF *fp, int size)
 388 {
 389         int ret;
 390         khint_t k;
 391         cache_t *p;
 392         khash_t(cache) *h = (khash_t(cache)*)fp->cache;
 393         if (MAX_BLOCK_SIZE >= fp->cache_size) return;
 394         if ((kh_size(h) + 1) * MAX_BLOCK_SIZE > fp->cache_size) {
 395                 /* A better way would be to remove the oldest block in the
 396                  * cache, but here we remove a random one for simplicity. This
 397                  * should not have a big impact on performance. */
 398                 for (k = kh_begin(h); k < kh_end(h); ++k)
 399                         if (kh_exist(h, k)) break;
 400                 if (k < kh_end(h)) {
 401                         free(kh_val(h, k).block);
 402                         kh_del(cache, h, k);
 403                 }
 404         }
 405         k = kh_put(cache, h, fp->block_address, &ret);
 406         if (ret == 0) return; // if this happens, a bug!
 407         p = &kh_val(h, k);
 408         p->size = fp->block_length;
 409         p->end_offset = fp->block_address + size;
 410         p->block = malloc(MAX_BLOCK_SIZE);
 411         memcpy(kh_val(h, k).block, fp->uncompressed_block, MAX_BLOCK_SIZE);
 412 }
 413
 414 static
 415 int
 416 read_block(BGZF* fp)
 417 {
 418     byte header[BLOCK_HEADER_LENGTH];
 419         int size = 0;
 420 #ifdef _USE_KNETFILE
 421     int64_t block_address = knet_tell(fp->x.fpr);
 422         if (load_block_from_cache(fp, block_address)) return 0;
 423     int count = knet_read(fp->x.fpr, header, sizeof(header));
 424 #else
 425     int64_t block_address = ftello(fp->file);
 426         if (load_block_from_cache(fp, block_address)) return 0;
 427     int count = fread(header, 1, sizeof(header), fp->file);
 428 #endif
 429     if (count == 0) {
 430         fp->block_length = 0;
 431         return 0;
 432     }
 433         size = count;
 434     if (count != sizeof(header)) {
 435         report_error(fp, "read failed");
 436         return -1;
 437     }
 438     if (!check_header(header)) {
 439         report_error(fp, "invalid block header");
 440         return -1;
 441     }
 442     int block_length = unpackInt16((uint8_t*)&header[16]) + 1;
 443     byte* compressed_block = (byte*) fp->compressed_block;
 444     memcpy(compressed_block, header, BLOCK_HEADER_LENGTH);
 445     int remaining = block_length - BLOCK_HEADER_LENGTH;
 446 #ifdef _USE_KNETFILE
 447     count = knet_read(fp->x.fpr, &compressed_block[BLOCK_HEADER_LENGTH], remaining);
 448 #else
 449     count = fread(&compressed_block[BLOCK_HEADER_LENGTH], 1, remaining, fp->file);
 450 #endif
 451     if (count != remaining) {
 452         report_error(fp, "read failed");
 453         return -1;
 454     }
 455         size += count;
 456     count = inflate_block(fp, block_length);
 457     if (count < 0) {
 458         return -1;
 459     }
 460     if (fp->block_length != 0) {
 461         // Do not reset offset if this read follows a seek.
 462         fp->block_offset = 0;
 463     }
 464     fp->block_address = block_address;
 465     fp->block_length = count;
 466         cache_block(fp, size);
 467     return 0;
 468 }
 469
 470 int
 471 bgzf_read(BGZF* fp, void* data, int length)
 472 {
 473     if (length <= 0) {
 474         return 0;
 475     }
 476     if (fp->open_mode != 'r') {
 477         report_error(fp, "file not open for reading");
 478         return -1;
 479     }
 480
 481     int bytes_read = 0;
 482     byte* output = data;
 483     while (bytes_read < length) {
 484         int available = fp->block_length - fp->block_offset;
 485         if (available <= 0) {
 486             if (read_block(fp) != 0) {
 487                 return -1;
 488             }
 489             available = fp->block_length - fp->block_offset;
 490             if (available <= 0) {
 491                 break;
 492             }
 493         }
 494         int copy_length = min(length-bytes_read, available);
 495         byte* buffer = fp->uncompressed_block;
 496         memcpy(output, buffer + fp->block_offset, copy_length);
 497         fp->block_offset += copy_length;
 498         output += copy_length;
 499         bytes_read += copy_length;
 500     }
 501     if (fp->block_offset == fp->block_length) {
 502 #ifdef _USE_KNETFILE
 503         fp->block_address = knet_tell(fp->x.fpr);
 504 #else
 505         fp->block_address = ftello(fp->file);
 506 #endif
 507         fp->block_offset = 0;
 508         fp->block_length = 0;
 509     }
 510     return bytes_read;
 511 }
 512
 513 static
 514 int
 515 flush_block(BGZF* fp)
 516 {
 517     while (fp->block_offset > 0) {
 518         int block_length = deflate_block(fp, fp->block_offset);
 519         if (block_length < 0) {
 520             return -1;
 521         }
 522 #ifdef _USE_KNETFILE
 523         int count = fwrite(fp->compressed_block, 1, block_length, fp->x.fpw);
 524 #else
 525         int count = fwrite(fp->compressed_block, 1, block_length, fp->file);
 526 #endif
 527         if (count != block_length) {
 528             report_error(fp, "write failed");
 529             return -1;
 530         }
 531         fp->block_address += block_length;
 532     }
 533     return 0;
 534 }
 535
 536 int
 537 bgzf_write(BGZF* fp, const void* data, int length)
 538 {
 539     if (fp->open_mode != 'w') {
 540         report_error(fp, "file not open for writing");
 541         return -1;
 542     }
 543
 544     if (fp->uncompressed_block == NULL) {
 545         fp->uncompressed_block = malloc(fp->uncompressed_block_size);
 546     }
 547
 548     const byte* input = data;
 549     int block_length = fp->uncompressed_block_size;
 550     int bytes_written = 0;
 551     while (bytes_written < length) {
 552         int copy_length = min(block_length - fp->block_offset, length - bytes_written);
 553         byte* buffer = fp->uncompressed_block;
 554         memcpy(buffer + fp->block_offset, input, copy_length);
 555         fp->block_offset += copy_length;
 556         input += copy_length;
 557         bytes_written += copy_length;
 558         if (fp->block_offset == block_length) {
 559             if (flush_block(fp) != 0) {
 560                 break;
 561             }
 562         }
 563     }
 564     return bytes_written;
 565 }
 566
 567 int
 568 bgzf_close(BGZF* fp)
 569 {
 570     if (fp->open_mode == 'w') {
 571         if (flush_block(fp) != 0) {
 572             return -1;
 573         }
 574 #ifdef _USE_KNETFILE
 575         if (fflush(fp->x.fpw) != 0) {
 576 #else
 577         if (fflush(fp->file) != 0) {
 578 #endif
 579             report_error(fp, "flush failed");
 580             return -1;
 581         }
 582     }
 583     if (fp->owned_file) {
 584 #ifdef _USE_KNETFILE
 585                 int ret;
 586                 if (fp->open_mode == 'w') ret = fclose(fp->x.fpw);
 587                 else ret = knet_close(fp->x.fpr);
 588         if (ret != 0) return -1;
 589 #else
 590         if (fclose(fp->file) != 0) {
 591             return -1;
 592         }
 593 #endif
 594     }
 595     free(fp->uncompressed_block);
 596     free(fp->compressed_block);
 597         free_cache(fp);
 598     free(fp);
 599     return 0;
 600 }
 601
 602 int64_t
 603 bgzf_tell(BGZF* fp)
 604 {
 605     return ((fp->block_address << 16) | (fp->block_offset & 0xFFFF));
 606 }
 607
 608 void bgzf_set_cache_size(BGZF *fp, int cache_size)
 609 {
 610         if (fp) fp->cache_size = cache_size;
 611 }
 612
 613 int64_t
 614 bgzf_seek(BGZF* fp, int64_t pos, int where)
 615 {
 616     if (fp->open_mode != 'r') {
 617         report_error(fp, "file not open for read");
 618         return -1;
 619     }
 620     if (where != SEEK_SET) {
 621         report_error(fp, "unimplemented seek option");
 622         return -1;
 623     }
 624     int block_offset = pos & 0xFFFF;
 625     int64_t block_address = (pos >> 16) & 0xFFFFFFFFFFFFLL;
 626 #ifdef _USE_KNETFILE
 627     if (knet_seek(fp->x.fpr, block_address, SEEK_SET) != 0) {
 628 #else
 629     if (fseeko(fp->file, block_address, SEEK_SET) != 0) {
 630 #endif
 631         report_error(fp, "seek failed");
 632         return -1;
 633     }
 634     fp->block_length = 0;  // indicates current block is not loaded
 635     fp->block_address = block_address;
 636     fp->block_offset = block_offset;
 637     return 0;
 638 }
 639