1 /*
2 * Copyright (c) 2019 LK Trusty Authors. All Rights Reserved.
3 *
4 * Permission is hereby granted, free of charge, to any person obtaining
5 * a copy of this software and associated documentation files
6 * (the "Software"), to deal in the Software without restriction,
7 * including without limitation the rights to use, copy, modify, merge,
8 * publish, distribute, sublicense, and/or sell copies of the Software,
9 * and to permit persons to whom the Software is furnished to do so,
10 * subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be
13 * included in all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
16 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
17 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
18 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
19 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
20 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
21 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
22 */
23
24 #pragma once
25
26 #include <assert.h>
27 #include <lk/compiler.h>
28 #include <lk/macros.h>
29 #include <stdbool.h>
30 #include <stddef.h>
31
32 __BEGIN_CDECLS
33
34 /*
35 * lk defines DEBUG_ASSERT for debug build asserts. Enable those checks for
36 * host test as well.
37 */
38 #ifndef DEBUG_ASSERT
39 #define DEBUG_ASSERT(e) assert(e)
40 #endif
41
42 /**
43 * struct bst_node - Node in binary search tree.
44 * @rank: Rank value used for rebalancing. 0 indicates node is not in a tree.
45 * @parent: Pointer to parent node or %NULL for root node.
46 * @child: Array of pointers to child nodes. @child[0] points to the left child
47 * and @child[1] points to the right child.
48 */
49 struct bst_node {
50 size_t rank;
51 struct bst_node *parent;
52 struct bst_node *child[2];
53 };
54
55 struct bst_root {
56 struct bst_node *root;
57 };
58
59 #define BST_NODE_INITIAL_VALUE {0, NULL, {NULL, NULL}}
60 #define BST_ROOT_INITIAL_VALUE {NULL}
61
bst_node_initialize(struct bst_node * node)62 static inline void bst_node_initialize(struct bst_node *node) {
63 /* Set rank to an invalid value to detect double insertion. */
64 node->rank = 0;
65 }
66
bst_root_initialize(struct bst_root * root)67 static inline void bst_root_initialize(struct bst_root *root) {
68 root->root = NULL;
69 }
70
71 /**
72 * bst_is_empty - Check if the binary search tree is empty.
73 * @root: Tree to check
74 *
75 * Return: %true if there are no nodes in the tree, %false otherwise.
76 */
bst_is_empty(struct bst_root * root)77 static inline bool bst_is_empty(struct bst_root *root) {
78 return !root->root;
79 }
80
81 /**
82 * bst_compare_t - Compare function provided by caller
83 * @a: First node to compare
84 * @b: Second node to compare
85 *
86 * Return: a positive number if @b should be after @a, 0 if @b is a match for
87 * @a, a negative otherwise.
88 */
89 typedef int (*bst_compare_t)(struct bst_node *a, struct bst_node *b);
90
91 /**
92 * bst_compare_key_t - Compare node to key function provided by caller
93 * @a: First node to compare
94 * @b: Opaque key pointer for second node to compare
95 *
96 * Return: a positive number if @b should be after @a, 0 if @b is a match for
97 * @a, a negative otherwise.
98 */
99 typedef int (*bst_compare_key_t)(const struct bst_node *a, const void *b);
100
101 /**
102 * bst_search - Find a node in a binary search tree.
103 * @root: Tree to search
104 * @node: Dummy node containing key to search for.
105 * @compare: Compare function.
106 *
107 * Find a node in a binary search tree. Use bst_search_type instead to get a
108 * pointer to the struct that contains @node.
109 *
110 * Note that if there are multiple matching nodes in the tree, the node returned
111 * may not be the leftmost matching node.
112 *
113 * Return: Node in @root matching @node, or %NULL if no matching node is found.
114 */
bst_search(const struct bst_root * root,struct bst_node * node,bst_compare_t compare)115 static inline struct bst_node *bst_search(const struct bst_root *root,
116 struct bst_node *node,
117 bst_compare_t compare) {
118 DEBUG_ASSERT(root);
119 DEBUG_ASSERT(node);
120 DEBUG_ASSERT(compare);
121
122 struct bst_node *tree_node = root->root;
123 while (tree_node) {
124 int cmp = compare(tree_node, node);
125 if (!cmp) {
126 return tree_node;
127 }
128 tree_node = tree_node->child[cmp > 0];
129 }
130 return NULL;
131 }
132
133 /**
134 * bst_search_key - Find a node in a binary search tree.
135 * @root: Tree to search
136 * @key: Pointer to a key to search for.
137 * @compare_key: Compare function. See &typedef bst_compare_key_t.
138 *
139 * Find a node in a binary search tree. Use bst_search_key_type instead to get a
140 * pointer to the struct that contains @node.
141 *
142 * Note that if there are multiple matching nodes in the tree, the node returned
143 * may not be the leftmost matching node.
144 *
145 * Return: Node in @root matching @node, or %NULL if no matching node is found.
146 */
bst_search_key(const struct bst_root * root,const void * key,bst_compare_key_t compare_key)147 static inline struct bst_node *bst_search_key(const struct bst_root *root,
148 const void *key,
149 bst_compare_key_t compare_key) {
150 DEBUG_ASSERT(root);
151 DEBUG_ASSERT(key);
152 DEBUG_ASSERT(compare_key);
153
154 struct bst_node *tree_node = root->root;
155 while (tree_node) {
156 int cmp = compare_key(tree_node, key);
157 if (!cmp) {
158 return tree_node;
159 }
160 tree_node = tree_node->child[cmp > 0];
161 }
162 return NULL;
163 }
164
165
166 /**
167 * bst_search_type - Find an item in a binary search tree.
168 * @root: Tree to search
169 * @item: Dummy item containing key to search for.
170 * @compare: Compare function.
171 * @type: Type of @item.
172 * @member: Name of struct bst_node embedded in @type.
173 *
174 * Return: Item in @root matching @item, or %NULL if no matching node is found.
175 */
176 #define bst_search_type(root, item, compare, type, member) ({ \
177 type *__item = (item); \
178 containerof_null_safe(bst_search(root, &__item->member, compare), type, \
179 member); \
180 })
181
182 /**
183 * bst_search_key_type - Find an item in a binary search tree.
184 * @root: Tree to search
185 * @key: Pointer to key to search for.
186 * @compare_key: Compare function.
187 * @type: Type to return.
188 * @member: Name of struct bst_node embedded in @type.
189 *
190 * Return: Item in @root matching @key, or %NULL if no matching node is found.
191 */
192 #define bst_search_key_type(root, key, compare_key, type, member) \
193 containerof_null_safe(bst_search_key(root, key, compare_key), type, member);
194
195 /* Internal helper. Don't call directly */
196 void bst_update_rank_insert(struct bst_root *root, struct bst_node *node);
197
198 /**
199 * bst_insert - Insert node in tree.
200 * @root: Tree.
201 * @node: Node to insert.
202 * @compare: Compare function.
203 *
204 * Insert @node in @root.
205 *
206 * Return: %true if @node was inserted. %false if a node matching @node is
207 * already in @root.
208 */
bst_insert(struct bst_root * root,struct bst_node * node,bst_compare_t compare)209 static inline bool bst_insert(struct bst_root *root, struct bst_node *node,
210 bst_compare_t compare) {
211 DEBUG_ASSERT(root);
212 DEBUG_ASSERT(node);
213 DEBUG_ASSERT(compare);
214 DEBUG_ASSERT(!node->rank);
215
216 struct bst_node *parent = NULL;
217 struct bst_node **parent_ptr = &root->root;
218 int diff;
219 bool is_right_child = false;
220 while (true) {
221 struct bst_node *tree_node = *parent_ptr;
222 if (!tree_node) {
223 node->rank = 1;
224 node->parent = parent;
225 node->child[0] = NULL;
226 node->child[1] = NULL;
227 *parent_ptr = node;
228 bst_update_rank_insert(root, node);
229 return true;
230 }
231 diff = compare(tree_node, node);
232 if (!diff) {
233 return false;
234 }
235 is_right_child = diff > 0;
236 parent_ptr = &tree_node->child[is_right_child];
237 parent = tree_node;
238 }
239 }
240
241 /**
242 * bst_delete - Remove node from tree.
243 * @root: Tree.
244 * @node: Node to delete
245 *
246 * Delete @node from @root.
247 */
248 void bst_delete(struct bst_root *root, struct bst_node *node);
249
250 /**
251 * bst_prev - Get previous node.
252 * @root: Tree.
253 * @node: Node to move from.
254 *
255 * Use bst_prev_type instead to use pointers to the struct that contains @node.
256 *
257 * Return: If @node is %NULL, right-most node in @root.
258 * If @node is not %NULL, right-most node to the left of @node.
259 * %NULL if the node described above does not exist.
260 */
261 struct bst_node *bst_prev(struct bst_root *root, struct bst_node *node);
262
263 /**
264 * bst_prev_type - Get previous item.
265 * @root: Tree.
266 * @item: Item to move from.
267 * @type: Type of @item.
268 * @member: Name of struct bst_node embedded in @type.
269 *
270 * Return: If @item is %NULL, right-most item in @root.
271 * If @item is not %NULL, right-most item to the left of @item.
272 * %NULL if the item described above does not exist.
273 */
274 #define bst_prev_type(root, item, type, member) \
275 containerof_null_safe(bst_prev(root, item), type, member)
276
277 /**
278 * bst_next - Get next node.
279 * @root: Tree.
280 * @node: Node to move from.
281 *
282 * Use bst_next_type instead to use pointers to the struct that contains @node.
283 *
284 * Return: If @node is %NULL, left-most node in @root.
285 * If @node is not %NULL, left-most node to the right of @node.
286 * %NULL if the node described above does not exist.
287 */
288 struct bst_node *bst_next(const struct bst_root *root, struct bst_node *node);
289
290 /**
291 * bst_next_type - Get previous item.
292 * @root: Tree.
293 * @item: Item to move from.
294 * @type: Type of @item.
295 * @member: Name of struct bst_node embedded in @type.
296 *
297 * Return: If @item is %NULL, left-most item in @root.
298 * If @item is not %NULL, left-most item to the right of @item.
299 * %NULL if the item described above does not exist.
300 */
301 #define bst_next_type(root, item, type, member) \
302 containerof_null_safe(bst_next(root, item), type, member)
303
304 /**
305 * bst_for_every_entry - Loop over every entry in a tree.
306 * @root: Tree.
307 * @entry: Entry variable used by loop body.
308 * @type: Type of @entry.
309 * @member: Name of struct bst_node embedded in @type.
310 *
311 * Loop over every node in @root, convert that node to @type and provide it as
312 * @entry to the loop body directly following this macro.
313 *
314 * It is safe to delete @entry from @root in the body if the loop, but it is not
315 * safe to delete any other nodes or insert any nodes.
316 */
317 #define bst_for_every_entry(root, entry, type, member) \
318 for (struct bst_node *_bst_for_every_cursor = bst_next(root, NULL); \
319 (_bst_for_every_cursor != NULL) && \
320 ((entry) = containerof(_bst_for_every_cursor, type, member)) && \
321 ((_bst_for_every_cursor = bst_next(root, _bst_for_every_cursor)) \
322 || true);)
323
324 /* Internal helper. Don't call directly */
325 void bst_delete_all_helper(struct bst_root *root, struct bst_node *node);
326
327 /**
328 * bst_for_every_entry_delete - Loop over tree and delete every entry.
329 * @root: Tree.
330 * @entry: Entry variable used by loop body.
331 * @type: Type of @entry.
332 * @member: Name of struct bst_node embedded in @type.
333 *
334 * Loop over every node in @root, convert that node to @type and provide it as
335 * @entry to the loop body directly following this macro.
336 *
337 * @entry will be removed from @root before entering the loop bode. It is not
338 * safe to delete any other nodes or insert any nodes.
339 */
340 #define bst_for_every_entry_delete(root, entry, type, member) \
341 for (struct bst_node *_bst_for_every_cursor = bst_next(root, NULL); \
342 (_bst_for_every_cursor != NULL) && ({\
343 (entry) = containerof(_bst_for_every_cursor, type, member); \
344 _bst_for_every_cursor = bst_next(root, _bst_for_every_cursor); \
345 bst_delete_all_helper(root, &(entry)->member); true;});)
346
347 __END_CDECLS
348