319 lines
8.8 KiB
C
319 lines
8.8 KiB
C
|
#include <stdlib.h>
|
||
|
|
#include <string.h>
|
||
|
|
#include <assert.h>
|
||
|
|
#include <stdio.h>
|
||
|
|
#include <zlib.h>
|
||
|
|
#include "rle.h"
|
||
|
|
#include "rope.h"
|
||
|
|
|
||
|
|
/*******************
|
||
|
|
*** Memory Pool ***
|
||
|
|
*******************/
|
||
|
|
|
||
|
|
#define MP_CHUNK_SIZE 0x100000 // 1MB per chunk
|
||
|
|
|
||
|
|
typedef struct { // memory pool for fast and compact memory allocation (no free)
|
||
|
|
int size, i, n_elems;
|
||
|
|
int64_t top, max;
|
||
|
|
uint8_t **mem;
|
||
|
|
} mempool_t;
|
||
|
|
|
||
|
|
static mempool_t *mp_init(int size)
|
||
|
|
{
|
||
|
|
mempool_t *mp;
|
||
|
|
mp = calloc(1, sizeof(mempool_t));
|
||
|
|
mp->size = size;
|
||
|
|
mp->i = mp->n_elems = MP_CHUNK_SIZE / size;
|
||
|
|
mp->top = -1;
|
||
|
|
return mp;
|
||
|
|
}
|
||
|
|
|
||
|
|
static void mp_destroy(mempool_t *mp)
|
||
|
|
{
|
||
|
|
int64_t i;
|
||
|
|
for (i = 0; i <= mp->top; ++i) free(mp->mem[i]);
|
||
|
|
free(mp->mem); free(mp);
|
||
|
|
}
|
||
|
|
|
||
|
|
static inline void *mp_alloc(mempool_t *mp)
|
||
|
|
{
|
||
|
|
if (mp->i == mp->n_elems) {
|
||
|
|
if (++mp->top == mp->max) {
|
||
|
|
mp->max = mp->max? mp->max<<1 : 1;
|
||
|
|
mp->mem = realloc(mp->mem, sizeof(void*) * mp->max);
|
||
|
|
}
|
||
|
|
mp->mem[mp->top] = calloc(mp->n_elems, mp->size);
|
||
|
|
mp->i = 0;
|
||
|
|
}
|
||
|
|
return mp->mem[mp->top] + (mp->i++) * mp->size;
|
||
|
|
}
|
||
|
|
|
||
|
|
/***************
|
||
|
|
*** B+ rope ***
|
||
|
|
***************/
|
||
|
|
|
||
|
|
rope_t *rope_init(int max_nodes, int block_len)
|
||
|
|
{
|
||
|
|
rope_t *rope;
|
||
|
|
rope = calloc(1, sizeof(rope_t));
|
||
|
|
if (block_len < 32) block_len = 32;
|
||
|
|
rope->max_nodes = (max_nodes+ 1)>>1<<1;
|
||
|
|
rope->block_len = (block_len + 7) >> 3 << 3;
|
||
|
|
rope->node = mp_init(sizeof(rpnode_t) * rope->max_nodes);
|
||
|
|
rope->leaf = mp_init(rope->block_len);
|
||
|
|
rope->root = mp_alloc(rope->node);
|
||
|
|
rope->root->n = 1;
|
||
|
|
rope->root->is_bottom = 1;
|
||
|
|
rope->root->p = mp_alloc(rope->leaf);
|
||
|
|
return rope;
|
||
|
|
}
|
||
|
|
|
||
|
|
void rope_destroy(rope_t *rope)
|
||
|
|
{
|
||
|
|
mp_destroy(rope->node);
|
||
|
|
mp_destroy(rope->leaf);
|
||
|
|
free(rope);
|
||
|
|
}
|
||
|
|
|
||
|
|
static inline rpnode_t *split_node(rope_t *rope, rpnode_t *u, rpnode_t *v)
|
||
|
|
{ // split $v's child. $u is the first node in the bucket. $v and $u are in the same bucket. IMPORTANT: there is always enough room in $u
|
||
|
|
int j, i = v - u;
|
||
|
|
rpnode_t *w; // $w is the sibling of $v
|
||
|
|
if (u == 0) { // only happens at the root; add a new root
|
||
|
|
u = v = mp_alloc(rope->node);
|
||
|
|
v->n = 1; v->p = rope->root; // the new root has the old root as the only child
|
||
|
|
memcpy(v->c, rope->c, 48);
|
||
|
|
for (j = 0; j < 6; ++j) v->l += v->c[j];
|
||
|
|
rope->root = v;
|
||
|
|
}
|
||
|
|
if (i != u->n - 1) // then make room for a new node
|
||
|
|
memmove(v + 2, v + 1, sizeof(rpnode_t) * (u->n - i - 1));
|
||
|
|
++u->n; w = v + 1;
|
||
|
|
memset(w, 0, sizeof(rpnode_t));
|
||
|
|
w->p = mp_alloc(u->is_bottom? rope->leaf : rope->node);
|
||
|
|
if (u->is_bottom) { // we are at the bottom level; $v->p is a string instead of a node
|
||
|
|
uint8_t *p = (uint8_t*)v->p, *q = (uint8_t*)w->p;
|
||
|
|
rle_split(p, q);
|
||
|
|
rle_count(q, w->c);
|
||
|
|
} else { // $v->p is a node, not a string
|
||
|
|
rpnode_t *p = v->p, *q = w->p; // $v and $w are siblings and thus $p and $q are cousins
|
||
|
|
p->n -= rope->max_nodes>>1;
|
||
|
|
memcpy(q, p + p->n, sizeof(rpnode_t) * (rope->max_nodes>>1));
|
||
|
|
q->n = rope->max_nodes>>1; // NB: this line must below memcpy() as $q->n and $q->is_bottom are modified by memcpy()
|
||
|
|
q->is_bottom = p->is_bottom;
|
||
|
|
for (i = 0; i < q->n; ++i)
|
||
|
|
for (j = 0; j < 6; ++j)
|
||
|
|
w->c[j] += q[i].c[j];
|
||
|
|
}
|
||
|
|
for (j = 0; j < 6; ++j) // compute $w->l and update $v->c
|
||
|
|
w->l += w->c[j], v->c[j] -= w->c[j];
|
||
|
|
v->l -= w->l; // update $v->c
|
||
|
|
return v;
|
||
|
|
}
|
||
|
|
|
||
|
|
int64_t rope_insert_run(rope_t *rope, int64_t x, int a, int64_t rl, rpcache_t *cache)
|
||
|
|
{ // insert $a after $x symbols in $rope and the returns rank(a, x)
|
||
|
|
rpnode_t *u = 0, *v = 0, *p = rope->root; // $v is the parent of $p; $u and $v are at the same level and $u is the first node in the bucket
|
||
|
|
int64_t y = 0, z = 0, cnt[6];
|
||
|
|
int n_runs;
|
||
|
|
do { // top-down update. Searching and node splitting are done together in one pass.
|
||
|
|
if (p->n == rope->max_nodes) { // node is full; split
|
||
|
|
v = split_node(rope, u, v); // $v points to the parent of $p; when a new root is added, $v points to the root
|
||
|
|
if (y + v->l < x) // if $v is not long enough after the split, we need to move both $p and its parent $v
|
||
|
|
y += v->l, z += v->c[a], ++v, p = v->p;
|
||
|
|
}
|
||
|
|
u = p;
|
||
|
|
if (v && x - y > v->l>>1) { // then search backwardly for the right node to descend
|
||
|
|
p += p->n - 1; y += v->l; z += v->c[a];
|
||
|
|
for (; y >= x; --p) y -= p->l, z -= p->c[a];
|
||
|
|
++p;
|
||
|
|
} else for (; y + p->l < x; ++p) y += p->l, z += p->c[a]; // then search forwardly
|
||
|
|
assert(p - u < u->n);
|
||
|
|
if (v) v->c[a] += rl, v->l += rl; // we should not change p->c[a] because this may cause troubles when p's child is split
|
||
|
|
v = p; p = p->p; // descend
|
||
|
|
} while (!u->is_bottom);
|
||
|
|
rope->c[a] += rl; // $rope->c should be updated after the loop as adding a new root needs the old $rope->c counts
|
||
|
|
if (cache) {
|
||
|
|
if (cache->p != (uint8_t*)p) memset(cache, 0, sizeof(rpcache_t));
|
||
|
|
n_runs = rle_insert_cached((uint8_t*)p, x - y, a, rl, cnt, v->c, &cache->beg, cache->bc);
|
||
|
|
cache->p = (uint8_t*)p;
|
||
|
|
} else n_runs = rle_insert((uint8_t*)p, x - y, a, rl, cnt, v->c);
|
||
|
|
z += cnt[a];
|
||
|
|
v->c[a] += rl; v->l += rl; // this should be after rle_insert(); otherwise rle_insert() won't work
|
||
|
|
if (n_runs + RLE_MIN_SPACE > rope->block_len) {
|
||
|
|
split_node(rope, u, v);
|
||
|
|
if (cache) memset(cache, 0, sizeof(rpcache_t));
|
||
|
|
}
|
||
|
|
return z;
|
||
|
|
}
|
||
|
|
|
||
|
|
static rpnode_t *rope_count_to_leaf(const rope_t *rope, int64_t x, int64_t cx[6], int64_t *rest)
|
||
|
|
{
|
||
|
|
rpnode_t *u, *v = 0, *p = rope->root;
|
||
|
|
int64_t y = 0;
|
||
|
|
int a;
|
||
|
|
|
||
|
|
memset(cx, 0, 48);
|
||
|
|
do {
|
||
|
|
u = p;
|
||
|
|
if (v && x - y > v->l>>1) {
|
||
|
|
p += p->n - 1; y += v->l;
|
||
|
|
for (a = 0; a != 6; ++a) cx[a] += v->c[a];
|
||
|
|
for (; y >= x; --p) {
|
||
|
|
y -= p->l;
|
||
|
|
for (a = 0; a != 6; ++a) cx[a] -= p->c[a];
|
||
|
|
}
|
||
|
|
++p;
|
||
|
|
} else {
|
||
|
|
for (; y + p->l < x; ++p) {
|
||
|
|
y += p->l;
|
||
|
|
for (a = 0; a != 6; ++a) cx[a] += p->c[a];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
v = p; p = p->p;
|
||
|
|
} while (!u->is_bottom);
|
||
|
|
*rest = x - y;
|
||
|
|
return v;
|
||
|
|
}
|
||
|
|
|
||
|
|
void rope_rank2a(const rope_t *rope, int64_t x, int64_t y, int64_t *cx, int64_t *cy)
|
||
|
|
{
|
||
|
|
rpnode_t *v;
|
||
|
|
int64_t rest;
|
||
|
|
v = rope_count_to_leaf(rope, x, cx, &rest);
|
||
|
|
if (y < x || cy == 0) {
|
||
|
|
rle_rank1a((const uint8_t*)v->p, rest, cx, v->c);
|
||
|
|
} else if (rest + (y - x) <= v->l) {
|
||
|
|
memcpy(cy, cx, 48);
|
||
|
|
rle_rank2a((const uint8_t*)v->p, rest, rest + (y - x), cx, cy, v->c);
|
||
|
|
} else {
|
||
|
|
rle_rank1a((const uint8_t*)v->p, rest, cx, v->c);
|
||
|
|
v = rope_count_to_leaf(rope, y, cy, &rest);
|
||
|
|
rle_rank1a((const uint8_t*)v->p, rest, cy, v->c);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/*********************
|
||
|
|
*** Rope iterator ***
|
||
|
|
*********************/
|
||
|
|
|
||
|
|
void rope_itr_first(const rope_t *rope, rpitr_t *i)
|
||
|
|
{
|
||
|
|
memset(i, 0, sizeof(rpitr_t));
|
||
|
|
i->rope = rope;
|
||
|
|
for (i->pa[i->d] = rope->root; !i->pa[i->d]->is_bottom;) // descend to the leftmost leaf
|
||
|
|
++i->d, i->pa[i->d] = i->pa[i->d - 1]->p;
|
||
|
|
}
|
||
|
|
|
||
|
|
const uint8_t *rope_itr_next_block(rpitr_t *i)
|
||
|
|
{
|
||
|
|
const uint8_t *ret;
|
||
|
|
assert(i->d < ROPE_MAX_DEPTH); // a B+ tree should not be that tall
|
||
|
|
if (i->d < 0) return 0;
|
||
|
|
ret = (uint8_t*)i->pa[i->d][i->ia[i->d]].p;
|
||
|
|
while (i->d >= 0 && ++i->ia[i->d] == i->pa[i->d]->n) i->ia[i->d--] = 0; // backtracking
|
||
|
|
if (i->d >= 0)
|
||
|
|
while (!i->pa[i->d]->is_bottom) // descend to the leftmost leaf
|
||
|
|
++i->d, i->pa[i->d] = i->pa[i->d - 1][i->ia[i->d - 1]].p;
|
||
|
|
return ret;
|
||
|
|
}
|
||
|
|
|
||
|
|
/***********
|
||
|
|
*** I/O ***
|
||
|
|
***********/
|
||
|
|
|
||
|
|
void rope_print_node(const rpnode_t *p)
|
||
|
|
{
|
||
|
|
if (p->is_bottom) {
|
||
|
|
int i;
|
||
|
|
putchar('(');
|
||
|
|
for (i = 0; i < p->n; ++i) {
|
||
|
|
uint8_t *block = (uint8_t*)p[i].p;
|
||
|
|
const uint8_t *q = block + 2, *end = block + 2 + *rle_nptr(block);
|
||
|
|
if (i) putchar(',');
|
||
|
|
while (q < end) {
|
||
|
|
int c = 0;
|
||
|
|
int64_t j, l;
|
||
|
|
rle_dec1(q, c, l);
|
||
|
|
for (j = 0; j < l; ++j) putchar("$ACGTN"[c]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
putchar(')');
|
||
|
|
} else {
|
||
|
|
int i;
|
||
|
|
putchar('(');
|
||
|
|
for (i = 0; i < p->n; ++i) {
|
||
|
|
if (i) putchar(',');
|
||
|
|
rope_print_node(p[i].p);
|
||
|
|
}
|
||
|
|
putchar(')');
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void rope_dump_node(const rpnode_t *p, FILE *fp)
|
||
|
|
{
|
||
|
|
int16_t i, n = p->n;
|
||
|
|
uint8_t is_bottom = p->is_bottom;
|
||
|
|
fwrite(&is_bottom, 1, 1, fp);
|
||
|
|
fwrite(&n, 2, 1, fp);
|
||
|
|
if (is_bottom) {
|
||
|
|
for (i = 0; i < n; ++i) {
|
||
|
|
fwrite(p[i].c, 8, 6, fp);
|
||
|
|
fwrite(p[i].p, 1, *rle_nptr(p[i].p) + 2, fp);
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
for (i = 0; i < p->n; ++i)
|
||
|
|
rope_dump_node(p[i].p, fp);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void rope_dump(const rope_t *r, FILE *fp)
|
||
|
|
{
|
||
|
|
fwrite(&r->max_nodes, 4, 1, fp);
|
||
|
|
fwrite(&r->block_len, 4, 1, fp);
|
||
|
|
rope_dump_node(r->root, fp);
|
||
|
|
}
|
||
|
|
|
||
|
|
rpnode_t *rope_restore_node(const rope_t *r, FILE *fp, int64_t c[6])
|
||
|
|
{
|
||
|
|
uint8_t is_bottom, a;
|
||
|
|
int16_t i, n;
|
||
|
|
rpnode_t *p;
|
||
|
|
fread(&is_bottom, 1, 1, fp);
|
||
|
|
fread(&n, 2, 1, fp);
|
||
|
|
p = mp_alloc(r->node);
|
||
|
|
p->is_bottom = is_bottom, p->n = n;
|
||
|
|
if (is_bottom) {
|
||
|
|
for (i = 0; i < n; ++i) {
|
||
|
|
uint16_t *q;
|
||
|
|
p[i].p = mp_alloc(r->leaf);
|
||
|
|
q = rle_nptr(p[i].p);
|
||
|
|
fread(p[i].c, 8, 6, fp);
|
||
|
|
fread(q, 2, 1, fp);
|
||
|
|
fread(q + 1, 1, *q, fp);
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
for (i = 0; i < n; ++i)
|
||
|
|
p[i].p = rope_restore_node(r, fp, p[i].c);
|
||
|
|
}
|
||
|
|
memset(c, 0, 48);
|
||
|
|
for (i = 0; i < n; ++i) {
|
||
|
|
p[i].l = 0;
|
||
|
|
for (a = 0; a < 6; ++a)
|
||
|
|
c[a] += p[i].c[a], p[i].l += p[i].c[a];
|
||
|
|
}
|
||
|
|
return p;
|
||
|
|
}
|
||
|
|
|
||
|
|
rope_t *rope_restore(FILE *fp)
|
||
|
|
{
|
||
|
|
rope_t *r;
|
||
|
|
r = calloc(1, sizeof(rope_t));
|
||
|
|
fread(&r->max_nodes, 4, 1, fp);
|
||
|
|
fread(&r->block_len, 4, 1, fp);
|
||
|
|
r->node = mp_init(sizeof(rpnode_t) * r->max_nodes);
|
||
|
|
r->leaf = mp_init(r->block_len);
|
||
|
|
r->root = rope_restore_node(r, fp, r->c);
|
||
|
|
return r;
|
||
|
|
}
|