universe@192: /*
universe@192:  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
universe@192:  *
universe@192:  * Copyright 2015 Olaf Wintermann. All rights reserved.
universe@192:  *
universe@192:  * Redistribution and use in source and binary forms, with or without
universe@192:  * modification, are permitted provided that the following conditions are met:
universe@192:  *
universe@192:  *   1. Redistributions of source code must retain the above copyright
universe@192:  *      notice, this list of conditions and the following disclaimer.
universe@192:  *
universe@192:  *   2. Redistributions in binary form must reproduce the above copyright
universe@192:  *      notice, this list of conditions and the following disclaimer in the
universe@192:  *      documentation and/or other materials provided with the distribution.
universe@192:  *
universe@192:  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
universe@192:  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
universe@192:  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
universe@192:  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
universe@192:  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
universe@192:  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
universe@192:  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
universe@192:  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
universe@192:  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
universe@192:  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
universe@192:  * POSSIBILITY OF SUCH DAMAGE.
universe@192:  */
universe@192: 
universe@192: #include "avl.h"
universe@192: 
universe@195: #define ptrcast(ptr) ((void*)(ptr))
universe@216: #define alloc_tree(al) (UcxAVLTree*) almalloc((al), sizeof(UcxAVLTree))
universe@216: #define alloc_node(al) (UcxAVLNode*) almalloc((al), sizeof(UcxAVLNode))
universe@195: 
universe@195: static void ucx_avl_connect(UcxAVLTree *tree,
universe@195:         UcxAVLNode *node, UcxAVLNode *child, intptr_t nullkey) {
universe@195:     if (child) {
universe@195:         child->parent = node;
universe@195:     }
universe@195:     // if child is NULL, nullkey decides if left or right pointer is cleared
universe@195:     if (tree->cmpfunc(
universe@195:         ptrcast(child ? child->key : nullkey),
universe@195:         ptrcast(node->key), tree->userdata) > 0) {
universe@195:       node->right = child;
universe@195:     } else {
universe@195:       node->left = child;
universe@195:     }
universe@195:     size_t lh = node->left ? node->left->height : 0;
universe@195:     size_t rh = node->right ? node->right->height : 0;
universe@195:     node->height = 1 + (lh > rh ? lh : rh);
universe@195: }
universe@195: 
universe@195: #define avlheight(node) ((node) ? (node)->height : 0)
universe@195: 
universe@195: static UcxAVLNode* avl_rotright(UcxAVLTree *tree, UcxAVLNode *l0) {
universe@195:     UcxAVLNode *p = l0->parent;
universe@195:     UcxAVLNode *l1 = l0->left;
universe@195:     if (p) {
universe@195:         ucx_avl_connect(tree, p, l1, 0);
universe@195:     } else {
universe@195:         l1->parent = NULL;
universe@195:     }
universe@195:     ucx_avl_connect(tree, l0, l1->right, l1->key);
universe@195:     ucx_avl_connect(tree, l1, l0, 0);
universe@195:     return l1;
universe@195: }
universe@195: 
universe@195: static UcxAVLNode* avl_rotleft(UcxAVLTree *tree, UcxAVLNode *l0) {
universe@195:     UcxAVLNode *p = l0->parent;
universe@195:     UcxAVLNode *l1 = l0->right;
universe@195:     if (p) {
universe@195:         ucx_avl_connect(tree, p, l1, 0);
universe@195:     } else {
universe@195:         l1->parent = NULL;
universe@195:     }
universe@195:     ucx_avl_connect(tree, l0, l1->left, l1->key);
universe@195:     ucx_avl_connect(tree, l1, l0, 0);
universe@195:     return l1;
universe@195: }
universe@195: 
universe@195: static void ucx_avl_balance(UcxAVLTree *tree, UcxAVLNode *n) {
universe@195:     int lh = avlheight(n->left);
universe@195:     int rh = avlheight(n->right);
universe@195:     n->height = 1 + (lh > rh ? lh : rh);
universe@195:     
universe@195:     if (lh - rh == 2) {
universe@195:       UcxAVLNode *c = n->left;
universe@195:       if (avlheight(c->right) - avlheight(c->left) == 1) {
universe@195:         avl_rotleft(tree, c);
universe@195:       }
universe@195:       n = avl_rotright(tree, n);
universe@195:     } else if (rh - lh == 2) {  
universe@195:       UcxAVLNode *c = n->right;
universe@195:       if (avlheight(c->left) - avlheight(c->right) == 1) {
universe@195:         avl_rotright(tree, c);
universe@195:       }
universe@195:       n = avl_rotleft(tree, n);
universe@195:     }
universe@195: 
universe@195:     if (n->parent) {
universe@195:       ucx_avl_balance(tree, n->parent);
universe@195:     } else {
universe@195:       tree->root = n;
universe@195:     }
universe@195: }
universe@195: 
universe@194: UcxAVLTree *ucx_avl_new(cmp_func cmpfunc) {
universe@194:     return ucx_avl_new_a(cmpfunc, ucx_default_allocator());
universe@194: }
universe@194: 
universe@194: UcxAVLTree *ucx_avl_new_a(cmp_func cmpfunc, UcxAllocator *allocator) {
universe@216:     UcxAVLTree* tree = alloc_tree(allocator);
universe@194:     if (tree) {
universe@194:         tree->allocator = allocator;
universe@194:         tree->cmpfunc = cmpfunc;
universe@194:         tree->root = NULL;
universe@194:         tree->userdata = NULL;
universe@194:     }
universe@194:     
universe@194:     return tree;
universe@194: }
universe@194: 
universe@196: static void ucx_avl_free_node(UcxAllocator *al, UcxAVLNode *node) {
universe@196:     if (node) {
universe@196:         ucx_avl_free_node(al, node->left);
universe@196:         ucx_avl_free_node(al, node->right);
universe@196:         alfree(al, node);
universe@196:     }
universe@196: }
universe@196: 
universe@196: void ucx_avl_free(UcxAVLTree *tree) {
universe@196:     UcxAllocator *al = tree->allocator;
universe@196:     ucx_avl_free_node(al, tree->root);
universe@196:     alfree(al, tree);
universe@196: }
universe@196: 
universe@205: UcxAVLNode *ucx_avl_get_node(UcxAVLTree *tree, intptr_t key) {
universe@195:     UcxAVLNode *n = tree->root;
universe@195:     int cmpresult;
universe@195:     while (n && (cmpresult = tree->cmpfunc(
universe@195:             ptrcast(key), ptrcast(n->key), tree->userdata))) {
universe@195:         n = cmpresult > 0 ? n->right : n->left;
universe@195:     }
universe@204:     return n;
universe@204: }
universe@204: 
universe@204: void *ucx_avl_get(UcxAVLTree *tree, intptr_t key) {
universe@205:     UcxAVLNode *n = ucx_avl_get_node(tree, key);
universe@195:     return n ? n->value : NULL;
universe@194: }
universe@194: 
universe@204: int ucx_avl_put(UcxAVLTree *tree, intptr_t key, void *value) {
universe@204:     return ucx_avl_put_s(tree, key, value, NULL);
universe@204: }
universe@204: 
universe@204: int ucx_avl_put_s(UcxAVLTree *tree, intptr_t key, void *value,
universe@204:         void **oldvalue) {
universe@195:     if (tree->root) {
universe@195:         UcxAVLNode *n = tree->root;
universe@195:         int cmpresult;
universe@197:         while ((cmpresult = tree->cmpfunc(
universe@197:                 ptrcast(key), ptrcast(n->key), tree->userdata))) {
universe@195:             UcxAVLNode *m = cmpresult > 0 ? n->right : n->left;
universe@195:             if (m) {
universe@195:                 n = m;
universe@195:             } else {
universe@195:                 break;
universe@195:             }
universe@195:         }
universe@195: 
universe@195:         if (cmpresult) {
universe@216:             UcxAVLNode* e = alloc_node(tree->allocator);
universe@204:             if (e) {
universe@204:                 e->key = key; e->value = value; e->height = 1;
universe@204:                 e->parent = e->left = e->right = NULL;
universe@204:                 ucx_avl_connect(tree, n, e, 0);
universe@204:                 ucx_avl_balance(tree, n);
universe@204:                 return 0;
universe@204:             } else {
universe@204:                 return 1;
universe@204:             }
universe@195:         } else {
universe@204:             if (oldvalue) {
universe@204:                 *oldvalue = n->value;
universe@204:             }
universe@195:             n->value = value;
universe@204:             return 0;
universe@195:         }
universe@195:     } else {
universe@216:         tree->root = alloc_node(tree->allocator);
universe@204:         if (tree->root) {
universe@204:             tree->root->key = key; tree->root->value = value;
universe@204:             tree->root->height = 1;
universe@204:             tree->root->parent = tree->root->left = tree->root->right = NULL;
universe@204:             
universe@204:             if (oldvalue) {
universe@204:                 *oldvalue = NULL;
universe@204:             }
universe@204:             
universe@204:             return 0;
universe@204:         } else {
universe@204:             return 1;
universe@204:         }
universe@195:     }
universe@194: }
universe@194: 
universe@204: int ucx_avl_remove(UcxAVLTree *tree, intptr_t key) {
universe@204:     return ucx_avl_remove_s(tree, key, NULL, NULL);
universe@204: }
universe@204:     
universe@205: int ucx_avl_remove_node(UcxAVLTree *tree, UcxAVLNode *node) {
universe@204:     return ucx_avl_remove_s(tree, node->key, NULL, NULL);
universe@204: }
universe@204: 
universe@204: int ucx_avl_remove_s(UcxAVLTree *tree, intptr_t key,
universe@204:         intptr_t *oldkey, void **oldvalue) {
universe@204:     
universe@195:     UcxAVLNode *n = tree->root;
universe@195:     int cmpresult;
universe@195:     while (n && (cmpresult = tree->cmpfunc(
universe@195:             ptrcast(key), ptrcast(n->key), tree->userdata))) {
universe@195:         n = cmpresult > 0 ? n->right : n->left;
universe@195:     }
universe@195:     if (n) {
universe@204:         if (oldkey) {
universe@204:             *oldkey = n->key;
universe@204:         }
universe@204:         if (oldvalue) {
universe@204:             *oldvalue = n->value;
universe@204:         }
universe@204:         
universe@195:         UcxAVLNode *p = n->parent;
universe@195:         if (n->left && n->right) {
universe@195:             UcxAVLNode *s = n->right;
universe@195:             while (s->left) {
universe@195:                 s = s->left;
universe@195:             }
universe@195:             ucx_avl_connect(tree, s->parent, s->right, s->key);
universe@195:             n->key = s->key; n->value = s->value;
universe@195:             p = s->parent;
universe@196:             alfree(tree->allocator, s);
universe@195:         } else {
universe@195:             if (p) {
universe@195:                 ucx_avl_connect(tree, p, n->right ? n->right:n->left, n->key);
universe@195:             } else {
universe@195:                 tree->root = n->right ? n->right : n->left;
olaf@203:                 if (tree->root) {
olaf@202:                     tree->root->parent = NULL;
olaf@202:                 }
universe@195:             }
universe@196:             alfree(tree->allocator, n);
universe@195:         }
universe@195: 
universe@195:         if (p) {
universe@195:             ucx_avl_balance(tree, p);
universe@195:         }
universe@204:         
universe@204:         return 0;
universe@195:     } else {
universe@204:         return 1;
universe@195:     }
universe@194: }
universe@194: 
universe@199: static size_t ucx_avl_countn(UcxAVLNode *node) {
universe@199:     if (node) {
universe@199:         return 1 + ucx_avl_countn(node->left) + ucx_avl_countn(node->right);
universe@199:     } else {
universe@199:         return 0;
universe@199:     }
universe@199: }
universe@199: 
universe@199: size_t ucx_avl_count(UcxAVLTree *tree) {
universe@199:     return ucx_avl_countn(tree->root);
universe@199: }