samtools/bam_index.c.pysam.c

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