shithub: ircs

ref: a90d1130f47c07132f1816366b05f18929911e60
dir: /trie.c/

#include <u.h>
#include <libc.h>
#include "dat.h"

int trieallocs;
int triedebug;

#define dprint	if(triedebug) print

Trie *
triealloc(void)
{
	Trie *t;

	t = mallocz(sizeof(Trie), 1);
	if(t == nil)
		sysfatal("mallocz: %r");
	trieallocs++;
	return t;
}

void
triefree(Trie *t)
{
	Trie *c, *n;

	for(c = t->child; c != nil; c = n){
		n = c->next;
		triefree(c);
	}
	free(t);
	trieallocs--;
}

void
triewalk(Trie *t, Triewalk *w, Rune *key, int keysize)
{
	w->root = t;
	w->curr = nil;
	w->key = key;
	w->keysize = keysize;
	w->depth = 0;
}

void *
trienext(Triewalk *w)
{
	Trie *t;
	int d;

	if(w->curr == nil){
		w->curr = w->root;
		w->depth = 0;
		w->key[0] = 0;
		if(w->root->value != nil)
			return w->root->value;
	}
	t = w->curr;
	d = w->depth;
	do{
		if(t->child != nil && d < w->keysize-1){
			t = t->child;
			d++;
		}
		else if(t->next != nil)
			t = t->next;
		else {
			while(t->next == nil && t != w->root){
				t = t->parent;
				d--;
			}
			if(t->next != nil)
				t = t->next;
			else {
				assert(t == w->root);
				w->curr = nil;
				return nil;
			}				
		}
		w->key[d-1] = t->rune;
	} while(t->value == nil);

	w->key[d] = 0;
	w->curr = t;
	w->depth = d;
	return t->value;
}

static Trie *
getchild(Trie *t, Rune r)
{
	dprint("  getchild %C: ", r);

	for(t = t->child; t != nil; t = t->next)
		if(t->rune == r)
			break;

	dprint("%s\n", t != nil ? "found" : "not found");
	return t;
}

static Trie *
addchild(Trie *t, Rune r)
{
	Trie *c, *n, *prev;

	dprint("  addchild %C\n", r);

	c = triealloc();
	c->rune = r;
	c->parent = t;

	prev = nil;
	for(n = t->child; n != nil; n = n->next){
		if((prev == nil || prev->rune < r) && r < n->rune)
			break;
		assert(r > n->rune);
		prev = n;
	}
	if(prev == nil)
		t->child = c;
	else
		prev->next = c;
	c->next = n;
	return c;
}

static int
delchild(Trie *t, Rune r)
{
	Trie *c, *prev;

	dprint("  delchild %C\n", r);

	prev = nil;
	for(c = t->child; c != nil; c = c->next){
		if(c->rune == r)
			break;
		prev = c;
	}
	if(c == nil)
		return 0;
	if(c->child != nil)
		return 0;
	if(prev == nil)
		t->child = c->next;
	else
		prev->next = c->next;
	triefree(c);
	return 1;
}

void *
trieadd(Trie *t, Rune *key, void *val)
{
	Trie *c;
	void *v;

	if(val == nil)
		return nil;

	while(*key){
		c = getchild(t, *key);
		if(c == nil)
			break;
		t = c;
		key++;
	}
	while(*key)
		t = addchild(t, *key++);

	v = t->value;
	t->value = val;
	return v;	
}

static Trie *
find(Trie *t, Rune *key)
{
	Trie *c;

	while(*key){
		c = getchild(t, *key);
		if(c == nil)
			break;
		t = c;
		key++;
	}
	if(*key)
		return nil;
	return t;
}

void *
triedel(Trie *t, Rune *key)
{
	Rune *r;
	void *v;

	t = find(t, key);
	if(t == nil)
		return nil;

	/* only non-leaves and root can have empty value */
	if(t->value == nil){
		assert(t->child != nil || t->parent == nil);
		return nil;
	}
	v = t->value;
	t->value = nil;

	if(t->child != nil || t->parent == nil)
		return v;

	r = runestrchr(key, 0);

	while(t->child == nil && t->value == nil && --r >= key){
		t = t->parent;
		assert(t != nil);
		assert(delchild(t, *r));
	}
	return v;
}

void *
trieget(Trie *t, Rune *key)
{
	t = find(t, key);
	if(t == nil)
		return nil;
	return t->value;	
}