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