0002-Add-missing-inequality-searches-to-rbtree-v2.patch

From 834a3c9743703c04e6326063d8f1210103730a0e Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavezast@gmail.com>
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)