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 == '\n') { // an empty line
  70                         if (state == 1) {
  71                                 offset = razf_tell(rz);
  72                                 continue;
  73                         } else if ((state == 0 && len < 0) || state == 2) continue;
  74                 }
  75                 if (c == '>') { // fasta header
  76                         if (len >= 0)
  77                                 fai_insert_index(idx, name, len, line_len, line_blen, offset);
  78                         l_name = 0;
  79                         while ((ret = razf_read(rz, &c, 1)) != 0 && !isspace(c)) {
  80                                 if (m_name < l_name + 2) {
  81                                         m_name = l_name + 2;
  82                                         kroundup32(m_name);
  83                                         name = (char*)realloc(name, m_name);
  84                                 }
  85                                 name[l_name++] = c;
  86                         }
  87                         name[l_name] = '\0';
  88                         assert(ret);
  89                         if (c != '\n') while (razf_read(rz, &c, 1) && c != '\n');
  90                         state = 1; len = 0;
  91                         offset = razf_tell(rz);
  92                 } else {
  93                         if (state == 3) {
  94                                 fprintf(stderr, "[fai_build_core] inlined empty line is not allowed in sequence '%s'. Abort!\n", name);
  95                                 exit(1);
  96                         }
  97                         if (state == 2) state = 3;
  98                         l1 = l2 = 0;
  99                         do {
 100                                 ++l1;
 101                                 if (isgraph(c)) ++l2;
 102                         } while ((ret = razf_read(rz, &c, 1)) && c != '\n');
 103                         if (state == 3 && l2) {
 104                                 fprintf(stderr, "[fai_build_core] different line length in sequence '%s'. Abort!\n", name);
 105                                 exit(1);
 106                         }
 107                         ++l1; len += l2;
 108                         if (l2 >= 0x10000) {
 109                                 fprintf(stderr, "[fai_build_core] line length exceeds 65535 in sequence '%s'. Abort!\n", name);
 110                                 exit(1);
 111                         }
 112                         if (state == 1) line_len = l1, line_blen = l2, state = 0;
 113                         else if (state == 0) {
 114                                 if (l1 != line_len || l2 != line_blen) state = 2;
 115                         }
 116                 }
 117         }
 118         fai_insert_index(idx, name, len, line_len, line_blen, offset);
 119         free(name);
 120         return idx;
 121 }
 122
 123 void fai_save(const faidx_t *fai, FILE *fp)
 124 {
 125         khint_t k;
 126         int i;
 127         for (i = 0; i < fai->n; ++i) {
 128                 faidx1_t x;
 129                 k = kh_get(s, fai->hash, fai->name[i]);
 130                 x = kh_value(fai->hash, k);
 131                 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);
 132         }
 133 }
 134
 135 faidx_t *fai_read(FILE *fp)
 136 {
 137         faidx_t *fai;
 138         char *buf, *p;
 139         int len, line_len, line_blen;
 140         long long offset;
 141         fai = (faidx_t*)calloc(1, sizeof(faidx_t));
 142         fai->hash = kh_init(s);
 143         buf = (char*)calloc(0x10000, 1);
 144         while (!feof(fp) && fgets(buf, 0x10000, fp)) {
 145                 for (p = buf; *p && isgraph(*p); ++p);
 146                 *p = 0; ++p;
 147                 sscanf(p, "%d%lld%d%d", &len, &offset, &line_blen, &line_len);
 148                 fai_insert_index(fai, buf, len, line_len, line_blen, offset);
 149         }
 150         free(buf);
 151         return fai;
 152 }
 153
 154 void fai_destroy(faidx_t *fai)
 155 {
 156         int i;
 157         for (i = 0; i < fai->n; ++i) free(fai->name[i]);
 158         free(fai->name);
 159         kh_destroy(s, fai->hash);
 160         if (fai->rz) razf_close(fai->rz);
 161         free(fai);
 162 }
 163
 164 void fai_build(const char *fn)
 165 {
 166         char *str;
 167         RAZF *rz;
 168         FILE *fp;
 169         faidx_t *fai;
 170         str = (char*)calloc(strlen(fn) + 5, 1);
 171         sprintf(str, "%s.fai", fn);
 172         rz = razf_open(fn, "r");
 173         assert(rz);
 174         fai = fai_build_core(rz);
 175         razf_close(rz);
 176         fp = fopen(str, "w");
 177         assert(fp);
 178         fai_save(fai, fp);
 179         fclose(fp);
 180         free(str);
 181         fai_destroy(fai);
 182 }
 183
 184 faidx_t *fai_load(const char *fn)
 185 {
 186         char *str;
 187         FILE *fp;
 188         faidx_t *fai;
 189         str = (char*)calloc(strlen(fn) + 5, 1);
 190         sprintf(str, "%s.fai", fn);
 191         fp = fopen(str, "r");
 192         if (fp == 0) {
 193                 fprintf(stderr, "[fai_load] build FASTA index.\n");
 194                 fai_build(fn);
 195                 fp = fopen(str, "r");
 196                 if (fp == 0) {
 197                         free(str);
 198                         return 0;
 199                 }
 200         }
 201         fai = fai_read(fp);
 202         fclose(fp);
 203         fai->rz = razf_open(fn, "r");
 204         if (fai->rz == 0) return 0;
 205         assert(fai->rz);
 206         free(str);
 207         return fai;
 208 }
 209
 210 char *fai_fetch(const faidx_t *fai, const char *str, int *len)
 211 {
 212         char *s, *p, c;
 213         int i, l, k;
 214         khiter_t iter;
 215         faidx1_t val;
 216         khash_t(s) *h;
 217         int beg, end;
 218
 219         beg = end = -1;
 220         h = fai->hash;
 221         l = strlen(str);
 222         p = s = (char*)malloc(l+1);
 223         /* squeeze out "," */
 224         for (i = k = 0; i != l; ++i)
 225                 if (str[i] != ',' && !isspace(str[i])) s[k++] = str[i];
 226         s[k] = 0;
 227         for (i = 0; i != k; ++i) if (s[i] == ':') break;
 228         s[i] = 0;
 229         iter = kh_get(s, h, s); /* get the ref_id */
 230         if (iter == kh_end(h)) {
 231                 *len = 0;
 232                 free(s); return 0;
 233         }
 234         val = kh_value(h, iter);
 235         if (i == k) { /* dump the whole sequence */
 236                 beg = 0; end = val.len;
 237         } else {
 238                 for (p = s + i + 1; i != k; ++i) if (s[i] == '-') break;
 239                 beg = atoi(p);
 240                 if (i < k) {
 241                         p = s + i + 1;
 242                         end = atoi(p);
 243                 } else end = val.len;
 244         }
 245         if (beg > 0) --beg;
 246         if (beg >= val.len) beg = val.len;
 247         if (end >= val.len) end = val.len;
 248         if (beg > end) beg = end;
 249         free(s);
 250
 251         // now retrieve the sequence
 252         l = 0;
 253         s = (char*)malloc(end - beg + 2);
 254         razf_seek(fai->rz, val.offset + beg / val.line_blen * val.line_len + beg % val.line_blen, SEEK_SET);
 255         while (razf_read(fai->rz, &c, 1) == 1 && l < end - beg)
 256                 if (isgraph(c)) s[l++] = c;
 257         s[l] = '\0';
 258         *len = l;
 259         return s;
 260 }
 261
 262 int faidx_main(int argc, char *argv[])
 263 {
 264         if (argc == 1) {
 265                 fprintf(stderr, "Usage: faidx <in.fasta> [<reg> [...]]\n");
 266                 return 1;
 267         } else {
 268                 if (argc == 2) fai_build(argv[1]);
 269                 else {
 270                         int i, j, k, l;
 271                         char *s;
 272                         faidx_t *fai;
 273                         fai = fai_load(argv[1]);
 274                         assert(fai);
 275                         for (i = 2; i != argc; ++i) {
 276                                 printf(">%s\n", argv[i]);
 277                                 s = fai_fetch(fai, argv[i], &l);
 278                                 for (j = 0; j < l; j += 60) {
 279                                         for (k = 0; k < 60 && k < l - j; ++k)
 280                                                 putchar(s[j + k]);
 281                                         putchar('\n');
 282                                 }
 283                                 free(s);
 284                         }
 285                         fai_destroy(fai);
 286                 }
 287         }
 288         return 0;
 289 }
 290
 291 #ifdef FAIDX_MAIN
 292 int main(int argc, char *argv[]) { return faidx_main(argc, argv); }
 293 #endif