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 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree) 101 { 102 int i, ob; 103 104 if(*tp == nil){ 105 r->bal = 0; 106 r->n[0] = nil; 107 r->n[1] = nil; 108 r->p = p; 109 *tp = r; 110 return 1; 111 } 112 ob = (*tp)->bal; 113 if((i = cmp(r, *tp)) != 0){ 114 (*tp)->bal += i * _insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp, 115 rfree); 116 balance(tp, p); 117 return ob == 0 && (*tp)->bal != 0; 118 } 119 120 /* install new entry */ 121 *rfree = *tp; /* save old node for freeing */ 122 *tp = r; /* insert new node */ 123 **tp = **rfree; /* copy old node's Avl contents */ 124 if(r->n[0]) /* fix node's children's parent pointers */ 125 r->n[0]->p = r; 126 if(r->n[1]) 127 r->n[1]->p = r; 128 129 return 0; 130 } 131 132 static int 133 successor(Avl **tp, Avl *p, Avl **r) 134 { 135 int ob; 136 137 if((*tp)->n[0] == nil){ 138 *r = *tp; 139 *tp = (*r)->n[1]; 140 if(*tp) 141 (*tp)->p = p; 142 return -1; 143 } 144 ob = (*tp)->bal; 145 (*tp)->bal -= successor(&(*tp)->n[0], *tp, r); 146 balance(tp, p); 147 return -(ob != 0 && (*tp)->bal == 0); 148 } 149 150 static int 151 _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del, 152 void (*predel)(Avl*, void*), void *arg) 153 { 154 int i, ob; 155 Avl *r, *or; 156 157 if(*tp == nil) 158 return 0; 159 160 ob = (*tp)->bal; 161 if((i=cmp(rx, *tp)) != 0){ 162 (*tp)->bal += i * _deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp, 163 del, predel, arg); 164 balance(tp, p); 165 return -(ob != 0 && (*tp)->bal == 0); 166 } 167 168 if(predel) 169 (*predel)(*tp, arg); 170 171 or = *tp; 172 if(or->n[i=0] == nil || or->n[i=1] == nil){ 173 *tp = or->n[1-i]; 174 if(*tp) 175 (*tp)->p = p; 176 *del = or; 177 return -1; 178 } 179 180 /* deleting node with two kids, find successor */ 181 or->bal += successor(&or->n[1], or, &r); 182 r->bal = or->bal; 183 r->n[0] = or->n[0]; 184 r->n[1] = or->n[1]; 185 *tp = r; 186 (*tp)->p = p; 187 /* node has changed; fix children's parent pointers */ 188 if(r->n[0]) 189 r->n[0]->p = r; 190 if(r->n[1]) 191 r->n[1]->p = r; 192 *del = or; 193 balance(tp, p); 194 return -(ob != 0 && (*tp)->bal == 0); 195 } 196 197 static void 198 checkparents(Avl *a, Avl *p) 199 { 200 if(a == nil) 201 return; 202 if(a->p != p) 203 print("bad parent\n"); 204 checkparents(a->n[0], a); 205 checkparents(a->n[1], a); 206 } 207 208 struct Avltree 209 { 210 Avl *root; 211 int (*cmp)(Avl*, Avl*); 212 Avlwalk *walks; 213 }; 214 struct Avlwalk 215 { 216 int started; 217 int moved; 218 Avlwalk *next; 219 Avltree *tree; 220 Avl *node; 221 }; 222 223 Avltree* 224 mkavltree(int (*cmp)(Avl*, Avl*)) 225 { 226 Avltree *t; 227 228 t = malloc(sizeof *t); 229 if(t == nil) 230 return nil; 231 memset(t, 0, sizeof *t); 232 t->cmp = cmp; 233 return t; 234 } 235 236 void 237 insertavl(Avltree *t, Avl *new, Avl **oldp) 238 { 239 *oldp = nil; 240 _insertavl(&t->root, nil, new, t->cmp, oldp); 241 } 242 243 static Avl* 244 findpredecessor(Avl *a) 245 { 246 if(a == nil) 247 return nil; 248 249 if(a->n[0] != nil){ 250 /* predecessor is rightmost descendant of left child */ 251 for(a = a->n[0]; a->n[1]; a = a->n[1]) 252 ; 253 return a; 254 }else{ 255 /* we're at a leaf, successor is a parent we enter from the right */ 256 while(a->p && a->p->n[0] == a) 257 a = a->p; 258 return a->p; 259 } 260 } 261 262 static Avl* 263 findsuccessor(Avl *a) 264 { 265 if(a == nil) 266 return nil; 267 268 if(a->n[1] != nil){ 269 /* successor is leftmost descendant of right child */ 270 for(a = a->n[1]; a->n[0]; a = a->n[0]) 271 ; 272 return a; 273 }else{ 274 /* we're at a leaf, successor is a parent we enter from the left going up */ 275 while(a->p && a->p->n[1] == a) 276 a = a->p; 277 return a->p; 278 } 279 } 280 281 static Avl* 282 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor) 283 { 284 int i; 285 Avl *p; 286 287 p = nil; 288 if(t == nil) 289 return nil; 290 do{ 291 assert(t->p == p); 292 if((i = cmp(r, t)) == 0) 293 return t; 294 p = t; 295 t = t->n[(i+1)/2]; 296 }while(t); 297 if(neighbor == 0) 298 return nil; 299 if(neighbor < 0) 300 return i > 0 ? p : findpredecessor(p); 301 return i < 0 ? p : findsuccessor(p); 302 } 303 304 Avl* 305 searchavl(Avltree *t, Avl *key, int neighbor) 306 { 307 return _lookupavl(t->root, key, t->cmp, neighbor); 308 } 309 310 Avl* 311 lookupavl(Avltree *t, Avl *key) 312 { 313 return _lookupavl(t->root, key, t->cmp, 0); 314 } 315 316 static void 317 walkdel(Avl *a, void *v) 318 { 319 Avl *p; 320 Avlwalk *w; 321 Avltree *t; 322 323 if(a == nil) 324 return; 325 326 p = findpredecessor(a); 327 t = v; 328 for(w = t->walks; w; w = w->next){ 329 if(w->node == a){ 330 /* back pointer to predecessor; not perfect but adequate */ 331 w->moved = 1; 332 w->node = p; 333 if(p == nil) 334 w->started = 0; 335 } 336 } 337 } 338 339 void 340 deleteavl(Avltree *t, Avl *key, Avl **oldp) 341 { 342 *oldp = nil; 343 _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t); 344 } 345 346 Avlwalk* 347 avlwalk(Avltree *t) 348 { 349 Avlwalk *w; 350 351 w = malloc(sizeof *w); 352 if(w == nil) 353 return nil; 354 memset(w, 0, sizeof *w); 355 w->tree = t; 356 w->next = t->walks; 357 t->walks = w; 358 return w; 359 } 360 361 Avl* 362 avlnext(Avlwalk *w) 363 { 364 Avl *a; 365 366 if(w->started==0){ 367 for(a = w->tree->root; a && a->n[0]; a = a->n[0]) 368 ; 369 w->node = a; 370 w->started = 1; 371 }else{ 372 a = findsuccessor(w->node); 373 if(a == w->node) 374 abort(); 375 w->node = a; 376 } 377 return w->node; 378 } 379 380 Avl* 381 avlprev(Avlwalk *w) 382 { 383 Avl *a; 384 385 if(w->started == 0){ 386 for(a = w->tree->root; a && a->n[1]; a = a->n[1]) 387 ; 388 w->node = a; 389 w->started = 1; 390 }else if(w->moved){ 391 w->moved = 0; 392 return w->node; 393 }else{ 394 a = findpredecessor(w->node); 395 if(a == w->node) 396 abort(); 397 w->node = a; 398 } 399 return w->node; 400 } 401 402 void 403 endwalk(Avlwalk *w) 404 { 405 Avltree *t; 406 Avlwalk **l; 407 408 t = w->tree; 409 for(l = &t->walks; *l; l = &(*l)->next){ 410 if(*l == w){ 411 *l = w->next; 412 break; 413 } 414 } 415 free(w); 416 } 417 418 static void 419 walkavl(Avl *t, void (*f)(Avl*, void*), void *v) 420 { 421 if(t == nil) 422 return; 423 walkavl(t->n[0], f, v); 424 f(t, v); 425 walkavl(t->n[1], f, v); 426 } 427