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 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
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
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
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 /*
208 static void
209 checkparents(Avl *a, Avl *p)
210 {
211 	if(a == nil)
212 		return;
213 	if(a->p != p)
214 		print("bad parent\n");
215 	checkparents(a->n[0], a);
216 	checkparents(a->n[1], a);
217 }
218 */
219 
220 struct Avltree
221 {
222 	Avl	*root;
223 	int	(*cmp)(Avl*, Avl*);
224 	Avlwalk	*walks;
225 };
226 struct Avlwalk
227 {
228 	int	started;
229 	int	moved;
230 	Avlwalk	*next;
231 	Avltree	*tree;
232 	Avl	*node;
233 };
234 
235 Avltree*
236 mkavltree(int (*cmp)(Avl*, Avl*))
237 {
238 	Avltree *t;
239 
240 	t = malloc(sizeof *t);
241 	if(t == nil)
242 		return nil;
243 	memset(t, 0, sizeof *t);
244 	t->cmp = cmp;
245 	return t;
246 }
247 
248 void
249 insertavl(Avltree *t, Avl *new, Avl **oldp)
250 {
251 	*oldp = nil;
252 	_insertavl(&t->root, nil, new, t->cmp, oldp);
253 }
254 
255 static Avl*
256 findpredecessor(Avl *a)
257 {
258 	if(a == nil)
259 		return nil;
260 
261 	if(a->n[0] != nil){
262 		/* predecessor is rightmost descendant of left child */
263 		for(a = a->n[0]; a->n[1]; a = a->n[1])
264 			;
265 		return a;
266 	}else{
267 		/* we're at a leaf, successor is a parent we enter from the right */
268 		while(a->p && a->p->n[0] == a)
269 			a = a->p;
270 		return a->p;
271 	}
272 }
273 
274 static Avl*
275 findsuccessor(Avl *a)
276 {
277 	if(a == nil)
278 		return nil;
279 
280 	if(a->n[1] != nil){
281 		/* successor is leftmost descendant of right child */
282 		for(a = a->n[1]; a->n[0]; a = a->n[0])
283 			;
284 		return a;
285 	}else{
286 		/* we're at a leaf, successor is a parent we enter from the left going up */
287 		while(a->p && a->p->n[1] == a)
288 			a = a->p;
289 		return a->p;
290 	}
291 }
292 
293 static Avl*
294 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*), int neighbor)
295 {
296 	int i;
297 	Avl *p;
298 
299 	p = nil;
300 	if(t == nil)
301 		return nil;
302 	do{
303 		assert(t->p == p);
304 		if((i = canoncmp(cmp(r, t))) == 0)
305 			return t;
306 		p = t;
307 		t = t->n[(i+1)/2];
308 	}while(t);
309 	if(neighbor == 0)
310 		return nil;
311 	if(neighbor < 0)
312 		return i > 0 ? p : findpredecessor(p);
313 	return i < 0 ? p : findsuccessor(p);
314 }
315 
316 Avl*
317 searchavl(Avltree *t, Avl *key, int neighbor)
318 {
319 	return _lookupavl(t->root, key, t->cmp, neighbor);
320 }
321 
322 Avl*
323 lookupavl(Avltree *t, Avl *key)
324 {
325 	return _lookupavl(t->root, key, t->cmp, 0);
326 }
327 
328 static void
329 walkdel(Avl *a, void *v)
330 {
331 	Avl *p;
332 	Avlwalk *w;
333 	Avltree *t;
334 
335 	if(a == nil)
336 		return;
337 
338 	p = findpredecessor(a);
339 	t = v;
340 	for(w = t->walks; w; w = w->next){
341 		if(w->node == a){
342 			/* back pointer to predecessor; not perfect but adequate */
343 			w->moved = 1;
344 			w->node = p;
345 			if(p == nil)
346 				w->started = 0;
347 		}
348 	}
349 }
350 
351 void
352 deleteavl(Avltree *t, Avl *key, Avl **oldp)
353 {
354 	*oldp = nil;
355 	_deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
356 }
357 
358 Avlwalk*
359 avlwalk(Avltree *t)
360 {
361 	Avlwalk *w;
362 
363 	w = malloc(sizeof *w);
364 	if(w == nil)
365 		return nil;
366 	memset(w, 0, sizeof *w);
367 	w->tree = t;
368 	w->next = t->walks;
369 	t->walks = w;
370 	return w;
371 }
372 
373 Avl*
374 avlnext(Avlwalk *w)
375 {
376 	Avl *a;
377 
378 	if(w->started==0){
379 		for(a = w->tree->root; a && a->n[0]; a = a->n[0])
380 			;
381 		w->node = a;
382 		w->started = 1;
383 	}else{
384 		a = findsuccessor(w->node);
385 		if(a == w->node)
386 			abort();
387 		w->node = a;
388 	}
389 	return w->node;
390 }
391 
392 Avl*
393 avlprev(Avlwalk *w)
394 {
395 	Avl *a;
396 
397 	if(w->started == 0){
398 		for(a = w->tree->root; a && a->n[1]; a = a->n[1])
399 			;
400 		w->node = a;
401 		w->started = 1;
402 	}else if(w->moved){
403 		w->moved = 0;
404 		return w->node;
405 	}else{
406 		a = findpredecessor(w->node);
407 		if(a == w->node)
408 			abort();
409 		w->node = a;
410 	}
411 	return w->node;
412 }
413 
414 void
415 endwalk(Avlwalk *w)
416 {
417 	Avltree *t;
418 	Avlwalk **l;
419 
420 	t = w->tree;
421 	for(l = &t->walks; *l; l = &(*l)->next){
422 		if(*l == w){
423 			*l = w->next;
424 			break;
425 		}
426 	}
427 	free(w);
428 }
429 
430 /*
431 static void
432 walkavl(Avl *t, void (*f)(Avl*, void*), void *v)
433 {
434 	if(t == nil)
435 		return;
436 	walkavl(t->n[0], f, v);
437 	f(t, v);
438 	walkavl(t->n[1], f, v);
439 }
440 */