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
rb_insert_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)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
rb_set_parent(struct rb_node * rb,struct rb_node * p)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
rb_set_parent_color(struct rb_node * rb,struct rb_node * p,int color)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
__rb_change_child(struct rb_node * old,struct rb_node * new,struct rb_node * parent,struct rb_root * root)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 *
__rb_erase_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)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
rb_erase_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)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