include/linux/rbtree_augmented.h - third_party/kernel - Git at Google

 /* SPDX-License-Identifier: GPL-2.0-or-later */
 /*
   Red Black Trees
   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   (C) 2002  David Woodhouse <dwmw2@infradead.org>
   (C) 2012  Michel Lespinasse <walken@google.com>


   linux/include/linux/rbtree_augmented.h
 */

 #ifndef _LINUX_RBTREE_AUGMENTED_H
 #define _LINUX_RBTREE_AUGMENTED_H

 #include <linux/compiler.h>
 #include <linux/rbtree.h>
 #include <linux/rcupdate.h>

 /*
  * Please note - only struct rb_augment_callbacks and the prototypes for
  * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
  * The rest are implementation details you are not expected to depend on.
  *
  * See Documentation/core-api/rbtree.rst for documentation and samples.
  */

 struct rb_augment_callbacks {
 	void (*propagate)(struct rb_node *node, struct rb_node *stop);
 	void (*copy)(struct rb_node *old, struct rb_node *new);
 	void (*rotate)(struct rb_node *old, struct rb_node *new);
 };

 extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));

 /*
  * Fixup the rbtree and update the augmented information when rebalancing.
  *
  * On insertion, the user must update the augmented information on the path
  * leading to the inserted node, then call rb_link_node() as usual and
  * rb_insert_augmented() instead of the usual rb_insert_color() call.
  * If rb_insert_augmented() rebalances the rbtree, it will callback into
  * a user provided function to update the augmented information on the
  * affected subtrees.
  */
 static inline void
 rb_insert_augmented(struct rb_node *node, struct rb_root *root,
 		    const struct rb_augment_callbacks *augment)
 {
 	__rb_insert_augmented(node, root, augment->rotate);
 }

 static inline void
 rb_insert_augmented_cached(struct rb_node *node,
 			   struct rb_root_cached *root, bool newleft,
 			   const struct rb_augment_callbacks *augment)
 {
 	if (newleft)
 		root->rb_leftmost = node;
 	rb_insert_augmented(node, &root->rb_root, augment);
 }

 /*
  * Template for declaring augmented rbtree callbacks (generic case)
  *
  * RBSTATIC:    'static' or empty
  * RBNAME:      name of the rb_augment_callbacks structure
  * RBSTRUCT:    struct type of the tree nodes
  * RBFIELD:     name of struct rb_node field within RBSTRUCT
  * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
  * RBCOMPUTE:   name of function that recomputes the RBAUGMENTED data
  */

 #define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,				\
 			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE)	\
 static inline void							\
 RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
 {									\
 	while (rb != stop) {						\
 		RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD);	\
 		if (RBCOMPUTE(node, true))				\
 			break;						\
 		rb = rb_parent(&node->RBFIELD);				\
 	}								\
 }									\
 static inline void							\
 RBNAME ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)		\
 {									\
 	RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);		\
 	RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);		\
 	new->RBAUGMENTED = old->RBAUGMENTED;				\
 }									\
 static void								\
 RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
 {									\
 	RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);		\
 	RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);		\
 	new->RBAUGMENTED = old->RBAUGMENTED;				\
 	RBCOMPUTE(old, false);						\
 }									\
 RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
 	.propagate = RBNAME ## _propagate,				\
 	.copy = RBNAME ## _copy,					\
 	.rotate = RBNAME ## _rotate					\
 };

 /*
  * Template for declaring augmented rbtree callbacks,
  * computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes.
  *
  * RBSTATIC:    'static' or empty
  * RBNAME:      name of the rb_augment_callbacks structure
  * RBSTRUCT:    struct type of the tree nodes
  * RBFIELD:     name of struct rb_node field within RBSTRUCT
  * RBTYPE:      type of the RBAUGMENTED field
  * RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree
  * RBCOMPUTE:   name of function that returns the per-node RBTYPE scalar
  */

 #define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD,	      \
 				 RBTYPE, RBAUGMENTED, RBCOMPUTE)	      \
 static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit)	      \
 {									      \
 	RBSTRUCT *child;						      \
 	RBTYPE max = RBCOMPUTE(node);					      \
 	if (node->RBFIELD.rb_left) {					      \
 		child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD);   \
 		if (child->RBAUGMENTED > max)				      \
 			max = child->RBAUGMENTED;			      \
 	}								      \
 	if (node->RBFIELD.rb_right) {					      \
 		child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD);  \
 		if (child->RBAUGMENTED > max)				      \
 			max = child->RBAUGMENTED;			      \
 	}								      \
 	if (exit && node->RBAUGMENTED == max)				      \
 		return true;						      \
 	node->RBAUGMENTED = max;					      \
 	return false;							      \
 }									      \
 RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,					      \
 		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)


 #define	RB_RED		0
 #define	RB_BLACK	1

 #define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))

 #define __rb_color(pc)     ((pc) & 1)
 #define __rb_is_black(pc)  __rb_color(pc)
 #define __rb_is_red(pc)    (!__rb_color(pc))
 #define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
 #define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
 #define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)

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

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

 static inline void
 __rb_change_child(struct rb_node *old, struct rb_node *new,
 		  struct rb_node *parent, struct rb_root *root)
 {
 	if (parent) {
 		if (parent->rb_left == old)
 			WRITE_ONCE(parent->rb_left, new);
 		else
 			WRITE_ONCE(parent->rb_right, new);
 	} else
 		WRITE_ONCE(root->rb_node, new);
 }

 static inline void
 __rb_change_child_rcu(struct rb_node *old, struct rb_node *new,
 		      struct rb_node *parent, struct rb_root *root)
 {
 	if (parent) {
 		if (parent->rb_left == old)
 			rcu_assign_pointer(parent->rb_left, new);
 		else
 			rcu_assign_pointer(parent->rb_right, new);
 	} else
 		rcu_assign_pointer(root->rb_node, new);
 }

 extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));

 static __always_inline struct rb_node *
 __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
 		     const struct rb_augment_callbacks *augment)
 {
 	struct rb_node *child = node->rb_right;
 	struct rb_node *tmp = node->rb_left;
 	struct rb_node *parent, *rebalance;
 	unsigned long pc;

 	if (!tmp) {
 		/*
 		 * Case 1: node to erase has no more than 1 child (easy!)
 		 *
 		 * Note that if there is one child it must be red due to 5)
 		 * and node must be black due to 4). We adjust colors locally
 		 * so as to bypass __rb_erase_color() later on.
 		 */
 		pc = node->__rb_parent_color;
 		parent = __rb_parent(pc);
 		__rb_change_child(node, child, parent, root);
 		if (child) {
 			child->__rb_parent_color = pc;
 			rebalance = NULL;
 		} else
 			rebalance = __rb_is_black(pc) ? parent : NULL;
 		tmp = parent;
 	} else if (!child) {
 		/* Still case 1, but this time the child is node->rb_left */
 		tmp->__rb_parent_color = pc = node->__rb_parent_color;
 		parent = __rb_parent(pc);
 		__rb_change_child(node, tmp, parent, root);
 		rebalance = NULL;
 		tmp = parent;
 	} else {
 		struct rb_node *successor = child, *child2;

 		tmp = child->rb_left;
 		if (!tmp) {
 			/*
 			 * Case 2: node's successor is its right child
 			 *
 			 *    (n)          (s)
 			 *    / \          / \
 			 *  (x) (s)  ->  (x) (c)
 			 *        \
 			 *        (c)
 			 */
 			parent = successor;
 			child2 = successor->rb_right;

 			augment->copy(node, successor);
 		} else {
 			/*
 			 * Case 3: node's successor is leftmost under
 			 * node's right child subtree
 			 *
 			 *    (n)          (s)
 			 *    / \          / \
 			 *  (x) (y)  ->  (x) (y)
 			 *      /            /
 			 *    (p)          (p)
 			 *    /            /
 			 *  (s)          (c)
 			 *    \
 			 *    (c)
 			 */
 			do {
 				parent = successor;
 				successor = tmp;
 				tmp = tmp->rb_left;
 			} while (tmp);
 			child2 = successor->rb_right;
 			WRITE_ONCE(parent->rb_left, child2);
 			WRITE_ONCE(successor->rb_right, child);
 			rb_set_parent(child, successor);

 			augment->copy(node, successor);
 			augment->propagate(parent, successor);
 		}

 		tmp = node->rb_left;
 		WRITE_ONCE(successor->rb_left, tmp);
 		rb_set_parent(tmp, successor);

 		pc = node->__rb_parent_color;
 		tmp = __rb_parent(pc);
 		__rb_change_child(node, successor, tmp, root);

 		if (child2) {
 			rb_set_parent_color(child2, parent, RB_BLACK);
 			rebalance = NULL;
 		} else {
 			rebalance = rb_is_black(successor) ? parent : NULL;
 		}
 		successor->__rb_parent_color = pc;
 		tmp = successor;
 	}

 	augment->propagate(tmp, NULL);
 	return rebalance;
 }

 static __always_inline void
 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
 		   const struct rb_augment_callbacks *augment)
 {
 	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
 	if (rebalance)
 		__rb_erase_color(rebalance, root, augment->rotate);
 }

 static __always_inline void
 rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
 			  const struct rb_augment_callbacks *augment)
 {
 	if (root->rb_leftmost == node)
 		root->rb_leftmost = rb_next(node);
 	rb_erase_augmented(node, &root->rb_root, augment);
 }

 #endif	/* _LINUX_RBTREE_AUGMENTED_H */
	/* SPDX-License-Identifier: GPL-2.0-or-later */
	/*
	Red Black Trees
	(C) 1999 Andrea Arcangeli <andrea@suse.de>
	(C) 2002 David Woodhouse <dwmw2@infradead.org>
	(C) 2012 Michel Lespinasse <walken@google.com>


	linux/include/linux/rbtree_augmented.h
	*/

	#ifndef _LINUX_RBTREE_AUGMENTED_H
	#define _LINUX_RBTREE_AUGMENTED_H

	#include <linux/compiler.h>
	#include <linux/rbtree.h>
	#include <linux/rcupdate.h>

	/*
	* Please note - only struct rb_augment_callbacks and the prototypes for
	* rb_insert_augmented() and rb_erase_augmented() are intended to be public.
	* The rest are implementation details you are not expected to depend on.
	*
	* See Documentation/core-api/rbtree.rst for documentation and samples.
	*/

	struct rb_augment_callbacks {
	void (propagate)(struct rb_node node, struct rb_node *stop);
	void (copy)(struct rb_node old, struct rb_node *new);
	void (rotate)(struct rb_node old, struct rb_node *new);
	};

	extern void __rb_insert_augmented(struct rb_node node, struct rb_root root,
	void (augment_rotate)(struct rb_node old, struct rb_node *new));

	/*
	* Fixup the rbtree and update the augmented information when rebalancing.
	*
	* On insertion, the user must update the augmented information on the path
	* leading to the inserted node, then call rb_link_node() as usual and
	* rb_insert_augmented() instead of the usual rb_insert_color() call.
	* If rb_insert_augmented() rebalances the rbtree, it will callback into
	* a user provided function to update the augmented information on the
	* affected subtrees.
	*/
	static inline void
	rb_insert_augmented(struct rb_node node, struct rb_root root,
	const struct rb_augment_callbacks *augment)
	{
	__rb_insert_augmented(node, root, augment->rotate);
	}

	static inline void
	rb_insert_augmented_cached(struct rb_node *node,
	struct rb_root_cached *root, bool newleft,
	const struct rb_augment_callbacks *augment)
	{
	if (newleft)
	root->rb_leftmost = node;
	rb_insert_augmented(node, &root->rb_root, augment);
	}

	/*
	* Template for declaring augmented rbtree callbacks (generic case)
	*
	* RBSTATIC: 'static' or empty
	* RBNAME: name of the rb_augment_callbacks structure
	* RBSTRUCT: struct type of the tree nodes
	* RBFIELD: name of struct rb_node field within RBSTRUCT
	* RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
	* RBCOMPUTE: name of function that recomputes the RBAUGMENTED data
	*/

	#define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \
	RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE) \
	static inline void \
	RBNAME ## _propagate(struct rb_node rb, struct rb_node stop) \
	{ \
	while (rb != stop) { \
	RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD); \
	if (RBCOMPUTE(node, true)) \
	break; \
	rb = rb_parent(&node->RBFIELD); \
	} \
	} \
	static inline void \
	RBNAME ## _copy(struct rb_node rb_old, struct rb_node rb_new) \
	{ \
	RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \
	RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \
	new->RBAUGMENTED = old->RBAUGMENTED; \
	} \
	static void \
	RBNAME ## _rotate(struct rb_node rb_old, struct rb_node rb_new) \
	{ \
	RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \
	RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \
	new->RBAUGMENTED = old->RBAUGMENTED; \
	RBCOMPUTE(old, false); \
	} \
	RBSTATIC const struct rb_augment_callbacks RBNAME = { \
	.propagate = RBNAME ## _propagate, \
	.copy = RBNAME ## _copy, \
	.rotate = RBNAME ## _rotate \
	};

	/*
	* Template for declaring augmented rbtree callbacks,
	* computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes.
	*
	* RBSTATIC: 'static' or empty
	* RBNAME: name of the rb_augment_callbacks structure
	* RBSTRUCT: struct type of the tree nodes
	* RBFIELD: name of struct rb_node field within RBSTRUCT
	* RBTYPE: type of the RBAUGMENTED field
	* RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree
	* RBCOMPUTE: name of function that returns the per-node RBTYPE scalar
	*/

	#define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD, \
	RBTYPE, RBAUGMENTED, RBCOMPUTE) \
	static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit) \
	{ \
	RBSTRUCT *child; \
	RBTYPE max = RBCOMPUTE(node); \
	if (node->RBFIELD.rb_left) { \
	child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD); \
	if (child->RBAUGMENTED > max) \
	max = child->RBAUGMENTED; \
	} \
	if (node->RBFIELD.rb_right) { \
	child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD); \
	if (child->RBAUGMENTED > max) \
	max = child->RBAUGMENTED; \
	} \
	if (exit && node->RBAUGMENTED == max) \
	return true; \
	node->RBAUGMENTED = max; \
	return false; \
	} \
	RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \
	RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)


	#define RB_RED 0
	#define RB_BLACK 1

	#define __rb_parent(pc) ((struct rb_node *)(pc & ~3))

	#define __rb_color(pc) ((pc) & 1)
	#define __rb_is_black(pc) __rb_color(pc)
	#define __rb_is_red(pc) (!__rb_color(pc))
	#define rb_color(rb) __rb_color((rb)->__rb_parent_color)
	#define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
	#define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)

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

	static inline void rb_set_parent_color(struct rb_node *rb,
	struct rb_node *p, int color)
	{
	rb->__rb_parent_color = (unsigned long)p \| color;
	}

	static inline void
	__rb_change_child(struct rb_node old, struct rb_node new,
	struct rb_node parent, struct rb_root root)
	{
	if (parent) {
	if (parent->rb_left == old)
	WRITE_ONCE(parent->rb_left, new);
	else
	WRITE_ONCE(parent->rb_right, new);
	} else
	WRITE_ONCE(root->rb_node, new);
	}

	static inline void
	__rb_change_child_rcu(struct rb_node old, struct rb_node new,
	struct rb_node parent, struct rb_root root)
	{
	if (parent) {
	if (parent->rb_left == old)
	rcu_assign_pointer(parent->rb_left, new);
	else
	rcu_assign_pointer(parent->rb_right, new);
	} else
	rcu_assign_pointer(root->rb_node, new);
	}

	extern void __rb_erase_color(struct rb_node parent, struct rb_root root,
	void (augment_rotate)(struct rb_node old, struct rb_node *new));

	static __always_inline struct rb_node *
	__rb_erase_augmented(struct rb_node node, struct rb_root root,
	const struct rb_augment_callbacks *augment)
	{
	struct rb_node *child = node->rb_right;
	struct rb_node *tmp = node->rb_left;
	struct rb_node parent, rebalance;
	unsigned long pc;

	if (!tmp) {
	/*
	* Case 1: node to erase has no more than 1 child (easy!)
	*
	* Note that if there is one child it must be red due to 5)
	* and node must be black due to 4). We adjust colors locally
	* so as to bypass __rb_erase_color() later on.
	*/
	pc = node->__rb_parent_color;
	parent = __rb_parent(pc);
	__rb_change_child(node, child, parent, root);
	if (child) {
	child->__rb_parent_color = pc;
	rebalance = NULL;
	} else
	rebalance = __rb_is_black(pc) ? parent : NULL;
	tmp = parent;
	} else if (!child) {
	/* Still case 1, but this time the child is node->rb_left */
	tmp->__rb_parent_color = pc = node->__rb_parent_color;
	parent = __rb_parent(pc);
	__rb_change_child(node, tmp, parent, root);
	rebalance = NULL;
	tmp = parent;
	} else {
	struct rb_node successor = child, child2;

	tmp = child->rb_left;
	if (!tmp) {
	/*
	* Case 2: node's successor is its right child
	*
	* (n) (s)
	* / \ / \
	* (x) (s) -> (x) (c)
	* \
	* (c)
	*/
	parent = successor;
	child2 = successor->rb_right;

	augment->copy(node, successor);
	} else {
	/*
	* Case 3: node's successor is leftmost under
	* node's right child subtree
	*
	* (n) (s)
	* / \ / \
	* (x) (y) -> (x) (y)
	* / /
	* (p) (p)
	* / /
	* (s) (c)
	* \
	* (c)
	*/
	do {
	parent = successor;
	successor = tmp;
	tmp = tmp->rb_left;
	} while (tmp);
	child2 = successor->rb_right;
	WRITE_ONCE(parent->rb_left, child2);
	WRITE_ONCE(successor->rb_right, child);
	rb_set_parent(child, successor);

	augment->copy(node, successor);
	augment->propagate(parent, successor);
	}

	tmp = node->rb_left;
	WRITE_ONCE(successor->rb_left, tmp);
	rb_set_parent(tmp, successor);

	pc = node->__rb_parent_color;
	tmp = __rb_parent(pc);
	__rb_change_child(node, successor, tmp, root);

	if (child2) {
	rb_set_parent_color(child2, parent, RB_BLACK);
	rebalance = NULL;
	} else {
	rebalance = rb_is_black(successor) ? parent : NULL;
	}
	successor->__rb_parent_color = pc;
	tmp = successor;
	}

	augment->propagate(tmp, NULL);
	return rebalance;
	}

	static __always_inline void
	rb_erase_augmented(struct rb_node node, struct rb_root root,
	const struct rb_augment_callbacks *augment)
	{
	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
	if (rebalance)
	__rb_erase_color(rebalance, root, augment->rotate);
	}

	static __always_inline void
	rb_erase_augmented_cached(struct rb_node node, struct rb_root_cached root,
	const struct rb_augment_callbacks *augment)
	{
	if (root->rb_leftmost == node)
	root->rb_leftmost = rb_next(node);
	rb_erase_augmented(node, &root->rb_root, augment);
	}

	#endif /* _LINUX_RBTREE_AUGMENTED_H */