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 #ifndef COMMON_RBTREE_H
13 #define COMMON_RBTREE_H
14
15 #include <common/types.h>
16 #include <common/macro.h>
17 #include <mm/kmalloc.h>
18 #include <common/kprint.h>
19
20 struct rb_node {
21 unsigned long __parent_color;
22 struct rb_node *right_child;
23 struct rb_node *left_child;
24 } __attribute__((aligned(sizeof(long))));
25 #define rb_entry(node, type, field) container_of(node, type, field)
26
27 /*
28 * wrapper type, distinguish between tree root and other nodes
29 * in rbtree interfaces.
30 * We have to differentiate rb_root from rb_node. For rb_node, users should
31 * embed it into themselves types, and retrieve pointer of their types from
32 * rb_node. However, users also need to hold a whole rbtree in some high-level
33 * types, e.g., mm_struct in Linux owns a rbtree storing vmr nodes. In the
34 * latter case, we should not let users embed a root rb_node in mm_struct,
35 * otherwise, rb_entry to that node would return a pointer of mm_struct. But
36 * users expect it to return a pointer of vmr. So we provide the rb_root type
37 * representing the whole rbtree, and users should embed this type into those
38 * high-level types like mm_struct.
39 */
40 struct rb_root {
41 struct rb_node *root_node;
42 };
43
44 #define RB_ROOT \
45 (struct rb_root) \
46 { \
47 NULL, \
48 }
49
init_rb_root(struct rb_root * root)50 static inline void init_rb_root(struct rb_root *root)
51 {
52 root->root_node = NULL;
53 }
54
55 /*
56 * compare functions defination
57 * We expect those functions have following semantics:
58 * cmp_xxx_func:
59 * return 0 if keys of left operand equal to right operand
60 * return <0 if left < right
61 * return >0 if left > right
62 * less_func:
63 * return true if left < right
64 * return false if left >= right
65 */
66 typedef int (*comp_node_func)(const struct rb_node *lhs,
67 const struct rb_node *rhs);
68 typedef int (*comp_key_func)(const void *key, const struct rb_node *node);
69 typedef bool (*less_func)(const struct rb_node *lhs, const struct rb_node *rhs);
70
71 /*
72 * search interfaces
73 */
74 /*
75 * rb_search - search for a node matching @key in rbtree
76 *
77 * @this: rbtree to be searched
78 * @key: pointer to user-specific key. We use void* to implement generality
79 * @cmp: user-provided compare function, comparing @key and rb_node in rbtree.
80 * Users should implement their logic to extract key from rb_node pointer.
81 *
82 * return:
83 * pointer of targeted node if @key exists in rbtree.
84 * NULL if no matching node is found.
85 */
86 struct rb_node *rb_search(struct rb_root *this, const void *key,
87 comp_key_func cmp);
88 /*
89 * rb_search - search for the **leftmost** node matching @key in rbtree
90 *
91 * @this: rbtree to be searched
92 * @key: pointer to user-specific key. We use void* to implement generality
93 * @cmp: user-provided compare function.
94 *
95 * return:
96 * pointer of targeted node if @key exists in rbtree.
97 * NULL if no matching node is found.
98 */
99 struct rb_node *rb_search_first(struct rb_root *this, const void *key,
100 comp_key_func cmp);
101
102 /*
103 * insert interfaces
104 */
105 /*
106 * rb_insert - insert a given @node into rbtree. @node
107 * will be inserted as right child of existing nodes if they have the same key.
108 *
109 * @this: rbtree to be inserted
110 * @node: node to be inserted. Users should embed rb_node in their custom types
111 * , allocate those types and provide pointers of rb_node in those types.
112 * @less: user-provided less function
113 *
114 */
115 void rb_insert(struct rb_root *this, struct rb_node *data, less_func less);
116
117 /*
118 * erase interfaces
119 */
120 /*
121 * rb_erase - remove @node from rbtree @this. Due to rb_node embedded in users'
122 * custom types, @node is not a pointer to head of a heap block. So rb_erase
123 * won't free memory of @node. User code should free memory using their types
124 * after erase rb_node contained in those types from @this.
125 *
126 * @this: rbtree to be erased
127 * @node: node to be erased
128 */
129 void rb_erase(struct rb_root *this, struct rb_node *node);
130
131 /*
132 * replace interfaces
133 */
134 /*
135 * rb_replace_node - replace @old with @new directly without rebalancing and
136 * recoloring. Users should make sure that keys of old and new are identical,
137 * otherwise the rbtree would be inconsistent.
138 *
139 * @this: rbtree to be replaced
140 * @old: old node in @this
141 * @new: new node to take over @old, should has identical key with @old.
142 */
143 void rb_replace_node(struct rb_root *this, struct rb_node *old,
144 struct rb_node *new);
145
146 /*
147 * traverse interfaces
148 */
149 /*
150 * rb_next - return the next node of @node in rbtree. If @node is the last node,
151 * return NULL.
152 */
153 struct rb_node *rb_next(const struct rb_node *node);
154 /*
155 * rb_prev - return the previous node of @node in rbtree. If @node is the first
156 * node, return NULL.
157 */
158 struct rb_node *rb_prev(const struct rb_node *node);
159 /*
160 * rb_first - return the first node of rbtree. If rbtree is empty, return NULL.
161 */
162 struct rb_node *rb_first(const struct rb_root *this);
163 /*
164 * rb_last - return the last node of rbtree. If rbtree is empty, return NULL.
165 */
166 struct rb_node *rb_last(const struct rb_root *this);
167 /*
168 * rb_next_match - return the next node of @node if it matches @key at the same
169 * time. Otherwise, return NULL. This function can be used to iterate over all
170 * nodes with a given key in rbtree.
171 */
172 struct rb_node *rb_next_match(const struct rb_node *node, const void *key,
173 comp_key_func cmp);
174
175 #define rb_for_each(root, node) \
176 for ((node) = rb_first(root); node; (node) = rb_next(node))
177
178 #define rb_key_for_each(root, node, key, cmp) \
179 for ((node) = rb_search_first(root, key, cmp); node; \
180 (node) = rb_next_match(node, key, cmp))
181
182 #endif /* COMMON_RBTREE_H */
183