1 /*
2  * Copyright (c) 2006-2022, RT-Thread Development Team
3  *
4  * SPDX-License-Identifier: Apache-2.0
5  *
6  * Change Logs:
7  * Date           Author       Notes
8  * 2019-10-12     Jesven       first version
9  * 2022-11-14     WangXiaoyao  Optimize footprint and performance
10  *                             Export as ADT for generic use case
11  */
12 #include <stddef.h>
13 #include <stdint.h>
14 
15 #include "avl.h"
16 
17 #define HEIGHT_OF(node) ((node) ? (node)->height : 0)
18 #define IS_RCHILD(node) (!((node) - ((node)->parent->avl_right)))
19 #define IS_LCHILD(node) (!((node) - ((node)->parent->avl_left)))
20 #define NODE_PLACE(node)                                                       \
21     IS_LCHILD(node) ? &(node)->parent->avl_left : &(node)->parent->avl_right
22 
rotate_right(struct util_avl_struct * axis,struct util_avl_struct * lchild,struct util_avl_struct * lrchild,struct util_avl_struct ** nodeplace,size_t lrheight)23 static inline void rotate_right(struct util_avl_struct *axis,
24                                 struct util_avl_struct *lchild,
25                                 struct util_avl_struct *lrchild,
26                                 struct util_avl_struct **nodeplace,
27                                 size_t lrheight)
28 {
29     axis->avl_left = lrchild;
30     lchild->avl_right = axis;
31 
32     axis->height = lrheight + 1;
33     lchild->height = axis->height + 1;
34 
35     lchild->parent = axis->parent;
36     axis->parent = lchild;
37 
38     *nodeplace = lchild;
39     if (lrchild != NULL)
40         lrchild->parent = axis;
41 }
42 
midmount_right(struct util_avl_struct * axis,struct util_avl_struct * lchild,struct util_avl_struct * lrchild,struct util_avl_struct ** nodeplace,size_t lrheight)43 static inline void midmount_right(struct util_avl_struct *axis,
44                                   struct util_avl_struct *lchild,
45                                   struct util_avl_struct *lrchild,
46                                   struct util_avl_struct **nodeplace,
47                                   size_t lrheight)
48 {
49     lchild->avl_right = lrchild->avl_left;
50     axis->avl_left = lrchild->avl_right;
51     lrchild->avl_left = lchild;
52     lrchild->avl_right = axis;
53 
54     lrchild->height = lchild->height;
55     lchild->height = lrheight;
56     axis->height = lrheight;
57 
58     lrchild->parent = axis->parent;
59     lchild->parent = lrchild;
60     axis->parent = lrchild;
61     if (lchild->avl_right != NULL)
62         lchild->avl_right->parent = lchild;
63     if (axis->avl_left != NULL)
64         axis->avl_left->parent = axis;
65     *nodeplace = lrchild;
66 }
67 
rotate_left(struct util_avl_struct * axis,struct util_avl_struct * rchild,struct util_avl_struct * rlchild,struct util_avl_struct ** nodeplace,size_t rlheight)68 static inline void rotate_left(struct util_avl_struct *axis,
69                                struct util_avl_struct *rchild,
70                                struct util_avl_struct *rlchild,
71                                struct util_avl_struct **nodeplace,
72                                size_t rlheight)
73 {
74     axis->avl_right = rlchild;
75     rchild->avl_left = axis;
76 
77     axis->height = rlheight + 1;
78     rchild->height = axis->height + 1;
79 
80     rchild->parent = axis->parent;
81     axis->parent = rchild;
82 
83     *nodeplace = rchild;
84     if (rlchild != NULL)
85         rlchild->parent = axis;
86 }
87 
midmount_left(struct util_avl_struct * axis,struct util_avl_struct * rchild,struct util_avl_struct * rlchild,struct util_avl_struct ** nodeplace,size_t rlheight)88 static inline void midmount_left(struct util_avl_struct *axis,
89                                  struct util_avl_struct *rchild,
90                                  struct util_avl_struct *rlchild,
91                                  struct util_avl_struct **nodeplace,
92                                  size_t rlheight)
93 {
94     rchild->avl_left = rlchild->avl_right;
95     axis->avl_right = rlchild->avl_left;
96     rlchild->avl_right = rchild;
97     rlchild->avl_left = axis;
98 
99     rlchild->height = rchild->height;
100     rchild->height = rlheight;
101     axis->height = rlheight;
102 
103     rlchild->parent = axis->parent;
104     rchild->parent = rlchild;
105     axis->parent = rlchild;
106     if (rchild->avl_left != NULL)
107         rchild->avl_left->parent = rchild;
108     if (axis->avl_right != NULL)
109         axis->avl_right->parent = axis;
110 
111     *nodeplace = rlchild;
112 }
113 
114 /**
115  * @brief avl insertion & delete conceptually contain 2 stage
116  * 1. insertion/delete of reference
117  * 2. rebalance
118  */
119 
util_avl_rebalance(struct util_avl_struct * node,struct util_avl_root * root)120 void util_avl_rebalance(struct util_avl_struct *node,
121                         struct util_avl_root *root)
122 {
123     if (!node)
124         return;
125 
126     struct util_avl_struct *axis = node;
127     struct util_avl_struct **nodeplace;
128     do
129     {
130         struct util_avl_struct *lchild = axis->avl_left;
131         struct util_avl_struct *rchild = axis->avl_right;
132         nodeplace = axis->parent ? NODE_PLACE(axis) : &root->root_node;
133         int lheight = HEIGHT_OF(lchild);
134         int rheight = HEIGHT_OF(rchild);
135         if (rheight + 1 < lheight)
136         {
137             struct util_avl_struct *lrchild = lchild->avl_right;
138             size_t lrheight = HEIGHT_OF(lrchild);
139             if (HEIGHT_OF(lchild->avl_left) >= lrheight)
140             {
141                 rotate_right(axis, lchild, lrchild, nodeplace, lrheight);
142                 axis = lchild->parent;
143             }
144             else
145             {
146                 midmount_right(axis, lchild, lrchild, nodeplace, lrheight);
147                 axis = lrchild->parent;
148             }
149         }
150         else if (lheight + 1 < rheight)
151         {
152             struct util_avl_struct *rlchild = rchild->avl_left;
153             size_t rlheight = HEIGHT_OF(rlchild);
154             if (HEIGHT_OF(rchild->avl_right) >= rlheight)
155             {
156                 rotate_left(axis, rchild, rlchild, nodeplace, rlheight);
157                 axis = rchild->parent;
158             }
159             else
160             {
161                 midmount_left(axis, rchild, rlchild, nodeplace, rlheight);
162                 axis = rlchild->parent;
163             }
164         }
165         else
166         {
167             int height = (lheight < rheight ? rheight : lheight) + 1;
168             if (height == axis->height)
169                 break;
170             axis->height = height;
171             axis = axis->parent;
172         }
173     } while (axis);
174 }
175 
util_avl_remove(struct util_avl_struct * node,struct util_avl_root * root)176 void util_avl_remove(struct util_avl_struct *node, struct util_avl_root *root)
177 {
178     struct util_avl_struct **nodeplace;
179 
180     if (root->root_node == NULL)
181         return;
182 
183     if (node->parent != NULL)
184     {
185         nodeplace = NODE_PLACE(node);
186     }
187     else
188     {
189         nodeplace = &root->root_node;
190     }
191 
192     /* deletion */
193     if (node->avl_right == NULL)
194     {
195         *nodeplace = node->avl_left;
196         if (node->avl_left != NULL)
197             node->avl_left->parent = node->parent;
198         node = node->parent;
199     }
200     else
201     {
202         struct util_avl_struct *rchild = node->avl_right;
203         if (rchild->avl_left == NULL)
204         {
205             *nodeplace = rchild;
206             rchild->avl_left = node->avl_left;
207             if (rchild->avl_left != NULL)
208                 rchild->avl_left->parent = rchild;
209             rchild->parent = node->parent;
210             util_avl_rebalance(rchild, root);
211             node = rchild->parent;
212         }
213         else
214         {
215             struct util_avl_struct *successor = rchild->avl_left;
216             struct util_avl_struct *sparent = rchild;
217             while (successor->avl_left != NULL)
218             {
219                 sparent = successor;
220                 successor = successor->avl_left;
221             }
222             *nodeplace = successor;
223             sparent->avl_left = successor->avl_right;
224             successor->avl_left = node->avl_left;
225             successor->avl_right = node->avl_right;
226 
227             if (successor->avl_left != NULL)
228                 successor->avl_left->parent = successor;
229             successor->avl_right->parent = successor;
230 
231             if (sparent->avl_left != NULL)
232                 sparent->avl_left->parent = sparent;
233             successor->parent = node->parent;
234             util_avl_rebalance(sparent, root);
235             node = successor;
236         }
237     }
238 
239     /* rebalance */
240     util_avl_rebalance(node, root);
241     return;
242 }
243