rbtree
view src/rbtree.c @ 7:2c0baa1fd419
del_tree would call 0 if the del func is null
| author | John Tsiombikas <nuclear@member.fsf.org> |
|---|---|
| date | Fri, 21 Oct 2011 11:32:38 +0300 |
| parents | 56a08d00bb41 |
| children | b45febfd2922 |
line source
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <string.h>
4 #include "rbtree.h"
6 #define INT2PTR(x) ((void*)(x))
7 #define PTR2INT(x) ((int)(x))
9 struct rbtree {
10 struct rbnode *root;
12 rb_alloc_func_t alloc;
13 rb_free_func_t free;
15 rb_cmp_func_t cmp;
16 rb_del_func_t del;
17 void *del_cls;
19 struct rbnode *rstack, *iter;
20 };
22 static int cmpaddr(void *ap, void *bp);
23 static int cmpint(void *ap, void *bp);
25 static int count_nodes(struct rbnode *node);
26 static void del_tree(struct rbnode *node, void (*delfunc)(struct rbnode*, void*), void *cls);
27 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data);
28 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key);
29 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key);*/
30 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls);
32 struct rbtree *rb_create(rb_cmp_func_t cmp_func)
33 {
34 struct rbtree *rb;
36 if(!(rb = malloc(sizeof *rb))) {
37 return 0;
38 }
39 if(rb_init(rb, cmp_func) == -1) {
40 free(rb);
41 return 0;
42 }
43 return rb;
44 }
46 void rb_free(struct rbtree *rb)
47 {
48 rb_destroy(rb);
49 free(rb);
50 }
53 int rb_init(struct rbtree *rb, rb_cmp_func_t cmp_func)
54 {
55 memset(rb, 0, sizeof *rb);
57 if(cmp_func == RB_KEY_INT) {
58 rb->cmp = cmpint;
59 } else if(cmp_func == RB_KEY_STRING) {
60 rb->cmp = (rb_cmp_func_t)strcmp;
61 } else {
62 rb->cmp = cmpaddr;
63 }
65 rb->alloc = malloc;
66 rb->free = free;
67 return 0;
68 }
70 void rb_destroy(struct rbtree *rb)
71 {
72 del_tree(rb->root, rb->del, rb->del_cls);
73 }
75 void rb_set_allocator(struct rbtree *rb, rb_alloc_func_t alloc, rb_free_func_t free)
76 {
77 rb->alloc = alloc;
78 rb->free = free;
79 }
82 void rb_set_compare_func(struct rbtree *rb, rb_cmp_func_t func)
83 {
84 rb->cmp = func;
85 }
87 void rb_set_delete_func(struct rbtree *rb, rb_del_func_t func, void *cls)
88 {
89 rb->del = func;
90 rb->del_cls = cls;
91 }
94 void rb_clear(struct rbtree *rb)
95 {
96 del_tree(rb->root, rb->del, rb->del_cls);
97 rb->root = 0;
98 }
100 int rb_copy(struct rbtree *dest, struct rbtree *src)
101 {
102 struct rbnode *node;
104 rb_clear(dest);
105 rb_begin(src);
106 while((node = rb_next(src))) {
107 if(rb_insert(dest, node->key, node->data) == -1) {
108 return -1;
109 }
110 }
111 return 0;
112 }
114 int rb_size(struct rbtree *rb)
115 {
116 return count_nodes(rb->root);
117 }
119 int rb_insert(struct rbtree *rb, void *key, void *data)
120 {
121 rb->root = insert(rb, rb->root, key, data);
122 rb->root->red = 0;
123 return 0;
124 }
126 int rb_inserti(struct rbtree *rb, int key, void *data)
127 {
128 rb->root = insert(rb, rb->root, INT2PTR(key), data);
129 rb->root->red = 0;
130 return 0;
131 }
134 int rb_delete(struct rbtree *rb, void *key)
135 {
136 rb->root = delete(rb, rb->root, key);
137 rb->root->red = 0;
138 return 0;
139 }
141 int rb_deletei(struct rbtree *rb, int key)
142 {
143 rb->root = delete(rb, rb->root, INT2PTR(key));
144 rb->root->red = 0;
145 return 0;
146 }
149 void *rb_find(struct rbtree *rb, void *key)
150 {
151 struct rbnode *node = rb->root;
153 while(node) {
154 int cmp = rb->cmp(key, node->key);
155 if(cmp == 0) {
156 return node;
157 }
158 node = cmp < 0 ? node->left : node->right;
159 }
160 return 0;
161 }
163 void *rb_findi(struct rbtree *rb, int key)
164 {
165 return rb_find(rb, INT2PTR(key));
166 }
169 void rb_foreach(struct rbtree *rb, void (*func)(struct rbnode*, void*), void *cls)
170 {
171 traverse(rb->root, func, cls);
172 }
175 struct rbnode *rb_root(struct rbtree *rb)
176 {
177 return rb->root;
178 }
180 void rb_begin(struct rbtree *rb)
181 {
182 rb->rstack = 0;
183 rb->iter = rb->root;
184 }
186 #define push(sp, x) ((x)->next = (sp), (sp) = (x))
187 #define pop(sp) ((sp) = (sp)->next)
188 #define top(sp) (sp)
190 struct rbnode *rb_next(struct rbtree *rb)
191 {
192 struct rbnode *res = 0;
194 while(rb->rstack || rb->iter) {
195 if(rb->iter) {
196 push(rb->rstack, rb->iter);
197 rb->iter = rb->iter->left;
198 } else {
199 rb->iter = top(rb->rstack);
200 pop(rb->rstack);
201 res = rb->iter;
202 rb->iter = rb->iter->right;
203 break;
204 }
205 }
206 return res;
207 }
209 void *rb_node_key(struct rbnode *node)
210 {
211 return node ? node->key : 0;
212 }
214 int rb_node_keyi(struct rbnode *node)
215 {
216 return node ? PTR2INT(node->key) : 0;
217 }
219 void *rb_node_data(struct rbnode *node)
220 {
221 return node ? node->data : 0;
222 }
224 static int cmpaddr(void *ap, void *bp)
225 {
226 return ap < bp ? -1 : (ap > bp ? 1 : 0);
227 }
229 static int cmpint(void *ap, void *bp)
230 {
231 return PTR2INT(ap) - PTR2INT(bp);
232 }
235 /* ---- left-leaning 2-3 red-black implementation ---- */
237 /* helper prototypes */
238 static int is_red(struct rbnode *tree);
239 static void color_flip(struct rbnode *tree);
240 static struct rbnode *rot_left(struct rbnode *a);
241 static struct rbnode *rot_right(struct rbnode *a);
242 static struct rbnode *find_min(struct rbnode *tree);
243 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree);
244 /*static struct rbnode *move_red_right(struct rbnode *tree);*/
245 static struct rbnode *move_red_left(struct rbnode *tree);
246 static struct rbnode *fix_up(struct rbnode *tree);
248 static int count_nodes(struct rbnode *node)
249 {
250 if(!node)
251 return 0;
253 return 1 + count_nodes(node->left) + count_nodes(node->right);
254 }
256 static void del_tree(struct rbnode *node, rb_del_func_t delfunc, void *cls)
257 {
258 if(!node)
259 return;
261 del_tree(node->left, delfunc, cls);
262 del_tree(node->right, delfunc, cls);
264 if(delfunc) {
265 delfunc(node, cls);
266 }
267 free(node);
268 }
270 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data)
271 {
272 int cmp;
274 if(!tree) {
275 struct rbnode *node = rb->alloc(sizeof *node);
276 node->red = 1;
277 node->key = key;
278 node->data = data;
279 node->left = node->right = 0;
280 return node;
281 }
283 cmp = rb->cmp(key, tree->key);
285 if(cmp < 0) {
286 tree->left = insert(rb, tree->left, key, data);
287 } else if(cmp > 0) {
288 tree->right = insert(rb, tree->right, key, data);
289 } else {
290 tree->data = data;
291 }
293 /* fix right-leaning reds */
294 if(is_red(tree->right)) {
295 tree = rot_left(tree);
296 }
297 /* fix two reds in a row */
298 if(is_red(tree->left) && is_red(tree->left->left)) {
299 tree = rot_right(tree);
300 }
302 /* if 4-node, split it by color inversion */
303 if(is_red(tree->left) && is_red(tree->right)) {
304 color_flip(tree);
305 }
307 return tree;
308 }
310 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key)
311 {
312 int cmp;
314 if(!tree) {
315 return 0;
316 }
318 cmp = rb->cmp(key, tree->key);
320 if(cmp < 0) {
321 if(!is_red(tree->left) && !is_red(tree->left->left)) {
322 tree = move_red_left(tree);
323 }
324 tree->left = delete(rb, tree->left, key);
325 } else {
326 /* need reds on the right */
327 if(is_red(tree->left)) {
328 tree = rot_right(tree);
329 }
331 /* found it at the bottom (XXX what certifies left is null?) */
332 if(cmp == 0 && !tree->right) {
333 if(rb->del) {
334 rb->del(tree, rb->del_cls);
335 }
336 rb->free(tree);
337 return 0;
338 }
340 if(!is_red(tree->right) && !is_red(tree->right->left)) {
341 tree = move_red_left(tree);
342 }
344 if(key == tree->key) {
345 struct rbnode *rmin = find_min(tree->right);
346 tree->key = rmin->key;
347 tree->data = rmin->data;
348 tree->right = del_min(rb, tree->right);
349 } else {
350 tree->right = delete(rb, tree->right, key);
351 }
352 }
354 return fix_up(tree);
355 }
357 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key)
358 {
359 int cmp;
361 if(!node)
362 return 0;
364 if((cmp = rb->cmp(key, node->key)) == 0) {
365 return node;
366 }
367 return find(rb, cmp < 0 ? node->left : node->right, key);
368 }*/
370 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls)
371 {
372 if(!node)
373 return;
375 traverse(node->left, func, cls);
376 func(node, cls);
377 traverse(node->right, func, cls);
378 }
380 /* helpers */
382 static int is_red(struct rbnode *tree)
383 {
384 return tree && tree->red;
385 }
387 static void color_flip(struct rbnode *tree)
388 {
389 tree->red = !tree->red;
390 tree->left->red = !tree->left->red;
391 tree->right->red = !tree->right->red;
392 }
394 static struct rbnode *rot_left(struct rbnode *a)
395 {
396 struct rbnode *b = a->right;
397 a->right = b->left;
398 b->left = a;
399 b->red = a->red;
400 a->red = 1;
401 return b;
402 }
404 static struct rbnode *rot_right(struct rbnode *a)
405 {
406 struct rbnode *b = a->left;
407 a->left = b->right;
408 b->right = a;
409 b->red = a->red;
410 a->red = 1;
411 return b;
412 }
414 static struct rbnode *find_min(struct rbnode *tree)
415 {
416 struct rbnode *node;
418 if(!tree)
419 return 0;
421 while(node->left) {
422 node = node->left;
423 }
424 return node;
425 }
427 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree)
428 {
429 if(!tree->left) {
430 if(rb->del) {
431 rb->del(tree->left, rb->del_cls);
432 }
433 rb->free(tree->left);
434 return 0;
435 }
437 /* make sure we've got red (3/4-nodes) at the left side so we can delete at the bottom */
438 if(!is_red(tree->left) && !is_red(tree->left->left)) {
439 tree = move_red_left(tree);
440 }
441 tree->left = del_min(rb, tree->left);
443 /* fix right-reds, red-reds, and split 4-nodes on the way up */
444 return fix_up(tree);
445 }
447 #if 0
448 /* push a red link on this node to the right */
449 static struct rbnode *move_red_right(struct rbnode *tree)
450 {
451 /* flipping it makes both children go red, so we have a red to the right */
452 color_flip(tree);
454 /* if after the flip we've got a red-red situation to the left, fix it */
455 if(is_red(tree->left->left)) {
456 tree = rot_right(tree);
457 color_flip(tree);
458 }
459 return tree;
460 }
461 #endif
463 /* push a red link on this node to the left */
464 static struct rbnode *move_red_left(struct rbnode *tree)
465 {
466 /* flipping it makes both children go red, so we have a red to the left */
467 color_flip(tree);
469 /* if after the flip we've got a red-red on the right-left, fix it */
470 if(is_red(tree->right->left)) {
471 tree->right = rot_right(tree->right);
472 tree = rot_left(tree);
473 color_flip(tree);
474 }
475 return tree;
476 }
478 static struct rbnode *fix_up(struct rbnode *tree)
479 {
480 /* fix right-leaning */
481 if(is_red(tree->right)) {
482 tree = rot_left(tree);
483 }
484 /* change invalid red-red pairs into a proper 4-node */
485 if(is_red(tree->left) && is_red(tree->left->left)) {
486 tree = rot_right(tree);
487 }
488 /* split 4-nodes */
489 if(is_red(tree->left) && is_red(tree->right)) {
490 color_flip(tree);
491 }
492 return tree;
493 }