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                 if ((l->list = (pair64_t*)calloc(l->m, 16)) == NULL) {
  84                         fprintf(stderr, "[%s] failed to allocate bam_binlist node.\n", __func__);
  85                         abort();
  86                 }
  87         }
  88         if (l->n == l->m) {
  89                 l->m <<= 1;
  90                 if ((l->list = (pair64_t*)realloc(l->list, l->m * 16)) == NULL) {
  91                         fprintf(stderr, "[%s] failed to resize bam_binlist node.\n", __func__);
  92                         abort();
  93                 }
  94         }
  95         l->list[l->n].u = beg; l->list[l->n++].v = end;
  96 }
  97
  98 static inline void insert_offset2(bam_lidx_t *index2, bam1_t *b, uint64_t offset)
  99 {
 100         int i, beg, end;
 101         beg = b->core.pos >> BAM_LIDX_SHIFT;
 102         end = (bam_calend(&b->core, bam1_cigar(b)) - 1) >> BAM_LIDX_SHIFT;
 103         if (index2->m < end + 1) {
 104                 int old_m = index2->m;
 105                 index2->m = end + 1;
 106                 kroundup32(index2->m);
 107                 if ((index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8)) == NULL) {
 108                         fprintf(stderr, "[%s] failed to resize offset array.\n", __func__);
 109                         abort();
 110                 }
 111                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
 112         }
 113         if (beg == end) {
 114                 if (index2->offset[beg] == 0) index2->offset[beg] = offset;
 115         } else {
 116                 for (i = beg; i <= end; ++i)
 117                         if (index2->offset[i] == 0) index2->offset[i] = offset;
 118         }
 119         index2->n = end + 1;
 120 }
 121
 122 static void merge_chunks(bam_index_t *idx)
 123 {
 124 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 125         khash_t(i) *index;
 126         int i, l, m;
 127         khint_t k;
 128         for (i = 0; i < idx->n; ++i) {
 129                 index = idx->index[i];
 130                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 131                         bam_binlist_t *p;
 132                         if (!kh_exist(index, k) || kh_key(index, k) == BAM_MAX_BIN) continue;
 133                         p = &kh_value(index, k);
 134                         m = 0;
 135                         for (l = 1; l < p->n; ++l) {
 136 #ifdef BAM_TRUE_OFFSET
 137                                 if (p->list[m].v + BAM_MIN_CHUNK_GAP > p->list[l].u) p->list[m].v = p->list[l].v;
 138 #else
 139                                 if (p->list[m].v>>16 == p->list[l].u>>16) p->list[m].v = p->list[l].v;
 140 #endif
 141                                 else p->list[++m] = p->list[l];
 142                         } // ~for(l)
 143                         p->n = m + 1;
 144                 } // ~for(k)
 145         } // ~for(i)
 146 #endif // defined(BAM_TRUE_OFFSET) || defined(BAM_BGZF)
 147 }
 148
 149 static void fill_missing(bam_index_t *idx)
 150 {
 151         int i, j;
 152         for (i = 0; i < idx->n; ++i) {
 153                 bam_lidx_t *idx2 = &idx->index2[i];
 154                 for (j = 1; j < idx2->n; ++j)
 155                         if (idx2->offset[j] == 0)
 156                                 idx2->offset[j] = idx2->offset[j-1];
 157         }
 158 }
 159
 160 bam_index_t *bam_index_core(bamFile fp)
 161 {
 162         bam1_t *b;
 163         bam_header_t *h;
 164         int i, ret;
 165         bam_index_t *idx;
 166         uint32_t last_bin, save_bin;
 167         int32_t last_coor, last_tid, save_tid;
 168         bam1_core_t *c;
 169         uint64_t save_off, last_off, n_mapped, n_unmapped, off_beg, off_end, n_no_coor;
 170
 171         if ((idx = (bam_index_t*)calloc(1, sizeof(bam_index_t))) == NULL) {
 172                 fprintf(stderr, "[%s] failed to allocate bam_index structure.\n", __func__);
 173                 return NULL;
 174         }
 175         if ((b = (bam1_t*)calloc(1, sizeof(bam1_t))) == NULL) {
 176                 fprintf(stderr, "[%s] failed to allocate bam1_t buffer.\n", __func__);
 177                 free(idx);
 178                 return NULL;
 179         }
 180         h = bam_header_read(fp);
 181         c = &b->core;
 182
 183         idx->n = h->n_targets;
 184         bam_header_destroy(h);
 185         if ((idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*))) == NULL) {
 186                 fprintf(stderr, "[%s] failed to allocate bam_index bin hash.\n", __func__);
 187                 free(b);
 188                 free(idx);
 189                 return NULL;
 190         }
 191         for (i = 0; i < idx->n; ++i) idx->index[i] = kh_init(i);
 192         if ((idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t))) == NULL) {
 193                 fprintf(stderr, "[%s] failed to allocate bam_index linear index.\n", __func__);
 194                 free(idx->index);
 195                 free(b);
 196                 free(idx);
 197                 return NULL;
 198         }
 199
 200         save_bin = save_tid = last_tid = last_bin = 0xffffffffu;
 201         save_off = last_off = bam_tell(fp); last_coor = 0xffffffffu;
 202         n_mapped = n_unmapped = n_no_coor = off_end = 0;
 203         off_beg = off_end = bam_tell(fp);
 204         while ((ret = bam_read1(fp, b)) >= 0) {
 205                 if (c->tid < 0) ++n_no_coor;
 206                 if (last_tid < c->tid || (last_tid >= 0 && c->tid < 0)) { // change of chromosomes
 207                         last_tid = c->tid;
 208                         last_bin = 0xffffffffu;
 209                 } else if ((uint32_t)last_tid > (uint32_t)c->tid) {
 210                         fprintf(stderr, "[bam_index_core] the alignment is not sorted (%s): %d-th chr > %d-th chr\n",
 211                                         bam1_qname(b), last_tid+1, c->tid+1);
 212                         return NULL;
 213                 } else if ((int32_t)c->tid >= 0 && last_coor > c->pos) {
 214                         fprintf(stderr, "[bam_index_core] the alignment is not sorted (%s): %u > %u in %d-th chr\n",
 215                                         bam1_qname(b), last_coor, c->pos, c->tid+1);
 216                         return NULL;
 217                 }
 218                 if (c->tid >= 0 && !(c->flag & BAM_FUNMAP)) insert_offset2(&idx->index2[b->core.tid], b, last_off);
 219                 if (c->bin != last_bin) { // then possibly write the binning index
 220                         if (save_bin != 0xffffffffu) // save_bin==0xffffffffu only happens to the first record
 221                                 insert_offset(idx->index[save_tid], save_bin, save_off, last_off);
 222                         if (last_bin == 0xffffffffu && save_tid != 0xffffffffu) { // write the meta element
 223                                 off_end = last_off;
 224                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, off_end);
 225                                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 226                                 n_mapped = n_unmapped = 0;
 227                                 off_beg = off_end;
 228                         }
 229                         save_off = last_off;
 230                         save_bin = last_bin = c->bin;
 231                         save_tid = c->tid;
 232                         if (save_tid < 0) break;
 233                 }
 234                 if (bam_tell(fp) <= last_off) {
 235                         fprintf(stderr, "[bam_index_core] bug in BGZF/RAZF: %llx < %llx\n",
 236                                         (unsigned long long)bam_tell(fp), (unsigned long long)last_off);
 237                         return NULL;
 238                 }
 239                 if (c->flag & BAM_FUNMAP) ++n_unmapped;
 240                 else ++n_mapped;
 241                 last_off = bam_tell(fp);
 242                 last_coor = b->core.pos;
 243         }
 244         if (save_tid >= 0) {
 245                 insert_offset(idx->index[save_tid], save_bin, save_off, bam_tell(fp));
 246                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, off_beg, bam_tell(fp));
 247                 insert_offset(idx->index[save_tid], BAM_MAX_BIN, n_mapped, n_unmapped);
 248         }
 249         merge_chunks(idx);
 250         fill_missing(idx);
 251         if (ret >= 0) {
 252                 while ((ret = bam_read1(fp, b)) >= 0) {
 253                         ++n_no_coor;
 254                         if (c->tid >= 0 && n_no_coor) {
 255                                 fprintf(stderr, "[bam_index_core] the alignment is not sorted: reads without coordinates prior to reads with coordinates.\n");
 256                                 return NULL;
 257                         }
 258                 }
 259         }
 260         if (ret < -1) fprintf(stderr, "[bam_index_core] truncated file? Continue anyway. (%d)\n", ret);
 261         free(b->data); free(b);
 262         idx->n_no_coor = n_no_coor;
 263         return idx;
 264 }
 265
 266 void bam_index_destroy(bam_index_t *idx)
 267 {
 268         khint_t k;
 269         int i;
 270         if (idx == 0) return;
 271         for (i = 0; i < idx->n; ++i) {
 272                 khash_t(i) *index = idx->index[i];
 273                 bam_lidx_t *index2 = idx->index2 + i;
 274                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 275                         if (kh_exist(index, k))
 276                                 free(kh_value(index, k).list);
 277                 }
 278                 kh_destroy(i, index);
 279                 free(index2->offset);
 280         }
 281         free(idx->index); free(idx->index2);
 282         free(idx);
 283 }
 284
 285 void bam_index_save(const bam_index_t *idx, FILE *fp)
 286 {
 287         int32_t i, size;
 288         khint_t k;
 289         if (fwrite("BAI\1", 1, 4, fp) < 4) {
 290                 fprintf(stderr, "[%s] failed to write magic number.\n", __func__);
 291                 return;
 292         }
 293         if (bam_is_be) {
 294                 uint32_t x = idx->n;
 295                 if (fwrite(bam_swap_endian_4p(&x), 4, 1, fp) < 1) {
 296                         fprintf(stderr, "[%s] failed to write n_ref.\n", __func__);
 297                         return;
 298                 }
 299         }
 300         else {
 301                 if (fwrite(&idx->n, 4, 1, fp) < 1) {
 302                         fprintf(stderr, "[%s] failed to write n_ref.\n", __func__);
 303                         return;
 304                 }
 305         }
 306         for (i = 0; i < idx->n; ++i) {
 307                 khash_t(i) *index = idx->index[i];
 308                 bam_lidx_t *index2 = idx->index2 + i;
 309                 // write binning index
 310                 size = kh_size(index);
 311                 if (bam_is_be) { // big endian
 312                         uint32_t x = size;
 313                         if (fwrite(bam_swap_endian_4p(&x), 4, 1, fp) < 1) {
 314                                 fprintf(stderr, "[%s] failed to write n_bin.\n", __func__);
 315                                 return;
 316                         }
 317                 }
 318                 else {
 319                         if (fwrite(&size, 4, 1, fp) < 1) {
 320                                 fprintf(stderr, "[%s] failed to write n_bin.\n", __func__);
 321                                 return;
 322                         }
 323                 }
 324                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 325                         if (kh_exist(index, k)) {
 326                                 bam_binlist_t *p = &kh_value(index, k);
 327                                 if (bam_is_be) { // big endian
 328                                         uint32_t x;
 329                                         x = kh_key(index, k);
 330                                         if (fwrite(bam_swap_endian_4p(&x), 4, 1, fp) < 1) {
 331                                                 fprintf(stderr, "[%s] failed to write bin.\n", __func__);
 332                                                 return;
 333                                         }
 334                                         x = p->n;
 335                                         if (fwrite(bam_swap_endian_4p(&x), 4, 1, fp) < 1) {
 336                                                 fprintf(stderr, "[%s] failed to write n_chunk.\n", __func__);
 337                                                 return;
 338                                         }
 339                                         for (x = 0; (int)x < p->n; ++x) {
 340                                                 bam_swap_endian_8p(&p->list[x].u);
 341                                                 bam_swap_endian_8p(&p->list[x].v);
 342                                         }
 343                                         if (fwrite(p->list, 16, p->n, fp) < p->n) {
 344                                                 fprintf(stderr, "[%s] failed to write %d chunks.\n", __func__, p->n);
 345                                                 return;
 346                                         }
 347                                         /*
 348                                         for (x = 0; (int)x < p->n; ++x) {
 349                                                 bam_swap_endian_8p(&p->list[x].u);
 350                                                 bam_swap_endian_8p(&p->list[x].v);
 351                                         } */
 352                                 }
 353                                 else {
 354                                         if (fwrite(&kh_key(index, k), 4, 1, fp) < 1) {
 355                                                 fprintf(stderr, "[%s] failed to write bin.\n", __func__);
 356                                                 return;
 357                                         }
 358                                         if (fwrite(&p->n, 4, 1, fp) < 1) {
 359                                                 fprintf(stderr, "[%s] failed to write n_chunk.\n", __func__);
 360                                                 return;
 361                                         }
 362                                         if (fwrite(p->list, 16, p->n, fp) < p->n) {
 363                                                 fprintf(stderr, "[%s] failed to write %d chunks.\n", __func__, p->n);
 364                                                 return;
 365                                         }
 366                                 }
 367                         }
 368                 }
 369                 // write linear index (index2)
 370                 if (bam_is_be) {
 371                         int x = index2->n;
 372                         if (fwrite(bam_swap_endian_4p(&x), 4, 1, fp) < 1) {
 373                                 fprintf(stderr, "[%s] failed to write n_intv.\n", __func__);
 374                                 return;
 375                         }
 376                 }
 377                 else {
 378                         if (fwrite(&index2->n, 4, 1, fp) < 1) {
 379                                 fprintf(stderr, "[%s] failed to write n_intv.\n", __func__);
 380                                 return;
 381                         }
 382                 }
 383                 if (bam_is_be) { // big endian
 384                         int x;
 385                         for (x = 0; (int)x < index2->n; ++x)
 386                                 bam_swap_endian_8p(&index2->offset[x]);
 387                         if (fwrite(index2->offset, 8, index2->n, fp) < index2->n) {
 388                                 fprintf(stderr, "[%s] failed to write ioffset.\n", __func__);
 389                                 return;
 390                         }
 391                         for (x = 0; (int)x < index2->n; ++x)
 392                                 bam_swap_endian_8p(&index2->offset[x]);
 393                 }
 394                 else {
 395                         if (fwrite(index2->offset, 8, index2->n, fp) < index2->n) {
 396                                 fprintf(stderr, "[%s] failed to write ioffset.\n", __func__);
 397                                 return;
 398                         }
 399                 }
 400         }
 401         { // write the number of reads coor-less records.
 402                 uint64_t x = idx->n_no_coor;
 403                 if (bam_is_be) bam_swap_endian_8p(&x);
 404                 if (fwrite(&x, 8, 1, fp) < 1) {
 405                         fprintf(stderr, "[%s] failed to write n_no_coor.\n", __func__);
 406                 }
 407         }
 408         fflush(fp);
 409 }
 410
 411 static bam_index_t *bam_index_load_core(FILE *fp)
 412 {
 413         int i;
 414         char magic[4];
 415         bam_index_t *idx;
 416         if (fp == 0) {
 417                 fprintf(stderr, "[bam_index_load_core] fail to load index.\n");
 418                 return 0;
 419         }
 420         if (fread(magic, 1, 4, fp) < 4) {
 421                 fprintf(stderr, "[%s] failed to read magic number.\n", __func__);
 422                 fclose(fp);
 423                 return 0;
 424         }
 425         if (strncmp(magic, "BAI\1", 4)) {
 426                 fprintf(stderr, "[bam_index_load] wrong magic number.\n");
 427                 fclose(fp);
 428                 return 0;
 429         }
 430         if ((idx = (bam_index_t*)calloc(1, sizeof(bam_index_t))) == NULL) {
 431                 fprintf(stderr, "[%s] failed to allocate bam_index structure.\n", __func__);
 432                 fclose(fp);
 433                 return NULL;
 434         }
 435         if (fread(&idx->n, 4, 1, fp) < 1) {
 436                 fprintf(stderr, "[%s] failed to read n_ref.\n", __func__);
 437                 free(idx);
 438                 fclose(fp);
 439                 return 0;
 440         }
 441         if (bam_is_be) bam_swap_endian_4p(&idx->n);
 442         if ((idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*))) == NULL) {
 443                 fprintf(stderr, "[%s] failed to allocate bam_index bin hash.\n", __func__);
 444                 free(idx);
 445                 fclose(fp);
 446                 return NULL;
 447         }
 448         if ((idx->index2 = (bam_lidx_t*)calloc(idx->n, sizeof(bam_lidx_t))) == NULL) {
 449                 fprintf(stderr, "[%s] failed to allocate bam_index linear index.\n", __func__);
 450                 free(idx->index);
 451                 free(idx);
 452                 fclose(fp);
 453                 return NULL;
 454         }
 455
 456         for (i = 0; i < idx->n; ++i) {
 457                 khash_t(i) *index;
 458                 bam_lidx_t *index2 = idx->index2 + i;
 459                 uint32_t key, size;
 460                 khint_t k;
 461                 int j, ret;
 462                 bam_binlist_t *p;
 463                 index = idx->index[i] = kh_init(i);
 464                 // load binning index
 465                 if (fread(&size, 4, 1, fp) < 1) {
 466                         fprintf(stderr, "[%s] failed to read n_bin.\n", __func__);
 467                         free(idx->index);
 468                         free(idx);
 469                         fclose(fp);
 470                         return 0;
 471                 }
 472                 if (bam_is_be) bam_swap_endian_4p(&size);
 473                 for (j = 0; j < (int)size; ++j) {
 474                         if (fread(&key, 4, 1, fp) < 1) {
 475                                 fprintf(stderr, "[%s] failed to read bin.\n", __func__);
 476                                 free(idx->index);
 477                                 free(idx);
 478                                 fclose(fp);
 479                                 return 0;
 480                         }
 481                         if (bam_is_be) bam_swap_endian_4p(&key);
 482                         k = kh_put(i, index, key, &ret);
 483                         p = &kh_value(index, k);
 484                         if (fread(&p->n, 4, 1, fp) < 1) {
 485                                 fprintf(stderr, "[%s] failed to read n_chunk.\n", __func__);
 486                                 free(idx->index);
 487                                 free(idx);
 488                                 fclose(fp);
 489                                 return 0;
 490                         }
 491                         if (bam_is_be) bam_swap_endian_4p(&p->n);
 492                         p->m = p->n;
 493                         if ((p->list = (pair64_t*)malloc(p->m * 16)) == NULL) {
 494                                 fprintf(stderr, "[%s] failed to allocate chunk array.\n", __func__);
 495                                 free(idx->index);
 496                                 free(idx);
 497                                 fclose(fp);
 498                                 return 0;
 499                         }
 500                         if (fread(p->list, 16, p->n, fp) < p->n) {
 501                                 fprintf(stderr, "[%s] failed to read %d chunks.\n", __func__, p->n);
 502                                 free(idx->index);
 503                                 free(idx);
 504                                 fclose(fp);
 505                                 return 0;
 506                         }
 507                         if (bam_is_be) {
 508                                 int x;
 509                                 for (x = 0; x < p->n; ++x) {
 510                                         bam_swap_endian_8p(&p->list[x].u);
 511                                         bam_swap_endian_8p(&p->list[x].v);
 512                                 }
 513                         }
 514                 }
 515                 // load linear index
 516                 if (fread(&index2->n, 4, 1, fp) < 1) {
 517                         fprintf(stderr, "[%s] failed to read n_intv", __func__);
 518                         free(idx->index);
 519                         free(idx);
 520                         fclose(fp);
 521                         return 0;
 522                 }
 523                 if (bam_is_be) bam_swap_endian_4p(&index2->n);
 524                 if (index2->n > 0) {
 525                         index2->m = index2->n;
 526                         if ((index2->offset = (uint64_t*)calloc(index2->m, 8)) == NULL) {
 527                                 fprintf(stderr, "[%s] falied to allocate interval array.\n", __func__);
 528                                 free(idx->index);
 529                                 free(idx);
 530                                 fclose(fp);
 531                                 return 0;
 532                         }
 533                         if (fread(index2->offset, index2->n, 8, fp) < index2->n) {
 534                                 fprintf(stderr, "[%s] failed to read %d intervals.", __func__, index2->n);
 535                                 free(index2->offset);
 536                                 free(idx->index);
 537                                 free(idx);
 538                                 fclose(fp);
 539                                 return 0;
 540                         }
 541                         if (bam_is_be) {
 542                                 for (j = 0; j < index2->n; ++j) {
 543                                         bam_swap_endian_8p(&index2->offset[j]);
 544                                 }
 545                         }
 546                 }
 547         }
 548         if (fread(&idx->n_no_coor, 8, 1, fp) == 0) idx->n_no_coor = 0;
 549         if (bam_is_be) bam_swap_endian_8p(&idx->n_no_coor);
 550         return idx;
 551 }
 552
 553 bam_index_t *bam_index_load_local(const char *_fn)
 554 {
 555         FILE *fp;
 556         char *fnidx, *fn;
 557
 558         if (strstr(_fn, "ftp://") == _fn || strstr(_fn, "http://") == _fn) {
 559                 const char *p;
 560                 int l = strlen(_fn);
 561                 for (p = _fn + l - 1; p >= _fn; --p)
 562                         if (*p == '/') break;
 563                 fn = strdup(p + 1);
 564         } else fn = strdup(_fn);
 565         if ((fnidx = (char*)calloc(strlen(fn) + 5, 1)) == NULL) {
 566                 fprintf(stderr, "[%s] failed to allocate space for index file name.\n", __func__);
 567                 free(fn);
 568                 return NULL;
 569         }
 570         strcpy(fnidx, fn); strcat(fnidx, ".bai");
 571         fp = fopen(fnidx, "rb");
 572         if (fp == 0) { // try "{base}.bai"
 573                 char *s = strstr(fn, "bam");
 574                 if (s == fn + strlen(fn) - 3) {
 575                         strcpy(fnidx, fn);
 576                         fnidx[strlen(fn)-1] = 'i';
 577                         fp = fopen(fnidx, "rb");
 578                 }
 579         }
 580         free(fnidx); free(fn);
 581         if (fp) {
 582                 bam_index_t *idx = bam_index_load_core(fp);
 583                 fclose(fp);
 584                 return idx;
 585         }
 586         else {
 587                 fprintf(stderr, "[%s] failed to open index file.\n", __func__);
 588                 return 0;
 589         }
 590 }
 591
 592 #ifdef _USE_KNETFILE
 593 static void download_from_remote(const char *url)
 594 {
 595         const int buf_size = 1 * 1024 * 1024;
 596         char *fn;
 597         FILE *fp;
 598         uint8_t *buf;
 599         knetFile *fp_remote;
 600         int l;
 601         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 602         l = strlen(url);
 603         for (fn = (char*)url + l - 1; fn >= url; --fn)
 604                 if (*fn == '/') break;
 605         ++fn; // fn now points to the file name
 606         fp_remote = knet_open(url, "r");
 607         if (fp_remote == 0) {
 608                 fprintf(stderr, "[%s] failed to open remote file.\n", __func__);
 609                 return;
 610         }
 611         if ((fp = fopen(fn, "wb")) == 0) {
 612                 fprintf(stderr, "[%s] failed to create file in the working directory.\n", __func__);
 613                 knet_close(fp_remote);
 614                 return;
 615         }
 616         if ((buf = (uint8_t*)calloc(buf_size, 1)) == NULL) {
 617                 fprintf(stderr, "[%s] failed to allocate buffer.\n", __func__);
 618                 fclose(fp);
 619                 return;
 620         }
 621         while ((l = knet_read(fp_remote, buf, buf_size)) != 0) {
 622                 if (fwrite(buf, 1, l, fp) < l) {
 623                         fprintf(stderr, "[%s] failed to write to destination file.\n", __func__);
 624                         break;
 625                 }
 626         }
 627         free(buf);
 628         fclose(fp);
 629         knet_close(fp_remote);
 630 }
 631 #else
 632 static void download_from_remote(const char *url)
 633 {
 634         return;
 635 }
 636 #endif
 637
 638 bam_index_t *bam_index_load(const char *fn)
 639 {
 640         bam_index_t *idx;
 641         idx = bam_index_load_local(fn);
 642         if (idx == 0 && (strstr(fn, "ftp://") == fn || strstr(fn, "http://") == fn)) {
 643                 char *fnidx = calloc(strlen(fn) + 5, 1);
 644                 if (fnidx == NULL) {
 645                         fprintf(stderr, "[%s] failed to allocate space for index filename.\n", __func__);
 646                         return NULL;
 647                 }
 648                 strcat(strcpy(fnidx, fn), ".bai");
 649                 fprintf(stderr, "[%s] attempting to download the remote index file.\n", __func__);
 650                 download_from_remote(fnidx);
 651                 idx = bam_index_load_local(fn);
 652         }
 653         if (idx == 0) fprintf(stderr, "[%s] fail to load BAM index.\n", __func__);
 654         return idx;
 655 }
 656
 657 int bam_index_build2(const char *fn, const char *_fnidx)
 658 {
 659         char *fnidx;
 660         FILE *fpidx;
 661         bamFile fp;
 662         bam_index_t *idx;
 663         if ((fp = bam_open(fn, "r")) == 0) {
 664                 fprintf(stderr, "[%s] failure to open the BAM file.\n", __func__);
 665                 return -1;
 666         }
 667         idx = bam_index_core(fp);
 668         bam_close(fp);
 669         if(idx == 0) {
 670                 fprintf(stderr, "[%s] fail to index the BAM file.\n", __func__);
 671                 return -1;
 672         }
 673         if (_fnidx == 0) {
 674                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 675                 strcpy(fnidx, fn); strcat(fnidx, ".bai");
 676         } else fnidx = strdup(_fnidx);
 677         fpidx = fopen(fnidx, "wb");
 678         if (fpidx == 0) {
 679                 fprintf(stderr, "[%s] fail to create the index file.\n", __func__);
 680                 free(fnidx);
 681                 return -1;
 682         }
 683         bam_index_save(idx, fpidx);
 684         bam_index_destroy(idx);
 685         fclose(fpidx);
 686         free(fnidx);
 687         return 0;
 688 }
 689
 690 int bam_index_build(const char *fn)
 691 {
 692         return bam_index_build2(fn, 0);
 693 }
 694
 695 int bam_index(int argc, char *argv[])
 696 {
 697         if (argc < 2) {
 698                 fprintf(stderr, "Usage: samtools index <in.bam> [out.index]\n");
 699                 return 1;
 700         }
 701         if (argc >= 3) bam_index_build2(argv[1], argv[2]);
 702         else bam_index_build(argv[1]);
 703         return 0;
 704 }
 705
 706 int bam_idxstats(int argc, char *argv[])
 707 {
 708         bam_index_t *idx;
 709         bam_header_t *header;
 710         bamFile fp;
 711         int i;
 712         if (argc < 2) {
 713                 fprintf(stderr, "Usage: samtools idxstats <in.bam>\n");
 714                 return 1;
 715         }
 716         fp = bam_open(argv[1], "r");
 717         if (fp == 0) { fprintf(stderr, "[%s] fail to open BAM.\n", __func__); return 1; }
 718         header = bam_header_read(fp);
 719         bam_close(fp);
 720         idx = bam_index_load(argv[1]);
 721         if (idx == 0) { fprintf(stderr, "[%s] fail to load the index.\n", __func__); return 1; }
 722         for (i = 0; i < idx->n; ++i) {
 723                 khint_t k;
 724                 khash_t(i) *h = idx->index[i];
 725                 printf("%s\t%d", header->target_name[i], header->target_len[i]);
 726                 k = kh_get(i, h, BAM_MAX_BIN);
 727                 if (k != kh_end(h))
 728                         printf("\t%llu\t%llu", (long long)kh_val(h, k).list[1].u, (long long)kh_val(h, k).list[1].v);
 729                 else printf("\t0\t0");
 730                 putchar('\n');
 731         }
 732         printf("*\t0\t0\t%llu\n", (long long)idx->n_no_coor);
 733         bam_header_destroy(header);
 734         bam_index_destroy(idx);
 735         return 0;
 736 }
 737
 738 static inline int reg2bins(uint32_t beg, uint32_t end, uint16_t list[BAM_MAX_BIN])
 739 {
 740         int i = 0, k;
 741         if (beg >= end) return 0;
 742         if (end >= 1u<<29) end = 1u<<29;
 743         --end;
 744         list[i++] = 0;
 745         for (k =    1 + (beg>>26); k <=    1 + (end>>26); ++k) list[i++] = k;
 746         for (k =    9 + (beg>>23); k <=    9 + (end>>23); ++k) list[i++] = k;
 747         for (k =   73 + (beg>>20); k <=   73 + (end>>20); ++k) list[i++] = k;
 748         for (k =  585 + (beg>>17); k <=  585 + (end>>17); ++k) list[i++] = k;
 749         for (k = 4681 + (beg>>14); k <= 4681 + (end>>14); ++k) list[i++] = k;
 750         return i;
 751 }
 752
 753 static inline int is_overlap(uint32_t beg, uint32_t end, const bam1_t *b)
 754 {
 755         uint32_t rbeg = b->core.pos;
 756         uint32_t rend = b->core.n_cigar? bam_calend(&b->core, bam1_cigar(b)) : b->core.pos + 1;
 757         return (rend > beg && rbeg < end);
 758 }
 759
 760 struct __bam_iter_t {
 761         int from_first; // read from the first record; no random access
 762         int tid, beg, end, n_off, i, finished;
 763         uint64_t curr_off;
 764         pair64_t *off;
 765 };
 766
 767 // bam_fetch helper function retrieves
 768 bam_iter_t bam_iter_query(const bam_index_t *idx, int tid, int beg, int end)
 769 {
 770         uint16_t *bins;
 771         int i, n_bins, n_off;
 772         pair64_t *off;
 773         khint_t k;
 774         khash_t(i) *index;
 775         uint64_t min_off;
 776         bam_iter_t iter = 0;
 777
 778         if (beg < 0) beg = 0;
 779         if (end < beg) return 0;
 780         // initialize iter
 781         if ((iter = calloc(1, sizeof(struct __bam_iter_t))) == NULL) {
 782                 fprintf(stderr, "[%s] failed to allocate bam iterator.\n", __func__);
 783                 abort();
 784         }
 785         iter->tid = tid, iter->beg = beg, iter->end = end; iter->i = -1;
 786         //
 787         if ((bins = (uint16_t*)calloc(BAM_MAX_BIN, 2)) == NULL) {
 788                 fprintf(stderr, "[%s] failed to allocate bin array.\n", __func__);
 789                 abort();
 790         }
 791         n_bins = reg2bins(beg, end, bins);
 792         index = idx->index[tid];
 793         if (idx->index2[tid].n > 0) {
 794                 min_off = (beg>>BAM_LIDX_SHIFT >= idx->index2[tid].n)? idx->index2[tid].offset[idx->index2[tid].n-1]
 795                         : idx->index2[tid].offset[beg>>BAM_LIDX_SHIFT];
 796                 if (min_off == 0) { // improvement for index files built by tabix prior to 0.1.4
 797                         int n = beg >> BAM_LIDX_SHIFT;
 798                         if (n > idx->index2[tid].n) n = idx->index2[tid].n;
 799                         for (i = n - 1; i >= 0; --i)
 800                                 if (idx->index2[tid].offset[i] != 0) break;
 801                         if (i >= 0) min_off = idx->index2[tid].offset[i];
 802                 }
 803         } else min_off = 0; // tabix 0.1.2 may produce such index files
 804         for (i = n_off = 0; i < n_bins; ++i) {
 805                 if ((k = kh_get(i, index, bins[i])) != kh_end(index))
 806                         n_off += kh_value(index, k).n;
 807         }
 808         if (n_off == 0) {
 809                 free(bins); return iter;
 810         }
 811         if ((off = (pair64_t*)calloc(n_off, 16)) == NULL) {
 812                 fprintf(stderr, "[%s] failed to allocate offset pair array.\n", __func__);
 813                 abort();
 814         }
 815         for (i = n_off = 0; i < n_bins; ++i) {
 816                 if ((k = kh_get(i, index, bins[i])) != kh_end(index)) {
 817                         int j;
 818                         bam_binlist_t *p = &kh_value(index, k);
 819                         for (j = 0; j < p->n; ++j)
 820                                 if (p->list[j].v > min_off) off[n_off++] = p->list[j];
 821                 }
 822         }
 823         free(bins);
 824         if (n_off == 0) {
 825                 free(off); return iter;
 826         }
 827         {
 828                 bam1_t *b = (bam1_t*)calloc(1, sizeof(bam1_t));
 829                 if (b == NULL) {
 830                         fprintf(stderr, "[%s] failure to allocate bam1_t buffer.\n", __func__);
 831                         abort();
 832                 }
 833                 int l;
 834                 ks_introsort(off, n_off, off);
 835                 // resolve completely contained adjacent blocks
 836                 for (i = 1, l = 0; i < n_off; ++i)
 837                         if (off[l].v < off[i].v)
 838                                 off[++l] = off[i];
 839                 n_off = l + 1;
 840                 // resolve overlaps between adjacent blocks; this may happen due to the merge in indexing
 841                 for (i = 1; i < n_off; ++i)
 842                         if (off[i-1].v >= off[i].u) off[i-1].v = off[i].u;
 843                 { // merge adjacent blocks
 844 #if defined(BAM_TRUE_OFFSET) || defined(BAM_VIRTUAL_OFFSET16)
 845                         for (i = 1, l = 0; i < n_off; ++i) {
 846 #ifdef BAM_TRUE_OFFSET
 847                                 if (off[l].v + BAM_MIN_CHUNK_GAP > off[i].u) off[l].v = off[i].v;
 848 #else
 849                                 if (off[l].v>>16 == off[i].u>>16) off[l].v = off[i].v;
 850 #endif
 851                                 else off[++l] = off[i];
 852                         }
 853                         n_off = l + 1;
 854 #endif
 855                 }
 856                 bam_destroy1(b);
 857         }
 858         iter->n_off = n_off; iter->off = off;
 859         return iter;
 860 }
 861
 862 pair64_t *get_chunk_coordinates(const bam_index_t *idx, int tid, int beg, int end, int *cnt_off)
 863 { // for pysam compatibility
 864         bam_iter_t iter;
 865         pair64_t *off;
 866         iter = bam_iter_query(idx, tid, beg, end);
 867         off = iter->off; *cnt_off = iter->n_off;
 868         free(iter);
 869         return off;
 870 }
 871
 872 void bam_iter_destroy(bam_iter_t iter)
 873 {
 874         if (iter) { free(iter->off); free(iter); }
 875 }
 876
 877 int bam_iter_read(bamFile fp, bam_iter_t iter, bam1_t *b)
 878 {
 879         int ret;
 880         if (iter && iter->finished) return -1;
 881         if (iter == 0 || iter->from_first) {
 882                 ret = bam_read1(fp, b);
 883                 if (ret < 0 && iter) iter->finished = 1;
 884                 return ret;
 885         }
 886         if (iter->off == 0) return -1;
 887         for (;;) {
 888                 if (iter->curr_off == 0 || iter->curr_off >= iter->off[iter->i].v) { // then jump to the next chunk
 889                         if (iter->i == iter->n_off - 1) { ret = -1; break; } // no more chunks
 890                         if (iter->i >= 0) assert(iter->curr_off == iter->off[iter->i].v); // otherwise bug
 891                         if (iter->i < 0 || iter->off[iter->i].v != iter->off[iter->i+1].u) { // not adjacent chunks; then seek
 892                                 bam_seek(fp, iter->off[iter->i+1].u, SEEK_SET);
 893                                 iter->curr_off = bam_tell(fp);
 894                         }
 895                         ++iter->i;
 896                 }
 897                 if ((ret = bam_read1(fp, b)) >= 0) {
 898                         iter->curr_off = bam_tell(fp);
 899                         if (b->core.tid != iter->tid || b->core.pos >= iter->end) { // no need to proceed
 900                                 ret = bam_validate1(NULL, b)? -1 : -5; // determine whether end of region or error
 901                                 break;
 902                         }
 903                         else if (is_overlap(iter->beg, iter->end, b)) return ret;
 904                 } else break; // end of file or error
 905         }
 906         iter->finished = 1;
 907         return ret;
 908 }
 909
 910 int bam_fetch(bamFile fp, const bam_index_t *idx, int tid, int beg, int end, void *data, bam_fetch_f func)
 911 {
 912         int ret;
 913         bam_iter_t iter;
 914         bam1_t *b;
 915         b = bam_init1();
 916         iter = bam_iter_query(idx, tid, beg, end);
 917         while ((ret = bam_iter_read(fp, iter, b)) >= 0) func(b, data);
 918         bam_iter_destroy(iter);
 919         bam_destroy1(b);
 920         return (ret == -1)? 0 : ret;
 921 }