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