strans/trie.c

#include <errno.h>
#include "dat.h"
#include "fn.h"

/*
 * A trie over the runes of a key.  Maps need prefix matches while
 * composing; dictionaries need exact matches; both are small enough for
 * one structure.  Node 0 is the root.
 */

static int
newnode(Trie *t)
{
	int i;

	if(t->n == t->cap){
		t->cap *= 2;
		t->nodes = erealloc(t->nodes, t->cap * sizeof(Tnode));
	}
	i = t->n++;
	memset(&t->nodes[i], 0, sizeof(Tnode));
	t->nodes[i].child = -1;
	t->nodes[i].sibling = -1;
	return i;
}

/* The children of a node are kept in rune order, so both end early. */
static int
find(Trie *t, int ni, Rune c)
{
	int pi;

	for(pi = t->nodes[ni].child; pi >= 0; pi = t->nodes[pi].sibling){
		if(t->nodes[pi].c == c)
			return pi;
		if(t->nodes[pi].c > c)
			break;
	}
	return -1;
}

static int
add(Trie *t, int ni, Rune c)
{
	int next, pi, prev;

	prev = -1;
	for(next = t->nodes[ni].child; next >= 0 && t->nodes[next].c < c;
	    next = t->nodes[next].sibling)
		prev = next;
	pi = newnode(t);
	t->nodes[pi].c = c;
	t->nodes[pi].sibling = next;
	if(prev < 0)
		t->nodes[ni].child = pi;
	else
		t->nodes[prev].sibling = pi;
	return pi;
}

Trie*
trienew(void)
{
	Trie *t;

	t = emalloc(sizeof(*t));
	t->cap = 1024;
	t->nodes = emalloc(t->cap * sizeof(Tnode));
	t->n = 0;
	newnode(t);
	return t;
}

void
trieput(Trie *t, char *key, int klen, char *val, int vlen)
{
	Str k;
	int ci, i, ni, same;

	if(!sinit(&k, key, klen))
		return;
	ni = 0;
	same = 1;
	for(i = 0; i < k.n; i++){
		if(same && i < t->last.n && t->last.r[i] == k.r[i]){
			ni = t->path[i];
			continue;
		}
		same = 0;
		ci = find(t, ni, k.r[i]);
		if(ci < 0)
			ci = add(t, ni, k.r[i]);
		ni = ci;
		t->path[i] = ni;
	}
	t->last = k;
	free(t->nodes[ni].val);
	t->nodes[ni].val = emalloc(vlen + 1);
	memmove(t->nodes[ni].val, val, vlen);
	t->nodes[ni].val[vlen] = '\0';
	t->nodes[ni].vlen = vlen;
}

static char*
readline(Biobuf *b, char *path)
{
	char *line;

	errno = 0;
	line = Brdstr(b, '\n', 1);
	if(errno != 0)
		die("can't read %s: %r", path);
	return line;
}

/*
 * A file holds "key<tab>value" lines; ';' starts a comment.  Keys and each
 * space-separated word of a value must fit a Str, which is what lookups
 * hand back.
 */
Trie*
trieopen(char *path)
{
	Trie *t;
	Biobuf *b;
	Str s;
	char *e, *line, *p, *tab;
	int len, lineno;

	b = Bopen(path, OREAD);
	if(b == nil)
		die("can't open %s: %r", path);
	t = trienew();
	for(lineno = 1; (line = readline(b, path)) != nil; lineno++){
		len = Blinelen(b);
		if(memchr(line, '\0', len) != nil)
			die("%s:%d: NUL byte", path, lineno);
		if(len > 0 && line[len-1] == '\r')
			line[--len] = '\0';
		if(len == 0 || line[0] == ';'){
			free(line);
			continue;
		}
		tab = memchr(line, '\t', len);
		if(tab == nil || tab == line || tab == line+len-1 ||
		   memchr(tab+1, '\t', line+len-(tab+1)) != nil)
			die("%s:%d: expected key, tab, value", path, lineno);
		if(!sinit(&s, line, tab-line))
			die("%s:%d: invalid or oversized key", path, lineno);
		for(p = tab+1; p < line+len; p = e+1){
			e = memchr(p, ' ', line+len-p);
			if(e == nil)
				e = line+len;
			if(!sinit(&s, p, e-p))
				die("%s:%d: invalid or oversized value", path, lineno);
		}
		trieput(t, line, tab-line, tab+1, line+len-(tab+1));
		free(line);
	}
	if(Bterm(b) < 0)
		die("can't close %s: %r", path);
	return t;
}

void
trieclose(Trie *t)
{
	int i;

	if(t == nil)
		return;
	for(i = 0; i < t->n; i++)
		free(t->nodes[i].val);
	free(t->nodes);
	free(t);
}

/* The node key leads to, or -1.  A nil trie is an unloaded map: every key misses. */
int
trienode(Trie *t, Str *key)
{
	int i, ni;

	if(t == nil)
		return -1;
	ni = 0;
	for(i = 0; i < key->n && ni >= 0; i++)
		ni = find(t, ni, key->r[i]);
	return ni;
}

int
trielookup(Trie *t, Str *key, char **val, int *vlen)
{
	int ni;

	*val = nil;
	*vlen = 0;
	ni = trienode(t, key);
	if(ni < 0)
		return TrieMiss;
	if(t->nodes[ni].val == nil)
		return TriePrefix;
	*val = t->nodes[ni].val;
	*vlen = t->nodes[ni].vlen;
	return TrieExact;
}