minibwa 0.0.0

data/ext/minibwa/minibwa/bwt.c

@@ -0,0 +1,715 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+#include <assert.h>
+#include <stdint.h>
+#include "kommon.h"
+#include "kalloc.h"
+#include "bwt.h"
+
+/********************
+ * Basic operations *
+ ********************/
+
+static void bwt_gen_cnt_table(uint32_t cnt[256])
+{
+	int i, j;
+	for (i = 0; i != 256; ++i) {
+		uint32_t x = 0;
+		for (j = 0; j != 4; ++j)
+			x |= (((i&3) == j) + ((i>>2&3) == j) + ((i>>4&3) == j) + (i>>6 == j)) << (j<<3);
+		cnt[i] = x;
+	}
+}
+
+mb_bwt_t *mb_bwt_init(void)
+{
+	mb_bwt_t *bwt;
+	bwt = kom_calloc(mb_bwt_t, 1);
+	bwt->sa_bit = (uint32_t)-1;
+	bwt_gen_cnt_table(bwt->cnt_table);
+	return bwt;
+}
+
+void mb_bwt_destroy(mb_bwt_t *bwt)
+{
+	if (bwt == 0) return;
+	free(bwt->pre); // pre is always heap-allocated (by mb_bwt_cache)
+	if (bwt->mmap) kom_munmap(bwt->mmap, bwt->mmap_len); // data/sa point into the mapped file
+	else { free(bwt->sa); free(bwt->data); }
+	free(bwt);
+}
+
+/******************
+ * Encode raw BWT *
+ ******************/
+
+#define raw_B00(b, k) ((b)[(k)>>4]>>((~(k)&0xf)<<1)&3)
+
+#define BWT_CNT_SHIFT 56
+#define BWT_CNT_MASK ((1ULL<<BWT_CNT_SHIFT) - 1)
+
+static uint64_t mb_bwt_data_len(uint64_t len)
+{
+	uint64_t bwt_len, occ_len;
+	bwt_len = (len + 127) / 128 * 4;
+	occ_len = ((len + 127) / 128 + 1) * 4; // +1 for the final counts
+	return bwt_len + occ_len;
+}
+
+/* BWT layout. Each block consists of u64[4]+u32[8], 64 bytes in total. The
+ * lower 56 bits of each u64[4] (see BWT_CNT_SHIFT) store the accumulative
+ * count of A/C/G/T bases. The higher 8 bits store the count of A/C/G/T in the
+ * next 64nt. u32[8] keeps a BWT substring of 128nt in length. Because it
+ * follows little endian, it can also be considered as u64[4] etc.
+ */
+mb_bwt_t *mb_bwt_init_from_raw(int is_byte, const void *raw_, uint64_t len, uint64_t primary)
+{
+	uint64_t c[4], x[4], i, k, *last_c = 0;
+	mb_bwt_t *bwt;
+	const uint32_t *raw32 = is_byte? 0 : (const uint32_t*)raw_;
+	const uint8_t *raw8 = is_byte? (const uint8_t*)raw_ : 0;
+
+	bwt = mb_bwt_init();
+	bwt->primary = primary;
+	bwt->seq_len = len;
+	bwt->data_len = mb_bwt_data_len(len);
+	bwt->data = kom_calloc(uint64_t, bwt->data_len);
+	if (len == 0) return bwt; // nothing to encode; avoid the last-block overflow
+
+	memset(c, 0, 32);
+	for (i = k = 0; i < len; ++i) {
+		uint8_t j, a = is_byte? raw8[i]&3 : raw_B00(raw32, i);
+		if ((i & 0x7f) == 0) { // hard coded: each block encodes 128bp
+			if (i > 0) {
+				memcpy(&bwt->data[k], x, 32);
+				k += 4;
+			}
+			last_c = &bwt->data[k];
+			memcpy(&bwt->data[k], c, 32);
+			k += 4;
+			memset(x, 0, 32);
+		} else if ((i & 0x3f) == 0 && last_c) {
+			for (j = 0; j < 4; ++j)
+				last_c[j] |= (c[j] - last_c[j]) << BWT_CNT_SHIFT;
+		}
+		++c[a];
+		x[(i&0x7f)>>5] |= (uint64_t)a << ((i&0x1f)<<1); // little endian
+	}
+	// the last block
+	memcpy(&bwt->data[k], x, 32);
+	k += 4;
+	memcpy(&bwt->data[k], c, 32);
+	k += 4;
+	assert(k == bwt->data_len);
+	for (i = 0, bwt->L2[0] = 0; i < 4; ++i)
+		bwt->L2[i+1] = bwt->L2[i] + c[i];
+	assert(bwt->L2[4] == len);
+	return bwt;
+}
+
+/********
+ * Rank *
+ ********/
+
+#define bwt_block(b, k) ((b)->data + ((k)>>7<<3))
+
+// retrieve a character from the $-removed BWT string. Note that mb_bwt_t::data is
+// not exactly the BWT string and therefore this macro is called bwt_B0 instead of bwt_B
+#define bwt_B0(b, k) ((b)->data[((k)>>7<<3) + 4 + (((k)&127)>>5)] >> (((k)&31)<<1) & 3)
+
+static inline void mb_bwt_block_prefetch(const mb_bwt_t *bwt, uint64_t k)
+{
+	if (k > 0) __builtin_prefetch(bwt_block(bwt, k - 1 - (k - 1 >= bwt->primary)));
+}
+
+static inline int rank_aux1(uint64_t y, uint8_t c)
+{
+	// reduce nucleotide counting to bits counting
+	y = ((c&2)? y : ~y) >> 1 & ((c&1)? y : ~y) & 0x5555555555555555ull;
+	// count the number of 1s in y
+	#if 0
+	y = (y & 0x3333333333333333ull) + (y >> 2 & 0x3333333333333333ull);
+	return ((y + (y >> 4)) & 0xf0f0f0f0f0f0f0full) * 0x101010101010101ull >> 56;
+	#else
+	return __builtin_popcountll(y);
+	#endif
+}
+
+uint64_t mb_bwt_rank11(const mb_bwt_t *bwt, uint64_t k, uint8_t c)
+{
+	const uint64_t *p, *end;
+	uint64_t n, mask;
+	if (k == 0) return 0;
+	if (k == bwt->seq_len + 1) return bwt->L2[c+1] - bwt->L2[c];
+	--k;
+	k -= (k >= bwt->primary); // because $ is not in bwt
+	mask = (k&0x7f) >= 64? (1ULL << (64 - BWT_CNT_SHIFT)) - 1 : 0;
+	p = bwt_block(bwt, k);
+	n = (p[c] & BWT_CNT_MASK) + ((p[c] >> BWT_CNT_SHIFT) & mask);
+	p += 4; // p points to 2-bit encoded BWT
+	end = p + ((k&0x7f) >> 5);
+	p += (k&0x7f) >= 64? 2 : 0;
+//	for (; p < end; ++p) n += rank_aux1(*p, c); // we go through this loop 0 or 1 time
+	if (p < end) n += rank_aux1(*p, c), ++p;
+	n += rank_aux1(*p << ((~k&0x1f) << 1), c);
+	if (c == 0) n -= ~k&0x1f; // "A" may be overcounted due to the shift above; this line corrects that
+	return n;
+}
+
+static inline const uint32_t *seek_block(const mb_bwt_t *bwt, uint64_t k, uint64_t cnt[4])
+{
+	const uint64_t *p = bwt_block(bwt, k);
+	uint64_t mask = (k&0x7f) >= 64? (1ULL << (64 - BWT_CNT_SHIFT)) - 1 : 0;
+	cnt[0] = (p[0] & BWT_CNT_MASK) + ((p[0] >> BWT_CNT_SHIFT) & mask);
+	cnt[1] = (p[1] & BWT_CNT_MASK) + ((p[1] >> BWT_CNT_SHIFT) & mask);
+	cnt[2] = (p[2] & BWT_CNT_MASK) + ((p[2] >> BWT_CNT_SHIFT) & mask);
+	cnt[3] = (p[3] & BWT_CNT_MASK) + ((p[3] >> BWT_CNT_SHIFT) & mask);
+	return (const uint32_t*)(p + 4);
+}
+
+static inline uint32_t rank_aux4(const mb_bwt_t *bwt, uint32_t x)
+{
+	return bwt->cnt_table[x&0xff] + bwt->cnt_table[x>>8&0xff] + bwt->cnt_table[x>>16&0xff] + bwt->cnt_table[x>>24];
+}
+
+void mb_bwt_rank1a(const mb_bwt_t *bwt, uint64_t k, uint64_t cnt[4])
+{
+	const uint32_t *q, *end;
+	uint32_t x, tmp;
+	if (k == 0) {
+		memset(cnt, 0, 4 * sizeof(uint64_t));
+		return;
+	}
+	--k;
+	k -= (k >= bwt->primary); // because $ is not in bwt
+	q = seek_block(bwt, k, cnt);
+	end = q + ((k&0x7f) >> 4);
+	q += (k&0x7f) >= 64? 4 : 0;
+	for (x = 0; q < end; ++q) x += rank_aux4(bwt, *q); // NB: this assumes little endian
+	tmp = *q << ((~k&0xf) << 1);
+	x += rank_aux4(bwt, tmp) - (~k&0xf);
+	cnt[0] += x&0xff, cnt[1] += x>>8&0xff, cnt[2] += x>>16&0xff, cnt[3] += x>>24;
+}
+
+void mb_bwt_rank2a(const mb_bwt_t *bwt, uint64_t k, uint64_t l, uint64_t cntk[4], uint64_t cntl[4])
+{
+	uint64_t k1 = k - 1, l1 = l - 1;
+	k1 -= (k1 >= bwt->primary);
+	l1 -= (l1 >= bwt->primary);
+	mb_bwt_block_prefetch(bwt, k);
+	if (k1>>7 != l1>>7 || k == 0 || l == 0) {
+		mb_bwt_block_prefetch(bwt, l);
+		mb_bwt_rank1a(bwt, k, cntk);
+		mb_bwt_rank1a(bwt, l, cntl);
+	} else if (l - k == 1) { // we can use a simpler procedure
+		uint64_t z = k - (k > bwt->primary);
+		mb_bwt_rank1a(bwt, k, cntk);
+		memcpy(cntl, cntk, 4 * sizeof(uint64_t));
+		++cntl[bwt_B0(bwt, z)];
+	} else {
+		const uint32_t *q, *endk, *endl;
+		uint32_t x, y, tmp;
+		k = k1, l = l1;
+		q = seek_block(bwt, k, cntk);
+		// prepare cntk[]
+		endk = q + ((k&0x7f) >> 4);
+		endl = q + ((l&0x7f) >> 4);
+		q += (k&0x7f) >= 64? 4 : 0;
+		for (x = 0; q < endk; ++q) x += rank_aux4(bwt, *q);
+		y = x;
+		tmp = *q << ((~k&0xf) << 1);
+		x += rank_aux4(bwt, tmp) - (~k&0xf);
+		// calculate cntl[] and finalize cntk[]
+		for (; q < endl; ++q) y += rank_aux4(bwt, *q);
+		tmp = *q << ((~l&0xf) << 1);
+		y += rank_aux4(bwt, tmp) - (~l&0xf);
+		memcpy(cntl, cntk, 4 * sizeof(uint64_t));
+		cntk[0] += x&0xff; cntk[1] += x>>8&0xff; cntk[2] += x>>16&0xff; cntk[3] += x>>24;
+		cntl[0] += y&0xff; cntl[1] += y>>8&0xff; cntl[2] += y>>16&0xff; cntl[3] += y>>24;
+	}
+}
+
+/*********************
+ * Bidirectional BWT *
+ *********************/
+
+void mb_bwt_extend(const mb_bwt_t *bwt, const mb_sai_t *ik, mb_sai_t ok[4], int is_back)
+{
+	uint64_t tk[4], tl[4];
+	int i;
+	mb_bwt_rank2a(bwt, ik->x[!is_back], ik->x[!is_back] + ik->size, tk, tl);
+	for (i = 0; i != 4; ++i) {
+		ok[i].x[!is_back] = bwt->L2[i] + 1 + tk[i]; // +1 for the missing sentinel
+		ok[i].size = (tl[i] -= tk[i]);
+	}
+	ok[3].x[is_back] = ik->x[is_back] + (ik->x[!is_back] <= bwt->primary && ik->x[!is_back] + ik->size > bwt->primary);
+	ok[2].x[is_back] = ok[3].x[is_back] + tl[3];
+	ok[1].x[is_back] = ok[2].x[is_back] + tl[2];
+	ok[0].x[is_back] = ok[1].x[is_back] + tl[1];
+}
+
+// backward search from pos
+static int64_t mb_bwt_back(const mb_bwt_t *f, const uint8_t *q, int64_t st, int64_t pos, int64_t min_occ, mb_sai_t *p)
+{
+	int64_t i = pos - 1;
+	mb_sai_t ok[4];
+	assert(q[pos] < 4); // the backward pass never involves N
+	if (f->pre && pos - st >= f->pre_len) { // then use precomputed k-mer index instead of base-by-base extension
+		uint64_t z = 0, l = 0;
+		for (i = pos; l < f->pre_len; --i, ++l) // get the k-mer
+			z = z << 2 | q[i];                  // NB: this loop doesn't check N
+		assert(z < 1<<f->pre_len*2);
+		*p = f->pre[z];
+	} else p->size = 0;
+	if (p->size < min_occ) { // then we need to use the standard procedure
+		mb_bwt_set_intv(f, q[pos], p);
+		i = pos - 1;
+	}
+	for (; i >= st; --i) { // backward extension
+		int c = q[i];
+		if (c > 3) break;
+		mb_bwt_extend(f, p, ok, 1);
+		if (ok[c].size < min_occ) break;
+		*p = ok[c];
+	}
+	return i;
+}
+
+// find super MEMs (SMEMs). See ropebwt3
+int64_t mb_bwt_smem(const mb_bwt_t *f, uint32_t len, const uint8_t *q, int64_t x, int64_t min_len, int64_t min_occ, mb_sai_t *p)
+{
+	int64_t i, j, xn;
+	mb_sai_t ik, ok[4];
+
+	assert(len <= INT32_MAX); // this can be relaxed if we define a new struct for mem
+	p->size = ik.size = 0;
+	if (len - x < min_len) return len;
+	for (i = x, xn = -1; i < x + min_len; ++i) // find the last N in [x,x+min_len)
+		if (q[i] > 3) xn = i;
+	if (xn >= 0) return xn + 1;
+	i = mb_bwt_back(f, q, x, x + min_len - 1, min_occ, &ik);
+	if (i >= x) return i + 1; // no MEM found
+	for (j = x + min_len; j < len; ++j) { // forward extension
+		int c = 3 - q[j];
+		if (q[j] > 3) break;
+		mb_bwt_extend(f, &ik, ok, 0);
+		if (ok[c].size < min_occ) break;
+		ik = ok[c];
+	}
+	*p = ik;
+	p->info = (uint64_t)x<<32 | j;
+	if (j == len) return len; // reaching end; no need to do another round
+	i = q[j] > 3? j : mb_bwt_back(f, q, x + 1, j, min_occ, &ik);
+	return i + 1;
+}
+
+/**************
+ * Batch SMEM *
+ **************/
+
+typedef struct { // a simplified version of kdq
+	int32_t front, count, cap;
+	int32_t *a;
+} tiny_queue_t;
+
+static void tq_init(void *km, tiny_queue_t *q, int32_t n)
+{
+	q->cap = n;
+	kom_roundup32(q->cap);
+	q->a = Kcalloc(km, int32_t, q->cap);
+	q->front = q->count = 0;
+}
+
+static inline void tq_push(tiny_queue_t *q, int32_t x)
+{
+	q->a[((q->count++) + q->front) & (q->cap - 1)] = x;
+}
+
+static inline int32_t tq_shift(tiny_queue_t *q)
+{
+	int32_t x;
+	if (q->count == 0) return -1;
+	x = q->a[q->front++];
+	q->front &= q->cap - 1;
+	--q->count;
+	return x;
+}
+
+static inline void se_one_step_back(const mb_bwt_t *bwt, mb_smem_entry_t *s)
+{
+	mb_sai_t ok[4];
+	int32_t c = s->q[s->i];
+	assert(c < 4); // shouldn't happen
+	mb_bwt_extend(bwt, &s->p, ok, 1);
+	if (ok[c].size < s->min_occ) { // move back to stage1
+		s->x = s->i + 1;
+		s->stage = 1;
+	} else { // stay in the two backward stages
+		s->p = ok[c];
+		s->i--;
+		mb_bwt_block_prefetch(bwt, s->p.x[0]); // prefetch for the next backward iteration
+		mb_bwt_block_prefetch(bwt, s->p.x[0] + s->p.size);
+	}
+}
+
+void mb_bwt_smem_batch(void *km, const mb_bwt_t *bwt, int32_t n, mb_smem_entry_t *a)
+{
+	int32_t i;
+	tiny_queue_t tq;
+
+	// initialize
+	tq_init(km, &tq, n);
+	for (i = 0; i < n; ++i) {
+		mb_smem_entry_t *s = &a[i];
+		tq_push(&tq, i);
+		s->stage = 1;
+		s->x = s->st;
+		if (s->v->m < 64) { // preallocate to avoid frequent krealloc(), which can be slow
+			s->v->m = 64;
+			s->v->a = Krealloc(km, mb_sai_t, s->v->a, s->v->m);
+		}
+	}
+
+	// core loop
+	while (tq.count > 0) {
+		int32_t idx;
+		mb_smem_entry_t *s;
+
+		idx = tq_shift(&tq);
+		s = &a[idx];
+		if (s->stage == 1) { // set interval for the first backward pass in smem; require ->x
+			int32_t i, xn;
+			if (s->en - s->x < s->min_len)
+				continue; // IMPORTANT: this skips the tq_push() at the end of this long while loop
+			for (i = s->x, xn = -1; i < s->x + s->min_len; ++i) // find the position of the last N
+				if (s->q[i] > 3) xn = i;
+			if (xn >= 0) { // skip N and stay in stage 1
+				s->x = xn + 1;
+			} else {
+				s->i = s->x + s->min_len - 1;
+				if (bwt->pre && s->min_len >= bwt->pre_len) { // get k-mer for prefetch
+					for (i = 0, s->kmer = 0; i < bwt->pre_len; ++i, s->i--)
+						s->kmer = s->kmer << 2 | s->q[s->i]; // backward pass shouldn't meet N
+					__builtin_prefetch(&bwt->pre[s->kmer]);
+					s->stage = 2;
+				} else { // skip stage 2
+					mb_bwt_set_intv(bwt, s->q[s->i--], &s->p);
+					s->stage = 3;
+				}
+			}
+		} else if (s->stage == 2 || s->stage == 5) { // k-mer lookup
+			s->p = bwt->pre[s->kmer];
+			if (s->p.size < s->min_occ) { // jumped too far with the k-mer cache; revert
+				s->i += bwt->pre_len;
+				mb_bwt_set_intv(bwt, s->q[s->i--], &s->p);
+			}
+			s->stage++;
+		} else if (s->stage == 3) { // first backward pass; require ->{i,p}
+			if (s->i < s->x) { // move to the next stage
+				mb_bwt_block_prefetch(bwt, s->p.x[1]); // prefetch for the forward pass
+				mb_bwt_block_prefetch(bwt, s->p.x[1] + s->p.size);
+				s->i = s->x + s->min_len;
+				s->stage = 4;
+			} else se_one_step_back(bwt, s);
+		} else if (s->stage == 4) { // forward pass; require ->{i,p}
+			if (s->i == s->en) {
+				s->p.info = (uint64_t)s->x << 32 | s->i;
+				Kgrow(km, mb_sai_t, s->v->a, s->v->n, s->v->m);
+				s->v->a[s->v->n++] = s->p; // save the interval
+				continue; // trigger termination as tq_push() at the end of the loop is skipped
+			} else {
+				int32_t i, c = 3 - (int32_t)s->q[s->i];
+				mb_sai_t ok[4];
+				if (c >= 0) mb_bwt_extend(bwt, &s->p, ok, 0);
+				if (c >= 0 && ok[c].size >= s->min_occ) { // stay in stage 4
+					s->p = ok[c];
+					s->i++;
+					mb_bwt_block_prefetch(bwt, s->p.x[1]);
+					mb_bwt_block_prefetch(bwt, s->p.x[1] + s->p.size);
+				} else {
+					s->p.info = (uint64_t)s->x << 32 | s->i;
+					Kgrow(km, mb_sai_t, s->v->a, s->v->n, s->v->m);
+					s->v->a[s->v->n++] = s->p; // save the interval
+					if (c < 0) { // if N, skip it and move back to stage 1
+						s->x = s->i + 1;
+						s->stage = 1;
+					} else if (bwt->pre && s->i - s->x - 1 >= bwt->pre_len) { // get k-mer
+						for (i = 0, s->kmer = 0; i < bwt->pre_len; ++i, s->i--)
+							s->kmer = s->kmer << 2 | s->q[s->i];
+						__builtin_prefetch(&bwt->pre[s->kmer]);
+						s->stage = 5;
+					} else { // skip stage 5
+						mb_bwt_set_intv(bwt, s->q[s->i--], &s->p);
+						s->stage = 6;
+					}
+				}
+			}
+		} else if (s->stage == 6) { // second backward pass
+			if (s->i < s->x + 1) {
+				s->x = s->i + 1;
+				s->stage = 1;
+			} else se_one_step_back(bwt, s);
+		}
+		tq_push(&tq, idx);
+	}
+	kfree(km, tq.a);
+}
+
+/***************************
+ * Suffix array operations *
+ ***************************/
+
+static inline uint64_t bwt_invPsi(const mb_bwt_t *bwt, uint64_t k) // compute inverse CSA
+{
+	uint64_t x = k - (k > bwt->primary);
+	int c = bwt_B0(bwt, x);
+	x = bwt->L2[c] + 1 + mb_bwt_rank11(bwt, k, c); // +1 to account for the sentinel
+	return k == bwt->primary? 0 : x;
+}
+
+// bwt->bwt and bwt->occ must be precalculated
+void mb_bwt_gen_sa(mb_bwt_t *bwt, uint32_t sa_bit)
+{
+	uint64_t isa, sa, i, mask; // S(isa) = sa
+
+	assert(bwt->data);
+	if (bwt->sa) free(bwt->sa);
+	bwt->sa_bit = sa_bit;
+	bwt->n_sa = (bwt->seq_len + (1<<sa_bit)) >> sa_bit;
+	bwt->sa = kom_calloc(uint64_t, bwt->n_sa);
+
+	// calculate SA value
+	isa = 0, sa = bwt->seq_len, mask = (1ULL<<sa_bit) - 1;
+	for (i = 0; i < bwt->seq_len; ++i) {
+		if ((isa & mask) == 0) bwt->sa[isa >> bwt->sa_bit] = sa;
+		--sa;
+		isa = bwt_invPsi(bwt, isa);
+	}
+	if ((isa & mask) == 0) bwt->sa[isa >> bwt->sa_bit] = sa;
+	bwt->sa[0] = (uint64_t)-1; // before this line, bwt->sa[0] = bwt->seq_len
+}
+
+uint64_t mb_bwt_sa(const mb_bwt_t *bwt, uint64_t k)
+{
+	uint64_t sa = 0, mask = (1ULL<<bwt->sa_bit) - 1;
+	while (k & mask) {
+		++sa;
+		k = bwt_invPsi(bwt, k);
+	}
+	// without setting bwt->sa[0] = -1, the following line should be
+	// changed to (sa + bwt->sa[k/bwt->sa_intv]) % (bwt->seq_len + 1)
+	return sa + bwt->sa[k >> bwt->sa_bit];
+}
+
+void mb_bwt_sa_batch(void *km, const mb_bwt_t *bwt, int64_t n, uint64_t *x)
+{
+	uint64_t mask = (1ULL<<bwt->sa_bit) - 1;
+	int64_t i, step = 0, r = n;
+	kom128_t *z;
+	if (n <= 0) return;
+	z = Kmalloc(km, kom128_t, n);
+	for (i = 0; i < n; ++i) {
+		z[i].x = x[i], z[i].y = i;
+		if ((z[i].x & mask) == 0)
+			__builtin_prefetch(&bwt->sa[z[i].x >> bwt->sa_bit]);
+		else
+			mb_bwt_block_prefetch(bwt, z[i].x);
+	}
+	for (step = 0; r > 0; ++step) {
+		int64_t r0 = r;
+		for (i = 0, r = 0; i < r0; ++i) {
+			if ((z[i].x & mask) == 0)
+				x[z[i].y] = step + bwt->sa[z[i].x >> bwt->sa_bit];
+			else z[r++] = z[i];
+		}
+		for (i = 0; i < r; ++i) {
+			z[i].x = bwt_invPsi(bwt, z[i].x);
+			if ((z[i].x & mask) == 0)
+				__builtin_prefetch(&bwt->sa[z[i].x >> bwt->sa_bit]);
+			else
+				mb_bwt_block_prefetch(bwt, z[i].x);
+		}
+	}
+	kfree(km, z);
+}
+
+/******************
+ * k-mer counting *
+ ******************/
+
+typedef struct {
+	mb_sai_t p;
+	int d, c;
+} count_pair64_t;
+
+void mb_bwt_count_kmer(const mb_bwt_t *bwt, int32_t depth, mb_sai_t *s) // adapted from kount in ropebwt3
+{
+	count_pair64_t *p, stack[64];
+	int32_t i, a, s_top = 0;
+	uint8_t str[16];
+	assert(depth <= 15);
+	for (a = 0; a < 4; ++a) {
+		p = &stack[s_top++];
+		mb_bwt_set_intv(bwt, a, &p->p);
+		p->d = 1, p->c = a;
+	}
+	while (s_top > 0) {
+		count_pair64_t top = stack[--s_top];
+		mb_sai_t ok[4];
+		if (top.d > 0) str[depth - top.d] = top.c;
+		mb_bwt_extend(bwt, &top.p, ok, 1);
+		for (a = 0; a < 4; ++a) {
+			str[depth - top.d - 1] = a;
+			if (top.d != depth - 1) {
+				p = &stack[s_top++];
+				p->p = ok[a];
+				p->d = top.d + 1;
+				p->c = a;
+			} else { // reaching the length; store in s[]
+				uint64_t x = 0;
+				for (i = 0; i < depth; ++i)
+					x |= (uint64_t)str[i] << i * 2;
+				s[x] = ok[a];
+			}
+		}
+	}
+}
+
+void mb_bwt_cache(mb_bwt_t *bwt, int32_t len)
+{
+	if (bwt->pre) free(bwt->pre);
+	bwt->pre_len = len;
+	bwt->pre = kom_calloc(mb_sai_t, 1 << len*2);
+	mb_bwt_count_kmer(bwt, len, bwt->pre);
+}
+
+/*************************
+ * Read/write BWT and SA *
+ *************************/
+
+static uint64_t read_huge(FILE *fp, uint64_t size, void *a)
+{ // Mac/Darwin has a bug when reading data longer than 2GB. This function fixes this issue by reading data in small chunks
+	const int bufsize = 0x1000000; // 16M block
+	uint64_t offset = 0;
+	while (size) {
+		int x = bufsize < size? bufsize : size;
+		if ((x = fread(a + offset, 1, x, fp)) == 0) break;
+		size -= x; offset += x;
+	}
+	return offset;
+}
+
+mb_bwt_t *mb_bwt_load_raw(const char *fn)
+{
+	mb_bwt_t *bwt;
+	uint32_t *raw;
+	uint64_t L2[5], primary, raw_size;
+	FILE *fp;
+
+	fp = fopen(fn, "rb");
+	fseek(fp, 0, SEEK_END);
+	raw_size = (ftell(fp) - sizeof(uint64_t) * 5) >> 2;
+	raw = kom_calloc(uint32_t, raw_size);
+	fseek(fp, 0, SEEK_SET);
+	fread(&primary, sizeof(uint64_t), 1, fp);
+	fread(L2 + 1, sizeof(uint64_t), 4, fp);
+	L2[0] = 0;
+	read_huge(fp, raw_size<<2, raw);
+	fclose(fp);
+	bwt = mb_bwt_init_from_raw(0, raw, L2[4], primary);
+	free(raw);
+	return bwt;
+}
+
+int mb_bwt_save(const char *fn, const mb_bwt_t *bwt)
+{
+	FILE *fp;
+	fp = fopen(fn, "wb");
+	if (fp == 0) return -1;
+	fwrite(MB_MAGIC, 1, 4, fp);
+	fwrite(&bwt->sa_bit, 4, 1, fp);
+	fwrite(&bwt->primary, 8, 1, fp);
+	fwrite(&bwt->L2[1], 8, 4, fp);
+	fwrite(bwt->data, 8, bwt->data_len, fp);
+	fwrite(&bwt->n_sa, 8, 1, fp);
+	if (bwt->sa_bit != (uint32_t)-1 && bwt->n_sa > 0 && bwt->sa)
+		fwrite(bwt->sa, 8, bwt->n_sa, fp);
+	fclose(fp);
+	return 0;
+}
+
+mb_bwt_t *mb_bwt_load(const char *fn)
+{
+	FILE *fp;
+	char magic[4];
+	uint64_t x[5];
+	mb_bwt_t *bwt;
+
+	fp = fopen(fn, "rb");
+	if (fp == 0) return 0;
+	fread(magic, 1, 4, fp);
+	if (strncmp(magic, MB_MAGIC, 4) != 0) {
+		fclose(fp);
+		return 0;
+	}
+	bwt = mb_bwt_init();
+	fread(&bwt->sa_bit, 4, 1, fp);
+	fread(x, 8, 5, fp);
+	bwt->primary = x[0];
+	memcpy(&bwt->L2[1], &x[1], 32);
+	bwt->seq_len = bwt->L2[4];
+	bwt->data_len = mb_bwt_data_len(bwt->seq_len);
+	bwt->data = kom_calloc(uint64_t, bwt->data_len);
+	read_huge(fp, bwt->data_len << 3, bwt->data);
+	fread(&bwt->n_sa, 8, 1, fp);
+	if (bwt->sa_bit != (uint32_t)-1 && bwt->n_sa > 0) {
+		uint64_t expected_n_sa = (bwt->seq_len + (1ULL << bwt->sa_bit)) >> bwt->sa_bit;
+		if (bwt->n_sa != expected_n_sa) {
+			mb_bwt_destroy(bwt);
+			return NULL;
+		}
+		bwt->sa = kom_malloc(uint64_t, bwt->n_sa);
+		fread(bwt->sa, 8, bwt->n_sa, fp);
+	}
+	fclose(fp);
+	return bwt;
+}
+
+mb_bwt_t *mb_bwt_load_mmap(const char *fn, int preload)
+{
+	uint8_t *base;
+	size_t map_len;
+	uint64_t data_off, sa_off, n_sa_off, min_len;
+	mb_bwt_t *bwt;
+
+	base = (uint8_t*)kom_mmap_file(fn, &map_len, preload);
+	if (base == 0) return 0;
+	if (map_len < 48 || strncmp((const char*)base, MB_MAGIC, 4) != 0) { kom_munmap(base, map_len); return 0; }
+
+	bwt = mb_bwt_init();
+	bwt->mmap = base;
+	bwt->mmap_len = map_len;
+	bwt->sa_bit = *(const uint32_t*)(base + 4);
+	bwt->primary = *(const uint64_t*)(base + 8);
+	memcpy(&bwt->L2[1], base + 16, 32);
+	bwt->seq_len = bwt->L2[4];
+	bwt->data_len = mb_bwt_data_len(bwt->seq_len);
+
+	data_off = 48;
+	n_sa_off = data_off + bwt->data_len * 8;
+	sa_off = n_sa_off + 8;
+	min_len = sa_off; // file must at least hold up to n_sa
+	if (map_len < min_len) { mb_bwt_destroy(bwt); return 0; }
+	bwt->data = (uint64_t*)(base + data_off);
+	bwt->n_sa = *(const uint64_t*)(base + n_sa_off);
+	if (bwt->sa_bit != (uint32_t)-1 && bwt->n_sa > 0) {
+		uint64_t expected_n_sa = (bwt->seq_len + (1ULL << bwt->sa_bit)) >> bwt->sa_bit;
+		if (bwt->n_sa != expected_n_sa || map_len < sa_off + bwt->n_sa * 8) {
+			mb_bwt_destroy(bwt);
+			return 0;
+		}
+		bwt->sa = (uint64_t*)(base + sa_off);
+	}
+	return bwt;
+}