1 /*
2 ** Copyright (C) 2014-2021 Cisco and/or its affiliates. All rights reserved.
3 ** Copyright (C) 2005-2013 Sourcefire, Inc.
4 **
5 ** This program is free software; you can redistribute it and/or modify
6 ** it under the terms of the GNU General Public License Version 2 as
7 ** published by the Free Software Foundation.  You may not use, modify or
8 ** distribute this program under any other version of the GNU General
9 ** Public License.
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 General Public License for more details.
15 **
16 ** You should have received a copy of the GNU General Public License
17 ** along with this program; if not, write to the Free Software
18 ** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
19 */
20 
21 #include <stdbool.h>
22 #include "fw_avltree.h"
23 #include <string.h>
24 #include <stdlib.h>
25 
is_root(struct FwAvlNode * node)26 static inline int is_root(struct FwAvlNode* node) { return (node->parent == NULL); }
get_balance(struct FwAvlNode * node)27 static inline int get_balance(struct FwAvlNode* node) { return node->balance; }
set_balance(int balance,struct FwAvlNode * node)28 static inline void set_balance(int balance, struct FwAvlNode* node) { node->balance = balance; }
inc_balance(struct FwAvlNode * node)29 static inline int inc_balance(struct FwAvlNode* node) { return ++node->balance; }
dec_balance(struct FwAvlNode * node)30 static inline int dec_balance(struct FwAvlNode* node) { return --node->balance; }
get_parent(struct FwAvlNode * node)31 static inline struct FwAvlNode* get_parent(struct FwAvlNode* node) { return node->parent; }
set_parent(struct FwAvlNode * parent,struct FwAvlNode * node)32 static inline void set_parent(struct FwAvlNode* parent, struct FwAvlNode* node) { node->parent = parent; }
33 
new_node(uint32_t key,void * data)34 static inline struct FwAvlNode* new_node(uint32_t key, void* data)
35 {
36     struct FwAvlNode* node = calloc(1, sizeof(struct FwAvlNode));
37     if(node != NULL)
38     {
39         node->key = key;
40         node->data = data;
41     }
42     return node;
43 }
44 
get_first(struct FwAvlNode * node)45 static inline struct FwAvlNode* get_first(struct FwAvlNode* node)
46 {
47     while(node->left != NULL)
48         node = node->left;
49     return node;
50 }
51 
get_last(struct FwAvlNode * node)52 static inline struct FwAvlNode* get_last(struct FwAvlNode* node)
53 {
54     while(node->right != NULL)
55         node = node->right;
56     return node;
57 }
58 
fwAvlFirst(const struct FwAvlTree * tree)59 struct FwAvlNode* fwAvlFirst(const struct FwAvlTree* tree)
60 {
61     if((tree != 0) && (tree->root != 0))
62         return get_first(tree->root);
63     else
64         return NULL;
65 }
66 
fwAvlLast(const struct FwAvlTree * tree)67 struct FwAvlNode* fwAvlLast(const struct FwAvlTree* tree)
68 {
69     if((tree != 0) && (tree->root != 0))
70         return get_last(tree->root);
71     else
72         return NULL;
73 }
74 
fwAvlNext(struct FwAvlNode * node)75 struct FwAvlNode* fwAvlNext(struct FwAvlNode* node)
76 {
77     struct FwAvlNode* parent = NULL;
78     struct FwAvlNode* tmp;
79 
80     if(node->right != NULL)
81     {
82         return get_first(node->right);
83     }
84     else
85     {
86         tmp = node;
87         while( ((parent = get_parent(tmp)) != NULL) && (parent->right == tmp) )
88             tmp = parent;
89     }
90     return parent;
91 }
92 
fwAvlPrev(struct FwAvlNode * node)93 struct FwAvlNode* fwAvlPrev(struct FwAvlNode* node)
94 {
95     struct FwAvlNode* parent;
96     struct FwAvlNode* tmp;
97 
98     if(node->left != NULL)
99     {
100         tmp = get_first(node->left);
101     }
102     else
103     {
104         tmp = node;
105         while( ((parent = get_parent(tmp)) != NULL) && (parent->left == tmp) )
106             tmp = parent;
107     }
108     return tmp;
109 }
110 
rotate_left(struct FwAvlNode * node,struct FwAvlTree * tree)111 static void rotate_left(struct FwAvlNode* node, struct FwAvlTree* tree)
112 {
113     struct FwAvlNode* p = node;
114     struct FwAvlNode* q = node->right;
115     struct FwAvlNode* parent = get_parent(node);
116 
117     if(!is_root(p))
118     {
119         if(parent->left == p)
120             parent->left = q;
121         else
122             parent->right = q;
123     }
124     else
125     {
126         tree->root = q;
127     }
128 
129     set_parent(parent, q);
130     set_parent(q, p);
131 
132     p->right = q->left;
133     if(p->right != NULL)
134     {
135         set_parent(p, p->right);
136     }
137     q->left = p;
138 }
139 
rotate_right(struct FwAvlNode * node,struct FwAvlTree * tree)140 static void rotate_right(struct FwAvlNode* node, struct FwAvlTree* tree)
141 {
142     struct FwAvlNode* p = node;
143     struct FwAvlNode* q = node->left;
144     struct FwAvlNode* parent = get_parent(node);
145 
146     if(!is_root(p))
147     {
148         if(parent->left == p)
149             parent->left = q;
150         else
151             parent->right = q;
152     }
153     else
154     {
155         tree->root = q;
156     }
157 
158     set_parent(parent, q);
159     set_parent(q, p);
160 
161     p->left = q->right;
162     if(p->left != NULL)
163     {
164         set_parent(p, p->left);
165     }
166     q->right = p;
167 }
168 
do_lookup(const uint32_t key,const struct FwAvlTree * tree,struct FwAvlNode ** pparent,struct FwAvlNode ** unbalanced,int * is_left)169 static inline struct FwAvlNode* do_lookup(const uint32_t key,
170         const struct FwAvlTree* tree, struct FwAvlNode** pparent,
171         struct FwAvlNode** unbalanced, int* is_left)
172 {
173     struct FwAvlNode* node = tree->root;
174 
175     *pparent = NULL;
176     *unbalanced = node;
177     *is_left = 0;
178 
179     while (node != NULL)
180     {
181         if(get_balance(node) != 0)
182         {
183             *unbalanced = node;
184         }
185 
186         *pparent = node;
187 
188         if(key == node->key)
189         {
190             return node;
191         }
192         else
193         {
194             if((*is_left = node->key > key) != 0)
195                 node = node->left;
196             else
197                 node = node->right;
198         }
199     }
200     return NULL;
201 }
202 
fwAvlLookup(const uint32_t key,const struct FwAvlTree * tree)203 void* fwAvlLookup(const uint32_t key, const struct FwAvlTree* tree)
204 {
205     struct FwAvlNode* node = NULL;
206     struct FwAvlNode* pparent;
207     struct FwAvlNode* unbalanced;
208     int is_left;
209 
210     if(tree == 0)
211     {
212         return NULL;
213     }
214 
215     node = do_lookup(key, tree, &pparent, &unbalanced, &is_left);
216 
217     if(node != NULL)
218     {
219         return node->data;
220     }
221     else
222     {
223         return NULL;
224     }
225 }
226 
set_child(struct FwAvlNode * child,struct FwAvlNode * node,int left)227 static inline void set_child(struct FwAvlNode* child, struct FwAvlNode* node,
228                         int left)
229 {
230     if(left != 0)
231         node->left = child;
232     else
233         node->right = child;
234 }
235 
fwAvlInsert(uint32_t key,void * data,struct FwAvlTree * tree)236 int fwAvlInsert(uint32_t key, void* data, struct FwAvlTree* tree)
237 {
238     int is_left;
239     struct FwAvlNode* parent;
240     struct FwAvlNode* right;
241     struct FwAvlNode* left;
242     struct FwAvlNode* unbalanced;
243     struct FwAvlNode* node;
244 
245     if(do_lookup(key, tree, &parent, &unbalanced, &is_left) != NULL)
246         return 1;
247 
248     if(!(node = new_node(key, data)))
249         return -1;
250 
251     tree->count++;
252     if(parent == NULL)
253     {
254         tree->root = node;
255         tree->first = node;
256         tree->last = node;
257         return 0;
258     }
259 
260     if(is_left != 0)
261     {
262         if(parent == tree->first)
263             tree->first = node;
264     }
265     else
266     {
267         if(parent == tree->last)
268             tree->last = node;
269     }
270     set_parent(parent, node);
271     set_child(node, parent, is_left);
272 
273     while(1)
274     {
275         if(parent->left == node)
276             dec_balance(parent);
277         else
278             inc_balance(parent);
279 
280         if(parent == unbalanced)
281             break;
282 
283         node = parent;
284         parent = get_parent(node);
285     }
286 
287     switch(get_balance(unbalanced))
288     {
289         case 1:
290         case -1:
291             tree->height++;
292             break;
293         case 0:
294             break;
295         case 2:
296             right = unbalanced->right;
297 
298             if(get_balance(right) == 1)
299             {
300                 set_balance(0, unbalanced);
301                 set_balance(0, right);
302             }
303             else
304             {
305                 switch(get_balance(right->left))
306                 {
307                     case 1:
308                         set_balance(-1, unbalanced);
309                         set_balance(0, right);
310                         break;
311                     case 0:
312                         set_balance(0, unbalanced);
313                         set_balance(0, right);
314                         break;
315                     case -1:
316                         set_balance(0, unbalanced);
317                         set_balance(1, right);
318                         break;
319                 }
320                 set_balance(0, right->left);
321                 rotate_right(right, tree);
322             }
323             rotate_left(unbalanced, tree);
324             break;
325         case -2:
326             left = unbalanced->left;
327 
328             if(get_balance(left) == -1)
329             {
330                 set_balance(0, unbalanced);
331                 set_balance(0, left);
332             }
333             else
334             {
335                 switch(get_balance(left->right))
336                 {
337                     case 1:
338                         set_balance(0, unbalanced);
339                         set_balance(-1, left);
340                         break;
341                     case 0:
342                         set_balance(0, unbalanced);
343                         set_balance(0, left);
344                         break;
345                     case -1:
346                         set_balance(1, unbalanced);
347                         set_balance(0, left);
348                         break;
349                 }
350                 set_balance(0, left->right);
351                 rotate_left(left, tree);
352             }
353             rotate_right(unbalanced, tree);
354             break;
355     }
356     return 0;
357 }
358 
fwAvlInit(void)359 struct FwAvlTree*  fwAvlInit(void)
360 {
361     struct FwAvlTree* tree = calloc(1, sizeof(struct FwAvlTree));
362     return tree; /* The calling function checks for NULL */
363 }
364 
newFwQNode(struct FwAvlNode * treeNode)365 struct FwQNode* newFwQNode(struct FwAvlNode* treeNode)
366 {
367     struct FwQNode* q_node = calloc(1, sizeof(struct FwQNode));
368 
369     if(q_node != NULL)
370     {
371         q_node->treeNode = treeNode;
372         q_node->next = NULL;
373     }
374     return(q_node);
375 }
376 
fwAvlSerialize(struct FwAvlTree * tree)377 struct FwQNode* fwAvlSerialize(struct FwAvlTree* tree)
378 {
379     struct FwQNode* head;
380     struct FwQNode* node;
381     struct FwQNode* tail;
382 
383     if((tree == NULL) || (tree->root == NULL)) return NULL;
384 
385     head = newFwQNode(tree->root);
386     node = head;
387     tail = head;
388 
389     while(node)
390     {
391         if(node->treeNode->left != NULL)
392         {
393             tail->next = newFwQNode(node->treeNode->left);
394             tail = tail->next;
395         }
396 
397         if(node->treeNode->right != NULL)
398         {
399             tail->next = newFwQNode(node->treeNode->right);
400             tail = tail->next;
401         }
402 
403         node = node->next;
404     }
405     return head;
406 }
407 
fwAvlDeleteTree(struct FwAvlTree * tree,void (* dataDelete)(void * data))408 void fwAvlDeleteTree(struct FwAvlTree* tree, void (*dataDelete)(void* data))
409 {
410     struct FwQNode* node = fwAvlSerialize(tree);
411     struct FwQNode* tmp;
412 
413     while(node != NULL)
414     {
415         if(dataDelete)
416             dataDelete(node->treeNode->data);
417         free(node->treeNode);
418         tmp = node;
419         node = node->next;
420         free(tmp);
421     }
422     free(tree);
423 }
424