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