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 } 437