1/*
2 * rbtree.c -- generic red black tree
3 *
4 * Copyright (c) 2001-2006, NLnet Labs. All rights reserved.
5 *
6 * See LICENSE for the license.
7 *
8 */
9
10#include "config.h"
11
12#include <assert.h>
13#include <stdlib.h>
14
15#include "rbtree.h"
16
17#define	BLACK	0
18#define	RED	1
19
20rbnode_type	rbtree_null_node = {
21	RBTREE_NULL,		/* Parent.  */
22	RBTREE_NULL,		/* Left.  */
23	RBTREE_NULL,		/* Right.  */
24	NULL,			/* Key.  */
25	BLACK			/* Color.  */
26};
27
28static void rbtree_rotate_left(rbtree_type *rbtree, rbnode_type *node);
29static void rbtree_rotate_right(rbtree_type *rbtree, rbnode_type *node);
30static void rbtree_insert_fixup(rbtree_type *rbtree, rbnode_type *node);
31static void rbtree_delete_fixup(rbtree_type* rbtree, rbnode_type* child, rbnode_type* child_parent);
32
33/*
34 * Creates a new red black tree, initializes and returns a pointer to it.
35 *
36 * Return NULL on failure.
37 *
38 */
39rbtree_type *
40rbtree_create (region_type *region, int (*cmpf)(const void *, const void *))
41{
42	rbtree_type *rbtree;
43
44	/* Allocate memory for it */
45	rbtree = (rbtree_type *) region_alloc(region, sizeof(rbtree_type));
46	if (!rbtree) {
47		return NULL;
48	}
49
50	/* Initialize it */
51	rbtree->root = RBTREE_NULL;
52	rbtree->count = 0;
53	rbtree->region = region;
54	rbtree->cmp = cmpf;
55
56	return rbtree;
57}
58
59/*
60 * Rotates the node to the left.
61 *
62 */
63static void
64rbtree_rotate_left(rbtree_type *rbtree, rbnode_type *node)
65{
66	rbnode_type *right = node->right;
67	node->right = right->left;
68	if (right->left != RBTREE_NULL)
69		right->left->parent = node;
70
71	right->parent = node->parent;
72
73	if (node->parent != RBTREE_NULL) {
74		if (node == node->parent->left) {
75			node->parent->left = right;
76		} else  {
77			node->parent->right = right;
78		}
79	} else {
80		rbtree->root = right;
81	}
82	right->left = node;
83	node->parent = right;
84}
85
86/*
87 * Rotates the node to the right.
88 *
89 */
90static void
91rbtree_rotate_right(rbtree_type *rbtree, rbnode_type *node)
92{
93	rbnode_type *left = node->left;
94	node->left = left->right;
95	if (left->right != RBTREE_NULL)
96		left->right->parent = node;
97
98	left->parent = node->parent;
99
100	if (node->parent != RBTREE_NULL) {
101		if (node == node->parent->right) {
102			node->parent->right = left;
103		} else  {
104			node->parent->left = left;
105		}
106	} else {
107		rbtree->root = left;
108	}
109	left->right = node;
110	node->parent = left;
111}
112
113static void
114rbtree_insert_fixup(rbtree_type *rbtree, rbnode_type *node)
115{
116	rbnode_type	*uncle;
117
118	/* While not at the root and need fixing... */
119	while (node != rbtree->root && node->parent->color == RED) {
120		/* If our parent is left child of our grandparent... */
121		if (node->parent == node->parent->parent->left) {
122			uncle = node->parent->parent->right;
123
124			/* If our uncle is red... */
125			if (uncle->color == RED) {
126				/* Paint the parent and the uncle black... */
127				node->parent->color = BLACK;
128				uncle->color = BLACK;
129
130				/* And the grandparent red... */
131				node->parent->parent->color = RED;
132
133				/* And continue fixing the grandparent */
134				node = node->parent->parent;
135			} else {				/* Our uncle is black... */
136				/* Are we the right child? */
137				if (node == node->parent->right) {
138					node = node->parent;
139					rbtree_rotate_left(rbtree, node);
140				}
141				/* Now we're the left child, repaint and rotate... */
142				node->parent->color = BLACK;
143				node->parent->parent->color = RED;
144				rbtree_rotate_right(rbtree, node->parent->parent);
145			}
146		} else {
147			uncle = node->parent->parent->left;
148
149			/* If our uncle is red... */
150			if (uncle->color == RED) {
151				/* Paint the parent and the uncle black... */
152				node->parent->color = BLACK;
153				uncle->color = BLACK;
154
155				/* And the grandparent red... */
156				node->parent->parent->color = RED;
157
158				/* And continue fixing the grandparent */
159				node = node->parent->parent;
160			} else {				/* Our uncle is black... */
161				/* Are we the right child? */
162				if (node == node->parent->left) {
163					node = node->parent;
164					rbtree_rotate_right(rbtree, node);
165				}
166				/* Now we're the right child, repaint and rotate... */
167				node->parent->color = BLACK;
168				node->parent->parent->color = RED;
169				rbtree_rotate_left(rbtree, node->parent->parent);
170			}
171		}
172	}
173	rbtree->root->color = BLACK;
174}
175
176
177/*
178 * Inserts a node into a red black tree.
179 *
180 * Returns NULL on failure or the pointer to the newly added node
181 * otherwise.
182 */
183rbnode_type *
184rbtree_insert (rbtree_type *rbtree, rbnode_type *data)
185{
186	/* XXX Not necessary, but keeps compiler quiet... */
187	int r = 0;
188
189	/* We start at the root of the tree */
190	rbnode_type	*node = rbtree->root;
191	rbnode_type	*parent = RBTREE_NULL;
192
193	/* Lets find the new parent... */
194	while (node != RBTREE_NULL) {
195		/* Compare two keys, do we have a duplicate? */
196		if ((r = rbtree->cmp(data->key, node->key)) == 0) {
197			return NULL;
198		}
199		parent = node;
200
201		if (r < 0) {
202			node = node->left;
203		} else {
204			node = node->right;
205		}
206	}
207
208	/* Initialize the new node */
209	data->parent = parent;
210	data->left = data->right = RBTREE_NULL;
211	data->color = RED;
212	rbtree->count++;
213
214	/* Insert it into the tree... */
215	if (parent != RBTREE_NULL) {
216		if (r < 0) {
217			parent->left = data;
218		} else {
219			parent->right = data;
220		}
221	} else {
222		rbtree->root = data;
223	}
224
225	/* Fix up the red-black properties... */
226	rbtree_insert_fixup(rbtree, data);
227
228	return data;
229}
230
231/*
232 * Searches the red black tree, returns the data if key is found or NULL otherwise.
233 *
234 */
235rbnode_type *
236rbtree_search (rbtree_type *rbtree, const void *key)
237{
238	rbnode_type *node;
239
240	if (rbtree_find_less_equal(rbtree, key, &node)) {
241		return node;
242	} else {
243		return NULL;
244	}
245}
246
247/* helpers for delete */
248static void swap_int8(uint8_t* x, uint8_t* y)
249{
250	uint8_t t = *x; *x = *y; *y = t;
251}
252
253static void swap_np(rbnode_type** x, rbnode_type** y)
254{
255	rbnode_type* t = *x; *x = *y; *y = t;
256}
257
258static void change_parent_ptr(rbtree_type* rbtree, rbnode_type* parent, rbnode_type* old, rbnode_type* new)
259{
260	if(parent == RBTREE_NULL)
261	{
262		assert(rbtree->root == old);
263		if(rbtree->root == old) rbtree->root = new;
264		return;
265	}
266	assert(parent->left == old || parent->right == old
267		|| parent->left == new || parent->right == new);
268	if(parent->left == old) parent->left = new;
269	if(parent->right == old) parent->right = new;
270}
271static void change_child_ptr(rbnode_type* child, rbnode_type* old, rbnode_type* new)
272{
273	if(child == RBTREE_NULL) return;
274	assert(child->parent == old || child->parent == new);
275	if(child->parent == old) child->parent = new;
276}
277
278rbnode_type*
279rbtree_delete(rbtree_type *rbtree, const void *key)
280{
281	rbnode_type *to_delete;
282	rbnode_type *child;
283	if((to_delete = rbtree_search(rbtree, key)) == 0) return 0;
284	rbtree->count--;
285
286	/* make sure we have at most one non-leaf child */
287	if(to_delete->left != RBTREE_NULL && to_delete->right != RBTREE_NULL)
288	{
289		/* swap with smallest from right subtree (or largest from left) */
290		rbnode_type *smright = to_delete->right;
291		while(smright->left != RBTREE_NULL)
292			smright = smright->left;
293		/* swap the smright and to_delete elements in the tree,
294		 * but the rbnode_type is first part of user data struct
295		 * so cannot just swap the keys and data pointers. Instead
296		 * readjust the pointers left,right,parent */
297
298		/* swap colors - colors are tied to the position in the tree */
299		swap_int8(&to_delete->color, &smright->color);
300
301		/* swap child pointers in parents of smright/to_delete */
302		change_parent_ptr(rbtree, to_delete->parent, to_delete, smright);
303		if(to_delete->right != smright)
304			change_parent_ptr(rbtree, smright->parent, smright, to_delete);
305
306		/* swap parent pointers in children of smright/to_delete */
307		change_child_ptr(smright->left, smright, to_delete);
308		change_child_ptr(smright->left, smright, to_delete);
309		change_child_ptr(smright->right, smright, to_delete);
310		change_child_ptr(smright->right, smright, to_delete);
311		change_child_ptr(to_delete->left, to_delete, smright);
312		if(to_delete->right != smright)
313			change_child_ptr(to_delete->right, to_delete, smright);
314		if(to_delete->right == smright)
315		{
316			/* set up so after swap they work */
317			to_delete->right = to_delete;
318			smright->parent = smright;
319		}
320
321		/* swap pointers in to_delete/smright nodes */
322		swap_np(&to_delete->parent, &smright->parent);
323		swap_np(&to_delete->left, &smright->left);
324		swap_np(&to_delete->right, &smright->right);
325
326		/* now delete to_delete (which is at the location where the smright previously was) */
327	}
328	assert(to_delete->left == RBTREE_NULL || to_delete->right == RBTREE_NULL);
329
330	if(to_delete->left != RBTREE_NULL) child = to_delete->left;
331	else child = to_delete->right;
332
333	/* unlink to_delete from the tree, replace to_delete with child */
334	change_parent_ptr(rbtree, to_delete->parent, to_delete, child);
335	change_child_ptr(child, to_delete, to_delete->parent);
336
337	if(to_delete->color == RED)
338	{
339		/* if node is red then the child (black) can be swapped in */
340	}
341	else if(child->color == RED)
342	{
343		/* change child to BLACK, removing a RED node is no problem */
344		if(child!=RBTREE_NULL) child->color = BLACK;
345	}
346	else rbtree_delete_fixup(rbtree, child, to_delete->parent);
347
348	/* unlink completely */
349	to_delete->parent = RBTREE_NULL;
350	to_delete->left = RBTREE_NULL;
351	to_delete->right = RBTREE_NULL;
352	to_delete->color = BLACK;
353	return to_delete;
354}
355
356static void rbtree_delete_fixup(rbtree_type* rbtree, rbnode_type* child, rbnode_type* child_parent)
357{
358	rbnode_type* sibling;
359	int go_up = 1;
360
361	/* determine sibling to the node that is one-black short */
362	if(child_parent->right == child) sibling = child_parent->left;
363	else sibling = child_parent->right;
364
365	while(go_up)
366	{
367		if(child_parent == RBTREE_NULL)
368		{
369			/* removed parent==black from root, every path, so ok */
370			return;
371		}
372
373		if(sibling->color == RED)
374		{	/* rotate to get a black sibling */
375			child_parent->color = RED;
376			sibling->color = BLACK;
377			if(child_parent->right == child)
378				rbtree_rotate_right(rbtree, child_parent);
379			else	rbtree_rotate_left(rbtree, child_parent);
380			/* new sibling after rotation */
381			if(child_parent->right == child) sibling = child_parent->left;
382			else sibling = child_parent->right;
383		}
384
385		if(child_parent->color == BLACK
386			&& sibling->color == BLACK
387			&& sibling->left->color == BLACK
388			&& sibling->right->color == BLACK)
389		{	/* fixup local with recolor of sibling */
390			if(sibling != RBTREE_NULL)
391				sibling->color = RED;
392
393			child = child_parent;
394			child_parent = child_parent->parent;
395			/* prepare to go up, new sibling */
396			if(child_parent->right == child) sibling = child_parent->left;
397			else sibling = child_parent->right;
398		}
399		else go_up = 0;
400	}
401
402	if(child_parent->color == RED
403		&& sibling->color == BLACK
404		&& sibling->left->color == BLACK
405		&& sibling->right->color == BLACK)
406	{
407		/* move red to sibling to rebalance */
408		if(sibling != RBTREE_NULL)
409			sibling->color = RED;
410		child_parent->color = BLACK;
411		return;
412	}
413	assert(sibling != RBTREE_NULL);
414
415	/* get a new sibling, by rotating at sibling. See which child
416	   of sibling is red */
417	if(child_parent->right == child
418		&& sibling->color == BLACK
419		&& sibling->right->color == RED
420		&& sibling->left->color == BLACK)
421	{
422		sibling->color = RED;
423		sibling->right->color = BLACK;
424		rbtree_rotate_left(rbtree, sibling);
425		/* new sibling after rotation */
426		if(child_parent->right == child) sibling = child_parent->left;
427		else sibling = child_parent->right;
428	}
429	else if(child_parent->left == child
430		&& sibling->color == BLACK
431		&& sibling->left->color == RED
432		&& sibling->right->color == BLACK)
433	{
434		sibling->color = RED;
435		sibling->left->color = BLACK;
436		rbtree_rotate_right(rbtree, sibling);
437		/* new sibling after rotation */
438		if(child_parent->right == child) sibling = child_parent->left;
439		else sibling = child_parent->right;
440	}
441
442	/* now we have a black sibling with a red child. rotate and exchange colors. */
443	sibling->color = child_parent->color;
444	child_parent->color = BLACK;
445	if(child_parent->right == child)
446	{
447		assert(sibling->left->color == RED);
448		sibling->left->color = BLACK;
449		rbtree_rotate_right(rbtree, child_parent);
450	}
451	else
452	{
453		assert(sibling->right->color == RED);
454		sibling->right->color = BLACK;
455		rbtree_rotate_left(rbtree, child_parent);
456	}
457}
458
459int
460rbtree_find_less_equal(rbtree_type *rbtree, const void *key, rbnode_type **result)
461{
462	int r;
463	rbnode_type *node;
464
465	assert(result);
466
467	/* We start at root... */
468	node = rbtree->root;
469
470	*result = NULL;
471
472	/* While there are children... */
473	while (node != RBTREE_NULL) {
474		r = rbtree->cmp(key, node->key);
475		if (r == 0) {
476			/* Exact match */
477			*result = node;
478			return 1;
479		}
480		if (r < 0) {
481			node = node->left;
482		} else {
483			/* Temporary match */
484			*result = node;
485			node = node->right;
486		}
487	}
488	return 0;
489}
490
491/*
492 * Finds the first element in the red black tree
493 *
494 */
495rbnode_type *
496rbtree_first (rbtree_type *rbtree)
497{
498	rbnode_type *node;
499
500	for (node = rbtree->root; node->left != RBTREE_NULL; node = node->left);
501	return node;
502}
503
504rbnode_type *
505rbtree_last (rbtree_type *rbtree)
506{
507	rbnode_type *node;
508
509	for (node = rbtree->root; node->right != RBTREE_NULL; node = node->right);
510	return node;
511}
512
513/*
514 * Returns the next node...
515 *
516 */
517rbnode_type *
518rbtree_next (rbnode_type *node)
519{
520	rbnode_type *parent;
521
522	if (node->right != RBTREE_NULL) {
523		/* One right, then keep on going left... */
524		for (node = node->right; node->left != RBTREE_NULL; node = node->left);
525	} else {
526		parent = node->parent;
527		while (parent != RBTREE_NULL && node == parent->right) {
528			node = parent;
529			parent = parent->parent;
530		}
531		node = parent;
532	}
533	return node;
534}
535
536rbnode_type *
537rbtree_previous(rbnode_type *node)
538{
539	rbnode_type *parent;
540
541	if (node->left != RBTREE_NULL) {
542		/* One left, then keep on going right... */
543		for (node = node->left; node->right != RBTREE_NULL; node = node->right);
544	} else {
545		parent = node->parent;
546		while (parent != RBTREE_NULL && node == parent->left) {
547			node = parent;
548			parent = parent->parent;
549		}
550		node = parent;
551	}
552	return node;
553}
554