btree_iterative.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <malloc.h>

//Iterative Relaxed B+Tree
//Just pedagogical. Based off Per Vognsen's B+Tree code
//https://gist.github.com/pervognsen/e7883b3de183fcd601c1edf7f7e9508b
// cl -Zi -Od /INCREMENTAL:NO persistent_btree.c

#define assert(expr) if(!(expr)) { *(char*)0 = 0; }
#define ArrayCount(a) (sizeof(a)/sizeof(a[0]))

//Fuck const forever
#pragma warning(disable : 4090)

//Shitty arena implementation
typedef struct Arena Arena;
struct Arena {
    uint8_t *base;
    uint64_t size;
    uint64_t capacity;
    uint32_t flags;
};

typedef enum ArenaFlags ArenaFlags;
enum ArenaFlags {
    ARENA_FLAGS_NONE = 0,
    ARENA_FLAGS_NO_ZERO = (1 << 0),
};

uint8_t *arena_push(Arena *arena, uint64_t size, uint32_t flags) {
    uint8_t *result = 0;
    uint64_t remaining = arena->capacity - arena->size;
    if(size <= remaining) {
        result = &arena->base[arena->size];
        if(!(flags & ARENA_FLAGS_NO_ZERO)) {
            memset(result, 0, size);
        }
    }
    arena->size += size;
    return result;
}
#define push_array(a, T, c) (T*)arena_push(a, (c)*sizeof(T), 0);
#define push_array_no_zero(a, T, c) (T*)arena_push(a, (c)*sizeof(T), ARENA_FLAGS_NO_ZERO);

Arena *arena_alloc(uint64_t capacity) {
    uint8_t *base = malloc(capacity);
    Arena *arena = (Arena*)base;
    arena->base = base;
    arena->size = sizeof(*arena);
    arena->capacity = capacity;
    arena->flags = 0;
    
    return arena;
}

static Arena *temp_arena = 0;

enum {
    BMAX = 4,
    BCUT = BMAX/2,
    
    BPATH_MAX = 32,
};

//just use a unified representation for both interior and leaf nodes
//for pedagogical simplicity
typedef struct BNode BNode;
struct BNode {
    uint8_t count;
    uint64_t keys[BMAX];
    union {
        BNode *children[BMAX];
        uint64_t vals[BMAX];
        BNode *next_free;
    };
};

typedef struct BTree BTree;
struct BTree {
    BNode *root;
    uint32_t height;
    
    Arena *node_arena;
    BNode *interior_free_list;
    
    Arena *leaf_arena;
    BNode *leaf_free_list;
};

typedef struct BPathSlot BPathSlot;
struct BPathSlot {
    BNode *parent;
    uint32_t child_index;
};

typedef struct BCursor BCursor;
struct BCursor {
    BTree *tree;
    uint32_t height;

    BNode *focus;
    uint32_t child_index;

    int32_t depth;
        
    union {
        BPathSlot path_stack[BPATH_MAX+1];
        struct {
            BPathSlot dummy;
            BPathSlot path[BPATH_MAX];
        };
    };
};

BNode *b_interior_node_alloc(BTree *tree) {
    BNode *node = tree->interior_free_list;
    if(!node) {
        node = push_array_no_zero(tree->node_arena, BNode, 1);
    } else {
        tree->interior_free_list = node->next_free;
    }
    memset(node, 0, sizeof(node[0]));
    return node;
}

BNode *b_leaf_node_alloc(BTree *tree) {
    BNode *node = tree->leaf_free_list;
    if(!node) {
        node = push_array_no_zero(tree->leaf_arena, BNode, 1);
    } else {
        tree->leaf_free_list = node->next_free;
    }
    memset(node, 0, sizeof(node[0]));
    return node;
}

static BTree *btree_alloc(void) {
    Arena *node_arena = arena_alloc(0x100000*sizeof(BNode));
    
    BTree *tree = push_array(node_arena, BTree, 1);
    tree->node_arena = node_arena;
    tree->leaf_arena = arena_alloc(0x100000*sizeof(BNode));
        
    BNode *root = b_leaf_node_alloc(tree);
    tree->root = root;
    
    return tree;
}

void b_copy_interior_node(BNode *dst, BNode *src, uint32_t cut, uint32_t count) {
    dst->count = count;
    memcpy(dst->keys, src->keys + cut, count * sizeof(src->keys[0]));
    memcpy(dst->children, src->children + cut, count * sizeof(src->children[0]));
}

void b_copy_leaf_node(BNode *dst, BNode *src, uint32_t cut, uint32_t count) {
    dst->count = count;
    memcpy(dst->keys, src->keys + cut, count * sizeof(src->keys[0]));
    memcpy(dst->vals, src->vals + cut, count * sizeof(src->vals[0]));
}

static bool go_down(BCursor *cursor, uint32_t child_index) {
    bool result = false;
    BNode *current = cursor->focus;
    if(cursor->depth != cursor->height && child_index < current->count) {
        assert(child_index < BMAX);
        assert(cursor->depth < BPATH_MAX);
        
        cursor->depth += 1;
        BPathSlot *path_slot = &cursor->path[cursor->depth-1];
        
        path_slot->parent = current;
        path_slot->child_index = child_index;
        
        BNode *child = current->children[child_index];

        cursor->focus = child;
        cursor->child_index = child_index;
        
        result = true;
    }
    return result;
}

static bool go_up(BCursor *cursor) {
    bool result = false;

    if(cursor->depth != 0) {
        BPathSlot *path_slot = &cursor->path[cursor->depth-1];
        cursor->depth -= 1;
        
        BNode *current = cursor->focus;
        BNode *parent = path_slot->parent;
        BNode *new_focus = parent;

        cursor->focus = new_focus;
        cursor->child_index = cursor->path[cursor->depth-1].child_index;
    }
    return result;
}

static void b_move_to_depth(BCursor *cursor, int32_t depth) {
    if(depth < cursor->depth) {
        cursor->depth = depth;
        cursor->focus = cursor->path[depth].parent;
        cursor->child_index = cursor->path[depth-1].child_index;
    }
}

static BCursor b_init_cursor(BTree *tree) {
    BCursor cursor = {0};
    cursor.tree = tree;
    cursor.focus = tree->root;
    cursor.height = tree->height;
    
    return cursor;
}

static uint64_t b_get_greatest_key(BNode *node) {
    assert(node->count > 0);
    uint64_t result = node->keys[node->count-1];
    return result;
}

//lower bound
static uint32_t array_search(uint64_t *arr, uint32_t count, uint64_t key) {
    uint32_t result = count;
    for(uint32_t i = 0; i < count; i += 1) {
        if(key <= arr[i]) {
            result = i;
            break;
        }
    }
    return result;
}

#define array_insert(arr, count, index, val) \
    memmove((arr) + (index) + 1, (arr) + (index), ((count) - (index)) * sizeof(arr[0])); \
    arr[index] = val;

#define array_delete(arr, count, index) \
    memmove((arr) + (index), (arr) + (index) + 1, ((count) - (index) - 1) * sizeof(arr[0]));


static bool b_insert(BCursor *cursor, uint64_t key, uint64_t val) {
    BTree *tree = cursor->tree;
    assert(cursor->depth == 0);
    for(uint32_t depth = 0; depth < cursor->height; depth += 1) {
        BNode *node = cursor->focus;
        uint32_t child_index = array_search(node->keys, node->count, key);
        assert(node->count > 0);
        if(child_index == node->count) {
            child_index -= 1;
            node->keys[child_index] = key;
        }
        go_down(cursor, child_index);
    }
    BNode *leaf = cursor->focus;
    uint32_t index = array_search(leaf->keys, leaf->count, key);
    bool found = index < leaf->count && leaf->keys[index] == key;
    
    BNode *new_sibling = 0;
    if(found) {
        leaf->vals[index] = val;
    } else {
        //split leaf node
        if(leaf->count == BMAX) {
            new_sibling = b_leaf_node_alloc(tree);
            b_copy_leaf_node(new_sibling, leaf, BCUT, BMAX - BCUT);
            leaf->count = BMAX - BCUT;
            if(index >= BCUT) {
                leaf = new_sibling;
                index -= BCUT;
            }
        }
        array_insert(leaf->keys, leaf->count, index, key);
        array_insert(leaf->vals, leaf->count, index, val);
        leaf->count += 1;
    }
    BNode *new_child = new_sibling;

    while(cursor->depth != 0) {
        uint32_t child_index = cursor->child_index;
        go_up(cursor);
                
        if(new_child != 0) {
            BNode *new_sibling = 0;
            BNode *node = cursor->focus;
            //split interior node
            if(node->count == BMAX) {
                new_sibling = b_interior_node_alloc(tree);
                b_copy_interior_node(new_sibling, node, BCUT, BMAX - BCUT);
                node->count = BMAX - BCUT;
                if(child_index >= BCUT) {
                    node = new_sibling;
                    child_index -= BCUT;
                }
            }
            node->keys[child_index] = b_get_greatest_key(node->children[child_index]);
            array_insert(node->keys, node->count, child_index + 1, b_get_greatest_key(new_child));
            array_insert(node->children, node->count, child_index + 1, new_child);
            node->count += 1;
            
            new_child = new_sibling;
        }
    }
    //new root
    if(new_child != 0) {
        BNode *old_root = cursor->focus;
        BNode *new_root = b_interior_node_alloc(tree);
        new_root->count = 2;
        new_root->keys[0] = b_get_greatest_key(old_root);
        new_root->keys[1] = b_get_greatest_key(new_child);
        new_root->children[0] = old_root;
        new_root->children[1] = new_child;
        
        cursor->focus = new_root;
        cursor->height += 1;
        
        tree->root = new_root;
        tree->height = cursor->height;
    }
    return found;
}


static bool b_delete(BCursor *cursor, uint64_t key) {
    bool found = false;
    assert(cursor->depth == 0);
    for(uint32_t depth = 0; depth < cursor->height; depth += 1) {
        BNode *node = cursor->focus;
        uint32_t child_index = array_search(node->keys, node->count, key);
        if(child_index == node->count) {
            break;
        }
        go_down(cursor, child_index);
    }
    bool node_empty = false;
    if(cursor->depth == cursor->height) {
        BNode *leaf = cursor->focus;
        uint32_t index = array_search(leaf->keys, leaf->count, key);
        found = index < leaf->count && leaf->keys[index] == key;
        if(found) {
            array_delete(leaf->keys, leaf->count, index);
            array_delete(leaf->vals, leaf->count, index);
            leaf->count -= 1;
            node_empty = leaf->count == 0;
            if(node_empty && cursor->depth != 0) {
                //release node...
            }
        }
    }
    if(found) {
        bool child_empty = node_empty;
        while(cursor->depth != 0) {
            uint32_t child_index = cursor->child_index;
            go_up(cursor);
            
            if(child_empty) {
                bool node_empty = false;

                BNode *node = cursor->focus;
                if(node->count == 1) {
                    if(cursor->depth != 0) {
                        //release node...
                    }
                    node_empty = true;
                } else {
                    array_delete(node->keys, node->count, child_index);
                    array_delete(node->children, node->count, child_index);
                    node->count -= 1;
                }
                child_empty = node_empty;
            }
        }
        bool root_empty = child_empty;
        if(root_empty) {
            BNode *node = cursor->focus;
            node->count = 0;
            cursor->height = 0;
            
            BTree *tree = cursor->tree;
            tree->height = 0;
        }
    }
    b_move_to_depth(cursor, 0);
    return found;
}

static bool b_search(BTree *tree, uint64_t key, uint64_t *val_out) {
    bool found = false;
    BNode *node = tree->root;
    uint32_t height = tree->height;
    for(uint32_t depth = 0; depth <= height; depth += 1) {
        uint32_t index = array_search(node->keys, node->count, key);
        if(index == node->count) {
            break;
        }
        if(depth == height) {
            if(node->keys[index] == key) {
                if(val_out) {
                    val_out[0] = node->vals[index];
                }
                found = true;
            }
            break;
        }
        node = node->children[index];
    }
    
    return found;
}
 
void print_btree(BNode *node) {
    
}

uint64_t murmur64(uint64_t h) {
    h ^= h >> 33;
    h *= 0xff51afd7ed558ccdL;
    h ^= h >> 33;
    h *= 0xc4ceb9fe1a85ec53L;
    h ^= h >> 33;
    return h;
}

int main(int argc, char **argv) {
    temp_arena = arena_alloc(0x4000); //would use scratch arena in practice

    BTree *tree = btree_alloc();
    BCursor cursor = b_init_cursor(tree);
    b_insert(&cursor, 1, 1);
    b_insert(&cursor, 2, 2);
    b_insert(&cursor, 3, 3);
    b_insert(&cursor, 4, 4);
    b_insert(&cursor, 5, 5);
    b_insert(&cursor, 6, 6);
    b_insert(&cursor, 7, 7);
    b_insert(&cursor, 8, 8);
    b_insert(&cursor, 9, 9);
    b_insert(&cursor, 10, 10);
    b_insert(&cursor, 11, 11);
    
    for(uint32_t i = 12; i <= 20; i += 1) {
        b_insert(&cursor, i, i);
    }
        
    b_delete(&cursor, 11);
    b_delete(&cursor, 10);
    b_delete(&cursor, 9);
    b_delete(&cursor, 8);
    
    b_insert(&cursor, 12, 12);
    b_insert(&cursor, 13, 13);
    
    
#if 0
    uint64_t *v = malloc(sizeof(uint64_t)*1000000);
    for(size_t i = 0; i < 1000000; i++) {
        uint64_t h = murmur64(i);
        v[i] = i;//h;
    }
    
    for(uint32_t i = 0; i < 1000000; i += 1) {
        b_insert(&cursor, v[i], v[i]);
    }

#else
    for(uint32_t i = 1; i <= 13; i += 1) {
        bool found = b_search(tree, i, 0);
        bool should_exist = !(i >= 8 && i <= 11);
        assert(found == should_exist);
    }
#endif

    return 0;
}