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