1 /*
2  * Copyright (C) 2015-2019 Alibaba Group Holding Limited
3  */
4 
5 #ifndef _K_RBTREE_H
6 #define _K_RBTREE_H
7 
8 #include <stdio.h>
9 #include <stddef.h>
10 #include <stdint.h>
11 #include <string.h>
12 #include <limits.h>
13 
14 struct k_rbtree_node_t {
15     unsigned long  rbt_parent_color;
16     struct k_rbtree_node_t *rbt_right;
17     struct k_rbtree_node_t *rbt_left;
18 } __attribute__((aligned(sizeof(long))));
19 
20 struct k_rbtree_root_t {
21     struct k_rbtree_node_t *rbt_node;
22 };
23 
24 struct k_rbtree_augment_callbacks {
25     void (*propagate)(struct k_rbtree_node_t *node, struct k_rbtree_node_t *stop);
26     void (*copy)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new_node);
27     void (*rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new_node);
28 };
29 
30 #define K_RBTREE_RED   0
31 #define K_RBTREE_BLACK 1
32 
33 #define RBT_ROOT (struct k_rbtree_root_t) { NULL, }
34 
35 #define __rbtree_parent(pcolor)    ((struct k_rbtree_node_t *)(pcolor & ~3))
36 #define __rbtree_color(pcolor)     ((pcolor) & 1)
37 #define __rbtree_is_black(pcolor)  __rbtree_color(pcolor)
38 #define __rbtree_is_red(pcolor)    (!__rbtree_color(pcolor))
39 
40 #define K_RBTREE_PARENT(rbt)      ((struct k_rbtree_node_t *)((rbt)->rbt_parent_color & ~3))
41 #define K_RBTREE_COLOR(rbt)       __rbtree_color((rbt)->rbt_parent_color)
42 #define K_RBTREE_IS_BLACK(rbt)    __rbtree_is_black((rbt)->rbt_parent_color)
43 #define K_RBTREE_IS_RED(rbt)      __rbtree_is_red((rbt)->rbt_parent_color)
44 
45 #define K_RBTREE_EMPTY_ROOT(root)  ((root)->rbt_node == NULL)
46 #define K_RBTREE_EMPTY_NODE(node)  \
47     ((node)->rbt_parent_color == (unsigned long)(node))
48 #define K_RBTREE_CLEAR_NODE(node)  \
49     ((node)->rbt_parent_color = (unsigned long)(node))
50 
51 #define k_rbtree_entry(ptr, type, member) ((type *)((uint8_t *)(ptr) - (size_t)(&((type *)0)->member)))
52 
53 #define k_rbtree_entry_safe(ptr, type, member) \
54     ({ ptr ? k_rbtree_entry(ptr, type, member) : NULL; })
55 
56 #define k_rbtree_postorder_for_each_entry_safe(pos, n, root, field) \
57     for (pos = k_rbtree_entry_safe(k_rbtree_first_postorder(root), typeof(*pos), field); \
58          pos && ({ n = k_rbtree_entry_safe(k_rbtree_next_postorder(&pos->field), \
59             typeof(*pos), field); 1; }); \
60          pos = n)
61 
62 extern void rbtree_insert_augmented(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
63                 void (*augment_rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new_node));
64 extern void rbtree_erase_color(struct k_rbtree_node_t *parent, struct k_rbtree_root_t *root,
65     void (*augment_rotate)(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new_node));
66 
rbtree_set_parent(struct k_rbtree_node_t * rbt,struct k_rbtree_node_t * p)67 static inline void rbtree_set_parent(struct k_rbtree_node_t *rbt, struct k_rbtree_node_t *p)
68 {
69     rbt->rbt_parent_color = K_RBTREE_COLOR(rbt) | (unsigned long)p;
70 }
71 
rbtree_set_parent_color(struct k_rbtree_node_t * rbt,struct k_rbtree_node_t * p,int color)72 static inline void rbtree_set_parent_color(struct k_rbtree_node_t *rbt,
73                        struct k_rbtree_node_t *p, int color)
74 {
75     rbt->rbt_parent_color = (unsigned long)p | color;
76 }
77 
rbtree_change_child(struct k_rbtree_node_t * old,struct k_rbtree_node_t * new_node,struct k_rbtree_node_t * parent,struct k_rbtree_root_t * root)78 static inline void rbtree_change_child(struct k_rbtree_node_t *old, struct k_rbtree_node_t *new_node,
79           struct k_rbtree_node_t *parent, struct k_rbtree_root_t *root)
80 {
81     if (parent) {
82         if (parent->rbt_left == old)
83             parent->rbt_left = new_node;
84         else
85             parent->rbt_right = new_node;
86     } else
87         root->rbt_node = new_node;
88 }
89 
rbtree_erase_augmented(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root,const struct k_rbtree_augment_callbacks * augment)90 static inline struct k_rbtree_node_t *rbtree_erase_augmented(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
91              const struct k_rbtree_augment_callbacks *augment)
92 {
93     struct k_rbtree_node_t *child = node->rbt_right, *tmp = node->rbt_left;
94     struct k_rbtree_node_t *parent, *rebalance;
95     unsigned long pc;
96 
97     if (!tmp) {
98         pc = node->rbt_parent_color;
99         parent = __rbtree_parent(pc);
100         rbtree_change_child(node, child, parent, root);
101         if (child) {
102             child->rbt_parent_color = pc;
103             rebalance = NULL;
104         } else
105             rebalance = __rbtree_is_black(pc) ? parent : NULL;
106         tmp = parent;
107     } else if (!child) {
108         tmp->rbt_parent_color = pc = node->rbt_parent_color;
109         parent = __rbtree_parent(pc);
110         rbtree_change_child(node, tmp, parent, root);
111         rebalance = NULL;
112         tmp = parent;
113     } else {
114         struct k_rbtree_node_t *successor = child, *child2;
115         tmp = child->rbt_left;
116         if (!tmp) {
117             parent = successor;
118             child2 = successor->rbt_right;
119             augment->copy(node, successor);
120         } else {
121             do {
122                 parent = successor;
123                 successor = tmp;
124                 tmp = tmp->rbt_left;
125             } while (tmp);
126             parent->rbt_left = child2 = successor->rbt_right;
127             successor->rbt_right = child;
128             rbtree_set_parent(child, successor);
129             augment->copy(node, successor);
130             augment->propagate(parent, successor);
131         }
132 
133         successor->rbt_left = tmp = node->rbt_left;
134         rbtree_set_parent(tmp, successor);
135 
136         pc = node->rbt_parent_color;
137         tmp = __rbtree_parent(pc);
138         rbtree_change_child(node, successor, tmp, root);
139         if (child2) {
140             successor->rbt_parent_color = pc;
141             rbtree_set_parent_color(child2, parent, K_RBTREE_BLACK);
142             rebalance = NULL;
143         } else {
144             unsigned long pc2 = successor->rbt_parent_color;
145             successor->rbt_parent_color = pc;
146             rebalance = __rbtree_is_black(pc2) ? parent : NULL;
147         }
148         tmp = successor;
149     }
150 
151     augment->propagate(tmp, NULL);
152     return rebalance;
153 }
154 
k_rbtree_erase_augmented(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root,const struct k_rbtree_augment_callbacks * augment)155 static inline void k_rbtree_erase_augmented(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
156            const struct k_rbtree_augment_callbacks *augment)
157 {
158     struct k_rbtree_node_t *rebalance = rbtree_erase_augmented(node, root, augment);
159     if (rebalance){
160         rbtree_erase_color(rebalance, root, augment->rotate);
161     }
162 }
163 
k_rbtree_insert_augmented(struct k_rbtree_node_t * node,struct k_rbtree_root_t * root,const struct k_rbtree_augment_callbacks * augment)164 static inline void k_rbtree_insert_augmented(struct k_rbtree_node_t *node, struct k_rbtree_root_t *root,
165             const struct k_rbtree_augment_callbacks *augment)
166 {
167     rbtree_insert_augmented(node, root, augment->rotate);
168 }
169 
k_rbtree_link_node(struct k_rbtree_node_t * node,struct k_rbtree_node_t * parent,struct k_rbtree_node_t ** rbt_link)170 static inline void k_rbtree_link_node(struct k_rbtree_node_t * node, struct k_rbtree_node_t * parent,
171                 struct k_rbtree_node_t ** rbt_link)
172 {
173     node->rbt_parent_color = (unsigned long)parent;
174     node->rbt_left = node->rbt_right = NULL;
175     *rbt_link = node;
176 }
177 
178 extern void k_rbtree_insert_color(struct k_rbtree_node_t *, struct k_rbtree_root_t *);
179 extern void k_rbtree_erase(struct k_rbtree_node_t *, struct k_rbtree_root_t *);
180 extern struct k_rbtree_node_t *k_rbtree_next(const struct k_rbtree_node_t *);
181 extern struct k_rbtree_node_t *k_rbtree_prev(const struct k_rbtree_node_t *);
182 extern struct k_rbtree_node_t *k_rbtree_first(const struct k_rbtree_root_t *);
183 extern struct k_rbtree_node_t *k_rbtree_last(const struct k_rbtree_root_t *);
184 extern struct k_rbtree_node_t *k_rbtree_first_postorder(const struct k_rbtree_root_t *);
185 extern struct k_rbtree_node_t *k_rbtree_next_postorder(const struct k_rbtree_node_t *);
186 extern void k_rbtree_replace_node(struct k_rbtree_node_t *victim, struct k_rbtree_node_t *new_node, struct k_rbtree_root_t *root);
187 
188 #endif /* _K_RBTREE_H */
189 
190