From 834a3c9743703c04e6326063d8f1210103730a0e Mon Sep 17 00:00:00 2001 From: steve-chavez Date: Wed, 29 Jun 2022 16:28:02 -0500 Subject: [PATCH 2/2] Add missing inequality searches to rbtree PostgreSQL contains the implementation of the red-black tree. The red-black tree is the ordered data structure, and one of its advantages is the ability to do inequality searches. This commit adds rbt_find_less() and rbt_find_great() functions implementing these searches. While these searches aren't yet used in the core code, they might be useful for extensions. --- src/backend/lib/rbtree.c | 62 ++++++++++++++ src/include/lib/rbtree.h | 2 + src/test/modules/test_rbtree/test_rbtree.c | 99 ++++++++++++++++++++++ 3 files changed, 163 insertions(+) diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c index e1cc4110bd4..427cec1ada8 100644 --- a/src/backend/lib/rbtree.c +++ b/src/backend/lib/rbtree.c @@ -161,6 +161,68 @@ 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 580a3e34397..13cf68415e2 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 7cb38759a2e..d770c9cca5f 100644 --- a/src/test/modules/test_rbtree/test_rbtree.c +++ b/src/test/modules/test_rbtree/test_rbtree.c @@ -278,6 +278,104 @@ testfind(int size) } } +/* + * Check the correctness of the rbt_find_less() and rbt_find_great() functions + * by searching for an equal key and iterating the lesser keys then the greater + * keys. + */ +static void +testfindltgt(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); + bool keyDeleted; + IntRBTreeNode searchNode = {.key = randomKey}; + IntRBTreeNode *lteNode; + IntRBTreeNode *gteNode; + IntRBTreeNode *node; + + /* 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 + */ + lteNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true); + gteNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true); + + if (lteNode == NULL || lteNode->key != searchNode.key) + elog(ERROR, "rbt_find_less() didn't find the equal key"); + + if (gteNode == NULL || gteNode->key != searchNode.key) + elog(ERROR, "rbt_find_great() didn't find the equal key"); + + if (lteNode != gteNode) + elog(ERROR, "rbt_find_less() and rbt_find_great() found different equal keys"); + + /* Find the rest of the naturals lesser than the search key */ + keyDeleted = false; + for (; searchNode.key > 0; searchNode.key--) + { + /* + * Find the next key. If the current key is deleted, we can pass + * equal_match == true and still find the next one. + */ + node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, + keyDeleted); + + /* ensure we find a lesser match */ + if (!node || !(node->key < searchNode.key)) + elog(ERROR, "rbt_find_less() didn't find a lesser key"); + + /* randomly delete the found key or leave it */ + keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1); + if (keyDeleted) + rbt_delete(tree, (RBTNode *) node); + } + + /* Find the rest of the naturals greater than the search key */ + keyDeleted = false; + for (searchNode.key = randomKey; searchNode.key < size - 1; searchNode.key++) + { + /* switch equal_match to false so we only find greater matches now */ + node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, + keyDeleted); + + /* ensure we find a greater match */ + if (!node || !(node->key > searchNode.key)) + elog(ERROR, "rbt_find_great() didn't find a greater key"); + + /* randomly delete the found key or leave it */ + keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1); + if (keyDeleted) + rbt_delete(tree, (RBTNode *) node); + } + + /* Check out of bounds searches find nothing */ + searchNode.key = -1; + node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true); + if (node != NULL) + elog(ERROR, "rbt_find_less() found non-inserted element"); + searchNode.key = 0; + node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, false); + if (node != NULL) + elog(ERROR, "rbt_find_less() found non-inserted element"); + searchNode.key = size; + node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true); + if (node != NULL) + elog(ERROR, "rbt_find_great() found non-inserted element"); + searchNode.key = size - 1; + node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, false); + if (node != NULL) + elog(ERROR, "rbt_find_great() found non-inserted element"); +} + /* * Check the correctness of the rbt_leftmost operation. * This operation should always return the smallest element of the tree. @@ -408,6 +506,7 @@ test_rb_tree(PG_FUNCTION_ARGS) testleftright(size); testrightleft(size); testfind(size); + testfindltgt(size); testleftmost(size); testdelete(size, Max(size / 10, 1)); PG_RETURN_VOID(); -- 2.24.3 (Apple Git-128)