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
singleleft(Avl ** tp,Avl * p)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
singleright(Avl ** tp,Avl * p)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
doublerightleft(Avl ** tp,Avl * p)62 doublerightleft(Avl **tp, Avl *p)
63 {
64 singleright(&(*tp)->n[1], *tp);
65 singleleft(tp, p);
66 }
67
68 static void
doubleleftright(Avl ** tp,Avl * p)69 doubleleftright(Avl **tp, Avl *p)
70 {
71 singleleft(&(*tp)->n[0], *tp);
72 singleright(tp, p);
73 }
74
75 static void
balance(Avl ** tp,Avl * p)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
canoncmp(int cmp)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
_insertavl(Avl ** tp,Avl * p,Avl * r,int (* cmp)(Avl *,Avl *),Avl ** rfree)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
successor(Avl ** tp,Avl * p,Avl ** r)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
_deleteavl(Avl ** tp,Avl * p,Avl * rx,int (* cmp)(Avl *,Avl *),Avl ** del,void (* predel)(Avl *,void *),void * arg)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
checkparents(Avl * a,Avl * p)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*
mkavltree(int (* cmp)(Avl *,Avl *))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
insertavl(Avltree * t,Avl * new,Avl ** oldp)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*
findpredecessor(Avl * a)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*
findsuccessor(Avl * a)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*
_lookupavl(Avl * t,Avl * r,int (* cmp)(Avl *,Avl *),int neighbor)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*
searchavl(Avltree * t,Avl * key,int neighbor)315 searchavl(Avltree *t, Avl *key, int neighbor)
316 {
317 return _lookupavl(t->root, key, t->cmp, neighbor);
318 }
319
320 Avl*
lookupavl(Avltree * t,Avl * key)321 lookupavl(Avltree *t, Avl *key)
322 {
323 return _lookupavl(t->root, key, t->cmp, 0);
324 }
325
326 static void
walkdel(Avl * a,void * v)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
deleteavl(Avltree * t,Avl * key,Avl ** oldp)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*
avlwalk(Avltree * t)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*
avlnext(Avlwalk * w)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*
avlprev(Avlwalk * w)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
endwalk(Avlwalk * w)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
walkavl(Avl * t,void (* f)(Avl *,void *),void * v)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