diff --git a/treetest.c b/treetest.c new file mode 100644 index 0000000..f4396c3 --- /dev/null +++ b/treetest.c @@ -0,0 +1,147 @@ +#include +#include +#include + +#define TREE_ERR 0 +#define TREE_OK 1 + +/* Potentially use a custom allocator on a tiny MCU */ +static struct { + void *(*node_alloc)(size_t); + void (*node_free)(void *); +} node_mm; + +typedef struct node { + struct node *left, *right; + int data; +} node_t; + +typedef struct { + struct node *root; +} tree_t; + +node_t *node_new(int data) { + node_t *node = node_mm.node_alloc(sizeof(node_t)); + if (!node) + return TREE_ERR; + + node->data = data; + node->left = node->right = NULL; + return node; +} + +int node_insert(node_t **n, int data) { + if (!*n) { + *n = node_new(data); + return TREE_OK; + } + + if ((*n)->data == data) + return TREE_ERR; // Already exists + + return node_insert(data < (*n)->data ? &(*n)->left : &(*n)->right, data); +} + +int tree_insert(tree_t *t, int data) { + return node_insert(&t->root, data); +} + +int tree_count(node_t *node) { + return node ? 1 + tree_count(node->left) + tree_count(node->right) : 0; +} + +int node_free(node_t **n) { + if (!*n) + return TREE_OK; + + node_free(&(*n)->left); + node_free(&(*n)->right); + + node_mm.node_free((*n)); + (*n) = NULL; + return TREE_OK; +} + +int node_erase(node_t **n, int data) { + if (!*n) + return TREE_ERR; // Not found + + if ((*n)->data < data) + return node_erase(&(*n)->left, data); + if ((*n)->data > data) + return node_erase(&(*n)->right, data); + + node_t *victim = *n; + if (victim->data != data) { + return TREE_ERR; + } + + if (!victim->left) { + *n = victim->right; + } else if (!victim->right) { + *n = victim->right; + } else { + /* This is ugly, seatbelt on */ + node_t **succ = &victim->right; + while ((*succ)->left) succ = &(*succ)->left; + + victim->data = (*succ)->data; + node_t *tmp = *succ; + *succ = (*succ)->right; + free(tmp); + return TREE_OK; + } + + node_mm.node_free(victim); + return TREE_OK; +} + +int tree_erase(tree_t *t, int data) { + return node_erase(&t->root, data); +} + +int tree_free(tree_t *t) { + return node_free(&t->root); +} + +int main(int argc, char *argv[]) { + (void)argc; + (void)argv; + + printf("Sizeof(node_t): %lu\n", sizeof(node_t)); + + node_mm.node_free = free; + node_mm.node_alloc = malloc; + + /* Check so this works */ + node_t *n = node_new(10); + assert(n); + + /* Insert some data and count it */ + tree_t t = {}; + assert(tree_insert(&t, 10)); + assert(tree_insert(&t, 12)); + assert(tree_insert(&t, 10) == TREE_ERR); + assert(tree_count(t.root) == 2); + assert(t.root); + + /* Free it all */ + tree_free(&t); + assert(t.root == NULL); + assert(tree_count(t.root) == 0); + + /* Re-insert some stuff and check */ + assert(tree_insert(&t, 120)); + assert(tree_count(t.root) == 1); + + assert(tree_erase(&t, 120) == TREE_OK); // Exists + assert(tree_erase(&t, 121) == TREE_ERR); // Does not exist + assert(tree_count(t.root) == 0); + assert(tree_insert(&t, 122)); + assert(tree_count(t.root) == 1); + assert(tree_erase(&t, 122) == TREE_OK); + assert(tree_count(t.root) == 0); + + printf("All tests passed!\n"); + return EXIT_SUCCESS; +}