[thirdparty/linux.git] / lib / rbtree.c

/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  (C) 2002  David Woodhouse <dwmw2@infradead.org>
  (C) 2012  Michel Lespinasse <walken@google.com>

  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_augmented.h>
#include <linux/export.h>

/*
 * 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
 *  nodes will be lowercase. Unknown color nodes shall be drawn as red within
 *  parentheses and have some accompanying text comment.
 */

/*
 * Notes on lockless lookups:
 *
 * All stores to the tree structure (rb_left and rb_right) must be done using
 * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
 * tree structure as seen in program order.
 *
 * These two requirements will allow lockless iteration of the tree -- not
 * correct iteration mind you, tree rotations are not atomic so a lookup might
 * miss entire subtrees.
 *
 * But they do guarantee that any such traversal will only see valid elements
 * and that it will indeed complete -- does not get stuck in a loop.
 *
 * It also guarantees that if the lookup returns an element it is the 'correct'
 * one. But not returning an element does _NOT_ mean it's not present.
 *
 * NOTE:
 *
 * Stores to __rb_parent_color are not important for simple lookups so those
 * are left undone as of now. Nor did I check for loops involving parent
 * pointers.
 */

static inline void rb_set_black(struct rb_node *rb)
{
	rb->__rb_parent_color |= RB_BLACK;
}

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

/*
 * 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);
	__rb_change_child(old, new, parent, root);
}

static __always_inline void
__rb_insert(struct rb_node *node, struct rb_root *root,
	    bool newleft, struct rb_node **leftmost,
	    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;

	if (newleft)
		*leftmost = node;

	while (true) {
		/*
		 * Loop invariant: node is red.
		 */
		if (unlikely(!parent)) {
			/*
			 * The inserted node is root. Either this is the
			 * first node, or we recursed at Case 1 below and
			 * are no longer violating 4).
			 */
			rb_set_parent_color(node, NULL, RB_BLACK);
			break;
		}

		/*
		 * If there is a black parent, we are done.
		 * Otherwise, take some corrective action as,
		 * per 4), we don't want a red root or two
		 * consecutive red nodes.
		 */
		if(rb_is_black(parent))
			break;

		gparent = rb_red_parent(parent);

		tmp = gparent->rb_right;
		if (parent != tmp) {	/* parent == gparent->rb_left */
			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;
			}

			tmp = parent->rb_right;
			if (node == tmp) {
				/*
				 * 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.
				 */
				tmp = node->rb_left;
				WRITE_ONCE(parent->rb_right, tmp);
				WRITE_ONCE(node->rb_left, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_right;
			}

			/*
			 * Case 3 - right rotate at gparent
			 *
			 *        G           P
			 *       / \         / \
			 *      p   U  -->  n   g
			 *     /                 \
			 *    n                   U
			 */
			WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
			WRITE_ONCE(parent->rb_right, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		} else {
			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;
			}

			tmp = parent->rb_left;
			if (node == tmp) {
				/* Case 2 - right rotate at parent */
				tmp = node->rb_right;
				WRITE_ONCE(parent->rb_left, tmp);
				WRITE_ONCE(node->rb_right, parent);
				if (tmp)
					rb_set_parent_color(tmp, parent,
							    RB_BLACK);
				rb_set_parent_color(parent, node, RB_RED);
				augment_rotate(parent, node);
				parent = node;
				tmp = node->rb_left;
			}

			/* Case 3 - left rotate at gparent */
			WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
			WRITE_ONCE(parent->rb_left, gparent);
			if (tmp)
				rb_set_parent_color(tmp, gparent, RB_BLACK);
			__rb_rotate_set_parents(gparent, parent, root, RB_RED);
			augment_rotate(gparent, parent);
			break;
		}
	}
}

/*
 * Inline version for rb_erase() use - we want to be able to inline
 * and eliminate the dummy_rotate callback there
 */
static __always_inline void
____rb_erase_color(struct rb_node *parent, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;

	while (true) {
		/*
		 * Loop invariants:
		 * - node is black (or NULL on first iteration)
		 * - node is not the root (parent is not NULL)
		 * - All leaf paths going through parent and node have a
		 *   black node count that is 1 lower than other leaf paths.
		 */
		sibling = parent->rb_right;
		if (node != sibling) {	/* node == parent->rb_left */
			if (rb_is_red(sibling)) {
				/*
				 * Case 1 - left rotate at parent
				 *
				 *     P               S
				 *    / \             / \
				 *   N   s    -->    p   Sr
				 *      / \         / \
				 *     Sl  Sr      N   Sl
				 */
				tmp1 = sibling->rb_left;
				WRITE_ONCE(parent->rb_right, tmp1);
				WRITE_ONCE(sibling->rb_left, parent);
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				augment_rotate(parent, sibling);
				sibling = tmp1;
			}
			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) which
					 * can be fixed by flipping p to black
					 * if it was red, or by recursing at p.
					 * p is red when coming from Case 1.
					 */
					rb_set_parent_color(sibling, parent,
							    RB_RED);
					if (rb_is_red(parent))
						rb_set_black(parent);
					else {
						node = parent;
						parent = rb_parent(node);
						if (parent)
							continue;
					}
					break;
				}
				/*
				 * Case 3 - right rotate at sibling
				 * (p could be either color here)
				 *
				 *   (p)           (p)
				 *   / \           / \
				 *  N   S    -->  N   sl
				 *     / \             \
				 *    sl  Sr            S
				 *                       \
				 *                        Sr
				 *
				 * Note: p might be red, and then both
				 * p and sl are red after rotation(which
				 * breaks property 4). This is fixed in
				 * Case 4 (in __rb_rotate_set_parents()
				 *         which set sl the color of p
				 *         and set p RB_BLACK)
				 *
				 *   (p)            (sl)
				 *   / \            /  \
				 *  N   sl   -->   P    S
				 *       \        /      \
				 *        S      N        Sr
				 *         \
				 *          Sr
				 */
				tmp1 = tmp2->rb_right;
				WRITE_ONCE(sibling->rb_left, tmp1);
				WRITE_ONCE(tmp2->rb_right, sibling);
				WRITE_ONCE(parent->rb_right, tmp2);
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				augment_rotate(sibling, tmp2);
				tmp1 = sibling;
				sibling = tmp2;
			}
			/*
			 * 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)
			 */
			tmp2 = sibling->rb_left;
			WRITE_ONCE(parent->rb_right, tmp2);
			WRITE_ONCE(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);
			augment_rotate(parent, sibling);
			break;
		} else {
			sibling = parent->rb_left;
			if (rb_is_red(sibling)) {
				/* Case 1 - right rotate at parent */
				tmp1 = sibling->rb_right;
				WRITE_ONCE(parent->rb_left, tmp1);
				WRITE_ONCE(sibling->rb_right, parent);
				rb_set_parent_color(tmp1, parent, RB_BLACK);
				__rb_rotate_set_parents(parent, sibling, root,
							RB_RED);
				augment_rotate(parent, sibling);
				sibling = tmp1;
			}
			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);
					if (rb_is_red(parent))
						rb_set_black(parent);
					else {
						node = parent;
						parent = rb_parent(node);
						if (parent)
							continue;
					}
					break;
				}
				/* Case 3 - left rotate at sibling */
				tmp1 = tmp2->rb_left;
				WRITE_ONCE(sibling->rb_right, tmp1);
				WRITE_ONCE(tmp2->rb_left, sibling);
				WRITE_ONCE(parent->rb_left, tmp2);
				if (tmp1)
					rb_set_parent_color(tmp1, sibling,
							    RB_BLACK);
				augment_rotate(sibling, tmp2);
				tmp1 = sibling;
				sibling = tmp2;
			}
			/* Case 4 - right rotate at parent + color flips */
			tmp2 = sibling->rb_right;
			WRITE_ONCE(parent->rb_left, tmp2);
			WRITE_ONCE(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);
			augment_rotate(parent, sibling);
			break;
		}
	}
}

/* Non-inline version for rb_erase_augmented() use */
void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	____rb_erase_color(parent, root, augment_rotate);
}
EXPORT_SYMBOL(__rb_erase_color);

/*
 * Non-augmented rbtree manipulation functions.
 *
 * We use dummy augmented callbacks here, and have the compiler optimize them
 * out of the rb_insert_color() and rb_erase() function definitions.
 */

static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}

static const struct rb_augment_callbacks dummy_callbacks = {
	.propagate = dummy_propagate,
	.copy = dummy_copy,
	.rotate = dummy_rotate
};

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
	__rb_insert(node, root, false, NULL, dummy_rotate);
}
EXPORT_SYMBOL(rb_insert_color);

void rb_erase(struct rb_node *node, struct rb_root *root)
{
	struct rb_node *rebalance;
	rebalance = __rb_erase_augmented(node, root,
					 NULL, &dummy_callbacks);
	if (rebalance)
		____rb_erase_color(rebalance, root, dummy_rotate);
}
EXPORT_SYMBOL(rb_erase);

void rb_insert_color_cached(struct rb_node *node,
			    struct rb_root_cached *root, bool leftmost)
{
	__rb_insert(node, &root->rb_root, leftmost,
		    &root->rb_leftmost, dummy_rotate);
}
EXPORT_SYMBOL(rb_insert_color_cached);

void rb_erase_cached(struct rb_node *node, struct rb_root_cached *root)
{
	struct rb_node *rebalance;
	rebalance = __rb_erase_augmented(node, &root->rb_root,
					 &root->rb_leftmost, &dummy_callbacks);
	if (rebalance)
		____rb_erase_color(rebalance, &root->rb_root, dummy_rotate);
}
EXPORT_SYMBOL(rb_erase_cached);

/*
 * Augmented rbtree manipulation functions.
 *
 * This instantiates the same __always_inline functions as in the non-augmented
 * case, but this time with user-defined callbacks.
 */

void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
			   bool newleft, struct rb_node **leftmost,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	__rb_insert(node, root, newleft, leftmost, augment_rotate);
}
EXPORT_SYMBOL(__rb_insert_augmented);

/*
 * This function returns the first node (in sort order) of the tree.
 */
struct rb_node *rb_first(const struct rb_root *root)
{
	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);

struct rb_node *rb_last(const struct rb_root *root)
{
	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);

struct rb_node *rb_next(const struct rb_node *node)
{
	struct rb_node *parent;

	if (RB_EMPTY_NODE(node))
		return NULL;

	/*
	 * If we have a right-hand child, go down and then left as far
	 * as we can.
	 */
	if (node->rb_right) {
		node = node->rb_right;
		while (node->rb_left)
			node=node->rb_left;
		return (struct rb_node *)node;
	}

	/*
	 * 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.
	 */
	while ((parent = rb_parent(node)) && node == parent->rb_right)
		node = parent;

	return parent;
}
EXPORT_SYMBOL(rb_next);

struct rb_node *rb_prev(const struct rb_node *node)
{
	struct rb_node *parent;

	if (RB_EMPTY_NODE(node))
		return NULL;

	/*
	 * If we have a left-hand child, go down and then right as far
	 * as we can.
	 */
	if (node->rb_left) {
		node = node->rb_left;
		while (node->rb_right)
			node=node->rb_right;
		return (struct rb_node *)node;
	}

	/*
	 * No left-hand children. Go up till we find an ancestor which
	 * is a right-hand child of its parent.
	 */
	while ((parent = rb_parent(node)) && node == parent->rb_left)
		node = parent;

	return parent;
}
EXPORT_SYMBOL(rb_prev);

void rb_replace_node(struct rb_node *victim, struct rb_node *new,
		     struct rb_root *root)
{
	struct rb_node *parent = rb_parent(victim);

	/* Copy the pointers/colour from the victim to the replacement */
	*new = *victim;

	/* Set the surrounding nodes to point to the replacement */
	if (victim->rb_left)
		rb_set_parent(victim->rb_left, new);
	if (victim->rb_right)
		rb_set_parent(victim->rb_right, new);
	__rb_change_child(victim, new, parent, root);
}
EXPORT_SYMBOL(rb_replace_node);

void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new,
			 struct rb_root *root)
{
	struct rb_node *parent = rb_parent(victim);

	/* Copy the pointers/colour from the victim to the replacement */
	*new = *victim;

	/* Set the surrounding nodes to point to the replacement */
	if (victim->rb_left)
		rb_set_parent(victim->rb_left, new);
	if (victim->rb_right)
		rb_set_parent(victim->rb_right, new);

	/* Set the parent's pointer to the new node last after an RCU barrier
	 * so that the pointers onwards are seen to be set correctly when doing
	 * an RCU walk over the tree.
	 */
	__rb_change_child_rcu(victim, new, parent, root);
}
EXPORT_SYMBOL(rb_replace_node_rcu);

static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
{
	for (;;) {
		if (node->rb_left)
			node = node->rb_left;
		else if (node->rb_right)
			node = node->rb_right;
		else
			return (struct rb_node *)node;
	}
}

struct rb_node *rb_next_postorder(const struct rb_node *node)
{
	const struct rb_node *parent;
	if (!node)
		return NULL;
	parent = rb_parent(node);

	/* If we're sitting on node, we've already seen our children */
	if (parent && node == parent->rb_left && parent->rb_right) {
		/* If we are the parent's left node, go to the parent's right
		 * node then all the way down to the left */
		return rb_left_deepest_node(parent->rb_right);
	} else
		/* Otherwise we are the parent's right node, and the parent
		 * should be next */
		return (struct rb_node *)parent;
}
EXPORT_SYMBOL(rb_next_postorder);

struct rb_node *rb_first_postorder(const struct rb_root *root)
{
	if (!root->rb_node)
		return NULL;

	return rb_left_deepest_node(root->rb_node);
}
EXPORT_SYMBOL(rb_first_postorder);
Commit	Line	Data
1da177e4 LT	1	/*
	2	Red Black Trees
	3	(C) 1999 Andrea Arcangeli <andrea@suse.de>
	4	(C) 2002 David Woodhouse <dwmw2@infradead.org>
46b6135a ML	5	(C) 2012 Michel Lespinasse <walken@google.com>
46b6135a ML	6
1da177e4 LT	7	This program is free software; you can redistribute it and/or modify
	8	it under the terms of the GNU General Public License as published by
	9	the Free Software Foundation; either version 2 of the License, or
	10	(at your option) any later version.
	11
	12	This program is distributed in the hope that it will be useful,
	13	but WITHOUT ANY WARRANTY; without even the implied warranty of
	14	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
	15	GNU General Public License for more details.
	16
	17	You should have received a copy of the GNU General Public License
	18	along with this program; if not, write to the Free Software
	19	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
	20
	21	linux/lib/rbtree.c
	22	*/
	23
9c079add	24	#include <linux/rbtree_augmented.h>
8bc3bcc9	25	#include <linux/export.h>
1da177e4	26
5bc9188a ML	27	/*
	28	* red-black trees properties: http://en.wikipedia.org/wiki/Rbtree
	29	*
	30	* 1) A node is either red or black
	31	* 2) The root is black
	32	* 3) All leaves (NULL) are black
	33	* 4) Both children of every red node are black
	34	* 5) Every simple path from root to leaves contains the same number
	35	* of black nodes.
	36	*
	37	* 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
	38	* consecutive red nodes in a path and every red node is therefore followed by
	39	* a black. So if B is the number of black nodes on every simple path (as per
	40	* 5), then the longest possible path due to 4 is 2B.
	41	*
	42	* We shall indicate color with case, where black nodes are uppercase and red
6280d235 ML	43	* nodes will be lowercase. Unknown color nodes shall be drawn as red within
6280d235 ML	44	* parentheses and have some accompanying text comment.
5bc9188a ML	45	*/
5bc9188a ML	46
d72da4a4 PZ	47	/*
	48	* Notes on lockless lookups:
	49	*
	50	* All stores to the tree structure (rb_left and rb_right) must be done using
	51	* WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
	52	* tree structure as seen in program order.
	53	*
	54	* These two requirements will allow lockless iteration of the tree -- not
	55	* correct iteration mind you, tree rotations are not atomic so a lookup might
	56	* miss entire subtrees.
	57	*
	58	* But they do guarantee that any such traversal will only see valid elements
	59	* and that it will indeed complete -- does not get stuck in a loop.
	60	*
	61	* It also guarantees that if the lookup returns an element it is the 'correct'
	62	* one. But not returning an element does _NOT_ mean it's not present.
	63	*
	64	* NOTE:
	65	*
	66	* Stores to __rb_parent_color are not important for simple lookups so those
	67	* are left undone as of now. Nor did I check for loops involving parent
	68	* pointers.
	69	*/
	70
46b6135a ML	71	static inline void rb_set_black(struct rb_node *rb)
	72	{
	73	rb->__rb_parent_color \|= RB_BLACK;
	74	}
	75
5bc9188a ML	76	static inline struct rb_node rb_red_parent(struct rb_node red)
	77	{
	78	return (struct rb_node *)red->__rb_parent_color;
	79	}
	80
5bc9188a ML	81	/*
	82	* Helper function for rotations:
	83	* - old's parent and color get assigned to new
	84	* - old gets assigned new as a parent and 'color' as a color.
	85	*/
	86	static inline void
	87	__rb_rotate_set_parents(struct rb_node old, struct rb_node new,
	88	struct rb_root *root, int color)
	89	{
	90	struct rb_node *parent = rb_parent(old);
	91	new->__rb_parent_color = old->__rb_parent_color;
	92	rb_set_parent_color(old, new, color);
7abc704a	93	__rb_change_child(old, new, parent, root);
5bc9188a ML	94	}
5bc9188a ML	95
14b94af0 ML	96	static __always_inline void
14b94af0 ML	97	__rb_insert(struct rb_node node, struct rb_root root,
cd9e61ed	98	bool newleft, struct rb_node **leftmost,
14b94af0	99	void (augment_rotate)(struct rb_node old, struct rb_node *new))
1da177e4	100	{
5bc9188a	101	struct rb_node parent = rb_red_parent(node), gparent, *tmp;
1da177e4	102
cd9e61ed DB	103	if (newleft)
	104	*leftmost = node;
	105
6d58452d ML	106	while (true) {
6d58452d ML	107	/*
2aadf7fc	108	* Loop invariant: node is red.
6d58452d	109	*/
2aadf7fc DB	110	if (unlikely(!parent)) {
	111	/*
	112	* The inserted node is root. Either this is the
	113	* first node, or we recursed at Case 1 below and
	114	* are no longer violating 4).
	115	*/
5bc9188a	116	rb_set_parent_color(node, NULL, RB_BLACK);
6d58452d	117	break;
2aadf7fc DB	118	}
	119
	120	/*
	121	* If there is a black parent, we are done.
	122	* Otherwise, take some corrective action as,
	123	* per 4), we don't want a red root or two
	124	* consecutive red nodes.
	125	*/
	126	if(rb_is_black(parent))
6d58452d ML	127	break;
6d58452d ML	128
5bc9188a ML	129	gparent = rb_red_parent(parent);
5bc9188a ML	130
59633abf ML	131	tmp = gparent->rb_right;
59633abf ML	132	if (parent != tmp) { /* parent == gparent->rb_left */
5bc9188a ML	133	if (tmp && rb_is_red(tmp)) {
	134	/*
	135	* Case 1 - color flips
	136	*
	137	* G g
	138	* / \ / \
	139	* p u --> P U
	140	* / /
1b9c53e8	141	* n n
5bc9188a ML	142	*
	143	* However, since g's parent might be red, and
	144	* 4) does not allow this, we need to recurse
	145	* at g.
	146	*/
	147	rb_set_parent_color(tmp, gparent, RB_BLACK);
	148	rb_set_parent_color(parent, gparent, RB_BLACK);
	149	node = gparent;
	150	parent = rb_parent(node);
	151	rb_set_parent_color(node, parent, RB_RED);
	152	continue;
1da177e4 LT	153	}
1da177e4 LT	154
59633abf ML	155	tmp = parent->rb_right;
59633abf ML	156	if (node == tmp) {
5bc9188a ML	157	/*
	158	* Case 2 - left rotate at parent
	159	*
	160	* G G
	161	* / \ / \
	162	* p U --> n U
	163	* \ /
	164	* n p
	165	*
	166	* This still leaves us in violation of 4), the
	167	* continuation into Case 3 will fix that.
	168	*/
d72da4a4 PZ	169	tmp = node->rb_left;
	170	WRITE_ONCE(parent->rb_right, tmp);
	171	WRITE_ONCE(node->rb_left, parent);
5bc9188a ML	172	if (tmp)
	173	rb_set_parent_color(tmp, parent,
	174	RB_BLACK);
	175	rb_set_parent_color(parent, node, RB_RED);
14b94af0	176	augment_rotate(parent, node);
1da177e4	177	parent = node;
59633abf	178	tmp = node->rb_right;
1da177e4 LT	179	}
1da177e4 LT	180
5bc9188a ML	181	/*
	182	* Case 3 - right rotate at gparent
	183	*
	184	* G P
	185	* / \ / \
	186	* p U --> n g
	187	* / \
	188	* n U
	189	*/
d72da4a4 PZ	190	WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
d72da4a4 PZ	191	WRITE_ONCE(parent->rb_right, gparent);
5bc9188a ML	192	if (tmp)
	193	rb_set_parent_color(tmp, gparent, RB_BLACK);
	194	__rb_rotate_set_parents(gparent, parent, root, RB_RED);
14b94af0	195	augment_rotate(gparent, parent);
1f052865	196	break;
1da177e4	197	} else {
5bc9188a ML	198	tmp = gparent->rb_left;
	199	if (tmp && rb_is_red(tmp)) {
	200	/* Case 1 - color flips */
	201	rb_set_parent_color(tmp, gparent, RB_BLACK);
	202	rb_set_parent_color(parent, gparent, RB_BLACK);
	203	node = gparent;
	204	parent = rb_parent(node);
	205	rb_set_parent_color(node, parent, RB_RED);
	206	continue;
1da177e4 LT	207	}
1da177e4 LT	208
59633abf ML	209	tmp = parent->rb_left;
59633abf ML	210	if (node == tmp) {
5bc9188a	211	/* Case 2 - right rotate at parent */
d72da4a4 PZ	212	tmp = node->rb_right;
	213	WRITE_ONCE(parent->rb_left, tmp);
	214	WRITE_ONCE(node->rb_right, parent);
5bc9188a ML	215	if (tmp)
	216	rb_set_parent_color(tmp, parent,
	217	RB_BLACK);
	218	rb_set_parent_color(parent, node, RB_RED);
14b94af0	219	augment_rotate(parent, node);
1da177e4	220	parent = node;
59633abf	221	tmp = node->rb_left;
1da177e4 LT	222	}
1da177e4 LT	223
5bc9188a	224	/* Case 3 - left rotate at gparent */
d72da4a4 PZ	225	WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
d72da4a4 PZ	226	WRITE_ONCE(parent->rb_left, gparent);
5bc9188a ML	227	if (tmp)
	228	rb_set_parent_color(tmp, gparent, RB_BLACK);
	229	__rb_rotate_set_parents(gparent, parent, root, RB_RED);
14b94af0	230	augment_rotate(gparent, parent);
1f052865	231	break;
1da177e4 LT	232	}
1da177e4 LT	233	}
1da177e4	234	}
1da177e4	235
3cb7a563 ML	236	/*
	237	* Inline version for rb_erase() use - we want to be able to inline
	238	* and eliminate the dummy_rotate callback there
	239	*/
	240	static __always_inline void
	241	____rb_erase_color(struct rb_node parent, struct rb_root root,
9c079add	242	void (augment_rotate)(struct rb_node old, struct rb_node *new))
1da177e4	243	{
46b6135a	244	struct rb_node node = NULL, sibling, tmp1, tmp2;
1da177e4	245
d6ff1273 ML	246	while (true) {
d6ff1273 ML	247	/*
46b6135a ML	248	* Loop invariants:
	249	* - node is black (or NULL on first iteration)
	250	* - node is not the root (parent is not NULL)
	251	* - All leaf paths going through parent and node have a
	252	* black node count that is 1 lower than other leaf paths.
d6ff1273	253	*/
59633abf ML	254	sibling = parent->rb_right;
59633abf ML	255	if (node != sibling) { /* node == parent->rb_left */
6280d235 ML	256	if (rb_is_red(sibling)) {
	257	/*
	258	* Case 1 - left rotate at parent
	259	*
	260	* P S
	261	* / \ / \
	262	* N s --> p Sr
	263	* / \ / \
	264	* Sl Sr N Sl
	265	*/
d72da4a4 PZ	266	tmp1 = sibling->rb_left;
	267	WRITE_ONCE(parent->rb_right, tmp1);
	268	WRITE_ONCE(sibling->rb_left, parent);
6280d235 ML	269	rb_set_parent_color(tmp1, parent, RB_BLACK);
	270	__rb_rotate_set_parents(parent, sibling, root,
	271	RB_RED);
9c079add	272	augment_rotate(parent, sibling);
6280d235	273	sibling = tmp1;
1da177e4	274	}
6280d235 ML	275	tmp1 = sibling->rb_right;
	276	if (!tmp1 \|\| rb_is_black(tmp1)) {
	277	tmp2 = sibling->rb_left;
	278	if (!tmp2 \|\| rb_is_black(tmp2)) {
	279	/*
	280	* Case 2 - sibling color flip
	281	* (p could be either color here)
	282	*
	283	* (p) (p)
	284	* / \ / \
	285	* N S --> N s
	286	* / \ / \
	287	* Sl Sr Sl Sr
	288	*
46b6135a ML	289	* This leaves us violating 5) which
	290	* can be fixed by flipping p to black
	291	* if it was red, or by recursing at p.
	292	* p is red when coming from Case 1.
6280d235 ML	293	*/
	294	rb_set_parent_color(sibling, parent,
	295	RB_RED);
46b6135a ML	296	if (rb_is_red(parent))
	297	rb_set_black(parent);
	298	else {
	299	node = parent;
	300	parent = rb_parent(node);
	301	if (parent)
	302	continue;
	303	}
	304	break;
1da177e4	305	}
6280d235 ML	306	/*
	307	* Case 3 - right rotate at sibling
	308	* (p could be either color here)
	309	*
	310	* (p) (p)
	311	* / \ / \
ce093a04	312	* N S --> N sl
6280d235	313	* / \ \
ce093a04	314	* sl Sr S
6280d235 ML	315	* \
6280d235 ML	316	* Sr
ce093a04 JC	317	*
	318	* Note: p might be red, and then both
	319	* p and sl are red after rotation(which
	320	* breaks property 4). This is fixed in
	321	* Case 4 (in __rb_rotate_set_parents()
	322	* which set sl the color of p
	323	* and set p RB_BLACK)
	324	*
	325	* (p) (sl)
	326	* / \ / \
	327	* N sl --> P S
	328	* \ / \
	329	* S N Sr
	330	* \
	331	* Sr
6280d235	332	*/
d72da4a4 PZ	333	tmp1 = tmp2->rb_right;
	334	WRITE_ONCE(sibling->rb_left, tmp1);
	335	WRITE_ONCE(tmp2->rb_right, sibling);
	336	WRITE_ONCE(parent->rb_right, tmp2);
6280d235 ML	337	if (tmp1)
	338	rb_set_parent_color(tmp1, sibling,
	339	RB_BLACK);
9c079add	340	augment_rotate(sibling, tmp2);
6280d235 ML	341	tmp1 = sibling;
6280d235 ML	342	sibling = tmp2;
1da177e4	343	}
6280d235 ML	344	/*
	345	* Case 4 - left rotate at parent + color flips
	346	* (p and sl could be either color here.
	347	* After rotation, p becomes black, s acquires
	348	* p's color, and sl keeps its color)
	349	*
	350	* (p) (s)
	351	* / \ / \
	352	* N S --> P Sr
	353	* / \ / \
	354	* (sl) sr N (sl)
	355	*/
d72da4a4 PZ	356	tmp2 = sibling->rb_left;
	357	WRITE_ONCE(parent->rb_right, tmp2);
	358	WRITE_ONCE(sibling->rb_left, parent);
6280d235 ML	359	rb_set_parent_color(tmp1, sibling, RB_BLACK);
	360	if (tmp2)
	361	rb_set_parent(tmp2, parent);
	362	__rb_rotate_set_parents(parent, sibling, root,
	363	RB_BLACK);
9c079add	364	augment_rotate(parent, sibling);
e125d147	365	break;
d6ff1273	366	} else {
6280d235 ML	367	sibling = parent->rb_left;
	368	if (rb_is_red(sibling)) {
	369	/* Case 1 - right rotate at parent */
d72da4a4 PZ	370	tmp1 = sibling->rb_right;
	371	WRITE_ONCE(parent->rb_left, tmp1);
	372	WRITE_ONCE(sibling->rb_right, parent);
6280d235 ML	373	rb_set_parent_color(tmp1, parent, RB_BLACK);
	374	__rb_rotate_set_parents(parent, sibling, root,
	375	RB_RED);
9c079add	376	augment_rotate(parent, sibling);
6280d235	377	sibling = tmp1;
1da177e4	378	}
6280d235 ML	379	tmp1 = sibling->rb_left;
	380	if (!tmp1 \|\| rb_is_black(tmp1)) {
	381	tmp2 = sibling->rb_right;
	382	if (!tmp2 \|\| rb_is_black(tmp2)) {
	383	/* Case 2 - sibling color flip */
	384	rb_set_parent_color(sibling, parent,
	385	RB_RED);
46b6135a ML	386	if (rb_is_red(parent))
	387	rb_set_black(parent);
	388	else {
	389	node = parent;
	390	parent = rb_parent(node);
	391	if (parent)
	392	continue;
	393	}
	394	break;
1da177e4	395	}
ce093a04	396	/* Case 3 - left rotate at sibling */
d72da4a4 PZ	397	tmp1 = tmp2->rb_left;
	398	WRITE_ONCE(sibling->rb_right, tmp1);
	399	WRITE_ONCE(tmp2->rb_left, sibling);
	400	WRITE_ONCE(parent->rb_left, tmp2);
6280d235 ML	401	if (tmp1)
	402	rb_set_parent_color(tmp1, sibling,
	403	RB_BLACK);
9c079add	404	augment_rotate(sibling, tmp2);
6280d235 ML	405	tmp1 = sibling;
6280d235 ML	406	sibling = tmp2;
1da177e4	407	}
ce093a04	408	/* Case 4 - right rotate at parent + color flips */
d72da4a4 PZ	409	tmp2 = sibling->rb_right;
	410	WRITE_ONCE(parent->rb_left, tmp2);
	411	WRITE_ONCE(sibling->rb_right, parent);
6280d235 ML	412	rb_set_parent_color(tmp1, sibling, RB_BLACK);
	413	if (tmp2)
	414	rb_set_parent(tmp2, parent);
	415	__rb_rotate_set_parents(parent, sibling, root,
	416	RB_BLACK);
9c079add	417	augment_rotate(parent, sibling);
e125d147	418	break;
1da177e4 LT	419	}
1da177e4 LT	420	}
1da177e4	421	}
3cb7a563 ML	422
	423	/* Non-inline version for rb_erase_augmented() use */
	424	void __rb_erase_color(struct rb_node parent, struct rb_root root,
	425	void (augment_rotate)(struct rb_node old, struct rb_node *new))
	426	{
	427	____rb_erase_color(parent, root, augment_rotate);
	428	}
9c079add	429	EXPORT_SYMBOL(__rb_erase_color);
14b94af0 ML	430
	431	/*
	432	* Non-augmented rbtree manipulation functions.
	433	*
	434	* We use dummy augmented callbacks here, and have the compiler optimize them
	435	* out of the rb_insert_color() and rb_erase() function definitions.
	436	*/
	437
	438	static inline void dummy_propagate(struct rb_node node, struct rb_node stop) {}
	439	static inline void dummy_copy(struct rb_node old, struct rb_node new) {}
	440	static inline void dummy_rotate(struct rb_node old, struct rb_node new) {}
	441
	442	static const struct rb_augment_callbacks dummy_callbacks = {
f231aebf KC	443	.propagate = dummy_propagate,
	444	.copy = dummy_copy,
	445	.rotate = dummy_rotate
14b94af0 ML	446	};
	447
	448	void rb_insert_color(struct rb_node node, struct rb_root root)
	449	{
cd9e61ed	450	__rb_insert(node, root, false, NULL, dummy_rotate);
14b94af0 ML	451	}
	452	EXPORT_SYMBOL(rb_insert_color);
	453
	454	void rb_erase(struct rb_node node, struct rb_root root)
	455	{
3cb7a563	456	struct rb_node *rebalance;
cd9e61ed DB	457	rebalance = __rb_erase_augmented(node, root,
cd9e61ed DB	458	NULL, &dummy_callbacks);
3cb7a563 ML	459	if (rebalance)
3cb7a563 ML	460	____rb_erase_color(rebalance, root, dummy_rotate);
1da177e4 LT	461	}
	462	EXPORT_SYMBOL(rb_erase);
	463
cd9e61ed DB	464	void rb_insert_color_cached(struct rb_node *node,
	465	struct rb_root_cached *root, bool leftmost)
	466	{
	467	__rb_insert(node, &root->rb_root, leftmost,
	468	&root->rb_leftmost, dummy_rotate);
	469	}
	470	EXPORT_SYMBOL(rb_insert_color_cached);
	471
	472	void rb_erase_cached(struct rb_node node, struct rb_root_cached root)
	473	{
	474	struct rb_node *rebalance;
	475	rebalance = __rb_erase_augmented(node, &root->rb_root,
	476	&root->rb_leftmost, &dummy_callbacks);
	477	if (rebalance)
	478	____rb_erase_color(rebalance, &root->rb_root, dummy_rotate);
	479	}
	480	EXPORT_SYMBOL(rb_erase_cached);
	481
14b94af0 ML	482	/*
	483	* Augmented rbtree manipulation functions.
	484	*
	485	* This instantiates the same __always_inline functions as in the non-augmented
	486	* case, but this time with user-defined callbacks.
	487	*/
	488
	489	void __rb_insert_augmented(struct rb_node node, struct rb_root root,
cd9e61ed	490	bool newleft, struct rb_node **leftmost,
14b94af0 ML	491	void (augment_rotate)(struct rb_node old, struct rb_node *new))
14b94af0 ML	492	{
cd9e61ed	493	__rb_insert(node, root, newleft, leftmost, augment_rotate);
14b94af0 ML	494	}
	495	EXPORT_SYMBOL(__rb_insert_augmented);
	496
1da177e4 LT	497	/*
	498	* This function returns the first node (in sort order) of the tree.
	499	*/
f4b477c4	500	struct rb_node rb_first(const struct rb_root root)
1da177e4 LT	501	{
	502	struct rb_node *n;
	503
	504	n = root->rb_node;
	505	if (!n)
	506	return NULL;
	507	while (n->rb_left)
	508	n = n->rb_left;
	509	return n;
	510	}
	511	EXPORT_SYMBOL(rb_first);
	512
f4b477c4	513	struct rb_node rb_last(const struct rb_root root)
1da177e4 LT	514	{
	515	struct rb_node *n;
	516
	517	n = root->rb_node;
	518	if (!n)
	519	return NULL;
	520	while (n->rb_right)
	521	n = n->rb_right;
	522	return n;
	523	}
	524	EXPORT_SYMBOL(rb_last);
	525
f4b477c4	526	struct rb_node rb_next(const struct rb_node node)
1da177e4	527	{
55a98102 DW	528	struct rb_node *parent;
55a98102 DW	529
4c199a93	530	if (RB_EMPTY_NODE(node))
10fd48f2 JA	531	return NULL;
10fd48f2 JA	532
7ce6ff9e ML	533	/*
	534	* If we have a right-hand child, go down and then left as far
	535	* as we can.
	536	*/
1da177e4	537	if (node->rb_right) {
cd9e61ed	538	node = node->rb_right;
1da177e4 LT	539	while (node->rb_left)
1da177e4 LT	540	node=node->rb_left;
f4b477c4	541	return (struct rb_node *)node;
1da177e4 LT	542	}
1da177e4 LT	543
7ce6ff9e ML	544	/*
	545	* No right-hand children. Everything down and left is smaller than us,
	546	* so any 'next' node must be in the general direction of our parent.
	547	* Go up the tree; any time the ancestor is a right-hand child of its
	548	* parent, keep going up. First time it's a left-hand child of its
	549	* parent, said parent is our 'next' node.
	550	*/
55a98102 DW	551	while ((parent = rb_parent(node)) && node == parent->rb_right)
55a98102 DW	552	node = parent;
1da177e4	553
55a98102	554	return parent;
1da177e4 LT	555	}
	556	EXPORT_SYMBOL(rb_next);
	557
f4b477c4	558	struct rb_node rb_prev(const struct rb_node node)
1da177e4	559	{
55a98102 DW	560	struct rb_node *parent;
55a98102 DW	561
4c199a93	562	if (RB_EMPTY_NODE(node))
10fd48f2 JA	563	return NULL;
10fd48f2 JA	564
7ce6ff9e ML	565	/*
	566	* If we have a left-hand child, go down and then right as far
	567	* as we can.
	568	*/
1da177e4	569	if (node->rb_left) {
cd9e61ed	570	node = node->rb_left;
1da177e4 LT	571	while (node->rb_right)
1da177e4 LT	572	node=node->rb_right;
f4b477c4	573	return (struct rb_node *)node;
1da177e4 LT	574	}
1da177e4 LT	575
7ce6ff9e ML	576	/*
	577	* No left-hand children. Go up till we find an ancestor which
	578	* is a right-hand child of its parent.
	579	*/
55a98102 DW	580	while ((parent = rb_parent(node)) && node == parent->rb_left)
55a98102 DW	581	node = parent;
1da177e4	582
55a98102	583	return parent;
1da177e4 LT	584	}
	585	EXPORT_SYMBOL(rb_prev);
	586
	587	void rb_replace_node(struct rb_node victim, struct rb_node new,
	588	struct rb_root *root)
	589	{
55a98102	590	struct rb_node *parent = rb_parent(victim);
1da177e4	591
c1adf200 DH	592	/* Copy the pointers/colour from the victim to the replacement */
	593	new = victim;
	594
1da177e4	595	/* Set the surrounding nodes to point to the replacement */
1da177e4	596	if (victim->rb_left)
55a98102	597	rb_set_parent(victim->rb_left, new);
1da177e4	598	if (victim->rb_right)
55a98102	599	rb_set_parent(victim->rb_right, new);
c1adf200 DH	600	__rb_change_child(victim, new, parent, root);
	601	}
	602	EXPORT_SYMBOL(rb_replace_node);
	603
	604	void rb_replace_node_rcu(struct rb_node victim, struct rb_node new,
	605	struct rb_root *root)
	606	{
	607	struct rb_node *parent = rb_parent(victim);
1da177e4 LT	608
	609	/* Copy the pointers/colour from the victim to the replacement */
	610	new = victim;
c1adf200 DH	611
	612	/* Set the surrounding nodes to point to the replacement */
	613	if (victim->rb_left)
	614	rb_set_parent(victim->rb_left, new);
	615	if (victim->rb_right)
	616	rb_set_parent(victim->rb_right, new);
	617
	618	/* Set the parent's pointer to the new node last after an RCU barrier
	619	* so that the pointers onwards are seen to be set correctly when doing
	620	* an RCU walk over the tree.
	621	*/
	622	__rb_change_child_rcu(victim, new, parent, root);
1da177e4	623	}
c1adf200	624	EXPORT_SYMBOL(rb_replace_node_rcu);
9dee5c51 CS	625
	626	static struct rb_node rb_left_deepest_node(const struct rb_node node)
	627	{
	628	for (;;) {
	629	if (node->rb_left)
	630	node = node->rb_left;
	631	else if (node->rb_right)
	632	node = node->rb_right;
	633	else
	634	return (struct rb_node *)node;
	635	}
	636	}
	637
	638	struct rb_node rb_next_postorder(const struct rb_node node)
	639	{
	640	const struct rb_node *parent;
	641	if (!node)
	642	return NULL;
	643	parent = rb_parent(node);
	644
	645	/* If we're sitting on node, we've already seen our children */
	646	if (parent && node == parent->rb_left && parent->rb_right) {
	647	/* If we are the parent's left node, go to the parent's right
	648	* node then all the way down to the left */
	649	return rb_left_deepest_node(parent->rb_right);
	650	} else
	651	/* Otherwise we are the parent's right node, and the parent
	652	* should be next */
	653	return (struct rb_node *)parent;
	654	}
	655	EXPORT_SYMBOL(rb_next_postorder);
	656
	657	struct rb_node rb_first_postorder(const struct rb_root root)
	658	{
	659	if (!root->rb_node)
	660	return NULL;
	661
	662	return rb_left_deepest_node(root->rb_node);
	663	}
	664	EXPORT_SYMBOL(rb_first_postorder);