/* Simple trie implementation
 * Copyright (C) 2008 Daniel Collins <solemnwarning@solemnwarning.net>
 * Version 1.0 (2008-03-18)
 *
 * I wrote this as a why-not thing, it's probably got a bug somewhere, if you
 * find/fix a bug please email me anyway.
 *
 * This is released to public domain, you may do anything you please with this
 * code, but if you submit it as your own in a class.... $DEITY will know!
*/

#include <stdlib.h>
#include <string.h>

#include "trie.h"

/* Allocate a new node and set all node/value pointers to NULL.
 *
 * Returns NULL if malloc() fails
*/
static trie_node_t *trie_alloc(void) {
	trie_node_t *node;
	int n;
	
	if(!(node = malloc(sizeof(trie_node_t)*256))) {
		return NULL;
	}
	
	for(n = 0; n < 256; n++) {
		node[n].node = NULL;
		node[n].value = NULL;
	}
	
	return node;
}

/* Create a new trie
 * Returns a trie* pointer on success, NULL on failure
*/
trie_t *trie_init(void) {
	trie_t *trie;
	
	if(!(trie = malloc(sizeof(trie_t)))) {
		return NULL;
	}
	if(!(trie->root = trie_alloc())) {
		free(trie);
		return NULL;
	}
	
	return trie;
}

/* Free a trie and all nodes from memory
 * Passing NULL is a no-op, same as free()
*/
void trie_free(trie_t *trie) {
	if(trie == NULL) {
		return;
	}
	
	trie_node_t *node = trie->root, *nnode;
	int n = 0;
	
	while(1) {
		if(n == 256) {
			free(node);
			if(node == trie->root) {
				break;
			}else{
				node = trie->root;
			}
			
			n = 0;
			continue;
		}
		
		if(node[n].node != NULL) {
			nnode = node[n].node;
			node[n].node = NULL;
			node = nnode;
			
			n = 0;
			continue;
		}
		
		n++;
	}
	
	free(trie);
}

/* Set a value in a trie
 * Returns 1 on success, zero on malloc() failure or zero-length key
*/
int trie_set(trie_t *trie, char *key, size_t klen, void *value) {
	size_t kpos = 0;
	trie_node_t *node = trie->root;
	unsigned char ukey;
	
	if(klen == 0) {
		return 0;
	}
	
	while((kpos+1) < klen) {
		ukey = key[kpos];
		if(node[ukey].node == NULL) {
			if(!(node[ukey].node = trie_alloc())) {
				return 0;
			}
			
			trie->size += (sizeof(trie_node_t)*256);
		}
		
		node = node[ukey].node;
		kpos++;
	}
	
	ukey = key[kpos];
	node[ukey].value = value;
	
	return 1;
}

/* Simple wrapper to trie_t_set() which takes a nil-terminated string as the
 * key, the resulting key includes the nil terminating character
*/
int trie_set_str(trie_t *trie, char *key, void *value) {
	return trie_set(trie, key, strlen(key)+1, value);
}

/* Search the supplied trie for a value
 * Returns the value if found, otherwise returns NULL
*/
void *trie_get(trie_t *trie, char *key, size_t klen) {
	size_t kpos = 0;
	trie_node_t *node = trie->root;
	unsigned char ukey;
	
	if(klen == 0) {
		return NULL;
	}
	
	while((kpos+1) < klen) {
		ukey = key[kpos];
		if(node[ukey].node == NULL) {
			return NULL;
		}
		
		node = node[ukey].node;
		kpos++;
	}
	
	ukey = key[kpos];
	return node[ukey].value;
}

/* Wrapper to trie_t_get(), using a nil-terminated string as the key, the
 * resulting key includes the nil terminating character
*/
void *trie_get_str(trie_t *trie, char *key) {
	return trie_get(trie, key, strlen(key)+1);
}