Plan 9 from Bell Labs’s /usr/web/sources/plan9/sys/src/libavl/avl.c

 ``` #include #include #include #include /* * In-memory database stored as self-balancing AVL tree. * See Lewis & Denenberg, Data Structures and Their Algorithms. */ static void singleleft(Avl **tp, Avl *p) { int l, r2; Avl *a, *c; a = *tp; c = a->n[1]; r2 = c->bal; l = (r2 > 0? r2: 0)+1 - a->bal; if((a->n[1] = c->n[0]) != nil) a->n[1]->p = a; if((c->n[0] = a) != nil) c->n[0]->p = c; if((*tp = c) != nil) (*tp)->p = p; a->bal = -l; c->bal = r2 - ((l > 0? l: 0)+1); } static void singleright(Avl **tp, Avl *p) { int l2, r; Avl *a, *c; a = *tp; c = a->n[0]; l2 = - c->bal; r = a->bal + ((l2 > 0? l2: 0)+1); if((a->n[0] = c->n[1]) != nil) a->n[0]->p = a; if((c->n[1] = a) != nil) c->n[1]->p = c; if((*tp = c) != nil) (*tp)->p = p; a->bal = r; c->bal = ((r > 0? r: 0)+1) - l2; } static void doublerightleft(Avl **tp, Avl *p) { singleright(&(*tp)->n[1], *tp); singleleft(tp, p); } static void doubleleftright(Avl **tp, Avl *p) { singleleft(&(*tp)->n[0], *tp); singleright(tp, p); } static void balance(Avl **tp, Avl *p) { switch((*tp)->bal){ case -2: if((*tp)->n[0]->bal <= 0) singleright(tp, p); else if((*tp)->n[0]->bal == 1) doubleleftright(tp, p); else assert(0); break; case 2: if((*tp)->n[1]->bal >= 0) singleleft(tp, p); else if((*tp)->n[1]->bal == -1) doublerightleft(tp, p); else assert(0); break; } } static int canoncmp(int cmp) { if(cmp < 0) return -1; else if(cmp > 0) return 1; return 0; } static int _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree) { int i, ob; if(*tp == nil){ r->bal = 0; r->n[0] = nil; r->n[1] = nil; r->p = p; *tp = r; return 1; } ob = (*tp)->bal; if((i = canoncmp(cmp(r, *tp))) != 0){ (*tp)->bal += i * _insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp, rfree); balance(tp, p); return ob == 0 && (*tp)->bal != 0; } /* install new entry */ *rfree = *tp; /* save old node for freeing */ *tp = r; /* insert new node */ **tp = **rfree; /* copy old node's Avl contents */ if(r->n[0]) /* fix node's children's parent pointers */ r->n[0]->p = r; if(r->n[1]) r->n[1]->p = r; return 0; } static int successor(Avl **tp, Avl *p, Avl **r) { int ob; if((*tp)->n[0] == nil){ *r = *tp; *tp = (*r)->n[1]; if(*tp) (*tp)->p = p; return -1; } ob = (*tp)->bal; (*tp)->bal -= successor(&(*tp)->n[0], *tp, r); balance(tp, p); return -(ob != 0 && (*tp)->bal == 0); } static int _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del, void (*predel)(Avl*, void*), void *arg) { int i, ob; Avl *r, *or; if(*tp == nil) return 0; ob = (*tp)->bal; if((i=canoncmp(cmp(rx, *tp))) != 0){ (*tp)->bal += i * _deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp, del, predel, arg); balance(tp, p); return -(ob != 0 && (*tp)->bal == 0); } if(predel) (*predel)(*tp, arg); or = *tp; if(or->n[i=0] == nil || or->n[i=1] == nil){ *tp = or->n[1-i]; if(*tp) (*tp)->p = p; *del = or; return -1; } /* deleting node with two kids, find successor */ or->bal += successor(&or->n[1], or, &r); r->bal = or->bal; r->n[0] = or->n[0]; r->n[1] = or->n[1]; *tp = r; (*tp)->p = p; /* node has changed; fix children's parent pointers */ if(r->n[0]) r->n[0]->p = r; if(r->n[1]) r->n[1]->p = r; *del = or; balance(tp, p); return -(ob != 0 && (*tp)->bal == 0); } static void checkparents(Avl *a, Avl *p) { if(a == nil) return; if(a->p != p) print("bad parent\n"); checkparents(a->n[0], a); checkparents(a->n[1], a); } struct Avltree { Avl *root; int (*cmp)(Avl*, Avl*); Avlwalk *walks; }; struct Avlwalk { int started; int moved; Avlwalk *next; Avltree *tree; Avl *node; }; Avltree* mkavltree(int (*cmp)(Avl*, Avl*)) { Avltree *t; t = malloc(sizeof *t); if(t == nil) return nil; memset(t, 0, sizeof *t); t->cmp = cmp; return t; } void insertavl(Avltree *t, Avl *new, Avl **oldp) { *oldp = nil; _insertavl(&t->root, nil, new, t->cmp, oldp); } static Avl* findpredecessor(Avl *a) { if(a == nil) return nil; if(a->n[0] != nil){ /* predecessor is rightmost descendant of left child */ for(a = a->n[0]; a->n[1]; a = a->n[1]) ; return a; }else{ /* we're at a leaf, successor is a parent we enter from the right */ while(a->p && a->p->n[0] == a) a = a->p; return a->p; } } static Avl* findsuccessor(Avl *a) { if(a == nil) return nil; if(a->n[1] != nil){ /* successor is leftmost descendant of right child */ for(a = a->n[1]; a->n[0]; a = a->n[0]) ; return a; }else{ /* we're at a leaf, successor is a parent we enter from the left going up */ while(a->p && a->p->n[1] == a) a = a->p; return a->p; } } static Avl* _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor) { int i; Avl *p; p = nil; if(t == nil) return nil; do{ assert(t->p == p); if((i = canoncmp(cmp(r, t))) == 0) return t; p = t; t = t->n[(i+1)/2]; }while(t); if(neighbor == 0) return nil; if(neighbor < 0) return i > 0 ? p : findpredecessor(p); return i < 0 ? p : findsuccessor(p); } Avl* searchavl(Avltree *t, Avl *key, int neighbor) { return _lookupavl(t->root, key, t->cmp, neighbor); } Avl* lookupavl(Avltree *t, Avl *key) { return _lookupavl(t->root, key, t->cmp, 0); } static void walkdel(Avl *a, void *v) { Avl *p; Avlwalk *w; Avltree *t; if(a == nil) return; p = findpredecessor(a); t = v; for(w = t->walks; w; w = w->next){ if(w->node == a){ /* back pointer to predecessor; not perfect but adequate */ w->moved = 1; w->node = p; if(p == nil) w->started = 0; } } } void deleteavl(Avltree *t, Avl *key, Avl **oldp) { *oldp = nil; _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t); } Avlwalk* avlwalk(Avltree *t) { Avlwalk *w; w = malloc(sizeof *w); if(w == nil) return nil; memset(w, 0, sizeof *w); w->tree = t; w->next = t->walks; t->walks = w; return w; } Avl* avlnext(Avlwalk *w) { Avl *a; if(w->started==0){ for(a = w->tree->root; a && a->n[0]; a = a->n[0]) ; w->node = a; w->started = 1; }else{ a = findsuccessor(w->node); if(a == w->node) abort(); w->node = a; } return w->node; } Avl* avlprev(Avlwalk *w) { Avl *a; if(w->started == 0){ for(a = w->tree->root; a && a->n[1]; a = a->n[1]) ; w->node = a; w->started = 1; }else if(w->moved){ w->moved = 0; return w->node; }else{ a = findpredecessor(w->node); if(a == w->node) abort(); w->node = a; } return w->node; } void endwalk(Avlwalk *w) { Avltree *t; Avlwalk **l; t = w->tree; for(l = &t->walks; *l; l = &(*l)->next){ if(*l == w){ *l = w->next; break; } } free(w); } static void walkavl(Avl *t, void (*f)(Avl*, void*), void *v) { if(t == nil) return; walkavl(t->n[0], f, v); f(t, v); walkavl(t->n[1], f, v); } ```