sys/src/libavl/avl.c

   1 #include <u.h>
   2 #include <libc.h>
   3 #include <bio.h>
   4 #include <avl.h>
   5
   6 /*
   7  * In-memory database stored as self-balancing AVL tree.
   8  * See Lewis & Denenberg, Data Structures and Their Algorithms.
   9  */
  10
  11 static void
  12 singleleft(Avl **tp, Avl *p)
  13 {
  14         int l, r2;
  15         Avl *a, *c;
  16
  17         a = *tp;
  18         c = a->n[1];
  19
  20         r2 = c->bal;
  21         l = (r2 > 0? r2: 0)+1 - a->bal;
  22
  23         if((a->n[1] = c->n[0]) != nil)
  24                 a->n[1]->p = a;
  25
  26         if((c->n[0] = a) != nil)
  27                 c->n[0]->p = c;
  28
  29         if((*tp = c) != nil)
  30                 (*tp)->p = p;
  31
  32         a->bal = -l;
  33         c->bal = r2 - ((l > 0? l: 0)+1);
  34
  35 }
  36
  37 static void
  38 singleright(Avl **tp, Avl *p)
  39 {
  40         int l2, r;
  41         Avl *a, *c;
  42
  43         a = *tp;
  44         c = a->n[0];
  45         l2 = - c->bal;
  46         r = a->bal + ((l2 > 0? l2: 0)+1);
  47
  48         if((a->n[0] = c->n[1]) != nil)
  49                 a->n[0]->p = a;
  50
  51         if((c->n[1] = a) != nil)
  52                 c->n[1]->p = c;
  53
  54         if((*tp = c) != nil)
  55                 (*tp)->p = p;
  56
  57         a->bal = r;
  58         c->bal = ((r > 0? r: 0)+1) - l2;
  59 }
  60
  61 static void
  62 doublerightleft(Avl **tp, Avl *p)
  63 {
  64         singleright(&(*tp)->n[1], *tp);
  65         singleleft(tp, p);
  66 }
  67
  68 static void
  69 doubleleftright(Avl **tp, Avl *p)
  70 {
  71         singleleft(&(*tp)->n[0], *tp);
  72         singleright(tp, p);
  73 }
  74
  75 static void
  76 balance(Avl **tp, Avl *p)
  77 {
  78         switch((*tp)->bal){
  79         case -2:
  80                 if((*tp)->n[0]->bal <= 0)
  81                         singleright(tp, p);
  82                 else if((*tp)->n[0]->bal == 1)
  83                         doubleleftright(tp, p);
  84                 else
  85                         assert(0);
  86                 break;
  87
  88         case 2:
  89                 if((*tp)->n[1]->bal >= 0)
  90                         singleleft(tp, p);
  91                 else if((*tp)->n[1]->bal == -1)
  92                         doublerightleft(tp, p);
  93                 else
  94                         assert(0);
  95                 break;
  96         }
  97 }
  98
  99 static int
 100 canoncmp(int cmp)
 101 {
 102         if(cmp < 0)
 103                 return -1;
 104         else if(cmp > 0)
 105                 return 1;
 106         return 0;
 107 }
 108
 109 static int
 110 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree)
 111 {
 112         int i, ob;
 113
 114         if(*tp == nil){
 115                 r->bal = 0;
 116                 r->n[0] = nil;
 117                 r->n[1] = nil;
 118                 r->p = p;
 119                 *tp = r;
 120                 return 1;
 121         }
 122         ob = (*tp)->bal;
 123         if((i = canoncmp(cmp(r, *tp))) != 0){
 124                 (*tp)->bal += i * _insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp,
 125                         rfree);
 126                 balance(tp, p);
 127                 return ob == 0 && (*tp)->bal != 0;
 128         }
 129
 130         /* install new entry */
 131         *rfree = *tp;           /* save old node for freeing */
 132         *tp = r;                /* insert new node */
 133         **tp = **rfree;         /* copy old node's Avl contents */
 134         if(r->n[0])             /* fix node's children's parent pointers */
 135                 r->n[0]->p = r;
 136         if(r->n[1])
 137                 r->n[1]->p = r;
 138
 139         return 0;
 140 }
 141
 142 static int
 143 successor(Avl **tp, Avl *p, Avl **r)
 144 {
 145         int ob;
 146
 147         if((*tp)->n[0] == nil){
 148                 *r = *tp;
 149                 *tp = (*r)->n[1];
 150                 if(*tp)
 151                         (*tp)->p = p;
 152                 return -1;
 153         }
 154         ob = (*tp)->bal;
 155         (*tp)->bal -= successor(&(*tp)->n[0], *tp, r);
 156         balance(tp, p);
 157         return -(ob != 0 && (*tp)->bal == 0);
 158 }
 159
 160 static int
 161 _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del,
 162         void (*predel)(Avl*, void*), void *arg)
 163 {
 164         int i, ob;
 165         Avl *r, *or;
 166
 167         if(*tp == nil)
 168                 return 0;
 169
 170         ob = (*tp)->bal;
 171         if((i=canoncmp(cmp(rx, *tp))) != 0){
 172                 (*tp)->bal += i * _deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp,
 173                         del, predel, arg);
 174                 balance(tp, p);
 175                 return -(ob != 0 && (*tp)->bal == 0);
 176         }
 177
 178         if(predel)
 179                 (*predel)(*tp, arg);
 180
 181         or = *tp;
 182         if(or->n[i=0] == nil || or->n[i=1] == nil){
 183                 *tp = or->n[1-i];
 184                 if(*tp)
 185                         (*tp)->p = p;
 186                 *del = or;
 187                 return -1;
 188         }
 189
 190         /* deleting node with two kids, find successor */
 191         or->bal += successor(&or->n[1], or, &r);
 192         r->bal = or->bal;
 193         r->n[0] = or->n[0];
 194         r->n[1] = or->n[1];
 195         *tp = r;
 196         (*tp)->p = p;
 197         /* node has changed; fix children's parent pointers */
 198         if(r->n[0])
 199                 r->n[0]->p = r;
 200         if(r->n[1])
 201                 r->n[1]->p = r;
 202         *del = or;
 203         balance(tp, p);
 204         return -(ob != 0 && (*tp)->bal == 0);
 205 }
 206
 207 struct Avltree
 208 {
 209         Avl     *root;
 210         int     (*cmp)(Avl*, Avl*);
 211         Avlwalk *walks;
 212 };
 213 struct Avlwalk
 214 {
 215         int     started;
 216         int     moved;
 217         Avlwalk *next;
 218         Avltree *tree;
 219         Avl     *node;
 220 };
 221
 222 Avltree*
 223 mkavltree(int (*cmp)(Avl*, Avl*))
 224 {
 225         Avltree *t;
 226
 227         t = malloc(sizeof *t);
 228         if(t == nil)
 229                 return nil;
 230         memset(t, 0, sizeof *t);
 231         t->cmp = cmp;
 232         return t;
 233 }
 234
 235 void
 236 insertavl(Avltree *t, Avl *new, Avl **oldp)
 237 {
 238         *oldp = nil;
 239         _insertavl(&t->root, nil, new, t->cmp, oldp);
 240 }
 241
 242 static Avl*
 243 findpredecessor(Avl *a)
 244 {
 245         if(a == nil)
 246                 return nil;
 247
 248         if(a->n[0] != nil){
 249                 /* predecessor is rightmost descendant of left child */
 250                 for(a = a->n[0]; a->n[1]; a = a->n[1])
 251                         ;
 252                 return a;
 253         }else{
 254                 /* we're at a leaf, successor is a parent we enter from the right */
 255                 while(a->p && a->p->n[0] == a)
 256                         a = a->p;
 257                 return a->p;
 258         }
 259 }
 260
 261 static Avl*
 262 findsuccessor(Avl *a)
 263 {
 264         if(a == nil)
 265                 return nil;
 266
 267         if(a->n[1] != nil){
 268                 /* successor is leftmost descendant of right child */
 269                 for(a = a->n[1]; a->n[0]; a = a->n[0])
 270                         ;
 271                 return a;
 272         }else{
 273                 /* we're at a leaf, successor is a parent we enter from the left going up */
 274                 while(a->p && a->p->n[1] == a)
 275                         a = a->p;
 276                 return a->p;
 277         }
 278 }
 279
 280 static Avl*
 281 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor)
 282 {
 283         int i;
 284         Avl *p;
 285
 286         p = nil;
 287         if(t == nil)
 288                 return nil;
 289         do{
 290                 assert(t->p == p);
 291                 if((i = canoncmp(cmp(r, t))) == 0)
 292                         return t;
 293                 p = t;
 294                 t = t->n[(i+1)/2];
 295         }while(t);
 296         if(neighbor == 0)
 297                 return nil;
 298         if(neighbor < 0)
 299                 return i > 0 ? p : findpredecessor(p);
 300         return i < 0 ? p : findsuccessor(p);
 301 }
 302
 303 Avl*
 304 searchavl(Avltree *t, Avl *key, int neighbor)
 305 {
 306         return _lookupavl(t->root, key, t->cmp, neighbor);
 307 }
 308
 309 Avl*
 310 lookupavl(Avltree *t, Avl *key)
 311 {
 312         return _lookupavl(t->root, key, t->cmp, 0);
 313 }
 314
 315 static void
 316 walkdel(Avl *a, void *v)
 317 {
 318         Avl *p;
 319         Avlwalk *w;
 320         Avltree *t;
 321
 322         if(a == nil)
 323                 return;
 324
 325         p = findpredecessor(a);
 326         t = v;
 327         for(w = t->walks; w; w = w->next){
 328                 if(w->node == a){
 329                         /* back pointer to predecessor; not perfect but adequate */
 330                         w->moved = 1;
 331                         w->node = p;
 332                         if(p == nil)
 333                                 w->started = 0;
 334                 }
 335         }
 336 }
 337
 338 void
 339 deleteavl(Avltree *t, Avl *key, Avl **oldp)
 340 {
 341         *oldp = nil;
 342         _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
 343 }
 344
 345 Avlwalk*
 346 avlwalk(Avltree *t)
 347 {
 348         Avlwalk *w;
 349
 350         w = malloc(sizeof *w);
 351         if(w == nil)
 352                 return nil;
 353         memset(w, 0, sizeof *w);
 354         w->tree = t;
 355         w->next = t->walks;
 356         t->walks = w;
 357         return w;
 358 }
 359
 360 Avl*
 361 avlnext(Avlwalk *w)
 362 {
 363         Avl *a;
 364
 365         if(w->started==0){
 366                 for(a = w->tree->root; a && a->n[0]; a = a->n[0])
 367                         ;
 368                 w->node = a;
 369                 w->started = 1;
 370         }else{
 371                 a = findsuccessor(w->node);
 372                 if(a == w->node)
 373                         abort();
 374                 w->node = a;
 375         }
 376         return w->node;
 377 }
 378
 379 Avl*
 380 avlprev(Avlwalk *w)
 381 {
 382         Avl *a;
 383
 384         if(w->started == 0){
 385                 for(a = w->tree->root; a && a->n[1]; a = a->n[1])
 386                         ;
 387                 w->node = a;
 388                 w->started = 1;
 389         }else if(w->moved){
 390                 w->moved = 0;
 391                 return w->node;
 392         }else{
 393                 a = findpredecessor(w->node);
 394                 if(a == w->node)
 395                         abort();
 396                 w->node = a;
 397         }
 398         return w->node;
 399 }
 400
 401 void
 402 endwalk(Avlwalk *w)
 403 {
 404         Avltree *t;
 405         Avlwalk **l;
 406
 407         t = w->tree;
 408         for(l = &t->walks; *l; l = &(*l)->next){
 409                 if(*l == w){
 410                         *l = w->next;
 411                         break;
 412                 }
 413         }
 414         free(w);
 415 }