zipper_per.c

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

//Just pedagogical. No proper memory managagement here, would use Arenas in practice

#define NODE_CHILD_COUNT 2

typedef struct Node Node;
struct Node {
    uint64_t key;
    uint64_t value;
    uint32_t ref_count;
    Node *children[NODE_CHILD_COUNT];
};

typedef struct Tree Tree;
struct Tree {
    Node *root;
    uint32_t node_count;
};

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

typedef struct Zipper Zipper;
struct Zipper {
    bool mutable;
    
    //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;
    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_clone(Node *src) {
    Node *node = malloc(sizeof(Node));
    node->key = src->key;
    node->value = src->value;
    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 Node **bst_search(Node **root_ptr, uint64_t key) {
    Node **node_ptr = root_ptr;
    
    while((*node_ptr)) {
        Node *node = *node_ptr;
        
        if(key < node->key) {
            node_ptr = &node->children[0];
        } else if(key > node->key) {
            node_ptr = &node->children[1];
        } else {
            return node_ptr;
        }
    }
    return node_ptr;
}

static Node *bst_insert(Node **root_ptr, uint64_t key, uint64_t value) {
    Node **best_match_ptr = bst_search(root_ptr, key);
    Node *best_match = *best_match_ptr;
    if(best_match && best_match->key == key) {
        return 0;
    }
    Node *node = calloc(1, sizeof(Node));
    node->ref_count = 1;
    node->key = key;
    node->value = value;
    
    *best_match_ptr = node;
        
    return node;
}

static void go_down(Zipper *zipper, uint32_t child_index) {
    Node *parent = zipper->focus;
    
    assert(child_index < NODE_CHILD_COUNT);
    Node *child = parent->children[child_index];
    if(zipper->mutable && child->ref_count != 1) {
        child = node_clone(child);
    }

    ZipperNode *path_node = calloc(1, sizeof(ZipperNode));
    path_node->parent = parent;
    path_node->child_index = child_index;

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

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;
        
        Node *original_child = parent->children[path_node->child_index];
        if(original_child != current) {            
            new_focus->children[path_node->child_index] = current;
        }
 
        zipper->focus = new_focus;
        
        free(path_node);
    }
}

static void make_mutable(Zipper *zipper, bool clone_focus) {
    if(clone_focus) {
        zipper->focus = node_clone(zipper->focus);
    }
    Node *child = zipper->focus;
    for(ZipperNode *path_node = zipper->path;
        path_node;
        path_node = path_node->prev) 
    {
        Node *new_parent = node_clone(path_node->parent);
        new_parent->children[path_node->child_index] = child;
        path_node->parent = new_parent;
        child = new_parent;
    }
}

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) {
    if(!zipper->mutable) {
        make_mutable(zipper, true);
        zipper->mutable = true;
    }
    
    Node *node = zipper->focus;
    Node *original_child = node->children[child_index];
    node_release(original_child);
    
    node->children[child_index] = child;
}

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

static void replace(Zipper *zipper, Node *new_node) {
    if(!zipper->mutable) {
        make_mutable(zipper, false);
        zipper->mutable = true;
    }
    
    Node *old_node = zipper->focus;
    node_release(old_node);
    zipper->focus = new_node;
}

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

int main() {
    Node *root = 0;
    bst_insert(&root, 50, 'A');
    bst_insert(&root, 25, 'B');
    bst_insert(&root, 75, 'C');
    bst_insert(&root, 12, 'D');
    bst_insert(&root, 83, 'E');
    bst_insert(&root, 30, 'F');
    
    //                         50
    //                 /                  \
    //                25                  75
    //           /         \                   \
    //          12         30                  83
    
    Zipper zipper = init_zipper(root);
    
    update(&zipper, 'W');
    update(&zipper, 'X');
    
    go_down(&zipper, 0);
    
    go_down(&zipper, 0);
    
    update(&zipper, 'Y');
    
    go_up(&zipper);
    
    go_down(&zipper, 1);
    
    update(&zipper, 'Z');
    
    Node *new_root = commit(&zipper);
    
    go_down(&zipper, 0);
    
    go_down(&zipper, 0);
    
    
    //                         2222  
    //                 /                  \
    //                25                  75
    //           /         \                   \
    //         666        777                  83


    return 0;
}