test_rbtree.c

Go to the documentation of this file.

/*--------------------------------------------------------------------------
 *
 * 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();
}