1 #include "all.h"
2
3 /*
4 * In-memory database stored as self-balancing AVL tree.
5 * See Lewis & Denenberg, Data Structures and Their Algorithms.
6 */
7 static void
singleleft(Avl ** tp,Avl * p)8 singleleft(Avl **tp, Avl *p)
9 {
10 Avl *a, *c;
11 int l, r2;
12
13 a = *tp;
14 c = a->n[1];
15
16 r2 = c->bal;
17 l = (r2 > 0 ? r2 : 0)+1 - a->bal;
18
19 if((a->n[1] = c->n[0]) != nil)
20 a->n[1]->p = a;
21
22 if((c->n[0] = a) != nil)
23 c->n[0]->p = c;
24
25 if((*tp = c) != nil)
26 (*tp)->p = p;
27
28 a->bal = -l;
29 c->bal = r2 - ((l > 0 ? l : 0)+1);
30
31 }
32
33 static void
singleright(Avl ** tp,Avl * p)34 singleright(Avl **tp, Avl *p)
35 {
36 Avl *a, *c;
37 int l2, r;
38
39 a = *tp;
40 c = a->n[0];
41 l2 = - c->bal;
42 r = a->bal + ((l2 > 0 ? l2 : 0)+1);
43
44 if((a->n[0] = c->n[1]) != nil)
45 a->n[0]->p = a;
46
47 if((c->n[1] = a) != nil)
48 c->n[1]->p = c;
49
50 if((*tp = c) != nil)
51 (*tp)->p = p;
52
53 a->bal = r;
54 c->bal = ((r > 0 ? r : 0)+1) - l2;
55 }
56
57 static void
doublerightleft(Avl ** tp,Avl * p)58 doublerightleft(Avl **tp, Avl *p)
59 {
60 singleright(&(*tp)->n[1], *tp);
61 singleleft(tp, p);
62 }
63
64 static void
doubleleftright(Avl ** tp,Avl * p)65 doubleleftright(Avl **tp, Avl *p)
66 {
67 singleleft(&(*tp)->n[0], *tp);
68 singleright(tp, p);
69 }
70
71 static void
balance(Avl ** tp,Avl * p)72 balance(Avl **tp, Avl *p)
73 {
74 switch((*tp)->bal){
75 case -2:
76 if((*tp)->n[0]->bal <= 0)
77 singleright(tp, p);
78 else if((*tp)->n[0]->bal == 1)
79 doubleleftright(tp, p);
80 else
81 assert(0);
82 break;
83
84 case 2:
85 if((*tp)->n[1]->bal >= 0)
86 singleleft(tp, p);
87 else if((*tp)->n[1]->bal == -1)
88 doublerightleft(tp, p);
89 else
90 assert(0);
91 break;
92 }
93 }
94
95 static int
canoncmp(int cmp)96 canoncmp(int cmp)
97 {
98 if(cmp < 0)
99 return -1;
100 else if(cmp > 0)
101 return 1;
102 return 0;
103 }
104
105 static int
_insertavl(Avl ** tp,Avl * p,Avl * r,int (* cmp)(Avl *,Avl *),Avl ** rfree)106 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree)
107 {
108 int i, ob;
109
110 if(*tp == nil){
111 r->bal = 0;
112 r->n[0] = nil;
113 r->n[1] = nil;
114 r->p = p;
115 *tp = r;
116 return 1;
117 }
118 ob = (*tp)->bal;
119 if((i=canoncmp(cmp(r, *tp))) != 0){
120 (*tp)->bal += i*_insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp, rfree);
121 balance(tp, p);
122 return ob==0 && (*tp)->bal != 0;
123 }
124
125 /* install new entry */
126 *rfree = *tp; /* save old node for freeing */
127 *tp = r; /* insert new node */
128 **tp = **rfree; /* copy old node's Avl contents */
129 if(r->n[0]) /* fix node's children's parent pointers */
130 r->n[0]->p = r;
131 if(r->n[1])
132 r->n[1]->p = r;
133
134 return 0;
135 }
136
137 static Avl*
_lookupavl(Avl * t,Avl * r,int (* cmp)(Avl *,Avl *))138 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*))
139 {
140 int i;
141 Avl *p;
142
143 p = nil;
144 while(t != nil){
145 assert(t->p == p);
146 if((i=canoncmp(cmp(r, t)))==0)
147 return t;
148 p = t;
149 t = t->n[(i+1)/2];
150 }
151 return nil;
152 }
153
154 static int
successor(Avl ** tp,Avl * p,Avl ** r)155 successor(Avl **tp, Avl *p, Avl **r)
156 {
157 int ob;
158
159 if((*tp)->n[0] == nil){
160 *r = *tp;
161 *tp = (*r)->n[1];
162 if(*tp)
163 (*tp)->p = p;
164 return -1;
165 }
166 ob = (*tp)->bal;
167 (*tp)->bal -= successor(&(*tp)->n[0], *tp, r);
168 balance(tp, p);
169 return -(ob!=0 && (*tp)->bal==0);
170 }
171
172 static int
_deleteavl(Avl ** tp,Avl * p,Avl * rx,int (* cmp)(Avl *,Avl *),Avl ** del,void (* predel)(Avl *,void *),void * arg)173 _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del, void (*predel)(Avl*, void*), void *arg)
174 {
175 int i, ob;
176 Avl *r, *or;
177
178 if(*tp == nil)
179 return 0;
180
181 ob = (*tp)->bal;
182 if((i=canoncmp(cmp(rx, *tp))) != 0){
183 (*tp)->bal += i*_deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp, del, predel, arg);
184 balance(tp, p);
185 return -(ob!=0 && (*tp)->bal==0);
186 }
187
188 if(predel)
189 (*predel)(*tp, arg);
190
191 or = *tp;
192 if(or->n[i=0]==nil || or->n[i=1]==nil){
193 *tp = or->n[1-i];
194 if(*tp)
195 (*tp)->p = p;
196 *del = or;
197 return -1;
198 }
199
200 /* deleting node with two kids, find successor */
201 or->bal += successor(&or->n[1], or, &r);
202 r->bal = or->bal;
203 r->n[0] = or->n[0];
204 r->n[1] = or->n[1];
205 *tp = r;
206 (*tp)->p = p;
207 /* node has changed; fix children's parent pointers */
208 if(r->n[0])
209 r->n[0]->p = r;
210 if(r->n[1])
211 r->n[1]->p = r;
212 *del = or;
213 balance(tp, p);
214 return -(ob!=0 && (*tp)->bal==0);
215 }
216
217 static void
checkparents(Avl * a,Avl * p)218 checkparents(Avl *a, Avl *p)
219 {
220 if(a==nil)
221 return;
222 if(a->p != p)
223 print("bad parent\n");
224 checkparents(a->n[0], a);
225 checkparents(a->n[1], a);
226 }
227
228 struct Avltree
229 {
230 Avl *root;
231 int (*cmp)(Avl*, Avl*);
232 Avlwalk *walks;
233 };
234 struct Avlwalk
235 {
236 int started;
237 int moved;
238 Avlwalk *next;
239 Avltree *tree;
240 Avl *node;
241 };
242
243
244 Avltree*
mkavltree(int (* cmp)(Avl *,Avl *))245 mkavltree(int (*cmp)(Avl*, Avl*))
246 {
247 Avltree *t;
248
249 t = emalloc(sizeof(*t));
250 t->cmp = cmp;
251 return t;
252 }
253
254 void
insertavl(Avltree * t,Avl * new,Avl ** oldp)255 insertavl(Avltree *t, Avl *new, Avl **oldp)
256 {
257 *oldp = nil;
258 _insertavl(&t->root, nil, new, t->cmp, oldp);
259 }
260
261 Avl*
lookupavl(Avltree * t,Avl * key)262 lookupavl(Avltree *t, Avl *key)
263 {
264 return _lookupavl(t->root, key, t->cmp);
265 }
266
267 static Avl*
findpredecessor(Avl * a)268 findpredecessor(Avl *a)
269 {
270
271 if(a == nil)
272 return nil;
273
274 if(a->n[0] != nil){
275 /* predecessor is rightmost descendant of left child */
276 for(a=a->n[0]; a->n[1]; a=a->n[1])
277 ;
278 return a;
279 }else{
280 /* we're at a leaf, successor is a parent we enter from the right */
281 while(a->p && a->p->n[0]==a)
282 a = a->p;
283 return a->p;
284 }
285 }
286
287 static Avl*
findsuccessor(Avl * a)288 findsuccessor(Avl *a)
289 {
290
291 if(a == nil)
292 return nil;
293
294 if(a->n[1] != nil){
295 /* successor is leftmost descendant of right child */
296 for(a=a->n[1]; a->n[0]; a=a->n[0])
297 ;
298 return a;
299 }else{
300 /* we're at a leaf, successor is a parent we enter from the left going up */
301 while(a->p && a->p->n[1] == a)
302 a = a->p;
303 return a->p;
304 }
305 }
306
307 static void
walkdel(Avl * a,void * v)308 walkdel(Avl *a, void *v)
309 {
310 Avl *p;
311 Avlwalk *w;
312 Avltree *t;
313
314 if(a == nil)
315 return;
316
317 p = findpredecessor(a);
318 t = v;
319 for(w=t->walks; w; w=w->next){
320 if(w->node == a){
321 /* back pointer to predecessor; not perfect but adequate */
322 w->moved = 1;
323 w->node = p;
324 if(p == nil)
325 w->started = 0;
326 }
327 }
328 }
329
330 void
deleteavl(Avltree * t,Avl * key,Avl ** oldp)331 deleteavl(Avltree *t, Avl *key, Avl **oldp)
332 {
333 *oldp = nil;
334 _deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
335 }
336
337 Avlwalk*
avlwalk(Avltree * t)338 avlwalk(Avltree *t)
339 {
340 Avlwalk *w;
341
342 w = emalloc(sizeof(*w));
343 w->tree = t;
344 w->next = t->walks;
345 t->walks = w;
346 return w;
347 }
348
349 Avl*
avlnext(Avlwalk * w)350 avlnext(Avlwalk *w)
351 {
352 Avl *a;
353
354 if(w->started==0){
355 for(a=w->tree->root; a && a->n[0]; a=a->n[0])
356 ;
357 w->node = a;
358 w->started = 1;
359 }else{
360 a = findsuccessor(w->node);
361 if(a == w->node)
362 abort();
363 w->node = a;
364 }
365 return w->node;
366 }
367
368 Avl*
avlprev(Avlwalk * w)369 avlprev(Avlwalk *w)
370 {
371 Avl *a;
372
373 if(w->started == 0){
374 for(a=w->tree->root; a && a->n[1]; a=a->n[1])
375 ;
376 w->node = a;
377 w->started = 1;
378 }else if(w->moved){
379 w->moved = 0;
380 return w->node;
381 }else{
382 a = findpredecessor(w->node);
383 if(a == w->node)
384 abort();
385 w->node = a;
386 }
387 return w->node;
388 }
389
390 void
endwalk(Avlwalk * w)391 endwalk(Avlwalk *w)
392 {
393 Avltree *t;
394 Avlwalk **l;
395
396 t = w->tree;
397 for(l=&t->walks; *l; l=&(*l)->next){
398 if(*l == w){
399 *l = w->next;
400 break;
401 }
402 }
403 free(w);
404 }
405
406 static void
walkavl(Avl * t,void (* f)(Avl *,void *),void * v)407 walkavl(Avl *t, void (*f)(Avl*, void*), void *v)
408 {
409 if(t == nil)
410 return;
411 walkavl(t->n[0], f, v);
412 f(t, v);
413 walkavl(t->n[1], f, v);
414 }
415
416