Add a copy function to the rb-tree implementation

If we use the rb-tree in libsquashfs objects, we need to be able
top copy an entire tree as part of the object.

Signed-off-by: David Oberhollenzer <david.oberhollenzer@sigma-star.at>

3 files changed, 111 insertions, 1 deletions

diff --git a/include/rbtree.h b/include/rbtree.h
index 7fe1752..bbba711 100644
--- a/include/rbtree.h
+++ b/include/rbtree.h
@@ -49,6 +49,8 @@ SQFS_INTERNAL int rbtree_init(rbtree_t *tree, size_t keysize, size_t valuesize,
 
 SQFS_INTERNAL void rbtree_cleanup(rbtree_t *tree);
 
+SQFS_INTERNAL int rbtree_copy(const rbtree_t *tree, rbtree_t *out);
+
 SQFS_INTERNAL int rbtree_insert(rbtree_t *tree, const void *key,
 				const void *value);
 
diff --git a/lib/util/rbtree.c b/lib/util/rbtree.c
index 4e96435..873add8 100644
--- a/lib/util/rbtree.c
+++ b/lib/util/rbtree.c
@@ -105,6 +105,39 @@ static rbtree_node_t *mknode(const rbtree_t *t, const void *key, const void *val
 	return node;
 }
 
+static rbtree_node_t *copy_node(const rbtree_t *t, const rbtree_node_t *n)
+{
+	rbtree_node_t *out;
+
+	out = calloc(1, sizeof(*out) + t->key_size_padded + t->value_size);
+	if (out == NULL)
+		return NULL;
+
+	memcpy(out, n, sizeof(*n) + t->key_size_padded + t->value_size);
+	out->left = NULL;
+	out->right = NULL;
+
+	if (n->left != NULL) {
+		out->left = copy_node(t, n->left);
+
+		if (out->left == NULL) {
+			destroy_nodes_dfs(out);
+			return NULL;
+		}
+	}
+
+	if (n->right != NULL) {
+		out->right = copy_node(t, n->right);
+
+		if (out->right == NULL) {
+			destroy_nodes_dfs(out);
+			return NULL;
+		}
+	}
+
+	return out;
+}
+
 int rbtree_init(rbtree_t *tree, size_t keysize, size_t valuesize,
 		int(*key_compare)(const void *, const void *))
 {
@@ -146,6 +179,23 @@ int rbtree_init(rbtree_t *tree, size_t keysize, size_t valuesize,
 	return 0;
 }
 
+int rbtree_copy(const rbtree_t *tree, rbtree_t *out)
+{
+	memcpy(out, tree, sizeof(*out));
+	out->root = NULL;
+
+	if (tree->root != NULL) {
+		out->root = copy_node(tree, tree->root);
+
+		if (out->root == NULL) {
+			memset(out, 0, sizeof(*out));
+			return SQFS_ERROR_ALLOC;
+		}
+	}
+
+	return 0;
+}
+
 void rbtree_cleanup(rbtree_t *tree)
 {
 	destroy_nodes_dfs(tree->root);
diff --git a/tests/libutil/rbtree.c b/tests/libutil/rbtree.c
index 173df1c..03c0025 100644
--- a/tests/libutil/rbtree.c
+++ b/tests/libutil/rbtree.c
@@ -109,13 +109,55 @@ static void check_black_depth_dfs(rbtree_node_t *n, size_t ref,
 		check_black_depth_dfs(n->right, ref, counter);
 }
 
+static int check_subtrees_equal(const rbtree_node_t *lhs,
+				const rbtree_node_t *rhs,
+				size_t datasize)
+{
+	if (lhs == rhs)
+		return -1;
+
+	if (lhs->value_offset != rhs->value_offset)
+		return -1;
+
+	if ((lhs->is_red && !rhs->is_red) || (!lhs->is_red && rhs->is_red))
+		return -1;
+
+	if (memcmp(lhs->data, rhs->data, datasize) != 0)
+		return -1;
+
+	if (lhs->left == NULL) {
+		if (rhs->left != NULL)
+			return -1;
+	} else {
+		if (rhs->left == NULL)
+			return -1;
+
+		if (check_subtrees_equal(lhs->left, rhs->left, datasize))
+			return -1;
+	}
+
+	if (lhs->right == NULL) {
+		if (rhs->right != NULL)
+			return -1;
+	} else {
+		if (rhs->right == NULL)
+			return -1;
+
+		if (check_subtrees_equal(lhs->right, rhs->right, datasize))
+			return -1;
+	}
+
+	return 0;
+}
+
 int main(void)
 {
 	size_t count, blkdepth, mind, maxd;
 	sqfs_s32 key, key2;
+	rbtree_t rb, copy;
 	rbtree_node_t *n;
 	sqfs_u64 value;
-	rbtree_t rb;
+	int ret;
 
 	TEST_ASSERT(rbtree_init(&rb, sizeof(sqfs_s32),
 				sizeof(sqfs_u64), key_compare) == 0);
@@ -168,6 +210,22 @@ int main(void)
 		TEST_EQUAL_UI((sqfs_u64)(key + 10000), value);
 	}
 
+	/* test if copy works */
+	ret = rbtree_copy(&rb, &copy);
+	TEST_EQUAL_I(ret, 0);
+
+	TEST_EQUAL_UI(rb.key_size, copy.key_size);
+	TEST_EQUAL_UI(rb.key_size_padded, copy.key_size_padded);
+	TEST_EQUAL_UI(rb.value_size, copy.value_size);
+	TEST_ASSERT(rb.key_compare == copy.key_compare);
+	TEST_ASSERT(rb.root != copy.root);
+
+	ret = check_subtrees_equal(rb.root, copy.root,
+				   rb.key_size_padded + rb.value_size);
+	TEST_EQUAL_I(ret, 0);
+
+	/* cleanup */
 	rbtree_cleanup(&rb);
+	rbtree_cleanup(&copy);
 	return EXIT_SUCCESS;
 }