bam_index.c

   1 #include <ctype.h>
   2 #include <assert.h>
   3 #include "bam.h"
   4 #include "khash.h"
   5 #include "ksort.h"
   6 #include "bam_endian.h"
   7 #ifdef _USE_KNETFILE
   8 #include "knetfile.h"
   9 #endif
  10
  11 /*!
  12   @header
  13
  14   Alignment indexing. Before indexing, BAM must be sorted based on the
  15   leftmost coordinate of alignments. In indexing, BAM uses two indices:
  16   a UCSC binning index and a simple linear index. The binning index is
  17   efficient for alignments spanning long distance, while the auxiliary
  18   linear index helps to reduce unnecessary seek calls especially for
  19   short alignments.
  20
  21   The UCSC binning scheme was suggested by Richard Durbin and Lincoln
  22   Stein and is explained by Kent et al. (2002). In this scheme, each bin
  23   represents a contiguous genomic region which can be fully contained in
  24   another bin; each alignment is associated with a bin which represents
  25   the smallest region containing the entire alignment. The binning
  26   scheme is essentially another representation of R-tree. A distinct bin
  27   uniquely corresponds to a distinct internal node in a R-tree. Bin A is
  28   a child of Bin B if region A is contained in B.
  29
  30   In BAM, each bin may span 2^29, 2^26, 2^23, 2^20, 2^17 or 2^14 bp. Bin
  31   0 spans a 512Mbp region, bins 1-8 span 64Mbp, 9-72 8Mbp, 73-584 1Mbp,
  32   585-4680 128Kbp and bins 4681-37449 span 16Kbp regions. If we want to
  33   find the alignments overlapped with a region [rbeg,rend), we need to
  34   calculate the list of bins that may be overlapped the region and test
  35   the alignments in the bins to confirm the overlaps. If the specified
  36   region is short, typically only a few alignments in six bins need to
  37   be retrieved. The overlapping alignments can be quickly fetched.
  38
  39  */
  40
  41 #define BAM_MIN_CHUNK_GAP 32768
  42 // 1<<14 is the size of minimum bin.
  43 #define BAM_LIDX_SHIFT    14
  44
  45 #define BAM_MAX_BIN 37450 // =(8^6-1)/7+1
  46
  47 typedef struct {
  48         uint64_t u, v;
  49 } pair64_t;
  50
  51 #define pair64_lt(a,b) ((a).u < (b).u)
  52 KSORT_INIT(off, pair64_t, pair64_lt)
  53
  54 typedef struct {
  55         uint32_t m, n;
  56         pair64_t *list;
  57 } bam_binlist_t;
  58
  59 typedef struct {
  60         int32_t n, m;
  61         uint64_t *offset;
  62 } bam_lidx_t;
  63
  64 KHASH_MAP_INIT_INT(i, bam_binlist_t)
  65
  66 struct __bam_index_t {
  67         int32_t n;
  68         uint64_t n_no_coor; // unmapped reads without coordinate
  69         khash_t(i) **index;
  70         bam_lidx_t *index2;
  71 };
  72
  73 // requirement: len <= LEN_MASK
  74 static inline void insert_offset(khash_t(i) *h, int bin, uint64_t beg, uint64_t end)
  75 {
  76         khint_t k;
  77         bam_binlist_t *l;
  78         int ret;
  79         k = kh_put(i, h, bin, &ret);
  80         l = &kh_value(h, k);
  81         if (ret) { // not present
  82                 l->m = 1; l->n = 0;
  83                 l->list = (pair64_t*)calloc(l->m, 16);
  84         }
  85         if (l->n == l->m) {
  86                 l->m <<= 1;
  87                 l->list = (pair64_t*)realloc(l->list, l->m * 16);
  88         }
  89         l->list[l->n].u = beg; l->list[l->n++].v = end;
  90 }
  91
  92 static inline void insert_offset2(bam_lidx_t *index2, bam1_t *b, uint64_t offset)
  93 {
  94         int i, beg, end;
  95         beg = b->core.pos >> BAM_LIDX_SHIFT;
  96         end = (bam_calend(&b->core, bam1_cigar(b)) - 1) >> BAM_LIDX_SHIFT;
  97         if (index2->m < end + 1) {
  98                 int old_m = index2->m;
  99                 index2->m = end + 1;
 100                 kroundup32(index2->m);
 101                 index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8);
 102                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
 103         }
 104         if (beg == end) {
 105                 if (index2->offset[beg] == 0) index2->offset[beg] = offset;
 106         } else {
 107                 for (i = beg; i <= end; ++i)
 108                         if (index2->offset[i] == 0) index2->offset[i] = offset;
 109         }
 110         index2->n = end + 1;
 111 }
 112
 113 static void merge_chunks(bam_index_t *idx)
 114 {
 115 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 116         khash_t(i) *index;
 117         int i, l, m;
 118         khint_t k;
 119         for (i = 0; i < idx->n; ++i) {
 120                 index = idx->index[i];
 121                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 122                         bam_binlist_t *p;
 123                         if (!kh_exist(index, k) || kh_key(index, k) == BAM_MAX_BIN) continue;
 124                         p = &kh_value(index, k);
 125                         m = 0;
 126                         for (l = 1; l < p->n; ++l) {
 127 #ifdef BAM_TRUE_OFFSET
 128                                 if (p->list[m].v + BAM_MIN_CHUNK_GAP > p->list[l].u) p->list[m].v = p->list[l].v;
 129 #else
 130                                 if (p->list[m].v>>16 == p->list[l].u>>16) p->list[m].v = p->list[l].v;
 131 #endif
 132                                 else p->list[++m] = p->list[l];
 133                         } // ~for(l)
 134                         p->n = m + 1;
 135                 } // ~for(k)
 136         } // ~for(i)
 137 #endif // defined(BAM_TRUE_OFFSET) || defined(BAM_BGZF)
 138 }
 139
 140 static void fill_missing(bam_index_t *idx)
 141 {
 142         int i, j;
 143         for (i = 0; i < idx->n; ++i) {
 144                 bam_lidx_t *idx2 = &idx->index2[i];
 145                 for (j = 1; j < idx2->n; ++j)
 146                         if (idx2->offset[j] == 0)
 147                                 idx2->offset[j] = idx2->offset[j-1];
 148         }
 149 }
 150
 151 bam_index_t *bam_index_core(bamFile fp)
 152 {
 153         bam1_t *b;
 154         bam_header_t *h;
 155         int i, ret;
 156         bam_index_t *idx;
 157         uint32_t last_bin, save_bin;
 158         int32_t last_coor, last_tid, save_tid;
 159         bam1_core_t *c;
 160         uint64_t save_off, last_off, n_mapped, n_unmapped, off_beg, off_end, n_no_coor;
 161
 162         idx = (bam_index_t*)calloc(1, sizeof(bam_index_t));
 163         b = (bam1_t*)calloc(1, sizeof(bam1_t));
 164         h = bam_header_read(fp);
 165         c = &b->core;
 166
 167         idx->n = h->n_targets;
 168         bam_header_destroy(h);
 169         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 170         for (i = 0; i < idx->n; ++i) idx->index[i] = kh_init(i);
 171         idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t));
 172
 173         save_bin = save_tid = last_tid = last_bin = 0xffffffffu;
 174         save_off = last_off = bam_tell(fp); last_coor = 0xffffffffu;
 175     n_mapped = n_unmapped = n_no_coor = off_end = 0;
 176         off_beg = off_end = bam_tell(fp);
 177         while ((ret = bam_read1(fp, b)) >= 0) {
 178                 if (c->tid < 0) ++n_no_coor;
 179                 if (last_tid != c->tid) { // change of chromosomes
 180                         last_tid = c->tid;
 181                         last_bin = 0xffffffffu;
 182                 } else if (last_coor > c->pos) {
 183                         fprintf(stderr, "[bam_index_core] the alignment is not sorted (%s): %u > %u in %d-th chr\n",
 184                                         bam1_qname(b), last_coor, c->pos, c->tid+1);
 185                         exit(1);
 186                 }
 187                 if (c->tid >= 0) insert_offset2(&idx->index2[b->core.tid], b, last_off);
 188                 if (c->bin != last_bin) { // then possibly write the binning index
 189                         if (save_bin != 0xffffffffu) // save_bin==0xffffffffu only happens to the first record
 190                                 insert_offset(idx->index[save_tid], save_bin, save_off, last_off);
 191                         if (last_bin == 0xffffffffu && save_tid != 0xffffffffu) { // write the meta element
 192                                 off_end = last_off;
 193                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, off_end);
 194                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 195                                 n_mapped = n_unmapped = 0;
 196                                 off_beg = off_end;
 197                         }
 198                         save_off = last_off;
 199                         save_bin = last_bin = c->bin;
 200                         save_tid = c->tid;
 201                         if (save_tid < 0) break;
 202                 }
 203                 if (bam_tell(fp) <= last_off) {
 204                         fprintf(stderr, "[bam_index_core] bug in BGZF/RAZF: %llx < %llx\n",
 205                                         (unsigned long long)bam_tell(fp), (unsigned long long)last_off);
 206                         exit(1);
 207                 }
 208                 if (c->flag & BAM_FUNMAP) ++n_unmapped;
 209                 else ++n_mapped;
 210                 last_off = bam_tell(fp);
 211                 last_coor = b->core.pos;
 212         }
 213         if (save_tid >= 0) {
 214                 insert_offset(idx->index[save_tid], save_bin, save_off, bam_tell(fp));
 215                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, off_end);
 216                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 217         }
 218         merge_chunks(idx);
 219         fill_missing(idx);
 220         if (ret >= 0)
 221                 while ((ret = bam_read1(fp, b)) >= 0) ++n_no_coor;
 222         if (ret < -1) fprintf(stderr, "[bam_index_core] truncated file? Continue anyway. (%d)\n", ret);
 223         free(b->data); free(b);
 224         idx->n_no_coor = n_no_coor;
 225         return idx;
 226 }
 227
 228 void bam_index_destroy(bam_index_t *idx)
 229 {
 230         khint_t k;
 231         int i;
 232         if (idx == 0) return;
 233         for (i = 0; i < idx->n; ++i) {
 234                 khash_t(i) *index = idx->index[i];
 235                 bam_lidx_t *index2 = idx->index2 + i;
 236                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 237                         if (kh_exist(index, k))
 238                                 free(kh_value(index, k).list);
 239                 }
 240                 kh_destroy(i, index);
 241                 free(index2->offset);
 242         }
 243         free(idx->index); free(idx->index2);
 244         free(idx);
 245 }
 246
 247 void bam_index_save(const bam_index_t *idx, FILE *fp)
 248 {
 249         int32_t i, size;
 250         khint_t k;
 251         fwrite("BAI\1", 1, 4, fp);
 252         if (bam_is_be) {
 253                 uint32_t x = idx->n;
 254                 fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 255         } else fwrite(&idx->n, 4, 1, fp);
 256         for (i = 0; i < idx->n; ++i) {
 257                 khash_t(i) *index = idx->index[i];
 258                 bam_lidx_t *index2 = idx->index2 + i;
 259                 // write binning index
 260                 size = kh_size(index);
 261                 if (bam_is_be) { // big endian
 262                         uint32_t x = size;
 263                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 264                 } else fwrite(&size, 4, 1, fp);
 265                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 266                         if (kh_exist(index, k)) {
 267                                 bam_binlist_t *p = &kh_value(index, k);
 268                                 if (bam_is_be) { // big endian
 269                                         uint32_t x;
 270                                         x = kh_key(index, k); fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 271                                         x = p->n; fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 272                                         for (x = 0; (int)x < p->n; ++x) {
 273                                                 bam_swap_endian_8p(&p->list[x].u);
 274                                                 bam_swap_endian_8p(&p->list[x].v);
 275                                         }
 276                                         fwrite(p->list, 16, p->n, fp);
 277                                         for (x = 0; (int)x < p->n; ++x) {
 278                                                 bam_swap_endian_8p(&p->list[x].u);
 279                                                 bam_swap_endian_8p(&p->list[x].v);
 280                                         }
 281                                 } else {
 282                                         fwrite(&kh_key(index, k), 4, 1, fp);
 283                                         fwrite(&p->n, 4, 1, fp);
 284                                         fwrite(p->list, 16, p->n, fp);
 285                                 }
 286                         }
 287                 }
 288                 // write linear index (index2)
 289                 if (bam_is_be) {
 290                         int x = index2->n;
 291                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 292                 } else fwrite(&index2->n, 4, 1, fp);
 293                 if (bam_is_be) { // big endian
 294                         int x;
 295                         for (x = 0; (int)x < index2->n; ++x)
 296                                 bam_swap_endian_8p(&index2->offset[x]);
 297                         fwrite(index2->offset, 8, index2->n, fp);
 298                         for (x = 0; (int)x < index2->n; ++x)
 299                                 bam_swap_endian_8p(&index2->offset[x]);
 300                 } else fwrite(index2->offset, 8, index2->n, fp);
 301         }
 302         { // write the number of reads coor-less records.
 303                 uint64_t x = idx->n_no_coor;
 304                 if (bam_is_be) bam_swap_endian_8p(&x);
 305                 fwrite(&x, 8, 1, fp);
 306         }
 307         fflush(fp);
 308 }
 309
 310 static bam_index_t *bam_index_load_core(FILE *fp)
 311 {
 312         int i;
 313         char magic[4];
 314         bam_index_t *idx;
 315         if (fp == 0) {
 316                 fprintf(stderr, "[bam_index_load_core] fail to load index.\n");
 317                 return 0;
 318         }
 319         fread(magic, 1, 4, fp);
 320         if (strncmp(magic, "BAI\1", 4)) {
 321                 fprintf(stderr, "[bam_index_load] wrong magic number.\n");
 322                 fclose(fp);
 323                 return 0;
 324         }
 325         idx = (bam_index_t*)calloc(1, sizeof(bam_index_t));
 326         fread(&idx->n, 4, 1, fp);
 327         if (bam_is_be) bam_swap_endian_4p(&idx->n);
 328         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 329         idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t));
 330         for (i = 0; i < idx->n; ++i) {
 331                 khash_t(i) *index;
 332                 bam_lidx_t *index2 = idx->index2 + i;
 333                 uint32_t key, size;
 334                 khint_t k;
 335                 int j, ret;
 336                 bam_binlist_t *p;
 337                 index = idx->index[i] = kh_init(i);
 338                 // load binning index
 339                 fread(&size, 4, 1, fp);
 340                 if (bam_is_be) bam_swap_endian_4p(&size);
 341                 for (j = 0; j < (int)size; ++j) {
 342                         fread(&key, 4, 1, fp);
 343                         if (bam_is_be) bam_swap_endian_4p(&key);
 344                         k = kh_put(i, index, key, &ret);
 345                         p = &kh_value(index, k);
 346                         fread(&p->n, 4, 1, fp);
 347                         if (bam_is_be) bam_swap_endian_4p(&p->n);
 348                         p->m = p->n;
 349                         p->list = (pair64_t*)malloc(p->m * 16);
 350                         fread(p->list, 16, p->n, fp);
 351                         if (bam_is_be) {
 352                                 int x;
 353                                 for (x = 0; x < p->n; ++x) {
 354                                         bam_swap_endian_8p(&p->list[x].u);
 355                                         bam_swap_endian_8p(&p->list[x].v);
 356                                 }
 357                         }
 358                 }
 359                 // load linear index
 360                 fread(&index2->n, 4, 1, fp);
 361                 if (bam_is_be) bam_swap_endian_4p(&index2->n);
 362                 index2->m = index2->n;
 363                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 364                 fread(index2->offset, index2->n, 8, fp);
 365                 if (bam_is_be)
 366                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 367         }
 368         if (fread(&idx->n_no_coor, 8, 1, fp) == 0) idx->n_no_coor = 0;
 369         if (bam_is_be) bam_swap_endian_8p(&idx->n_no_coor);
 370         return idx;
 371 }
 372
 373 bam_index_t *bam_index_load_local(const char *_fn)
 374 {
 375         FILE *fp;
 376         char *fnidx, *fn;
 377
 378         if (strstr(_fn, "ftp://") == _fn || strstr(_fn, "http://") == _fn) {
 379                 const char *p;
 380                 int l = strlen(_fn);
 381                 for (p = _fn + l - 1; p >= _fn; --p)
 382                         if (*p == '/') break;
 383                 fn = strdup(p + 1);
 384         } else fn = strdup(_fn);
 385         fnidx = (char*)calloc(strlen(fn) + 5, 1);
 386         strcpy(fnidx, fn); strcat(fnidx, ".bai");
 387         fp = fopen(fnidx, "rb");
 388         if (fp == 0) { // try "{base}.bai"
 389                 char *s = strstr(fn, "bam");
 390                 if (s == fn + strlen(fn) - 3) {
 391                         strcpy(fnidx, fn);
 392                         fnidx[strlen(fn)-1] = 'i';
 393                         fp = fopen(fnidx, "rb");
 394                 }
 395         }
 396         free(fnidx); free(fn);
 397         if (fp) {
 398                 bam_index_t *idx = bam_index_load_core(fp);
 399                 fclose(fp);
 400                 return idx;
 401         } else return 0;
 402 }
 403
 404 #ifdef _USE_KNETFILE
 405 static void download_from_remote(const char *url)
 406 {
 407         const int buf_size = 1 * 1024 * 1024;
 408         char *fn;
 409         FILE *fp;
 410         uint8_t *buf;
 411         knetFile *fp_remote;
 412         int l;
 413         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 414         l = strlen(url);
 415         for (fn = (char*)url + l - 1; fn >= url; --fn)
 416                 if (*fn == '/') break;
 417         ++fn; // fn now points to the file name
 418         fp_remote = knet_open(url, "r");
 419         if (fp_remote == 0) {
 420                 fprintf(stderr, "[download_from_remote] fail to open remote file.\n");
 421                 return;
 422         }
 423         if ((fp = fopen(fn, "wb")) == 0) {
 424                 fprintf(stderr, "[download_from_remote] fail to create file in the working directory.\n");
 425                 knet_close(fp_remote);
 426                 return;
 427         }
 428         buf = (uint8_t*)calloc(buf_size, 1);
 429         while ((l = knet_read(fp_remote, buf, buf_size)) != 0)
 430                 fwrite(buf, 1, l, fp);
 431         free(buf);
 432         fclose(fp);
 433         knet_close(fp_remote);
 434 }
 435 #else
 436 static void download_from_remote(const char *url)
 437 {
 438         return;
 439 }
 440 #endif
 441
 442 bam_index_t *bam_index_load(const char *fn)
 443 {
 444         bam_index_t *idx;
 445         idx = bam_index_load_local(fn);
 446         if (idx == 0 && (strstr(fn, "ftp://") == fn || strstr(fn, "http://") == fn)) {
 447                 char *fnidx = calloc(strlen(fn) + 5, 1);
 448                 strcat(strcpy(fnidx, fn), ".bai");
 449                 fprintf(stderr, "[bam_index_load] attempting to download the remote index file.\n");
 450                 download_from_remote(fnidx);
 451                 idx = bam_index_load_local(fn);
 452         }
 453         if (idx == 0) fprintf(stderr, "[bam_index_load] fail to load BAM index.\n");
 454         return idx;
 455 }
 456
 457 int bam_index_build2(const char *fn, const char *_fnidx)
 458 {
 459         char *fnidx;
 460         FILE *fpidx;
 461         bamFile fp;
 462         bam_index_t *idx;
 463         if ((fp = bam_open(fn, "r")) == 0) {
 464                 fprintf(stderr, "[bam_index_build2] fail to open the BAM file.\n");
 465                 return -1;
 466         }
 467         idx = bam_index_core(fp);
 468         bam_close(fp);
 469         if (_fnidx == 0) {
 470                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 471                 strcpy(fnidx, fn); strcat(fnidx, ".bai");
 472         } else fnidx = strdup(_fnidx);
 473         fpidx = fopen(fnidx, "wb");
 474         if (fpidx == 0) {
 475                 fprintf(stderr, "[bam_index_build2] fail to create the index file.\n");
 476                 free(fnidx);
 477                 return -1;
 478         }
 479         bam_index_save(idx, fpidx);
 480         bam_index_destroy(idx);
 481         fclose(fpidx);
 482         free(fnidx);
 483         return 0;
 484 }
 485
 486 int bam_index_build(const char *fn)
 487 {
 488         return bam_index_build2(fn, 0);
 489 }
 490
 491 int bam_index(int argc, char *argv[])
 492 {
 493         if (argc < 2) {
 494                 fprintf(stderr, "Usage: samtools index <in.bam> [out.index]\n");
 495                 return 1;
 496         }
 497         if (argc >= 3) bam_index_build2(argv[1], argv[2]);
 498         else bam_index_build(argv[1]);
 499         return 0;
 500 }
 501
 502 int bam_idxstats(int argc, char *argv[])
 503 {
 504         bam_index_t *idx;
 505         bam_header_t *header;
 506         bamFile fp;
 507         int i;
 508         if (argc < 2) {
 509                 fprintf(stderr, "Usage: samtools idxstats <in.bam>\n");
 510                 return 1;
 511         }
 512         fp = bam_open(argv[1], "r");
 513         if (fp == 0) { fprintf(stderr, "[%s] fail to open BAM.\n", __func__); return 1; }
 514         header = bam_header_read(fp);
 515         bam_close(fp);
 516         idx = bam_index_load(argv[1]);
 517         if (idx == 0) { fprintf(stderr, "[%s] fail to load the index.\n", __func__); return 1; }
 518         for (i = 0; i < idx->n; ++i) {
 519                 khint_t k;
 520                 khash_t(i) *h = idx->index[i];
 521                 printf("%s\t%d", header->target_name[i], header->target_len[i]);
 522                 k = kh_get(i, h, BAM_MAX_BIN);
 523                 if (k != kh_end(h))
 524                         printf("\t%llu\t%llu", (long long)kh_val(h, k).list[1].u, (long long)kh_val(h, k).list[1].v);
 525                 else printf("\t0\t0");
 526                 putchar('\n');
 527         }
 528         printf("*\t0\t0\t%llu\n", (long long)idx->n_no_coor);
 529         bam_header_destroy(header);
 530         bam_index_destroy(idx);
 531         return 0;
 532 }
 533
 534 static inline int reg2bins(uint32_t beg, uint32_t end, uint16_t list[BAM_MAX_BIN])
 535 {
 536         int i = 0, k;
 537         if (beg >= end) return 0;
 538         if (end >= 1u<<29) end = 1u<<29;
 539         --end;
 540         list[i++] = 0;
 541         for (k =    1 + (beg>>26); k <=    1 + (end>>26); ++k) list[i++] = k;
 542         for (k =    9 + (beg>>23); k <=    9 + (end>>23); ++k) list[i++] = k;
 543         for (k =   73 + (beg>>20); k <=   73 + (end>>20); ++k) list[i++] = k;
 544         for (k =  585 + (beg>>17); k <=  585 + (end>>17); ++k) list[i++] = k;
 545         for (k = 4681 + (beg>>14); k <= 4681 + (end>>14); ++k) list[i++] = k;
 546         return i;
 547 }
 548
 549 static inline int is_overlap(uint32_t beg, uint32_t end, const bam1_t *b)
 550 {
 551         uint32_t rbeg = b->core.pos;
 552         uint32_t rend = b->core.n_cigar? bam_calend(&b->core, bam1_cigar(b)) : b->core.pos + 1;
 553         return (rend > beg && rbeg < end);
 554 }
 555
 556 struct __bam_iter_t {
 557         int from_first; // read from the first record; no random access
 558         int tid, beg, end, n_off, i, finished;
 559         uint64_t curr_off;
 560         pair64_t *off;
 561 };
 562
 563 // bam_fetch helper function retrieves
 564 bam_iter_t bam_iter_query(const bam_index_t *idx, int tid, int beg, int end)
 565 {
 566         uint16_t *bins;
 567         int i, n_bins, n_off;
 568         pair64_t *off;
 569         khint_t k;
 570         khash_t(i) *index;
 571         uint64_t min_off;
 572         bam_iter_t iter = 0;
 573
 574         if (beg < 0) beg = 0;
 575         if (end < beg) return 0;
 576         // initialize iter
 577         iter = calloc(1, sizeof(struct __bam_iter_t));
 578         iter->tid = tid, iter->beg = beg, iter->end = end; iter->i = -1;
 579         //
 580         bins = (uint16_t*)calloc(BAM_MAX_BIN, 2);
 581         n_bins = reg2bins(beg, end, bins);
 582         index = idx->index[tid];
 583         if (idx->index2[tid].n > 0) {
 584                 min_off = (beg>>BAM_LIDX_SHIFT >= idx->index2[tid].n)? idx->index2[tid].offset[idx->index2[tid].n-1]
 585                         : idx->index2[tid].offset[beg>>BAM_LIDX_SHIFT];
 586                 if (min_off == 0) { // improvement for index files built by tabix prior to 0.1.4
 587                         int n = beg>>BAM_LIDX_SHIFT;
 588                         if (n > idx->index2[tid].n) n = idx->index2[tid].n;
 589                         for (i = n - 1; i >= 0; --i)
 590                                 if (idx->index2[tid].offset[i] != 0) break;
 591                         if (i >= 0) min_off = idx->index2[tid].offset[i];
 592                 }
 593         } else min_off = 0; // tabix 0.1.2 may produce such index files
 594         for (i = n_off = 0; i < n_bins; ++i) {
 595                 if ((k = kh_get(i, index, bins[i])) != kh_end(index))
 596                         n_off += kh_value(index, k).n;
 597         }
 598         if (n_off == 0) {
 599                 free(bins); return iter;
 600         }
 601         off = (pair64_t*)calloc(n_off, 16);
 602         for (i = n_off = 0; i < n_bins; ++i) {
 603                 if ((k = kh_get(i, index, bins[i])) != kh_end(index)) {
 604                         int j;
 605                         bam_binlist_t *p = &kh_value(index, k);
 606                         for (j = 0; j < p->n; ++j)
 607                                 if (p->list[j].v > min_off) off[n_off++] = p->list[j];
 608                 }
 609         }
 610         free(bins);
 611         {
 612                 bam1_t *b = (bam1_t*)calloc(1, sizeof(bam1_t));
 613                 int l;
 614                 ks_introsort(off, n_off, off);
 615                 // resolve completely contained adjacent blocks
 616                 for (i = 1, l = 0; i < n_off; ++i)
 617                         if (off[l].v < off[i].v)
 618                                 off[++l] = off[i];
 619                 n_off = l + 1;
 620                 // resolve overlaps between adjacent blocks; this may happen due to the merge in indexing
 621                 for (i = 1; i < n_off; ++i)
 622                         if (off[i-1].v >= off[i].u) off[i-1].v = off[i].u;
 623                 { // merge adjacent blocks
 624 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 625                         for (i = 1, l = 0; i < n_off; ++i) {
 626 #ifdef BAM_TRUE_OFFSET
 627                                 if (off[l].v + BAM_MIN_CHUNK_GAP > off[i].u) off[l].v = off[i].v;
 628 #else
 629                                 if (off[l].v>>16 == off[i].u>>16) off[l].v = off[i].v;
 630 #endif
 631                                 else off[++l] = off[i];
 632                         }
 633                         n_off = l + 1;
 634 #endif
 635                 }
 636                 bam_destroy1(b);
 637         }
 638         iter->n_off = n_off; iter->off = off;
 639         return iter;
 640 }
 641
 642 pair64_t *get_chunk_coordinates(const bam_index_t *idx, int tid, int beg, int end, int *cnt_off)
 643 { // for pysam compatibility
 644         bam_iter_t iter;
 645         pair64_t *off;
 646         iter = bam_iter_query(idx, tid, beg, end);
 647         off = iter->off; *cnt_off = iter->n_off;
 648         free(iter);
 649         return off;
 650 }
 651
 652 void bam_iter_destroy(bam_iter_t iter)
 653 {
 654         if (iter) { free(iter->off); free(iter); }
 655 }
 656
 657 int bam_iter_read(bamFile fp, bam_iter_t iter, bam1_t *b)
 658 {
 659         if (iter && iter->finished) return -1;
 660         if (iter == 0 || iter->from_first) {
 661                 int ret = bam_read1(fp, b);
 662                 if (ret < 0 && iter) iter->finished = 1;
 663                 return ret;
 664         }
 665         if (iter->off == 0) return -1;
 666         for (;;) {
 667                 int ret;
 668                 if (iter->curr_off == 0 || iter->curr_off >= iter->off[iter->i].v) { // then jump to the next chunk
 669                         if (iter->i == iter->n_off - 1) break; // no more chunks
 670                         if (iter->i >= 0) assert(iter->curr_off == iter->off[iter->i].v); // otherwise bug
 671                         if (iter->i < 0 || iter->off[iter->i].v != iter->off[iter->i+1].u) { // not adjacent chunks; then seek
 672                                 bam_seek(fp, iter->off[iter->i+1].u, SEEK_SET);
 673                                 iter->curr_off = bam_tell(fp);
 674                         }
 675                         ++iter->i;
 676                 }
 677                 if ((ret = bam_read1(fp, b)) > 0) {
 678                         iter->curr_off = bam_tell(fp);
 679                         if (b->core.tid != iter->tid || b->core.pos >= iter->end) break; // no need to proceed
 680                         else if (is_overlap(iter->beg, iter->end, b)) return ret;
 681                 } else break; // end of file
 682         }
 683         iter->finished = 1;
 684         return -1;
 685 }
 686
 687 int bam_fetch(bamFile fp, const bam_index_t *idx, int tid, int beg, int end, void *data, bam_fetch_f func)
 688 {
 689         bam_iter_t iter;
 690         bam1_t *b;
 691         b = bam_init1();
 692         iter = bam_iter_query(idx, tid, beg, end);
 693         while (bam_iter_read(fp, iter, b) >= 0) func(b, data);
 694         bam_destroy1(b);
 695         return 0;
 696 }