1 /* SPDX-License-Identifier: GPL-2.0+ */ 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 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