300 lines
8.0 KiB
C
300 lines
8.0 KiB
C
#include "ktest.h"
|
|
#include "rbt.h"
|
|
#include "heap.h"
|
|
|
|
typedef struct IntNode {
|
|
RbtNode node;
|
|
int32 key;
|
|
} IntNode;
|
|
|
|
static int compareIntNode(const RbtNode* a, const RbtNode* b)
|
|
{
|
|
const IntNode* ia = rbtEntry(a, IntNode, node);
|
|
const IntNode* ib = rbtEntry(b, IntNode, node);
|
|
if (ia->key < ib->key)
|
|
return -1;
|
|
if (ia->key > ib->key)
|
|
return 1;
|
|
return 0;
|
|
}
|
|
|
|
static int lookupIntKey(const void* key, const RbtNode* node)
|
|
{
|
|
int32 k = *(const int32*)key;
|
|
const IntNode* in = rbtEntry(node, IntNode, node);
|
|
if (k < in->key)
|
|
return -1;
|
|
if (k > in->key)
|
|
return 1;
|
|
return 0;
|
|
}
|
|
|
|
/* Fisher-Yates over a small LCG so insertion/removal order is scrambled but
|
|
reproducible across runs -- the point is to exercise every rotation case,
|
|
not to be a real source of randomness. */
|
|
static void shuffle(int32* arr, int32 n, uint32 seed)
|
|
{
|
|
uint32 state = seed;
|
|
for (int32 i = n - 1; i > 0; i--) {
|
|
state = state * 1664525u + 1013904223u;
|
|
int32 j = (int32)(state % (uint32)(i + 1));
|
|
int32 tmp = arr[i];
|
|
arr[i] = arr[j];
|
|
arr[j] = tmp;
|
|
}
|
|
}
|
|
|
|
/* Returns the black-height of node (counting nullptr leaves as height 1) and
|
|
clears *valid if either red-black rule is broken anywhere in the subtree:
|
|
a red node with a red child, or unequal black-height between the two
|
|
sides of some node. */
|
|
static int32 checkSubtree(RbtNode* node, bool* valid)
|
|
{
|
|
if (node == nullptr)
|
|
return 1;
|
|
|
|
if (node->red) {
|
|
if ((node->left != nullptr && node->left->red) ||
|
|
(node->right != nullptr && node->right->red))
|
|
*valid = false;
|
|
}
|
|
if (node->left != nullptr && node->left->parent != node)
|
|
*valid = false;
|
|
if (node->right != nullptr && node->right->parent != node)
|
|
*valid = false;
|
|
|
|
int32 leftHeight = checkSubtree(node->left, valid);
|
|
int32 rightHeight = checkSubtree(node->right, valid);
|
|
if (leftHeight != rightHeight)
|
|
*valid = false;
|
|
|
|
return leftHeight + (node->red ? 0 : 1);
|
|
}
|
|
|
|
static bool checkInvariants(RbTree* tree)
|
|
{
|
|
if (tree->root != nullptr && tree->root->red)
|
|
return false;
|
|
bool valid = true;
|
|
checkSubtree(tree->root, &valid);
|
|
return valid;
|
|
}
|
|
|
|
KTEST(rbtEmptyTree)
|
|
{
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
|
|
KT_TRUE(t, rbtIsEmpty(&tree));
|
|
KT_EQ(t, rbtCount(&tree), 0);
|
|
KT_NULL(t, rbtFirst(&tree));
|
|
KT_NULL(t, rbtLast(&tree));
|
|
|
|
int32 key = 42;
|
|
KT_NULL(t, rbtFind(&tree, &key, lookupIntKey));
|
|
}
|
|
|
|
KTEST(rbtBasicInsertFindOrder)
|
|
{
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
|
|
int32 values[] = { 5, 3, 8, 1, 4, 7, 9 };
|
|
usize n = sizeof(values) / sizeof(values[0]);
|
|
IntNode nodes[7];
|
|
|
|
for (usize i = 0; i < n; i++) {
|
|
nodes[i].key = values[i];
|
|
KT_TRUE(t, rbtInsert(&tree, &nodes[i].node, compareIntNode));
|
|
}
|
|
KT_EQ(t, rbtCount(&tree), n);
|
|
KT_TRUE(t, checkInvariants(&tree));
|
|
|
|
int32 expected[] = { 1, 3, 4, 5, 7, 8, 9 };
|
|
usize i = 0;
|
|
for (RbtNode* c = rbtFirst(&tree); c != nullptr; c = rbtNext(c), i++)
|
|
KT_EQ(t, rbtEntry(c, IntNode, node)->key, expected[i]);
|
|
KT_EQ(t, i, n);
|
|
|
|
/* walking backwards from the max visits the same keys in reverse */
|
|
i = n;
|
|
for (RbtNode* c = rbtLast(&tree); c != nullptr; c = rbtPrev(c))
|
|
KT_EQ(t, rbtEntry(c, IntNode, node)->key, expected[--i]);
|
|
KT_EQ(t, i, 0);
|
|
|
|
int32 hitKey = 8;
|
|
RbtNode* found = rbtFind(&tree, &hitKey, lookupIntKey);
|
|
KT_NOTNULL(t, found);
|
|
if (found)
|
|
KT_EQ(t, rbtEntry(found, IntNode, node)->key, 8);
|
|
|
|
int32 missKey = 100;
|
|
KT_NULL(t, rbtFind(&tree, &missKey, lookupIntKey));
|
|
|
|
KT_EQ(t, rbtEntry(rbtFirst(&tree), IntNode, node)->key, 1);
|
|
KT_EQ(t, rbtEntry(rbtLast(&tree), IntNode, node)->key, 9);
|
|
}
|
|
|
|
KTEST(rbtRejectsDuplicateInsert)
|
|
{
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
|
|
IntNode a, b;
|
|
a.key = 5;
|
|
b.key = 5;
|
|
|
|
KT_TRUE(t, rbtInsert(&tree, &a.node, compareIntNode));
|
|
KT_FALSE(t, rbtInsert(&tree, &b.node, compareIntNode));
|
|
KT_EQ(t, rbtCount(&tree), 1);
|
|
}
|
|
|
|
KTEST(rbtSingleRemoveRelinksNeighbors)
|
|
{
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
|
|
IntNode nodes[5];
|
|
int32 values[] = { 10, 20, 30, 40, 50 };
|
|
for (int i = 0; i < 5; i++) {
|
|
nodes[i].key = values[i];
|
|
rbtInsert(&tree, &nodes[i].node, compareIntNode);
|
|
}
|
|
|
|
rbtRemove(&tree, &nodes[2].node); /* remove 30, the middle */
|
|
KT_EQ(t, rbtCount(&tree), 4);
|
|
KT_TRUE(t, checkInvariants(&tree));
|
|
|
|
int32 key30 = 30;
|
|
KT_NULL(t, rbtFind(&tree, &key30, lookupIntKey));
|
|
|
|
int32 expected[] = { 10, 20, 40, 50 };
|
|
usize i = 0;
|
|
for (RbtNode* c = rbtFirst(&tree); c != nullptr; c = rbtNext(c), i++)
|
|
KT_EQ(t, rbtEntry(c, IntNode, node)->key, expected[i]);
|
|
KT_EQ(t, i, 4);
|
|
}
|
|
|
|
#define rbtStressN 500
|
|
|
|
KTEST(rbtStressInsertMaintainsInvariants)
|
|
{
|
|
IntNode* pool = (IntNode*)kmalloc(sizeof(IntNode) * rbtStressN);
|
|
int32* order = (int32*)kmalloc(sizeof(int32) * rbtStressN);
|
|
KT_NOTNULL(t, pool);
|
|
KT_NOTNULL(t, order);
|
|
if (!pool || !order) {
|
|
kfree(pool);
|
|
kfree(order);
|
|
return;
|
|
}
|
|
|
|
for (int32 i = 0; i < rbtStressN; i++)
|
|
order[i] = i;
|
|
shuffle(order, rbtStressN, 0xC0FFEEu);
|
|
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
for (int32 i = 0; i < rbtStressN; i++) {
|
|
int32 key = order[i];
|
|
pool[key].key = key;
|
|
KT_TRUE(t, rbtInsert(&tree, &pool[key].node, compareIntNode));
|
|
}
|
|
KT_EQ(t, rbtCount(&tree), rbtStressN);
|
|
KT_TRUE(t, checkInvariants(&tree));
|
|
|
|
int32 expect = 0;
|
|
usize visited = 0;
|
|
for (RbtNode* c = rbtFirst(&tree); c != nullptr; c = rbtNext(c)) {
|
|
KT_EQ(t, rbtEntry(c, IntNode, node)->key, expect);
|
|
expect++;
|
|
visited++;
|
|
}
|
|
KT_EQ(t, visited, rbtStressN);
|
|
|
|
kfree(pool);
|
|
kfree(order);
|
|
}
|
|
|
|
KTEST(rbtStressRemoveMaintainsInvariants)
|
|
{
|
|
IntNode* pool = (IntNode*)kmalloc(sizeof(IntNode) * rbtStressN);
|
|
int32* order = (int32*)kmalloc(sizeof(int32) * rbtStressN);
|
|
int32* removeOrder = (int32*)kmalloc(sizeof(int32) * rbtStressN);
|
|
KT_NOTNULL(t, pool);
|
|
KT_NOTNULL(t, order);
|
|
KT_NOTNULL(t, removeOrder);
|
|
if (!pool || !order || !removeOrder) {
|
|
kfree(pool);
|
|
kfree(order);
|
|
kfree(removeOrder);
|
|
return;
|
|
}
|
|
|
|
for (int32 i = 0; i < rbtStressN; i++) {
|
|
order[i] = i;
|
|
removeOrder[i] = i;
|
|
}
|
|
shuffle(order, rbtStressN, 0xDEADBEEFu);
|
|
shuffle(removeOrder, rbtStressN, 0x1234567u);
|
|
|
|
RbTree tree;
|
|
rbtInit(&tree);
|
|
for (int32 i = 0; i < rbtStressN; i++) {
|
|
int32 key = order[i];
|
|
pool[key].key = key;
|
|
rbtInsert(&tree, &pool[key].node, compareIntNode);
|
|
}
|
|
|
|
/* remove every even key, in scrambled order */
|
|
int32 removed = 0;
|
|
for (int32 i = 0; i < rbtStressN; i++) {
|
|
int32 key = removeOrder[i];
|
|
if (key % 2 != 0)
|
|
continue;
|
|
RbtNode* found = rbtFind(&tree, &key, lookupIntKey);
|
|
KT_NOTNULL(t, found);
|
|
if (found) {
|
|
rbtRemove(&tree, found);
|
|
removed++;
|
|
}
|
|
}
|
|
KT_EQ(t, rbtCount(&tree), rbtStressN - removed);
|
|
KT_TRUE(t, checkInvariants(&tree));
|
|
|
|
for (int32 key = 0; key < rbtStressN; key++) {
|
|
RbtNode* found = rbtFind(&tree, &key, lookupIntKey);
|
|
if (key % 2 == 0)
|
|
KT_NULL(t, found);
|
|
else
|
|
KT_NOTNULL(t, found);
|
|
}
|
|
|
|
int32 expect = 1;
|
|
usize visited = 0;
|
|
for (RbtNode* c = rbtFirst(&tree); c != nullptr; c = rbtNext(c)) {
|
|
KT_EQ(t, rbtEntry(c, IntNode, node)->key, expect);
|
|
expect += 2;
|
|
visited++;
|
|
}
|
|
KT_EQ(t, visited, rbtStressN - removed);
|
|
|
|
/* drain the rest, in the same scrambled order, down to nothing */
|
|
for (int32 i = 0; i < rbtStressN; i++) {
|
|
int32 key = removeOrder[i];
|
|
if (key % 2 == 0)
|
|
continue;
|
|
RbtNode* found = rbtFind(&tree, &key, lookupIntKey);
|
|
if (found)
|
|
rbtRemove(&tree, found);
|
|
}
|
|
KT_TRUE(t, rbtIsEmpty(&tree));
|
|
KT_EQ(t, rbtCount(&tree), 0);
|
|
KT_NULL(t, rbtFirst(&tree));
|
|
KT_NULL(t, rbtLast(&tree));
|
|
|
|
kfree(pool);
|
|
kfree(order);
|
|
kfree(removeOrder);
|
|
}
|