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