Treetest tree exercise

treetest.c

@@ -0,0 +1,147 @@
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#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;
+}