rbtree
view src/rbtree.c @ 3:53afe96233f2
changed the tail-recursive find and find_min to loops
| author | John Tsiombikas <nuclear@mutantstargoat.com> |
|---|---|
| date | Sun, 09 Oct 2011 08:30:54 +0300 |
| parents | 6621337b6378 |
| children | 56a08d00bb41 |
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 }
93 int rb_size(struct rbtree *rb)
94 {
95 return count_nodes(rb->root);
96 }
98 int rb_insert(struct rbtree *rb, void *key, void *data)
99 {
100 rb->root = insert(rb, rb->root, key, data);
101 rb->root->red = 0;
102 return 0;
103 }
105 int rb_inserti(struct rbtree *rb, int key, void *data)
106 {
107 rb->root = insert(rb, rb->root, INT2PTR(key), data);
108 rb->root->red = 0;
109 return 0;
110 }
113 int rb_delete(struct rbtree *rb, void *key)
114 {
115 rb->root = delete(rb, rb->root, key);
116 rb->root->red = 0;
117 return 0;
118 }
120 int rb_deletei(struct rbtree *rb, int key)
121 {
122 rb->root = delete(rb, rb->root, INT2PTR(key));
123 rb->root->red = 0;
124 return 0;
125 }
128 void *rb_find(struct rbtree *rb, void *key)
129 {
130 struct rbnode *node = rb->root;
132 while(node) {
133 int cmp = rb->cmp(key, node->key);
134 if(cmp == 0) {
135 return node;
136 }
137 node = cmp < 0 ? node->left : node->right;
138 }
139 return 0;
140 }
142 void *rb_findi(struct rbtree *rb, int key)
143 {
144 return rb_find(rb, INT2PTR(key));
145 }
148 void rb_foreach(struct rbtree *rb, void (*func)(struct rbnode*, void*), void *cls)
149 {
150 traverse(rb->root, func, cls);
151 }
154 struct rbnode *rb_root(struct rbtree *rb)
155 {
156 return rb->root;
157 }
159 void rb_begin(struct rbtree *rb)
160 {
161 rb->rstack = 0;
162 rb->iter = rb->root;
163 }
165 #define push(sp, x) ((x)->next = (sp), (sp) = (x))
166 #define pop(sp) ((sp) = (sp)->next)
167 #define top(sp) (sp)
169 struct rbnode *rb_next(struct rbtree *rb)
170 {
171 struct rbnode *res = 0;
173 while(rb->rstack || rb->iter) {
174 if(rb->iter) {
175 push(rb->rstack, rb->iter);
176 rb->iter = rb->iter->left;
177 } else {
178 rb->iter = top(rb->rstack);
179 pop(rb->rstack);
180 res = rb->iter;
181 rb->iter = rb->iter->right;
182 break;
183 }
184 }
185 return res;
186 }
188 void *rb_node_key(struct rbnode *node)
189 {
190 return node ? node->key : 0;
191 }
193 int rb_node_keyi(struct rbnode *node)
194 {
195 return node ? PTR2INT(node->key) : 0;
196 }
198 void *rb_node_data(struct rbnode *node)
199 {
200 return node ? node->data : 0;
201 }
203 static int cmpaddr(void *ap, void *bp)
204 {
205 return ap < bp ? -1 : (ap > bp ? 1 : 0);
206 }
208 static int cmpint(void *ap, void *bp)
209 {
210 return PTR2INT(ap) - PTR2INT(bp);
211 }
214 /* ---- left-leaning 2-3 red-black implementation ---- */
216 /* helper prototypes */
217 static int is_red(struct rbnode *tree);
218 static void color_flip(struct rbnode *tree);
219 static struct rbnode *rot_left(struct rbnode *a);
220 static struct rbnode *rot_right(struct rbnode *a);
221 static struct rbnode *find_min(struct rbnode *tree);
222 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree);
223 /*static struct rbnode *move_red_right(struct rbnode *tree);*/
224 static struct rbnode *move_red_left(struct rbnode *tree);
225 static struct rbnode *fix_up(struct rbnode *tree);
227 static int count_nodes(struct rbnode *node)
228 {
229 if(!node)
230 return 0;
232 return 1 + count_nodes(node->left) + count_nodes(node->right);
233 }
235 static void del_tree(struct rbnode *node, rb_del_func_t delfunc, void *cls)
236 {
237 if(!node)
238 return;
240 del_tree(node->left, delfunc, cls);
241 del_tree(node->right, delfunc, cls);
243 delfunc(node, cls);
244 free(node);
245 }
247 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data)
248 {
249 int cmp;
251 if(!tree) {
252 struct rbnode *node = rb->alloc(sizeof *node);
253 node->red = 1;
254 node->key = key;
255 node->data = data;
256 node->left = node->right = 0;
257 return node;
258 }
260 cmp = rb->cmp(key, tree->key);
262 if(cmp < 0) {
263 tree->left = insert(rb, tree->left, key, data);
264 } else if(cmp > 0) {
265 tree->right = insert(rb, tree->right, key, data);
266 } else {
267 tree->data = data;
268 }
270 /* fix right-leaning reds */
271 if(is_red(tree->right)) {
272 tree = rot_left(tree);
273 }
274 /* fix two reds in a row */
275 if(is_red(tree->left) && is_red(tree->left->left)) {
276 tree = rot_right(tree);
277 }
279 /* if 4-node, split it by color inversion */
280 if(is_red(tree->left) && is_red(tree->right)) {
281 color_flip(tree);
282 }
284 return tree;
285 }
287 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key)
288 {
289 int cmp;
291 if(!tree) {
292 return 0;
293 }
295 cmp = rb->cmp(key, tree->key);
297 if(cmp < 0) {
298 if(!is_red(tree->left) && !is_red(tree->left->left)) {
299 tree = move_red_left(tree);
300 }
301 tree->left = delete(rb, tree->left, key);
302 } else {
303 /* need reds on the right */
304 if(is_red(tree->left)) {
305 tree = rot_right(tree);
306 }
308 /* found it at the bottom (XXX what certifies left is null?) */
309 if(cmp == 0 && !tree->right) {
310 if(rb->del) {
311 rb->del(tree, rb->del_cls);
312 }
313 rb->free(tree);
314 return 0;
315 }
317 if(!is_red(tree->right) && !is_red(tree->right->left)) {
318 tree = move_red_left(tree);
319 }
321 if(key == tree->key) {
322 struct rbnode *rmin = find_min(tree->right);
323 tree->key = rmin->key;
324 tree->data = rmin->data;
325 tree->right = del_min(rb, tree->right);
326 } else {
327 tree->right = delete(rb, tree->right, key);
328 }
329 }
331 return fix_up(tree);
332 }
334 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key)
335 {
336 int cmp;
338 if(!node)
339 return 0;
341 if((cmp = rb->cmp(key, node->key)) == 0) {
342 return node;
343 }
344 return find(rb, cmp < 0 ? node->left : node->right, key);
345 }*/
347 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls)
348 {
349 if(!node)
350 return;
352 traverse(node->left, func, cls);
353 func(node, cls);
354 traverse(node->right, func, cls);
355 }
357 /* helpers */
359 static int is_red(struct rbnode *tree)
360 {
361 return tree && tree->red;
362 }
364 static void color_flip(struct rbnode *tree)
365 {
366 tree->red = !tree->red;
367 tree->left->red = !tree->left->red;
368 tree->right->red = !tree->right->red;
369 }
371 static struct rbnode *rot_left(struct rbnode *a)
372 {
373 struct rbnode *b = a->right;
374 a->right = b->left;
375 b->left = a;
376 b->red = a->red;
377 a->red = 1;
378 return b;
379 }
381 static struct rbnode *rot_right(struct rbnode *a)
382 {
383 struct rbnode *b = a->left;
384 a->left = b->right;
385 b->right = a;
386 b->red = a->red;
387 a->red = 1;
388 return b;
389 }
391 static struct rbnode *find_min(struct rbnode *tree)
392 {
393 struct rbnode *node;
395 if(!tree)
396 return 0;
398 while(node->left) {
399 node = node->left;
400 }
401 return node;
402 }
404 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree)
405 {
406 if(!tree->left) {
407 if(rb->del) {
408 rb->del(tree->left, rb->del_cls);
409 }
410 rb->free(tree->left);
411 return 0;
412 }
414 /* make sure we've got red (3/4-nodes) at the left side so we can delete at the bottom */
415 if(!is_red(tree->left) && !is_red(tree->left->left)) {
416 tree = move_red_left(tree);
417 }
418 tree->left = del_min(rb, tree->left);
420 /* fix right-reds, red-reds, and split 4-nodes on the way up */
421 return fix_up(tree);
422 }
424 #if 0
425 /* push a red link on this node to the right */
426 static struct rbnode *move_red_right(struct rbnode *tree)
427 {
428 /* flipping it makes both children go red, so we have a red to the right */
429 color_flip(tree);
431 /* if after the flip we've got a red-red situation to the left, fix it */
432 if(is_red(tree->left->left)) {
433 tree = rot_right(tree);
434 color_flip(tree);
435 }
436 return tree;
437 }
438 #endif
440 /* push a red link on this node to the left */
441 static struct rbnode *move_red_left(struct rbnode *tree)
442 {
443 /* flipping it makes both children go red, so we have a red to the left */
444 color_flip(tree);
446 /* if after the flip we've got a red-red on the right-left, fix it */
447 if(is_red(tree->right->left)) {
448 tree->right = rot_right(tree->right);
449 tree = rot_left(tree);
450 color_flip(tree);
451 }
452 return tree;
453 }
455 static struct rbnode *fix_up(struct rbnode *tree)
456 {
457 /* fix right-leaning */
458 if(is_red(tree->right)) {
459 tree = rot_left(tree);
460 }
461 /* change invalid red-red pairs into a proper 4-node */
462 if(is_red(tree->left) && is_red(tree->left->left)) {
463 tree = rot_right(tree);
464 }
465 /* split 4-nodes */
466 if(is_red(tree->left) && is_red(tree->right)) {
467 color_flip(tree);
468 }
469 return tree;
470 }