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 static void
 208 checkparents(Avl *a, Avl *p)
 209 {
 210         if(a == nil)
 211                 return;
 212         if(a->p != p)
 213                 print("bad parent\n");
 214         checkparents(a->n[0], a);
 215         checkparents(a->n[1], a);
 216 }
 217
 218 struct Avltree
 219 {
 220         Avl     *root;
 221         int     (*cmp)(Avl*, Avl*);
 222         Avlwalk *walks;
 223 };
 224 struct Avlwalk
 225 {
 226         int     started;
 227         int     moved;
 228         Avlwalk *next;
 229         Avltree *tree;
 230         Avl     *node;
 231 };
 232
 233 Avltree*
 234 mkavltree(int (*cmp)(Avl*, Avl*))
 235 {
 236         Avltree *t;
 237
 238         t = malloc(sizeof *t);
 239         if(t == nil)
 240                 return nil;
 241         memset(t, 0, sizeof *t);
 242         t->cmp = cmp;
 243         return t;
 244 }
 245
 246 void
 247 insertavl(Avltree *t, Avl *new, Avl **oldp)
 248 {
 249         *oldp = nil;
 250         _insertavl(&t->root, nil, new, t->cmp, oldp);
 251 }
 252
 253 static Avl*
 254 findpredecessor(Avl *a)
 255 {
 256         if(a == nil)
 257                 return nil;
 258
 259         if(a->n[0] != nil){
 260                 /* predecessor is rightmost descendant of left child */
 261                 for(a = a->n[0]; a->n[1]; a = a->n[1])
 262                         ;
 263                 return a;
 264         }else{
 265                 /* we're at a leaf, successor is a parent we enter from the right */
 266                 while(a->p && a->p->n[0] == a)
 267                         a = a->p;
 268                 return a->p;
 269         }
 270 }
 271
 272 static Avl*
 273 findsuccessor(Avl *a)
 274 {
 275         if(a == nil)
 276                 return nil;
 277
 278         if(a->n[1] != nil){
 279                 /* successor is leftmost descendant of right child */
 280                 for(a = a->n[1]; a->n[0]; a = a->n[0])
 281                         ;
 282                 return a;
 283         }else{
 284                 /* we're at a leaf, successor is a parent we enter from the left going up */
 285                 while(a->p && a->p->n[1] == a)
 286                         a = a->p;
 287                 return a->p;
 288         }
 289 }
 290
 291 static Avl*
 292 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor)
 293 {
 294         int i;
 295         Avl *p;
 296
 297         p = nil;
 298         if(t == nil)
 299                 return nil;
 300         do{
 301                 assert(t->p == p);
 302                 if((i = canoncmp(cmp(r, t))) == 0)
 303                         return t;
 304                 p = t;
 305                 t = t->n[(i+1)/2];
 306         }while(t);
 307         if(neighbor == 0)
 308                 return nil;
 309         if(neighbor < 0)
 310                 return i > 0 ? p : findpredecessor(p);
 311         return i < 0 ? p : findsuccessor(p);
 312 }
 313
 314 Avl*
 315 searchavl(Avltree *t, Avl *key, int neighbor)
 316 {
 317         return _lookupavl(t->root, key, t->cmp, neighbor);
 318 }
 319
 320 Avl*
 321 lookupavl(Avltree *t, Avl *key)
 322 {
 323         return _lookupavl(t->root, key, t->cmp, 0);
 324 }
 325
 326 static void
 327 walkdel(Avl *a, void *v)
 328 {
 329         Avl *p;
 330         Avlwalk *w;
 331         Avltree *t;
 332
 333         if(a == nil)
 334                 return;
 335
 336         p = findpredecessor(a);
 337         t = v;
 338         for(w = t->walks; w; w = w->next){
 339                 if(w->node == a){
 340                         /* back pointer to predecessor; not perfect but adequate */
 341                         w->moved = 1;
 342                         w->node = p;
 343                         if(p == nil)
 344                                 w->started = 0;
 345                 }
 346         }
 347 }
 348
 349 void
 350 deleteavl(Avltree *t, Avl *key, Avl **oldp)
 351 {
 352         *oldp = nil;
 353         _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
 354 }
 355
 356 Avlwalk*
 357 avlwalk(Avltree *t)
 358 {
 359         Avlwalk *w;
 360
 361         w = malloc(sizeof *w);
 362         if(w == nil)
 363                 return nil;
 364         memset(w, 0, sizeof *w);
 365         w->tree = t;
 366         w->next = t->walks;
 367         t->walks = w;
 368         return w;
 369 }
 370
 371 Avl*
 372 avlnext(Avlwalk *w)
 373 {
 374         Avl *a;
 375
 376         if(w->started==0){
 377                 for(a = w->tree->root; a && a->n[0]; a = a->n[0])
 378                         ;
 379                 w->node = a;
 380                 w->started = 1;
 381         }else{
 382                 a = findsuccessor(w->node);
 383                 if(a == w->node)
 384                         abort();
 385                 w->node = a;
 386         }
 387         return w->node;
 388 }
 389
 390 Avl*
 391 avlprev(Avlwalk *w)
 392 {
 393         Avl *a;
 394
 395         if(w->started == 0){
 396                 for(a = w->tree->root; a && a->n[1]; a = a->n[1])
 397                         ;
 398                 w->node = a;
 399                 w->started = 1;
 400         }else if(w->moved){
 401                 w->moved = 0;
 402                 return w->node;
 403         }else{
 404                 a = findpredecessor(w->node);
 405                 if(a == w->node)
 406                         abort();
 407                 w->node = a;
 408         }
 409         return w->node;
 410 }
 411
 412 void
 413 endwalk(Avlwalk *w)
 414 {
 415         Avltree *t;
 416         Avlwalk **l;
 417
 418         t = w->tree;
 419         for(l = &t->walks; *l; l = &(*l)->next){
 420                 if(*l == w){
 421                         *l = w->next;
 422                         break;
 423                 }
 424         }
 425         free(w);
 426 }
 427
 428 static void
 429 walkavl(Avl *t, void (*f)(Avl*, void*), void *v)
 430 {
 431         if(t == nil)
 432                 return;
 433         walkavl(t->n[0], f, v);
 434         f(t, v);
 435         walkavl(t->n[1], f, v);
 436 }