ucx/avl.c

   1 /*
   2  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS HEADER.
   3  *
   4  * Copyright 2017 Mike Becker, Olaf Wintermann All rights reserved.
   5  *
   6  * Redistribution and use in source and binary forms, with or without
   7  * modification, are permitted provided that the following conditions are met:
   8  *
   9  *   1. Redistributions of source code must retain the above copyright
  10  *      notice, this list of conditions and the following disclaimer.
  11  *
  12  *   2. Redistributions in binary form must reproduce the above copyright
  13  *      notice, this list of conditions and the following disclaimer in the
  14  *      documentation and/or other materials provided with the distribution.
  15  *
  16  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  17  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  18  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  19  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
  20  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  21  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  22  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  23  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  24  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  25  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  26  * POSSIBILITY OF SUCH DAMAGE.
  27  */
  28
  29 #include "ucx/avl.h"
  30
  31 #include <limits.h>
  32
  33 #define ptrcast(ptr) ((void*)(ptr))
  34 #define alloc_tree(al) (UcxAVLTree*) almalloc((al), sizeof(UcxAVLTree))
  35 #define alloc_node(al) (UcxAVLNode*) almalloc((al), sizeof(UcxAVLNode))
  36
  37 static void ucx_avl_connect(UcxAVLTree *tree,
  38         UcxAVLNode *node, UcxAVLNode *child, intptr_t nullkey) {
  39     if (child) {
  40         child->parent = node;
  41     }
  42     // if child is NULL, nullkey decides if left or right pointer is cleared
  43     if (tree->cmpfunc(
  44         ptrcast(child ? child->key : nullkey),
  45         ptrcast(node->key), tree->userdata) > 0) {
  46       node->right = child;
  47     } else {
  48       node->left = child;
  49     }
  50     size_t lh = node->left ? node->left->height : 0;
  51     size_t rh = node->right ? node->right->height : 0;
  52     node->height = 1 + (lh > rh ? lh : rh);
  53 }
  54
  55 #define avlheight(node) ((node) ? (node)->height : 0)
  56
  57 static UcxAVLNode* avl_rotright(UcxAVLTree *tree, UcxAVLNode *l0) {
  58     UcxAVLNode *p = l0->parent;
  59     UcxAVLNode *l1 = l0->left;
  60     if (p) {
  61         ucx_avl_connect(tree, p, l1, 0);
  62     } else {
  63         l1->parent = NULL;
  64     }
  65     ucx_avl_connect(tree, l0, l1->right, l1->key);
  66     ucx_avl_connect(tree, l1, l0, 0);
  67     return l1;
  68 }
  69
  70 static UcxAVLNode* avl_rotleft(UcxAVLTree *tree, UcxAVLNode *l0) {
  71     UcxAVLNode *p = l0->parent;
  72     UcxAVLNode *l1 = l0->right;
  73     if (p) {
  74         ucx_avl_connect(tree, p, l1, 0);
  75     } else {
  76         l1->parent = NULL;
  77     }
  78     ucx_avl_connect(tree, l0, l1->left, l1->key);
  79     ucx_avl_connect(tree, l1, l0, 0);
  80     return l1;
  81 }
  82
  83 static void ucx_avl_balance(UcxAVLTree *tree, UcxAVLNode *n) {
  84     int lh = avlheight(n->left);
  85     int rh = avlheight(n->right);
  86     n->height = 1 + (lh > rh ? lh : rh);
  87
  88     if (lh - rh == 2) {
  89       UcxAVLNode *c = n->left;
  90       if (avlheight(c->right) - avlheight(c->left) == 1) {
  91         avl_rotleft(tree, c);
  92       }
  93       n = avl_rotright(tree, n);
  94     } else if (rh - lh == 2) {
  95       UcxAVLNode *c = n->right;
  96       if (avlheight(c->left) - avlheight(c->right) == 1) {
  97         avl_rotright(tree, c);
  98       }
  99       n = avl_rotleft(tree, n);
 100     }
 101
 102     if (n->parent) {
 103       ucx_avl_balance(tree, n->parent);
 104     } else {
 105       tree->root = n;
 106     }
 107 }
 108
 109 UcxAVLTree *ucx_avl_new(cmp_func cmpfunc) {
 110     return ucx_avl_new_a(cmpfunc, ucx_default_allocator());
 111 }
 112
 113 UcxAVLTree *ucx_avl_new_a(cmp_func cmpfunc, UcxAllocator *allocator) {
 114     UcxAVLTree* tree = alloc_tree(allocator);
 115     if (tree) {
 116         tree->allocator = allocator;
 117         tree->cmpfunc = cmpfunc;
 118         tree->root = NULL;
 119         tree->userdata = NULL;
 120     }
 121
 122     return tree;
 123 }
 124
 125 static void ucx_avl_free_node(UcxAllocator *al, UcxAVLNode *node) {
 126     if (node) {
 127         ucx_avl_free_node(al, node->left);
 128         ucx_avl_free_node(al, node->right);
 129         alfree(al, node);
 130     }
 131 }
 132
 133 void ucx_avl_free(UcxAVLTree *tree) {
 134     UcxAllocator *al = tree->allocator;
 135     ucx_avl_free_node(al, tree->root);
 136     alfree(al, tree);
 137 }
 138
 139 static void ucx_avl_free_content_node(UcxAllocator *al, UcxAVLNode *node,
 140         ucx_destructor destr) {
 141     if (node) {
 142         ucx_avl_free_content_node(al, node->left, destr);
 143         ucx_avl_free_content_node(al, node->right, destr);
 144         if (destr) {
 145             destr(node->value);
 146         } else {
 147             alfree(al, node->value);
 148         }
 149     }
 150 }
 151
 152 void ucx_avl_free_content(UcxAVLTree *tree, ucx_destructor destr) {
 153     ucx_avl_free_content_node(tree->allocator, tree->root, destr);
 154 }
 155
 156 UcxAVLNode *ucx_avl_get_node(UcxAVLTree *tree, intptr_t key) {
 157     UcxAVLNode *n = tree->root;
 158     int cmpresult;
 159     while (n && (cmpresult = tree->cmpfunc(
 160             ptrcast(key), ptrcast(n->key), tree->userdata))) {
 161         n = cmpresult > 0 ? n->right : n->left;
 162     }
 163     return n;
 164 }
 165
 166 void *ucx_avl_get(UcxAVLTree *tree, intptr_t key) {
 167     UcxAVLNode *n = ucx_avl_get_node(tree, key);
 168     return n ? n->value : NULL;
 169 }
 170
 171 UcxAVLNode *ucx_avl_find_node(UcxAVLTree *tree, intptr_t key,
 172         distance_func dfnc, int mode) {
 173     UcxAVLNode *n = tree->root;
 174     UcxAVLNode *closest = NULL;
 175
 176     intmax_t cmpresult;
 177     intmax_t closest_dist;
 178     closest_dist = mode == UCX_AVL_FIND_LOWER_BOUNDED ? INTMAX_MIN : INTMAX_MAX;
 179
 180     while (n && (cmpresult = dfnc(
 181             ptrcast(key), ptrcast(n->key), tree->userdata))) {
 182         if (mode == UCX_AVL_FIND_CLOSEST) {
 183             intmax_t dist = cmpresult;
 184             if (dist < 0) dist *= -1;
 185             if (dist < closest_dist) {
 186                 closest_dist = dist;
 187                 closest = n;
 188             }
 189         } else if (mode == UCX_AVL_FIND_LOWER_BOUNDED && cmpresult <= 0) {
 190             if (cmpresult > closest_dist) {
 191                 closest_dist = cmpresult;
 192                 closest = n;
 193             }
 194         } else if (mode == UCX_AVL_FIND_UPPER_BOUNDED && cmpresult >= 0) {
 195             if (cmpresult < closest_dist) {
 196                 closest_dist = cmpresult;
 197                 closest = n;
 198             }
 199         }
 200         n = cmpresult > 0 ? n->right : n->left;
 201     }
 202     return n ? n : closest;
 203 }
 204
 205 void *ucx_avl_find(UcxAVLTree *tree, intptr_t key,
 206         distance_func dfnc, int mode) {
 207     UcxAVLNode *n = ucx_avl_find_node(tree, key, dfnc, mode);
 208     return n ? n->value : NULL;
 209 }
 210
 211 int ucx_avl_put(UcxAVLTree *tree, intptr_t key, void *value) {
 212     return ucx_avl_put_s(tree, key, value, NULL);
 213 }
 214
 215 int ucx_avl_put_s(UcxAVLTree *tree, intptr_t key, void *value,
 216         void **oldvalue) {
 217     if (tree->root) {
 218         UcxAVLNode *n = tree->root;
 219         int cmpresult;
 220         while ((cmpresult = tree->cmpfunc(
 221                 ptrcast(key), ptrcast(n->key), tree->userdata))) {
 222             UcxAVLNode *m = cmpresult > 0 ? n->right : n->left;
 223             if (m) {
 224                 n = m;
 225             } else {
 226                 break;
 227             }
 228         }
 229
 230         if (cmpresult) {
 231             UcxAVLNode* e = alloc_node(tree->allocator);
 232             if (e) {
 233                 e->key = key; e->value = value; e->height = 1;
 234                 e->parent = e->left = e->right = NULL;
 235                 ucx_avl_connect(tree, n, e, 0);
 236                 ucx_avl_balance(tree, n);
 237                 return 0;
 238             } else {
 239                 return 1;
 240             }
 241         } else {
 242             if (oldvalue) {
 243                 *oldvalue = n->value;
 244             }
 245             n->value = value;
 246             return 0;
 247         }
 248     } else {
 249         tree->root = alloc_node(tree->allocator);
 250         if (tree->root) {
 251             tree->root->key = key; tree->root->value = value;
 252             tree->root->height = 1;
 253             tree->root->parent = tree->root->left = tree->root->right = NULL;
 254
 255             if (oldvalue) {
 256                 *oldvalue = NULL;
 257             }
 258
 259             return 0;
 260         } else {
 261             return 1;
 262         }
 263     }
 264 }
 265
 266 int ucx_avl_remove(UcxAVLTree *tree, intptr_t key) {
 267     return ucx_avl_remove_s(tree, key, NULL, NULL);
 268 }
 269
 270 int ucx_avl_remove_node(UcxAVLTree *tree, UcxAVLNode *node) {
 271     return ucx_avl_remove_s(tree, node->key, NULL, NULL);
 272 }
 273
 274 int ucx_avl_remove_s(UcxAVLTree *tree, intptr_t key,
 275         intptr_t *oldkey, void **oldvalue) {
 276
 277     UcxAVLNode *n = tree->root;
 278     int cmpresult;
 279     while (n && (cmpresult = tree->cmpfunc(
 280             ptrcast(key), ptrcast(n->key), tree->userdata))) {
 281         n = cmpresult > 0 ? n->right : n->left;
 282     }
 283     if (n) {
 284         if (oldkey) {
 285             *oldkey = n->key;
 286         }
 287         if (oldvalue) {
 288             *oldvalue = n->value;
 289         }
 290
 291         UcxAVLNode *p = n->parent;
 292         if (n->left && n->right) {
 293             UcxAVLNode *s = n->right;
 294             while (s->left) {
 295                 s = s->left;
 296             }
 297             ucx_avl_connect(tree, s->parent, s->right, s->key);
 298             n->key = s->key; n->value = s->value;
 299             p = s->parent;
 300             alfree(tree->allocator, s);
 301         } else {
 302             if (p) {
 303                 ucx_avl_connect(tree, p, n->right ? n->right:n->left, n->key);
 304             } else {
 305                 tree->root = n->right ? n->right : n->left;
 306                 if (tree->root) {
 307                     tree->root->parent = NULL;
 308                 }
 309             }
 310             alfree(tree->allocator, n);
 311         }
 312
 313         if (p) {
 314             ucx_avl_balance(tree, p);
 315         }
 316
 317         return 0;
 318     } else {
 319         return 1;
 320     }
 321 }
 322
 323 static size_t ucx_avl_countn(UcxAVLNode *node) {
 324     if (node) {
 325         return 1 + ucx_avl_countn(node->left) + ucx_avl_countn(node->right);
 326     } else {
 327         return 0;
 328     }
 329 }
 330
 331 size_t ucx_avl_count(UcxAVLTree *tree) {
 332     return ucx_avl_countn(tree->root);
 333 }
 334
 335 UcxAVLNode* ucx_avl_pred(UcxAVLNode* node) {
 336     if (node->left) {
 337         UcxAVLNode* n = node->left;
 338         while (n->right) {
 339             n = n->right;
 340         }
 341         return n;
 342     } else {
 343         UcxAVLNode* n = node;
 344         while (n->parent) {
 345             if (n->parent->right == n) {
 346                 return n->parent;
 347             } else {
 348                 n = n->parent;
 349             }
 350         }
 351         return NULL;
 352     }
 353 }
 354
 355 UcxAVLNode* ucx_avl_succ(UcxAVLNode* node) {
 356     if (node->right) {
 357         UcxAVLNode* n = node->right;
 358         while (n->left) {
 359             n = n->left;
 360         }
 361         return n;
 362     } else {
 363         UcxAVLNode* n = node;
 364         while (n->parent) {
 365             if (n->parent->left == n) {
 366                 return n->parent;
 367             } else {
 368                 n = n->parent;
 369             }
 370         }
 371         return NULL;
 372     }
 373 }