#ifndef _BTREE_H_
#define _BTREE_H_

typedef struct btree_node
{
    unsigned int level;
    unsigned int index;
    struct btree_node *link[2];
} btree_node;


btree_node *btree_create();

btree_node *btree_insert(btree_node *root, int index);
btree_node *btree_find(btree_node *root, int index);
void btree_remove(btree_node *root, int index);

#endif

#include "btree.h"

// Because Lite-C doesn't support static inline, we use macros instead
#define btree_skew(root) do{if(root->link[0]->level == root->level && root->level != 0){btree_node *save = root->link[0]; root->link[0] = save->link[1]; save->link[1] = root; root = save;}}while(0)
#define btree_split(root) do{if(root->link[1]->link[1]->level == root->level && root->level != 0){btree_node *save = root->link[1]; root->link[1] = save->link[0]; save->link[0] = root; save->level ++; root = save;}}while(0)

#if 0
// The non inlined version of the macros for better understanding
btree_node *btree_skew(btree_node *root)
{
    if(root->link[0]->level == root->level && root->level != 0) 
    {
        btree_node *save = root->link[0];
        
        root->link[0] = save->link[1];
        save->link[1] = root;
        
        root = save;
    }
    
    return root;
}

btree_node *btree_split(btree_node *root)
{
    if(root->link[1]->link[1]->level == root->level && root->level != 0)
    {
        btree_node *save = root->link[1];
        
        root->link[1] = save->link[0];
        save->link[0] = root;
        save->level ++;
        
        root = save;
    }
    
    return root;
}
#endif


btree_node *btree_create()
{
    btree_node *node = (btree_node *)malloc(sizeof(btree_node));
    if(node)
    {
        node->level = 0;
        node->index = 0;
        node->link[0] = node->link[1] = NULL;
    }
    
    return node;
}

btree_node *btree_insert(btree_node *root, int index)
{
    if(!root)
        return NULL;
    
    btree_node *it = root;
    btree_node *up[32];
    int top = 0, direction = 0;
    
    while(1)
    {
        up[top] = it;
        direction = (it->index < index);
        
        top++;
        if(it->link[direction] == NULL)
            break;
        
        it = it->link[direction];
    }
    
    btree_node *insertion = btree_create();
    if(!insertion)
        return NULL;
    
    it->link[direction] = insertion;
    it->link[direction]->index = index;
    it->link[direction]->level = 1;
    
    while((top - 1) >= 0)
    {
        top --;
        if(top != 0)
            direction = (up[top - 1]->link[1] == up[top]);
        
        btree_node *node = up[top];
        btree_skew(node);
        btree_split(node);
        
        if(top != 0)
        {
            up[top - 1]->link[direction] = up[top];
        }
        else
        {
            root = up[top];
        }
    }
    
    return insertion;
}

btree_node *btree_find(btree_node *root, int index)
{
    if(!root)
        return NULL;
    
    btree_node *it = root;
    int direction = 0;
    
    while(1)
    {
        if(it->index == index)
            return it;
        
        direction = (it->index < index);
        
        if(!it->link[direction])
            break;
        
        it = it->link[direction];
    }
    
    return NULL;
}

// Remark, you still have to call free on the removed node!
void btree_remove(btree_node *root, int index)
{
    if(!root)
        return;
    
    btree_node *it = root;
    btree_node *up[32];
    int top = 0, direction = 0;
    
    while(1)
    {
        up[top] = it;
        top ++;
        
        if(!it)
        {
            return;
        }
        else if(it->index == index)
        {
            break;
        }
        
        direction = (it->index < index);
        it = it->link[direction];
    }
    
    if(it->link[0] == NULL || it->link[1] == NULL) 
    {
        int direction2 = (it->link[0] == NULL);
        
        if(top - 1 != 0)
        {
            up[top - 1]->link[direction] = it->link[direction2];
        }
        else
        {
            root = it->link[1];
        }
    }
    else 
    {
        btree_node *heir = it->link[1];
        btree_node *prev = it;
        
        while(heir->link[0] != NULL)
        {
            up[top] = prev = heir;
            heir = heir->link[0];
            
            top ++;
        }
        
        it->index = heir->index;
        prev->link[(prev == it)] = heir->link[1];
    }
    
    while((top - 1) >= 0 ) 
    {
        top --;
        if(top != 0)
            direction = (up[top - 1]->link[1] == up[top]);
        
        if(up[top]->link[0]->level < up[top]->level - 1 || up[top]->link[1]->level < up[top]->level - 1)
        {
            if(up[top]->link[1]->level > up[top]->level - 1)
            {
                up[top]->level --;
                up[top]->link[1]->level = up[top]->level;
            }
            
            btree_node *node = up[top];
            btree_skew(node);
            btree_skew(node->link[1]);
            btree_skew(node->link[1]->link[1]);
            
            btree_split(node);
            btree_split(node->link[1]);
        }
        
        if(top != 0)
        {
            up[top - 1]->link[direction] = up[top];
        }
        else
        {
            root = up[top];
        }
    }
}