From 02db0ae8ff83a42913b1b4224ccd8377f5fc5323 Mon Sep 17 00:00:00 2001 From: David Oberhollenzer Date: Fri, 5 Mar 2021 17:41:53 +0100 Subject: 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 --- include/rbtree.h | 2 ++ lib/util/rbtree.c | 50 +++++++++++++++++++++++++++++++++++++++++ tests/libutil/rbtree.c | 60 +++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 111 insertions(+), 1 deletion(-) 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, ©); + 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(©); return EXIT_SUCCESS; } -- cgit v1.2.3