/*--------------------------------------------------------------------------
 *
 * test_rbtree.c
 *  Test correctness of red-black tree operations.
 *
 * Copyright (c) 2009-2025, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *  src/test/modules/test_rbtree/test_rbtree.c
 *
 * -------------------------------------------------------------------------
 */

#include "postgres.h"

#include "common/pg_prng.h"
#include "fmgr.h"
#include "lib/rbtree.h"
#include "utils/memutils.h"

PG_MODULE_MAGIC;


/*
 * Our test trees store an integer key, and nothing else.
 */
typedef struct IntRBTreeNode
{
 RBTNode  rbtnode;
 int   key;
} IntRBTreeNode;


/*
 * Node comparator.  We don't worry about overflow in the subtraction,
 * since none of our test keys are negative.
 */
static int
irbt_cmp(const RBTNode *a, const RBTNode *b, void *arg)
{
 const IntRBTreeNode *ea = (const IntRBTreeNode *) a;
 const IntRBTreeNode *eb = (const IntRBTreeNode *) b;

 return ea->key - eb->key;
}

/*
 * Node combiner.  For testing purposes, just check that library doesn't
 * try to combine unequal keys.
 */
static void
irbt_combine(RBTNode *existing, const RBTNode *newdata, void *arg)
{
 const IntRBTreeNode *eexist = (const IntRBTreeNode *) existing;
 const IntRBTreeNode *enew = (const IntRBTreeNode *) newdata;

 if (eexist->key != enew->key)
  elog(ERROR, "red-black tree combines %d into %d",
    enew->key, eexist->key);
}

/* Node allocator */
static RBTNode *
irbt_alloc(void *arg)
{
 return (RBTNode *) palloc(sizeof(IntRBTreeNode));
}

/* Node freer */
static void
irbt_free(RBTNode *node, void *arg)
{
 pfree(node);
}

/*
 * Create a red-black tree using our support functions
 */
static RBTree *
create_int_rbtree(void)
{
 return rbt_create(sizeof(IntRBTreeNode),
       irbt_cmp,
       irbt_combine,
       irbt_alloc,
       irbt_free,
       NULL);
}

/*
 * Generate a random permutation of the integers 0..size-1
 */
static int *
GetPermutation(int size)
{
 int     *permutation;
 int   i;

 permutation = (int *) palloc(size * sizeof(int));

 permutation[0] = 0;

 /*
  * This is the "inside-out" variant of the Fisher-Yates shuffle algorithm.
  * Notionally, we append each new value to the array and then swap it with
  * a randomly-chosen array element (possibly including itself, else we
  * fail to generate permutations with the last integer last).  The swap
  * step can be optimized by combining it with the insertion.
 */
 for (i = 1; i < size; i++)
 {
  int   j = pg_prng_uint64_range(&pg_global_prng_state, 0, i);

  if (j < i)    /* avoid fetching undefined data if j=i */
   permutation[i] = permutation[j];
  permutation[j] = i;
 }

 return permutation;
}

/*
 * Populate an empty RBTree with "size" integers having the values
 * 0, step, 2*step, 3*step, ..., inserting them in random order
 */
static void
rbt_populate(RBTree *tree, int size, int step)
{
 int     *permutation = GetPermutation(size);
 IntRBTreeNode node;
 bool  isNew;
 int   i;

 /* Insert values.  We don't expect any collisions. */
 for (i = 0; i < size; i++)
 {
  node.key = step * permutation[i];
  rbt_insert(tree, (RBTNode *) &node, &isNew);
  if (!isNew)
   elog(ERROR, "unexpected !isNew result from rbt_insert");
 }

 /*
  * Re-insert the first value to make sure collisions work right.  It's
  * probably not useful to test that case over again for all the values.
 */
 if (size > 0)
 {
  node.key = step * permutation[0];
  rbt_insert(tree, (RBTNode *) &node, &isNew);
  if (isNew)
   elog(ERROR, "unexpected isNew result from rbt_insert");
 }

 pfree(permutation);
}

/*
 * Check the correctness of left-right traversal.
 * Left-right traversal is correct if all elements are
 * visited in increasing order.
 */
static void
testleftright(int size)
{
 RBTree    *tree = create_int_rbtree();
 IntRBTreeNode *node;
 RBTreeIterator iter;
 int   lastKey = -1;
 int   count = 0;

 /* check iteration over empty tree */
 rbt_begin_iterate(tree, LeftRightWalk, &iter);
 if (rbt_iterate(&iter) != NULL)
  elog(ERROR, "left-right walk over empty tree produced an element");

 /* fill tree with consecutive natural numbers */
 rbt_populate(tree, size, 1);

 /* iterate over the tree */
 rbt_begin_iterate(tree, LeftRightWalk, &iter);

 while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
 {
  /* check that order is increasing */
  if (node->key <= lastKey)
   elog(ERROR, "left-right walk gives elements not in sorted order");
  lastKey = node->key;
  count++;
 }

 if (lastKey != size - 1)
  elog(ERROR, "left-right walk did not reach end");
 if (count != size)
  elog(ERROR, "left-right walk missed some elements");
}

/*
 * Check the correctness of right-left traversal.
 * Right-left traversal is correct if all elements are
 * visited in decreasing order.
 */
static void
testrightleft(int size)
{
 RBTree    *tree = create_int_rbtree();
 IntRBTreeNode *node;
 RBTreeIterator iter;
 int   lastKey = size;
 int   count = 0;

 /* check iteration over empty tree */
 rbt_begin_iterate(tree, RightLeftWalk, &iter);
 if (rbt_iterate(&iter) != NULL)
  elog(ERROR, "right-left walk over empty tree produced an element");

 /* fill tree with consecutive natural numbers */
 rbt_populate(tree, size, 1);

 /* iterate over the tree */
 rbt_begin_iterate(tree, RightLeftWalk, &iter);

 while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
 {
  /* check that order is decreasing */
  if (node->key >= lastKey)
   elog(ERROR, "right-left walk gives elements not in sorted order");
  lastKey = node->key;
  count++;
 }

 if (lastKey != 0)
  elog(ERROR, "right-left walk did not reach end");
 if (count != size)
  elog(ERROR, "right-left walk missed some elements");
}

/*
 * Check the correctness of the rbt_find operation by searching for
 * both elements we inserted and elements we didn't.
 */
static void
testfind(int size)
{
 RBTree    *tree = create_int_rbtree();
 int   i;

 /* Insert even integers from 0 to 2 * (size-1) */
 rbt_populate(tree, size, 2);

 /* Check that all inserted elements can be found */
 for (i = 0; i < size; i++)
 {
  IntRBTreeNode node;
  IntRBTreeNode *resultNode;

  node.key = 2 * i;
  resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
  if (resultNode == NULL)
   elog(ERROR, "inserted element was not found");
  if (node.key != resultNode->key)
   elog(ERROR, "find operation in rbtree gave wrong result");
 }

 /*
  * Check that not-inserted elements can not be found, being sure to try
  * values before the first and after the last element.
 */
 for (i = -1; i <= 2 * size; i += 2)
 {
  IntRBTreeNode node;
  IntRBTreeNode *resultNode;

  node.key = i;
  resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
  if (resultNode != NULL)
   elog(ERROR, "not-inserted element was found");
 }
}

/*
 * 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.
 */
static void
testleftmost(int size)
{
 RBTree    *tree = create_int_rbtree();
 IntRBTreeNode *result;

 /* Check that empty tree has no leftmost element */
 if (rbt_leftmost(tree) != NULL)
  elog(ERROR, "leftmost node of empty tree is not NULL");

 /* fill tree with consecutive natural numbers */
 rbt_populate(tree, size, 1);

 /* Check that leftmost element is the smallest one */
 result = (IntRBTreeNode *) rbt_leftmost(tree);
 if (result == NULL || result->key != 0)
  elog(ERROR, "rbt_leftmost gave wrong result");
}

/*
 * Check the correctness of the rbt_delete operation.
 */
static void
testdelete(int size, int delsize)
{
 RBTree    *tree = create_int_rbtree();
 int     *deleteIds;
 bool    *chosen;
 int   i;

 /* fill tree with consecutive natural numbers */
 rbt_populate(tree, size, 1);

 /* Choose unique ids to delete */
 deleteIds = (int *) palloc(delsize * sizeof(int));
 chosen = (bool *) palloc0(size * sizeof(bool));

 for (i = 0; i < delsize; i++)
 {
  int   k = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);

  while (chosen[k])
   k = (k + 1) % size;
  deleteIds[i] = k;
  chosen[k] = true;
 }

 /* Delete elements */
 for (i = 0; i < delsize; i++)
 {
  IntRBTreeNode find;
  IntRBTreeNode *node;

  find.key = deleteIds[i];
  /* Locate the node to be deleted */
  node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
  if (node == NULL || node->key != deleteIds[i])
   elog(ERROR, "expected element was not found during deleting");
  /* Delete it */
  rbt_delete(tree, (RBTNode *) node);
 }

 /* Check that deleted elements are deleted */
 for (i = 0; i < size; i++)
 {
  IntRBTreeNode node;
  IntRBTreeNode *result;

  node.key = i;
  result = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
  if (chosen[i])
  {
   /* Deleted element should be absent */
   if (result != NULL)
    elog(ERROR, "deleted element still present in the rbtree");
  }
  else
  {
   /* Else it should be present */
   if (result == NULL || result->key != i)
    elog(ERROR, "delete operation removed wrong rbtree value");
  }
 }

 /* Delete remaining elements, so as to exercise reducing tree to empty */
 for (i = 0; i < size; i++)
 {
  IntRBTreeNode find;
  IntRBTreeNode *node;

  if (chosen[i])
   continue;
  find.key = i;
  /* Locate the node to be deleted */
  node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
  if (node == NULL || node->key != i)
   elog(ERROR, "expected element was not found during deleting");
  /* Delete it */
  rbt_delete(tree, (RBTNode *) node);
 }

 /* Tree should now be empty */
 if (rbt_leftmost(tree) != NULL)
  elog(ERROR, "deleting all elements failed");

 pfree(deleteIds);
 pfree(chosen);
}

/*
 * SQL-callable entry point to perform all tests
 *
 * Argument is the number of entries to put in the trees
 */
PG_FUNCTION_INFO_V1(test_rb_tree);

Datum
test_rb_tree(PG_FUNCTION_ARGS)
{
 int   size = PG_GETARG_INT32(0);

 if (size <= 0 || size > MaxAllocSize / sizeof(int))
  elog(ERROR, "invalid size for test_rb_tree: %d", size);
 testleftright(size);
 testrightleft(size);
 testfind(size);
 testfindltgt(size);
 testleftmost(size);
 testdelete(size, Max(size / 10, 1));
 PG_RETURN_VOID();
}