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.

Discussion: https://postgr.es/m/CAGRrpzYE8-7GCoaPjOiL9T_HY605MRax-2jgTtLq236uksZ1Sw%40mail.gmail.com
Author: Steve Chavez, Alexander Korotkov
Reviewed-by: Alexander Korotkov

diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c
index e1cc4110bd447d1deb3046276ee40eb9feed7e27..427cec1ada83a174af34da5b4cb1402061cf5bb9 100644 (file)
--- 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 580a3e3439720aa3e7aa9a054ee3d3b9dc51db2c..13cf68415e2578448c05335b41046b4a91955b30 100644 (file)
--- 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 7cb38759a2ed7e2831c9e4779c1c25456d4ce621..3e638112a47e0fec2e9ead804f25e3344a596314 100644 (file)
--- a/src/test/modules/test_rbtree/test_rbtree.c
+++ b/src/test/modules/test_rbtree/test_rbtree.c
@@ -278,6 +278,107 @@ 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++)
+       {
+               /*
+                * 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_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 +509,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();