Blob
1 #include <u.h>2 #include <libc.h>3 #include <bio.h>4 #include <avl.h>6 /*7 * In-memory database stored as self-balancing AVL tree.8 * See Lewis & Denenberg, Data Structures and Their Algorithms.9 */11 static void12 singleleft(Avl **tp, Avl *p)13 {14 int l, r2;15 Avl *a, *c;17 a = *tp;18 c = a->n[1];20 r2 = c->bal;21 l = (r2 > 0? r2: 0)+1 - a->bal;23 if((a->n[1] = c->n[0]) != nil)24 a->n[1]->p = a;26 if((c->n[0] = a) != nil)27 c->n[0]->p = c;29 if((*tp = c) != nil)30 (*tp)->p = p;32 a->bal = -l;33 c->bal = r2 - ((l > 0? l: 0)+1);35 }37 static void38 singleright(Avl **tp, Avl *p)39 {40 int l2, r;41 Avl *a, *c;43 a = *tp;44 c = a->n[0];45 l2 = - c->bal;46 r = a->bal + ((l2 > 0? l2: 0)+1);48 if((a->n[0] = c->n[1]) != nil)49 a->n[0]->p = a;51 if((c->n[1] = a) != nil)52 c->n[1]->p = c;54 if((*tp = c) != nil)55 (*tp)->p = p;57 a->bal = r;58 c->bal = ((r > 0? r: 0)+1) - l2;59 }61 static void62 doublerightleft(Avl **tp, Avl *p)63 {64 singleright(&(*tp)->n[1], *tp);65 singleleft(tp, p);66 }68 static void69 doubleleftright(Avl **tp, Avl *p)70 {71 singleleft(&(*tp)->n[0], *tp);72 singleright(tp, p);73 }75 static void76 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 else85 assert(0);86 break;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 else94 assert(0);95 break;96 }97 }99 static int100 canoncmp(int cmp)101 {102 if(cmp < 0)103 return -1;104 else if(cmp > 0)105 return 1;106 return 0;107 }109 static int110 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree)111 {112 int i, ob;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 }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;139 return 0;140 }142 static int143 successor(Avl **tp, Avl *p, Avl **r)144 {145 int ob;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 }160 static int161 _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;167 if(*tp == nil)168 return 0;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 }178 if(predel)179 (*predel)(*tp, arg);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 }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 }207 /*208 static void209 checkparents(Avl *a, Avl *p)210 {211 if(a == nil)212 return;213 if(a->p != p)214 print("bad parent\n");215 checkparents(a->n[0], a);216 checkparents(a->n[1], a);217 }218 */220 struct Avltree221 {222 Avl *root;223 int (*cmp)(Avl*, Avl*);224 Avlwalk *walks;225 };226 struct Avlwalk227 {228 int started;229 int moved;230 Avlwalk *next;231 Avltree *tree;232 Avl *node;233 };235 Avltree*236 mkavltree(int (*cmp)(Avl*, Avl*))237 {238 Avltree *t;240 t = malloc(sizeof *t);241 if(t == nil)242 return nil;243 memset(t, 0, sizeof *t);244 t->cmp = cmp;245 return t;246 }248 void249 insertavl(Avltree *t, Avl *new, Avl **oldp)250 {251 *oldp = nil;252 _insertavl(&t->root, nil, new, t->cmp, oldp);253 }255 static Avl*256 findpredecessor(Avl *a)257 {258 if(a == nil)259 return nil;261 if(a->n[0] != nil){262 /* predecessor is rightmost descendant of left child */263 for(a = a->n[0]; a->n[1]; a = a->n[1])264 ;265 return a;266 }else{267 /* we're at a leaf, successor is a parent we enter from the right */268 while(a->p && a->p->n[0] == a)269 a = a->p;270 return a->p;271 }272 }274 static Avl*275 findsuccessor(Avl *a)276 {277 if(a == nil)278 return nil;280 if(a->n[1] != nil){281 /* successor is leftmost descendant of right child */282 for(a = a->n[1]; a->n[0]; a = a->n[0])283 ;284 return a;285 }else{286 /* we're at a leaf, successor is a parent we enter from the left going up */287 while(a->p && a->p->n[1] == a)288 a = a->p;289 return a->p;290 }291 }293 static Avl*294 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor)295 {296 int i;297 Avl *p;299 p = nil;300 if(t == nil)301 return nil;302 do{303 assert(t->p == p);304 if((i = canoncmp(cmp(r, t))) == 0)305 return t;306 p = t;307 t = t->n[(i+1)/2];308 }while(t);309 if(neighbor == 0)310 return nil;311 if(neighbor < 0)312 return i > 0 ? p : findpredecessor(p);313 return i < 0 ? p : findsuccessor(p);314 }316 Avl*317 searchavl(Avltree *t, Avl *key, int neighbor)318 {319 return _lookupavl(t->root, key, t->cmp, neighbor);320 }322 Avl*323 lookupavl(Avltree *t, Avl *key)324 {325 return _lookupavl(t->root, key, t->cmp, 0);326 }328 static void329 walkdel(Avl *a, void *v)330 {331 Avl *p;332 Avlwalk *w;333 Avltree *t;335 if(a == nil)336 return;338 p = findpredecessor(a);339 t = v;340 for(w = t->walks; w; w = w->next){341 if(w->node == a){342 /* back pointer to predecessor; not perfect but adequate */343 w->moved = 1;344 w->node = p;345 if(p == nil)346 w->started = 0;347 }348 }349 }351 void352 deleteavl(Avltree *t, Avl *key, Avl **oldp)353 {354 *oldp = nil;355 _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);356 }358 Avlwalk*359 avlwalk(Avltree *t)360 {361 Avlwalk *w;363 w = malloc(sizeof *w);364 if(w == nil)365 return nil;366 memset(w, 0, sizeof *w);367 w->tree = t;368 w->next = t->walks;369 t->walks = w;370 return w;371 }373 Avl*374 avlnext(Avlwalk *w)375 {376 Avl *a;378 if(w->started==0){379 for(a = w->tree->root; a && a->n[0]; a = a->n[0])380 ;381 w->node = a;382 w->started = 1;383 }else{384 a = findsuccessor(w->node);385 if(a == w->node)386 abort();387 w->node = a;388 }389 return w->node;390 }392 Avl*393 avlprev(Avlwalk *w)394 {395 Avl *a;397 if(w->started == 0){398 for(a = w->tree->root; a && a->n[1]; a = a->n[1])399 ;400 w->node = a;401 w->started = 1;402 }else if(w->moved){403 w->moved = 0;404 return w->node;405 }else{406 a = findpredecessor(w->node);407 if(a == w->node)408 abort();409 w->node = a;410 }411 return w->node;412 }414 void415 endwalk(Avlwalk *w)416 {417 Avltree *t;418 Avlwalk **l;420 t = w->tree;421 for(l = &t->walks; *l; l = &(*l)->next){422 if(*l == w){423 *l = w->next;424 break;425 }426 }427 free(w);428 }430 /*431 static void432 walkavl(Avl *t, void (*f)(Avl*, void*), void *v)433 {434 if(t == nil)435 return;436 walkavl(t->n[0], f, v);437 f(t, v);438 walkavl(t->n[1], f, v);439 }440 */