1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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#include <linux/rcupdate.h>
30
31
32
33
34
35
36
37
38
39struct rb_augment_callbacks {
40 void (*propagate)(struct rb_node *node, struct rb_node *stop);
41 void (*copy)(struct rb_node *old, struct rb_node *new);
42 void (*rotate)(struct rb_node *old, struct rb_node *new);
43};
44
45extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
46 void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
47static inline void
48rb_insert_augmented(struct rb_node *node, struct rb_root *root,
49 const struct rb_augment_callbacks *augment)
50{
51 __rb_insert_augmented(node, root, augment->rotate);
52}
53
54#define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield, \
55 rbtype, rbaugmented, rbcompute) \
56static inline void \
57rbname ## _propagate(struct rb_node *rb, struct rb_node *stop) \
58{ \
59 while (rb != stop) { \
60 rbstruct *node = rb_entry(rb, rbstruct, rbfield); \
61 rbtype augmented = rbcompute(node); \
62 if (node->rbaugmented == augmented) \
63 break; \
64 node->rbaugmented = augmented; \
65 rb = rb_parent(&node->rbfield); \
66 } \
67} \
68static inline void \
69rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new) \
70{ \
71 rbstruct *old = rb_entry(rb_old, rbstruct, rbfield); \
72 rbstruct *new = rb_entry(rb_new, rbstruct, rbfield); \
73 new->rbaugmented = old->rbaugmented; \
74} \
75static void \
76rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new) \
77{ \
78 rbstruct *old = rb_entry(rb_old, rbstruct, rbfield); \
79 rbstruct *new = rb_entry(rb_new, rbstruct, rbfield); \
80 new->rbaugmented = old->rbaugmented; \
81 old->rbaugmented = rbcompute(old); \
82} \
83rbstatic const struct rb_augment_callbacks rbname = { \
84 rbname ## _propagate, rbname ## _copy, rbname ## _rotate \
85};
86
87
88#define RB_RED 0
89#define RB_BLACK 1
90
91#define __rb_parent(pc) ((struct rb_node *)(pc & ~3))
92
93#define __rb_color(pc) ((pc) & 1)
94#define __rb_is_black(pc) __rb_color(pc)
95#define __rb_is_red(pc) (!__rb_color(pc))
96#define rb_color(rb) __rb_color((rb)->__rb_parent_color)
97#define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
98#define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
99
100static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
101{
102 rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
103}
104
105static inline void rb_set_parent_color(struct rb_node *rb,
106 struct rb_node *p, int color)
107{
108 rb->__rb_parent_color = (unsigned long)p | color;
109}
110
111static inline void
112__rb_change_child(struct rb_node *old, struct rb_node *new,
113 struct rb_node *parent, struct rb_root *root)
114{
115 if (parent) {
116 if (parent->rb_left == old)
117 WRITE_ONCE(parent->rb_left, new);
118 else
119 WRITE_ONCE(parent->rb_right, new);
120 } else
121 WRITE_ONCE(root->rb_node, new);
122}
123
124extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
125 void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
126
127static __always_inline struct rb_node *
128__rb_erase_augmented(struct rb_node *node, struct rb_root *root,
129 const struct rb_augment_callbacks *augment)
130{
131 struct rb_node *child = node->rb_right;
132 struct rb_node *tmp = node->rb_left;
133 struct rb_node *parent, *rebalance;
134 unsigned long pc;
135
136 if (!tmp) {
137
138
139
140
141
142
143
144 pc = node->__rb_parent_color;
145 parent = __rb_parent(pc);
146 __rb_change_child(node, child, parent, root);
147 if (child) {
148 child->__rb_parent_color = pc;
149 rebalance = NULL;
150 } else
151 rebalance = __rb_is_black(pc) ? parent : NULL;
152 tmp = parent;
153 } else if (!child) {
154
155 tmp->__rb_parent_color = pc = node->__rb_parent_color;
156 parent = __rb_parent(pc);
157 __rb_change_child(node, tmp, parent, root);
158 rebalance = NULL;
159 tmp = parent;
160 } else {
161 struct rb_node *successor = child, *child2;
162
163 tmp = child->rb_left;
164 if (!tmp) {
165
166
167
168
169
170
171
172
173
174 parent = successor;
175 child2 = successor->rb_right;
176
177 augment->copy(node, successor);
178 } else {
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193 do {
194 parent = successor;
195 successor = tmp;
196 tmp = tmp->rb_left;
197 } while (tmp);
198 child2 = successor->rb_right;
199 WRITE_ONCE(parent->rb_left, child2);
200 WRITE_ONCE(successor->rb_right, child);
201 rb_set_parent(child, successor);
202
203 augment->copy(node, successor);
204 augment->propagate(parent, successor);
205 }
206
207 tmp = node->rb_left;
208 WRITE_ONCE(successor->rb_left, tmp);
209 rb_set_parent(tmp, successor);
210
211 pc = node->__rb_parent_color;
212 tmp = __rb_parent(pc);
213 __rb_change_child(node, successor, tmp, root);
214
215 if (child2) {
216 successor->__rb_parent_color = pc;
217 rb_set_parent_color(child2, parent, RB_BLACK);
218 rebalance = NULL;
219 } else {
220 unsigned long pc2 = successor->__rb_parent_color;
221 successor->__rb_parent_color = pc;
222 rebalance = __rb_is_black(pc2) ? parent : NULL;
223 }
224 tmp = successor;
225 }
226
227 augment->propagate(tmp, NULL);
228 return rebalance;
229}
230
231static __always_inline void
232rb_erase_augmented(struct rb_node *node, struct rb_root *root,
233 const struct rb_augment_callbacks *augment)
234{
235 struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
236 if (rebalance)
237 __rb_erase_color(rebalance, root, augment->rotate);
238}
239
240#endif
241