rbtree.c 15.2 KB
Newer Older
L
Linus Torvalds 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  (C) 2002  David Woodhouse <dwmw2@infradead.org>
  
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/lib/rbtree.c
*/

#include <linux/rbtree.h>
24
#include <linux/export.h>
L
Linus Torvalds 已提交
25

26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
/*
 * red-black trees properties:  http://en.wikipedia.org/wiki/Rbtree
 *
 *  1) A node is either red or black
 *  2) The root is black
 *  3) All leaves (NULL) are black
 *  4) Both children of every red node are black
 *  5) Every simple path from root to leaves contains the same number
 *     of black nodes.
 *
 *  4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
 *  consecutive red nodes in a path and every red node is therefore followed by
 *  a black. So if B is the number of black nodes on every simple path (as per
 *  5), then the longest possible path due to 4 is 2B.
 *
 *  We shall indicate color with case, where black nodes are uppercase and red
42 43
 *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
 *  parentheses and have some accompanying text comment.
44 45
 */

46 47 48 49 50 51 52 53 54 55 56 57
#define	RB_RED		0
#define	RB_BLACK	1

#define rb_color(r)   ((r)->__rb_parent_color & 1)
#define rb_is_red(r)   (!rb_color(r))
#define rb_is_black(r) rb_color(r)

static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
{
	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
}

58 59 60 61 62 63 64 65 66 67 68
static inline void rb_set_parent_color(struct rb_node *rb,
				       struct rb_node *p, int color)
{
	rb->__rb_parent_color = (unsigned long)p | color;
}

static inline struct rb_node *rb_red_parent(struct rb_node *red)
{
	return (struct rb_node *)red->__rb_parent_color;
}

69 70 71 72 73 74 75 76 77 78 79 80 81
static inline void
__rb_change_child(struct rb_node *old, struct rb_node *new,
		  struct rb_node *parent, struct rb_root *root)
{
	if (parent) {
		if (parent->rb_left == old)
			parent->rb_left = new;
		else
			parent->rb_right = new;
	} else
		root->rb_node = new;
}

82 83 84 85 86 87 88 89 90 91 92 93
/*
 * Helper function for rotations:
 * - old's parent and color get assigned to new
 * - old gets assigned new as a parent and 'color' as a color.
 */
static inline void
__rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
			struct rb_root *root, int color)
{
	struct rb_node *parent = rb_parent(old);
	new->__rb_parent_color = old->__rb_parent_color;
	rb_set_parent_color(old, new, color);
94
	__rb_change_child(old, new, parent, root);
95 96
}

L
Linus Torvalds 已提交
97 98
void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
99
	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
L
Linus Torvalds 已提交
100

101 102 103 104 105 106 107 108 109
	while (true) {
		/*
		 * Loop invariant: node is red
		 *
		 * If there is a black parent, we are done.
		 * Otherwise, take some corrective action as we don't
		 * want a red root or two consecutive red nodes.
		 */
		if (!parent) {
110
			rb_set_parent_color(node, NULL, RB_BLACK);
111 112 113 114
			break;
		} else if (rb_is_black(parent))
			break;

115 116
		gparent = rb_red_parent(parent);

117 118
		tmp = gparent->rb_right;
		if (parent != tmp) {	/* parent == gparent->rb_left */
119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138
			if (tmp && rb_is_red(tmp)) {
				/*
				 * Case 1 - color flips
				 *
				 *       G            g
				 *      / \          / \
				 *     p   u  -->   P   U
				 *    /            /
				 *   n            N
				 *
				 * However, since g's parent might be red, and
				 * 4) does not allow this, we need to recurse
				 * at g.
				 */
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);
				rb_set_parent_color(node, parent, RB_RED);
				continue;
L
Linus Torvalds 已提交
139 140
			}

141 142
			tmp = parent->rb_right;
			if (node == tmp) {
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160
				/*
				 * Case 2 - left rotate at parent
				 *
				 *      G             G
				 *     / \           / \
				 *    p   U  -->    n   U
				 *     \           /
				 *      n         p
				 *
				 * This still leaves us in violation of 4), the
				 * continuation into Case 3 will fix that.
				 */
				parent->rb_right = tmp = node->rb_left;
				node->rb_left = parent;
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
L
Linus Torvalds 已提交
161
				parent = node;
162
				tmp = node->rb_right;
L
Linus Torvalds 已提交
163 164
			}

165 166 167 168 169 170 171 172 173
			/*
			 * Case 3 - right rotate at gparent
			 *
			 *        G           P
			 *       / \         / \
			 *      p   U  -->  n   g
			 *     /                 \
			 *    n                   U
			 */
174
			gparent->rb_left = tmp;  /* == parent->rb_right */
175 176 177 178
			parent->rb_right = gparent;
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
179
			break;
L
Linus Torvalds 已提交
180
		} else {
181 182 183 184 185 186 187 188 189
			tmp = gparent->rb_left;
			if (tmp && rb_is_red(tmp)) {
				/* Case 1 - color flips */
				rb_set_parent_color(tmp, gparent, RB_BLACK);
				rb_set_parent_color(parent, gparent, RB_BLACK);
				node = gparent;
				parent = rb_parent(node);
				rb_set_parent_color(node, parent, RB_RED);
				continue;
L
Linus Torvalds 已提交
190 191
			}

192 193
			tmp = parent->rb_left;
			if (node == tmp) {
194 195 196 197 198 199 200
				/* Case 2 - right rotate at parent */
				parent->rb_left = tmp = node->rb_right;
				node->rb_right = parent;
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
L
Linus Torvalds 已提交
201
				parent = node;
202
				tmp = node->rb_left;
L
Linus Torvalds 已提交
203 204
			}

205
			/* Case 3 - left rotate at gparent */
206
			gparent->rb_right = tmp;  /* == parent->rb_left */
207 208 209 210
			parent->rb_left = gparent;
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
211
			break;
L
Linus Torvalds 已提交
212 213 214 215 216 217 218 219
		}
	}
}
EXPORT_SYMBOL(rb_insert_color);

static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
			     struct rb_root *root)
{
220
	struct rb_node *sibling, *tmp1, *tmp2;
L
Linus Torvalds 已提交
221

222 223 224 225 226 227 228 229 230 231 232
	while (true) {
		/*
		 * Loop invariant: all leaf paths going through node have a
		 * black node count that is 1 lower than other leaf paths.
		 *
		 * If node is red, we can flip it to black to adjust.
		 * If node is the root, all leaf paths go through it.
		 * Otherwise, we need to adjust the tree through color flips
		 * and tree rotations as per one of the 4 cases below.
		 */
		if (node && rb_is_red(node)) {
233
			rb_set_parent_color(node, parent, RB_BLACK);
234 235 236
			break;
		} else if (!parent) {
			break;
237 238 239
		}
		sibling = parent->rb_right;
		if (node != sibling) {	/* node == parent->rb_left */
240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255
			if (rb_is_red(sibling)) {
				/*
				 * Case 1 - left rotate at parent
				 *
				 *     P               S
				 *    / \             / \
				 *   N   s    -->    p   Sr
				 *      / \         / \
				 *     Sl  Sr      N   Sl
				 */
				parent->rb_right = tmp1 = sibling->rb_left;
				sibling->rb_left = parent;
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				sibling = tmp1;
L
Linus Torvalds 已提交
256
			}
257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
			tmp1 = sibling->rb_right;
			if (!tmp1 || rb_is_black(tmp1)) {
				tmp2 = sibling->rb_left;
				if (!tmp2 || rb_is_black(tmp2)) {
					/*
					 * Case 2 - sibling color flip
					 * (p could be either color here)
					 *
					 *    (p)           (p)
					 *    / \           / \
					 *   N   S    -->  N   s
					 *      / \           / \
					 *     Sl  Sr        Sl  Sr
					 *
					 * This leaves us violating 5), so
					 * recurse at p. If p is red, the
					 * recursion will just flip it to black
					 * and exit. If coming from Case 1,
					 * p is known to be red.
					 */
					rb_set_parent_color(sibling, parent,
							    RB_RED);
279 280 281
					node = parent;
					parent = rb_parent(node);
					continue;
L
Linus Torvalds 已提交
282
				}
283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302
				/*
				 * Case 3 - right rotate at sibling
				 * (p could be either color here)
				 *
				 *   (p)           (p)
				 *   / \           / \
				 *  N   S    -->  N   Sl
				 *     / \             \
				 *    sl  Sr            s
				 *                       \
				 *                        Sr
				 */
				sibling->rb_left = tmp1 = tmp2->rb_right;
				tmp2->rb_right = sibling;
				parent->rb_right = tmp2;
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				tmp1 = sibling;
				sibling = tmp2;
L
Linus Torvalds 已提交
303
			}
304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322
			/*
			 * Case 4 - left rotate at parent + color flips
			 * (p and sl could be either color here.
			 *  After rotation, p becomes black, s acquires
			 *  p's color, and sl keeps its color)
			 *
			 *      (p)             (s)
			 *      / \             / \
			 *     N   S     -->   P   Sr
			 *        / \         / \
			 *      (sl) sr      N  (sl)
			 */
			parent->rb_right = tmp2 = sibling->rb_left;
			sibling->rb_left = parent;
			rb_set_parent_color(tmp1, sibling, RB_BLACK);
			if (tmp2)
				rb_set_parent(tmp2, parent);
			__rb_rotate_set_parents(parent, sibling, root,
						RB_BLACK);
323
			break;
324
		} else {
325 326 327 328 329 330 331 332 333
			sibling = parent->rb_left;
			if (rb_is_red(sibling)) {
				/* Case 1 - right rotate at parent */
				parent->rb_left = tmp1 = sibling->rb_right;
				sibling->rb_right = parent;
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				sibling = tmp1;
L
Linus Torvalds 已提交
334
			}
335 336 337 338 339 340 341
			tmp1 = sibling->rb_left;
			if (!tmp1 || rb_is_black(tmp1)) {
				tmp2 = sibling->rb_right;
				if (!tmp2 || rb_is_black(tmp2)) {
					/* Case 2 - sibling color flip */
					rb_set_parent_color(sibling, parent,
							    RB_RED);
342 343 344
					node = parent;
					parent = rb_parent(node);
					continue;
L
Linus Torvalds 已提交
345
				}
346 347 348 349 350 351 352 353 354
				/* Case 3 - right rotate at sibling */
				sibling->rb_right = tmp1 = tmp2->rb_left;
				tmp2->rb_left = sibling;
				parent->rb_left = tmp2;
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				tmp1 = sibling;
				sibling = tmp2;
L
Linus Torvalds 已提交
355
			}
356 357 358 359 360 361 362 363
			/* Case 4 - left rotate at parent + color flips */
			parent->rb_left = tmp2 = sibling->rb_right;
			sibling->rb_right = parent;
			rb_set_parent_color(tmp1, sibling, RB_BLACK);
			if (tmp2)
				rb_set_parent(tmp2, parent);
			__rb_rotate_set_parents(parent, sibling, root,
						RB_BLACK);
364
			break;
L
Linus Torvalds 已提交
365 366 367 368 369 370
		}
	}
}

void rb_erase(struct rb_node *node, struct rb_root *root)
{
371 372
	struct rb_node *child = node->rb_right, *tmp = node->rb_left;
	struct rb_node *parent;
L
Linus Torvalds 已提交
373 374
	int color;

375 376 377 378 379 380 381 382 383 384 385 386 387 388 389
	if (!tmp) {
	case1:
		/* Case 1: node to erase has no more than 1 child (easy!) */

		parent = rb_parent(node);
		color = rb_color(node);

		if (child)
			rb_set_parent(child, parent);
		__rb_change_child(node, child, parent, root);
	} else if (!child) {
		/* Still case 1, but this time the child is node->rb_left */
		child = tmp;
		goto case1;
	} else {
L
Linus Torvalds 已提交
390 391
		struct rb_node *old = node, *left;

392
		node = child;
L
Linus Torvalds 已提交
393 394
		while ((left = node->rb_left) != NULL)
			node = left;
395

396
		__rb_change_child(old, node, rb_parent(old), root);
397

L
Linus Torvalds 已提交
398
		child = node->rb_right;
399
		parent = rb_parent(node);
400
		color = rb_color(node);
L
Linus Torvalds 已提交
401

402
		if (parent == old) {
L
Linus Torvalds 已提交
403
			parent = node;
404 405 406
		} else {
			if (child)
				rb_set_parent(child, parent);
407
			parent->rb_left = child;
408 409 410

			node->rb_right = old->rb_right;
			rb_set_parent(old->rb_right, node);
411
		}
412

413
		node->__rb_parent_color = old->__rb_parent_color;
L
Linus Torvalds 已提交
414
		node->rb_left = old->rb_left;
415
		rb_set_parent(old->rb_left, node);
L
Linus Torvalds 已提交
416 417 418 419 420 421 422
	}

	if (color == RB_BLACK)
		__rb_erase_color(child, parent, root);
}
EXPORT_SYMBOL(rb_erase);

423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454
static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
{
	struct rb_node *parent;

up:
	func(node, data);
	parent = rb_parent(node);
	if (!parent)
		return;

	if (node == parent->rb_left && parent->rb_right)
		func(parent->rb_right, data);
	else if (parent->rb_left)
		func(parent->rb_left, data);

	node = parent;
	goto up;
}

/*
 * after inserting @node into the tree, update the tree to account for
 * both the new entry and any damage done by rebalance
 */
void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
{
	if (node->rb_left)
		node = node->rb_left;
	else if (node->rb_right)
		node = node->rb_right;

	rb_augment_path(node, func, data);
}
455
EXPORT_SYMBOL(rb_augment_insert);
456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480

/*
 * before removing the node, find the deepest node on the rebalance path
 * that will still be there after @node gets removed
 */
struct rb_node *rb_augment_erase_begin(struct rb_node *node)
{
	struct rb_node *deepest;

	if (!node->rb_right && !node->rb_left)
		deepest = rb_parent(node);
	else if (!node->rb_right)
		deepest = node->rb_left;
	else if (!node->rb_left)
		deepest = node->rb_right;
	else {
		deepest = rb_next(node);
		if (deepest->rb_right)
			deepest = deepest->rb_right;
		else if (rb_parent(deepest) != node)
			deepest = rb_parent(deepest);
	}

	return deepest;
}
481
EXPORT_SYMBOL(rb_augment_erase_begin);
482 483 484 485 486 487 488 489 490 491

/*
 * after removal, update the tree to account for the removed entry
 * and any rebalance damage.
 */
void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
{
	if (node)
		rb_augment_path(node, func, data);
}
492
EXPORT_SYMBOL(rb_augment_erase_end);
493

L
Linus Torvalds 已提交
494 495 496
/*
 * This function returns the first node (in sort order) of the tree.
 */
497
struct rb_node *rb_first(const struct rb_root *root)
L
Linus Torvalds 已提交
498 499 500 501 502 503 504 505 506 507 508 509
{
	struct rb_node	*n;

	n = root->rb_node;
	if (!n)
		return NULL;
	while (n->rb_left)
		n = n->rb_left;
	return n;
}
EXPORT_SYMBOL(rb_first);

510
struct rb_node *rb_last(const struct rb_root *root)
L
Linus Torvalds 已提交
511 512 513 514 515 516 517 518 519 520 521 522
{
	struct rb_node	*n;

	n = root->rb_node;
	if (!n)
		return NULL;
	while (n->rb_right)
		n = n->rb_right;
	return n;
}
EXPORT_SYMBOL(rb_last);

523
struct rb_node *rb_next(const struct rb_node *node)
L
Linus Torvalds 已提交
524
{
525 526
	struct rb_node *parent;

527
	if (RB_EMPTY_NODE(node))
528 529
		return NULL;

530 531 532 533
	/*
	 * If we have a right-hand child, go down and then left as far
	 * as we can.
	 */
L
Linus Torvalds 已提交
534 535 536 537
	if (node->rb_right) {
		node = node->rb_right; 
		while (node->rb_left)
			node=node->rb_left;
538
		return (struct rb_node *)node;
L
Linus Torvalds 已提交
539 540
	}

541 542 543 544 545 546 547
	/*
	 * No right-hand children. Everything down and left is smaller than us,
	 * so any 'next' node must be in the general direction of our parent.
	 * Go up the tree; any time the ancestor is a right-hand child of its
	 * parent, keep going up. First time it's a left-hand child of its
	 * parent, said parent is our 'next' node.
	 */
548 549
	while ((parent = rb_parent(node)) && node == parent->rb_right)
		node = parent;
L
Linus Torvalds 已提交
550

551
	return parent;
L
Linus Torvalds 已提交
552 553 554
}
EXPORT_SYMBOL(rb_next);

555
struct rb_node *rb_prev(const struct rb_node *node)
L
Linus Torvalds 已提交
556
{
557 558
	struct rb_node *parent;

559
	if (RB_EMPTY_NODE(node))
560 561
		return NULL;

562 563 564 565
	/*
	 * If we have a left-hand child, go down and then right as far
	 * as we can.
	 */
L
Linus Torvalds 已提交
566 567 568 569
	if (node->rb_left) {
		node = node->rb_left; 
		while (node->rb_right)
			node=node->rb_right;
570
		return (struct rb_node *)node;
L
Linus Torvalds 已提交
571 572
	}

573 574 575 576
	/*
	 * No left-hand children. Go up till we find an ancestor which
	 * is a right-hand child of its parent.
	 */
577 578
	while ((parent = rb_parent(node)) && node == parent->rb_left)
		node = parent;
L
Linus Torvalds 已提交
579

580
	return parent;
L
Linus Torvalds 已提交
581 582 583 584 585 586
}
EXPORT_SYMBOL(rb_prev);

void rb_replace_node(struct rb_node *victim, struct rb_node *new,
		     struct rb_root *root)
{
587
	struct rb_node *parent = rb_parent(victim);
L
Linus Torvalds 已提交
588 589

	/* Set the surrounding nodes to point to the replacement */
590
	__rb_change_child(victim, new, parent, root);
L
Linus Torvalds 已提交
591
	if (victim->rb_left)
592
		rb_set_parent(victim->rb_left, new);
L
Linus Torvalds 已提交
593
	if (victim->rb_right)
594
		rb_set_parent(victim->rb_right, new);
L
Linus Torvalds 已提交
595 596 597 598 599

	/* Copy the pointers/colour from the victim to the replacement */
	*new = *victim;
}
EXPORT_SYMBOL(rb_replace_node);