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