zipper.c

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

//https://sedgewick.io/wp-content/themes/sedgewick/papers/2008LLRB.pdf
//Partially Persistent Left-leaning Red-Black Tree

//Just pedagogical. No proper memory managagement here, would use Arenas in practice
// cl -Zi -Od /INCREMENTAL:NO zipper.c
enum {
    LEFT = 0,
    RIGHT = 1,
    NODE_CHILD_COUNT = 2,
};

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

typedef struct Node Node;
struct Node {
    uint64_t key;
    union {
        //char value_u8;
        uint64_t value;
    };
    Color color;
    uint32_t ref_count;
    Node *children[NODE_CHILD_COUNT];
};

typedef struct Tree Tree;
struct Tree {
    Node **roots;
    uint32_t roots_count;
    //Arena *roots_arena;
    //Arena *nodes_arena;
    //Node **free_list;
};

typedef struct ZipperNode ZipperNode;
struct ZipperNode {
    Node *parent;
    uint32_t child_index;
        
    bool dirty;
    
    ZipperNode *prev;
};

typedef struct Zipper Zipper;
struct Zipper {
    bool dirty;
    
    //context
    Node *focus;
    
    ZipperNode *path;
    uint32_t path_count;
};

static Node *node_acquire(Node *node) {
    if(node == 0) return 0;
    node->ref_count++;
    return node;
}

static void node_release(Node *node) {
    if(node == 0) return;
    //assert(node->ref_count != 0);
    
    if(--node->ref_count == 0) {
        for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
            node_release(node->children[i]);
        }
        //free(node);
    }
}

static Node node_nil = {0};

static Node *node_clone(Node *src) {
    Node *node = malloc(sizeof(Node));
    node->key = src->key;
    node->value = src->value;
    node->ref_count = 1;//0;
    node->color = src->color;
    for(uint32_t i = 0; i < NODE_CHILD_COUNT; i++) {
        node->children[i] = node_acquire(src->children[i]);
    }
    return node;
}

static bool is_red(Node *node) {
    bool result = node && node->color == RED;
    return result;
}

static void go_down(Zipper *zipper, uint32_t child_index) {
    Node *current = zipper->focus;
    
    bool dirty = false;
    
    assert(child_index < NODE_CHILD_COUNT);
    Node *child = current ? current->children[child_index] : 0;
    if(zipper->dirty && child && child->ref_count == 1) {
        dirty = true;
    }
    
    ZipperNode *path_node = calloc(1, sizeof(ZipperNode));
    path_node->parent = current;
    path_node->child_index = child_index;
    path_node->dirty = zipper->dirty;

    path_node->prev = zipper->path;
    zipper->path = path_node; 
    zipper->path_count++;
    
    zipper->focus = child;
    zipper->dirty = dirty;
}

static void go_up(Zipper *zipper) {
    if(zipper->path != 0) {
        ZipperNode *path_node = zipper->path;
        zipper->path = path_node->prev;
        zipper->path_count--;
        
        Node *current = zipper->focus;
        Node *parent = path_node->parent;
        Node *new_focus = parent;
        uint32_t child_index = path_node->child_index;
        bool dirty = path_node->dirty;
        
        Node *original_child = parent ? parent->children[child_index] : 0;
        if(original_child != current) {
            if(zipper->dirty && !dirty) {
                Node *new_parent = node_clone(parent);
                new_focus = new_parent;
                dirty = true;
            }
            new_focus->children[child_index] = current;
        }
 
        zipper->focus = new_focus;
        zipper->dirty = dirty;
        
        free(path_node);
    }
}

static Node *commit(Zipper *zipper) {
    while(zipper->path != 0) {
        go_up(zipper);
    }
    Node *root = zipper->focus;
    return root;
}

static void update_child(Zipper *zipper, uint32_t child_index, Node *child) {
    Node *node = zipper->focus;
    if(!zipper->dirty) {
        node = node_clone(node);
        zipper->focus = node;
        zipper->dirty = true;
    }
    node->children[child_index] = child;
}

static void update_color(Zipper *zipper, Color color) {
    Node *node = zipper->focus;
    if(!zipper->dirty) {
        node = node_clone(node);
        zipper->focus = node;
        zipper->dirty = true;
    }
    node->color = color;
}

static void update(Zipper *zipper, uint64_t value) {
    Node *node = zipper->focus;
    if(!zipper->dirty) {
        node = node_clone(node);
        zipper->focus = node;
        zipper->dirty = true;
    }
    node->value = value;
}

static void rotate_left(Zipper *zipper) {
    Node *h = zipper->focus;
    Color color = h->color;
    Node *x = h->children[RIGHT];
    update_child(zipper, RIGHT, x->children[LEFT]);
    update_color(zipper, RED);
    zipper->focus = x;
    zipper->dirty = x->ref_count == 1;
    update_child(zipper, LEFT, h);
    update_color(zipper, color);
}

static void rotate_right(Zipper *zipper) {  
    Node *h = zipper->focus;
    Color color = h->color;
    Node *x = h->children[LEFT];
    update_child(zipper, LEFT, x->children[RIGHT]);
    update_color(zipper, RED);
    zipper->focus = x;
    zipper->dirty = x->ref_count == 1;
    update_child(zipper, RIGHT, h);
    update_color(zipper, color);
}

static void flip_color(Zipper *zipper) {
    go_down(zipper, LEFT);
    update_color(zipper, !zipper->focus->color);
    go_up(zipper);
    go_down(zipper, RIGHT);
    update_color(zipper, !zipper->focus->color);
    go_up(zipper);
    update_color(zipper, !zipper->focus->color);
}

static void insert(Zipper *zipper, uint64_t key, uint64_t value) {
    while(zipper->focus != 0) {
        Node *node = zipper->focus;
        
        if(is_red(node->children[LEFT]) && is_red(node->children[RIGHT])) {
            flip_color(zipper);
        }
        
        if(key < node->key) {
            go_down(zipper, LEFT);
        } else if(key > node->key) {
            go_down(zipper, RIGHT);
        } else {
            assert(!"moo");
            return;
        }
    }
    Node *new_node = calloc(1, sizeof(Node));
    new_node->key = key;
    new_node->value = value;
    new_node->color = RED;
    new_node->ref_count = 1;
    
    zipper->focus = new_node;
    zipper->dirty = true;
    
    while(zipper->path != 0) {
        go_up(zipper);
        Node *node = zipper->focus;
        
        //balancing
        if(is_red(node->children[RIGHT]) && !is_red(node->children[LEFT])) {
            rotate_left(zipper);
        }
        
        node = zipper->focus;
        if(is_red(node->children[LEFT]) && is_red(node->children[LEFT]->children[LEFT])) {
            rotate_right(zipper);
        }
    }
    update_color(zipper, BLACK);
}

static Zipper init_zipper(Node *root) {
    Zipper zipper = { .focus = root };
    return zipper;
}


static void rb_insert(Tree *tree, uint64_t key, uint64_t value) {


}


int main() {
    //                         A
    //                 /                  \
    //                B                     F
    //           /         \               /    
    //          D           E            C      
    
    Node *root = 0;
    Zipper zipper = init_zipper(0); //A
    
#if 1
    for(uint32_t i = 0; i < 10; i++) {
        insert(&zipper, i, i);
    }
#else
    
    insert(&zipper, 50, 'A');
    insert(&zipper, 25, 'B');
    insert(&zipper, 75, 'C');
    insert(&zipper, 12, 'D');
    insert(&zipper, 30, 'E');
    insert(&zipper, 83, 'F');
#endif
    root = commit(&zipper); //A
    
    
    {
        Zipper zipper = init_zipper(root);
        
        for(uint32_t i = 10; i < 20; i++) {
            insert(&zipper, i, i);
        }
        Node *root2 = commit(&zipper);
        
        int lol = 3;
    }
    
    
    int lol = 3;
    
/*
    Zipper zipper = init_zipper(root); //A
    
    //update(&zipper, 'W'); //A->W
    //update(&zipper, 'X'); //W->X
    
    go_down(&zipper, LEFT); //B
    
    update(&zipper, 'P');
    
    go_down(&zipper, LEFT); //D
    
    update(&zipper, 'Y');   //D->Y
    
    go_up(&zipper); //B'
    
    go_down(&zipper, RIGHT); //E
    
    update(&zipper, 'Z'); //Z
    
    Node *new_root = commit(&zipper); //X
*/
    

    //                         X
    //                 /                  \
    //                B'                    C
    //           /         \                   \
    //          Y           Z                   F


    return 0;
}