tdb2: copy tdb1's changed expansion logic.
[ccan] / ccan / avl / avl.c
1 /*
2  * Copyright (C) 2010 Joseph Adams <joeyadams3.14159@gmail.com>
3  * 
4  * Permission is hereby granted, free of charge, to any person obtaining a copy
5  * of this software and associated documentation files (the "Software"), to deal
6  * in the Software without restriction, including without limitation the rights
7  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8  * copies of the Software, and to permit persons to whom the Software is
9  * furnished to do so, subject to the following conditions:
10  * 
11  * The above copyright notice and this permission notice shall be included in
12  * all copies or substantial portions of the Software.
13  * 
14  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
20  * THE SOFTWARE.
21  */
22
23 #include "avl.h"
24
25 #include <assert.h>
26 #include <stdlib.h>
27
28 static AvlNode *mkNode(const void *key, const void *value);
29 static void freeNode(AvlNode *node);
30
31 static AvlNode *lookup(const AVL *avl, AvlNode *node, const void *key);
32
33 static bool insert(AVL *avl, AvlNode **p, const void *key, const void *value);
34 static bool remove(AVL *avl, AvlNode **p, const void *key, AvlNode **ret);
35 static bool removeExtremum(AvlNode **p, int side, AvlNode **ret);
36
37 static int sway(AvlNode **p, int sway);
38 static void balance(AvlNode **p, int side);
39
40 static bool checkBalances(AvlNode *node, int *height);
41 static bool checkOrder(AVL *avl);
42 static size_t countNode(AvlNode *node);
43
44 /*
45  * Utility macros for converting between
46  * "balance" values (-1 or 1) and "side" values (0 or 1).
47  *
48  * bal(0)   == -1
49  * bal(1)   == +1
50  * side(-1) == 0
51  * side(+1) == 1
52  */
53 #define bal(side) ((side) == 0 ? -1 : 1)
54 #define side(bal) ((bal)  == 1 ?  1 : 0)
55
56 static int sign(int cmp)
57 {
58         if (cmp < 0)
59                 return -1;
60         if (cmp == 0)
61                 return 0;
62         return 1;
63 }
64
65 AVL *avl_new(AvlCompare compare)
66 {
67         AVL *avl = malloc(sizeof(*avl));
68         
69         assert(avl != NULL);
70         
71         avl->compare = compare;
72         avl->root = NULL;
73         avl->count = 0;
74         return avl;
75 }
76
77 void avl_free(AVL *avl)
78 {
79         freeNode(avl->root);
80         free(avl);
81 }
82
83 void *avl_lookup(const AVL *avl, const void *key)
84 {
85         AvlNode *found = lookup(avl, avl->root, key);
86         return found ? (void*) found->value : NULL;
87 }
88
89 AvlNode *avl_lookup_node(const AVL *avl, const void *key)
90 {
91         return lookup(avl, avl->root, key);
92 }
93
94 size_t avl_count(const AVL *avl)
95 {
96         return avl->count;
97 }
98
99 bool avl_insert(AVL *avl, const void *key, const void *value)
100 {
101         size_t old_count = avl->count;
102         insert(avl, &avl->root, key, value);
103         return avl->count != old_count;
104 }
105
106 bool avl_remove(AVL *avl, const void *key)
107 {
108         AvlNode *node = NULL;
109         
110         remove(avl, &avl->root, key, &node);
111         
112         if (node == NULL) {
113                 return false;
114         } else {
115                 free(node);
116                 return true;
117         }
118 }
119
120 static AvlNode *mkNode(const void *key, const void *value)
121 {
122         AvlNode *node = malloc(sizeof(*node));
123         
124         assert(node != NULL);
125         
126         node->key = key;
127         node->value = value;
128         node->lr[0] = NULL;
129         node->lr[1] = NULL;
130         node->balance = 0;
131         return node;
132 }
133
134 static void freeNode(AvlNode *node)
135 {
136         if (node) {
137                 freeNode(node->lr[0]);
138                 freeNode(node->lr[1]);
139                 free(node);
140         }
141 }
142
143 static AvlNode *lookup(const AVL *avl, AvlNode *node, const void *key)
144 {
145         int cmp;
146         
147         if (node == NULL)
148                 return NULL;
149         
150         cmp = avl->compare(key, node->key);
151         
152         if (cmp < 0)
153                 return lookup(avl, node->lr[0], key);
154         if (cmp > 0)
155                 return lookup(avl, node->lr[1], key);
156         return node;
157 }
158
159 /*
160  * Insert a key/value into a subtree, rebalancing if necessary.
161  *
162  * Return true if the subtree's height increased.
163  */
164 static bool insert(AVL *avl, AvlNode **p, const void *key, const void *value)
165 {
166         if (*p == NULL) {
167                 *p = mkNode(key, value);
168                 avl->count++;
169                 return true;
170         } else {
171                 AvlNode *node = *p;
172                 int      cmp  = sign(avl->compare(key, node->key));
173                 
174                 if (cmp == 0) {
175                         node->key = key;
176                         node->value = value;
177                         return false;
178                 }
179                 
180                 if (!insert(avl, &node->lr[side(cmp)], key, value))
181                         return false;
182                 
183                 /* If tree's balance became -1 or 1, it means the tree's height grew due to insertion. */
184                 return sway(p, cmp) != 0;
185         }
186 }
187
188 /*
189  * Remove the node matching the given key.
190  * If present, return the removed node through *ret .
191  * The returned node's lr and balance are meaningless.
192  *
193  * Return true if the subtree's height decreased.
194  */
195 static bool remove(AVL *avl, AvlNode **p, const void *key, AvlNode **ret)
196 {
197         if (*p == NULL) {
198                 return false;
199         } else {
200                 AvlNode *node = *p;
201                 int      cmp  = sign(avl->compare(key, node->key));
202                 
203                 if (cmp == 0) {
204                         *ret = node;
205                         avl->count--;
206                         
207                         if (node->lr[0] != NULL && node->lr[1] != NULL) {
208                                 AvlNode *replacement;
209                                 int      side;
210                                 bool     shrunk;
211                                 
212                                 /* Pick a subtree to pull the replacement from such that
213                                  * this node doesn't have to be rebalanced. */
214                                 side = node->balance <= 0 ? 0 : 1;
215                                 
216                                 shrunk = removeExtremum(&node->lr[side], 1 - side, &replacement);
217                                 
218                                 replacement->lr[0]   = node->lr[0];
219                                 replacement->lr[1]   = node->lr[1];
220                                 replacement->balance = node->balance;
221                                 *p = replacement;
222                                 
223                                 if (!shrunk)
224                                         return false;
225                                 
226                                 replacement->balance -= bal(side);
227                                 
228                                 /* If tree's balance became 0, it means the tree's height shrank due to removal. */
229                                 return replacement->balance == 0;
230                         }
231                         
232                         if (node->lr[0] != NULL)
233                                 *p = node->lr[0];
234                         else
235                                 *p = node->lr[1];
236                         
237                         return true;
238                         
239                 } else {
240                         if (!remove(avl, &node->lr[side(cmp)], key, ret))
241                                 return false;
242                         
243                         /* If tree's balance became 0, it means the tree's height shrank due to removal. */
244                         return sway(p, -cmp) == 0;
245                 }
246         }
247 }
248
249 /*
250  * Remove either the left-most (if side == 0) or right-most (if side == 1)
251  * node in a subtree, returning the removed node through *ret .
252  * The returned node's lr and balance are meaningless.
253  *
254  * The subtree must not be empty (i.e. *p must not be NULL).
255  *
256  * Return true if the subtree's height decreased.
257  */
258 static bool removeExtremum(AvlNode **p, int side, AvlNode **ret)
259 {
260         AvlNode *node = *p;
261         
262         if (node->lr[side] == NULL) {
263                 *ret = node;
264                 *p = node->lr[1 - side];
265                 return true;
266         }
267         
268         if (!removeExtremum(&node->lr[side], side, ret))
269                 return false;
270         
271         /* If tree's balance became 0, it means the tree's height shrank due to removal. */
272         return sway(p, -bal(side)) == 0;
273 }
274
275 /*
276  * Rebalance a node if necessary.  Think of this function
277  * as a higher-level interface to balance().
278  *
279  * sway must be either -1 or 1, and indicates what was added to
280  * the balance of this node by a prior operation.
281  *
282  * Return the new balance of the subtree.
283  */
284 static int sway(AvlNode **p, int sway)
285 {
286         if ((*p)->balance != sway)
287                 (*p)->balance += sway;
288         else
289                 balance(p, side(sway));
290         
291         return (*p)->balance;
292 }
293
294 /*
295  * Perform tree rotations on an unbalanced node.
296  *
297  * side == 0 means the node's balance is -2 .
298  * side == 1 means the node's balance is +2 .
299  */
300 static void balance(AvlNode **p, int side)
301 {
302         AvlNode  *node  = *p,
303                  *child = node->lr[side];
304         int opposite    = 1 - side;
305         int bal         = bal(side);
306         
307         if (child->balance != -bal) {
308                 /* Left-left (side == 0) or right-right (side == 1) */
309                 node->lr[side]      = child->lr[opposite];
310                 child->lr[opposite] = node;
311                 *p = child;
312                 
313                 child->balance -= bal;
314                 node->balance = -child->balance;
315                 
316         } else {
317                 /* Left-right (side == 0) or right-left (side == 1) */
318                 AvlNode *grandchild = child->lr[opposite];
319                 
320                 node->lr[side]           = grandchild->lr[opposite];
321                 child->lr[opposite]      = grandchild->lr[side];
322                 grandchild->lr[side]     = child;
323                 grandchild->lr[opposite] = node;
324                 *p = grandchild;
325                 
326                 node->balance       = 0;
327                 child->balance      = 0;
328                 
329                 if (grandchild->balance == bal)
330                         node->balance  = -bal;
331                 else if (grandchild->balance == -bal)
332                         child->balance = bal;
333                 
334                 grandchild->balance = 0;
335         }
336 }
337
338
339 /************************* avl_check_invariants() *************************/
340
341 bool avl_check_invariants(AVL *avl)
342 {
343         int    dummy;
344         
345         return checkBalances(avl->root, &dummy)
346             && checkOrder(avl)
347             && countNode(avl->root) == avl->count;
348 }
349
350 static bool checkBalances(AvlNode *node, int *height)
351 {
352         if (node) {
353                 int h0, h1;
354                 
355                 if (!checkBalances(node->lr[0], &h0))
356                         return false;
357                 if (!checkBalances(node->lr[1], &h1))
358                         return false;
359                 
360                 if (node->balance != h1 - h0 || node->balance < -1 || node->balance > 1)
361                         return false;
362                 
363                 *height = (h0 > h1 ? h0 : h1) + 1;
364                 return true;
365         } else {
366                 *height = 0;
367                 return true;
368         }
369 }
370
371 static bool checkOrder(AVL *avl)
372 {
373         AvlIter     i;
374         const void *last     = NULL;
375         bool        last_set = false;
376         
377         avl_foreach(i, avl) {
378                 if (last_set && avl->compare(last, i.key) >= 0)
379                         return false;
380                 last     = i.key;
381                 last_set = true;
382         }
383         
384         return true;
385 }
386
387 static size_t countNode(AvlNode *node)
388 {
389         if (node)
390                 return 1 + countNode(node->lr[0]) + countNode(node->lr[1]);
391         else
392                 return 0;
393 }
394
395
396 /************************* Traversal *************************/
397
398 void avl_iter_begin(AvlIter *iter, AVL *avl, AvlDirection dir)
399 {
400         AvlNode *node = avl->root;
401         
402         iter->stack_index = 0;
403         iter->direction   = dir;
404         
405         if (node == NULL) {
406                 iter->key      = NULL;
407                 iter->value    = NULL;
408                 iter->node     = NULL;
409                 return;
410         }
411         
412         while (node->lr[dir] != NULL) {
413                 iter->stack[iter->stack_index++] = node;
414                 node = node->lr[dir];
415         }
416         
417         iter->key   = (void*) node->key;
418         iter->value = (void*) node->value;
419         iter->node  = node;
420 }
421
422 void avl_iter_next(AvlIter *iter)
423 {
424         AvlNode     *node = iter->node;
425         AvlDirection dir  = iter->direction;
426         
427         if (node == NULL)
428                 return;
429         
430         node = node->lr[1 - dir];
431         if (node != NULL) {
432                 while (node->lr[dir] != NULL) {
433                         iter->stack[iter->stack_index++] = node;
434                         node = node->lr[dir];
435                 }
436         } else if (iter->stack_index > 0) {
437                 node = iter->stack[--iter->stack_index];
438         } else {
439                 iter->key      = NULL;
440                 iter->value    = NULL;
441                 iter->node     = NULL;
442                 return;
443         }
444         
445         iter->node  = node;
446         iter->key   = (void*) node->key;
447         iter->value = (void*) node->value;
448 }