B-tree

btree.c

#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

#define BTREE_ORDER 15

typedef int key_t;

struct btree_node {
    int num_keys;
    struct btree_node *parent;
    struct btree_node *children[BTREE_ORDER];
    key_t keys[BTREE_ORDER - 1];    
};

struct btree {
    struct btree_node *root;
};

static struct btree_node *btree_node_create(struct btree_node *parent)
{
    struct btree_node *node = malloc(sizeof(struct btree_node));
    node->num_keys = 0;
    node->parent = parent;
    return node;
}

static inline bool btree_node_is_leaf(struct btree_node *node)
{
    return node->children[0] == NULL;
}

static void btree_node_split(struct btree *btree, struct btree_node *lhs_node)
{
    const int mid = (BTREE_ORDER - 1) / 2;
    const key_t median = lhs_node->keys[mid];

    struct btree_node *rhs_node = btree_node_create(lhs_node->parent);

    // Copy upper half of keys and children to new node
    memcpy(rhs_node->keys, lhs_node->keys + mid + 1, (BTREE_ORDER - 1 - mid - 1) * sizeof(key_t));
    memcpy(rhs_node->children, lhs_node->children + mid + 1, (BTREE_ORDER - 1 - mid) * sizeof(struct btree_node *));

    lhs_node->num_keys = mid;
    rhs_node->num_keys = BTREE_ORDER - 1 - mid - 1;

    // Update parent pointers for children
    if (!btree_node_is_leaf(rhs_node)) {
        for (int i = 0; i <= rhs_node->num_keys; i++) {
            rhs_node->children[i]->parent = rhs_node;
        }
    }

    // Handle root split
    if (lhs_node->parent == NULL) {
        struct btree_node *new_root = btree_node_create(NULL);
        new_root->keys[0] = median;
        new_root->children[0] = lhs_node;
        new_root->children[1] = rhs_node;
        new_root->num_keys = 1;
        lhs_node->parent = new_root;
        rhs_node->parent = new_root;
        btree->root = new_root;
        return;
    }

    // Insert median into parent
    struct btree_node *parent = lhs_node->parent;
    int insert_pos = 0;
    while (insert_pos < parent->num_keys && parent->keys[insert_pos] < median) {
        insert_pos++;
    }

    // Shift keys and children to make room
    memmove(parent->keys + insert_pos + 1, parent->keys + insert_pos, (parent->num_keys - insert_pos) * sizeof(key_t));
    memmove(parent->children + insert_pos + 2, parent->children + insert_pos + 1, (parent->num_keys - insert_pos) * sizeof(struct btree_node *));

    // Insert median and new node
    parent->keys[insert_pos] = median;
    parent->children[insert_pos + 1] = rhs_node;
    parent->num_keys++;
    rhs_node->parent = parent;

    // If parent is now full, split it recursively
    if (parent->num_keys == BTREE_ORDER - 1) {
        btree_node_split(btree, parent);
    }
}

void btree_insert(struct btree *btree, int key)
{
    // If tree is empty, create root
    if (btree->root == NULL) {
        btree->root = btree_node_create(NULL);
        btree->root->keys[0] = key;
        btree->root->num_keys = 1;
        return;
    }

    // Find leaf node where key should be inserted
    struct btree_node *node = btree->root;
    while (!btree_node_is_leaf(node)) {
        int pos = 0;
        while (pos < node->num_keys && node->keys[pos] < key) {
            pos++;
        }
        node = node->children[pos];
    }

    // Find insert position in leaf
    int pos = 0;
    while (pos < node->num_keys && node->keys[pos] < key) {
        pos++;
    }

    // Shift existing keys to make room
    memmove(node->keys + pos + 1, node->keys + pos, (node->num_keys - pos) * sizeof(key_t));

    // Insert new key
    node->keys[pos] = key;
    node->num_keys++;

    // Split if node is full
    if (node->num_keys == BTREE_ORDER - 1) {
        btree_node_split(btree, node);
    }
}

bool btree_contains(struct btree *btree, int key)
{
    if (btree->root == NULL) {
        return false;
    }

    struct btree_node *node = btree->root;
    while (node != NULL) {
        int pos = 0;
        while (pos < node->num_keys && node->keys[pos] < key) {
            pos++;
        }
        if (pos < node->num_keys && node->keys[pos] == key) {
            return true;
        }
        if (btree_node_is_leaf(node)) {
            return false;
        }
        node = node->children[pos];
    }
    return false;
}

static void indent(int depth)
{
    for (int i = 0; i < depth; i++) {
        printf("  ");
    }
}

void btree_node_print(struct btree_node *node, int depth)
{
    if (node == NULL) {
        return;
    }

    indent(depth);
    printf("\\\n");
    for (int i = 0; i < node->num_keys; i++) {
        if (node->children[i] != NULL) {
            btree_node_print(node->children[i], depth + 1);
        }
        indent(depth + 1);
        printf("%d\n", node->keys[i]);
    }
    if (node->children[node->num_keys] != NULL) {
        btree_node_print(node->children[node->num_keys], depth + 1);
    }
    indent(depth);
    printf("/\n");
}

void btree_print(struct btree *btree)
{
    btree_node_print(btree->root, 0);
}

#define N 10000000

int main()
{
    struct btree btree = {0};
    static const int random_seed = 762959903;
    //const int random_seed = time(NULL);
    
    srand(random_seed);
    for (int i = 0; i < N; i++) {
        int random_num = rand();
        assert(!btree_contains(&btree, random_num));
        btree_insert(&btree, random_num);
        assert(btree_contains(&btree, random_num));
    }

    srand(random_seed);
    for (int i = 0; i < N; i++) {
        int random_num = rand();
        assert(btree_contains(&btree, random_num));
    }

    //btree_print(&btree);

    return 0;
}