1/*
2 * rbtree.c -- generic red black tree
3 *
4 * Taken from Unbound, modified for ldns
5 *
6 * Copyright (c) 2001-2008, NLnet Labs. All rights reserved.
7 *
8 * This software is open source.
9 *
10 * Redistribution and use in source and binary forms, with or without
11 * modification, are permitted provided that the following conditions
12 * are met:
13 *
14 * Redistributions of source code must retain the above copyright notice,
15 * this list of conditions and the following disclaimer.
16 *
17 * Redistributions in binary form must reproduce the above copyright notice,
18 * this list of conditions and the following disclaimer in the documentation
19 * and/or other materials provided with the distribution.
20 *
21 * Neither the name of the NLNET LABS nor the names of its contributors may
22 * be used to endorse or promote products derived from this software without
23 * specific prior written permission.
24 *
25 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
26 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
27 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
28 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
29 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
30 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
31 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
32 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
33 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
34 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
35 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
36 *
37 */
38
39/**
40 * \file
41 * Implementation of a redblack tree.
42 */
43
44#include <ldns/config.h>
45#include <ldns/rbtree.h>
46#include <ldns/util.h>
47#include <stdlib.h>
48
49/** Node colour black */
50#define	BLACK	0
51/** Node colour red */
52#define	RED	1
53
54/** the NULL node, global alloc */
55ldns_rbnode_t	ldns_rbtree_null_node = {
56	LDNS_RBTREE_NULL,	/* Parent.  */
57	LDNS_RBTREE_NULL,	/* Left.  */
58	LDNS_RBTREE_NULL,	/* Right.  */
59	NULL,			/* Key.  */
60	NULL,               /* Data. */
61	BLACK		     /* Color.  */
62};
63
64/** rotate subtree left (to preserve redblack property) */
65static void ldns_rbtree_rotate_left(ldns_rbtree_t *rbtree, ldns_rbnode_t *node);
66/** rotate subtree right (to preserve redblack property) */
67static void ldns_rbtree_rotate_right(ldns_rbtree_t *rbtree, ldns_rbnode_t *node);
68/** Fixup node colours when insert happened */
69static void ldns_rbtree_insert_fixup(ldns_rbtree_t *rbtree, ldns_rbnode_t *node);
70/** Fixup node colours when delete happened */
71static void ldns_rbtree_delete_fixup(ldns_rbtree_t* rbtree, ldns_rbnode_t* child, ldns_rbnode_t* child_parent);
72
73/*
74 * Creates a new red black tree, intializes and returns a pointer to it.
75 *
76 * Return NULL on failure.
77 *
78 */
79ldns_rbtree_t *
80ldns_rbtree_create (int (*cmpf)(const void *, const void *))
81{
82	ldns_rbtree_t *rbtree;
83
84	/* Allocate memory for it */
85	rbtree = (ldns_rbtree_t *) LDNS_MALLOC(ldns_rbtree_t);
86	if (!rbtree) {
87		return NULL;
88	}
89
90	/* Initialize it */
91	ldns_rbtree_init(rbtree, cmpf);
92
93	return rbtree;
94}
95
96void
97ldns_rbtree_init(ldns_rbtree_t *rbtree, int (*cmpf)(const void *, const void *))
98{
99	/* Initialize it */
100	rbtree->root = LDNS_RBTREE_NULL;
101	rbtree->count = 0;
102	rbtree->cmp = cmpf;
103}
104
105void
106ldns_rbtree_free(ldns_rbtree_t *rbtree)
107{
108	LDNS_FREE(rbtree);
109}
110
111/*
112 * Rotates the node to the left.
113 *
114 */
115static void
116ldns_rbtree_rotate_left(ldns_rbtree_t *rbtree, ldns_rbnode_t *node)
117{
118	ldns_rbnode_t *right = node->right;
119	node->right = right->left;
120	if (right->left != LDNS_RBTREE_NULL)
121		right->left->parent = node;
122
123	right->parent = node->parent;
124
125	if (node->parent != LDNS_RBTREE_NULL) {
126		if (node == node->parent->left) {
127			node->parent->left = right;
128		} else  {
129			node->parent->right = right;
130		}
131	} else {
132		rbtree->root = right;
133	}
134	right->left = node;
135	node->parent = right;
136}
137
138/*
139 * Rotates the node to the right.
140 *
141 */
142static void
143ldns_rbtree_rotate_right(ldns_rbtree_t *rbtree, ldns_rbnode_t *node)
144{
145	ldns_rbnode_t *left = node->left;
146	node->left = left->right;
147	if (left->right != LDNS_RBTREE_NULL)
148		left->right->parent = node;
149
150	left->parent = node->parent;
151
152	if (node->parent != LDNS_RBTREE_NULL) {
153		if (node == node->parent->right) {
154			node->parent->right = left;
155		} else  {
156			node->parent->left = left;
157		}
158	} else {
159		rbtree->root = left;
160	}
161	left->right = node;
162	node->parent = left;
163}
164
165static void
166ldns_rbtree_insert_fixup(ldns_rbtree_t *rbtree, ldns_rbnode_t *node)
167{
168	ldns_rbnode_t	*uncle;
169
170	/* While not at the root and need fixing... */
171	while (node != rbtree->root && node->parent->color == RED) {
172		/* If our parent is left child of our grandparent... */
173		if (node->parent == node->parent->parent->left) {
174			uncle = node->parent->parent->right;
175
176			/* If our uncle is red... */
177			if (uncle->color == RED) {
178				/* Paint the parent and the uncle black... */
179				node->parent->color = BLACK;
180				uncle->color = BLACK;
181
182				/* And the grandparent red... */
183				node->parent->parent->color = RED;
184
185				/* And continue fixing the grandparent */
186				node = node->parent->parent;
187			} else {				/* Our uncle is black... */
188				/* Are we the right child? */
189				if (node == node->parent->right) {
190					node = node->parent;
191					ldns_rbtree_rotate_left(rbtree, node);
192				}
193				/* Now we're the left child, repaint and rotate... */
194				node->parent->color = BLACK;
195				node->parent->parent->color = RED;
196				ldns_rbtree_rotate_right(rbtree, node->parent->parent);
197			}
198		} else {
199			uncle = node->parent->parent->left;
200
201			/* If our uncle is red... */
202			if (uncle->color == RED) {
203				/* Paint the parent and the uncle black... */
204				node->parent->color = BLACK;
205				uncle->color = BLACK;
206
207				/* And the grandparent red... */
208				node->parent->parent->color = RED;
209
210				/* And continue fixing the grandparent */
211				node = node->parent->parent;
212			} else {				/* Our uncle is black... */
213				/* Are we the right child? */
214				if (node == node->parent->left) {
215					node = node->parent;
216					ldns_rbtree_rotate_right(rbtree, node);
217				}
218				/* Now we're the right child, repaint and rotate... */
219				node->parent->color = BLACK;
220				node->parent->parent->color = RED;
221				ldns_rbtree_rotate_left(rbtree, node->parent->parent);
222			}
223		}
224	}
225	rbtree->root->color = BLACK;
226}
227
228void
229ldns_rbtree_insert_vref(ldns_rbnode_t *data, void *rbtree)
230{
231	(void) ldns_rbtree_insert((ldns_rbtree_t *) rbtree,
232						 data);
233}
234
235/*
236 * Inserts a node into a red black tree.
237 *
238 * Returns NULL on failure or the pointer to the newly added node
239 * otherwise.
240 */
241ldns_rbnode_t *
242ldns_rbtree_insert (ldns_rbtree_t *rbtree, ldns_rbnode_t *data)
243{
244	/* XXX Not necessary, but keeps compiler quiet... */
245	int r = 0;
246
247	/* We start at the root of the tree */
248	ldns_rbnode_t	*node = rbtree->root;
249	ldns_rbnode_t	*parent = LDNS_RBTREE_NULL;
250
251	/* Lets find the new parent... */
252	while (node != LDNS_RBTREE_NULL) {
253		/* Compare two keys, do we have a duplicate? */
254		if ((r = rbtree->cmp(data->key, node->key)) == 0) {
255			return NULL;
256		}
257		parent = node;
258
259		if (r < 0) {
260			node = node->left;
261		} else {
262			node = node->right;
263		}
264	}
265
266	/* Initialize the new node */
267	data->parent = parent;
268	data->left = data->right = LDNS_RBTREE_NULL;
269	data->color = RED;
270	rbtree->count++;
271
272	/* Insert it into the tree... */
273	if (parent != LDNS_RBTREE_NULL) {
274		if (r < 0) {
275			parent->left = data;
276		} else {
277			parent->right = data;
278		}
279	} else {
280		rbtree->root = data;
281	}
282
283	/* Fix up the red-black properties... */
284	ldns_rbtree_insert_fixup(rbtree, data);
285
286	return data;
287}
288
289/*
290 * Searches the red black tree, returns the data if key is found or NULL otherwise.
291 *
292 */
293ldns_rbnode_t *
294ldns_rbtree_search (ldns_rbtree_t *rbtree, const void *key)
295{
296	ldns_rbnode_t *node;
297
298	if (ldns_rbtree_find_less_equal(rbtree, key, &node)) {
299		return node;
300	} else {
301		return NULL;
302	}
303}
304
305/** helpers for delete: swap node colours */
306static void swap_int8(uint8_t* x, uint8_t* y)
307{
308	uint8_t t = *x; *x = *y; *y = t;
309}
310
311/** helpers for delete: swap node pointers */
312static void swap_np(ldns_rbnode_t** x, ldns_rbnode_t** y)
313{
314	ldns_rbnode_t* t = *x; *x = *y; *y = t;
315}
316
317/** Update parent pointers of child trees of 'parent' */
318static void change_parent_ptr(ldns_rbtree_t* rbtree, ldns_rbnode_t* parent, ldns_rbnode_t* old, ldns_rbnode_t* new)
319{
320	if(parent == LDNS_RBTREE_NULL)
321	{
322		if(rbtree->root == old) rbtree->root = new;
323		return;
324	}
325	if(parent->left == old) parent->left = new;
326	if(parent->right == old) parent->right = new;
327}
328/** Update parent pointer of a node 'child' */
329static void change_child_ptr(ldns_rbnode_t* child, ldns_rbnode_t* old, ldns_rbnode_t* new)
330{
331	if(child == LDNS_RBTREE_NULL) return;
332	if(child->parent == old) child->parent = new;
333}
334
335ldns_rbnode_t*
336ldns_rbtree_delete(ldns_rbtree_t *rbtree, const void *key)
337{
338	ldns_rbnode_t *to_delete;
339	ldns_rbnode_t *child;
340	if((to_delete = ldns_rbtree_search(rbtree, key)) == 0) return 0;
341	rbtree->count--;
342
343	/* make sure we have at most one non-leaf child */
344	if(to_delete->left != LDNS_RBTREE_NULL &&
345	   to_delete->right != LDNS_RBTREE_NULL)
346	{
347		/* swap with smallest from right subtree (or largest from left) */
348		ldns_rbnode_t *smright = to_delete->right;
349		while(smright->left != LDNS_RBTREE_NULL)
350			smright = smright->left;
351		/* swap the smright and to_delete elements in the tree,
352		 * but the ldns_rbnode_t is first part of user data struct
353		 * so cannot just swap the keys and data pointers. Instead
354		 * readjust the pointers left,right,parent */
355
356		/* swap colors - colors are tied to the position in the tree */
357		swap_int8(&to_delete->color, &smright->color);
358
359		/* swap child pointers in parents of smright/to_delete */
360		change_parent_ptr(rbtree, to_delete->parent, to_delete, smright);
361		if(to_delete->right != smright)
362			change_parent_ptr(rbtree, smright->parent, smright, to_delete);
363
364		/* swap parent pointers in children of smright/to_delete */
365		change_child_ptr(smright->left, smright, to_delete);
366		change_child_ptr(smright->left, smright, to_delete);
367		change_child_ptr(smright->right, smright, to_delete);
368		change_child_ptr(smright->right, smright, to_delete);
369		change_child_ptr(to_delete->left, to_delete, smright);
370		if(to_delete->right != smright)
371			change_child_ptr(to_delete->right, to_delete, smright);
372		if(to_delete->right == smright)
373		{
374			/* set up so after swap they work */
375			to_delete->right = to_delete;
376			smright->parent = smright;
377		}
378
379		/* swap pointers in to_delete/smright nodes */
380		swap_np(&to_delete->parent, &smright->parent);
381		swap_np(&to_delete->left, &smright->left);
382		swap_np(&to_delete->right, &smright->right);
383
384		/* now delete to_delete (which is at the location where the smright previously was) */
385	}
386
387	if(to_delete->left != LDNS_RBTREE_NULL) child = to_delete->left;
388	else child = to_delete->right;
389
390	/* unlink to_delete from the tree, replace to_delete with child */
391	change_parent_ptr(rbtree, to_delete->parent, to_delete, child);
392	change_child_ptr(child, to_delete, to_delete->parent);
393
394	if(to_delete->color == RED)
395	{
396		/* if node is red then the child (black) can be swapped in */
397	}
398	else if(child->color == RED)
399	{
400		/* change child to BLACK, removing a RED node is no problem */
401		if(child!=LDNS_RBTREE_NULL) child->color = BLACK;
402	}
403	else ldns_rbtree_delete_fixup(rbtree, child, to_delete->parent);
404
405	/* unlink completely */
406	to_delete->parent = LDNS_RBTREE_NULL;
407	to_delete->left = LDNS_RBTREE_NULL;
408	to_delete->right = LDNS_RBTREE_NULL;
409	to_delete->color = BLACK;
410	return to_delete;
411}
412
413static void ldns_rbtree_delete_fixup(ldns_rbtree_t* rbtree, ldns_rbnode_t* child, ldns_rbnode_t* child_parent)
414{
415	ldns_rbnode_t* sibling;
416	int go_up = 1;
417
418	/* determine sibling to the node that is one-black short */
419	if(child_parent->right == child) sibling = child_parent->left;
420	else sibling = child_parent->right;
421
422	while(go_up)
423	{
424		if(child_parent == LDNS_RBTREE_NULL)
425		{
426			/* removed parent==black from root, every path, so ok */
427			return;
428		}
429
430		if(sibling->color == RED)
431		{	/* rotate to get a black sibling */
432			child_parent->color = RED;
433			sibling->color = BLACK;
434			if(child_parent->right == child)
435				ldns_rbtree_rotate_right(rbtree, child_parent);
436			else	ldns_rbtree_rotate_left(rbtree, child_parent);
437			/* new sibling after rotation */
438			if(child_parent->right == child) sibling = child_parent->left;
439			else sibling = child_parent->right;
440		}
441
442		if(child_parent->color == BLACK
443			&& sibling->color == BLACK
444			&& sibling->left->color == BLACK
445			&& sibling->right->color == BLACK)
446		{	/* fixup local with recolor of sibling */
447			if(sibling != LDNS_RBTREE_NULL)
448				sibling->color = RED;
449
450			child = child_parent;
451			child_parent = child_parent->parent;
452			/* prepare to go up, new sibling */
453			if(child_parent->right == child) sibling = child_parent->left;
454			else sibling = child_parent->right;
455		}
456		else go_up = 0;
457	}
458
459	if(child_parent->color == RED
460		&& sibling->color == BLACK
461		&& sibling->left->color == BLACK
462		&& sibling->right->color == BLACK)
463	{
464		/* move red to sibling to rebalance */
465		if(sibling != LDNS_RBTREE_NULL)
466			sibling->color = RED;
467		child_parent->color = BLACK;
468		return;
469	}
470
471	/* get a new sibling, by rotating at sibling. See which child
472	   of sibling is red */
473	if(child_parent->right == child
474		&& sibling->color == BLACK
475		&& sibling->right->color == RED
476		&& sibling->left->color == BLACK)
477	{
478		sibling->color = RED;
479		sibling->right->color = BLACK;
480		ldns_rbtree_rotate_left(rbtree, sibling);
481		/* new sibling after rotation */
482		if(child_parent->right == child) sibling = child_parent->left;
483		else sibling = child_parent->right;
484	}
485	else if(child_parent->left == child
486		&& sibling->color == BLACK
487		&& sibling->left->color == RED
488		&& sibling->right->color == BLACK)
489	{
490		sibling->color = RED;
491		sibling->left->color = BLACK;
492		ldns_rbtree_rotate_right(rbtree, sibling);
493		/* new sibling after rotation */
494		if(child_parent->right == child) sibling = child_parent->left;
495		else sibling = child_parent->right;
496	}
497
498	/* now we have a black sibling with a red child. rotate and exchange colors. */
499	sibling->color = child_parent->color;
500	child_parent->color = BLACK;
501	if(child_parent->right == child)
502	{
503		sibling->left->color = BLACK;
504		ldns_rbtree_rotate_right(rbtree, child_parent);
505	}
506	else
507	{
508		sibling->right->color = BLACK;
509		ldns_rbtree_rotate_left(rbtree, child_parent);
510	}
511}
512
513int
514ldns_rbtree_find_less_equal(ldns_rbtree_t *rbtree, const void *key, ldns_rbnode_t **result)
515{
516	int r;
517	ldns_rbnode_t *node;
518
519	/* We start at root... */
520	node = rbtree->root;
521
522	*result = NULL;
523
524	/* While there are children... */
525	while (node != LDNS_RBTREE_NULL) {
526		r = rbtree->cmp(key, node->key);
527		if (r == 0) {
528			/* Exact match */
529			*result = node;
530			return 1;
531		}
532		if (r < 0) {
533			node = node->left;
534		} else {
535			/* Temporary match */
536			*result = node;
537			node = node->right;
538		}
539	}
540	return 0;
541}
542
543/*
544 * Finds the first element in the red black tree
545 *
546 */
547ldns_rbnode_t *
548ldns_rbtree_first(const ldns_rbtree_t *rbtree)
549{
550	ldns_rbnode_t *node = rbtree->root;
551
552	if (rbtree->root != LDNS_RBTREE_NULL) {
553		for (node = rbtree->root; node->left != LDNS_RBTREE_NULL; node = node->left);
554	}
555	return node;
556}
557
558ldns_rbnode_t *
559ldns_rbtree_last(const ldns_rbtree_t *rbtree)
560{
561	ldns_rbnode_t *node = rbtree->root;
562
563	if (rbtree->root != LDNS_RBTREE_NULL) {
564		for (node = rbtree->root; node->right != LDNS_RBTREE_NULL; node = node->right);
565	}
566	return node;
567}
568
569/*
570 * Returns the next node...
571 *
572 */
573ldns_rbnode_t *
574ldns_rbtree_next(ldns_rbnode_t *node)
575{
576	ldns_rbnode_t *parent;
577
578	if (node->right != LDNS_RBTREE_NULL) {
579		/* One right, then keep on going left... */
580		for (node = node->right;
581			node->left != LDNS_RBTREE_NULL;
582			node = node->left);
583	} else {
584		parent = node->parent;
585		while (parent != LDNS_RBTREE_NULL && node == parent->right) {
586			node = parent;
587			parent = parent->parent;
588		}
589		node = parent;
590	}
591	return node;
592}
593
594ldns_rbnode_t *
595ldns_rbtree_previous(ldns_rbnode_t *node)
596{
597	ldns_rbnode_t *parent;
598
599	if (node->left != LDNS_RBTREE_NULL) {
600		/* One left, then keep on going right... */
601		for (node = node->left;
602			node->right != LDNS_RBTREE_NULL;
603			node = node->right);
604	} else {
605		parent = node->parent;
606		while (parent != LDNS_RBTREE_NULL && node == parent->left) {
607			node = parent;
608			parent = parent->parent;
609		}
610		node = parent;
611	}
612	return node;
613}
614
615/**
616 * split off elements number of elements from the start
617 * of the name tree and return a new tree
618 */
619ldns_rbtree_t *
620ldns_rbtree_split(ldns_rbtree_t *tree,
621			   size_t elements)
622{
623	ldns_rbtree_t *new_tree;
624	ldns_rbnode_t *cur_node;
625	ldns_rbnode_t *move_node;
626	size_t count = 0;
627
628	new_tree = ldns_rbtree_create(tree->cmp);
629
630	cur_node = ldns_rbtree_first(tree);
631	while (count < elements && cur_node != LDNS_RBTREE_NULL) {
632		move_node = ldns_rbtree_delete(tree, cur_node->key);
633		(void)ldns_rbtree_insert(new_tree, move_node);
634		cur_node = ldns_rbtree_first(tree);
635		count++;
636	}
637
638	return new_tree;
639}
640
641/*
642 * add all node from the second tree to the first (removing them from the
643 * second), and fix up nsec(3)s if present
644 */
645void
646ldns_rbtree_join(ldns_rbtree_t *tree1, ldns_rbtree_t *tree2)
647{
648	ldns_traverse_postorder(tree2, ldns_rbtree_insert_vref, tree1);
649}
650
651/** recursive descent traverse */
652static void
653traverse_post(void (*func)(ldns_rbnode_t*, void*), void* arg,
654	ldns_rbnode_t* node)
655{
656	if(!node || node == LDNS_RBTREE_NULL)
657		return;
658	/* recurse */
659	traverse_post(func, arg, node->left);
660	traverse_post(func, arg, node->right);
661	/* call user func */
662	(*func)(node, arg);
663}
664
665void
666ldns_traverse_postorder(ldns_rbtree_t* tree,
667	void (*func)(ldns_rbnode_t*, void*), void* arg)
668{
669	traverse_post(func, arg, tree->root);
670}
671