1 /**
2  * @file lv_rb.c
3  *
4  */
5 
6 /*********************
7  *      INCLUDES
8  *********************/
9 #include "lv_rb_private.h"
10 #include "../stdlib/lv_string.h"
11 
12 /*********************
13  *      DEFINES
14  *********************/
15 
16 /**********************
17  *      TYPEDEFS
18  **********************/
19 
20 /**********************
21  *  STATIC PROTOTYPES
22  **********************/
23 
24 static lv_rb_node_t * rb_create_node(lv_rb_t * tree);
25 static lv_rb_node_t * rb_find_leaf_parent(lv_rb_t * tree, lv_rb_node_t * node);
26 static void rb_right_rotate(lv_rb_t * tree, lv_rb_node_t * node);
27 static void rb_left_rotate(lv_rb_t * tree, lv_rb_node_t * node);
28 static void rb_insert_color(lv_rb_t * tree, lv_rb_node_t * node);
29 static void rb_delete_color(lv_rb_t * tree, lv_rb_node_t * node1, lv_rb_node_t * node2);
30 
31 /**********************
32  *  GLOBAL VARIABLES
33  **********************/
34 
35 /**********************
36  *  STATIC VARIABLES
37  **********************/
38 
39 /**********************
40  *      MACROS
41  **********************/
42 
43 /**********************
44  *   GLOBAL FUNCTIONS
45  **********************/
46 
lv_rb_init(lv_rb_t * tree,lv_rb_compare_t compare,size_t node_size)47 bool lv_rb_init(lv_rb_t * tree, lv_rb_compare_t compare, size_t node_size)
48 {
49     LV_ASSERT_NULL(tree);
50     LV_ASSERT_NULL(compare);
51     LV_ASSERT(node_size > 0);
52 
53     if(tree == NULL || compare == NULL || node_size == 0) {
54         return false;
55     }
56 
57     lv_memzero(tree, sizeof(lv_rb_t));
58 
59     tree->root = NULL;
60     tree->compare = compare;
61     tree->size = node_size;
62 
63     return true;
64 }
65 
lv_rb_insert(lv_rb_t * tree,void * key)66 lv_rb_node_t * lv_rb_insert(lv_rb_t * tree, void * key)
67 {
68     LV_ASSERT_NULL(tree);
69     if(tree == NULL) {
70         return NULL;
71     }
72 
73     lv_rb_node_t * node = lv_rb_find(tree, key);
74     if(node) return node;
75     else {
76         node = rb_create_node(tree);
77         if(node == NULL) return NULL;
78 
79         if(tree->root == NULL) {
80             tree->root = node;
81             node->parent = NULL;
82             node->color = LV_RB_COLOR_BLACK;
83             return node;
84         }
85     }
86 
87     void * new_data = node->data;
88     node->data = key;
89     lv_rb_node_t * parent = rb_find_leaf_parent(tree, node);
90 
91     node->parent = parent;
92     node->color = LV_RB_COLOR_RED;
93 
94     if(tree->compare(key, parent->data) < 0) parent->left = node;
95     else parent->right = node;
96 
97     rb_insert_color(tree, node);
98 
99     node->data = new_data;
100     return node;
101 }
102 
lv_rb_find(lv_rb_t * tree,const void * key)103 lv_rb_node_t * lv_rb_find(lv_rb_t * tree, const void * key)
104 {
105     LV_ASSERT_NULL(tree);
106     if(tree == NULL) {
107         return NULL;
108     }
109 
110     lv_rb_node_t * current = tree->root;
111 
112     while(current != NULL) {
113         lv_rb_compare_res_t cmp = tree->compare(key, current->data);
114 
115         if(cmp == 0) {
116             return current;
117         }
118         else if(cmp < 0) {
119             current = current->left;
120         }
121         else {
122             current = current->right;
123         }
124     }
125 
126     return NULL;
127 }
128 
lv_rb_remove_node(lv_rb_t * tree,lv_rb_node_t * node)129 void * lv_rb_remove_node(lv_rb_t * tree, lv_rb_node_t * node)
130 {
131     lv_rb_node_t * child = NULL;
132     lv_rb_node_t * parent = NULL;
133     lv_rb_color_t color = LV_RB_COLOR_BLACK;
134 
135     if(node->left != NULL && node->right != NULL) {
136         lv_rb_node_t * replace = node;
137         replace = lv_rb_minimum_from(replace->right);
138 
139         if(node->parent != NULL) {
140             if(node->parent->left == node) {
141                 node->parent->left = replace;
142             }
143             else {
144                 node->parent->right = replace;
145             }
146         }
147         else {
148             tree->root = replace;
149         }
150 
151         child = replace->right;
152         parent = replace->parent;
153         color = replace->color;
154 
155         if(parent == node) {
156             parent = replace;
157         }
158         else {
159             if(child != NULL) {
160                 child->parent = parent;
161             }
162             parent->left = child;
163             replace->right = node->right;
164             node->right->parent = replace;
165         }
166 
167         replace->parent = node->parent;
168         replace->color = node->color;
169         replace->left = node->left;
170         node->left->parent = replace;
171 
172         if(color == LV_RB_COLOR_BLACK) {
173             rb_delete_color(tree, child, parent);
174         }
175 
176         void * data = node->data;
177         lv_free(node);
178         return data;
179     }
180 
181     child = node->right != NULL ? node->right : node->left;
182     parent = node->parent;
183     color = node->color;
184 
185     if(child != NULL) {
186         child->parent = parent;
187     }
188 
189     if(parent != NULL) {
190         if(parent->left == node) {
191             parent->left = child;
192         }
193         else {
194             parent->right = child;
195         }
196     }
197     else {
198         tree->root = child;
199     }
200 
201     if(color == LV_RB_COLOR_BLACK) {
202         rb_delete_color(tree, child, parent);
203     }
204 
205     void * data = node->data;
206     lv_free(node);
207     return data;
208 }
209 
lv_rb_remove(lv_rb_t * tree,const void * key)210 void * lv_rb_remove(lv_rb_t * tree, const void * key)
211 {
212     LV_ASSERT_NULL(tree);
213     if(tree == NULL) {
214         return NULL;
215     }
216 
217     lv_rb_node_t * node = lv_rb_find(tree, key);
218     LV_ASSERT_NULL(node);
219     if(node == NULL) {
220         LV_LOG_WARN("rb delete %d not found", (int)(uintptr_t)key);
221         return NULL;
222     }
223 
224     return lv_rb_remove_node(tree, node);
225 }
226 
lv_rb_drop_node(lv_rb_t * tree,lv_rb_node_t * node)227 bool lv_rb_drop_node(lv_rb_t * tree, lv_rb_node_t * node)
228 {
229     LV_ASSERT_NULL(tree);
230     if(tree == NULL) {
231         return false;
232     }
233 
234     void * data = lv_rb_remove_node(tree, node);
235     if(data) {
236         lv_free(data);
237         return true;
238     }
239     return false;
240 }
241 
lv_rb_drop(lv_rb_t * tree,const void * key)242 bool lv_rb_drop(lv_rb_t * tree, const void * key)
243 {
244     LV_ASSERT_NULL(tree);
245     if(tree == NULL) {
246         return false;
247     }
248 
249     void * data = lv_rb_remove(tree, key);
250     if(data) {
251         lv_free(data);
252         return true;
253     }
254     return false;
255 }
256 
lv_rb_destroy(lv_rb_t * tree)257 void lv_rb_destroy(lv_rb_t * tree)
258 {
259     LV_ASSERT_NULL(tree);
260 
261     if(tree == NULL) {
262         return;
263     }
264 
265     lv_rb_node_t * node = tree->root;
266     lv_rb_node_t * parent = NULL;
267 
268     while(node != NULL) {
269         if(node->left != NULL) {
270             node = node->left;
271         }
272         else if(node->right != NULL) {
273             node = node->right;
274         }
275         else {
276             parent = node->parent;
277             if(parent != NULL) {
278                 if(parent->left == node) {
279                     parent->left = NULL;
280                 }
281                 else {
282                     parent->right = NULL;
283                 }
284             }
285             lv_free(node->data);
286             lv_free(node);
287             node = parent;
288         }
289     }
290     tree->root = NULL;
291 }
292 
lv_rb_minimum(lv_rb_t * tree)293 lv_rb_node_t * lv_rb_minimum(lv_rb_t * tree)
294 {
295     LV_ASSERT_NULL(tree);
296     if(tree == NULL) {
297         return NULL;
298     }
299     return lv_rb_minimum_from(tree->root);
300 }
301 
lv_rb_maximum(lv_rb_t * tree)302 lv_rb_node_t * lv_rb_maximum(lv_rb_t * tree)
303 {
304     LV_ASSERT_NULL(tree);
305     if(tree == NULL) {
306         return NULL;
307     }
308     return lv_rb_maximum_from(tree->root);
309 }
310 
lv_rb_minimum_from(lv_rb_node_t * node)311 lv_rb_node_t * lv_rb_minimum_from(lv_rb_node_t * node)
312 {
313     while(node->left != NULL) {
314         node = node->left;
315     }
316 
317     return node;
318 }
319 
lv_rb_maximum_from(lv_rb_node_t * node)320 lv_rb_node_t * lv_rb_maximum_from(lv_rb_node_t * node)
321 {
322     while(node->right != NULL) {
323         node = node->right;
324     }
325 
326     return node;
327 }
328 
329 /**********************
330  *   STATIC FUNCTIONS
331  **********************/
332 
rb_create_node(lv_rb_t * tree)333 static lv_rb_node_t * rb_create_node(lv_rb_t * tree)
334 {
335     lv_rb_node_t * node = lv_malloc_zeroed(sizeof(lv_rb_node_t));
336     LV_ASSERT_MALLOC(node);
337     if(node == NULL) {
338         return NULL;
339     }
340 
341     node->data = lv_malloc_zeroed(tree->size);
342     LV_ASSERT_MALLOC(node->data);
343     if(node->data == NULL) {
344         lv_free(node);
345         return NULL;
346     }
347 
348     node->color = LV_RB_COLOR_RED;
349     node->left = NULL;
350     node->right = NULL;
351 
352     return node;
353 }
354 
rb_find_leaf_parent(lv_rb_t * tree,lv_rb_node_t * node)355 static lv_rb_node_t * rb_find_leaf_parent(lv_rb_t * tree, lv_rb_node_t * node)
356 {
357     lv_rb_node_t * current = tree->root;
358     lv_rb_node_t * parent = current;
359 
360     while(current != NULL) {
361         parent = current;
362 
363         if(tree->compare(node->data, current->data) < 0) {
364             current = current->left;
365         }
366         else {
367             current = current->right;
368         }
369     }
370 
371     return parent;
372 }
373 
rb_right_rotate(lv_rb_t * tree,lv_rb_node_t * node)374 static void rb_right_rotate(lv_rb_t * tree, lv_rb_node_t * node)
375 {
376     lv_rb_node_t * left = node->left;
377     node->left = left->right;
378 
379     if(left->right != NULL) {
380         left->right->parent = node;
381     }
382 
383     left->parent = node->parent;
384 
385     if(node->parent == NULL) {
386         tree->root = left;
387     }
388     else if(node == node->parent->right) {
389         node->parent->right = left;
390     }
391     else {
392         node->parent->left = left;
393     }
394 
395     left->right = node;
396     node->parent = left;
397 }
398 
rb_left_rotate(lv_rb_t * tree,lv_rb_node_t * node)399 static void rb_left_rotate(lv_rb_t * tree, lv_rb_node_t * node)
400 {
401     lv_rb_node_t * right = node->right;
402     node->right = right->left;
403 
404     if(right->left != NULL) {
405         right->left->parent = node;
406     }
407 
408     right->parent = node->parent;
409 
410     if(node->parent == NULL) {
411         tree->root = right;
412     }
413     else if(node == node->parent->left) {
414         node->parent->left = right;
415     }
416     else {
417         node->parent->right = right;
418     }
419 
420     right->left = node;
421     node->parent = right;
422 }
423 
rb_insert_color(lv_rb_t * tree,lv_rb_node_t * node)424 static void rb_insert_color(lv_rb_t * tree, lv_rb_node_t * node)
425 {
426     lv_rb_node_t * parent = NULL;
427     lv_rb_node_t * gparent = NULL;
428 
429     while((parent = node->parent) && parent->color == LV_RB_COLOR_RED) {
430         gparent = parent->parent;
431 
432         if(parent == gparent->left) {
433             {
434                 lv_rb_node_t * uncle = gparent->right;
435                 if(uncle && uncle->color == LV_RB_COLOR_RED) {
436                     uncle->color = LV_RB_COLOR_BLACK;
437                     parent->color = LV_RB_COLOR_BLACK;
438                     gparent->color = LV_RB_COLOR_RED;
439                     node = gparent;
440                     continue;
441                 }
442             }
443 
444             if(parent->right == node) {
445                 lv_rb_node_t * tmp;
446                 rb_left_rotate(tree, parent);
447                 tmp = parent;
448                 parent = node;
449                 node = tmp;
450             }
451 
452             parent->color = LV_RB_COLOR_BLACK;
453             gparent->color = LV_RB_COLOR_RED;
454             rb_right_rotate(tree, gparent);
455         }
456         else {
457             {
458                 lv_rb_node_t * uncle = gparent->left;
459                 if(uncle && uncle->color == LV_RB_COLOR_RED) {
460                     uncle->color = LV_RB_COLOR_BLACK;
461                     parent->color = LV_RB_COLOR_BLACK;
462                     gparent->color = LV_RB_COLOR_RED;
463                     node = gparent;
464                     continue;
465                 }
466             }
467 
468             if(parent->left == node) {
469                 lv_rb_node_t * tmp;
470                 rb_right_rotate(tree, parent);
471                 tmp = parent;
472                 parent = node;
473                 node = tmp;
474             }
475 
476             parent->color = LV_RB_COLOR_BLACK;
477             gparent->color = LV_RB_COLOR_RED;
478             rb_left_rotate(tree, gparent);
479         }
480     }
481 
482     tree->root->color = LV_RB_COLOR_BLACK;
483 }
484 
rb_delete_color(lv_rb_t * tree,lv_rb_node_t * node1,lv_rb_node_t * node2)485 static void rb_delete_color(lv_rb_t * tree, lv_rb_node_t * node1, lv_rb_node_t * node2)
486 {
487     LV_ASSERT_NULL(tree);
488     if(tree == NULL) {
489         return;
490     }
491 
492     while((node1 == NULL || node1->color == LV_RB_COLOR_BLACK) && node1 != tree->root) {
493         if(node2->left == node1) {
494             lv_rb_node_t * pNode2 = node2->right;
495             if(pNode2->color == LV_RB_COLOR_RED) {
496                 pNode2->color = LV_RB_COLOR_BLACK;
497                 node2->color = LV_RB_COLOR_RED;
498                 rb_left_rotate(tree, node2);
499                 pNode2 = node2->right;
500             }
501 
502             if((pNode2->left == NULL || pNode2->left->color == LV_RB_COLOR_BLACK)
503                && (pNode2->right == NULL || pNode2->right->color == LV_RB_COLOR_BLACK)) {
504                 pNode2->color = LV_RB_COLOR_RED;
505                 node1 = node2;
506                 node2 = node2->parent;
507             }
508             else {
509                 if(pNode2->right == NULL || pNode2->right->color == LV_RB_COLOR_BLACK) {
510                     pNode2->left->color = LV_RB_COLOR_BLACK;
511                     pNode2->color = LV_RB_COLOR_RED;
512                     rb_right_rotate(tree, pNode2);
513                     pNode2 = node2->right;
514                 }
515                 pNode2->color = node2->color;
516                 node2->color = LV_RB_COLOR_BLACK;
517                 pNode2->right->color = LV_RB_COLOR_BLACK;
518                 rb_left_rotate(tree, node2);
519                 node1 = tree->root;
520                 break;
521             }
522         }
523         else {
524             lv_rb_node_t * pNode2 = node2->left;
525             if(pNode2->color == LV_RB_COLOR_RED) {
526                 pNode2->color = LV_RB_COLOR_BLACK;
527                 node2->color = LV_RB_COLOR_RED;
528                 rb_right_rotate(tree, node2);
529                 pNode2 = node2->left;
530             }
531 
532             if((pNode2->left == NULL || pNode2->left->color == LV_RB_COLOR_BLACK)
533                && (pNode2->right == NULL || pNode2->right->color == LV_RB_COLOR_BLACK)) {
534                 pNode2->color = LV_RB_COLOR_RED;
535                 node1 = node2;
536                 node2 = node2->parent;
537             }
538             else {
539                 if(pNode2->left == NULL || pNode2->left->color == LV_RB_COLOR_BLACK) {
540                     pNode2->right->color = LV_RB_COLOR_BLACK;
541                     pNode2->color = LV_RB_COLOR_RED;
542                     rb_left_rotate(tree, pNode2);
543                     pNode2 = node2->left;
544                 }
545                 pNode2->color = node2->color;
546                 node2->color = LV_RB_COLOR_BLACK;
547                 pNode2->left->color = LV_RB_COLOR_BLACK;
548                 rb_right_rotate(tree, node2);
549                 node1 = tree->root;
550                 break;
551             }
552         }
553     }
554     if(node1 != NULL)
555         node1->color = LV_RB_COLOR_BLACK;
556 }
557