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