btree.c

Go to the documentation of this file.

/*
 * Copyright (C) 2006 Jakub Jermar
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * - Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 * - The name of the author may not be used to endorse or promote products
 *   derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <adt/btree.h>
#include <adt/list.h>
#include <mm/slab.h>
#include <debug.h>
#include <panic.h>
#include <typedefs.h>
#include <print.h>

static void btree_destroy_subtree(btree_node_t *root);
static void _btree_insert(btree_t *t, btree_key_t key, void *value, btree_node_t *rsubtree, btree_node_t *node);
static void _btree_remove(btree_t *t, btree_key_t key, btree_node_t *node);
static void node_initialize(btree_node_t *node);
static void node_insert_key_and_lsubtree(btree_node_t *node, btree_key_t key, void *value, btree_node_t *lsubtree);
static void node_insert_key_and_rsubtree(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree);
static void node_remove_key_and_lsubtree(btree_node_t *node, btree_key_t key);
static void node_remove_key_and_rsubtree(btree_node_t *node, btree_key_t key);
static btree_node_t *node_split(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree, btree_key_t *median);
static btree_node_t *node_combine(btree_node_t *node);
static index_t find_key_by_subtree(btree_node_t *node, btree_node_t *subtree, bool right);
static void rotate_from_right(btree_node_t *lnode, btree_node_t *rnode, index_t idx);
static void rotate_from_left(btree_node_t *lnode, btree_node_t *rnode, index_t idx);
static bool try_insert_by_rotation_to_left(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree);
static bool try_insert_by_rotation_to_right(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree);
static bool try_rotation_from_left(btree_node_t *rnode);
static bool try_rotation_from_right(btree_node_t *lnode);

#define ROOT_NODE(n)            (!(n)->parent)
#define INDEX_NODE(n)           ((n)->subtree[0] != NULL)
#define LEAF_NODE(n)            ((n)->subtree[0] == NULL)

#define FILL_FACTOR             ((BTREE_M-1)/2)

#define MEDIAN_LOW_INDEX(n)     (((n)->keys-1)/2)
#define MEDIAN_HIGH_INDEX(n)    ((n)->keys/2)
#define MEDIAN_LOW(n)           ((n)->key[MEDIAN_LOW_INDEX((n))]);
#define MEDIAN_HIGH(n)          ((n)->key[MEDIAN_HIGH_INDEX((n))]);

static slab_cache_t *btree_node_slab;

void btree_init(void)
{
        btree_node_slab = slab_cache_create("btree_node_slab", sizeof(btree_node_t), 0, NULL, NULL, SLAB_CACHE_MAGDEFERRED);
}

void btree_create(btree_t *t)
{
        list_initialize(&t->leaf_head);
        t->root = (btree_node_t *) slab_alloc(btree_node_slab, 0);
        node_initialize(t->root);
        list_append(&t->root->leaf_link, &t->leaf_head);
}

void btree_destroy(btree_t *t)
{
        btree_destroy_subtree(t->root);
}

void btree_insert(btree_t *t, btree_key_t key, void *value, btree_node_t *leaf_node)
{
        btree_node_t *lnode;
        
        ASSERT(value);
        
        lnode = leaf_node;
        if (!lnode) {
                if (btree_search(t, key, &lnode)) {
                        panic("B-tree %p already contains key %d\n", t, key);
                }
        }
        
        _btree_insert(t, key, value, NULL, lnode);
}

void btree_destroy_subtree(btree_node_t *root)
{
        int i;

        if (root->keys) {
                for (i = 0; i < root->keys + 1; i++) { 
                        if (root->subtree[i])
                                btree_destroy_subtree(root->subtree[i]);
                }
        }
        slab_free(btree_node_slab, root);
}

void _btree_insert(btree_t *t, btree_key_t key, void *value, btree_node_t *rsubtree, btree_node_t *node)
{
        if (node->keys < BTREE_MAX_KEYS) {
                /*
                 * Node conatins enough space, the key can be stored immediately.
                 */
                node_insert_key_and_rsubtree(node, key, value, rsubtree);
        } else if (try_insert_by_rotation_to_left(node, key, value, rsubtree)) {
                /*
                 * The key-value-rsubtree triplet has been inserted because
                 * some keys could have been moved to the left sibling.
                 */
        } else if (try_insert_by_rotation_to_right(node, key, value, rsubtree)) {
                /*
                 * The key-value-rsubtree triplet has been inserted because
                 * some keys could have been moved to the right sibling.
                 */
        } else {
                btree_node_t *rnode;
                btree_key_t median;
                
                /*
                 * Node is full and both siblings (if both exist) are full too.
                 * Split the node and insert the smallest key from the node containing
                 * bigger keys (i.e. the new node) into its parent.
                 */

                rnode = node_split(node, key, value, rsubtree, &median);

                if (LEAF_NODE(node)) {
                        list_prepend(&rnode->leaf_link, &node->leaf_link);
                }
                
                if (ROOT_NODE(node)) {
                        /*
                         * We split the root node. Create new root.
                         */
                        t->root = (btree_node_t *) slab_alloc(btree_node_slab, 0);
                        node->parent = t->root;
                        rnode->parent = t->root;
                        node_initialize(t->root);
                        
                        /*
                         * Left-hand side subtree will be the old root (i.e. node).
                         * Right-hand side subtree will be rnode.
                         */                     
                        t->root->subtree[0] = node;

                        t->root->depth = node->depth + 1;
                }
                _btree_insert(t, median, NULL, rnode, node->parent);
        }       
                
}

void btree_remove(btree_t *t, btree_key_t key, btree_node_t *leaf_node)
{
        btree_node_t *lnode;
        
        lnode = leaf_node;
        if (!lnode) {
                if (!btree_search(t, key, &lnode)) {
                        panic("B-tree %p does not contain key %d\n", t, key);
                }
        }
        
        _btree_remove(t, key, lnode);
}

void _btree_remove(btree_t *t, btree_key_t key, btree_node_t *node)
{
        if (ROOT_NODE(node)) {
                if (node->keys == 1 && node->subtree[0]) {
                        /*
                         * Free the current root and set new root.
                         */
                        t->root = node->subtree[0];
                        t->root->parent = NULL;
                        slab_free(btree_node_slab, node);
                } else {
                        /*
                         * Remove the key from the root node.
                         * Note that the right subtree is removed because when
                         * combining two nodes, the left-side sibling is preserved
                         * and the right-side sibling is freed.
                         */
                        node_remove_key_and_rsubtree(node, key);
                }
                return;
        }
        
        if (node->keys <= FILL_FACTOR) {
                /*
                 * If the node is below the fill factor,
                 * try to borrow keys from left or right sibling.
                 */
                if (!try_rotation_from_left(node))
                        try_rotation_from_right(node);
        }
        
        if (node->keys > FILL_FACTOR) {
                int i;

                /*
                 * The key can be immediatelly removed.
                 *
                 * Note that the right subtree is removed because when
                 * combining two nodes, the left-side sibling is preserved
                 * and the right-side sibling is freed.
                 */
                node_remove_key_and_rsubtree(node, key);
                for (i = 0; i < node->parent->keys; i++) {
                        if (node->parent->key[i] == key)
                                node->parent->key[i] = node->key[0];
                }
                
        } else {
                index_t idx;
                btree_node_t *rnode, *parent;

                /*
                 * The node is below the fill factor as well as its left and right sibling.
                 * Resort to combining the node with one of its siblings.
                 * The node which is on the left is preserved and the node on the right is
                 * freed.
                 */
                parent = node->parent;
                node_remove_key_and_rsubtree(node, key);
                rnode = node_combine(node);
                if (LEAF_NODE(rnode))
                        list_remove(&rnode->leaf_link);
                idx = find_key_by_subtree(parent, rnode, true);
                ASSERT((int) idx != -1);
                slab_free(btree_node_slab, rnode);
                _btree_remove(t, parent->key[idx], parent);
        }
}

void *btree_search(btree_t *t, btree_key_t key, btree_node_t **leaf_node)
{
        btree_node_t *cur, *next;
        
        /*
         * Iteratively descend to the leaf that can contain the searched key.
         */
        for (cur = t->root; cur; cur = next) {

                /* Last iteration will set this with proper leaf node address. */
                *leaf_node = cur;
                
                /*
                 * The key can be in the leftmost subtree.
                 * Test it separately.
                 */
                if (key < cur->key[0]) {
                        next = cur->subtree[0];
                        continue;
                } else {
                        void *val;
                        int i;
                
                        /*
                         * Now if the key is smaller than cur->key[i]
                         * it can only mean that the value is in cur->subtree[i]
                         * or it is not in the tree at all.
                         */
                        for (i = 1; i < cur->keys; i++) {
                                if (key < cur->key[i]) {
                                        next = cur->subtree[i];
                                        val = cur->value[i - 1];

                                        if (LEAF_NODE(cur))
                                                return key == cur->key[i - 1] ? val : NULL;

                                        goto descend;
                                } 
                        }
                        
                        /*
                         * Last possibility is that the key is in the rightmost subtree.
                         */
                        next = cur->subtree[i];
                        val = cur->value[i - 1];
                        if (LEAF_NODE(cur))
                                return key == cur->key[i - 1] ? val : NULL;
                }
                descend:
                        ;
        }

        /*
         * The key was not found in the *leaf_node and is smaller than any of its keys.
         */
        return NULL;
}

btree_node_t *btree_leaf_node_left_neighbour(btree_t *t, btree_node_t *node)
{
        ASSERT(LEAF_NODE(node));
        if (node->leaf_link.prev != &t->leaf_head)
                return list_get_instance(node->leaf_link.prev, btree_node_t, leaf_link);
        else
                return NULL;
}

btree_node_t *btree_leaf_node_right_neighbour(btree_t *t, btree_node_t *node)
{
        ASSERT(LEAF_NODE(node));
        if (node->leaf_link.next != &t->leaf_head)
                return list_get_instance(node->leaf_link.next, btree_node_t, leaf_link);
        else
                return NULL;
}

void node_initialize(btree_node_t *node)
{
        int i;

        node->keys = 0;
        
        /* Clean also space for the extra key. */
        for (i = 0; i < BTREE_MAX_KEYS + 1; i++) {
                node->key[i] = 0;
                node->value[i] = NULL;
                node->subtree[i] = NULL;
        }
        node->subtree[i] = NULL;
        
        node->parent = NULL;
        
        link_initialize(&node->leaf_link);

        link_initialize(&node->bfs_link);
        node->depth = 0;
}

void node_insert_key_and_lsubtree(btree_node_t *node, btree_key_t key, void *value, btree_node_t *lsubtree)
{
        int i;

        for (i = 0; i < node->keys; i++) {
                if (key < node->key[i]) {
                        int j;
                
                        for (j = node->keys; j > i; j--) {
                                node->key[j] = node->key[j - 1];
                                node->value[j] = node->value[j - 1];
                                node->subtree[j + 1] = node->subtree[j];
                        }
                        node->subtree[j + 1] = node->subtree[j];
                        break;  
                }
        }
        node->key[i] = key;
        node->value[i] = value;
        node->subtree[i] = lsubtree;
                        
        node->keys++;
}

void node_insert_key_and_rsubtree(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree)
{
        int i;

        for (i = 0; i < node->keys; i++) {
                if (key < node->key[i]) {
                        int j;
                
                        for (j = node->keys; j > i; j--) {
                                node->key[j] = node->key[j - 1];
                                node->value[j] = node->value[j - 1];
                                node->subtree[j + 1] = node->subtree[j];
                        }
                        break;  
                }
        }
        node->key[i] = key;
        node->value[i] = value;
        node->subtree[i + 1] = rsubtree;
                        
        node->keys++;
}

void node_remove_key_and_lsubtree(btree_node_t *node, btree_key_t key)
{
        int i, j;
        
        for (i = 0; i < node->keys; i++) {
                if (key == node->key[i]) {
                        for (j = i + 1; j < node->keys; j++) {
                                node->key[j - 1] = node->key[j];
                                node->value[j - 1] = node->value[j];
                                node->subtree[j - 1] = node->subtree[j];
                        }
                        node->subtree[j - 1] = node->subtree[j];
                        node->keys--;
                        return;
                }
        }
        panic("node %p does not contain key %d\n", node, key);
}

void node_remove_key_and_rsubtree(btree_node_t *node, btree_key_t key)
{
        int i, j;
        
        for (i = 0; i < node->keys; i++) {
                if (key == node->key[i]) {
                        for (j = i + 1; j < node->keys; j++) {
                                node->key[j - 1] = node->key[j];
                                node->value[j - 1] = node->value[j];
                                node->subtree[j] = node->subtree[j + 1];
                        }
                        node->keys--;
                        return;
                }
        }
        panic("node %p does not contain key %d\n", node, key);
}

btree_node_t *node_split(btree_node_t *node, btree_key_t key, void *value, btree_node_t *rsubtree, btree_key_t *median)
{
        btree_node_t *rnode;
        int i, j;

        ASSERT(median);
        ASSERT(node->keys == BTREE_MAX_KEYS);

        /*
         * Use the extra space to store the extra node.
         */
        node_insert_key_and_rsubtree(node, key, value, rsubtree);

        /*
         * Compute median of keys.
         */
        *median = MEDIAN_HIGH(node);
                
        /*
         * Allocate and initialize new right sibling.
         */
        rnode = (btree_node_t *) slab_alloc(btree_node_slab, 0);
        node_initialize(rnode);
        rnode->parent = node->parent;
        rnode->depth = node->depth;
        
        /*
         * Copy big keys, values and subtree pointers to the new right sibling.
         * If this is an index node, do not copy the median.
         */
        i = (int) INDEX_NODE(node);
        for (i += MEDIAN_HIGH_INDEX(node), j = 0; i < node->keys; i++, j++) {
                rnode->key[j] = node->key[i];
                rnode->value[j] = node->value[i];
                rnode->subtree[j] = node->subtree[i];
                
                /*
                 * Fix parent links in subtrees.
                 */
                if (rnode->subtree[j])
                        rnode->subtree[j]->parent = rnode;
                        
        }
        rnode->subtree[j] = node->subtree[i];
        if (rnode->subtree[j])
                rnode->subtree[j]->parent = rnode;

        rnode->keys = j;        /* Set number of keys of the new node. */
        node->keys /= 2;        /* Shrink the old node. */
                
        return rnode;
}

btree_node_t *node_combine(btree_node_t *node)
{
        index_t idx;
        btree_node_t *rnode;
        int i;

        ASSERT(!ROOT_NODE(node));
        
        idx = find_key_by_subtree(node->parent, node, false);
        if (idx == node->parent->keys) {
                /*
                 * Rightmost subtree of its parent, combine with the left sibling.
                 */
                idx--;
                rnode = node;
                node = node->parent->subtree[idx];
        } else {
                rnode = node->parent->subtree[idx + 1];
        }

        /* Index nodes need to insert parent node key in between left and right node. */
        if (INDEX_NODE(node))
                node->key[node->keys++] = node->parent->key[idx];
        
        /* Copy the key-value-subtree triplets from the right node. */
        for (i = 0; i < rnode->keys; i++) {
                node->key[node->keys + i] = rnode->key[i];
                node->value[node->keys + i] = rnode->value[i];
                if (INDEX_NODE(node)) {
                        node->subtree[node->keys + i] = rnode->subtree[i];
                        rnode->subtree[i]->parent = node;
                }
        }
        if (INDEX_NODE(node)) {
                node->subtree[node->keys + i] = rnode->subtree[i];
                rnode->subtree[i]->parent = node;
        }

        node->keys += rnode->keys;

        return rnode;
}

index_t find_key_by_subtree(btree_node_t *node, btree_node_t *subtree, bool right)
{
        int i;
        
        for (i = 0; i < node->keys + 1; i++) {
                if (subtree == node->subtree[i])
                        return i - (int) (right != false);
        }
        panic("node %p does not contain subtree %p\n", node, subtree);
}

void rotate_from_left(btree_node_t *lnode, btree_node_t *rnode, index_t idx)
{
        btree_key_t key;

        key = lnode->key[lnode->keys - 1];
                
        if (LEAF_NODE(lnode)) {
                void *value;

                value = lnode->value[lnode->keys - 1];
                node_remove_key_and_rsubtree(lnode, key);
                node_insert_key_and_lsubtree(rnode, key, value, NULL);
                lnode->parent->key[idx] = key;
        } else {
                btree_node_t *rsubtree;

                rsubtree = lnode->subtree[lnode->keys];
                node_remove_key_and_rsubtree(lnode, key);
                node_insert_key_and_lsubtree(rnode, lnode->parent->key[idx], NULL, rsubtree);
                lnode->parent->key[idx] = key;

                /* Fix parent link of the reconnected right subtree. */
                rsubtree->parent = rnode;
        }

}

void rotate_from_right(btree_node_t *lnode, btree_node_t *rnode, index_t idx)
{
        btree_key_t key;

        key = rnode->key[0];
                
        if (LEAF_NODE(rnode)) {
                void *value;

                value = rnode->value[0];
                node_remove_key_and_lsubtree(rnode, key);
                node_insert_key_and_rsubtree(lnode, key, value, NULL);
                rnode->parent->key[idx] = rnode->key[0];
        } else {
                btree_node_t *lsubtree;

                lsubtree = rnode->subtree[0];
                node_remove_key_and_lsubtree(rnode, key);
                node_insert_key_and_rsubtree(lnode, rnode->parent->key[idx], NULL, lsubtree);
                rnode->parent->key[idx] = key;

                /* Fix parent link of the reconnected left subtree. */
                lsubtree->parent = lnode;
        }

}

bool try_insert_by_rotation_to_left(btree_node_t *node, btree_key_t inskey, void *insvalue, btree_node_t *rsubtree)
{
        index_t idx;
        btree_node_t *lnode;

        /*
         * If this is root node, the rotation can not be done.
         */
        if (ROOT_NODE(node))
                return false;
        
        idx = find_key_by_subtree(node->parent, node, true);
        if ((int) idx == -1) {
                /*
                 * If this node is the leftmost subtree of its parent,
                 * the rotation can not be done.
                 */
                return false;
        }
                
        lnode = node->parent->subtree[idx];
        if (lnode->keys < BTREE_MAX_KEYS) {
                /*
                 * The rotaion can be done. The left sibling has free space.
                 */
                node_insert_key_and_rsubtree(node, inskey, insvalue, rsubtree);
                rotate_from_right(lnode, node, idx);
                return true;
        }

        return false;
}

bool try_insert_by_rotation_to_right(btree_node_t *node, btree_key_t inskey, void *insvalue, btree_node_t *rsubtree)
{
        index_t idx;
        btree_node_t *rnode;

        /*
         * If this is root node, the rotation can not be done.
         */
        if (ROOT_NODE(node))
                return false;
        
        idx = find_key_by_subtree(node->parent, node, false);
        if (idx == node->parent->keys) {
                /*
                 * If this node is the rightmost subtree of its parent,
                 * the rotation can not be done.
                 */
                return false;
        }
                
        rnode = node->parent->subtree[idx + 1];
        if (rnode->keys < BTREE_MAX_KEYS) {
                /*
                 * The rotaion can be done. The right sibling has free space.
                 */
                node_insert_key_and_rsubtree(node, inskey, insvalue, rsubtree);
                rotate_from_left(node, rnode, idx);
                return true;
        }

        return false;
}

bool try_rotation_from_left(btree_node_t *rnode)
{
        index_t idx;
        btree_node_t *lnode;

        /*
         * If this is root node, the rotation can not be done.
         */
        if (ROOT_NODE(rnode))
                return false;
        
        idx = find_key_by_subtree(rnode->parent, rnode, true);
        if ((int) idx == -1) {
                /*
                 * If this node is the leftmost subtree of its parent,
                 * the rotation can not be done.
                 */
                return false;
        }
                
        lnode = rnode->parent->subtree[idx];
        if (lnode->keys > FILL_FACTOR) {
                rotate_from_left(lnode, rnode, idx);
                return true;
        }
        
        return false;
}

bool try_rotation_from_right(btree_node_t *lnode)
{
        index_t idx;
        btree_node_t *rnode;

        /*
         * If this is root node, the rotation can not be done.
         */
        if (ROOT_NODE(lnode))
                return false;
        
        idx = find_key_by_subtree(lnode->parent, lnode, false);
        if (idx == lnode->parent->keys) {
                /*
                 * If this node is the rightmost subtree of its parent,
                 * the rotation can not be done.
                 */
                return false;
        }
                
        rnode = lnode->parent->subtree[idx + 1];
        if (rnode->keys > FILL_FACTOR) {
                rotate_from_right(lnode, rnode, idx);
                return true;
        }       

        return false;
}

void btree_print(btree_t *t)
{
        int i, depth = t->root->depth;
        link_t head, *cur;

        printf("Printing B-tree:\n");
        list_initialize(&head);
        list_append(&t->root->bfs_link, &head);

        /*
         * Use BFS search to print out the tree.
         * Levels are distinguished from one another by node->depth.
         */     
        while (!list_empty(&head)) {
                link_t *hlp;
                btree_node_t *node;
                
                hlp = head.next;
                ASSERT(hlp != &head);
                node = list_get_instance(hlp, btree_node_t, bfs_link);
                list_remove(hlp);
                
                ASSERT(node);
                
                if (node->depth != depth) {
                        printf("\n");
                        depth = node->depth;
                }

                printf("(");
                for (i = 0; i < node->keys; i++) {
                        printf("%lld%s", node->key[i], i < node->keys - 1 ? "," : "");
                        if (node->depth && node->subtree[i]) {
                                list_append(&node->subtree[i]->bfs_link, &head);
                        }
                }
                if (node->depth && node->subtree[i]) {
                        list_append(&node->subtree[i]->bfs_link, &head);
                }
                printf(")");
        }
        printf("\n");
        
        printf("Printing list of leaves:\n");
        for (cur = t->leaf_head.next; cur != &t->leaf_head; cur = cur->next) {
                btree_node_t *node;
                
                node = list_get_instance(cur, btree_node_t, leaf_link);
                
                ASSERT(node);

                printf("(");
                for (i = 0; i < node->keys; i++)
                        printf("%lld%s", node->key[i], i < node->keys - 1 ? "," : "");
                printf(")");
        }
        printf("\n");
}

---