BTree.cpp

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

typedef uint8_t BIndex;

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

struct BLeafNode {
    BIndex length;
    Key keys[BMAX];
    union {
        Value values[BMAX];
        BNode *alignment_padding;
    };
};

struct BInteriorNode {
    BIndex length;
    Key keys[BMAX];
    union {
        BNode *children[BMAX];
        Value alignment_padding;
    };
}; 

static BNode *BLeafNode_Create(BIndex length, Key *keys, Value *values) {
    BNode *leaf = (BNode *)Allocate(sizeof(BLeafNode));
    leaf->length = length;
    CopyMemory(leaf->keys, keys, length * sizeof(Key));
    CopyMemory(leaf->values, values, length * sizeof(Value));
    return leaf;
}

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

static BNode *BInteriorNode_Create(BIndex length, Key *keys, BNode **children) {
    BNode *node = (BNode *)Allocate(sizeof(BInteriorNode));
    BInteriorNode_Initialize(node, length, keys, children);
    return node;
}

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

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

static BNode *BLeafNode_Add(BNode *leaf, Key key, Value value) {
    BIndex index = Array_Search(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 = BLeafNode_Create(BMAX - BCUT, leaf->keys + BCUT, leaf->values + BCUT);
        leaf->length = BCUT;
        if (index >= BCUT) {
            leaf = new_sibling;
            index -= BCUT;
        }
    }
    Array_Add(leaf->keys, leaf->length, index, key);
    Array_Add(leaf->values, leaf->length, index, value);
    leaf->length++;
    return new_sibling;
}

static BNode *BInteriorNode_Add(BNode *node, Key key, Value value, BIndex height) {
    Assert(height > 0);
    Assert(node->length > 0);
    BIndex index = Array_Search(node->keys, node->length, key);
    if (index == node->length) {
        index--;
        node->keys[index] = key;
    }
    BNode *new_child;
    if (height == 1) {
        new_child = BLeafNode_Add(node->children[index], key, value);
    } else {
        new_child = BInteriorNode_Add(node->children[index], key, value, height - 1);
    }
    BNode *new_sibling = 0;
    if (new_child) {
        if (node->length == BMAX) {
            new_sibling = BInteriorNode_Create(BMAX - BCUT, node->keys + BCUT, node->children + BCUT);
            node->length = BCUT;
            if (index >= BCUT) {
                node = new_sibling;
                index -= BCUT;
            }
        }
        node->keys[index] = BNode_GetGreatestKey(node->children[index]);
        Array_Add(node->keys, node->length, index + 1, BNode_GetGreatestKey(new_child));
        Array_Add(node->children, node->length, index + 1, new_child);
        node->length++;
    }
    return new_sibling;
}

static bool BLeafNode_Remove(BNode *leaf, Key key) {
    BIndex index = Array_Search(leaf->keys, leaf->length, key);
    if (index < leaf->length && leaf->keys[index] == key) {
        Array_Remove(leaf->keys, leaf->length, index);
        Array_Remove(leaf->values, leaf->length, index);
        leaf->length--;
        return leaf->length == 0;
    }
    return false;
}

static bool BInteriorNode_Remove(BNode *node, Key key, BIndex height) {
    Assert(height > 0);
    BIndex index = Array_Search(node->keys, node->length, key);
    if (index == node->length) {
        return false;
    }
    bool child_empty;
    if (height == 1) {
        child_empty = BLeafNode_Remove(node->children[index], key);
    } else {
        child_empty = BInteriorNode_Remove(node->children[index], key, height - 1);
    }
    if (child_empty) {
        if (node->length == 1) {
            return true;
        } else {
            if (height > 1) {
                BInteriorNode_Destroy(node->children[index], height - 1);
            }
            Free(node->children[index]);
            Array_Remove(node->keys, node->length, index);
            Array_Remove(node->children, node->length, index);
            node->length--;
        }
    }
    return false;
}

struct BTree {
    BIndex height;
    BNode root;
};

void BTree_Initialize(BTree *tree) {
    Assert(BMAX >= 4);
    tree->height = 0;
    tree->root.length = 0;
}

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

Value *BTree_Get(BTree *tree, Key key) {
    BIndex height = tree->height;
    BNode *node = &tree->root;
    for (;;) {
        BIndex index = Array_Search(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_Add(BTree *tree, Key key, Value value) {
    BNode *root = &tree->root;
    BNode *new_root_sibling;
    if (tree->height == 0) {
        new_root_sibling = BLeafNode_Add(root, key, value);
    } else {
        new_root_sibling = BInteriorNode_Add(root, key, value, tree->height);
    }
    if (new_root_sibling) {
        Assert(tree->height != BHEIGHT);
        BNode *old_root = BInteriorNode_Create(root->length, root->keys, root->children);
        Key keys[2] = {BNode_GetGreatestKey(old_root), BNode_GetGreatestKey(new_root_sibling)};
        BNode *children[2] = {old_root, new_root_sibling};
        BInteriorNode_Initialize(root, 2, keys, children);
        tree->height++;
    }
}

void BTree_Remove(BTree *tree, Key key) {
    if (tree->height == 0) {
        BLeafNode_Remove(&tree->root, key);
    } else {
        if (BInteriorNode_Remove(&tree->root, key, tree->height)) {
            BInteriorNode_Destroy(&tree->root, tree->height);
            tree->height = 0;
            tree->root.length = 0;
        }
    }
}

struct BPathNode {
    BNode *node;
    BIndex index;
};

struct BIterator {
    Key *key;
    Key *key_start;
    Key *key_end;
    Value *value;
    Value *value_start;
    Value *value_end;
    BIndex height;
    BPathNode path[BHEIGHT];
    BPathNode *head;
};

static void BIterator_SetLeafIndex(BIterator *iterator, BIndex index) {
    Assert(iterator->height == 0);
    BIndex length = iterator->head->node->length;
    Assert(index <= length);
    iterator->key = iterator->head->node->keys + index;
    iterator->key_start = iterator->head->node->keys;
    iterator->key_end = iterator->head->node->keys + length;
    iterator->value = iterator->head->node->values + index;
    iterator->value_start = iterator->head->node->values;
    iterator->value_end = iterator->head->node->values + length;
}

void BIterator_Copy(BIterator *iterator, BIterator *source) {
    CopyMemory(iterator, source, sizeof(BIterator));
}

bool BIterator_Equals(BIterator *iterator1, BIterator *iterator2) {
    return iterator1->key == iterator2->key;
}

bool BIterator_HasItem(BIterator *iterator) {
    return iterator->key != iterator->key_end;
}

bool BIterator_HasPreviousItem(BIterator *iterator) {
    return iterator->key != iterator->key_start;
}

void BIterator_FindStart(BIterator *iterator) {
    while (iterator->height > 0) {
        BPathNode *parent = iterator->head;
        iterator->head++;
        iterator->head->node = parent->node->children[parent->index];
        iterator->head->index = 0;
        iterator->height--;
    }
    BIterator_SetLeafIndex(iterator, 0);
}

void BIterator_FindEnd(BIterator *iterator) {
    while (iterator->height > 0) {
        BPathNode *parent = iterator->head;
        parent->index = parent->node->length - 1;
        iterator->head++;
        iterator->head->node = parent->node->children[parent->index];
        iterator->height--;
    }
    BIterator_SetLeafIndex(iterator, iterator->head->node->length);
}

void BIterator_FindPrevious(BIterator *iterator) {
    while (iterator->head->index == 0) {
        if (iterator->head == iterator->path) {
            BIterator_FindStart(iterator);
            return;
        }
        iterator->head--;
        iterator->height++;
    }
    iterator->head->index--;
    BNode *node = iterator->head->node->children[iterator->head->index];
    iterator->head++;
    iterator->head->node = node;
    iterator->height--;
    BIterator_FindEnd(iterator);
}

void BIterator_FindNext(BIterator *iterator) {
    while (iterator->head->index == iterator->head->node->length) {
        if (iterator->head == iterator->path) {
            BIterator_FindEnd(iterator);
            return;
        }
        iterator->head--;
        iterator->head->index++;
        iterator->height++;
    }
    BIterator_FindStart(iterator);
}

void BIterator_FindNextLeaf(BIterator *iterator) {
    iterator->head->index = iterator->head->node->length;
    BIterator_FindNext(iterator);
}

void BIterator_FindPreviousLeaf(BIterator *iterator) {
    iterator->head->index = 0;
    BIterator_FindPrevious(iterator);
}

void BIterator_FindNextItem(BIterator *iterator) {
    iterator->key++;
    iterator->value++;
    if (!BIterator_HasItem(iterator)) {
        BIterator_FindNextLeaf(iterator);
    }
}

void BIterator_FindPreviousItem(BIterator *iterator) {
    iterator->key--;
    iterator->value--;
    if (!BIterator_HasPreviousItem(iterator)) {
        BIterator_FindPreviousLeaf(iterator);
    }
}

void BIterator_FindEqualOrGreaterItem(BIterator *iterator, Key key) {
    for (;;) {
        BNode *node = iterator->head->node;
        BIndex index = Array_Search(node->keys, node->length, key);
        iterator->head->index = index;
        if (index == node->length) {
            BIterator_FindNext(iterator);
            break;
        }
        if (iterator->height == 0) {
            BIterator_SetLeafIndex(iterator, index);
            break;
        }
        iterator->height--;
        iterator->head++;
        iterator->head->node = node->children[index];
    }
    Assert(!BIterator_HasItem(iterator) || key <= *iterator->key);
}

void BIterator_SetToRoot(BIterator *iterator, BTree *tree) {
    iterator->head = iterator->path;
    iterator->head->node = &tree->root;
    iterator->head->index = 0;
    iterator->height = tree->height;
}

void BTree_FindEnd(BTree *tree, BIterator *iterator) {
    BIterator_SetToRoot(iterator, tree);
    BIterator_FindEnd(iterator);
}

void BTree_FindStart(BTree *tree, BIterator *iterator) {
    BIterator_SetToRoot(iterator, tree);
    BIterator_FindStart(iterator);
}

void BTree_FindEqualOrGreaterItem(BTree *tree, BIterator *iterator, Key key) {
    BIterator_SetToRoot(iterator, tree);
    BIterator_FindEqualOrGreaterItem(iterator, key);
}

void BTree_FindItemsInRange(BTree *tree, BIterator *iterator_start, BIterator *iterator_end, Key key_start, Key key_end) {
    BTree_FindEqualOrGreaterItem(tree, iterator_start, key_start);
    if (key_start < key_end) {
        BTree_FindEqualOrGreaterItem(tree, iterator_end, key_end);
    } else {
        BIterator_Copy(iterator_end, iterator_start);
    }
}