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