From 5d49ddd764b05083e1c276df0b50d43ec6896931 Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Wed, 29 Jun 2022 16:28:02 -0500 Subject: [PATCH 2/2] Add rbtree missing comparison searches * Add find great/less or equal by adding an argument --- src/backend/lib/rbtree.c | 61 +++++++++++++ src/include/lib/rbtree.h | 2 + src/test/modules/test_rbtree/test_rbtree.c | 101 +++++++++++++++++++++ 3 files changed, 164 insertions(+) diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c index e1cc4110bd..92f6b99164 100644 --- a/src/backend/lib/rbtree.c +++ b/src/backend/lib/rbtree.c @@ -161,6 +161,67 @@ rbt_find(RBTree *rbt, const RBTNode *data) return NULL; } +/* + * rbt_find_great: search for a greater value in an RBTree + * + * If equal_match is true, this will be a great or equal search. + * + * Returns the matching tree entry, or NULL if no match is found. + */ +RBTNode * +rbt_find_great(RBTree *rbt, const RBTNode *data, bool equal_match) +{ + RBTNode *node = rbt->root; + RBTNode *greater = NULL; + + while (node != RBTNIL) + { + int cmp = rbt->comparator(data, node, rbt->arg); + + if (equal_match && cmp == 0) + return node; + else if (cmp < 0) + { + greater = node; + node = node->left; + } + else + node = node->right; + } + + return greater; +} + +/* + * rbt_find_less: search for a lesser value in an RBTree + * + * If equal_match is true, this will be a less or equal search. + * + * Returns the matching tree entry, or NULL if no match is found. + */ +RBTNode * +rbt_find_less(RBTree *rbt, const RBTNode *data, bool equal_match) +{ + RBTNode *node = rbt->root; + RBTNode *lesser = NULL; + + while (node != RBTNIL) + { + int cmp = rbt->comparator(data, node, rbt->arg); + + if (equal_match && cmp == 0) + return node; + else if (cmp > 0){ + lesser = node; + node = node->right; + } + else + node = node->left; + } + + return lesser; +} + /* * rbt_leftmost: fetch the leftmost (smallest-valued) tree node. * Returns NULL if tree is empty. diff --git a/src/include/lib/rbtree.h b/src/include/lib/rbtree.h index 580a3e3439..13cf68415e 100644 --- a/src/include/lib/rbtree.h +++ b/src/include/lib/rbtree.h @@ -67,6 +67,8 @@ extern RBTree *rbt_create(Size node_size, void *arg); extern RBTNode *rbt_find(RBTree *rbt, const RBTNode *data); +extern RBTNode *rbt_find_great(RBTree *rbt, const RBTNode *data, bool equal_match); +extern RBTNode *rbt_find_less(RBTree *rbt, const RBTNode *data, bool equal_match); extern RBTNode *rbt_leftmost(RBTree *rbt); extern RBTNode *rbt_insert(RBTree *rbt, const RBTNode *data, bool *isNew); diff --git a/src/test/modules/test_rbtree/test_rbtree.c b/src/test/modules/test_rbtree/test_rbtree.c index 7cb38759a2..c17d524e01 100644 --- a/src/test/modules/test_rbtree/test_rbtree.c +++ b/src/test/modules/test_rbtree/test_rbtree.c @@ -278,6 +278,105 @@ testfind(int size) } } +/* + * Check the correctness of the rbt_find_great function by searching for + * an equal key and all of the greater keys. + */ +static void +testfindgte(int size) +{ + RBTree *tree = create_int_rbtree(); + + /* + * Using the size as the random key to search wouldn't allow us to get at + * least one greater match, so we do size - 1 + */ + int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1); + IntRBTreeNode searchNode = {.key = randomKey}; + IntRBTreeNode *eqNode; + IntRBTreeNode *gtNode; + + /* Insert natural numbers */ + rbt_populate(tree, size, 1); + + /* + * Since the search key is included in the naturals of the tree, we're + * sure to find an equal match + */ + eqNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true); + + if (eqNode == NULL) + elog(ERROR, "key was not found"); + + /* ensure we find an equal match */ + if (!(eqNode->key == searchNode.key)) + elog(ERROR, "rbt_find_great operation in rbtree didn't find an equal key"); + + /* Find the rest of the naturals greater than the search key */ + for (int i = 1; i < size - searchNode.key; i++) + { + // switch equal_match to false so we only find greater matches now + gtNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, false); + + /* ensure we find a greater match */ + if (!(gtNode->key > searchNode.key)) + elog(ERROR, "rbt_find_great operation in rbtree didn't find a greater key"); + + /* delete the previous match to find a greater one */ + rbt_delete(tree, (RBTNode *) gtNode); + } +} + +/* + * Check the correctness of the rbt_find_less operation by searching for + * an equal key and all of the lesser keys. + */ +static void +testfindlte(int size) +{ + RBTree *tree = create_int_rbtree(); + + /* + * Using 0 as the random key to search wouldn't allow us to get at least + * one lesser match, so we start from 1 + */ + int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 1, size); + IntRBTreeNode searchNode = {.key = randomKey}; + IntRBTreeNode *eqNode; + IntRBTreeNode *ltNode; + + /* Insert natural numbers */ + rbt_populate(tree, size, 1); + + /* + * Since the search key is included in the naturals of the tree, we're + * sure to find an equal match + */ + eqNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true); + + if (eqNode == NULL) + elog(ERROR, "key was not found"); + + /* ensure we find an equal match */ + if (!(eqNode->key == searchNode.key)){ + elog(ERROR, "rbt_find_less operation in rbtree didn't find an equal key"); + } + + /* Find the rest of the naturals lesser than the search key */ + for (int i = searchNode.key; i > 0; i--) + { + // switch equal_match to false so we only find lesser matches now + ltNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, false); + + /* ensure we find a lesser match */ + if (!(ltNode->key < searchNode.key)) + elog(ERROR, "rbt_find_less operation in rbtree didn't find a lesser key"); + + /* delete the previous match to find a lesser one */ + rbt_delete(tree, (RBTNode *) ltNode); + } +} + /* * Check the correctness of the rbt_leftmost operation. * This operation should always return the smallest element of the tree. @@ -408,6 +507,8 @@ test_rb_tree(PG_FUNCTION_ARGS) testleftright(size); testrightleft(size); testfind(size); + testfindgte(size); + testfindlte(size); testleftmost(size); testdelete(size, Max(size / 10, 1)); PG_RETURN_VOID(); -- 2.34.1