kern
view src/rbtree.c @ 80:4db99a52863e
fixed the "endianess" of the text messages in the ATA identify info block. this
is the first time I've seen wrong byteorder in ascii text, the ATA committee
should be commended.
| author | John Tsiombikas <nuclear@member.fsf.org> |
|---|---|
| date | Tue, 06 Dec 2011 13:35:39 +0200 |
| parents | b45e2d5f0ae1 |
| children |
line source
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <string.h>
4 #include "rbtree.h"
5 #include "panic.h"
7 #define INT2PTR(x) ((void*)(x))
8 #define PTR2INT(x) ((int)(x))
10 static int cmpaddr(void *ap, void *bp);
11 static int cmpint(void *ap, void *bp);
13 static int count_nodes(struct rbnode *node);
14 static void del_tree(struct rbnode *node, void (*delfunc)(struct rbnode*, void*), void *cls);
15 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data);
16 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key);
17 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls);
19 struct rbtree *rb_create(rb_cmp_func_t cmp_func)
20 {
21 struct rbtree *rb;
23 if(!(rb = malloc(sizeof *rb))) {
24 return 0;
25 }
26 if(rb_init(rb, cmp_func) == -1) {
27 free(rb);
28 return 0;
29 }
30 return rb;
31 }
33 void rb_free(struct rbtree *rb)
34 {
35 rb_destroy(rb);
36 free(rb);
37 }
40 int rb_init(struct rbtree *rb, rb_cmp_func_t cmp_func)
41 {
42 memset(rb, 0, sizeof *rb);
44 if(cmp_func == RB_KEY_INT) {
45 rb->cmp = cmpint;
46 } else if(cmp_func == RB_KEY_STRING) {
47 rb->cmp = (rb_cmp_func_t)strcmp;
48 } else {
49 rb->cmp = cmpaddr;
50 }
52 rb->alloc = malloc;
53 rb->free = free;
54 return 0;
55 }
57 void rb_destroy(struct rbtree *rb)
58 {
59 del_tree(rb->root, rb->del, rb->del_cls);
60 }
62 void rb_clear(struct rbtree *rb)
63 {
64 del_tree(rb->root, rb->del, rb->del_cls);
65 rb->root = 0;
66 }
68 int rb_copy(struct rbtree *dest, struct rbtree *src)
69 {
70 struct rbnode *node;
72 rb_clear(dest);
74 rb_begin(src);
75 while((node = rb_next(src))) {
76 if(rb_insert(dest, node->key, node->data) == -1) {
77 return -1;
78 }
79 }
80 return 0;
81 }
83 void rb_set_allocator(struct rbtree *rb, rb_alloc_func_t alloc, rb_free_func_t free)
84 {
85 rb->alloc = alloc;
86 rb->free = free;
87 }
90 void rb_set_compare_func(struct rbtree *rb, rb_cmp_func_t func)
91 {
92 rb->cmp = func;
93 }
95 void rb_set_delete_func(struct rbtree *rb, rb_del_func_t func, void *cls)
96 {
97 rb->del = func;
98 rb->del_cls = cls;
99 }
101 int rb_size(struct rbtree *rb)
102 {
103 return count_nodes(rb->root);
104 }
106 int rb_insert(struct rbtree *rb, void *key, void *data)
107 {
108 rb->root = insert(rb, rb->root, key, data);
109 rb->root->red = 0;
110 return 0;
111 }
113 int rb_inserti(struct rbtree *rb, int key, void *data)
114 {
115 rb->root = insert(rb, rb->root, INT2PTR(key), data);
116 rb->root->red = 0;
117 return 0;
118 }
121 int rb_delete(struct rbtree *rb, void *key)
122 {
123 rb->root = delete(rb, rb->root, key);
124 rb->root->red = 0;
125 return 0;
126 }
128 int rb_deletei(struct rbtree *rb, int key)
129 {
130 rb->root = delete(rb, rb->root, INT2PTR(key));
131 rb->root->red = 0;
132 return 0;
133 }
136 void *rb_find(struct rbtree *rb, void *key)
137 {
138 struct rbnode *node = rb->root;
140 while(node) {
141 int cmp = rb->cmp(key, node->key);
142 if(cmp == 0) {
143 return node;
144 }
145 node = cmp < 0 ? node->left : node->right;
146 }
147 return 0;
148 }
150 void *rb_findi(struct rbtree *rb, int key)
151 {
152 return rb_find(rb, INT2PTR(key));
153 }
156 void rb_foreach(struct rbtree *rb, void (*func)(struct rbnode*, void*), void *cls)
157 {
158 traverse(rb->root, func, cls);
159 }
162 struct rbnode *rb_root(struct rbtree *rb)
163 {
164 return rb->root;
165 }
167 void rb_begin(struct rbtree *rb)
168 {
169 rb->rstack = 0;
170 rb->iter = rb->root;
171 }
173 #define push(sp, x) ((x)->next = (sp), (sp) = (x))
174 #define pop(sp) ((sp) = (sp)->next)
175 #define top(sp) (sp)
177 struct rbnode *rb_next(struct rbtree *rb)
178 {
179 struct rbnode *res = 0;
181 while(rb->rstack || rb->iter) {
182 if(rb->iter) {
183 push(rb->rstack, rb->iter);
184 rb->iter = rb->iter->left;
185 } else {
186 rb->iter = top(rb->rstack);
187 pop(rb->rstack);
188 res = rb->iter;
189 rb->iter = rb->iter->right;
190 break;
191 }
192 }
193 return res;
194 }
196 void *rb_node_key(struct rbnode *node)
197 {
198 return node ? node->key : 0;
199 }
201 int rb_node_keyi(struct rbnode *node)
202 {
203 return node ? PTR2INT(node->key) : 0;
204 }
206 void *rb_node_data(struct rbnode *node)
207 {
208 return node ? node->data : 0;
209 }
211 static int cmpaddr(void *ap, void *bp)
212 {
213 return ap < bp ? -1 : (ap > bp ? 1 : 0);
214 }
216 static int cmpint(void *ap, void *bp)
217 {
218 return PTR2INT(ap) - PTR2INT(bp);
219 }
222 /* ---- left-leaning 2-3 red-black implementation ---- */
224 /* helper prototypes */
225 static int is_red(struct rbnode *tree);
226 static void color_flip(struct rbnode *tree);
227 static struct rbnode *rot_left(struct rbnode *a);
228 static struct rbnode *rot_right(struct rbnode *a);
229 static struct rbnode *find_min(struct rbnode *tree);
230 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree);
231 /*static struct rbnode *move_red_right(struct rbnode *tree);*/
232 static struct rbnode *move_red_left(struct rbnode *tree);
233 static struct rbnode *fix_up(struct rbnode *tree);
235 static int count_nodes(struct rbnode *node)
236 {
237 if(!node)
238 return 0;
240 return 1 + count_nodes(node->left) + count_nodes(node->right);
241 }
243 static void del_tree(struct rbnode *node, rb_del_func_t delfunc, void *cls)
244 {
245 if(!node)
246 return;
248 del_tree(node->left, delfunc, cls);
249 del_tree(node->right, delfunc, cls);
251 if(delfunc) {
252 delfunc(node, cls);
253 }
254 free(node);
255 }
257 static struct rbnode *insert(struct rbtree *rb, struct rbnode *tree, void *key, void *data)
258 {
259 int cmp;
261 if(!tree) {
262 struct rbnode *node = rb->alloc(sizeof *node);
263 if(!node) {
264 panic("failed to allocate tree node\n");
265 }
266 node->red = 1;
267 node->key = key;
268 node->data = data;
269 node->left = node->right = 0;
270 return node;
271 }
273 cmp = rb->cmp(key, tree->key);
275 if(cmp < 0) {
276 tree->left = insert(rb, tree->left, key, data);
277 } else if(cmp > 0) {
278 tree->right = insert(rb, tree->right, key, data);
279 } else {
280 tree->data = data;
281 }
283 /* fix right-leaning reds */
284 if(is_red(tree->right)) {
285 tree = rot_left(tree);
286 }
287 /* fix two reds in a row */
288 if(is_red(tree->left) && is_red(tree->left->left)) {
289 tree = rot_right(tree);
290 }
292 /* if 4-node, split it by color inversion */
293 if(is_red(tree->left) && is_red(tree->right)) {
294 color_flip(tree);
295 }
297 return tree;
298 }
300 static struct rbnode *delete(struct rbtree *rb, struct rbnode *tree, void *key)
301 {
302 int cmp;
304 if(!tree) {
305 return 0;
306 }
308 cmp = rb->cmp(key, tree->key);
310 if(cmp < 0) {
311 if(!is_red(tree->left) && !is_red(tree->left->left)) {
312 tree = move_red_left(tree);
313 }
314 tree->left = delete(rb, tree->left, key);
315 } else {
316 /* need reds on the right */
317 if(is_red(tree->left)) {
318 tree = rot_right(tree);
319 }
321 /* found it at the bottom (XXX what certifies left is null?) */
322 if(cmp == 0 && !tree->right) {
323 if(rb->del) {
324 rb->del(tree, rb->del_cls);
325 }
326 rb->free(tree);
327 return 0;
328 }
330 if(!is_red(tree->right) && !is_red(tree->right->left)) {
331 tree = move_red_left(tree);
332 }
334 if(key == tree->key) {
335 struct rbnode *rmin = find_min(tree->right);
336 tree->key = rmin->key;
337 tree->data = rmin->data;
338 tree->right = del_min(rb, tree->right);
339 } else {
340 tree->right = delete(rb, tree->right, key);
341 }
342 }
344 return fix_up(tree);
345 }
347 /*static struct rbnode *find(struct rbtree *rb, struct rbnode *node, void *key)
348 {
349 int cmp;
351 if(!node)
352 return 0;
354 if((cmp = rb->cmp(key, node->key)) == 0) {
355 return node;
356 }
357 return find(rb, cmp < 0 ? node->left : node->right, key);
358 }*/
360 static void traverse(struct rbnode *node, void (*func)(struct rbnode*, void*), void *cls)
361 {
362 if(!node)
363 return;
365 traverse(node->left, func, cls);
366 func(node, cls);
367 traverse(node->right, func, cls);
368 }
370 /* helpers */
372 static int is_red(struct rbnode *tree)
373 {
374 return tree && tree->red;
375 }
377 static void color_flip(struct rbnode *tree)
378 {
379 tree->red = !tree->red;
380 tree->left->red = !tree->left->red;
381 tree->right->red = !tree->right->red;
382 }
384 static struct rbnode *rot_left(struct rbnode *a)
385 {
386 struct rbnode *b = a->right;
387 a->right = b->left;
388 b->left = a;
389 b->red = a->red;
390 a->red = 1;
391 return b;
392 }
394 static struct rbnode *rot_right(struct rbnode *a)
395 {
396 struct rbnode *b = a->left;
397 a->left = b->right;
398 b->right = a;
399 b->red = a->red;
400 a->red = 1;
401 return b;
402 }
404 static struct rbnode *find_min(struct rbnode *tree)
405 {
406 struct rbnode *node;
408 if(!tree)
409 return 0;
411 while(node->left) {
412 node = node->left;
413 }
414 return node;
415 }
417 static struct rbnode *del_min(struct rbtree *rb, struct rbnode *tree)
418 {
419 if(!tree->left) {
420 if(rb->del) {
421 rb->del(tree->left, rb->del_cls);
422 }
423 rb->free(tree->left);
424 return 0;
425 }
427 /* make sure we've got red (3/4-nodes) at the left side so we can delete at the bottom */
428 if(!is_red(tree->left) && !is_red(tree->left->left)) {
429 tree = move_red_left(tree);
430 }
431 tree->left = del_min(rb, tree->left);
433 /* fix right-reds, red-reds, and split 4-nodes on the way up */
434 return fix_up(tree);
435 }
437 #if 0
438 /* push a red link on this node to the right */
439 static struct rbnode *move_red_right(struct rbnode *tree)
440 {
441 /* flipping it makes both children go red, so we have a red to the right */
442 color_flip(tree);
444 /* if after the flip we've got a red-red situation to the left, fix it */
445 if(is_red(tree->left->left)) {
446 tree = rot_right(tree);
447 color_flip(tree);
448 }
449 return tree;
450 }
451 #endif
453 /* push a red link on this node to the left */
454 static struct rbnode *move_red_left(struct rbnode *tree)
455 {
456 /* flipping it makes both children go red, so we have a red to the left */
457 color_flip(tree);
459 /* if after the flip we've got a red-red on the right-left, fix it */
460 if(is_red(tree->right->left)) {
461 tree->right = rot_right(tree->right);
462 tree = rot_left(tree);
463 color_flip(tree);
464 }
465 return tree;
466 }
468 static struct rbnode *fix_up(struct rbnode *tree)
469 {
470 /* fix right-leaning */
471 if(is_red(tree->right)) {
472 tree = rot_left(tree);
473 }
474 /* change invalid red-red pairs into a proper 4-node */
475 if(is_red(tree->left) && is_red(tree->left->left)) {
476 tree = rot_right(tree);
477 }
478 /* split 4-nodes */
479 if(is_red(tree->left) && is_red(tree->right)) {
480 color_flip(tree);
481 }
482 return tree;
483 }
485 void rb_dbg_print_tree(struct rbtree *tree)
486 {
487 struct rbnode *node;
489 rb_begin(tree);
490 while((node = rb_next(tree))) {
491 printf("%d ", rb_node_keyi(node));
492 }
493 printf("\n");
494 }