fix bug where replaced nodes did not have the same level

README.rst

@@ -1,4 +1,4 @@
-Intrusive, single file, imperative AA tree that works on freestanding
+Intrusive, single file, C89+ imperative AA tree that works on freestanding
 systems.
 
 The implementation is about 300 lines.
@@ -38,7 +38,9 @@ own structure and search function like ::
 		return strcmp(n1->key, n2->key);
 	}
 
-This is standard, portable C99.
+This is standard, portable C89. Intrusive code has the advantage of
+requiring only one allocation instead of two, allowing for easier
+memory management (i.e. arena memory allocation).
 
 -----------------
 Using the Library

aatree.h

@@ -36,22 +36,56 @@ struct aatree {
 	unsigned level;
 };
 
+/**
+ * Function signature for node comparison.
+ *
+ * This functions returns
+ * * negative if the key of the first argument is less than the second,
+ * * 0 if the keys compare equal,
+ * * and a positive number if the first key is greater than the second.
+ */
 typedef int (*aatree_cmp_f)(struct aatree *, struct aatree *);
 
+/**
+ * Insert a node into the tree with root at ``root``. Returns the
+ * new root of the tree. If there was an element with the same key
+ * in the tree, the node is placed and the old element is returned
+ * into ``old``.
+ *
+ * Old may be NULL.
+ */
 AATREE_PREFIX struct aatree *aatree_ins(struct aatree *root,
                                         struct aatree *new_t,
                                         struct aatree **old,
                                         aatree_cmp_f cmp);
 
+/**
+ * Search the tree for a node that compares equal to "key". The
+ * "key" node does not need to have any of the values in "struct
+ * aatree" populated, but it does need to have a valid key.
+ *
+ * Returns the node, or NULL if not found.
+ */
 AATREE_PREFIX struct aatree *aatree_src(struct aatree *root,
                                         struct aatree *key,
                                         aatree_cmp_f cmp);
 
+/**
+ * Delete a node from the tree starting at "root" that compares equal
+ * to "key". See "aatree_src" documentation for information on how to
+ * fill the "key" node. The new root of the tree is returned.
+ *
+ * The deleted node is stored in "*match". Match must be a valid pointer
+ * because it is used in the delete algorithm.
+ */
 AATREE_PREFIX static struct aatree *aatree_del(struct aatree *root,
                                                struct aatree *key,
                                                struct aatree **match,
                                                aatree_cmp_f cmp);
 
+/**
+ * Initialize a node.
+ */
 static void aatree_init(struct aatree *t)
 {
 	t->left = t->right = t->parent = NULL;
@@ -87,9 +121,17 @@ static void aatree_replace(struct aatree *old, struct aatree *nw)
 {
 	aatree_relink(old, nw);
 	if (nw) {
-		nw->left = old->left;
+		nw->level = old->level;
-		nw->right = old->right;
+
 		nw->parent = old->parent;
+
+		nw->left = old->left;
+		if (nw->left)
+			nw->left->parent = nw;
+
+		nw->right = old->right;
+		if (nw->right)
+			nw->right->parent = nw;
 	}
 }
 
@@ -188,7 +230,8 @@ AATREE_PREFIX struct aatree *aatree_ins(struct aatree *root,
 			else
 				break;
 		} else {
-			*old = root;
+			if (old)
+				*old = root;
 			aatree_replace(root, new_t);
 			return orig_root;
 		}
@@ -262,7 +305,8 @@ AATREE_PREFIX struct aatree *aatree_del(struct aatree *root,
 		} else {
 			*match = root;
 			if (!root->right) {
-				repl = root->left;
+				assert(!root->left);
+				repl = NULL;
 			} else {
 				repl = root->right;
 				while (repl->left)
@@ -321,6 +365,7 @@ AATREE_PREFIX struct aatree *aatree_del(struct aatree *root,
 }
 #ifdef AATREE_TEST
 #include <stdio.h>
+#include <stdlib.h>
 #define NODE_NUM 64
 
 typedef unsigned long (*aatree_print_fun)(struct aatree *);
@@ -339,17 +384,13 @@ static void aatree_print(struct aatree *root, aatree_print_fun f)
 		f = def_print;
 
 	if (root->left) {
-		printf("%lu -> %lu;\n", f(root), f(root->left));
+		printf("\\\"%lu-%d\\\" -> \\\"%lu-%d\\\"; ", f(root), root->level, f(root->left), root->left->level);
 		aatree_print(root->left, f);
-	} else {
-		printf("%lu -> \"%lu-l\";\n", f(root), f(root));
 	}
 
 	if (root->right) {
-		printf("%lu -> %lu;\n", f(root), f(root->right));
+		printf("\\\"%lu-%d\\\" -> \\\"%lu-%d\\\"; ", f(root), root->level, f(root->right), root->right->level);
 		aatree_print(root->right, f);
-	} else {
-		printf("%lu -> \"%lu-r\";\n", f(root), f(root));
 	}
 }
 
@@ -420,6 +461,20 @@ static void test_multiple_insert(void)
 	}
 }
 
+static void test_single_delete(void)
+{
+	struct aatree *root = NULL, *match = NULL, *old;
+	struct test_node test_node;
+
+	init();
+
+	root = aatree_ins(root, &nodes[0].node, &old, test_cmp);
+
+	test_node.key = 0;
+	assert(aatree_del(root, &test_node.node, &match, test_cmp) == NULL);
+	assert(match == &nodes[0].node);
+}
+
 static void test_delete(void)
 {
 	struct aatree *root = NULL, *match = NULL, *old;
@@ -445,10 +500,6 @@ static void test_delete(void)
 
 	for (i = 0; i < NODE_NUM/2; i++) {
 		/*
-		printf("digraph t%d {", i);
-		aatree_print(root, tnp);
-		printf("}\n");
-		fflush(stdout);
 		*/
 
 		test_node.key = i;
@@ -461,14 +512,92 @@ static void test_delete(void)
 	}
 }
 
+static void shuf(int alen, int arr[])
+{
+	int i;
+	int repl;
+	int tmp;
+
+	for (i = 0; i < alen; i++) {
+		tmp = arr[i];
+		repl = rand() % alen;
+		arr[i] = arr[repl];
+		arr[repl] = tmp;
+	}
+}
+
+static void test_shuffle_insert_delete(void)
+{
+	struct aatree *root = NULL, *match = NULL, *old;
+	struct test_node test_node;
+	int i, j;
+	int order[NODE_NUM];
+
+	init();
+
+	for (i = 0; i < NODE_NUM; i++)
+		order[i] = i;
+
+	shuf(NODE_NUM, order);
+
+	for (i = 0; i < NODE_NUM; i++) {
+		root = aatree_ins(root, &nodes[order[i]].node, &old, test_cmp);
+		assert(!old);
+	}
+
+	for (i = 0; i < NODE_NUM; i++) {
+		test_node.key = i;
+		assert(aatree_src(root, &test_node.node, test_cmp) == &nodes[i].node);
+	}
+
+	shuf(NODE_NUM, order);
+
+	/*
+	printf("{\"t000\": \"");
+	aatree_print(root, tnp);
+	fflush(stdout);
+	printf("\"\n");
+	*/
+
+	for (i = 0; i < NODE_NUM; i++) {
+		test_node.key = order[i];
+		root = aatree_del(root, &test_node.node, &match, test_cmp);
+		assert(match == &nodes[order[i]].node);
+
+		/*
+		printf(", \"t%03d\": \"", i + 1);
+		aatree_print(root, tnp);
+		fflush(stdout);
+		printf("\"\n");
+		*/
+
+		for (j = 0; j <= i; j++) {
+			test_node.key = order[j];
+			assert(!aatree_src(root, &test_node.node, test_cmp));
+		}
+
+		for (j = i+1; j < NODE_NUM; j++) {
+			fflush(stdout);
+
+			test_node.key = order[j];
+			assert(aatree_src(root, &test_node.node, test_cmp) == &nodes[order[j]].node);
+		}
+	}
+
+	assert(!root);
+}
+
 int main(void)
 {
 	test_single_insert();
 	test_multiple_insert();
+	test_single_delete();
 	test_delete();
+	test_shuffle_insert_delete();
 
 	return 0;
 }
+
 #endif /* AATREE_TEST */
 #endif /* AATREE_IMPL */
 #endif /* AATREE_H */