tabix/index.c.pysam.c

   1 #include "pysam.h"
   2
   3 #include <ctype.h>
   4 #include <assert.h>
   5 #include <sys/stat.h>
   6 #include "khash.h"
   7 #include "ksort.h"
   8 #include "kstring.h"
   9 #include "bam_endian.h"
  10 #ifdef _USE_KNETFILE
  11 #include "knetfile.h"
  12 #endif
  13 #include "tabix.h"
  14
  15 #define TAD_MIN_CHUNK_GAP 32768
  16 // 1<<14 is the size of minimum bin.
  17 #define TAD_LIDX_SHIFT    14
  18
  19 typedef struct {
  20         uint64_t u, v;
  21 } pair64_t;
  22
  23 #define pair64_lt(a,b) ((a).u < (b).u)
  24 KSORT_INIT(offt, pair64_t, pair64_lt)
  25
  26 typedef struct {
  27         uint32_t m, n;
  28         pair64_t *list;
  29 } ti_binlist_t;
  30
  31 typedef struct {
  32         int32_t n, m;
  33         uint64_t *offset;
  34 } ti_lidx_t;
  35
  36 KHASH_MAP_INIT_INT(i, ti_binlist_t)
  37 KHASH_MAP_INIT_STR(s, int)
  38
  39 struct __ti_index_t {
  40         ti_conf_t conf;
  41         int32_t n, max;
  42         khash_t(s) *tname;
  43         khash_t(i) **index;
  44         ti_lidx_t *index2;
  45 };
  46
  47 struct __ti_iter_t {
  48         int from_first; // read from the first record; no random access
  49         int tid, beg, end, n_off, i, finished;
  50         uint64_t curr_off;
  51         kstring_t str;
  52         const ti_index_t *idx;
  53         pair64_t *off;
  54 };
  55
  56 typedef struct {
  57         int tid, beg, end, bin;
  58 } ti_intv_t;
  59
  60 ti_conf_t ti_conf_gff = { 0, 1, 4, 5, '#', 0 };
  61 ti_conf_t ti_conf_bed = { TI_FLAG_UCSC, 1, 2, 3, '#', 0 };
  62 ti_conf_t ti_conf_psltbl = { TI_FLAG_UCSC, 15, 17, 18, '#', 0 };
  63 ti_conf_t ti_conf_sam = { TI_PRESET_SAM, 3, 4, 0, '@', 0 };
  64 ti_conf_t ti_conf_vcf = { TI_PRESET_VCF, 1, 2, 0, '#', 0 };
  65
  66 /***************
  67  * read a line *
  68  ***************/
  69
  70 /*
  71 int ti_readline(BGZF *fp, kstring_t *str)
  72 {
  73         int c, l = 0;
  74         str->l = 0;
  75         while ((c = bgzf_getc(fp)) >= 0 && c != '\n') {
  76                 ++l;
  77                 if (c != '\r') kputc(c, str);
  78         }
  79         if (c < 0 && l == 0) return -1; // end of file
  80         return str->l;
  81 }
  82 */
  83
  84 /* Below is a faster implementation largely equivalent to the one
  85  * commented out above. */
  86 int ti_readline(BGZF *fp, kstring_t *str)
  87 {
  88         return bgzf_getline(fp, '\n', str);
  89 }
  90
  91 /*************************************
  92  * get the interval from a data line *
  93  *************************************/
  94
  95 static inline int ti_reg2bin(uint32_t beg, uint32_t end)
  96 {
  97         --end;
  98         if (beg>>14 == end>>14) return 4681 + (beg>>14);
  99         if (beg>>17 == end>>17) return  585 + (beg>>17);
 100         if (beg>>20 == end>>20) return   73 + (beg>>20);
 101         if (beg>>23 == end>>23) return    9 + (beg>>23);
 102         if (beg>>26 == end>>26) return    1 + (beg>>26);
 103         return 0;
 104 }
 105
 106 static int get_tid(ti_index_t *idx, const char *ss)
 107 {
 108         khint_t k;
 109         int tid;
 110         k = kh_get(s, idx->tname, ss);
 111         if (k == kh_end(idx->tname)) { // a new target sequence
 112                 int ret, size;
 113                 // update idx->n, ->max, ->index and ->index2
 114                 if (idx->n == idx->max) {
 115                         idx->max = idx->max? idx->max<<1 : 8;
 116                         idx->index = realloc(idx->index, idx->max * sizeof(void*));
 117                         idx->index2 = realloc(idx->index2, idx->max * sizeof(ti_lidx_t));
 118                 }
 119                 memset(&idx->index2[idx->n], 0, sizeof(ti_lidx_t));
 120                 idx->index[idx->n++] = kh_init(i);
 121                 // update ->tname
 122                 tid = size = kh_size(idx->tname);
 123                 k = kh_put(s, idx->tname, strdup(ss), &ret);
 124                 kh_value(idx->tname, k) = size;
 125                 assert(idx->n == kh_size(idx->tname));
 126         } else tid = kh_value(idx->tname, k);
 127         return tid;
 128 }
 129
 130 int ti_get_intv(const ti_conf_t *conf, int len, char *line, ti_interval_t *intv)
 131 {
 132         int i, b = 0, id = 1, ncols = 0;
 133         char *s;
 134         intv->ss = intv->se = 0; intv->beg = intv->end = -1;
 135         for (i = 0; i <= len; ++i) {
 136                 if (line[i] == '\t' || line[i] == 0) {
 137             ++ncols;
 138                         if (id == conf->sc) {
 139                                 intv->ss = line + b; intv->se = line + i;
 140                         } else if (id == conf->bc) {
 141                                 // here ->beg is 0-based.
 142                                 intv->beg = intv->end = strtol(line + b, &s, 0);
 143                                 if (!(conf->preset&TI_FLAG_UCSC)) --intv->beg;
 144                                 else ++intv->end;
 145                                 if (intv->beg < 0) intv->beg = 0;
 146                                 if (intv->end < 1) intv->end = 1;
 147                         } else {
 148                                 if ((conf->preset&0xffff) == TI_PRESET_GENERIC) {
 149                                         if (id == conf->ec) intv->end = strtol(line + b, &s, 0);
 150                                 } else if ((conf->preset&0xffff) == TI_PRESET_SAM) {
 151                                         if (id == 6) { // CIGAR
 152                                                 int l = 0, op;
 153                                                 char *t;
 154                                                 for (s = line + b; s < line + i;) {
 155                                                         long x = strtol(s, &t, 10);
 156                                                         op = toupper(*t);
 157                                                         if (op == 'M' || op == 'D' || op == 'N') l += x;
 158                                                         s = t + 1;
 159                                                 }
 160                                                 if (l == 0) l = 1;
 161                                                 intv->end = intv->beg + l;
 162                                         }
 163                                 } else if ((conf->preset&0xffff) == TI_PRESET_VCF) {
 164                                         // FIXME: the following is NOT tested and is likely to be buggy
 165                                         if (id == 4) {
 166                                                 if (b < i) intv->end = intv->beg + (i - b);
 167                                         } else if (id == 8) { // look for "END="
 168                                                 int c = line[i];
 169                                                 line[i] = 0;
 170                                                 s = strstr(line + b, "END=");
 171                                                 if (s == line + b) s += 4;
 172                                                 else if (s) {
 173                                                         s = strstr(line + b, ";END=");
 174                                                         if (s) s += 5;
 175                                                 }
 176                                                 if (s) intv->end = strtol(s, &s, 0);
 177                                                 line[i] = c;
 178                                         }
 179                                 }
 180                         }
 181                         b = i + 1;
 182                         ++id;
 183                 }
 184         }
 185 /*
 186         if (ncols < conf->sc || ncols < conf->bc || ncols < conf->ec) {
 187                 if (ncols == 1) fprintf(pysamerr,"[get_intv] Is the file tab-delimited? The line has %d field only: %s\n", ncols, line);
 188                 else fprintf(pysamerr,"[get_intv] The line has %d field(s) only: %s\n", ncols, line);
 189                 exit(1);
 190         }
 191 */
 192         if (intv->ss == 0 || intv->se == 0 || intv->beg < 0 || intv->end < 0) return -1;
 193         return 0;
 194 }
 195
 196 static int get_intv(ti_index_t *idx, kstring_t *str, ti_intv_t *intv)
 197 {
 198         ti_interval_t x;
 199         intv->tid = intv->beg = intv->end = intv->bin = -1;
 200         if (ti_get_intv(&idx->conf, str->l, str->s, &x) == 0) {
 201                 int c = *x.se;
 202                 *x.se = '\0'; intv->tid = get_tid(idx, x.ss); *x.se = c;
 203                 intv->beg = x.beg; intv->end = x.end;
 204                 intv->bin = ti_reg2bin(intv->beg, intv->end);
 205                 return (intv->tid >= 0 && intv->beg >= 0 && intv->end >= 0)? 0 : -1;
 206         } else {
 207                 fprintf(pysamerr, "[%s] the following line cannot be parsed and skipped: %s\n", __func__, str->s);
 208                 return -1;
 209         }
 210 }
 211
 212 /************
 213  * indexing *
 214  ************/
 215
 216 // requirement: len <= LEN_MASK
 217 static inline void insert_offset(khash_t(i) *h, int bin, uint64_t beg, uint64_t end)
 218 {
 219         khint_t k;
 220         ti_binlist_t *l;
 221         int ret;
 222         k = kh_put(i, h, bin, &ret);
 223         l = &kh_value(h, k);
 224         if (ret) { // not present
 225                 l->m = 1; l->n = 0;
 226                 l->list = (pair64_t*)calloc(l->m, 16);
 227         }
 228         if (l->n == l->m) {
 229                 l->m <<= 1;
 230                 l->list = (pair64_t*)realloc(l->list, l->m * 16);
 231         }
 232         l->list[l->n].u = beg; l->list[l->n++].v = end;
 233 }
 234
 235 static inline uint64_t insert_offset2(ti_lidx_t *index2, int _beg, int _end, uint64_t offset)
 236 {
 237         int i, beg, end;
 238         beg = _beg >> TAD_LIDX_SHIFT;
 239         end = (_end - 1) >> TAD_LIDX_SHIFT;
 240         if (index2->m < end + 1) {
 241                 int old_m = index2->m;
 242                 index2->m = end + 1;
 243                 kroundup32(index2->m);
 244                 index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8);
 245                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
 246         }
 247         if (beg == end) {
 248                 if (index2->offset[beg] == 0) index2->offset[beg] = offset;
 249         } else {
 250                 for (i = beg; i <= end; ++i)
 251                         if (index2->offset[i] == 0) index2->offset[i] = offset;
 252         }
 253         if (index2->n < end + 1) index2->n = end + 1;
 254         return (uint64_t)beg<<32 | end;
 255 }
 256
 257 static void merge_chunks(ti_index_t *idx)
 258 {
 259         khash_t(i) *index;
 260         int i, l, m;
 261         khint_t k;
 262         for (i = 0; i < idx->n; ++i) {
 263                 index = idx->index[i];
 264                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 265                         ti_binlist_t *p;
 266                         if (!kh_exist(index, k)) continue;
 267                         p = &kh_value(index, k);
 268                         m = 0;
 269                         for (l = 1; l < p->n; ++l) {
 270                                 if (p->list[m].v>>16 == p->list[l].u>>16) p->list[m].v = p->list[l].v;
 271                                 else p->list[++m] = p->list[l];
 272                         } // ~for(l)
 273                         p->n = m + 1;
 274                 } // ~for(k)
 275         } // ~for(i)
 276 }
 277
 278 static void fill_missing(ti_index_t *idx)
 279 {
 280         int i, j;
 281         for (i = 0; i < idx->n; ++i) {
 282                 ti_lidx_t *idx2 = &idx->index2[i];
 283                 for (j = 1; j < idx2->n; ++j)
 284                         if (idx2->offset[j] == 0)
 285                                 idx2->offset[j] = idx2->offset[j-1];
 286         }
 287 }
 288
 289 ti_index_t *ti_index_core(BGZF *fp, const ti_conf_t *conf)
 290 {
 291         int ret;
 292         ti_index_t *idx;
 293         uint32_t last_bin, save_bin;
 294         int32_t last_coor, last_tid, save_tid;
 295         uint64_t save_off, last_off, lineno = 0, offset0 = (uint64_t)-1, tmp;
 296         kstring_t *str;
 297
 298         str = calloc(1, sizeof(kstring_t));
 299
 300         idx = (ti_index_t*)calloc(1, sizeof(ti_index_t));
 301         idx->conf = *conf;
 302         idx->n = idx->max = 0;
 303         idx->tname = kh_init(s);
 304         idx->index = 0;
 305         idx->index2 = 0;
 306
 307         save_bin = save_tid = last_tid = last_bin = 0xffffffffu;
 308         save_off = last_off = bgzf_tell(fp); last_coor = 0xffffffffu;
 309         while ((ret = ti_readline(fp, str)) >= 0) {
 310                 ti_intv_t intv;
 311                 ++lineno;
 312                 if (lineno <= idx->conf.line_skip || str->s[0] == idx->conf.meta_char) {
 313                         last_off = bgzf_tell(fp);
 314                         continue;
 315                 }
 316                 get_intv(idx, str, &intv);
 317         if ( intv.beg<0 || intv.end<0 )
 318         {
 319             fprintf(pysamerr,"[ti_index_core] the indexes overlap or are out of bounds\n");
 320             exit(1);
 321         }
 322                 if (last_tid != intv.tid) { // change of chromosomes
 323             if (last_tid>intv.tid )
 324             {
 325                 fprintf(pysamerr,"[ti_index_core] the chromosome blocks not continuous at line %llu, is the file sorted? [pos %d]\n",(unsigned long long)lineno,intv.beg+1);
 326                 exit(1);
 327             }
 328                         last_tid = intv.tid;
 329                         last_bin = 0xffffffffu;
 330                 } else if (last_coor > intv.beg) {
 331                         fprintf(pysamerr, "[ti_index_core] the file out of order at line %llu\n", (unsigned long long)lineno);
 332                         exit(1);
 333                 }
 334                 tmp = insert_offset2(&idx->index2[intv.tid], intv.beg, intv.end, last_off);
 335                 if (last_off == 0) offset0 = tmp;
 336                 if (intv.bin != last_bin) { // then possibly write the binning index
 337                         if (save_bin != 0xffffffffu) // save_bin==0xffffffffu only happens to the first record
 338                                 insert_offset(idx->index[save_tid], save_bin, save_off, last_off);
 339                         save_off = last_off;
 340                         save_bin = last_bin = intv.bin;
 341                         save_tid = intv.tid;
 342                         if (save_tid < 0) break;
 343                 }
 344                 if (bgzf_tell(fp) <= last_off) {
 345                         fprintf(pysamerr, "[ti_index_core] bug in BGZF: %llx < %llx\n",
 346                                         (unsigned long long)bgzf_tell(fp), (unsigned long long)last_off);
 347                         exit(1);
 348                 }
 349                 last_off = bgzf_tell(fp);
 350                 last_coor = intv.beg;
 351         }
 352         if (save_tid >= 0) insert_offset(idx->index[save_tid], save_bin, save_off, bgzf_tell(fp));
 353         merge_chunks(idx);
 354         fill_missing(idx);
 355         if (offset0 != (uint64_t)-1 && idx->n && idx->index2[0].offset) {
 356                 int i, beg = offset0>>32, end = offset0&0xffffffffu;
 357                 for (i = beg; i <= end; ++i) idx->index2[0].offset[i] = 0;
 358         }
 359
 360         free(str->s); free(str);
 361         return idx;
 362 }
 363
 364 void ti_index_destroy(ti_index_t *idx)
 365 {
 366         khint_t k;
 367         int i;
 368         if (idx == 0) return;
 369         // destroy the name hash table
 370         for (k = kh_begin(idx->tname); k != kh_end(idx->tname); ++k) {
 371                 if (kh_exist(idx->tname, k))
 372                         free((char*)kh_key(idx->tname, k));
 373         }
 374         kh_destroy(s, idx->tname);
 375         // destroy the binning index
 376         for (i = 0; i < idx->n; ++i) {
 377                 khash_t(i) *index = idx->index[i];
 378                 ti_lidx_t *index2 = idx->index2 + i;
 379                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 380                         if (kh_exist(index, k))
 381                                 free(kh_value(index, k).list);
 382                 }
 383                 kh_destroy(i, index);
 384                 free(index2->offset);
 385         }
 386         free(idx->index);
 387         // destroy the linear index
 388         free(idx->index2);
 389         free(idx);
 390 }
 391
 392 /******************
 393  * index file I/O *
 394  ******************/
 395
 396 void ti_index_save(const ti_index_t *idx, BGZF *fp)
 397 {
 398         int32_t i, size, ti_is_be;
 399         khint_t k;
 400         ti_is_be = bam_is_big_endian();
 401         bgzf_write(fp, "TBI\1", 4);
 402         if (ti_is_be) {
 403                 uint32_t x = idx->n;
 404                 bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 405         } else bgzf_write(fp, &idx->n, 4);
 406         assert(sizeof(ti_conf_t) == 24);
 407         if (ti_is_be) { // write ti_conf_t;
 408                 uint32_t x[6];
 409                 memcpy(x, &idx->conf, 24);
 410                 for (i = 0; i < 6; ++i) bgzf_write(fp, bam_swap_endian_4p(&x[i]), 4);
 411         } else bgzf_write(fp, &idx->conf, sizeof(ti_conf_t));
 412         { // write target names
 413                 char **name;
 414                 int32_t l = 0;
 415                 name = calloc(kh_size(idx->tname), sizeof(void*));
 416                 for (k = kh_begin(idx->tname); k != kh_end(idx->tname); ++k)
 417                         if (kh_exist(idx->tname, k))
 418                                 name[kh_value(idx->tname, k)] = (char*)kh_key(idx->tname, k);
 419                 for (i = 0; i < kh_size(idx->tname); ++i)
 420                         l += strlen(name[i]) + 1;
 421                 if (ti_is_be) bgzf_write(fp, bam_swap_endian_4p(&l), 4);
 422                 else bgzf_write(fp, &l, 4);
 423                 for (i = 0; i < kh_size(idx->tname); ++i)
 424                         bgzf_write(fp, name[i], strlen(name[i]) + 1);
 425                 free(name);
 426         }
 427         for (i = 0; i < idx->n; ++i) {
 428                 khash_t(i) *index = idx->index[i];
 429                 ti_lidx_t *index2 = idx->index2 + i;
 430                 // write binning index
 431                 size = kh_size(index);
 432                 if (ti_is_be) { // big endian
 433                         uint32_t x = size;
 434                         bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 435                 } else bgzf_write(fp, &size, 4);
 436                 for (k = kh_begin(index); k != kh_end(index); ++k) {
 437                         if (kh_exist(index, k)) {
 438                                 ti_binlist_t *p = &kh_value(index, k);
 439                                 if (ti_is_be) { // big endian
 440                                         uint32_t x;
 441                                         x = kh_key(index, k); bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 442                                         x = p->n; bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 443                                         for (x = 0; (int)x < p->n; ++x) {
 444                                                 bam_swap_endian_8p(&p->list[x].u);
 445                                                 bam_swap_endian_8p(&p->list[x].v);
 446                                         }
 447                                         bgzf_write(fp, p->list, 16 * p->n);
 448                                         for (x = 0; (int)x < p->n; ++x) {
 449                                                 bam_swap_endian_8p(&p->list[x].u);
 450                                                 bam_swap_endian_8p(&p->list[x].v);
 451                                         }
 452                                 } else {
 453                                         bgzf_write(fp, &kh_key(index, k), 4);
 454                                         bgzf_write(fp, &p->n, 4);
 455                                         bgzf_write(fp, p->list, 16 * p->n);
 456                                 }
 457                         }
 458                 }
 459                 // write linear index (index2)
 460                 if (ti_is_be) {
 461                         int x = index2->n;
 462                         bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 463                 } else bgzf_write(fp, &index2->n, 4);
 464                 if (ti_is_be) { // big endian
 465                         int x;
 466                         for (x = 0; (int)x < index2->n; ++x)
 467                                 bam_swap_endian_8p(&index2->offset[x]);
 468                         bgzf_write(fp, index2->offset, 8 * index2->n);
 469                         for (x = 0; (int)x < index2->n; ++x)
 470                                 bam_swap_endian_8p(&index2->offset[x]);
 471                 } else bgzf_write(fp, index2->offset, 8 * index2->n);
 472         }
 473 }
 474
 475 static ti_index_t *ti_index_load_core(BGZF *fp)
 476 {
 477         int i, ti_is_be;
 478         char magic[4];
 479         ti_index_t *idx;
 480         ti_is_be = bam_is_big_endian();
 481         if (fp == 0) {
 482                 fprintf(pysamerr, "[ti_index_load_core] fail to load index.\n");
 483                 return 0;
 484         }
 485         bgzf_read(fp, magic, 4);
 486         if (strncmp(magic, "TBI\1", 4)) {
 487                 fprintf(pysamerr, "[ti_index_load] wrong magic number.\n");
 488                 return 0;
 489         }
 490         idx = (ti_index_t*)calloc(1, sizeof(ti_index_t));
 491         bgzf_read(fp, &idx->n, 4);
 492         if (ti_is_be) bam_swap_endian_4p(&idx->n);
 493         idx->tname = kh_init(s);
 494         idx->index = (khash_t(i)**)calloc(idx->n, sizeof(void*));
 495         idx->index2 = (ti_lidx_t*)calloc(idx->n, sizeof(ti_lidx_t));
 496         // read idx->conf
 497         bgzf_read(fp, &idx->conf, sizeof(ti_conf_t));
 498         if (ti_is_be) {
 499                 bam_swap_endian_4p(&idx->conf.preset);
 500                 bam_swap_endian_4p(&idx->conf.sc);
 501                 bam_swap_endian_4p(&idx->conf.bc);
 502                 bam_swap_endian_4p(&idx->conf.ec);
 503                 bam_swap_endian_4p(&idx->conf.meta_char);
 504                 bam_swap_endian_4p(&idx->conf.line_skip);
 505         }
 506         { // read target names
 507                 int j, ret;
 508                 kstring_t *str;
 509                 int32_t l;
 510                 uint8_t *buf;
 511                 bgzf_read(fp, &l, 4);
 512                 if (ti_is_be) bam_swap_endian_4p(&l);
 513                 buf = calloc(l, 1);
 514                 bgzf_read(fp, buf, l);
 515                 str = calloc(1, sizeof(kstring_t));
 516                 for (i = j = 0; i < l; ++i) {
 517                         if (buf[i] == 0) {
 518                                 khint_t k = kh_put(s, idx->tname, strdup(str->s), &ret);
 519                                 kh_value(idx->tname, k) = j++;
 520                                 str->l = 0;
 521                         } else kputc(buf[i], str);
 522                 }
 523                 free(str->s); free(str); free(buf);
 524         }
 525         for (i = 0; i < idx->n; ++i) {
 526                 khash_t(i) *index;
 527                 ti_lidx_t *index2 = idx->index2 + i;
 528                 uint32_t key, size;
 529                 khint_t k;
 530                 int j, ret;
 531                 ti_binlist_t *p;
 532                 index = idx->index[i] = kh_init(i);
 533                 // load binning index
 534                 bgzf_read(fp, &size, 4);
 535                 if (ti_is_be) bam_swap_endian_4p(&size);
 536                 for (j = 0; j < (int)size; ++j) {
 537                         bgzf_read(fp, &key, 4);
 538                         if (ti_is_be) bam_swap_endian_4p(&key);
 539                         k = kh_put(i, index, key, &ret);
 540                         p = &kh_value(index, k);
 541                         bgzf_read(fp, &p->n, 4);
 542                         if (ti_is_be) bam_swap_endian_4p(&p->n);
 543                         p->m = p->n;
 544                         p->list = (pair64_t*)malloc(p->m * 16);
 545                         bgzf_read(fp, p->list, 16 * p->n);
 546                         if (ti_is_be) {
 547                                 int x;
 548                                 for (x = 0; x < p->n; ++x) {
 549                                         bam_swap_endian_8p(&p->list[x].u);
 550                                         bam_swap_endian_8p(&p->list[x].v);
 551                                 }
 552                         }
 553                 }
 554                 // load linear index
 555                 bgzf_read(fp, &index2->n, 4);
 556                 if (ti_is_be) bam_swap_endian_4p(&index2->n);
 557                 index2->m = index2->n;
 558                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 559                 bgzf_read(fp, index2->offset, index2->n * 8);
 560                 if (ti_is_be)
 561                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 562         }
 563         return idx;
 564 }
 565
 566 ti_index_t *ti_index_load_local(const char *fnidx)
 567 {
 568         BGZF *fp;
 569         fp = bgzf_open(fnidx, "r");
 570         if (fp) {
 571                 ti_index_t *idx = ti_index_load_core(fp);
 572                 bgzf_close(fp);
 573                 return idx;
 574         } else return 0;
 575 }
 576
 577 #ifdef _USE_KNETFILE
 578 static void download_from_remote(const char *url)
 579 {
 580         const int buf_size = 1 * 1024 * 1024;
 581         char *fn;
 582         FILE *fp;
 583         uint8_t *buf;
 584         knetFile *fp_remote;
 585         int l;
 586         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 587         l = strlen(url);
 588         for (fn = (char*)url + l - 1; fn >= url; --fn)
 589                 if (*fn == '/') break;
 590         ++fn; // fn now points to the file name
 591         fp_remote = knet_open(url, "r");
 592         if (fp_remote == 0) {
 593                 fprintf(pysamerr, "[download_from_remote] fail to open remote file.\n");
 594                 return;
 595         }
 596         if ((fp = fopen(fn, "w")) == 0) {
 597                 fprintf(pysamerr, "[download_from_remote] fail to create file in the working directory.\n");
 598                 knet_close(fp_remote);
 599                 return;
 600         }
 601         buf = (uint8_t*)calloc(buf_size, 1);
 602         while ((l = knet_read(fp_remote, buf, buf_size)) != 0)
 603                 fwrite(buf, 1, l, fp);
 604         free(buf);
 605         fclose(fp);
 606         knet_close(fp_remote);
 607 }
 608 #else
 609 static void download_from_remote(const char *url)
 610 {
 611         return;
 612 }
 613 #endif
 614
 615 static char *get_local_version(const char *fn)
 616 {
 617     struct stat sbuf;
 618         char *fnidx = (char*)calloc(strlen(fn) + 5, 1);
 619         strcat(strcpy(fnidx, fn), ".tbi");
 620         if ((strstr(fnidx, "ftp://") == fnidx || strstr(fnidx, "http://") == fnidx)) {
 621                 char *p, *url;
 622                 int l = strlen(fnidx);
 623                 for (p = fnidx + l - 1; p >= fnidx; --p)
 624                         if (*p == '/') break;
 625                 url = fnidx; fnidx = strdup(p + 1);
 626                 if (stat(fnidx, &sbuf) == 0) {
 627                         free(url);
 628                         return fnidx;
 629                 }
 630                 fprintf(pysamerr, "[%s] downloading the index file...\n", __func__);
 631                 download_from_remote(url);
 632                 free(url);
 633         }
 634     if (stat(fnidx, &sbuf) == 0) return fnidx;
 635         free(fnidx); return 0;
 636 }
 637
 638 const char **ti_seqname(const ti_index_t *idx, int *n)
 639 {
 640         const char **names;
 641         khint_t k;
 642         *n = idx->n;
 643         names = calloc(idx->n, sizeof(void*));
 644         for (k = kh_begin(idx->tname); k < kh_end(idx->tname); ++k)
 645                 if (kh_exist(idx->tname, k))
 646                         names[kh_val(idx->tname, k)] = kh_key(idx->tname, k);
 647         return names;
 648 }
 649
 650 ti_index_t *ti_index_load(const char *fn)
 651 {
 652         ti_index_t *idx;
 653     char *fname = get_local_version(fn);
 654         if (fname == 0) return 0;
 655         idx = ti_index_load_local(fname);
 656         if (idx == 0) fprintf(pysamerr, "[ti_index_load] fail to load the index: %s\n", fname);
 657     free(fname);
 658         return idx;
 659 }
 660
 661 int ti_index_build2(const char *fn, const ti_conf_t *conf, const char *_fnidx)
 662 {
 663         char *fnidx;
 664         BGZF *fp, *fpidx;
 665         ti_index_t *idx;
 666         if ((fp = bgzf_open(fn, "r")) == 0) {
 667                 fprintf(pysamerr, "[ti_index_build2] fail to open the file: %s\n", fn);
 668                 return -1;
 669         }
 670         idx = ti_index_core(fp, conf);
 671         bgzf_close(fp);
 672         if (_fnidx == 0) {
 673                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 674                 strcpy(fnidx, fn); strcat(fnidx, ".tbi");
 675         } else fnidx = strdup(_fnidx);
 676         fpidx = bgzf_open(fnidx, "w");
 677         if (fpidx == 0) {
 678                 fprintf(pysamerr, "[ti_index_build2] fail to create the index file.\n");
 679                 free(fnidx);
 680                 return -1;
 681         }
 682         ti_index_save(idx, fpidx);
 683         ti_index_destroy(idx);
 684         bgzf_close(fpidx);
 685         free(fnidx);
 686         return 0;
 687 }
 688
 689 int ti_index_build(const char *fn, const ti_conf_t *conf)
 690 {
 691         return ti_index_build2(fn, conf, 0);
 692 }
 693
 694 /********************************************
 695  * parse a region in the format chr:beg-end *
 696  ********************************************/
 697
 698 int ti_get_tid(const ti_index_t *idx, const char *name)
 699 {
 700         khiter_t iter;
 701         const khash_t(s) *h = idx->tname;
 702         iter = kh_get(s, h, name); /* get the tid */
 703         if (iter == kh_end(h)) return -1;
 704         return kh_value(h, iter);
 705 }
 706
 707 int ti_parse_region(const ti_index_t *idx, const char *str, int *tid, int *begin, int *end)
 708 {
 709         char *s, *p;
 710         int i, l, k;
 711         l = strlen(str);
 712         p = s = (char*)malloc(l+1);
 713         /* squeeze out "," */
 714         for (i = k = 0; i != l; ++i)
 715                 if (str[i] != ',' && !isspace(str[i])) s[k++] = str[i];
 716         s[k] = 0;
 717         for (i = 0; i != k; ++i) if (s[i] == ':') break;
 718         s[i] = 0;
 719         if ((*tid = ti_get_tid(idx, s)) < 0) {
 720                 free(s);
 721                 return -1;
 722         }
 723         if (i == k) { /* dump the whole sequence */
 724                 *begin = 0; *end = 1<<29; free(s);
 725                 return 0;
 726         }
 727         for (p = s + i + 1; i != k; ++i) if (s[i] == '-') break;
 728         *begin = atoi(p);
 729         if (i < k) {
 730                 p = s + i + 1;
 731                 *end = atoi(p);
 732         } else *end = 1<<29;
 733         if (*begin > 0) --*begin;
 734         free(s);
 735         if (*begin > *end) return -1;
 736         return 0;
 737 }
 738
 739 /*******************************
 740  * retrieve a specified region *
 741  *******************************/
 742
 743 #define MAX_BIN 37450 // =(8^6-1)/7+1
 744
 745 static inline int reg2bins(uint32_t beg, uint32_t end, uint16_t list[MAX_BIN])
 746 {
 747         int i = 0, k;
 748         if (beg >= end) return 0;
 749         if (end >= 1u<<29) end = 1u<<29;
 750         --end;
 751         list[i++] = 0;
 752         for (k =    1 + (beg>>26); k <=    1 + (end>>26); ++k) list[i++] = k;
 753         for (k =    9 + (beg>>23); k <=    9 + (end>>23); ++k) list[i++] = k;
 754         for (k =   73 + (beg>>20); k <=   73 + (end>>20); ++k) list[i++] = k;
 755         for (k =  585 + (beg>>17); k <=  585 + (end>>17); ++k) list[i++] = k;
 756         for (k = 4681 + (beg>>14); k <= 4681 + (end>>14); ++k) list[i++] = k;
 757         return i;
 758 }
 759
 760 ti_iter_t ti_iter_first()
 761 {
 762         ti_iter_t iter;
 763         iter = calloc(1, sizeof(struct __ti_iter_t));
 764         iter->from_first = 1;
 765         return iter;
 766 }
 767
 768 ti_iter_t ti_iter_query(const ti_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         ti_iter_t iter = 0;
 777
 778         if (beg < 0) beg = 0;
 779         if (end < beg) return 0;
 780         // initialize the iterator
 781         iter = calloc(1, sizeof(struct __ti_iter_t));
 782         iter->idx = idx; iter->tid = tid; iter->beg = beg; iter->end = end; iter->i = -1;
 783         // random access
 784         bins = (uint16_t*)calloc(MAX_BIN, 2);
 785         n_bins = reg2bins(beg, end, bins);
 786         index = idx->index[tid];
 787         if (idx->index2[tid].n > 0) {
 788                 min_off = (beg>>TAD_LIDX_SHIFT >= idx->index2[tid].n)? idx->index2[tid].offset[idx->index2[tid].n-1]
 789                         : idx->index2[tid].offset[beg>>TAD_LIDX_SHIFT];
 790                 if (min_off == 0) { // improvement for index files built by tabix prior to 0.1.4
 791                         int n = beg>>TAD_LIDX_SHIFT;
 792                         if (n > idx->index2[tid].n) n = idx->index2[tid].n;
 793                         for (i = n - 1; i >= 0; --i)
 794                                 if (idx->index2[tid].offset[i] != 0) break;
 795                         if (i >= 0) min_off = idx->index2[tid].offset[i];
 796                 }
 797         } else min_off = 0; // tabix 0.1.2 may produce such index files
 798         for (i = n_off = 0; i < n_bins; ++i) {
 799                 if ((k = kh_get(i, index, bins[i])) != kh_end(index))
 800                         n_off += kh_value(index, k).n;
 801         }
 802         if (n_off == 0) {
 803                 free(bins); return iter;
 804         }
 805         off = (pair64_t*)calloc(n_off, 16);
 806         for (i = n_off = 0; i < n_bins; ++i) {
 807                 if ((k = kh_get(i, index, bins[i])) != kh_end(index)) {
 808                         int j;
 809                         ti_binlist_t *p = &kh_value(index, k);
 810                         for (j = 0; j < p->n; ++j)
 811                                 if (p->list[j].v > min_off) off[n_off++] = p->list[j];
 812                 }
 813         }
 814         if (n_off == 0) {
 815                 free(bins); free(off); return iter;
 816         }
 817         free(bins);
 818         {
 819                 int l;
 820                 ks_introsort(offt, n_off, off);
 821                 // resolve completely contained adjacent blocks
 822                 for (i = 1, l = 0; i < n_off; ++i)
 823                         if (off[l].v < off[i].v)
 824                                 off[++l] = off[i];
 825                 n_off = l + 1;
 826                 // resolve overlaps between adjacent blocks; this may happen due to the merge in indexing
 827                 for (i = 1; i < n_off; ++i)
 828                         if (off[i-1].v >= off[i].u) off[i-1].v = off[i].u;
 829                 { // merge adjacent blocks
 830                         for (i = 1, l = 0; i < n_off; ++i) {
 831                                 if (off[l].v>>16 == off[i].u>>16) off[l].v = off[i].v;
 832                                 else off[++l] = off[i];
 833                         }
 834                         n_off = l + 1;
 835                 }
 836         }
 837         iter->n_off = n_off; iter->off = off;
 838         return iter;
 839 }
 840
 841 const char *ti_iter_read(BGZF *fp, ti_iter_t iter, int *len)
 842 {
 843         if (iter->finished) return 0;
 844         if (iter->from_first) {
 845                 int ret;
 846                 if ((ret = ti_readline(fp, &iter->str)) < 0) {
 847                         iter->finished = 1;
 848                         return 0;
 849                 } else {
 850                         if (len) *len = iter->str.l;
 851                         return iter->str.s;
 852                 }
 853         }
 854         if (iter->n_off == 0) return 0;
 855         while (1) {
 856                 int ret;
 857                 if (iter->curr_off == 0 || iter->curr_off >= iter->off[iter->i].v) { // then jump to the next chunk
 858                         if (iter->i == iter->n_off - 1) break; // no more chunks
 859                         if (iter->i >= 0) assert(iter->curr_off == iter->off[iter->i].v); // otherwise bug
 860                         if (iter->i < 0 || iter->off[iter->i].v != iter->off[iter->i+1].u) { // not adjacent chunks; then seek
 861                                 bgzf_seek(fp, iter->off[iter->i+1].u, SEEK_SET);
 862                                 iter->curr_off = bgzf_tell(fp);
 863                         }
 864                         ++iter->i;
 865                 }
 866                 if ((ret = ti_readline(fp, &iter->str)) >= 0) {
 867                         ti_intv_t intv;
 868                         iter->curr_off = bgzf_tell(fp);
 869                         if (iter->str.s[0] == iter->idx->conf.meta_char) continue;
 870                         get_intv((ti_index_t*)iter->idx, &iter->str, &intv);
 871                         if (intv.tid != iter->tid || intv.beg >= iter->end) break; // no need to proceed
 872                         else if (intv.end > iter->beg && iter->end > intv.beg) {
 873                                 if (len) *len = iter->str.l;
 874                                 return iter->str.s;
 875                         }
 876                 } else break; // end of file
 877         }
 878         iter->finished = 1;
 879         return 0;
 880 }
 881
 882 void ti_iter_destroy(ti_iter_t iter)
 883 {
 884         if (iter) {
 885                 free(iter->str.s); free(iter->off);
 886                 free(iter);
 887         }
 888 }
 889
 890 int ti_fetch(BGZF *fp, const ti_index_t *idx, int tid, int beg, int end, void *data, ti_fetch_f func)
 891 {
 892         ti_iter_t iter;
 893         const char *s;
 894         int len;
 895         iter = ti_iter_query(idx, tid, beg, end);
 896         while ((s = ti_iter_read(fp, iter, &len)) != 0)
 897                 func(len, s, data);
 898         ti_iter_destroy(iter);
 899         return 0;
 900 }
 901
 902 const ti_conf_t *ti_get_conf(ti_index_t *idx) { return idx? &idx->conf : 0; }
 903
 904 /*******************
 905  * High-level APIs *
 906  *******************/
 907
 908 tabix_t *ti_open(const char *fn, const char *fnidx)
 909 {
 910         tabix_t *t;
 911         BGZF *fp;
 912         if ((fp = bgzf_open(fn, "r")) == 0) return 0;
 913         t = calloc(1, sizeof(tabix_t));
 914         t->fn = strdup(fn);
 915         if (fnidx) t->fnidx = strdup(fnidx);
 916         t->fp = fp;
 917         return t;
 918 }
 919
 920 void ti_close(tabix_t *t)
 921 {
 922         if (t) {
 923                 bgzf_close(t->fp);
 924                 if (t->idx) ti_index_destroy(t->idx);
 925                 free(t->fn); free(t->fnidx);
 926                 free(t);
 927         }
 928 }
 929
 930 int ti_lazy_index_load(tabix_t *t)
 931 {
 932         if (t->idx == 0) { // load index
 933                 if (t->fnidx) t->idx = ti_index_load_local(t->fnidx);
 934                 else t->idx = ti_index_load(t->fn);
 935                 if (t->idx == 0) return -1; // fail to load index
 936         }
 937         return 0;
 938 }
 939
 940 ti_iter_t ti_queryi(tabix_t *t, int tid, int beg, int end)
 941 {
 942         if (tid < 0) return ti_iter_first();
 943         if (ti_lazy_index_load(t) != 0) return 0;
 944         return ti_iter_query(t->idx, tid, beg, end);
 945 }
 946
 947 ti_iter_t ti_querys(tabix_t *t, const char *reg)
 948 {
 949         int tid, beg, end;
 950         if (reg == 0) return ti_iter_first();
 951         if (ti_lazy_index_load(t) != 0) return 0;
 952         if (ti_parse_region(t->idx, reg, &tid, &beg, &end) < 0) return 0;
 953         return ti_iter_query(t->idx, tid, beg, end);
 954 }
 955
 956 ti_iter_t ti_query(tabix_t *t, const char *name, int beg, int end)
 957 {
 958         int tid;
 959         if (name == 0) return ti_iter_first();
 960         // then need to load the index
 961         if (ti_lazy_index_load(t) != 0) return 0;
 962         if ((tid = ti_get_tid(t->idx, name)) < 0) return 0;
 963         return ti_iter_query(t->idx, tid, beg, end);
 964 }
 965
 966 const char *ti_read(tabix_t *t, ti_iter_t iter, int *len)
 967 {
 968         return ti_iter_read(t->fp, iter, len);
 969 }