a4c2142f44d71464917a0dbbdc02d218a777f5db
[monolithium.git] / kernel / src / avl_tree.c
1 /*
2  * avl_tree.c
3  *
4  * Copyright (C) 2016 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
5  *
6  * This program is free software: you can redistribute it and/or modify
7  * it under the terms of the GNU Affero General Public License as
8  * published by the Free Software Foundation, either version 3 of the
9  * License, or (at your option) any later version.
10  *
11  * This program is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14  * GNU Affero General Public License for more details.
15  *
16  * You should have received a copy of the GNU Affero General Public License
17  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
18  */
19
20 #include <avl_tree.h>
21 #include <exception.h>
22
23 static int avl_tree_height(avl_tree_t *tree)
24 {
25     if (tree == NULL) return 0;
26     return 1 + MAX(avl_tree_height(tree->left), avl_tree_height(tree->right));
27 }
28
29 static int avl_tree_balance_factor(avl_tree_t *tree)
30 {
31     int left_height = avl_tree_height(tree->left);
32     int right_height = avl_tree_height(tree->right);
33
34     return left_height - right_height;
35 }
36
37 static void avl_tree_rotation(avl_tree_t **tree, bool_t clockwise)
38 {
39     avl_tree_t *root = *tree;
40     if (root == NULL) return;
41
42     if (clockwise)
43     {
44         avl_tree_t *pivot = root->left;
45         root->left = pivot->right;
46         if (root->left) root->left->parent = root;
47         pivot->right = root;
48         pivot->parent = root->parent;
49         root->parent = pivot;
50         *tree = pivot;
51     }
52     else
53     {
54         avl_tree_t *pivot = root->right;
55         root->right = pivot->left;
56         if (root->right) root->right->parent = root;
57         pivot->left = root;
58         pivot->parent = root->parent;
59         root->parent = pivot;
60         *tree = pivot;
61     }
62 }
63
64 static void avl_tree_balance(avl_tree_t **tree)
65 {
66     avl_tree_t *root = *tree;
67     if (root == NULL) return;
68
69     int balance_factor = avl_tree_balance_factor(root);
70     if (balance_factor < -1)
71     {
72         if (avl_tree_balance_factor(root->right) > 0) avl_tree_rotation(&root->right, TRUE);
73         avl_tree_rotation(tree, FALSE);
74     }
75     else if (balance_factor > 1)
76     {
77         if (avl_tree_balance_factor(root->left) < 0) avl_tree_rotation(&root->left, FALSE);
78         avl_tree_rotation(tree, TRUE);
79     }
80 }
81
82 static void avl_tree_insert_internal(avl_tree_t **tree, avl_tree_t *node, avl_tree_t *parent)
83 {
84     avl_tree_t *root = *tree;
85
86     if (root == NULL)
87     {
88         node->parent = parent;
89         node->left = node->right = node->next_equal = NULL;
90         *tree = node;
91         return;
92     }
93
94     if (node->key < root->key) avl_tree_insert_internal(&root->left, node, root);
95     else avl_tree_insert_internal(&root->right, node, root);
96
97     avl_tree_balance(tree);
98 }
99
100 void avl_tree_insert(avl_tree_t **tree, avl_tree_t *node)
101 {
102     avl_tree_t *existing = avl_tree_lookup(*tree, node->key);
103
104     if (existing != NULL)
105     {
106         node->next_equal = existing->next_equal;
107         existing->next_equal = node;
108     }
109     else
110     {
111         avl_tree_insert_internal(tree, node, NULL);
112     }
113 }
114
115 bool_t avl_tree_remove(avl_tree_t **tree, avl_tree_t *node)
116 {
117     bool_t result;
118     avl_tree_t *root = *tree;
119     if (root == NULL) return FALSE;
120
121     if (root->key == node->key)
122     {
123         if (root->next_equal != NULL)
124         {
125             if (root == node)
126             {
127                 root->next_equal->parent = root->parent;
128                 root->next_equal->left = root->left;
129                 root->next_equal->right = root->right;
130
131                 if (root->parent)
132                 {
133                     if (root == root->parent->left) root->parent->left = root->next_equal;
134                     else if (root == root->parent->right) root->parent->right = root->next_equal;
135                 }
136
137                 if (root->left) root->left->parent = root->next_equal;
138                 if (root->right) root->right->parent = root->next_equal;
139
140                 *tree = root->next_equal;
141             }
142             else
143             {
144                 avl_tree_t *ptr = root;
145                 while (ptr->next_equal != node) ptr = ptr->next_equal;
146                 ptr->next_equal = ptr->next_equal->next_equal;
147             }
148         }
149         else if (root->left == NULL && root->right == NULL)
150         {
151             *tree = NULL;
152         }
153         else if (root->left != NULL && root->right == NULL)
154         {
155             root->left->parent = root->parent;
156             *tree = root->left;
157         }
158         else if (root->right != NULL && root->left == NULL)
159         {
160             root->right->parent = root->parent;
161             *tree = root->right;
162         }
163         else
164         {
165             avl_tree_t *replacement = root->left;
166             while (replacement->right != NULL) replacement = replacement->right;
167
168             avl_tree_t *next_replacement = NULL;
169
170             if (replacement->next_equal)
171             {
172                 next_replacement = replacement->next_equal;
173                 next_replacement->parent = replacement->parent;
174                 next_replacement->left = replacement->left;
175                 next_replacement->right = replacement->right;
176             }
177             else if (replacement->left)
178             {
179                 next_replacement = replacement->left;
180             }
181
182             if (replacement->parent->right == replacement) replacement->parent->right = next_replacement;
183             else replacement->parent->left = next_replacement;
184
185             replacement->parent = root->parent;
186             replacement->left = root->left;
187             replacement->right = root->right;
188
189             if (root->left) root->left->parent = replacement;
190             if (root->right) root->right->parent = replacement;
191
192             *tree = replacement;
193         }
194
195         return TRUE;
196     }
197     else if (node->key < root->key) result = avl_tree_remove(&root->left, node);
198     else result = avl_tree_remove(&root->right, node);
199
200     avl_tree_balance(tree);
201     return result;
202 }
203
204 avl_tree_t *avl_tree_lookup(avl_tree_t *tree, qword_t key)
205 {
206     if (tree == NULL) return NULL;
207
208     if (tree->key == key) return tree;
209     else if (key < tree->key) return avl_tree_lookup(tree->left, key);
210     else return avl_tree_lookup(tree->right, key);
211 }
212
213 avl_tree_t *avl_tree_lower_bound(avl_tree_t *tree, qword_t key)
214 {
215     while (tree != NULL && tree->key > key) tree = tree->left;
216     if (tree == NULL) return NULL;
217
218     while (tree->right != NULL && tree->right->key <= key) tree = tree->right;
219     return tree;
220 }
221
222 avl_tree_t *avl_tree_upper_bound(avl_tree_t *tree, qword_t key)
223 {
224     while (tree != NULL && tree->key < key) tree = tree->right;
225     if (tree == NULL) return NULL;
226
227     while (tree->left != NULL && tree->left->key >= key) tree = tree->left;
228     return tree;
229 }
230
231 bool_t avl_tree_change_key(avl_tree_t **tree, avl_tree_t *node, qword_t new_key)
232 {
233     if (!avl_tree_remove(tree, node)) return FALSE;
234
235     node->key = new_key;
236     avl_tree_insert(tree, node);
237
238     return TRUE;
239 }
240
241 avl_tree_t *avl_get_next_node(avl_tree_t *node)
242 {
243     if (node->right)
244     {
245         node = node->right;
246         while (node->left) node = node->left;
247     }
248     else
249     {
250         while (node->parent && node->parent->right == node) node = node->parent;
251         node = node->parent;
252     }
253
254     return node;
255 }
256
257 avl_tree_t *avl_get_previous_node(avl_tree_t *node)
258 {
259     if (node->left)
260     {
261         node = node->left;
262         while (node->right) node = node->right;
263     }
264     else
265     {
266         while (node->parent && node->parent->left == node) node = node->parent;
267         node = node->parent;
268     }
269
270     return node;
271 }