[mirror_ubuntu-zesty-kernel.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,
	    void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;

	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) {
			rb_set_parent_color(node, NULL, RB_BLACK);
			break;
		} else 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 = {
	dummy_propagate, dummy_copy, dummy_rotate
};

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
	__rb_insert(node, root, 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, &dummy_callbacks);
	if (rebalance)
		____rb_erase_color(rebalance, root, dummy_rotate);
}
EXPORT_SYMBOL(rb_erase);

/*
 * 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,
	void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
{
	__rb_insert(node, root, 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
	97	__rb_insert(struct rb_node node, struct rb_root root,
	98	void (augment_rotate)(struct rb_node old, struct rb_node *new))
1da177e4	99	{
5bc9188a	100	struct rb_node parent = rb_red_parent(node), gparent, *tmp;
1da177e4	101
6d58452d ML	102	while (true) {
	103	/*
	104	* Loop invariant: node is red
	105	*
	106	* If there is a black parent, we are done.
	107	* Otherwise, take some corrective action as we don't
	108	* want a red root or two consecutive red nodes.
	109	*/
6d58452d	110	if (!parent) {
5bc9188a	111	rb_set_parent_color(node, NULL, RB_BLACK);
6d58452d ML	112	break;
	113	} else if (rb_is_black(parent))
	114	break;
	115
5bc9188a ML	116	gparent = rb_red_parent(parent);
5bc9188a ML	117
59633abf ML	118	tmp = gparent->rb_right;
59633abf ML	119	if (parent != tmp) { /* parent == gparent->rb_left */
5bc9188a ML	120	if (tmp && rb_is_red(tmp)) {
	121	/*
	122	* Case 1 - color flips
	123	*
	124	* G g
	125	* / \ / \
	126	* p u --> P U
	127	* / /
1b9c53e8	128	* n n
5bc9188a ML	129	*
	130	* However, since g's parent might be red, and
	131	* 4) does not allow this, we need to recurse
	132	* at g.
	133	*/
	134	rb_set_parent_color(tmp, gparent, RB_BLACK);
	135	rb_set_parent_color(parent, gparent, RB_BLACK);
	136	node = gparent;
	137	parent = rb_parent(node);
	138	rb_set_parent_color(node, parent, RB_RED);
	139	continue;
1da177e4 LT	140	}
1da177e4 LT	141
59633abf ML	142	tmp = parent->rb_right;
59633abf ML	143	if (node == tmp) {
5bc9188a ML	144	/*
	145	* Case 2 - left rotate at parent
	146	*
	147	* G G
	148	* / \ / \
	149	* p U --> n U
	150	* \ /
	151	* n p
	152	*
	153	* This still leaves us in violation of 4), the
	154	* continuation into Case 3 will fix that.
	155	*/
d72da4a4 PZ	156	tmp = node->rb_left;
	157	WRITE_ONCE(parent->rb_right, tmp);
	158	WRITE_ONCE(node->rb_left, parent);
5bc9188a ML	159	if (tmp)
	160	rb_set_parent_color(tmp, parent,
	161	RB_BLACK);
	162	rb_set_parent_color(parent, node, RB_RED);
14b94af0	163	augment_rotate(parent, node);
1da177e4	164	parent = node;
59633abf	165	tmp = node->rb_right;
1da177e4 LT	166	}
1da177e4 LT	167
5bc9188a ML	168	/*
	169	* Case 3 - right rotate at gparent
	170	*
	171	* G P
	172	* / \ / \
	173	* p U --> n g
	174	* / \
	175	* n U
	176	*/
d72da4a4 PZ	177	WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
d72da4a4 PZ	178	WRITE_ONCE(parent->rb_right, gparent);
5bc9188a ML	179	if (tmp)
	180	rb_set_parent_color(tmp, gparent, RB_BLACK);
	181	__rb_rotate_set_parents(gparent, parent, root, RB_RED);
14b94af0	182	augment_rotate(gparent, parent);
1f052865	183	break;
1da177e4	184	} else {
5bc9188a ML	185	tmp = gparent->rb_left;
	186	if (tmp && rb_is_red(tmp)) {
	187	/* Case 1 - color flips */
	188	rb_set_parent_color(tmp, gparent, RB_BLACK);
	189	rb_set_parent_color(parent, gparent, RB_BLACK);
	190	node = gparent;
	191	parent = rb_parent(node);
	192	rb_set_parent_color(node, parent, RB_RED);
	193	continue;
1da177e4 LT	194	}
1da177e4 LT	195
59633abf ML	196	tmp = parent->rb_left;
59633abf ML	197	if (node == tmp) {
5bc9188a	198	/* Case 2 - right rotate at parent */
d72da4a4 PZ	199	tmp = node->rb_right;
	200	WRITE_ONCE(parent->rb_left, tmp);
	201	WRITE_ONCE(node->rb_right, parent);
5bc9188a ML	202	if (tmp)
	203	rb_set_parent_color(tmp, parent,
	204	RB_BLACK);
	205	rb_set_parent_color(parent, node, RB_RED);
14b94af0	206	augment_rotate(parent, node);
1da177e4	207	parent = node;
59633abf	208	tmp = node->rb_left;
1da177e4 LT	209	}
1da177e4 LT	210
5bc9188a	211	/* Case 3 - left rotate at gparent */
d72da4a4 PZ	212	WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
d72da4a4 PZ	213	WRITE_ONCE(parent->rb_left, gparent);
5bc9188a ML	214	if (tmp)
	215	rb_set_parent_color(tmp, gparent, RB_BLACK);
	216	__rb_rotate_set_parents(gparent, parent, root, RB_RED);
14b94af0	217	augment_rotate(gparent, parent);
1f052865	218	break;
1da177e4 LT	219	}
1da177e4 LT	220	}
1da177e4	221	}
1da177e4	222
3cb7a563 ML	223	/*
	224	* Inline version for rb_erase() use - we want to be able to inline
	225	* and eliminate the dummy_rotate callback there
	226	*/
	227	static __always_inline void
	228	____rb_erase_color(struct rb_node parent, struct rb_root root,
9c079add	229	void (augment_rotate)(struct rb_node old, struct rb_node *new))
1da177e4	230	{
46b6135a	231	struct rb_node node = NULL, sibling, tmp1, tmp2;
1da177e4	232
d6ff1273 ML	233	while (true) {
d6ff1273 ML	234	/*
46b6135a ML	235	* Loop invariants:
	236	* - node is black (or NULL on first iteration)
	237	* - node is not the root (parent is not NULL)
	238	* - All leaf paths going through parent and node have a
	239	* black node count that is 1 lower than other leaf paths.
d6ff1273	240	*/
59633abf ML	241	sibling = parent->rb_right;
59633abf ML	242	if (node != sibling) { /* node == parent->rb_left */
6280d235 ML	243	if (rb_is_red(sibling)) {
	244	/*
	245	* Case 1 - left rotate at parent
	246	*
	247	* P S
	248	* / \ / \
	249	* N s --> p Sr
	250	* / \ / \
	251	* Sl Sr N Sl
	252	*/
d72da4a4 PZ	253	tmp1 = sibling->rb_left;
	254	WRITE_ONCE(parent->rb_right, tmp1);
	255	WRITE_ONCE(sibling->rb_left, parent);
6280d235 ML	256	rb_set_parent_color(tmp1, parent, RB_BLACK);
	257	__rb_rotate_set_parents(parent, sibling, root,
	258	RB_RED);
9c079add	259	augment_rotate(parent, sibling);
6280d235	260	sibling = tmp1;
1da177e4	261	}
6280d235 ML	262	tmp1 = sibling->rb_right;
	263	if (!tmp1 \|\| rb_is_black(tmp1)) {
	264	tmp2 = sibling->rb_left;
	265	if (!tmp2 \|\| rb_is_black(tmp2)) {
	266	/*
	267	* Case 2 - sibling color flip
	268	* (p could be either color here)
	269	*
	270	* (p) (p)
	271	* / \ / \
	272	* N S --> N s
	273	* / \ / \
	274	* Sl Sr Sl Sr
	275	*
46b6135a ML	276	* This leaves us violating 5) which
	277	* can be fixed by flipping p to black
	278	* if it was red, or by recursing at p.
	279	* p is red when coming from Case 1.
6280d235 ML	280	*/
	281	rb_set_parent_color(sibling, parent,
	282	RB_RED);
46b6135a ML	283	if (rb_is_red(parent))
	284	rb_set_black(parent);
	285	else {
	286	node = parent;
	287	parent = rb_parent(node);
	288	if (parent)
	289	continue;
	290	}
	291	break;
1da177e4	292	}
6280d235 ML	293	/*
	294	* Case 3 - right rotate at sibling
	295	* (p could be either color here)
	296	*
	297	* (p) (p)
	298	* / \ / \
ce093a04	299	* N S --> N sl
6280d235	300	* / \ \
ce093a04	301	* sl Sr S
6280d235 ML	302	* \
6280d235 ML	303	* Sr
ce093a04 JC	304	*
	305	* Note: p might be red, and then both
	306	* p and sl are red after rotation(which
	307	* breaks property 4). This is fixed in
	308	* Case 4 (in __rb_rotate_set_parents()
	309	* which set sl the color of p
	310	* and set p RB_BLACK)
	311	*
	312	* (p) (sl)
	313	* / \ / \
	314	* N sl --> P S
	315	* \ / \
	316	* S N Sr
	317	* \
	318	* Sr
6280d235	319	*/
d72da4a4 PZ	320	tmp1 = tmp2->rb_right;
	321	WRITE_ONCE(sibling->rb_left, tmp1);
	322	WRITE_ONCE(tmp2->rb_right, sibling);
	323	WRITE_ONCE(parent->rb_right, tmp2);
6280d235 ML	324	if (tmp1)
	325	rb_set_parent_color(tmp1, sibling,
	326	RB_BLACK);
9c079add	327	augment_rotate(sibling, tmp2);
6280d235 ML	328	tmp1 = sibling;
6280d235 ML	329	sibling = tmp2;
1da177e4	330	}
6280d235 ML	331	/*
	332	* Case 4 - left rotate at parent + color flips
	333	* (p and sl could be either color here.
	334	* After rotation, p becomes black, s acquires
	335	* p's color, and sl keeps its color)
	336	*
	337	* (p) (s)
	338	* / \ / \
	339	* N S --> P Sr
	340	* / \ / \
	341	* (sl) sr N (sl)
	342	*/
d72da4a4 PZ	343	tmp2 = sibling->rb_left;
	344	WRITE_ONCE(parent->rb_right, tmp2);
	345	WRITE_ONCE(sibling->rb_left, parent);
6280d235 ML	346	rb_set_parent_color(tmp1, sibling, RB_BLACK);
	347	if (tmp2)
	348	rb_set_parent(tmp2, parent);
	349	__rb_rotate_set_parents(parent, sibling, root,
	350	RB_BLACK);
9c079add	351	augment_rotate(parent, sibling);
e125d147	352	break;
d6ff1273	353	} else {
6280d235 ML	354	sibling = parent->rb_left;
	355	if (rb_is_red(sibling)) {
	356	/* Case 1 - right rotate at parent */
d72da4a4 PZ	357	tmp1 = sibling->rb_right;
	358	WRITE_ONCE(parent->rb_left, tmp1);
	359	WRITE_ONCE(sibling->rb_right, parent);
6280d235 ML	360	rb_set_parent_color(tmp1, parent, RB_BLACK);
	361	__rb_rotate_set_parents(parent, sibling, root,
	362	RB_RED);
9c079add	363	augment_rotate(parent, sibling);
6280d235	364	sibling = tmp1;
1da177e4	365	}
6280d235 ML	366	tmp1 = sibling->rb_left;
	367	if (!tmp1 \|\| rb_is_black(tmp1)) {
	368	tmp2 = sibling->rb_right;
	369	if (!tmp2 \|\| rb_is_black(tmp2)) {
	370	/* Case 2 - sibling color flip */
	371	rb_set_parent_color(sibling, parent,
	372	RB_RED);
46b6135a ML	373	if (rb_is_red(parent))
	374	rb_set_black(parent);
	375	else {
	376	node = parent;
	377	parent = rb_parent(node);
	378	if (parent)
	379	continue;
	380	}
	381	break;
1da177e4	382	}
ce093a04	383	/* Case 3 - left rotate at sibling */
d72da4a4 PZ	384	tmp1 = tmp2->rb_left;
	385	WRITE_ONCE(sibling->rb_right, tmp1);
	386	WRITE_ONCE(tmp2->rb_left, sibling);
	387	WRITE_ONCE(parent->rb_left, tmp2);
6280d235 ML	388	if (tmp1)
	389	rb_set_parent_color(tmp1, sibling,
	390	RB_BLACK);
9c079add	391	augment_rotate(sibling, tmp2);
6280d235 ML	392	tmp1 = sibling;
6280d235 ML	393	sibling = tmp2;
1da177e4	394	}
ce093a04	395	/* Case 4 - right rotate at parent + color flips */
d72da4a4 PZ	396	tmp2 = sibling->rb_right;
	397	WRITE_ONCE(parent->rb_left, tmp2);
	398	WRITE_ONCE(sibling->rb_right, parent);
6280d235 ML	399	rb_set_parent_color(tmp1, sibling, RB_BLACK);
	400	if (tmp2)
	401	rb_set_parent(tmp2, parent);
	402	__rb_rotate_set_parents(parent, sibling, root,
	403	RB_BLACK);
9c079add	404	augment_rotate(parent, sibling);
e125d147	405	break;
1da177e4 LT	406	}
1da177e4 LT	407	}
1da177e4	408	}
3cb7a563 ML	409
	410	/* Non-inline version for rb_erase_augmented() use */
	411	void __rb_erase_color(struct rb_node parent, struct rb_root root,
	412	void (augment_rotate)(struct rb_node old, struct rb_node *new))
	413	{
	414	____rb_erase_color(parent, root, augment_rotate);
	415	}
9c079add	416	EXPORT_SYMBOL(__rb_erase_color);
14b94af0 ML	417
	418	/*
	419	* Non-augmented rbtree manipulation functions.
	420	*
	421	* We use dummy augmented callbacks here, and have the compiler optimize them
	422	* out of the rb_insert_color() and rb_erase() function definitions.
	423	*/
	424
	425	static inline void dummy_propagate(struct rb_node node, struct rb_node stop) {}
	426	static inline void dummy_copy(struct rb_node old, struct rb_node new) {}
	427	static inline void dummy_rotate(struct rb_node old, struct rb_node new) {}
	428
	429	static const struct rb_augment_callbacks dummy_callbacks = {
	430	dummy_propagate, dummy_copy, dummy_rotate
	431	};
	432
	433	void rb_insert_color(struct rb_node node, struct rb_root root)
	434	{
	435	__rb_insert(node, root, dummy_rotate);
	436	}
	437	EXPORT_SYMBOL(rb_insert_color);
	438
	439	void rb_erase(struct rb_node node, struct rb_root root)
	440	{
3cb7a563 ML	441	struct rb_node *rebalance;
	442	rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
	443	if (rebalance)
	444	____rb_erase_color(rebalance, root, dummy_rotate);
1da177e4 LT	445	}
	446	EXPORT_SYMBOL(rb_erase);
	447
14b94af0 ML	448	/*
	449	* Augmented rbtree manipulation functions.
	450	*
	451	* This instantiates the same __always_inline functions as in the non-augmented
	452	* case, but this time with user-defined callbacks.
	453	*/
	454
	455	void __rb_insert_augmented(struct rb_node node, struct rb_root root,
	456	void (augment_rotate)(struct rb_node old, struct rb_node *new))
	457	{
	458	__rb_insert(node, root, augment_rotate);
	459	}
	460	EXPORT_SYMBOL(__rb_insert_augmented);
	461
1da177e4 LT	462	/*
	463	* This function returns the first node (in sort order) of the tree.
	464	*/
f4b477c4	465	struct rb_node rb_first(const struct rb_root root)
1da177e4 LT	466	{
	467	struct rb_node *n;
	468
	469	n = root->rb_node;
	470	if (!n)
	471	return NULL;
	472	while (n->rb_left)
	473	n = n->rb_left;
	474	return n;
	475	}
	476	EXPORT_SYMBOL(rb_first);
	477
f4b477c4	478	struct rb_node rb_last(const struct rb_root root)
1da177e4 LT	479	{
	480	struct rb_node *n;
	481
	482	n = root->rb_node;
	483	if (!n)
	484	return NULL;
	485	while (n->rb_right)
	486	n = n->rb_right;
	487	return n;
	488	}
	489	EXPORT_SYMBOL(rb_last);
	490
f4b477c4	491	struct rb_node rb_next(const struct rb_node node)
1da177e4	492	{
55a98102 DW	493	struct rb_node *parent;
55a98102 DW	494
4c199a93	495	if (RB_EMPTY_NODE(node))
10fd48f2 JA	496	return NULL;
10fd48f2 JA	497
7ce6ff9e ML	498	/*
	499	* If we have a right-hand child, go down and then left as far
	500	* as we can.
	501	*/
1da177e4 LT	502	if (node->rb_right) {
	503	node = node->rb_right;
	504	while (node->rb_left)
	505	node=node->rb_left;
f4b477c4	506	return (struct rb_node *)node;
1da177e4 LT	507	}
1da177e4 LT	508
7ce6ff9e ML	509	/*
	510	* No right-hand children. Everything down and left is smaller than us,
	511	* so any 'next' node must be in the general direction of our parent.
	512	* Go up the tree; any time the ancestor is a right-hand child of its
	513	* parent, keep going up. First time it's a left-hand child of its
	514	* parent, said parent is our 'next' node.
	515	*/
55a98102 DW	516	while ((parent = rb_parent(node)) && node == parent->rb_right)
55a98102 DW	517	node = parent;
1da177e4	518
55a98102	519	return parent;
1da177e4 LT	520	}
	521	EXPORT_SYMBOL(rb_next);
	522
f4b477c4	523	struct rb_node rb_prev(const struct rb_node node)
1da177e4	524	{
55a98102 DW	525	struct rb_node *parent;
55a98102 DW	526
4c199a93	527	if (RB_EMPTY_NODE(node))
10fd48f2 JA	528	return NULL;
10fd48f2 JA	529
7ce6ff9e ML	530	/*
	531	* If we have a left-hand child, go down and then right as far
	532	* as we can.
	533	*/
1da177e4 LT	534	if (node->rb_left) {
	535	node = node->rb_left;
	536	while (node->rb_right)
	537	node=node->rb_right;
f4b477c4	538	return (struct rb_node *)node;
1da177e4 LT	539	}
1da177e4 LT	540
7ce6ff9e ML	541	/*
	542	* No left-hand children. Go up till we find an ancestor which
	543	* is a right-hand child of its parent.
	544	*/
55a98102 DW	545	while ((parent = rb_parent(node)) && node == parent->rb_left)
55a98102 DW	546	node = parent;
1da177e4	547
55a98102	548	return parent;
1da177e4 LT	549	}
	550	EXPORT_SYMBOL(rb_prev);
	551
	552	void rb_replace_node(struct rb_node victim, struct rb_node new,
	553	struct rb_root *root)
	554	{
55a98102	555	struct rb_node *parent = rb_parent(victim);
1da177e4	556
c1adf200 DH	557	/* Copy the pointers/colour from the victim to the replacement */
	558	new = victim;
	559
1da177e4	560	/* Set the surrounding nodes to point to the replacement */
1da177e4	561	if (victim->rb_left)
55a98102	562	rb_set_parent(victim->rb_left, new);
1da177e4	563	if (victim->rb_right)
55a98102	564	rb_set_parent(victim->rb_right, new);
c1adf200 DH	565	__rb_change_child(victim, new, parent, root);
	566	}
	567	EXPORT_SYMBOL(rb_replace_node);
	568
	569	void rb_replace_node_rcu(struct rb_node victim, struct rb_node new,
	570	struct rb_root *root)
	571	{
	572	struct rb_node *parent = rb_parent(victim);
1da177e4 LT	573
	574	/* Copy the pointers/colour from the victim to the replacement */
	575	new = victim;
c1adf200 DH	576
	577	/* Set the surrounding nodes to point to the replacement */
	578	if (victim->rb_left)
	579	rb_set_parent(victim->rb_left, new);
	580	if (victim->rb_right)
	581	rb_set_parent(victim->rb_right, new);
	582
	583	/* Set the parent's pointer to the new node last after an RCU barrier
	584	* so that the pointers onwards are seen to be set correctly when doing
	585	* an RCU walk over the tree.
	586	*/
	587	__rb_change_child_rcu(victim, new, parent, root);
1da177e4	588	}
c1adf200	589	EXPORT_SYMBOL(rb_replace_node_rcu);
9dee5c51 CS	590
	591	static struct rb_node rb_left_deepest_node(const struct rb_node node)
	592	{
	593	for (;;) {
	594	if (node->rb_left)
	595	node = node->rb_left;
	596	else if (node->rb_right)
	597	node = node->rb_right;
	598	else
	599	return (struct rb_node *)node;
	600	}
	601	}
	602
	603	struct rb_node rb_next_postorder(const struct rb_node node)
	604	{
	605	const struct rb_node *parent;
	606	if (!node)
	607	return NULL;
	608	parent = rb_parent(node);
	609
	610	/* If we're sitting on node, we've already seen our children */
	611	if (parent && node == parent->rb_left && parent->rb_right) {
	612	/* If we are the parent's left node, go to the parent's right
	613	* node then all the way down to the left */
	614	return rb_left_deepest_node(parent->rb_right);
	615	} else
	616	/* Otherwise we are the parent's right node, and the parent
	617	* should be next */
	618	return (struct rb_node *)parent;
	619	}
	620	EXPORT_SYMBOL(rb_next_postorder);
	621
	622	struct rb_node rb_first_postorder(const struct rb_root root)
	623	{
	624	if (!root->rb_node)
	625	return NULL;
	626
	627	return rb_left_deepest_node(root->rb_node);
	628	}
	629	EXPORT_SYMBOL(rb_first_postorder);