bcftools/index.c

   1 #include <assert.h>
   2 #include <ctype.h>
   3 #include <sys/stat.h>
   4 #include "bam_endian.h"
   5 #include "kstring.h"
   6 #include "bcf.h"
   7 #ifdef _USE_KNETFILE
   8 #include "knetfile.h"
   9 #endif
  10
  11 #define TAD_LIDX_SHIFT 13
  12
  13 typedef struct {
  14         int32_t n, m;
  15         uint64_t *offset;
  16 } bcf_lidx_t;
  17
  18 struct __bcf_idx_t {
  19         int32_t n;
  20         bcf_lidx_t *index2;
  21 };
  22
  23 /************
  24  * indexing *
  25  ************/
  26
  27 static inline void insert_offset2(bcf_lidx_t *index2, int _beg, int _end, uint64_t offset)
  28 {
  29         int i, beg, end;
  30         beg = _beg >> TAD_LIDX_SHIFT;
  31         end = (_end - 1) >> TAD_LIDX_SHIFT;
  32         if (index2->m < end + 1) {
  33                 int old_m = index2->m;
  34                 index2->m = end + 1;
  35                 kroundup32(index2->m);
  36                 index2->offset = (uint64_t*)realloc(index2->offset, index2->m * 8);
  37                 memset(index2->offset + old_m, 0, 8 * (index2->m - old_m));
  38         }
  39         if (beg == end) {
  40                 if (index2->offset[beg] == 0) index2->offset[beg] = offset;
  41         } else {
  42                 for (i = beg; i <= end; ++i)
  43                         if (index2->offset[i] == 0) index2->offset[i] = offset;
  44         }
  45         if (index2->n < end + 1) index2->n = end + 1;
  46 }
  47
  48 static void fill_missing(bcf_idx_t *idx)
  49 {
  50         int i, j;
  51         for (i = 0; i < idx->n; ++i) {
  52                 bcf_lidx_t *idx2 = &idx->index2[i];
  53                 for (j = 1; j < idx2->n; ++j)
  54                         if (idx2->offset[j] == 0)
  55                                 idx2->offset[j] = idx2->offset[j-1];
  56         }
  57 }
  58
  59 bcf_idx_t *bcf_idx_core(bcf_t *bp)
  60 {
  61         bcf_idx_t *idx;
  62         int32_t last_coor, last_tid;
  63         uint64_t last_off;
  64         kstring_t *str;
  65         BGZF *fp = bp->fp;
  66         bcf1_t *b;
  67         int ret;
  68
  69         b = calloc(1, sizeof(bcf1_t));
  70         str = calloc(1, sizeof(kstring_t));
  71         idx = (bcf_idx_t*)calloc(1, sizeof(bcf_idx_t));
  72         idx->n = bp->h.n_ref;
  73         idx->index2 = calloc(bp->h.n_ref, sizeof(bcf_lidx_t));
  74
  75         last_tid = 0xffffffffu;
  76         last_off = bgzf_tell(fp); last_coor = 0xffffffffu;
  77         while ((ret = bcf_read(bp, b)) > 0) {
  78                 int end, tmp;
  79                 if (last_tid != b->tid) { // change of chromosomes
  80                         last_tid = b->tid;
  81                 } else if (last_coor > b->pos) {
  82                         fprintf(stderr, "[bcf_idx_core] the input is out of order\n");
  83                         free(str->s); free(str); free(idx); bcf_destroy(b);
  84                         return 0;
  85                 }
  86                 tmp = strlen(b->ref);
  87                 end = b->pos + (tmp > 0? tmp : 1);
  88                 insert_offset2(&idx->index2[b->tid], b->pos, end, last_off);
  89                 last_off = bgzf_tell(fp);
  90                 last_coor = b->pos;
  91         }
  92         fill_missing(idx);
  93         free(str->s); free(str); bcf_destroy(b);
  94         return idx;
  95 }
  96
  97 void bcf_idx_destroy(bcf_idx_t *idx)
  98 {
  99         int i;
 100         if (idx == 0) return;
 101         for (i = 0; i < idx->n; ++i) free(idx->index2[i].offset);
 102         free(idx->index2);
 103         free(idx);
 104 }
 105
 106 /******************
 107  * index file I/O *
 108  ******************/
 109
 110 void bcf_idx_save(const bcf_idx_t *idx, BGZF *fp)
 111 {
 112         int32_t i, ti_is_be;
 113         ti_is_be = bam_is_big_endian();
 114         bgzf_write(fp, "BCI\4", 4);
 115         if (ti_is_be) {
 116                 uint32_t x = idx->n;
 117                 bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 118         } else bgzf_write(fp, &idx->n, 4);
 119         for (i = 0; i < idx->n; ++i) {
 120                 bcf_lidx_t *index2 = idx->index2 + i;
 121                 // write linear index (index2)
 122                 if (ti_is_be) {
 123                         int x = index2->n;
 124                         bgzf_write(fp, bam_swap_endian_4p(&x), 4);
 125                 } else bgzf_write(fp, &index2->n, 4);
 126                 if (ti_is_be) { // big endian
 127                         int x;
 128                         for (x = 0; (int)x < index2->n; ++x)
 129                                 bam_swap_endian_8p(&index2->offset[x]);
 130                         bgzf_write(fp, index2->offset, 8 * index2->n);
 131                         for (x = 0; (int)x < index2->n; ++x)
 132                                 bam_swap_endian_8p(&index2->offset[x]);
 133                 } else bgzf_write(fp, index2->offset, 8 * index2->n);
 134         }
 135 }
 136
 137 static bcf_idx_t *bcf_idx_load_core(BGZF *fp)
 138 {
 139         int i, ti_is_be;
 140         char magic[4];
 141         bcf_idx_t *idx;
 142         ti_is_be = bam_is_big_endian();
 143         if (fp == 0) {
 144                 fprintf(stderr, "[%s] fail to load index.\n", __func__);
 145                 return 0;
 146         }
 147         bgzf_read(fp, magic, 4);
 148         if (strncmp(magic, "BCI\4", 4)) {
 149                 fprintf(stderr, "[%s] wrong magic number.\n", __func__);
 150                 return 0;
 151         }
 152         idx = (bcf_idx_t*)calloc(1, sizeof(bcf_idx_t));
 153         bgzf_read(fp, &idx->n, 4);
 154         if (ti_is_be) bam_swap_endian_4p(&idx->n);
 155         idx->index2 = (bcf_lidx_t*)calloc(idx->n, sizeof(bcf_lidx_t));
 156         for (i = 0; i < idx->n; ++i) {
 157                 bcf_lidx_t *index2 = idx->index2 + i;
 158                 int j;
 159                 bgzf_read(fp, &index2->n, 4);
 160                 if (ti_is_be) bam_swap_endian_4p(&index2->n);
 161                 index2->m = index2->n;
 162                 index2->offset = (uint64_t*)calloc(index2->m, 8);
 163                 bgzf_read(fp, index2->offset, index2->n * 8);
 164                 if (ti_is_be)
 165                         for (j = 0; j < index2->n; ++j) bam_swap_endian_8p(&index2->offset[j]);
 166         }
 167         return idx;
 168 }
 169
 170 bcf_idx_t *bcf_idx_load_local(const char *fnidx)
 171 {
 172         BGZF *fp;
 173         fp = bgzf_open(fnidx, "r");
 174         if (fp) {
 175                 bcf_idx_t *idx = bcf_idx_load_core(fp);
 176                 bgzf_close(fp);
 177                 return idx;
 178         } else return 0;
 179 }
 180
 181 #ifdef _USE_KNETFILE
 182 static void download_from_remote(const char *url)
 183 {
 184         const int buf_size = 1 * 1024 * 1024;
 185         char *fn;
 186         FILE *fp;
 187         uint8_t *buf;
 188         knetFile *fp_remote;
 189         int l;
 190         if (strstr(url, "ftp://") != url && strstr(url, "http://") != url) return;
 191         l = strlen(url);
 192         for (fn = (char*)url + l - 1; fn >= url; --fn)
 193                 if (*fn == '/') break;
 194         ++fn; // fn now points to the file name
 195         fp_remote = knet_open(url, "r");
 196         if (fp_remote == 0) {
 197                 fprintf(stderr, "[download_from_remote] fail to open remote file.\n");
 198                 return;
 199         }
 200         if ((fp = fopen(fn, "w")) == 0) {
 201                 fprintf(stderr, "[download_from_remote] fail to create file in the working directory.\n");
 202                 knet_close(fp_remote);
 203                 return;
 204         }
 205         buf = (uint8_t*)calloc(buf_size, 1);
 206         while ((l = knet_read(fp_remote, buf, buf_size)) != 0)
 207                 fwrite(buf, 1, l, fp);
 208         free(buf);
 209         fclose(fp);
 210         knet_close(fp_remote);
 211 }
 212 #else
 213 static void download_from_remote(const char *url)
 214 {
 215         return;
 216 }
 217 #endif
 218
 219 static char *get_local_version(const char *fn)
 220 {
 221     struct stat sbuf;
 222         char *fnidx = (char*)calloc(strlen(fn) + 5, 1);
 223         strcat(strcpy(fnidx, fn), ".bci");
 224         if ((strstr(fnidx, "ftp://") == fnidx || strstr(fnidx, "http://") == fnidx)) {
 225                 char *p, *url;
 226                 int l = strlen(fnidx);
 227                 for (p = fnidx + l - 1; p >= fnidx; --p)
 228                         if (*p == '/') break;
 229                 url = fnidx; fnidx = strdup(p + 1);
 230                 if (stat(fnidx, &sbuf) == 0) {
 231                         free(url);
 232                         return fnidx;
 233                 }
 234                 fprintf(stderr, "[%s] downloading the index file...\n", __func__);
 235                 download_from_remote(url);
 236                 free(url);
 237         }
 238     if (stat(fnidx, &sbuf) == 0) return fnidx;
 239         free(fnidx); return 0;
 240 }
 241
 242 bcf_idx_t *bcf_idx_load(const char *fn)
 243 {
 244         bcf_idx_t *idx;
 245     char *fname = get_local_version(fn);
 246         if (fname == 0) return 0;
 247         idx = bcf_idx_load_local(fname);
 248     free(fname);
 249         return idx;
 250 }
 251
 252 int bcf_idx_build2(const char *fn, const char *_fnidx)
 253 {
 254         char *fnidx;
 255         BGZF *fpidx;
 256         bcf_t *bp;
 257         bcf_idx_t *idx;
 258         if ((bp = bcf_open(fn, "r")) == 0) {
 259                 fprintf(stderr, "[bcf_idx_build2] fail to open the BAM file.\n");
 260                 return -1;
 261         }
 262         idx = bcf_idx_core(bp);
 263         bcf_close(bp);
 264         if (_fnidx == 0) {
 265                 fnidx = (char*)calloc(strlen(fn) + 5, 1);
 266                 strcpy(fnidx, fn); strcat(fnidx, ".bci");
 267         } else fnidx = strdup(_fnidx);
 268         fpidx = bgzf_open(fnidx, "w");
 269         if (fpidx == 0) {
 270                 fprintf(stderr, "[bcf_idx_build2] fail to create the index file.\n");
 271                 free(fnidx);
 272                 return -1;
 273         }
 274         bcf_idx_save(idx, fpidx);
 275         bcf_idx_destroy(idx);
 276         bgzf_close(fpidx);
 277         free(fnidx);
 278         return 0;
 279 }
 280
 281 int bcf_idx_build(const char *fn)
 282 {
 283         return bcf_idx_build2(fn, 0);
 284 }
 285
 286 /********************************************
 287  * parse a region in the format chr:beg-end *
 288  ********************************************/
 289
 290 int bcf_parse_region(void *str2id, const char *str, int *tid, int *begin, int *end)
 291 {
 292         char *s, *p;
 293         int i, l, k;
 294         l = strlen(str);
 295         p = s = (char*)malloc(l+1);
 296         /* squeeze out "," */
 297         for (i = k = 0; i != l; ++i)
 298                 if (str[i] != ',' && !isspace(str[i])) s[k++] = str[i];
 299         s[k] = 0;
 300         for (i = 0; i != k; ++i) if (s[i] == ':') break;
 301         s[i] = 0;
 302         if ((*tid = bcf_str2id(str2id, s)) < 0) {
 303                 free(s);
 304                 return -1;
 305         }
 306         if (i == k) { /* dump the whole sequence */
 307                 *begin = 0; *end = 1<<29; free(s);
 308                 return 0;
 309         }
 310         for (p = s + i + 1; i != k; ++i) if (s[i] == '-') break;
 311         *begin = atoi(p);
 312         if (i < k) {
 313                 p = s + i + 1;
 314                 *end = atoi(p);
 315         } else *end = 1<<29;
 316         if (*begin > 0) --*begin;
 317         free(s);
 318         if (*begin > *end) return -1;
 319         return 0;
 320 }
 321
 322 /*******************************
 323  * retrieve a specified region *
 324  *******************************/
 325
 326 uint64_t bcf_idx_query(const bcf_idx_t *idx, int tid, int beg, int end)
 327 {
 328         uint64_t min_off;
 329         if (beg < 0) beg = 0;
 330         if (end < beg) return 0xffffffffffffffffull;
 331         min_off = (beg>>TAD_LIDX_SHIFT >= idx->index2[tid].n)? idx->index2[tid].offset[idx->index2[tid].n-1]
 332                 : idx->index2[tid].offset[beg>>TAD_LIDX_SHIFT];
 333         assert(min_off);
 334         return min_off;
 335 }
 336
 337 int bcf_main_index(int argc, char *argv[])
 338 {
 339         if (argc == 1) {
 340                 fprintf(stderr, "Usage: bcftools index <in.bcf>\n");
 341                 return 1;
 342         }
 343         bcf_idx_build(argv[1]);
 344         return 0;
 345 }