faidx.c

   1 #include <ctype.h>
   2 #include <assert.h>
   3 #include <string.h>
   4 #include <stdlib.h>
   5 #include <stdio.h>
   6 #include "faidx.h"
   7 #include "khash.h"
   8
   9 typedef struct {
  10         uint64_t len:32, line_len:16, line_blen:16;
  11         uint64_t offset;
  12 } faidx1_t;
  13 KHASH_MAP_INIT_STR(s, faidx1_t)
  14
  15 #ifndef _NO_RAZF
  16 #include "razf.h"
  17 #else
  18 extern off_t ftello(FILE *stream);
  19 extern int fseeko(FILE *stream, off_t offset, int whence);
  20 #define RAZF FILE
  21 #define razf_read(fp, buf, size) fread(buf, 1, size, fp)
  22 #define razf_open(fn, mode) fopen(fn, mode)
  23 #define razf_close(fp) fclose(fp)
  24 #define razf_seek(fp, offset, whence) fseeko(fp, offset, whence)
  25 #define razf_tell(fp) ftello(fp)
  26 #endif
  27
  28 struct __faidx_t {
  29         RAZF *rz;
  30         int n, m;
  31         char **name;
  32         khash_t(s) *hash;
  33 };
  34
  35 #ifndef kroundup32
  36 #define kroundup32(x) (--(x), (x)|=(x)>>1, (x)|=(x)>>2, (x)|=(x)>>4, (x)|=(x)>>8, (x)|=(x)>>16, ++(x))
  37 #endif
  38
  39 static inline void fai_insert_index(faidx_t *idx, const char *name, int len, int line_len, int line_blen, uint64_t offset)
  40 {
  41         khint_t k;
  42         int ret;
  43         faidx1_t t;
  44         if (idx->n == idx->m) {
  45                 idx->m = idx->m? idx->m<<1 : 16;
  46                 idx->name = (char**)realloc(idx->name, sizeof(void*) * idx->m);
  47         }
  48         idx->name[idx->n] = strdup(name);
  49         k = kh_put(s, idx->hash, idx->name[idx->n], &ret);
  50         t.len = len; t.line_len = line_len; t.line_blen = line_blen; t.offset = offset;
  51         kh_value(idx->hash, k) = t;
  52         ++idx->n;
  53 }
  54
  55 faidx_t *fai_build_core(RAZF *rz)
  56 {
  57         char c, *name;
  58         int l_name, m_name, ret;
  59         int len, line_len, line_blen, state;
  60         int l1, l2;
  61         faidx_t *idx;
  62         uint64_t offset;
  63
  64         idx = (faidx_t*)calloc(1, sizeof(faidx_t));
  65         idx->hash = kh_init(s);
  66         name = 0; l_name = m_name = 0;
  67         len = line_len = line_blen = -1; state = 0; l1 = l2 = -1; offset = 0;
  68         while (razf_read(rz, &c, 1)) {
  69                 if (c == '>') { // fasta header
  70                         if (len >= 0)
  71                                 fai_insert_index(idx, name, len, line_len, line_blen, offset);
  72                         l_name = 0;
  73                         while ((ret = razf_read(rz, &c, 1)) != 0 && !isspace(c)) {
  74                                 if (m_name < l_name + 2) {
  75                                         m_name = l_name + 2;
  76                                         kroundup32(m_name);
  77                                         name = (char*)realloc(name, m_name);
  78                                 }
  79                                 name[l_name++] = c;
  80                         }
  81                         name[l_name] = '\0';
  82                         assert(ret);
  83                         if (c != '\n') while (razf_read(rz, &c, 1) && c != '\n');
  84                         state = 1; len = 0;
  85                         offset = razf_tell(rz);
  86                 } else {
  87                         if (state == 3) {
  88                                 fprintf(stderr, "[fai_build_core] inlined empty line is not allowed in sequence '%s'. Abort!\n", name);
  89                                 exit(1);
  90                         }
  91                         if (state == 2) state = 3;
  92                         l1 = l2 = 0;
  93                         do {
  94                                 ++l1;
  95                                 if (isgraph(c)) ++l2;
  96                         } while ((ret = razf_read(rz, &c, 1)) && c != '\n');
  97                         if (state == 3 && l2) {
  98                                 fprintf(stderr, "[fai_build_core] different line length in sequence '%s'. Abort!\n", name);
  99                                 exit(1);
 100                         }
 101                         ++l1; len += l2;
 102                         if (l2 >= 0x10000) {
 103                                 fprintf(stderr, "[fai_build_core] line length exceeds 65535 in sequence '%s'. Abort!\n", name);
 104                                 exit(1);
 105                         }
 106                         if (state == 1) line_len = l1, line_blen = l2, state = 0;
 107                         else if (state == 0) {
 108                                 if (l1 != line_len || l2 != line_blen) state = 2;
 109                         }
 110                 }
 111         }
 112         fai_insert_index(idx, name, len, line_len, line_blen, offset);
 113         free(name);
 114         return idx;
 115 }
 116
 117 void fai_save(const faidx_t *fai, FILE *fp)
 118 {
 119         khint_t k;
 120         int i;
 121         for (i = 0; i < fai->n; ++i) {
 122                 faidx1_t x;
 123                 k = kh_get(s, fai->hash, fai->name[i]);
 124                 x = kh_value(fai->hash, k);
 125                 fprintf(fp, "%s\t%d\t%lld\t%d\t%d\n", fai->name[i], (int)x.len, (long long)x.offset, (int)x.line_blen, (int)x.line_len);
 126         }
 127 }
 128
 129 faidx_t *fai_read(FILE *fp)
 130 {
 131         faidx_t *fai;
 132         char *buf, *p;
 133         int len, line_len, line_blen;
 134         long long offset;
 135         fai = (faidx_t*)calloc(1, sizeof(faidx_t));
 136         fai->hash = kh_init(s);
 137         buf = (char*)calloc(0x10000, 1);
 138         while (!feof(fp) && fgets(buf, 0x10000, fp)) {
 139                 for (p = buf; *p && isgraph(*p); ++p);
 140                 *p = 0; ++p;
 141                 sscanf(p, "%d%lld%d%d", &len, &offset, &line_blen, &line_len);
 142                 fai_insert_index(fai, buf, len, line_len, line_blen, offset);
 143         }
 144         free(buf);
 145         return fai;
 146 }
 147
 148 void fai_destroy(faidx_t *fai)
 149 {
 150         int i;
 151         for (i = 0; i < fai->n; ++i) free(fai->name[i]);
 152         free(fai->name);
 153         kh_destroy(s, fai->hash);
 154         if (fai->rz) razf_close(fai->rz);
 155         free(fai);
 156 }
 157
 158 void fai_build(const char *fn)
 159 {
 160         char *str;
 161         RAZF *rz;
 162         FILE *fp;
 163         faidx_t *fai;
 164         str = (char*)calloc(strlen(fn) + 5, 1);
 165         sprintf(str, "%s.fai", fn);
 166         rz = razf_open(fn, "r");
 167         assert(rz);
 168         fai = fai_build_core(rz);
 169         razf_close(rz);
 170         fp = fopen(str, "w");
 171         assert(fp);
 172         fai_save(fai, fp);
 173         fclose(fp);
 174         free(str);
 175         fai_destroy(fai);
 176 }
 177
 178 faidx_t *fai_load(const char *fn)
 179 {
 180         char *str;
 181         FILE *fp;
 182         faidx_t *fai;
 183         str = (char*)calloc(strlen(fn) + 5, 1);
 184         sprintf(str, "%s.fai", fn);
 185         fp = fopen(str, "r");
 186         if (fp == 0) {
 187                 fprintf(stderr, "[fai_load] build FASTA index.\n");
 188                 fai_build(fn);
 189                 fp = fopen(str, "r");
 190                 if (fp == 0) {
 191                         free(str);
 192                         return 0;
 193                 }
 194         }
 195         fai = fai_read(fp);
 196         fclose(fp);
 197         fai->rz = razf_open(fn, "r");
 198         if (fai->rz == 0) return 0;
 199         assert(fai->rz);
 200         free(str);
 201         return fai;
 202 }
 203
 204 char *fai_fetch(const faidx_t *fai, const char *str, int *len)
 205 {
 206         char *s, *p, c;
 207         int i, l, k;
 208         khiter_t iter;
 209         faidx1_t val;
 210         khash_t(s) *h;
 211         int beg, end;
 212
 213         beg = end = -1;
 214         h = fai->hash;
 215         l = strlen(str);
 216         p = s = (char*)malloc(l+1);
 217         /* squeeze out "," */
 218         for (i = k = 0; i != l; ++i)
 219                 if (str[i] != ',' && !isspace(str[i])) s[k++] = str[i];
 220         s[k] = 0;
 221         for (i = 0; i != k; ++i) if (s[i] == ':') break;
 222         s[i] = 0;
 223         iter = kh_get(s, h, s); /* get the ref_id */
 224         if (iter == kh_end(h)) {
 225                 *len = 0;
 226                 free(s); return 0;
 227         }
 228         val = kh_value(h, iter);
 229         if (i == k) { /* dump the whole sequence */
 230                 beg = 0; end = val.len;
 231         } else {
 232                 for (p = s + i + 1; i != k; ++i) if (s[i] == '-') break;
 233                 beg = atoi(p);
 234                 if (i < k) {
 235                         p = s + i + 1;
 236                         end = atoi(p);
 237                 } else end = val.len;
 238         }
 239         if (beg > 0) --beg;
 240         if (beg >= val.len) beg = val.len;
 241         if (end >= val.len) end = val.len;
 242         if (beg > end) beg = end;
 243         free(s);
 244
 245         // now retrieve the sequence
 246         l = 0;
 247         s = (char*)malloc(end - beg + 2);
 248         razf_seek(fai->rz, val.offset + beg / val.line_blen * val.line_len + beg % val.line_blen, SEEK_SET);
 249         while (razf_read(fai->rz, &c, 1) == 1 && l < end - beg)
 250                 if (isgraph(c)) s[l++] = c;
 251         s[l] = '\0';
 252         *len = l;
 253         return s;
 254 }
 255
 256 int faidx_main(int argc, char *argv[])
 257 {
 258         if (argc == 1) {
 259                 fprintf(stderr, "Usage: faidx <in.fasta> [<reg> [...]]\n");
 260                 return 1;
 261         } else {
 262                 if (argc == 2) fai_build(argv[1]);
 263                 else {
 264                         int i, j, k, l;
 265                         char *s;
 266                         faidx_t *fai;
 267                         fai = fai_load(argv[1]);
 268                         assert(fai);
 269                         for (i = 2; i != argc; ++i) {
 270                                 printf(">%s\n", argv[i]);
 271                                 s = fai_fetch(fai, argv[i], &l);
 272                                 for (j = 0; j < l; j += 60) {
 273                                         for (k = 0; k < 60 && k < l - j; ++k)
 274                                                 putchar(s[j + k]);
 275                                         putchar('\n');
 276                                 }
 277                                 free(s);
 278                         }
 279                         fai_destroy(fai);
 280                 }
 281         }
 282         return 0;
 283 }
 284
 285 #ifdef FAIDX_MAIN
 286 int main(int argc, char *argv[]) { return faidx_main(argc, argv); }
 287 #endif