Home | History | Annotate | Download | only in linux
      1 /* SPDX-License-Identifier: GPL-2.0+ */
      2 /*
      3   Red Black Trees
      4   (C) 1999  Andrea Arcangeli <andrea (at) suse.de>
      5   (C) 2002  David Woodhouse <dwmw2 (at) infradead.org>
      6   (C) 2012  Michel Lespinasse <walken (at) google.com>
      7 
      8   linux/include/linux/rbtree_augmented.h
      9 */
     10 
     11 #ifndef _LINUX_RBTREE_AUGMENTED_H
     12 #define _LINUX_RBTREE_AUGMENTED_H
     13 
     14 #include <linux/compiler.h>
     15 #include <linux/rbtree.h>
     16 
     17 /*
     18  * Please note - only struct rb_augment_callbacks and the prototypes for
     19  * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
     20  * The rest are implementation details you are not expected to depend on.
     21  *
     22  * See Documentation/rbtree.txt for documentation and samples.
     23  */
     24 
     25 struct rb_augment_callbacks {
     26 	void (*propagate)(struct rb_node *node, struct rb_node *stop);
     27 	void (*copy)(struct rb_node *old, struct rb_node *new);
     28 	void (*rotate)(struct rb_node *old, struct rb_node *new);
     29 };
     30 
     31 extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
     32 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
     33 static inline void
     34 rb_insert_augmented(struct rb_node *node, struct rb_root *root,
     35 		    const struct rb_augment_callbacks *augment)
     36 {
     37 	__rb_insert_augmented(node, root, augment->rotate);
     38 }
     39 
     40 #define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield,	\
     41 			     rbtype, rbaugmented, rbcompute)		\
     42 static inline void							\
     43 rbname ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
     44 {									\
     45 	while (rb != stop) {						\
     46 		rbstruct *node = rb_entry(rb, rbstruct, rbfield);	\
     47 		rbtype augmented = rbcompute(node);			\
     48 		if (node->rbaugmented == augmented)			\
     49 			break;						\
     50 		node->rbaugmented = augmented;				\
     51 		rb = rb_parent(&node->rbfield);				\
     52 	}								\
     53 }									\
     54 static inline void							\
     55 rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)		\
     56 {									\
     57 	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
     58 	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
     59 	new->rbaugmented = old->rbaugmented;				\
     60 }									\
     61 static void								\
     62 rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
     63 {									\
     64 	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
     65 	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
     66 	new->rbaugmented = old->rbaugmented;				\
     67 	old->rbaugmented = rbcompute(old);				\
     68 }									\
     69 rbstatic const struct rb_augment_callbacks rbname = {			\
     70 	rbname ## _propagate, rbname ## _copy, rbname ## _rotate	\
     71 };
     72 
     73 
     74 #define	RB_RED		0
     75 #define	RB_BLACK	1
     76 
     77 #define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))
     78 
     79 #define __rb_color(pc)     ((pc) & 1)
     80 #define __rb_is_black(pc)  __rb_color(pc)
     81 #define __rb_is_red(pc)    (!__rb_color(pc))
     82 #define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
     83 #define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
     84 #define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
     85 
     86 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
     87 {
     88 	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
     89 }
     90 
     91 static inline void rb_set_parent_color(struct rb_node *rb,
     92 				       struct rb_node *p, int color)
     93 {
     94 	rb->__rb_parent_color = (unsigned long)p | color;
     95 }
     96 
     97 static inline void
     98 __rb_change_child(struct rb_node *old, struct rb_node *new,
     99 		  struct rb_node *parent, struct rb_root *root)
    100 {
    101 	if (parent) {
    102 		if (parent->rb_left == old)
    103 			parent->rb_left = new;
    104 		else
    105 			parent->rb_right = new;
    106 	} else
    107 		root->rb_node = new;
    108 }
    109 
    110 extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
    111 	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
    112 
    113 static __always_inline struct rb_node *
    114 __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
    115 		     const struct rb_augment_callbacks *augment)
    116 {
    117 	struct rb_node *child = node->rb_right, *tmp = node->rb_left;
    118 	struct rb_node *parent, *rebalance;
    119 	unsigned long pc;
    120 
    121 	if (!tmp) {
    122 		/*
    123 		 * Case 1: node to erase has no more than 1 child (easy!)
    124 		 *
    125 		 * Note that if there is one child it must be red due to 5)
    126 		 * and node must be black due to 4). We adjust colors locally
    127 		 * so as to bypass __rb_erase_color() later on.
    128 		 */
    129 		pc = node->__rb_parent_color;
    130 		parent = __rb_parent(pc);
    131 		__rb_change_child(node, child, parent, root);
    132 		if (child) {
    133 			child->__rb_parent_color = pc;
    134 			rebalance = NULL;
    135 		} else
    136 			rebalance = __rb_is_black(pc) ? parent : NULL;
    137 		tmp = parent;
    138 	} else if (!child) {
    139 		/* Still case 1, but this time the child is node->rb_left */
    140 		tmp->__rb_parent_color = pc = node->__rb_parent_color;
    141 		parent = __rb_parent(pc);
    142 		__rb_change_child(node, tmp, parent, root);
    143 		rebalance = NULL;
    144 		tmp = parent;
    145 	} else {
    146 		struct rb_node *successor = child, *child2;
    147 		tmp = child->rb_left;
    148 		if (!tmp) {
    149 			/*
    150 			 * Case 2: node's successor is its right child
    151 			 *
    152 			 *    (n)          (s)
    153 			 *    / \          / \
    154 			 *  (x) (s)  ->  (x) (c)
    155 			 *        \
    156 			 *        (c)
    157 			 */
    158 			parent = successor;
    159 			child2 = successor->rb_right;
    160 			augment->copy(node, successor);
    161 		} else {
    162 			/*
    163 			 * Case 3: node's successor is leftmost under
    164 			 * node's right child subtree
    165 			 *
    166 			 *    (n)          (s)
    167 			 *    / \          / \
    168 			 *  (x) (y)  ->  (x) (y)
    169 			 *      /            /
    170 			 *    (p)          (p)
    171 			 *    /            /
    172 			 *  (s)          (c)
    173 			 *    \
    174 			 *    (c)
    175 			 */
    176 			do {
    177 				parent = successor;
    178 				successor = tmp;
    179 				tmp = tmp->rb_left;
    180 			} while (tmp);
    181 			parent->rb_left = child2 = successor->rb_right;
    182 			successor->rb_right = child;
    183 			rb_set_parent(child, successor);
    184 			augment->copy(node, successor);
    185 			augment->propagate(parent, successor);
    186 		}
    187 
    188 		successor->rb_left = tmp = node->rb_left;
    189 		rb_set_parent(tmp, successor);
    190 
    191 		pc = node->__rb_parent_color;
    192 		tmp = __rb_parent(pc);
    193 		__rb_change_child(node, successor, tmp, root);
    194 		if (child2) {
    195 			successor->__rb_parent_color = pc;
    196 			rb_set_parent_color(child2, parent, RB_BLACK);
    197 			rebalance = NULL;
    198 		} else {
    199 			unsigned long pc2 = successor->__rb_parent_color;
    200 			successor->__rb_parent_color = pc;
    201 			rebalance = __rb_is_black(pc2) ? parent : NULL;
    202 		}
    203 		tmp = successor;
    204 	}
    205 
    206 	augment->propagate(tmp, NULL);
    207 	return rebalance;
    208 }
    209 
    210 static __always_inline void
    211 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
    212 		   const struct rb_augment_callbacks *augment)
    213 {
    214 	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
    215 	if (rebalance)
    216 		__rb_erase_color(rebalance, root, augment->rotate);
    217 }
    218 
    219 #endif	/* _LINUX_RBTREE_AUGMENTED_H */
    220