0001-Add-red-black-tree-missing-comparison-searches.patch

From 85435fe0fad92d593940f98a493d1acd973ccda2 Mon Sep 17 00:00:00 2001
From: steve-chavez <stevechavezast@gmail.com>
Date: Tue, 28 Jun 2022 23:46:38 -0500
Subject: [PATCH] Add red-black tree missing comparison searches

* Added find great/less or equal
* Change rbtree sentinel to C99 designator
---
 src/backend/lib/rbtree.c                   |  60 +++++++++++-
 src/include/lib/rbtree.h                   |   2 +
 src/test/modules/test_rbtree/test_rbtree.c | 104 +++++++++++++++++++++
 3 files changed, 165 insertions(+), 1 deletion(-)

diff --git a/src/backend/lib/rbtree.c b/src/backend/lib/rbtree.c
index a9981dbada..af3990d01c 100644
--- a/src/backend/lib/rbtree.c
+++ b/src/backend/lib/rbtree.c
@@ -62,7 +62,7 @@ struct RBTree
 
 static RBTNode sentinel =
 {
-	RBTBLACK, RBTNIL, RBTNIL, NULL
+	.color = RBTBLACK,.left = RBTNIL,.right = RBTNIL,.parent = NULL
 };
 
 
@@ -161,6 +161,64 @@ rbt_find(RBTree *rbt, const RBTNode *data)
 	return NULL;
 }
 
+/*
+ * rbt_find_great_equal: search for an equal or greater value in an RBTree
+ *
+ * Returns the matching tree entry, or NULL if no match is found.
+ */
+RBTNode *
+rbt_find_great_equal(RBTree *rbt, const RBTNode *data)
+{
+	RBTNode    *node = rbt->root;
+	RBTNode    *greater = NULL;
+
+	while (node != RBTNIL)
+	{
+		int			cmp = rbt->comparator(data, node, rbt->arg);
+
+		if (cmp == 0)
+			return node;
+		else if (cmp < 0)
+		{
+			greater = node;
+			node = node->left;
+		}
+		else
+			node = node->right;
+	}
+
+	return greater;
+}
+
+/*
+ * rbt_find_less_equal: search for an equal or lesser value in an RBTree
+ *
+ * Returns the matching tree entry, or NULL if no match is found.
+ */
+RBTNode *
+rbt_find_less_equal(RBTree *rbt, const RBTNode *data)
+{
+	RBTNode    *node = rbt->root;
+	RBTNode    *lesser = NULL;
+
+	while (node != RBTNIL)
+	{
+		int			cmp = rbt->comparator(data, node, rbt->arg);
+
+		if (cmp == 0)
+			return node;
+		else if (cmp < 0)
+			node = node->left;
+		else
+		{
+			lesser = node;
+			node = node->right;
+		}
+	}
+
+	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..2e8f5b0a8f 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_equal(RBTree *rbt, const RBTNode *data);
+extern RBTNode *rbt_find_less_equal(RBTree *rbt, const RBTNode *data);
 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..6a5a8abee2 100644
--- a/src/test/modules/test_rbtree/test_rbtree.c
+++ b/src/test/modules/test_rbtree/test_rbtree.c
@@ -278,6 +278,108 @@ testfind(int size)
 	}
 }
 
+/*
+ * Check the correctness of the rbt_find_great_equal operation 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_equal(tree, (RBTNode *) &searchNode);
+
+	if (eqNode == NULL)
+		elog(ERROR, "key was not found");
+
+	/* ensure we find an equal match */
+	if (!(eqNode->key == searchNode.key))
+		elog(ERROR, "find_great_equal operation in rbtree didn't find an equal key");
+
+	/* Delete the equal match so we can find greater matches */
+	rbt_delete(tree, (RBTNode *) eqNode);
+
+	/* Find the rest of the naturals greater than the search key */
+	for (int i = 1; i < size - searchNode.key; i++)
+	{
+		gtNode = (IntRBTreeNode *) rbt_find_great_equal(tree, (RBTNode *) &searchNode);
+
+		/* ensure we find a greater match */
+		if (!(gtNode->key > searchNode.key))
+			elog(ERROR, "find_great_equal 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_equal 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_equal(tree, (RBTNode *) &searchNode);
+
+	if (eqNode == NULL)
+		elog(ERROR, "key was not found");
+
+	/* ensure we find an equal match */
+	if (!(eqNode->key == searchNode.key))
+		elog(ERROR, "find_less_equal operation in rbtree didn't find an equal key");
+
+	/* Delete the equal match so we can find lesser matches */
+	rbt_delete(tree, (RBTNode *) eqNode);
+
+	/* Find the rest of the naturals lesser than the search key */
+	for (int i = searchNode.key; i > 0; i--)
+	{
+		ltNode = (IntRBTreeNode *) rbt_find_less_equal(tree, (RBTNode *) &searchNode);
+
+		/* ensure we find a lesser match */
+		if (!(ltNode->key < searchNode.key))
+			elog(ERROR, "find_less_equal 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 +510,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