xref: /plan9-contrib/sys/src/cmd/replica/avl.c (revision ec59a3ddbfceee0efe34584c2c9981a5e5ff1ec4)
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
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
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
58 doublerightleft(Avl **tp, Avl *p)
59 {
60 	singleright(&(*tp)->n[1], *tp);
61 	singleleft(tp, p);
62 }
63 
64 static void
65 doubleleftright(Avl **tp, Avl *p)
66 {
67 	singleleft(&(*tp)->n[0], *tp);
68 	singleright(tp, p);
69 }
70 
71 static void
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
96 _insertavl(Avl **tp, Avl *p, Avl *r, int (*cmp)(Avl*,Avl*), Avl **rfree)
97 {
98 	int i, ob;
99 
100 	if(*tp == nil){
101 		r->bal = 0;
102 		r->n[0] = nil;
103 		r->n[1] = nil;
104 		r->p = p;
105 		*tp = r;
106 		return 1;
107 	}
108 	ob = (*tp)->bal;
109 	if((i=cmp(r, *tp)) != 0){
110 		(*tp)->bal += i*_insertavl(&(*tp)->n[(i+1)/2], *tp, r, cmp, rfree);
111 		balance(tp, p);
112 		return ob==0 && (*tp)->bal != 0;
113 	}
114 
115 	/* install new entry */
116 	*rfree = *tp;	/* save old node for freeing */
117 	*tp = r;		/* insert new node */
118 	**tp = **rfree;	/* copy old node's Avl contents */
119 	if(r->n[0])		/* fix node's children's parent pointers */
120 		r->n[0]->p = r;
121 	if(r->n[1])
122 		r->n[1]->p = r;
123 
124 	return 0;
125 }
126 
127 static Avl*
128 _lookupavl(Avl *t, Avl *r, int (*cmp)(Avl*,Avl*))
129 {
130 	int i;
131 	Avl *p;
132 
133 	p = nil;
134 	while(t != nil){
135 		assert(t->p == p);
136 		if((i=cmp(r, t))==0)
137 			return t;
138 		p = t;
139 		t = t->n[(i+1)/2];
140 	}
141 	return nil;
142 }
143 
144 static int
145 successor(Avl **tp, Avl *p, Avl **r)
146 {
147 	int ob;
148 
149 	if((*tp)->n[0] == nil){
150 		*r = *tp;
151 		*tp = (*r)->n[1];
152 		if(*tp)
153 			(*tp)->p = p;
154 		return -1;
155 	}
156 	ob = (*tp)->bal;
157 	(*tp)->bal -= successor(&(*tp)->n[0], *tp, r);
158 	balance(tp, p);
159 	return -(ob!=0 && (*tp)->bal==0);
160 }
161 
162 static int
163 _deleteavl(Avl **tp, Avl *p, Avl *rx, int(*cmp)(Avl*,Avl*), Avl **del, void (*predel)(Avl*, void*), void *arg)
164 {
165 	int i, ob;
166 	Avl *r, *or;
167 
168 	if(*tp == nil)
169 		return 0;
170 
171 	ob = (*tp)->bal;
172 	if((i=cmp(rx, *tp)) != 0){
173 		(*tp)->bal += i*_deleteavl(&(*tp)->n[(i+1)/2], *tp, rx, cmp, 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
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 
234 Avltree*
235 mkavltree(int (*cmp)(Avl*, Avl*))
236 {
237 	Avltree *t;
238 
239 	t = emalloc(sizeof(*t));
240 	t->cmp = cmp;
241 	return t;
242 }
243 
244 void
245 insertavl(Avltree *t, Avl *new, Avl **oldp)
246 {
247 	*oldp = nil;
248 	_insertavl(&t->root, nil, new, t->cmp, oldp);
249 }
250 
251 Avl*
252 lookupavl(Avltree *t, Avl *key)
253 {
254 	return _lookupavl(t->root, key, t->cmp);
255 }
256 
257 static Avl*
258 findpredecessor(Avl *a)
259 {
260 
261 	if(a == nil)
262 		return nil;
263 
264 	if(a->n[0] != nil){
265 		/* predecessor is rightmost descendant of left child */
266 		for(a=a->n[0]; a->n[1]; a=a->n[1])
267 			;
268 		return a;
269 	}else{
270 		/* we're at a leaf, successor is a parent we enter from the right */
271 		while(a->p && a->p->n[0]==a)
272 			a = a->p;
273 		return a->p;
274 	}
275 }
276 
277 static Avl*
278 findsuccessor(Avl *a)
279 {
280 
281 	if(a == nil)
282 		return nil;
283 
284 	if(a->n[1] != nil){
285 		/* successor is leftmost descendant of right child */
286 		for(a=a->n[1]; a->n[0]; a=a->n[0])
287 			;
288 		return a;
289 	}else{
290 		/* we're at a leaf, successor is a parent we enter from the left going up */
291 		while(a->p && a->p->n[1] == a)
292 			a = a->p;
293 		return a->p;
294 	}
295 }
296 
297 static void
298 walkdel(Avl *a, void *v)
299 {
300 	Avl *p;
301 	Avlwalk *w;
302 	Avltree *t;
303 
304 	if(a == nil)
305 		return;
306 
307 	p = findpredecessor(a);
308 	t = v;
309 	for(w=t->walks; w; w=w->next){
310 		if(w->node == a){
311 			/* back pointer to predecessor; not perfect but adequate */
312 			w->moved = 1;
313 			w->node = p;
314 			if(p == nil)
315 				w->started = 0;
316 		}
317 	}
318 }
319 
320 void
321 deleteavl(Avltree *t, Avl *key, Avl **oldp)
322 {
323 	*oldp = nil;
324 	_deleteavl(&t->root, nil, key, t->cmp, oldp, walkdel, t);
325 }
326 
327 Avlwalk*
328 avlwalk(Avltree *t)
329 {
330 	Avlwalk *w;
331 
332 	w = emalloc(sizeof(*w));
333 	w->tree = t;
334 	w->next = t->walks;
335 	t->walks = w;
336 	return w;
337 }
338 
339 Avl*
340 avlnext(Avlwalk *w)
341 {
342 	Avl *a;
343 
344 	if(w->started==0){
345 		for(a=w->tree->root; a && a->n[0]; a=a->n[0])
346 			;
347 		w->node = a;
348 		w->started = 1;
349 	}else{
350 		a = findsuccessor(w->node);
351 		if(a == w->node)
352 			abort();
353 		w->node = a;
354 	}
355 	return w->node;
356 }
357 
358 Avl*
359 avlprev(Avlwalk *w)
360 {
361 	Avl *a;
362 
363 	if(w->started == 0){
364 		for(a=w->tree->root; a && a->n[1]; a=a->n[1])
365 			;
366 		w->node = a;
367 		w->started = 1;
368 	}else if(w->moved){
369 		w->moved = 0;
370 		return w->node;
371 	}else{
372 		a = findpredecessor(w->node);
373 		if(a == w->node)
374 			abort();
375 		w->node = a;
376 	}
377 	return w->node;
378 }
379 
380 void
381 endwalk(Avlwalk *w)
382 {
383 	Avltree *t;
384 	Avlwalk **l;
385 
386 	t = w->tree;
387 	for(l=&t->walks; *l; l=&(*l)->next){
388 		if(*l == w){
389 			*l = w->next;
390 			break;
391 		}
392 	}
393 	free(w);
394 }
395 
396 static void
397 walkavl(Avl *t, void (*f)(Avl*, void*), void *v)
398 {
399 	if(t == nil)
400 		return;
401 	walkavl(t->n[0], f, v);
402 	f(t, v);
403 	walkavl(t->n[1], f, v);
404 }
405 
406