B-tree

From vegard.wiki
Revision as of 17:17, 29 October 2020 by Vegard (talk | contribs) (make it extra clear that I'm not the author)
(diff) ← Older revision | Latest revision (diff) | Newer revision → (diff)
Jump to navigation Jump to search

By Per Vognsen:

#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define INLINE inline
#define CopyMemory memcpy
#define Allocate malloc
#define Free free
#define Assert assert

typedef unsigned int Key;
typedef unsigned long Value;

// find where to insert the given key to maintain the sorted array
uint32_t SearchKeys(Key *keys, uint32_t length, Key key)
{
    for (uint32_t i = 0; i < length; ++i) {
        if (key <= keys[i])
            return i;
    }

    return length;
}

template<typename T>
void Array_Insert(T *array, uint32_t length, uint32_t index, T value)
{
    // shift everything up
    for (uint32_t i = length ; i-- > index; )
        array[i + 1] = array[i];

    array[index] = value;
}

template<typename T>
void Array_Delete(T *array, uint32_t length, uint32_t index)
{
    // shift everything down
    for (uint32_t i = index + 1; i < length; ++i)
        array[i - 1] = array[i];
}

// --- https://gist.github.com/pervognsen/2d48ef9757ee3fd579179239febc817e
// Per Vognsen, public domain

enum { BMAX = 32, BMIN = BMAX / 2, BHEIGHT = 6 };

struct BNode {
    uint32_t length;
    Key keys[BMAX];
    union {
        BNode *children[BMAX];
        Value values[BMAX];
    };
};

static void BNode_Initialize(BNode *node, uint32_t length, Key *keys, void *children) {
    node->length = length;
    CopyMemory(node->keys, keys, length * sizeof(Key));
    CopyMemory(node->children, children, length * sizeof(BNode *));
}

static BNode *BNode_Create(uint32_t length, Key *keys, void *children) {
    BNode *node = (BNode *)Allocate(sizeof(BNode));
    BNode_Initialize(node, length, keys, children);
    return node;
}

static void BNode_Destroy(BNode *node, uint32_t height) {
    for (uint32_t index = 0; index < node->length; index++) {
        if (height > 1) {
            BNode_Destroy(node->children[index], height - 1);
        }
        Free(node->children[index]);
    }
}

static INLINE Key BNode_GetMaxKey(BNode *node) {
    Assert(node->length > 0);
    return node->keys[node->length - 1];
}

static BNode *BNodeLeaf_Insert(BNode *leaf, Key key, Value value) {
    uint32_t index = SearchKeys(leaf->keys, leaf->length, key);
    if (index < leaf->length && leaf->keys[index] == key) {
        leaf->values[index] = value;
        return 0;
    }
    BNode *new_sibling = 0;
    if (leaf->length == BMAX) {
        new_sibling = BNode_Create(BMIN, leaf->keys + BMIN, leaf->values + BMIN);
        leaf->length = BMIN;
        if (index >= BMIN) {
            leaf = new_sibling;
            index -= BMIN;
        }
    }
    Array_Insert(leaf->keys, leaf->length, index, key);
    Array_Insert(leaf->values, leaf->length, index, value);
    leaf->length++;
    return new_sibling;
}

static BNode *BNode_Insert(BNode *node, Key key, Value value, uint32_t height) {
    Assert(height > 0);
    Assert(node->length > 0);
    uint32_t index = SearchKeys(node->keys, node->length, key);
    if (index == node->length) {
        index--;
        node->keys[index] = key;
    }
    BNode *new_child;
    if (height == 1) {
        new_child = BNodeLeaf_Insert(node->children[index], key, value);
    } else {
        new_child = BNode_Insert(node->children[index], key, value, height - 1);
    }
    BNode *new_sibling = 0;
    if (new_child) {
        if (node->length == BMAX) {
            new_sibling = BNode_Create(BMIN, node->keys + BMIN, node->children + BMIN);
            node->length = BMIN;
            if (index >= BMIN) {
                node = new_sibling;
                index -= BMIN;
            }
        }
        node->keys[index] = BNode_GetMaxKey(node->children[index]);
        Array_Insert(node->keys, node->length, index + 1, BNode_GetMaxKey(new_child));
        Array_Insert(node->children, node->length, index + 1, new_child);
        node->length++;
    }
    return new_sibling;
}

static bool BNodeLeaf_Delete(BNode *leaf, Key key) {
    uint32_t index = SearchKeys(leaf->keys, leaf->length, key);
    if (index < leaf->length && leaf->keys[index] == key) {
        Array_Delete(leaf->keys, leaf->length, index);
        Array_Delete(leaf->values, leaf->length, index);
        leaf->length--;
        return leaf->length == 0;
    }
    return false;
}

static void BNode_Delete(BNode *node, Key key, uint32_t height) {
    Assert(height > 0);
    uint32_t index = SearchKeys(node->keys, node->length, key);
    if (index < node->length) {
        if (height == 1) {
            if (BNodeLeaf_Delete(node->children[index], key) && node->length > 1) {
                Free(node->children[index]);
                Array_Delete(node->keys, node->length, index);
                Array_Delete(node->children, node->length, index);
                node->length--;
            }
        } else {
            BNode_Delete(node->children[index], key, height - 1);
        }
    }
}

struct BTree {
    uint32_t height;
    BNode root;
};

void BTree_Initialize(BTree *tree) {
    Assert(BMAX == 2 * BMIN);
    Assert(sizeof(BNode *) == sizeof(Value));
    tree->height = 0;
    tree->root.length = 0;
}

void BTree_Destroy(BTree *tree) {
    if (tree->height > 0) {
        BNode_Destroy(&tree->root, tree->height);
    }
}

Value *BTree_Find(BTree *tree, Key key) {
    uint32_t height = tree->height;
    BNode *node = &tree->root;
    for (;;) {
        uint32_t index = SearchKeys(node->keys, node->length, key);
        if (index == node->length) {
            return 0;
        }
        if (height == 0) {
            return (node->keys[index] == key) ? (node->values + index) : 0;
        }
        height--;
        node = node->children[index];
    }
}

void BTree_Insert(BTree *tree, Key key, Value value) {
    BNode *root = &tree->root;
    BNode *new_root_sibling;
    if (tree->height == 0) {
        new_root_sibling = BNodeLeaf_Insert(root, key, value);
    } else {
        new_root_sibling = BNode_Insert(root, key, value, tree->height);
    }
    if (new_root_sibling) {
        BNode *old_root = BNode_Create(root->length, root->keys, root->children);
        Key keys[2] = {BNode_GetMaxKey(old_root), BNode_GetMaxKey(new_root_sibling)};
        BNode *children[2] = {old_root, new_root_sibling};
        BNode_Initialize(root, 2, keys, children);
        tree->height++;
    }
}

void BTree_Delete(BTree *tree, Key key) {
    if (tree->height == 0) {
        BNodeLeaf_Delete(&tree->root, key);
    } else {
        BNode_Delete(&tree->root, key, tree->height);
    }
}

See also