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) { // change of chromosomes
 182                         last_tid = c->tid;
 183                         last_bin = 0xffffffffu;
 184                 } else if (last_coor > c->pos) {
 185                         fprintf(pysamerr, "[bam_index_core] the alignment is not sorted (%s): %u > %u in %d-th chr\n",
 186                                         bam1_qname(b), last_coor, c->pos, c->tid+1);
 187                         exit(1);
 188                 }
 189                 if (c->tid >= 0) insert_offset2(&idx->index2[b->core.tid], b, last_off);
 190                 if (c->bin != last_bin) { // then possibly write the binning index
 191                         if (save_bin != 0xffffffffu) // save_bin==0xffffffffu only happens to the first record
 192                                 insert_offset(idx->index[save_tid], save_bin, save_off, last_off);
 193                         if (last_bin == 0xffffffffu && save_tid != 0xffffffffu) { // write the meta element
 194                                 off_end = last_off;
 195                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, off_end);
 196                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 197                                 n_mapped = n_unmapped = 0;
 198                                 off_beg = off_end;
 199                         }
 200                         save_off = last_off;
 201                         save_bin = last_bin = c->bin;
 202                         save_tid = c->tid;
 203                         if (save_tid < 0) break;
 204                 }
 205                 if (bam_tell(fp) <= last_off) {
 206                         fprintf(pysamerr, "[bam_index_core] bug in BGZF/RAZF: %llx < %llx\n",
 207                                         (unsigned long long)bam_tell(fp), (unsigned long long)last_off);
 208                         exit(1);
 209                 }
 210                 if (c->flag & BAM_FUNMAP) ++n_unmapped;
 211                 else ++n_mapped;
 212                 last_off = bam_tell(fp);
 213                 last_coor = b->core.pos;
 214         }
 215         if (save_tid >= 0) {
 216                 insert_offset(idx->index[save_tid], save_bin, save_off, bam_tell(fp));
 217                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, bam_tell(fp));
 218                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 219         }
 220         merge_chunks(idx);
 221         fill_missing(idx);
 222         if (ret >= 0) {
 223                 while ((ret = bam_read1(fp, b)) >= 0) {
 224                         ++n_no_coor;
 225                         if (c->tid >= 0 && n_no_coor) {
 226                                 fprintf(pysamerr, "[bam_index_core] the alignment is not sorted: reads without coordinates prior to reads with coordinates.\n");
 227                                 exit(1);
 228                         }
 229                 }
 230         }
 231         if (ret < -1) fprintf(pysamerr, "[bam_index_core] truncated file? Continue anyway. (%d)\n", ret);
 232         free(b->data); free(b);
 233         idx->n_no_coor = n_no_coor;
 234         return idx;
 235 }
 236
 237 void bam_index_destroy(bam_index_t *idx)
 238 {
 239         khint_t k;
 240         int i;
 241         if (idx == 0) return;
 242         for (i = 0; i < idx->n; ++i) {
 243                 khash_t(i) *index = idx->index[i];
 244                 bam_lidx_t *index2 = idx->index2 + i;
 245                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 246                         if (kh_exist(index, k))
 247                                 free(kh_value(index, k).list);
 248                 }
 249                 kh_destroy(i, index);
 250                 free(index2->offset);
 251         }
 252         free(idx->index); free(idx->index2);
 253         free(idx);
 254 }
 255
 256 void bam_index_save(const bam_index_t *idx, FILE *fp)
 257 {
 258         int32_t i, size;
 259         khint_t k;
 260         fwrite("BAI\1", 1, 4, fp);
 261         if (bam_is_be) {
 262                 uint32_t x = idx->n;
 263                 fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 264         } else fwrite(&idx->n, 4, 1, fp);
 265         for (i = 0; i < idx->n; ++i) {
 266                 khash_t(i) *index = idx->index[i];
 267                 bam_lidx_t *index2 = idx->index2 + i;
 268                 // write binning index
 269                 size = kh_size(index);
 270                 if (bam_is_be) { // big endian
 271                         uint32_t x = size;
 272                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 273                 } else fwrite(&size, 4, 1, fp);
 274                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 275                         if (kh_exist(index, k)) {
 276                                 bam_binlist_t *p = &kh_value(index, k);
 277                                 if (bam_is_be) { // big endian
 278                                         uint32_t x;
 279                                         x = kh_key(index, k); fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 280                                         x = p->n; fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 281                                         for (x = 0; (int)x < p->n; ++x) {
 282                                                 bam_swap_endian_8p(&p->list[x].u);
 283                                                 bam_swap_endian_8p(&p->list[x].v);
 284                                         }
 285                                         fwrite(p->list, 16, p->n, fp);
 286                                         for (x = 0; (int)x < p->n; ++x) {
 287                                                 bam_swap_endian_8p(&p->list[x].u);
 288                                                 bam_swap_endian_8p(&p->list[x].v);
 289                                         }
 290                                 } else {
 291                                         fwrite(&kh_key(index, k), 4, 1, fp);
 292                                         fwrite(&p->n, 4, 1, fp);
 293                                         fwrite(p->list, 16, p->n, fp);
 294                                 }
 295                         }
 296                 }
 297                 // write linear index (index2)
 298                 if (bam_is_be) {
 299                         int x = index2->n;
 300                         fwrite(bam_swap_endian_4p(&x), 4, 1, fp);
 301                 } else fwrite(&index2->n, 4, 1, fp);
 302                 if (bam_is_be) { // big endian
 303                         int x;
 304                         for (x = 0; (int)x < index2->n; ++x)
 305                                 bam_swap_endian_8p(&index2->offset[x]);
 306                         fwrite(index2->offset, 8, index2->n, fp);
 307                         for (x = 0; (int)x < index2->n; ++x)
 308                                 bam_swap_endian_8p(&index2->offset[x]);
 309                 } else fwrite(index2->offset, 8, index2->n, fp);
 310         }
 311         { // write the number of reads coor-less records.
 312                 uint64_t x = idx->n_no_coor;
 313                 if (bam_is_be) bam_swap_endian_8p(&x);
 314                 fwrite(&x, 8, 1, fp);
 315         }
 316         fflush(fp);
 317 }
 318
 319 static bam_index_t *bam_index_load_core(FILE *fp)
 320 {
 321         int i;
 322         char magic[4];
 323         bam_index_t *idx;
 324         if (fp == 0) {
 325                 fprintf(pysamerr, "[bam_index_load_core] fail to load index.\n");
 326                 return 0;
 327         }
 328         fread(magic, 1, 4, fp);
 329         if (strncmp(magic, "BAI\1", 4)) {
 330                 fprintf(pysamerr, "[bam_index_load] wrong magic number.\n");
 331                 fclose(fp);
 332                 return 0;
 333         }
 334         idx = (bam_index_t*)calloc(1, sizeof(bam_index_t));
 335         fread(&idx->n, 4, 1, fp);
 336         if (bam_is_be) bam_swap_endian_4p(&idx->n);
 337         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 338         idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t));
 339         for (i = 0; i < idx->n; ++i) {
 340                 khash_t(i) *index;
 341                 bam_lidx_t *index2 = idx->index2 + i;
 342                 uint32_t key, size;
 343                 khint_t k;
 344                 int j, ret;
 345                 bam_binlist_t *p;
 346                 index = idx->index[i] = kh_init(i);
 347                 // load binning index
 348                 fread(&size, 4, 1, fp);
 349                 if (bam_is_be) bam_swap_endian_4p(&size);
 350                 for (j = 0; j < (int)size; ++j) {
 351                         fread(&key, 4, 1, fp);
 352                         if (bam_is_be) bam_swap_endian_4p(&key);
 353                         k = kh_put(i, index, key, &ret);
 354                         p = &kh_value(index, k);
 355                         fread(&p->n, 4, 1, fp);
 356                         if (bam_is_be) bam_swap_endian_4p(&p->n);
 357                         p->m = p->n;
 358                         p->list = (pair64_t*)malloc(p->m * 16);
 359                         fread(p->list, 16, p->n, fp);
 360                         if (bam_is_be) {
 361                                 int x;
 362                                 for (x = 0; x < p->n; ++x) {
 363                                         bam_swap_endian_8p(&p->list[x].u);
 364                                         bam_swap_endian_8p(&p->list[x].v);
 365                                 }
 366                         }
 367                 }
 368                 // load linear index
 369                 fread(&index2->n, 4, 1, fp);
 370                 if (bam_is_be) bam_swap_endian_4p(&index2->n);
 371                 index2->m = index2->n;
 372                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 373                 fread(index2->offset, index2->n, 8, fp);
 374                 if (bam_is_be)
 375                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 376         }
 377         if (fread(&idx->n_no_coor, 8, 1, fp) == 0) idx->n_no_coor = 0;
 378         if (bam_is_be) bam_swap_endian_8p(&idx->n_no_coor);
 379         return idx;
 380 }
 381
 382 bam_index_t *bam_index_load_local(const char *_fn)
 383 {
 384         FILE *fp;
 385         char *fnidx, *fn;
 386
 387         if (strstr(_fn, "ftp://") == _fn || strstr(_fn, "http://") == _fn) {
 388                 const char *p;
 389                 int l = strlen(_fn);
 390                 for (p = _fn + l - 1; p >= _fn; --p)
 391                         if (*p == '/') break;
 392                 fn = strdup(p + 1);
 393         } else fn = strdup(_fn);
 394         fnidx = (char*)calloc(strlen(fn) + 5, 1);
 395         strcpy(fnidx, fn); strcat(fnidx, ".bai");
 396         fp = fopen(fnidx, "rb");
 397         if (fp == 0) { // try "{base}.bai"
 398                 char *s = strstr(fn, "bam");
 399                 if (s == fn + strlen(fn) - 3) {
 400                         strcpy(fnidx, fn);
 401                         fnidx[strlen(fn)-1] = 'i';
 402                         fp = fopen(fnidx, "rb");
 403                 }
 404         }
 405         free(fnidx); free(fn);
 406         if (fp) {
 407                 bam_index_t *idx = bam_index_load_core(fp);
 408                 fclose(fp);
 409                 return idx;
 410         } else return 0;
 411 }
 412
 413 #ifdef _USE_KNETFILE
 414 static void download_from_remote(const char *url)
 415 {
 416         const int buf_size = 1 * 1024 * 1024;
 417         char *fn;
 418         FILE *fp;
 419         uint8_t *buf;
 420         knetFile *fp_remote;
 421         int l;
 422         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 423         l = strlen(url);
 424         for (fn = (char*)url + l - 1; fn >= url; --fn)
 425                 if (*fn == '/') break;
 426         ++fn; // fn now points to the file name
 427         fp_remote = knet_open(url, "r");
 428         if (fp_remote == 0) {
 429                 fprintf(pysamerr, "[download_from_remote] fail to open remote file.\n");
 430                 return;
 431         }
 432         if ((fp = fopen(fn, "wb")) == 0) {
 433                 fprintf(pysamerr, "[download_from_remote] fail to create file in the working directory.\n");
 434                 knet_close(fp_remote);
 435                 return;
 436         }
 437         buf = (uint8_t*)calloc(buf_size, 1);
 438         while ((l = knet_read(fp_remote, buf, buf_size)) != 0)
 439                 fwrite(buf, 1, l, fp);
 440         free(buf);
 441         fclose(fp);
 442         knet_close(fp_remote);
 443 }
 444 #else
 445 static void download_from_remote(const char *url)
 446 {
 447         return;
 448 }
 449 #endif
 450
 451 bam_index_t *bam_index_load(const char *fn)
 452 {
 453         bam_index_t *idx;
 454         idx = bam_index_load_local(fn);
 455         if (idx == 0 && (strstr(fn, "ftp://") == fn || strstr(fn, "http://") == fn)) {
 456                 char *fnidx = calloc(strlen(fn) + 5, 1);
 457                 strcat(strcpy(fnidx, fn), ".bai");
 458                 fprintf(pysamerr, "[bam_index_load] attempting to download the remote index file.\n");
 459                 download_from_remote(fnidx);
 460                 idx = bam_index_load_local(fn);
 461         }
 462         if (idx == 0) fprintf(pysamerr, "[bam_index_load] fail to load BAM index.\n");
 463         return idx;
 464 }
 465
 466 int bam_index_build2(const char *fn, const char *_fnidx)
 467 {
 468         char *fnidx;
 469         FILE *fpidx;
 470         bamFile fp;
 471         bam_index_t *idx;
 472         if ((fp = bam_open(fn, "r")) == 0) {
 473                 fprintf(pysamerr, "[bam_index_build2] fail to open the BAM file.\n");
 474                 return -1;
 475         }
 476         idx = bam_index_core(fp);
 477         bam_close(fp);
 478         if (_fnidx == 0) {
 479                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 480                 strcpy(fnidx, fn); strcat(fnidx, ".bai");
 481         } else fnidx = strdup(_fnidx);
 482         fpidx = fopen(fnidx, "wb");
 483         if (fpidx == 0) {
 484                 fprintf(pysamerr, "[bam_index_build2] fail to create the index file.\n");
 485                 free(fnidx);
 486                 return -1;
 487         }
 488         bam_index_save(idx, fpidx);
 489         bam_index_destroy(idx);
 490         fclose(fpidx);
 491         free(fnidx);
 492         return 0;
 493 }
 494
 495 int bam_index_build(const char *fn)
 496 {
 497         return bam_index_build2(fn, 0);
 498 }
 499
 500 int bam_index(int argc, char *argv[])
 501 {
 502         if (argc < 2) {
 503                 fprintf(pysamerr, "Usage: samtools index <in.bam> [out.index]\n");
 504                 return 1;
 505         }
 506         if (argc >= 3) bam_index_build2(argv[1], argv[2]);
 507         else bam_index_build(argv[1]);
 508         return 0;
 509 }
 510
 511 int bam_idxstats(int argc, char *argv[])
 512 {
 513         bam_index_t *idx;
 514         bam_header_t *header;
 515         bamFile fp;
 516         int i;
 517         if (argc < 2) {
 518                 fprintf(pysamerr, "Usage: samtools idxstats <in.bam>\n");
 519                 return 1;
 520         }
 521         fp = bam_open(argv[1], "r");
 522         if (fp == 0) { fprintf(pysamerr, "[%s] fail to open BAM.\n", __func__); return 1; }
 523         header = bam_header_read(fp);
 524         bam_close(fp);
 525         idx = bam_index_load(argv[1]);
 526         if (idx == 0) { fprintf(pysamerr, "[%s] fail to load the index.\n", __func__); return 1; }
 527         for (i = 0; i < idx->n; ++i) {
 528                 khint_t k;
 529                 khash_t(i) *h = idx->index[i];
 530                 printf("%s\t%d", header->target_name[i], header->target_len[i]);
 531                 k = kh_get(i, h, BAM_MAX_BIN);
 532                 if (k != kh_end(h))
 533                         printf("\t%llu\t%llu", (long long)kh_val(h, k).list[1].u, (long long)kh_val(h, k).list[1].v);
 534                 else printf("\t0\t0");
 535                 putchar('\n');
 536         }
 537         printf("*\t0\t0\t%llu\n", (long long)idx->n_no_coor);
 538         bam_header_destroy(header);
 539         bam_index_destroy(idx);
 540         return 0;
 541 }
 542
 543 static inline int reg2bins(uint32_t beg, uint32_t end, uint16_t list[BAM_MAX_BIN])
 544 {
 545         int i = 0, k;
 546         if (beg >= end) return 0;
 547         if (end >= 1u<<29) end = 1u<<29;
 548         --end;
 549         list[i++] = 0;
 550         for (k =    1 + (beg>>26); k <=    1 + (end>>26); ++k) list[i++] = k;
 551         for (k =    9 + (beg>>23); k <=    9 + (end>>23); ++k) list[i++] = k;
 552         for (k =   73 + (beg>>20); k <=   73 + (end>>20); ++k) list[i++] = k;
 553         for (k =  585 + (beg>>17); k <=  585 + (end>>17); ++k) list[i++] = k;
 554         for (k = 4681 + (beg>>14); k <= 4681 + (end>>14); ++k) list[i++] = k;
 555         return i;
 556 }
 557
 558 static inline int is_overlap(uint32_t beg, uint32_t end, const bam1_t *b)
 559 {
 560         uint32_t rbeg = b->core.pos;
 561         uint32_t rend = b->core.n_cigar? bam_calend(&b->core, bam1_cigar(b)) : b->core.pos + 1;
 562         return (rend > beg && rbeg < end);
 563 }
 564
 565 struct __bam_iter_t {
 566         int from_first; // read from the first record; no random access
 567         int tid, beg, end, n_off, i, finished;
 568         uint64_t curr_off;
 569         pair64_t *off;
 570 };
 571
 572 // bam_fetch helper function retrieves
 573 bam_iter_t bam_iter_query(const bam_index_t *idx, int tid, int beg, int end)
 574 {
 575         uint16_t *bins;
 576         int i, n_bins, n_off;
 577         pair64_t *off;
 578         khint_t k;
 579         khash_t(i) *index;
 580         uint64_t min_off;
 581         bam_iter_t iter = 0;
 582
 583         if (beg < 0) beg = 0;
 584         if (end < beg) return 0;
 585         // initialize iter
 586         iter = calloc(1, sizeof(struct __bam_iter_t));
 587         iter->tid = tid, iter->beg = beg, iter->end = end; iter->i = -1;
 588         //
 589         bins = (uint16_t*)calloc(BAM_MAX_BIN, 2);
 590         n_bins = reg2bins(beg, end, bins);
 591         index = idx->index[tid];
 592         if (idx->index2[tid].n > 0) {
 593                 min_off = (beg>>BAM_LIDX_SHIFT >= idx->index2[tid].n)? idx->index2[tid].offset[idx->index2[tid].n-1]
 594                         : idx->index2[tid].offset[beg>>BAM_LIDX_SHIFT];
 595                 if (min_off == 0) { // improvement for index files built by tabix prior to 0.1.4
 596                         int n = beg>>BAM_LIDX_SHIFT;
 597                         if (n > idx->index2[tid].n) n = idx->index2[tid].n;
 598                         for (i = n - 1; i >= 0; --i)
 599                                 if (idx->index2[tid].offset[i] != 0) break;
 600                         if (i >= 0) min_off = idx->index2[tid].offset[i];
 601                 }
 602         } else min_off = 0; // tabix 0.1.2 may produce such index files
 603         for (i = n_off = 0; i < n_bins; ++i) {
 604                 if ((k = kh_get(i, index, bins[i])) != kh_end(index))
 605                         n_off += kh_value(index, k).n;
 606         }
 607         if (n_off == 0) {
 608                 free(bins); return iter;
 609         }
 610         off = (pair64_t*)calloc(n_off, 16);
 611         for (i = n_off = 0; i < n_bins; ++i) {
 612                 if ((k = kh_get(i, index, bins[i])) != kh_end(index)) {
 613                         int j;
 614                         bam_binlist_t *p = &kh_value(index, k);
 615                         for (j = 0; j < p->n; ++j)
 616                                 if (p->list[j].v > min_off) off[n_off++] = p->list[j];
 617                 }
 618         }
 619         free(bins);
 620         if (n_off == 0) {
 621                 free(off); return iter;
 622         }
 623         {
 624                 bam1_t *b = (bam1_t*)calloc(1, sizeof(bam1_t));
 625                 int l;
 626                 ks_introsort(off, n_off, off);
 627                 // resolve completely contained adjacent blocks
 628                 for (i = 1, l = 0; i < n_off; ++i)
 629                         if (off[l].v < off[i].v)
 630                                 off[++l] = off[i];
 631                 n_off = l + 1;
 632                 // resolve overlaps between adjacent blocks; this may happen due to the merge in indexing
 633                 for (i = 1; i < n_off; ++i)
 634                         if (off[i-1].v >= off[i].u) off[i-1].v = off[i].u;
 635                 { // merge adjacent blocks
 636 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 637                         for (i = 1, l = 0; i < n_off; ++i) {
 638 #ifdef BAM_TRUE_OFFSET
 639                                 if (off[l].v + BAM_MIN_CHUNK_GAP > off[i].u) off[l].v = off[i].v;
 640 #else
 641                                 if (off[l].v>>16 == off[i].u>>16) off[l].v = off[i].v;
 642 #endif
 643                                 else off[++l] = off[i];
 644                         }
 645                         n_off = l + 1;
 646 #endif
 647                 }
 648                 bam_destroy1(b);
 649         }
 650         iter->n_off = n_off; iter->off = off;
 651         return iter;
 652 }
 653
 654 pair64_t *get_chunk_coordinates(const bam_index_t *idx, int tid, int beg, int end, int *cnt_off)
 655 { // for pysam compatibility
 656         bam_iter_t iter;
 657         pair64_t *off;
 658         iter = bam_iter_query(idx, tid, beg, end);
 659         off = iter->off; *cnt_off = iter->n_off;
 660         free(iter);
 661         return off;
 662 }
 663
 664 void bam_iter_destroy(bam_iter_t iter)
 665 {
 666         if (iter) { free(iter->off); free(iter); }
 667 }
 668
 669 int bam_iter_read(bamFile fp, bam_iter_t iter, bam1_t *b)
 670 {
 671         int ret;
 672         if (iter && iter->finished) return -1;
 673         if (iter == 0 || iter->from_first) {
 674                 ret = bam_read1(fp, b);
 675                 if (ret < 0 && iter) iter->finished = 1;
 676                 return ret;
 677         }
 678         if (iter->off == 0) return -1;
 679         for (;;) {
 680                 if (iter->curr_off == 0 || iter->curr_off >= iter->off[iter->i].v) { // then jump to the next chunk
 681                         if (iter->i == iter->n_off - 1) { ret = -1; break; } // no more chunks
 682                         if (iter->i >= 0) assert(iter->curr_off == iter->off[iter->i].v); // otherwise bug
 683                         if (iter->i < 0 || iter->off[iter->i].v != iter->off[iter->i+1].u) { // not adjacent chunks; then seek
 684                                 bam_seek(fp, iter->off[iter->i+1].u, SEEK_SET);
 685                                 iter->curr_off = bam_tell(fp);
 686                         }
 687                         ++iter->i;
 688                 }
 689                 if ((ret = bam_read1(fp, b)) >= 0) {
 690                         iter->curr_off = bam_tell(fp);
 691                         if (b->core.tid != iter->tid || b->core.pos >= iter->end) { // no need to proceed
 692                                 ret = bam_validate1(NULL, b)? -1 : -5; // determine whether end of region or error
 693                                 break;
 694                         }
 695                         else if (is_overlap(iter->beg, iter->end, b)) return ret;
 696                 } else break; // end of file or error
 697         }
 698         iter->finished = 1;
 699         return ret;
 700 }
 701
 702 int bam_fetch(bamFile fp, const bam_index_t *idx, int tid, int beg, int end, void *data, bam_fetch_f func)
 703 {
 704         int ret;
 705         bam_iter_t iter;
 706         bam1_t *b;
 707         b = bam_init1();
 708         iter = bam_iter_query(idx, tid, beg, end);
 709         while ((ret = bam_iter_read(fp, iter, b)) >= 0) func(b, data);
 710         bam_iter_destroy(iter);
 711         bam_destroy1(b);
 712         return (ret == -1)? 0 : ret;
 713 }