rayzor
view src/rbtree.c @ 20:6b11a3f8706e
foo
| author | John Tsiombikas <nuclear@member.fsf.org> |
|---|---|
| date | Tue, 15 Apr 2014 01:36:03 +0300 |
| parents | |
| children |
line source
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <string.h>
4 #include "inttypes.h"
5 #include "rbtree.h"
7 #define INT2PTR(x) ((void*)(intptr_t)(x))
8 #define PTR2INT(x) ((int)(intptr_t)(x))
10 struct rbtree {
11 struct rbnode *root;
13 rb_alloc_func_t alloc;
14 rb_free_func_t free;
16 rb_cmp_func_t cmp;
17 rb_del_func_t del;
18 void *del_cls;
20 struct rbnode *rstack, *iter;
21 };
23 static int cmpaddr(const void *ap, const void *bp);
24 static int cmpint(const void *ap, const void *bp);
26 static int count_nodes(struct rbnode *node);
27 static void del_tree(struct rbnode *node, void (*delfunc)(struct rbnode*, void*), void *cls);
28 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data);
29 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key);
30 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key);*/
31 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls);
33 struct rbtree *rb_create(rb_cmp_func_t cmp_func)
34 {
35 struct rbtree *rb;
37 if(!(rb = malloc(sizeof *rb))) {
38 return 0;
39 }
40 if(rb_init(rb, cmp_func) == -1) {
41 free(rb);
42 return 0;
43 }
44 return rb;
45 }
47 void rb_free(struct rbtree *rb)
48 {
49 rb_destroy(rb);
50 free(rb);
51 }
54 int rb_init(struct rbtree *rb, rb_cmp_func_t cmp_func)
55 {
56 memset(rb, 0, sizeof *rb);
58 if(cmp_func == RB_KEY_INT) {
59 rb->cmp = cmpint;
60 } else if(cmp_func == RB_KEY_STRING) {
61 rb->cmp = (rb_cmp_func_t)strcmp;
62 } else {
63 rb->cmp = cmpaddr;
64 }
66 rb->alloc = malloc;
67 rb->free = free;
68 return 0;
69 }
71 void rb_destroy(struct rbtree *rb)
72 {
73 del_tree(rb->root, rb->del, rb->del_cls);
74 }
76 void rb_set_allocator(struct rbtree *rb, rb_alloc_func_t alloc, rb_free_func_t free)
77 {
78 rb->alloc = alloc;
79 rb->free = free;
80 }
83 void rb_set_compare_func(struct rbtree *rb, rb_cmp_func_t func)
84 {
85 rb->cmp = func;
86 }
88 void rb_set_delete_func(struct rbtree *rb, rb_del_func_t func, void *cls)
89 {
90 rb->del = func;
91 rb->del_cls = cls;
92 }
95 void rb_clear(struct rbtree *rb)
96 {
97 del_tree(rb->root, rb->del, rb->del_cls);
98 rb->root = 0;
99 }
101 int rb_copy(struct rbtree *dest, struct rbtree *src)
102 {
103 struct rbnode *node;
105 rb_clear(dest);
106 rb_begin(src);
107 while((node = rb_next(src))) {
108 if(rb_insert(dest, node->key, node->data) == -1) {
109 return -1;
110 }
111 }
112 return 0;
113 }
115 int rb_size(struct rbtree *rb)
116 {
117 return count_nodes(rb->root);
118 }
120 int rb_insert(struct rbtree *rb, void *key, void *data)
121 {
122 rb->root = insert(rb, rb->root, key, data);
123 rb->root->red = 0;
124 return 0;
125 }
127 int rb_inserti(struct rbtree *rb, int key, void *data)
128 {
129 rb->root = insert(rb, rb->root, INT2PTR(key), data);
130 rb->root->red = 0;
131 return 0;
132 }
135 int rb_delete(struct rbtree *rb, void *key)
136 {
137 rb->root = delete(rb, rb->root, key);
138 rb->root->red = 0;
139 return 0;
140 }
142 int rb_deletei(struct rbtree *rb, int key)
143 {
144 rb->root = delete(rb, rb->root, INT2PTR(key));
145 rb->root->red = 0;
146 return 0;
147 }
150 struct rbnode *rb_find(struct rbtree *rb, void *key)
151 {
152 struct rbnode *node = rb->root;
154 while(node) {
155 int cmp = rb->cmp(key, node->key);
156 if(cmp == 0) {
157 return node;
158 }
159 node = cmp < 0 ? node->left : node->right;
160 }
161 return 0;
162 }
164 struct rbnode *rb_findi(struct rbtree *rb, int key)
165 {
166 return rb_find(rb, INT2PTR(key));
167 }
170 void rb_foreach(struct rbtree *rb, void (*func)(struct rbnode*, void*), void *cls)
171 {
172 traverse(rb->root, func, cls);
173 }
176 struct rbnode *rb_root(struct rbtree *rb)
177 {
178 return rb->root;
179 }
181 void rb_begin(struct rbtree *rb)
182 {
183 rb->rstack = 0;
184 rb->iter = rb->root;
185 }
187 #define push(sp, x) ((x)->next = (sp), (sp) = (x))
188 #define pop(sp) ((sp) = (sp)->next)
189 #define top(sp) (sp)
191 struct rbnode *rb_next(struct rbtree *rb)
192 {
193 struct rbnode *res = 0;
195 while(rb->rstack || rb->iter) {
196 if(rb->iter) {
197 push(rb->rstack, rb->iter);
198 rb->iter = rb->iter->left;
199 } else {
200 rb->iter = top(rb->rstack);
201 pop(rb->rstack);
202 res = rb->iter;
203 rb->iter = rb->iter->right;
204 break;
205 }
206 }
207 return res;
208 }
210 void *rb_node_key(struct rbnode *node)
211 {
212 return node ? node->key : 0;
213 }
215 int rb_node_keyi(struct rbnode *node)
216 {
217 return node ? PTR2INT(node->key) : 0;
218 }
220 void *rb_node_data(struct rbnode *node)
221 {
222 return node ? node->data : 0;
223 }
225 static int cmpaddr(const void *ap, const void *bp)
226 {
227 return ap < bp ? -1 : (ap > bp ? 1 : 0);
228 }
230 static int cmpint(const void *ap, const void *bp)
231 {
232 return PTR2INT(ap) - PTR2INT(bp);
233 }
236 /* ---- left-leaning 2-3 red-black implementation ---- */
238 /* helper prototypes */
239 static int is_red(struct rbnode *tree);
240 static void color_flip(struct rbnode *tree);
241 static struct rbnode *rot_left(struct rbnode *a);
242 static struct rbnode *rot_right(struct rbnode *a);
243 static struct rbnode *find_min(struct rbnode *tree);
244 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree);
245 /*static struct rbnode *move_red_right(struct rbnode *tree);*/
246 static struct rbnode *move_red_left(struct rbnode *tree);
247 static struct rbnode *fix_up(struct rbnode *tree);
249 static int count_nodes(struct rbnode *node)
250 {
251 if(!node)
252 return 0;
254 return 1 + count_nodes(node->left) + count_nodes(node->right);
255 }
257 static void del_tree(struct rbnode *node, rb_del_func_t delfunc, void *cls)
258 {
259 if(!node)
260 return;
262 del_tree(node->left, delfunc, cls);
263 del_tree(node->right, delfunc, cls);
265 if(delfunc) {
266 delfunc(node, cls);
267 }
268 free(node);
269 }
271 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data)
272 {
273 int cmp;
275 if(!tree) {
276 struct rbnode *node = rb->alloc(sizeof *node);
277 node->red = 1;
278 node->key = key;
279 node->data = data;
280 node->left = node->right = 0;
281 return node;
282 }
284 cmp = rb->cmp(key, tree->key);
286 if(cmp < 0) {
287 tree->left = insert(rb, tree->left, key, data);
288 } else if(cmp > 0) {
289 tree->right = insert(rb, tree->right, key, data);
290 } else {
291 tree->data = data;
292 }
294 /* fix right-leaning reds */
295 if(is_red(tree->right)) {
296 tree = rot_left(tree);
297 }
298 /* fix two reds in a row */
299 if(is_red(tree->left) && is_red(tree->left->left)) {
300 tree = rot_right(tree);
301 }
303 /* if 4-node, split it by color inversion */
304 if(is_red(tree->left) && is_red(tree->right)) {
305 color_flip(tree);
306 }
308 return tree;
309 }
311 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key)
312 {
313 int cmp;
315 if(!tree) {
316 return 0;
317 }
319 cmp = rb->cmp(key, tree->key);
321 if(cmp < 0) {
322 if(!is_red(tree->left) && !is_red(tree->left->left)) {
323 tree = move_red_left(tree);
324 }
325 tree->left = delete(rb, tree->left, key);
326 } else {
327 /* need reds on the right */
328 if(is_red(tree->left)) {
329 tree = rot_right(tree);
330 }
332 /* found it at the bottom (XXX what certifies left is null?) */
333 if(cmp == 0 && !tree->right) {
334 if(rb->del) {
335 rb->del(tree, rb->del_cls);
336 }
337 rb->free(tree);
338 return 0;
339 }
341 if(!is_red(tree->right) && !is_red(tree->right->left)) {
342 tree = move_red_left(tree);
343 }
345 if(key == tree->key) {
346 struct rbnode *rmin = find_min(tree->right);
347 tree->key = rmin->key;
348 tree->data = rmin->data;
349 tree->right = del_min(rb, tree->right);
350 } else {
351 tree->right = delete(rb, tree->right, key);
352 }
353 }
355 return fix_up(tree);
356 }
358 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key)
359 {
360 int cmp;
362 if(!node)
363 return 0;
365 if((cmp = rb->cmp(key, node->key)) == 0) {
366 return node;
367 }
368 return find(rb, cmp < 0 ? node->left : node->right, key);
369 }*/
371 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls)
372 {
373 if(!node)
374 return;
376 traverse(node->left, func, cls);
377 func(node, cls);
378 traverse(node->right, func, cls);
379 }
381 /* helpers */
383 static int is_red(struct rbnode *tree)
384 {
385 return tree && tree->red;
386 }
388 static void color_flip(struct rbnode *tree)
389 {
390 tree->red = !tree->red;
391 tree->left->red = !tree->left->red;
392 tree->right->red = !tree->right->red;
393 }
395 static struct rbnode *rot_left(struct rbnode *a)
396 {
397 struct rbnode *b = a->right;
398 a->right = b->left;
399 b->left = a;
400 b->red = a->red;
401 a->red = 1;
402 return b;
403 }
405 static struct rbnode *rot_right(struct rbnode *a)
406 {
407 struct rbnode *b = a->left;
408 a->left = b->right;
409 b->right = a;
410 b->red = a->red;
411 a->red = 1;
412 return b;
413 }
415 static struct rbnode *find_min(struct rbnode *tree)
416 {
417 if(!tree)
418 return 0;
420 while(tree->left) {
421 tree = tree->left;
422 }
423 return tree;
424 }
426 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree)
427 {
428 if(!tree->left) {
429 if(rb->del) {
430 rb->del(tree->left, rb->del_cls);
431 }
432 rb->free(tree->left);
433 return 0;
434 }
436 /* make sure we've got red (3/4-nodes) at the left side so we can delete at the bottom */
437 if(!is_red(tree->left) && !is_red(tree->left->left)) {
438 tree = move_red_left(tree);
439 }
440 tree->left = del_min(rb, tree->left);
442 /* fix right-reds, red-reds, and split 4-nodes on the way up */
443 return fix_up(tree);
444 }
446 #if 0
447 /* push a red link on this node to the right */
448 static struct rbnode *move_red_right(struct rbnode *tree)
449 {
450 /* flipping it makes both children go red, so we have a red to the right */
451 color_flip(tree);
453 /* if after the flip we've got a red-red situation to the left, fix it */
454 if(is_red(tree->left->left)) {
455 tree = rot_right(tree);
456 color_flip(tree);
457 }
458 return tree;
459 }
460 #endif
462 /* push a red link on this node to the left */
463 static struct rbnode *move_red_left(struct rbnode *tree)
464 {
465 /* flipping it makes both children go red, so we have a red to the left */
466 color_flip(tree);
468 /* if after the flip we've got a red-red on the right-left, fix it */
469 if(is_red(tree->right->left)) {
470 tree->right = rot_right(tree->right);
471 tree = rot_left(tree);
472 color_flip(tree);
473 }
474 return tree;
475 }
477 static struct rbnode *fix_up(struct rbnode *tree)
478 {
479 /* fix right-leaning */
480 if(is_red(tree->right)) {
481 tree = rot_left(tree);
482 }
483 /* change invalid red-red pairs into a proper 4-node */
484 if(is_red(tree->left) && is_red(tree->left->left)) {
485 tree = rot_right(tree);
486 }
487 /* split 4-nodes */
488 if(is_red(tree->left) && is_red(tree->right)) {
489 color_flip(tree);
490 }
491 return tree;
492 }