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_compare_t - Compare function provided by caller
73  * @a: First node to compare
74  * @b: Second node to compare
75  *
76  * Return: a positive number if @b should be after @a, 0 if @b is a match for
77  * @a, a negative otherwise.
78  */
79 typedef int (*bst_compare_t)(struct bst_node *a, struct bst_node *b);
80 
81 /**
82  * bst_compare_key_t - Compare node to key function provided by caller
83  * @a: First node to compare
84  * @b: Opaque key pointer for 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_key_t)(const struct bst_node *a, const void *b);
90 
91 /**
92  * bst_search - Find a node in a binary search tree.
93  * @root:       Tree to search
94  * @node:       Dummy node containing key to search for.
95  * @compare:    Compare function.
96  *
97  * Find a node in a binary search tree. Use bst_search_type instead to get a
98  * pointer to the struct that contains @node.
99  *
100  * Note that if there are multiple matching nodes in the tree, the node returned
101  * may not be the leftmost matching node.
102  *
103  * Return: Node in @root matching @node, or %NULL if no matching node is found.
104  */
bst_search(const struct bst_root * root,struct bst_node * node,bst_compare_t compare)105 static inline struct bst_node *bst_search(const struct bst_root *root,
106                                           struct bst_node *node,
107                                           bst_compare_t compare) {
108     DEBUG_ASSERT(root);
109     DEBUG_ASSERT(node);
110     DEBUG_ASSERT(compare);
111 
112     struct bst_node *tree_node = root->root;
113     while (tree_node) {
114         int cmp = compare(tree_node, node);
115         if (!cmp) {
116             return tree_node;
117         }
118         tree_node = tree_node->child[cmp > 0];
119     }
120     return NULL;
121 }
122 
123 /**
124  * bst_search_key - Find a node in a binary search tree.
125  * @root:           Tree to search
126  * @key:            Pointer to a key to search for.
127  * @compare_key:    Compare function. See &typedef bst_compare_key_t.
128  *
129  * Find a node in a binary search tree. Use bst_search_key_type instead to get a
130  * pointer to the struct that contains @node.
131  *
132  * Note that if there are multiple matching nodes in the tree, the node returned
133  * may not be the leftmost matching node.
134  *
135  * Return: Node in @root matching @node, or %NULL if no matching node is found.
136  */
bst_search_key(const struct bst_root * root,const void * key,bst_compare_key_t compare_key)137 static inline struct bst_node *bst_search_key(const struct bst_root *root,
138                                               const void *key,
139                                               bst_compare_key_t compare_key) {
140     DEBUG_ASSERT(root);
141     DEBUG_ASSERT(key);
142     DEBUG_ASSERT(compare_key);
143 
144     struct bst_node *tree_node = root->root;
145     while (tree_node) {
146         int cmp = compare_key(tree_node, key);
147         if (!cmp) {
148             return tree_node;
149         }
150         tree_node = tree_node->child[cmp > 0];
151     }
152     return NULL;
153 }
154 
155 
156 /**
157  * bst_search_type - Find an item in a binary search tree.
158  * @root:       Tree to search
159  * @item:       Dummy item containing key to search for.
160  * @compare:    Compare function.
161  * @type:       Type of @item.
162  * @member:     Name of struct bst_node embedded in @type.
163  *
164  * Return: Item in @root matching @item, or %NULL if no matching node is found.
165  */
166 #define bst_search_type(root, item, compare, type, member) ({ \
167     type *__item = (item); \
168     containerof_null_safe(bst_search(root, &__item->member, compare), type, \
169                           member); \
170 })
171 
172 /**
173  * bst_search_key_type - Find an item in a binary search tree.
174  * @root:           Tree to search
175  * @key:            Pointer to key to search for.
176  * @compare_key:    Compare function.
177  * @type:           Type to return.
178  * @member:         Name of struct bst_node embedded in @type.
179  *
180  * Return: Item in @root matching @key, or %NULL if no matching node is found.
181  */
182 #define bst_search_key_type(root, key, compare_key, type, member) \
183     containerof_null_safe(bst_search_key(root, key, compare_key), type, member);
184 
185 /* Internal helper. Don't call directly */
186 void bst_update_rank_insert(struct bst_root *root, struct bst_node *node);
187 
188 /**
189  * bst_insert - Insert node in tree.
190  * @root:       Tree.
191  * @node:       Node to insert.
192  * @compare:    Compare function.
193  *
194  * Insert @node in @root.
195  *
196  * Return: %true if @node was inserted. %false if a node matching @node is
197  * already in @root.
198  */
bst_insert(struct bst_root * root,struct bst_node * node,bst_compare_t compare)199 static inline bool bst_insert(struct bst_root *root, struct bst_node *node,
200                               bst_compare_t compare) {
201     DEBUG_ASSERT(root);
202     DEBUG_ASSERT(node);
203     DEBUG_ASSERT(compare);
204     DEBUG_ASSERT(!node->rank);
205 
206     struct bst_node *parent = NULL;
207     struct bst_node **parent_ptr = &root->root;
208     int diff;
209     bool is_right_child = false;
210     while (true) {
211         struct bst_node *tree_node = *parent_ptr;
212         if (!tree_node) {
213             node->rank = 1;
214             node->parent = parent;
215             node->child[0] = NULL;
216             node->child[1] = NULL;
217             *parent_ptr = node;
218             bst_update_rank_insert(root, node);
219             return true;
220         }
221         diff = compare(tree_node, node);
222         if (!diff) {
223             return false;
224         }
225         is_right_child = diff > 0;
226         parent_ptr = &tree_node->child[is_right_child];
227         parent = tree_node;
228     }
229 }
230 
231 /**
232  * bst_delete - Remove node from tree.
233  * @root:   Tree.
234  * @node:   Node to delete
235  *
236  * Delete @node from @root.
237  */
238 void bst_delete(struct bst_root *root, struct bst_node *node);
239 
240 /**
241  * bst_prev - Get previous node.
242  * @root:       Tree.
243  * @node:       Node to move from.
244  *
245  * Use bst_prev_type instead to use pointers to the struct that contains @node.
246  *
247  * Return: If @node is %NULL, right-most node in @root.
248  *         If @node is not %NULL, right-most node to the left of @node.
249  *         %NULL if the node described above does not exist.
250  */
251 struct bst_node *bst_prev(struct bst_root *root, struct bst_node *node);
252 
253 /**
254  * bst_prev_type - Get previous item.
255  * @root:       Tree.
256  * @item:       Item to move from.
257  * @type:       Type of @item.
258  * @member:     Name of struct bst_node embedded in @type.
259  *
260  * Return: If @item is %NULL, right-most item in @root.
261  *         If @item is not %NULL, right-most item to the left of @item.
262  *         %NULL if the item described above does not exist.
263  */
264 #define bst_prev_type(root, item, type, member) \
265     containerof_null_safe(bst_prev(root, item), type, member)
266 
267 /**
268  * bst_next - Get next node.
269  * @root:       Tree.
270  * @node:       Node to move from.
271  *
272  * Use bst_next_type instead to use pointers to the struct that contains @node.
273  *
274  * Return: If @node is %NULL, left-most node in @root.
275  *         If @node is not %NULL, left-most node to the right of @node.
276  *         %NULL if the node described above does not exist.
277  */
278 struct bst_node *bst_next(const struct bst_root *root, struct bst_node *node);
279 
280 /**
281  * bst_next_type - Get previous item.
282  * @root:       Tree.
283  * @item:       Item to move from.
284  * @type:       Type of @item.
285  * @member:     Name of struct bst_node embedded in @type.
286  *
287  * Return: If @item is %NULL, left-most item in @root.
288  *         If @item is not %NULL, left-most item to the right of @item.
289  *         %NULL if the item described above does not exist.
290  */
291 #define bst_next_type(root, item, type, member) \
292     containerof_null_safe(bst_next(root, item), type, member)
293 
294 /**
295  * bst_for_every_entry - Loop over every entry in a tree.
296  * @root:       Tree.
297  * @entry:      Entry variable used by loop body.
298  * @type:       Type of @entry.
299  * @member:     Name of struct bst_node embedded in @type.
300  *
301  * Loop over every node in @root, convert that node to @type and provide it as
302  * @entry to the loop body directly following this macro.
303  *
304  * It is safe to delete @entry from @root in the body if the loop, but it is not
305  * safe to delete any other nodes or insert any nodes.
306  */
307 #define bst_for_every_entry(root, entry, type, member) \
308     for (struct bst_node *_bst_for_every_cursor = bst_next(root, NULL); \
309             (_bst_for_every_cursor != NULL) && \
310             ((entry) = containerof(_bst_for_every_cursor, type, member)) && \
311             ((_bst_for_every_cursor = bst_next(root, _bst_for_every_cursor)) \
312              || true);)
313 
314 /* Internal helper. Don't call directly */
315 void bst_delete_all_helper(struct bst_root *root, struct bst_node *node);
316 
317 /**
318  * bst_for_every_entry_delete - Loop over tree and delete every entry.
319  * @root:       Tree.
320  * @entry:      Entry variable used by loop body.
321  * @type:       Type of @entry.
322  * @member:     Name of struct bst_node embedded in @type.
323  *
324  * Loop over every node in @root, convert that node to @type and provide it as
325  * @entry to the loop body directly following this macro.
326  *
327  * @entry will be removed from @root before entering the loop bode. It is not
328  * safe to delete any other nodes or insert any nodes.
329  */
330 #define bst_for_every_entry_delete(root, entry, type, member) \
331     for (struct bst_node *_bst_for_every_cursor = bst_next(root, NULL); \
332             (_bst_for_every_cursor != NULL) && ({\
333             (entry) = containerof(_bst_for_every_cursor, type, member); \
334             _bst_for_every_cursor = bst_next(root, _bst_for_every_cursor); \
335             bst_delete_all_helper(root, &(entry)->member); true;});)
336 
337 __END_CDECLS
338