1 /*
2 * Copyright (c) 2023 Institute of Parallel And Distributed Systems (IPADS), Shanghai Jiao Tong University (SJTU)
3 * Licensed under the Mulan PSL v2.
4 * You can use this software according to the terms and conditions of the Mulan PSL v2.
5 * You may obtain a copy of Mulan PSL v2 at:
6 * http://license.coscl.org.cn/MulanPSL2
7 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
8 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
9 * PURPOSE.
10 * See the Mulan PSL v2 for more details.
11 */
12
13 #include <stdbool.h>
14 #include <chcore/container/rbtree_plus.h>
15 #include <chcore/container/rbtree_impl.h>
16
17 struct comp_rb_key_wrapper_struct {
18 const void *user_key;
19 comp_rbp_key_func user_comp;
20 };
21
22 /**
23 * @brief Wrapper around user provided @key and compare function of type
24 * `comp_rbp_key_func`. Users just need to implement their compare function
25 * comparing their custom key and a rbp_node. But internally, we leverage
26 * existing rbtree search function to search a rbp_node containing @key given by
27 * user, which requires a compare function of type `comp_rb_key_func`. So we
28 * implement this wrapper to convert the user provided compare function to the
29 * required `comp_rb_key_func` using data structure comp_rb_key_wrapper_struct.
30 *
31 * @param key A pointer to a comp_rb_key_wrapper_struct struct, which contains
32 * user provided real key and compare function.
33 * @param rb_node
34 * @return int
35 */
comp_rb_key_wrapper(const void * key,const struct rb_node * rb_node)36 static int comp_rb_key_wrapper(const void *key, const struct rb_node *rb_node)
37 {
38 struct comp_rb_key_wrapper_struct *wrapper =
39 (struct comp_rb_key_wrapper_struct *)key;
40 struct rbp_node *rbp_node = rb_entry(rb_node, struct rbp_node, rbnode);
41
42 return wrapper->user_comp(wrapper->user_key, rbp_node);
43 }
44
init_rbtree_plus(struct rbtree_plus * this)45 void init_rbtree_plus(struct rbtree_plus *this)
46 {
47 init_rb_root(&this->root);
48 init_list_head(&this->head);
49 }
50
rbp_search(struct rbtree_plus * this,const void * key,comp_rbp_key_func cmp)51 struct rbp_node *rbp_search(struct rbtree_plus *this, const void *key,
52 comp_rbp_key_func cmp)
53 {
54 struct rb_node *ret;
55 /**
56 * Construct a wrapper struct with user provided key and compare
57 * function. See comp_rbp_key_func for more details.
58 */
59 struct comp_rb_key_wrapper_struct wrapper = {
60 .user_key = key,
61 .user_comp = cmp,
62 };
63
64 ret = rb_search(&this->root, &wrapper, comp_rb_key_wrapper);
65
66 /**
67 * !! Attention: if ret is NULL, it means the key is not found in the
68 * tree. We should return NULL too instead of invoke rb_entry on ret,
69 * which would generate a non-NULL but invalid (negative in fact)
70 * pointer.
71 */
72 if (ret == NULL) {
73 return NULL;
74 }
75
76 return rb_entry(ret, struct rbp_node, rbnode);
77 }
78
rbp_search_first(struct rbtree_plus * this,const void * key,comp_rbp_key_func cmp)79 struct rbp_node *rbp_search_first(struct rbtree_plus *this, const void *key,
80 comp_rbp_key_func cmp)
81 {
82 struct rb_node *ret;
83 struct comp_rb_key_wrapper_struct wrapper = {
84 .user_key = key,
85 .user_comp = cmp,
86 };
87
88 ret = rb_search_first(&this->root, &wrapper, comp_rb_key_wrapper);
89
90 /**
91 * !! Attention: if ret is NULL, it means the key is not found in the
92 * tree. We should return NULL too instead of invoke rb_entry on ret,
93 * which would generate a non-NULL but invalid (negative in fact)
94 * pointer.
95 */
96 if (ret == NULL) {
97 return NULL;
98 }
99
100 return rb_entry(ret, struct rbp_node, rbnode);
101 }
102
rbp_search_nearest(struct rbtree_plus * this,const void * key,comp_rbp_key_func cmp)103 struct rbp_node *rbp_search_nearest(struct rbtree_plus *this, const void *key,
104 comp_rbp_key_func cmp)
105 {
106 struct rbp_node *cur_rbp_node, *ret = NULL;
107 struct rb_node *cur = this->root.root_node;
108 int cmp_ret;
109 while (cur) {
110 cur_rbp_node = rb_entry(cur, struct rbp_node, rbnode);
111 cmp_ret = cmp(key, cur_rbp_node);
112 ret = cur_rbp_node;
113 if (cmp_ret < 0) {
114 cur = cur->left_child;
115 } else if (cmp_ret > 0) {
116 cur = cur->right_child;
117 } else {
118 break;
119 }
120 }
121
122 return ret;
123 }
124
rbp_insert(struct rbtree_plus * this,struct rbp_node * data,less_rbp_func less)125 void rbp_insert(struct rbtree_plus *this, struct rbp_node *data,
126 less_rbp_func less)
127 {
128 struct rb_node **new_link = &this->root.root_node;
129 struct rb_node *parent = NULL, *cur = this->root.root_node;
130 bool insert_at_left = false;
131 struct rbp_node *parent_rbp_node;
132
133 while (cur) {
134 if (less(data, rb_entry(cur, struct rbp_node, rbnode))) {
135 parent = cur;
136 new_link = &cur->left_child;
137 cur = cur->left_child;
138 insert_at_left = true;
139 } else {
140 parent = cur;
141 new_link = &cur->right_child;
142 cur = cur->right_child;
143 insert_at_left = false;
144 }
145 }
146
147 /**
148 * After the binary search logic above, we have the following
149 * invariants:
150 *
151 * @parent points to parent node of the new rbp_node(if the tree is not
152 * empty).
153 * @insert_at_left indicates whether the new rbp_node should be inserted
154 * as the left child of @parent or right child of @parent.
155 */
156 if (list_empty(&this->head)) {
157 /**
158 * If the tree is empty, the new rbp_node is the root node of
159 * rbtree plus. We just insert it as the first node of the list.
160 */
161 list_add(&data->list_node, &this->head);
162 } else {
163 /**
164 * If the tree is not empty, we need to insert the new rbp_node
165 * into an appropriate position in the list. Choosing the
166 * position is rely on
167 * @insert_at_left. If @insert_at_left is true, the key in @data
168 * is smaller than @parent and greater than the previous node of
169 * @parent in the list. So we just insert @data after
170 * @parent.prev. So as if @insert_at_left is false.
171 *
172 * Because there is a special list_head(this->head) in the list,
173 * the logic above is always correct even though there is only 1
174 * node in the rbtree plus.
175 */
176 parent_rbp_node = rb_entry(parent, struct rbp_node, rbnode);
177 if (insert_at_left) {
178 list_add(&data->list_node, parent_rbp_node->list_node.prev);
179 } else {
180 list_add(&data->list_node, &parent_rbp_node->list_node);
181 }
182 }
183
184 __rb_link_node(&data->rbnode, parent, new_link);
185 __rb_insert_color(&this->root, &data->rbnode);
186 }
187
rbp_erase(struct rbtree_plus * this,struct rbp_node * node)188 void rbp_erase(struct rbtree_plus *this, struct rbp_node *node)
189 {
190 list_del(&node->list_node);
191 rb_erase(&this->root, &node->rbnode);
192 }
193
rbp_replace_node(struct rbtree_plus * this,struct rbp_node * old,struct rbp_node * new)194 void rbp_replace_node(struct rbtree_plus *this, struct rbp_node *old,
195 struct rbp_node *new)
196 {
197 struct list_head *old_prev = old->list_node.prev;
198 rb_replace_node(&this->root, &old->rbnode, &new->rbnode);
199 list_del(&old->list_node);
200 list_add(&new->list_node, old_prev);
201 }
202
rbp_next(const struct rbtree_plus * this,const struct rbp_node * node)203 struct rbp_node *rbp_next(const struct rbtree_plus *this,
204 const struct rbp_node *node)
205 {
206 if (!node) {
207 return NULL;
208 }
209 if (node->list_node.next == &this->head) {
210 return NULL;
211 }
212 return list_entry(node->list_node.next, struct rbp_node, list_node);
213 }
214
rbp_prev(const struct rbtree_plus * this,const struct rbp_node * node)215 struct rbp_node *rbp_prev(const struct rbtree_plus *this,
216 const struct rbp_node *node)
217 {
218 if (!node) {
219 return NULL;
220 }
221 if (node->list_node.prev == &this->head) {
222 return NULL;
223 }
224 return list_entry(node->list_node.prev, struct rbp_node, list_node);
225 }
226
rbp_first(const struct rbtree_plus * this)227 struct rbp_node *rbp_first(const struct rbtree_plus *this)
228 {
229 if (list_empty(&this->head)) {
230 return NULL;
231 }
232 return list_entry(this->head.next, struct rbp_node, list_node);
233 }
234
rbp_last(const struct rbtree_plus * this)235 struct rbp_node *rbp_last(const struct rbtree_plus *this)
236 {
237 if (list_empty(&this->head)) {
238 return NULL;
239 }
240 return list_entry(this->head.prev, struct rbp_node, list_node);
241 }
242
rbp_next_match(const struct rbtree_plus * this,const struct rbp_node * node,const void * key,comp_rbp_key_func cmp)243 struct rbp_node *rbp_next_match(const struct rbtree_plus *this,
244 const struct rbp_node *node, const void *key,
245 comp_rbp_key_func cmp)
246 {
247 struct rbp_node *ret = rbp_next(this, node);
248 if (ret && cmp(key, ret) != 0) {
249 ret = NULL;
250 }
251 return ret;
252 }