1 /*
2  * Copyright (C) 2015-2019 Alibaba Group Holding Limited
3  */
4 #include "k_rbtree.h"
5 
rbtree_set_black(struct k_rbtree_node_t * rbt)6 static inline void rbtree_set_black(struct k_rbtree_node_t *rbt)
7 {
8     rbt->rbt_parent_color |= K_RBTREE_BLACK;
9 }
10 
rbtree_rotate_set_parents(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new,struct k_rbtree_root_t * root,int color)11 static inline void rbtree_rotate_set_parents(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new,
12             struct k_rbtree_root_t *root, int color)
13 {
14     struct k_rbtree_node_t *parent = K_RBTREE_PARENT(old);
15     new->rbt_parent_color = old->rbt_parent_color;
16     rbtree_set_parent_color(old, new, color);
17     rbtree_change_child(old, new, parent, root);
18 }
19 
rbtree_insert(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root,void (* augment_rotate)(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new))20 static void rbtree_insert(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
21         void (*augment_rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new))
22 {
23     struct k_rbtree_node_t *parent = (struct k_rbtree_node_t *)node->rbt_parent_color, *gparent, *tmp;
24 
25     while (1) {
26         if (!parent) {
27             rbtree_set_parent_color(node, NULL, K_RBTREE_BLACK);
28             break;
29         } else if (K_RBTREE_IS_BLACK(parent)){
30             break;
31         }
32 
33         gparent = (struct k_rbtree_node_t *)parent->rbt_parent_color;
34 
35         tmp = gparent->rbt_right;
36         if (parent != tmp) {
37             if (tmp && K_RBTREE_IS_RED(tmp)) {
38                 rbtree_set_parent_color(tmp, gparent, K_RBTREE_BLACK);
39                 rbtree_set_parent_color(parent, gparent, K_RBTREE_BLACK);
40                 node = gparent;
41                 parent = K_RBTREE_PARENT(node);
42                 rbtree_set_parent_color(node, parent, K_RBTREE_RED);
43                 continue;
44             }
45 
46             tmp = parent->rbt_right;
47             if (node == tmp) {
48                 parent->rbt_right = tmp = node->rbt_left;
49                 node->rbt_left = parent;
50                 if (tmp){
51                     rbtree_set_parent_color(tmp, parent,
52                                 K_RBTREE_BLACK);
53                 }
54                 rbtree_set_parent_color(parent, node, K_RBTREE_RED);
55                 augment_rotate(parent, node);
56                 parent = node;
57                 tmp = node->rbt_right;
58             }
59 
60             gparent->rbt_left = tmp;
61             parent->rbt_right = gparent;
62             if (tmp){
63                 rbtree_set_parent_color(tmp, gparent, K_RBTREE_BLACK);
64             }
65             rbtree_rotate_set_parents(gparent, parent, root, K_RBTREE_RED);
66             augment_rotate(gparent, parent);
67             break;
68         } else {
69             tmp = gparent->rbt_left;
70             if (tmp && K_RBTREE_IS_RED(tmp)) {
71                 rbtree_set_parent_color(tmp, gparent, K_RBTREE_BLACK);
72                 rbtree_set_parent_color(parent, gparent, K_RBTREE_BLACK);
73                 node = gparent;
74                 parent = K_RBTREE_PARENT(node);
75                 rbtree_set_parent_color(node, parent, K_RBTREE_RED);
76                 continue;
77             }
78 
79             tmp = parent->rbt_left;
80             if (node == tmp) {
81                 parent->rbt_left = tmp = node->rbt_right;
82                 node->rbt_right = parent;
83                 if (tmp){
84                     rbtree_set_parent_color(tmp, parent,
85                                 K_RBTREE_BLACK);
86                 }
87                 rbtree_set_parent_color(parent, node, K_RBTREE_RED);
88                 augment_rotate(parent, node);
89                 parent = node;
90                 tmp = node->rbt_left;
91             }
92 
93             gparent->rbt_right = tmp;
94             parent->rbt_left = gparent;
95             if (tmp){
96                 rbtree_set_parent_color(tmp, gparent, K_RBTREE_BLACK);
97             }
98             rbtree_rotate_set_parents(gparent, parent, root, K_RBTREE_RED);
99             augment_rotate(gparent, parent);
100             break;
101         }
102     }
103 }
104 
rbtree_erase_color(struct k_rbtree_node_t * parent,struct k_rbtree_root_t * root,void (* augment_rotate)(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new))105 void rbtree_erase_color(struct k_rbtree_node_t *parent, struct k_rbtree_root_t *root,
106                  void (*augment_rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new))
107 {
108     struct k_rbtree_node_t *node = NULL, *sibling, *tmp1, *tmp2;
109 
110     while (1) {
111         sibling = parent->rbt_right;
112         if (node != sibling) {
113             if (K_RBTREE_IS_RED(sibling)) {
114                 parent->rbt_right = tmp1 = sibling->rbt_left;
115                 sibling->rbt_left = parent;
116                 rbtree_set_parent_color(tmp1, parent, K_RBTREE_BLACK);
117                 rbtree_rotate_set_parents(parent, sibling, root,
118                             K_RBTREE_RED);
119                 augment_rotate(parent, sibling);
120                 sibling = tmp1;
121             }
122             tmp1 = sibling->rbt_right;
123             if (!tmp1 || K_RBTREE_IS_BLACK(tmp1)) {
124                 tmp2 = sibling->rbt_left;
125                 if (!tmp2 || K_RBTREE_IS_BLACK(tmp2)) {
126                     rbtree_set_parent_color(sibling, parent,
127                                 K_RBTREE_RED);
128                     if (K_RBTREE_IS_RED(parent)){
129                         rbtree_set_black(parent);
130                     }
131                     else {
132                         node = parent;
133                         parent = K_RBTREE_PARENT(node);
134                         if (parent){
135                             continue;
136                         }
137                     }
138                     break;
139                 }
140                 sibling->rbt_left = tmp1 = tmp2->rbt_right;
141                 tmp2->rbt_right = sibling;
142                 parent->rbt_right = tmp2;
143                 if (tmp1){
144                     rbtree_set_parent_color(tmp1, sibling,
145                                 K_RBTREE_BLACK);
146                 }
147                 augment_rotate(sibling, tmp2);
148                 tmp1 = sibling;
149                 sibling = tmp2;
150             }
151             parent->rbt_right = tmp2 = sibling->rbt_left;
152             sibling->rbt_left = parent;
153             rbtree_set_parent_color(tmp1, sibling, K_RBTREE_BLACK);
154             if (tmp2){
155                 rbtree_set_parent(tmp2, parent);
156             }
157             rbtree_rotate_set_parents(parent, sibling, root,
158                         K_RBTREE_BLACK);
159             augment_rotate(parent, sibling);
160             break;
161         } else {
162             sibling = parent->rbt_left;
163             if (K_RBTREE_IS_RED(sibling)) {
164                 parent->rbt_left = tmp1 = sibling->rbt_right;
165                 sibling->rbt_right = parent;
166                 rbtree_set_parent_color(tmp1, parent, K_RBTREE_BLACK);
167                 rbtree_rotate_set_parents(parent, sibling, root,
168                             K_RBTREE_RED);
169                 augment_rotate(parent, sibling);
170                 sibling = tmp1;
171             }
172             tmp1 = sibling->rbt_left;
173             if (!tmp1 || K_RBTREE_IS_BLACK(tmp1)) {
174                 tmp2 = sibling->rbt_right;
175                 if (!tmp2 || K_RBTREE_IS_BLACK(tmp2)) {
176                     rbtree_set_parent_color(sibling, parent,
177                                 K_RBTREE_RED);
178                     if (K_RBTREE_IS_RED(parent)){
179                         rbtree_set_black(parent);
180                     }
181                     else {
182                         node = parent;
183                         parent = K_RBTREE_PARENT(node);
184                         if (parent){
185                             continue;
186                         }
187                     }
188                     break;
189                 }
190                 sibling->rbt_right = tmp1 = tmp2->rbt_left;
191                 tmp2->rbt_left = sibling;
192                 parent->rbt_left = tmp2;
193                 if (tmp1){
194                     rbtree_set_parent_color(tmp1, sibling,
195                                 K_RBTREE_BLACK);
196                 }
197                 augment_rotate(sibling, tmp2);
198                 tmp1 = sibling;
199                 sibling = tmp2;
200             }
201             parent->rbt_left = tmp2 = sibling->rbt_right;
202             sibling->rbt_right = parent;
203             rbtree_set_parent_color(tmp1, sibling, K_RBTREE_BLACK);
204             if (tmp2){
205                 rbtree_set_parent(tmp2, parent);
206             }
207             rbtree_rotate_set_parents(parent, sibling, root,
208                         K_RBTREE_BLACK);
209             augment_rotate(parent, sibling);
210             break;
211         }
212     }
213 }
214 
rbtree_dummy_propagate(struct k_rbtree_node_t * node,struct k_rbtree_node_t * stop)215 static inline void rbtree_dummy_propagate(struct k_rbtree_node_t *node, struct k_rbtree_node_t *stop) {}
rbtree_dummy_copy(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new)216 static inline void rbtree_dummy_copy(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new) {}
rbtree_dummy_rotate(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new)217 static inline void rbtree_dummy_rotate(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new) {}
218 
219 static const struct k_rbtree_augment_callbacks dummy_callbacks = {
220     rbtree_dummy_propagate, rbtree_dummy_copy, rbtree_dummy_rotate
221 };
222 
rbtree_insert_augmented(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root,void (* augment_rotate)(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new))223 void rbtree_insert_augmented(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
224          void (*augment_rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new))
225 {
226     rbtree_insert(node, root, augment_rotate);
227 }
228 
k_rbtree_insert_color(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root)229 void k_rbtree_insert_color(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root)
230 {
231     rbtree_insert(node, root, rbtree_dummy_rotate);
232 }
233 
k_rbtree_erase(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root)234 void k_rbtree_erase(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root)
235 {
236     struct k_rbtree_node_t *rebalance;
237     rebalance = rbtree_erase_augmented(node, root, &dummy_callbacks);
238     if (rebalance){
239         rbtree_erase_color(rebalance, root, rbtree_dummy_rotate);
240     }
241 }
242 
k_rbtree_first(const struct k_rbtree_root_t * root)243 struct k_rbtree_node_t *k_rbtree_first(const struct k_rbtree_root_t *root)
244 {
245     struct k_rbtree_node_t *n;
246 
247     n = root->rbt_node;
248     if (!n){
249         return NULL;
250     }
251     while (n->rbt_left){
252         n = n->rbt_left;
253     }
254     return n;
255 }
256 
k_rbtree_last(const struct k_rbtree_root_t * root)257 struct k_rbtree_node_t *k_rbtree_last(const struct k_rbtree_root_t *root)
258 {
259     struct k_rbtree_node_t    *n;
260 
261     n = root->rbt_node;
262     if (!n){
263         return NULL;
264     }
265     while (n->rbt_right){
266         n = n->rbt_right;
267     }
268     return n;
269 }
270 
k_rbtree_next(const struct k_rbtree_node_t * node)271 struct k_rbtree_node_t *k_rbtree_next(const struct k_rbtree_node_t *node)
272 {
273     struct k_rbtree_node_t *parent;
274 
275     if (K_RBTREE_EMPTY_NODE(node)){
276         return NULL;
277     }
278 
279     if (node->rbt_right) {
280         node = node->rbt_right;
281         while (node->rbt_left){
282             node=node->rbt_left;
283         }
284         return (struct k_rbtree_node_t *)node;
285     }
286 
287     while ((parent = K_RBTREE_PARENT(node)) && node == parent->rbt_right){
288         node = parent;
289     }
290 
291     return parent;
292 }
293 
k_rbtree_prev(const struct k_rbtree_node_t * node)294 struct k_rbtree_node_t *k_rbtree_prev(const struct k_rbtree_node_t *node)
295 {
296     struct k_rbtree_node_t *parent;
297 
298     if (K_RBTREE_EMPTY_NODE(node)){
299         return NULL;
300     }
301 
302     if (node->rbt_left) {
303         node = node->rbt_left;
304         while (node->rbt_right){
305             node=node->rbt_right;
306         }
307         return (struct k_rbtree_node_t *)node;
308     }
309 
310     while ((parent = K_RBTREE_PARENT(node)) && node == parent->rbt_left){
311         node = parent;
312     }
313 
314     return parent;
315 }
316 
k_rbtree_replace_node(struct k_rbtree_node_t * victim,struct k_rbtree_node_t * new,struct k_rbtree_root_t * root)317 void k_rbtree_replace_node(struct k_rbtree_node_t *victim, struct k_rbtree_node_t *new,
318              struct k_rbtree_root_t *root)
319 {
320     struct k_rbtree_node_t *parent = K_RBTREE_PARENT(victim);
321 
322     rbtree_change_child(victim, new, parent, root);
323     if (victim->rbt_left){
324         rbtree_set_parent(victim->rbt_left, new);
325     }
326     if (victim->rbt_right){
327         rbtree_set_parent(victim->rbt_right, new);
328     }
329 
330     *new = *victim;
331 }
332 
k_rbtree_left_deepest_node(const struct k_rbtree_node_t * node)333 static struct k_rbtree_node_t *k_rbtree_left_deepest_node(const struct k_rbtree_node_t *node)
334 {
335     for (;;) {
336         if (node->rbt_left){
337             node = node->rbt_left;
338         }
339         else if (node->rbt_right){
340             node = node->rbt_right;
341         }
342         else{
343             return (struct k_rbtree_node_t *)node;
344         }
345     }
346 }
347 
k_rbtree_next_postorder(const struct k_rbtree_node_t * node)348 struct k_rbtree_node_t *k_rbtree_next_postorder(const struct k_rbtree_node_t *node)
349 {
350     const struct k_rbtree_node_t *parent;
351     if (!node){
352         return NULL;
353     }
354 
355     parent = K_RBTREE_PARENT(node);
356 
357     if (parent && node == parent->rbt_left && parent->rbt_right) {
358         return k_rbtree_left_deepest_node(parent->rbt_right);
359     } else{
360         return (struct k_rbtree_node_t *)parent;
361     }
362 }
363 
k_rbtree_first_postorder(const struct k_rbtree_root_t * root)364 struct k_rbtree_node_t *k_rbtree_first_postorder(const struct k_rbtree_root_t *root)
365 {
366     if (!root->rbt_node){
367         return NULL;
368     }
369 
370     return k_rbtree_left_deepest_node(root->rbt_node);
371 }
372 
373