zipper.c

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

//Partially Persistent Red-Black Tree
//Imperative implementation of "Faster, Simpler Red-Black Trees" by Cameron Moy
//Implemented using Zippers
//https://ccs.neu.edu/~camoy/pub/red-black-tree.pdf
//
//https://github.com/zarif98sjs/RedBlackTree-An-Intuitive-Approach
//https://ranger.uta.edu/~weems/NOTES5311/sigcse05.pdf

//Just pedagogical. No proper memory managagement here, would use Arenas in practice
// cl -Zi -Od /INCREMENTAL:NO zipper.c

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

enum {
    LEFT = 0,
    RIGHT = 1,
    NODE_CHILD_COUNT = 2,
};

typedef enum RBColor RBColor;
enum RBColor {
    RED = 0,
    BLACK = 1,
};

typedef struct RBNode RBNode;
struct RBNode {
    uint64_t key;
    uint64_t value;
    RBColor color;
    uint32_t ref_count;
    union {
        RBNode *children[NODE_CHILD_COUNT];
        RBNode *next_free;
    };
};

typedef struct RBTree RBTree;
struct RBTree {
    RBNode **roots;
    uint32_t roots_count;
    uint32_t roots_capacity;
    //Arena *roots_arena;
    //Arena *nodes_arena;
    RBNode *free_list;
};

typedef struct RBPathNode RBPathNode;
struct RBPathNode {
    RBNode *parent;
    uint32_t child_index;
    uint32_t dir;
        
    bool mutable;
    bool do_release;
    
    RBPathNode *prev;
};

typedef struct RBCursor RBCursor;
struct RBCursor {
    bool mutable;
    bool do_release;
    
    //context
    RBNode *focus;
    uint32_t dir;
    
    RBPathNode *path;
    uint32_t depth;
    RBPathNode *free_list;
    //Arena *path_arena;
    
    RBTree *tree;
};

static RBNode rb_nil = { .color = BLACK, .children[LEFT] = &rb_nil, .children[RIGHT] = &rb_nil };

static RBNode *node_acquire(RBNode *node) {
    if(node != &rb_nil) {
        node->ref_count++;
    }
    return node;
}

static void node_release(RBNode **free_list, RBNode *node) {
    if(node != &rb_nil) {
        assert(node->ref_count != 0);
        
        if(--node->ref_count == 0) {
            node->next_free = free_list[0];
            free_list[0] = node;
        }
    }
}

static RBTree *rb_init_tree(void) {
    RBTree *tree = calloc(1, sizeof(RBTree));
    
    return tree;
}

static RBNode *rb_alloc_node(RBTree *tree) {
    RBNode *node = tree->free_list;
    if(!node) {
        node = malloc(sizeof(RBNode));
    } else {
        tree->free_list = node->next_free;
    }
    memset(node, 0, sizeof(*node));
    return node;
}

static RBNode *node_clone(RBTree *tree, RBNode *src) {
    RBNode *node = rb_alloc_node(tree);
    node->key = src->key;
    node->value = src->value;
    node->color = src->color;
    node->ref_count = 1;
    for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
        node->children[i] = node_acquire(src->children[i]);
    }
    return node;
}

static bool go_down(RBCursor *cursor, uint32_t child_index) {
    bool result = false;
    RBNode *current = cursor->focus;
        
    if(current != &rb_nil) {
        assert(child_index < NODE_CHILD_COUNT);
        RBNode *child = current->children[child_index];
        bool mutable = cursor->mutable && child->ref_count == 1;
        
        RBPathNode *path_node = cursor->free_list;
        if(!path_node) {
            path_node = malloc(sizeof(RBPathNode));
        } else {
            cursor->free_list = path_node->prev;
        }
        memset(path_node, 0, sizeof(*path_node));
        
        path_node->parent = current;
        path_node->child_index = child_index;
        path_node->dir = cursor->dir;
        path_node->mutable = cursor->mutable;
        path_node->do_release = cursor->do_release;

        path_node->prev = cursor->path;
        cursor->path = path_node; 
        cursor->depth += 1;
        
        cursor->focus = child;
        cursor->mutable = mutable;
        cursor->dir = child_index;
        cursor->do_release = true;
        
        result = true;
    }
    return result;
}

static bool go_up(RBCursor *cursor) {
    bool result = false;
    if(cursor->depth != 0) {
        RBPathNode *path_node = cursor->path;
        cursor->path = path_node->prev;
        cursor->depth -= 1;
        
        RBNode *current = cursor->focus;
        RBNode *parent = path_node->parent;
        RBNode *new_focus = parent;
        uint32_t child_index = path_node->child_index;
        bool mutable = path_node->mutable;
        
        RBNode *original_child = parent->children[child_index];
        if(original_child != current) {
            RBTree *tree = cursor->tree;
            if(cursor->mutable && !mutable) {
                RBNode *new_parent = node_clone(tree, parent);
                new_focus = new_parent;
                mutable = true;
            }
            
            new_focus->children[child_index] = current;
            if(cursor->do_release) {
                node_release(&tree->free_list, original_child);
            }
        }
 
        cursor->focus = new_focus;
        cursor->mutable = mutable;
        cursor->dir = path_node->dir;
        cursor->do_release = path_node->do_release;

        path_node->prev = cursor->free_list;
        cursor->free_list = path_node;
        
        result = true;
    }
    return result;
}

static RBCursor rb_init_cursor(RBTree *tree) {
    RBCursor cursor = { .tree = tree };
    RBNode *root = &rb_nil;
    if(tree->roots_count != 0) {
        root = tree->roots[tree->roots_count-1];
    }        
    cursor.focus = root;
    
    return cursor;
}

static RBNode *rb_commit(RBCursor *cursor) {
    while(cursor->depth != 0) {
        go_up(cursor);
    }
    RBNode *root = cursor->focus;
    //tree->roots[tree->roots_count] = root;
    //tree->roots_count += 1;
    cursor->mutable = false;
    return root;
}

static RBNode *mutable_focus(RBCursor *cursor) {
    RBNode *node = cursor->focus;
    if(!cursor->mutable) {
        assert(node != &rb_nil);
        node = node_clone(cursor->tree, node);
        cursor->focus = node;
        cursor->mutable = true;
    }
    return node;
}

static void update_child(RBCursor *cursor, uint32_t child_index, RBNode *child, bool do_release) {
    RBNode *node = mutable_focus(cursor);
    RBNode *original_child = node->children[child_index];
    
    if(child != original_child) {
        node->children[child_index] = child;
        if(do_release) {
            RBTree *tree = cursor->tree;
            node_release(&tree->free_list, original_child);
        }
    }
}

static void update_color(RBCursor *cursor, RBColor color) {
    RBNode *node = mutable_focus(cursor);
    node->color = color;
}

static void update_value(RBCursor *cursor, uint64_t value) {
    RBNode *node = mutable_focus(cursor);
    node->value = value;
}

static void update_key(RBCursor *cursor, uint64_t key) {
    RBNode *node = mutable_focus(cursor);
    node->key = key;
}

static void rotate_left(RBCursor *cursor) {
    RBNode *h = mutable_focus(cursor);
    RBColor color = h->color;
    RBNode *x = h->children[RIGHT];
    bool x_mut = x->ref_count == 1;
    update_child(cursor, RIGHT, x->children[LEFT], !x_mut);
    update_color(cursor, RED);
    cursor->focus = x;
    cursor->mutable = x_mut;
    cursor->do_release = false;
    update_child(cursor, LEFT, h, false);
    update_color(cursor, color);
}

static void rotate_right(RBCursor *cursor) {  
    RBNode *h = mutable_focus(cursor);
    RBColor color = h->color;
    RBNode *x = h->children[LEFT];
    bool x_mut = x->ref_count == 1;
    update_child(cursor, LEFT, x->children[RIGHT], !x_mut);
    update_color(cursor, RED);
    cursor->focus = x;
    cursor->mutable = x_mut;
    cursor->do_release = false;
    update_child(cursor, RIGHT, h, false);
    update_color(cursor, color);
}

static void 
rb_balance_node(RBCursor *cursor, RBNode *x, RBNode *y, RBNode *z, RBNode *b, RBNode *c, RBColor root_color, bool mutable) {
    cursor->focus = y;
    cursor->mutable = mutable;
    assert(mutable); //???
    bool y_mut = y->ref_count == 1;
    assert(mutable == y_mut);
    update_color(cursor, root_color);
    update_child(cursor, LEFT, x, false);
    update_child(cursor, RIGHT, z, false);
    cursor->do_release = !mutable;
    go_down(cursor, LEFT);
    bool x_mut = x->ref_count == 1;
    update_color(cursor, BLACK);
    update_child(cursor, RIGHT, b, false);
    cursor->do_release = !x_mut;
    go_up(cursor);
    go_down(cursor, RIGHT);
    bool z_mut = z->ref_count == 1;
    update_color(cursor, BLACK);
    update_child(cursor, LEFT, c, false);
    cursor->do_release = !z_mut;
    go_up(cursor);    
}

static bool rb_balance(RBCursor *cursor, RBColor root_color) {
    RBNode *node = cursor->focus;

    bool result = false;
    if(node->children[LEFT]->color == RED && node->children[LEFT]->children[LEFT]->color == RED) {
        RBNode *z = node;
        RBNode *y = z->children[LEFT];
        RBNode *x = y->children[LEFT];
        RBNode *b = x->children[RIGHT];
        RBNode *c = y->children[RIGHT];
        bool mutable = cursor->mutable && y->ref_count == 1;
        
        rb_balance_node(cursor, x, y, z, b, c, root_color, mutable);
    } else if(node->children[LEFT]->color == RED && node->children[LEFT]->children[RIGHT]->color == RED) {
        RBNode *z = node;
        RBNode *x = z->children[LEFT];
        RBNode *y = x->children[RIGHT];
        RBNode *b = y->children[LEFT];
        RBNode *c = y->children[RIGHT];
        bool mutable = cursor->mutable && x->ref_count == 1 && y->ref_count == 1;
        
        rb_balance_node(cursor, x, y, z, b, c, root_color, mutable);
    } else if(node->children[RIGHT]->color == RED && node->children[RIGHT]->children[LEFT]->color == RED) {
        RBNode *x = node;
        RBNode *z = x->children[RIGHT];
        RBNode *y = z->children[LEFT];
        RBNode *b = y->children[LEFT];
        RBNode *c = y->children[RIGHT];
        bool mutable = cursor->mutable && z->ref_count == 1 && y->ref_count == 1;
        
        rb_balance_node(cursor, x, y, z, b, c, root_color, mutable);
    } else if(node->children[RIGHT]->color == RED && node->children[RIGHT]->children[RIGHT]->color == RED) {
        RBNode *x = node;
        RBNode *y = x->children[RIGHT];
        RBNode *z = y->children[RIGHT];
        RBNode *b = y->children[LEFT];
        RBNode *c = z->children[LEFT];
        bool mutable = cursor->mutable && y->ref_count == 1;
        
        rb_balance_node(cursor, x, y, z, b, c, root_color, mutable);
    } else {
        result = true;
    }
    return result;
}

static bool rb_insert(RBCursor *cursor, uint64_t key, uint64_t value) {
    bool found = false;
    while(cursor->focus != &rb_nil) {
        RBNode *node = cursor->focus;
        
        if(key < node->key) {
            go_down(cursor, LEFT);
        } else if(key > node->key) {
            go_down(cursor, RIGHT);
        } else {
            update_value(cursor, value);
            found = true;
            break;
        }
    }
    if(!found) {
        RBNode *new_node = rb_alloc_node(cursor->tree);
        new_node->key = key;
        new_node->value = value;
        new_node->color = RED;
        new_node->ref_count = 1;
        
        for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
            new_node->children[i] = &rb_nil;
        }
        
        cursor->focus = new_node;
        cursor->mutable = true;
    }
    bool do_balance = true;
    while(cursor->depth != 0) {
        go_up(cursor);
        if(do_balance) {
            bool no_match = rb_balance(cursor, RED);
            do_balance = !(no_match && cursor->focus->color == BLACK);
        }
    }
    update_color(cursor, BLACK);
    
    return found;
}

static bool rb_equalize(RBCursor *cursor, uint32_t dir) {
    bool deficit = false;
    uint32_t begin_d = cursor->depth;
    uint32_t opposite = !dir;
    while(true) {
        RBNode *n = cursor->focus;
        if(n->children[opposite]->color == RED) {
            if(dir == LEFT) {
                rotate_left(cursor);
            } else {
                rotate_right(cursor);
            }
            update_color(cursor, BLACK);
            go_down(cursor, dir);
        } else {
            break;
        }
    }
    RBNode *n = cursor->focus;
    if(n->children[opposite] != &rb_nil) {
        go_down(cursor, opposite);
        update_color(cursor, RED);
        go_up(cursor);
    }
    RBColor root_color = n->color;
    bool no_match = rb_balance(cursor, root_color);
    if(no_match) {
        if(n->color == RED) {
            update_color(cursor, BLACK);
        } else {
            deficit = true;
        }
    }
    for(uint32_t d = cursor->depth; d > begin_d; d--) {
        go_up(cursor);
    }
    return deficit;
}

static bool rb_del_min(RBCursor *cursor, uint64_t *min_key_out, uint64_t *min_val_out) {
    bool deficit = false;
    uint32_t begin_d = cursor->depth;
    while(true) {
        RBNode *n = cursor->focus;
        if(n->children[LEFT] != &rb_nil) {
            go_down(cursor, LEFT);
        } else {
            break;
        }
    }
    RBNode *n = cursor->focus;
    *min_key_out = n->key;
    *min_val_out = n->value;
    RBNode *replacement = n->children[RIGHT];
    cursor->focus = replacement;
    cursor->mutable &= replacement->ref_count == 1;
    if(n->color == BLACK) {
        if(replacement->color == RED) {
            update_color(cursor, BLACK);
        } else {
            deficit = true;
        }
    }
    cursor->mutable = true;
    for(uint32_t d = cursor->depth; d > begin_d; d--) {
        go_up(cursor);
        if(deficit) {
            deficit = rb_equalize(cursor, LEFT);
        }
    }
    return deficit;
}

static bool rb_delete(RBCursor *cursor, uint64_t key) {
    bool found = false;
    while(cursor->focus != &rb_nil) {
        RBNode *node = cursor->focus;
        
        if(key < node->key) {
            go_down(cursor, LEFT);
        } else if(key > node->key) {
            go_down(cursor, RIGHT);
        } else {
            found = true;
            break;
        }
    }
    bool deficit = false;
    if(found) {
        RBNode *node = cursor->focus;
        if(node->children[RIGHT] == &rb_nil) {
            RBNode *replacement = node->children[LEFT];
            cursor->focus = replacement;
            cursor->mutable &= replacement->ref_count == 1;
            if(node->color == BLACK) {
                if(replacement->color == RED) {
                    update_color(cursor, BLACK);
                } else {
                    deficit = true;
                }
            }
            cursor->mutable = true;
        } else {
            uint64_t successor_key = 0;
            uint64_t successor_value = 0;
            go_down(cursor, RIGHT);
            deficit = rb_del_min(cursor, &successor_key, &successor_value);
            go_up(cursor);
            update_key(cursor, successor_key);
            update_value(cursor, successor_value);
            if(deficit) {
                deficit = rb_equalize(cursor, RIGHT);
            }
        }
    }
   while(cursor->depth != 0) {
        uint32_t dir = cursor->dir;
        go_up(cursor);
        if(deficit) {
            deficit = rb_equalize(cursor, dir);
        }
    }
    if(cursor->focus != &rb_nil) {
        update_color(cursor, BLACK);
    }
    return found;
}


static bool rb_search(RBNode *root, uint64_t key, uint64_t *val_out) {
    bool found = false;
    RBNode *n = root;
    while(n != &rb_nil) {
        if(key < n->key) {
            n = n->children[LEFT];
        } else if(key > n->key) {
            n = n->children[RIGHT];
        } else {
            found = true;
            if(val_out) {
                *val_out = n->value;
            }
            break;
        }
    }
    return found;
}



static void print_rb_tree_recurse(RBNode *node) {
    if(node != &rb_nil) {
        char *extra = node->color == RED ? ", shape=doublecircle, color=red" : "";
        printf("\tn_%lld_%p [label=\"%lld\\nrc: %d\"%s];\n", node->key, node, node->key, node->ref_count, extra);
        for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
            RBNode *child = node->children[i];        
            if(child != &rb_nil) {
                char *extra = child->color == RED ? ", splines=ortho, minlen=0" : "";
                char *color = child->color == RED ? "red" : "black";
                uint32_t penwidth = child->color == RED ? 5 : 1;
                printf("\t" "n_%lld_%p -> n_%lld_%p [color=\"%s\", penwidth=%d, dir=none%s];\n", 
                    node->key, node, child->key, child, color, penwidth, extra);
                //print_rb_tree_recurse(child);
            }
        }
        for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
            RBNode *child = node->children[i];        
            if(child != &rb_nil) {
                print_rb_tree_recurse(child);
            }
        }
    }
}

static void print_rb_tree(RBNode *node) {
    printf("strict digraph BST {\n");
    print_rb_tree_recurse(node);
    printf("}\n");
}

int main() {
    RBTree *tree = rb_init_tree();
    
    RBCursor cursor = rb_init_cursor(tree);
    
    RBNode *root1 = &rb_nil;
    RBNode *root2 = &rb_nil;
    
    
#define TEST 3
#if TEST==0
    for(size_t i = 0; i < 100; i++) {
        rb_insert(&cursor, i, i);
    }

    for(size_t i = 20; i < 80; i++) {
        rb_delete(&cursor, i);
    }
    root1 = rb_commit(&cursor);
    
    for(uint32_t i = 0; i < 100; i++) {
        bool should_exist = !(i >= 20 && i < 80);
        bool found = rb_search(root1, i, 0);
        assert(found == should_exist);
    }
    
#elif TEST==1
    uint32_t mid = 100;
    uint32_t end = 200;
    
    for(size_t i = 0; i < 100; i++) {
        rb_insert(&cursor, i, i);
    }
    root1 = rb_commit(&cursor);
    
    for(uint32_t i = 20; i < 80; i++) {
        bool del_found = rb_delete(&cursor, i);
        assert(del_found);
    }
    root2 = rb_commit(&cursor);
    
    for(uint32_t i = 0; i < 100; i++) {
        bool should_exist = true;
        bool found = rb_search(root1, i, 0);
        
        if(found != should_exist) {
            printf("NOT FOUND = %d\n", i);
        }
        assert(found == should_exist);
    }
    
    for(uint32_t i = 0; i < 100; i++) {
        bool should_exist = !(i >= 20 && i < 80);
        bool found = rb_search(root2, i, 0);
        assert(found == should_exist);
    }
#elif TEST==2
    
    uint64_t items[] = {55,33,4,878,2323,1,585,12,0};

    for(size_t i = 0; i < ArrayCount(items); i++) {
        rb_insert(&cursor, items[i], 0);
    }


    root1 = rb_commit(&cursor);

    rb_insert(&cursor, 8, 8);
    
    root2 = rb_commit(&cursor);
    
    for(size_t i = 0; i < ArrayCount(items); i++) {
        bool r1_found = rb_search(root1, items[i], 0);
        bool r2_found = rb_search(root2, items[i], 0);
        
        assert(r1_found);
        assert(r2_found);
    }
    assert(rb_search(root2, 8, 0));
    
#elif TEST==3
    
    //for(size_t i = 0; i < 100; i++) {
    //    rb_insert(&cursor, i*2+1, 0);
    //}
    //
    //for(size_t i = 100; i != 0; i--) {
    //    rb_insert(&cursor, i*2-1, 0);
    //}
    
    for(size_t i = 0; i < 100; i++) {
        rb_insert(&cursor, i, 0);
    }
    
    root1 = rb_commit(&cursor);

    for(size_t i = 100; i != 50; i--) {
        rb_delete(&cursor, i);
    }
    
    root2 = rb_commit(&cursor);
    
#endif
#undef TEST

    bool cluster = false;

    printf("strict digraph BST {\n");
    if(cluster) printf("subgraph cluster_0 {\n");
    print_rb_tree_recurse(root1);
    if(cluster) printf("}\n");
    if(cluster) printf("subgraph cluster_2 {\n");
    print_rb_tree_recurse(root2);
    if(cluster) printf("}\n");
    printf("}\n");

    return 0;
}