Joey's btree module.
[ccan] / ccan / btree / btree.c
1 /*
2         Copyright (c) 2010  Joseph A. Adams
3         All rights reserved.
4         
5         Redistribution and use in source and binary forms, with or without
6         modification, are permitted provided that the following conditions
7         are met:
8         1. Redistributions of source code must retain the above copyright
9            notice, this list of conditions and the following disclaimer.
10         2. Redistributions in binary form must reproduce the above copyright
11            notice, this list of conditions and the following disclaimer in the
12            documentation and/or other materials provided with the distribution.
13         3. The name of the author may not be used to endorse or promote products
14            derived from this software without specific prior written permission.
15         
16         THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
17         IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
18         OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
19         IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
20         INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
21         NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
22         DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
23         THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24         (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
25         THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26 */
27
28 #include "btree.h"
29
30 #include <assert.h>
31 #include <stdlib.h>
32 #include <stdio.h>
33
34 #define MAX (BTREE_ITEM_MAX)
35 #define MIN (BTREE_ITEM_MAX >> 1)
36
37 static struct btree_node *node_alloc(int internal);
38 static void node_delete(struct btree_node *node, struct btree *btree);
39
40 static void branch_begin(btree_iterator iter);
41 static void branch_end(btree_iterator iter);
42 static void begin_end_lr(btree_iterator iter, struct btree_node *node, int lr);
43
44 /*
45  * If iter->node has parent, returns 1 and ascends the iterator such that
46  * iter->node->branch[iter->k] will be what iter->node was.
47  *
48  * If iter->node does not have a parent (is a root), returns 0 and leaves the
49  * iterator untouched.
50  */
51 #define ascend(iter) ((iter)->node->parent \
52         ? (iter)->k = (iter)->node->k, (iter)->node = (iter)->node->parent, 1 \
53         : 0)
54
55 static void node_insert(const void *x, struct btree_node *xr,
56                                 struct btree_node *p, unsigned int k);
57 static void node_split(const void **x, struct btree_node **xr,
58                                 struct btree_node *p, unsigned int k);
59
60 static void node_remove_leaf_item(struct btree_node *node, unsigned int k);
61 void node_restore(struct btree_node *node, unsigned int k);
62
63 static int node_walk_backward(const struct btree_node *node,
64                                 btree_action_t action, void *ctx);
65 static int node_walk_forward(const struct btree_node *node,
66                                 btree_action_t action, void *ctx);
67
68
69 /************************* Public functions *************************/
70
71 struct btree *btree_new(btree_search_t search)
72 {
73         struct btree *btree = calloc(1, sizeof(struct btree));
74         struct btree_node *node = node_alloc(0);
75                 node->parent = NULL;
76                 node->count = 0;
77                 node->depth = 0;
78         btree->root = node;
79         btree->search = search;
80         return btree;
81 }
82
83 void btree_delete(struct btree *btree)
84 {
85         node_delete(btree->root, btree);
86         free(btree);
87 }
88
89 int btree_begin_end_lr(const struct btree *btree, btree_iterator iter, int lr)
90 {
91         struct btree_node *node;
92         
93         iter->btree = (struct btree *)btree;
94         begin_end_lr(iter, btree->root, lr);
95         
96         /* Set iter->item if any items exist. */
97         node = iter->node;
98         if (node->count) {
99                 iter->item = (void*)node->item[iter->k - lr];
100                 return 1;
101         }
102         
103         return 0;
104 }
105
106 int btree_deref(btree_iterator iter)
107 {
108         if (iter->k >= iter->node->count) {
109                 struct btree_iterator_s tmp = *iter;
110                 do {
111                         if (!ascend(iter)) {
112                                 *iter = tmp;
113                                 return 0;
114                         }
115                 } while (iter->k >= iter->node->count);
116         }
117         
118         iter->item = (void*)iter->node->item[iter->k];
119         return 1;
120 }
121
122 int btree_prev(btree_iterator iter)
123 {
124         if (iter->node->depth) {
125                 branch_end(iter);
126         } else if (iter->k == 0) {
127                 struct btree_iterator_s tmp = *iter;
128                 do {
129                         if (!ascend(iter)) {
130                                 *iter = tmp;
131                                 return 0;
132                         }
133                 } while (iter->k == 0);
134         }
135         
136         iter->item = (void*)iter->node->item[--iter->k];
137         return 1;
138 }
139
140 int btree_next(btree_iterator iter)
141 {
142         int ret = btree_deref(iter);
143         if (ret) {
144                 iter->k++;
145                 if (iter->node->depth)
146                         branch_begin(iter);
147         }
148         return ret;
149 }
150
151 int btree_find_lr(const struct btree *btree, const void *key,
152                                 btree_iterator iter, int lr)
153 {
154         struct btree_node *node = btree->root;
155         unsigned int k;
156         unsigned int depth;
157         int found = 0;
158         
159         iter->btree = (struct btree *)btree;
160         iter->item = NULL;
161         
162         depth = node->depth;
163         for (;;) {
164                 int f = 0;
165                 k = btree->search(key, node->item, node->count, lr, &f);
166                 
167                 if (f) {
168                         iter->item = (void*)node->item[k - lr];
169                         found = 1;
170                 }
171                 if (!depth--)
172                         break;
173                 
174                 node = node->branch[k];
175         }
176         
177         iter->node = node;
178         iter->k = k;
179         
180         return found;
181 }
182
183 int btree_walk_backward(const struct btree *btree,
184                                 btree_action_t action, void *ctx)
185 {
186         return node_walk_backward(btree->root, action, ctx);
187 }
188
189 int btree_walk_forward(const struct btree *btree,
190                                 btree_action_t action, void *ctx)
191 {
192         return node_walk_forward(btree->root, action, ctx);
193 }
194
195 void btree_insert_at(btree_iterator iter, const void *item)
196 {
197         const void *x = item;
198         struct btree_node *xr = NULL;
199         struct btree_node *p;
200         struct btree *btree = iter->btree;
201         
202         /* btree_insert_at always sets iter->item to item. */
203         iter->item = (void*)item;
204         
205         /*
206          * If node is not a leaf, fall to the end of the left branch of item[k]
207          * so that it will be a leaf. This does not modify the iterator's logical
208          * position.
209          */
210         if (iter->node->depth)
211                 branch_end(iter);
212         
213         /*
214          * First try inserting item into this node.
215          * If it's too big, split it, and repeat by
216          * trying to insert the median and right subtree into parent.
217          */
218         if (iter->node->count < MAX) {
219                 node_insert(x, xr, iter->node, iter->k);
220                 goto finished;
221         } else {
222                 for (;;) {
223                         node_split(&x, &xr, iter->node, iter->k);
224                         
225                         if (!ascend(iter))
226                                 break;
227                         
228                         if (iter->node->count < MAX) {
229                                 node_insert(x, xr, iter->node, iter->k);
230                                 goto finished;
231                         }
232                 }
233                 
234                 /*
235                  * If splitting came all the way up to the root, create a new root whose
236                  * left branch is the current root, median is x, and right branch is the
237                  * half split off from the root.
238                  */
239                 assert(iter->node == btree->root);
240                 p = node_alloc(1);
241                 p->parent = NULL;
242                 p->count = 1;
243                 p->depth = btree->root->depth + 1;
244                 p->item[0] = x;
245                 p->branch[0] = btree->root;
246                         btree->root->parent = p;
247                         btree->root->k = 0;
248                 p->branch[1] = xr;
249                         xr->parent = p;
250                         xr->k = 1;
251                 btree->root = p;
252         }
253         
254 finished:
255         btree->count++;
256         iter->node = NULL;
257 }
258
259 int btree_remove_at(btree_iterator iter)
260 {
261         struct btree *btree = iter->btree;
262         struct btree_node *root;
263         
264         if (!btree_deref(iter))
265                 return 0;
266         
267         if (!iter->node->depth) {
268                 node_remove_leaf_item(iter->node, iter->k);
269                 if (iter->node->count >= MIN || !iter->node->parent)
270                         goto finished;
271         } else {
272                 /*
273                  * We can't remove an item from an internal node, so we'll replace it
274                  * with its successor (which will always be in a leaf), then remove
275                  * the original copy of the successor.
276                  */
277                 
278                 /* Save pointer to condemned item. */
279                 const void **x = &iter->node->item[iter->k];
280                 
281                 /* Descend to successor. */
282                 iter->k++;
283                 branch_begin(iter);
284                 
285                 /* Replace condemned item with successor. */
286                 *x = iter->node->item[0];
287                 
288                 /* Remove successor. */
289                 node_remove_leaf_item(iter->node, 0);
290         }
291         
292         /*
293          * Restore nodes that fall under their minimum count.  This may
294          * propagate all the way up to the root.
295          */
296         for (;;) {
297                 if (iter->node->count >= MIN)
298                         goto finished;
299                 if (!ascend(iter))
300                         break;
301                 node_restore(iter->node, iter->k);
302         }
303         
304         /*
305          * If combining came all the way up to the root, and it has no more
306          * dividers, delete it and make its only branch the root.
307          */
308         root = iter->node;
309         assert(root == btree->root);
310         assert(root->depth > 0);
311         if (root->count == 0) {
312                 btree->root = root->branch[0];
313                 btree->root->parent = NULL;
314                 free(root);
315         }
316         
317 finished:
318         btree->count--;
319         iter->node = NULL;
320         return 1;
321 }
322
323 /*
324  * ascends iterator a until it matches iterator b's depth.
325  *
326  * Returns -1 if they end up on the same k (meaning a < b).
327  * Returns 0 otherwise.
328  */
329 static int elevate(btree_iterator a, btree_iterator b)
330 {
331         while (a->node->depth < b->node->depth)
332                 ascend(a);
333         
334         if (a->k == b->k)
335                 return -1;
336         return 0;
337 }
338
339 int btree_cmp_iters(const btree_iterator iter_a, const btree_iterator iter_b)
340 {
341         btree_iterator a = {*iter_a}, b = {*iter_b};
342         int ad, bd;
343         
344         ad = btree_deref(a);
345         bd = btree_deref(b);
346         
347         /* Check cases where one or both iterators are at the end. */
348         if (!ad)
349                 return bd ? 1 : 0;
350         if (!bd)
351                 return ad ? -1 : 0;
352         
353         /* Bring iterators to the same depth. */
354         if (a->node->depth < b->node->depth) {
355                 if (elevate(a, b))
356                         return -1;
357         } else if (a->node->depth > b->node->depth) {
358                 if (elevate(b, a))
359                         return 1;
360         }
361         
362         /* Bring iterators to the same node. */
363         while (a->node != b->node) {
364                 ascend(a);
365                 ascend(b);
366         }
367         
368         /* Now we can compare by k directly. */
369         if (a->k < b->k)
370                 return -1;
371         if (a->k > b->k)
372                 return 1;
373         
374         return 0;
375 }
376
377
378 /************************* Private functions *************************/
379
380 static struct btree_node *node_alloc(int internal)
381 {
382         struct btree_node *node;
383         size_t isize = internal
384                 ? sizeof(struct btree_node*) * (BTREE_ITEM_MAX+1)
385                 : 0;
386         node = malloc(sizeof(struct btree_node) + isize);
387         return node;
388 }
389
390 static void node_delete(struct btree_node *node, struct btree *btree)
391 {
392         unsigned int i, count = node->count;
393         
394         if (!node->depth) {
395                 if (btree->destroy) {
396                         for (i=0; i<count; i++)
397                                 btree->destroy((void*)node->item[i], btree->destroy_ctx);
398                 }
399         } else {
400                 for (i=0; i<count; i++) {
401                         node_delete(node->branch[i], btree);
402                         if (btree->destroy)
403                                 btree->destroy((void*)node->item[i], btree->destroy_ctx);
404                 }
405                 node_delete(node->branch[count], btree);
406         }
407         
408         free(node);
409 }
410
411 /* Set iter to beginning of branch pointed to by iter. */
412 static void branch_begin(btree_iterator iter)
413 {
414         struct btree_node *node = iter->node->branch[iter->k];
415         unsigned int depth = node->depth;
416         while (depth--)
417                 node = node->branch[0];
418         iter->node = node;
419         iter->k = 0;
420 }
421
422 /* Set iter to end of branch pointed to by iter. */
423 static void branch_end(btree_iterator iter)
424 {
425         struct btree_node *node = iter->node->branch[iter->k];
426         unsigned int depth = node->depth;
427         while (depth--)
428                 node = node->branch[node->count];
429         iter->node = node;
430         iter->k = node->count;
431 }
432
433 /* Traverse to the beginning or end of node, depending on lr. */
434 static void begin_end_lr(btree_iterator iter, struct btree_node *node, int lr)
435 {
436         iter->node = node;
437         iter->k = lr ? node->count : 0;
438         if (node->depth)
439                 (lr ? branch_end : branch_begin)(iter);
440 }
441
442 /*
443  * Inserts item x and right branch xr into node p at position k.
444  *
445  * Assumes p exists and has enough room.
446  * Ignores xr if p is a leaf.
447  */
448 static void node_insert(const void *x, struct btree_node *xr,
449                                 struct btree_node *p, unsigned int k)
450 {
451         unsigned int i;
452         
453         for (i = p->count; i-- > k;)
454                 p->item[i+1] = p->item[i];
455         p->item[k] = x;
456         
457         if (p->depth) {
458                 k++;
459                 for (i = p->count+1; i-- > k;) {
460                         p->branch[i+1] = p->branch[i];
461                         p->branch[i+1]->k = i+1;
462                 }
463                 p->branch[k] = xr;
464                 xr->parent = p;
465                 xr->k = k;
466         }
467         
468         p->count++;
469 }
470
471 /*
472  * Inserts item *x and subtree *xr into node p at position k, splitting it into
473  * nodes p and *xr with median item *x.
474  *
475  * Assumes p->count == MAX.
476  * Ignores original *xr if p is a leaf, but always sets it.
477  */
478 static void node_split(const void **x, struct btree_node **xr,
479                                 struct btree_node *p, unsigned int k)
480 {
481         unsigned int i, split;
482         struct btree_node *l = p, *r;
483         
484         /*
485          * If k <= MIN, item will be inserted into left subtree, so give l
486          * fewer items initially.
487          * Otherwise, item will be inserted into right subtree, so give r
488          * fewer items initially.
489          */
490         if (k <= MIN)
491                 split = MIN;
492         else
493                 split = MIN + 1;
494         
495         /*
496          * If l->depth is 0, allocate a leaf node.
497          * Otherwise, allocate an internal node.
498          */
499         r = node_alloc(l->depth);
500         
501         /* l and r will be siblings, so they will have the same parent and depth. */
502         r->parent = l->parent;
503         r->depth = l->depth;
504         
505         /*
506          * Initialize items/branches of right side.
507          * Do not initialize r's leftmost branch yet because we don't know
508          * whether it will be l's current rightmost branch or if *xr will
509          * take its place.
510          */
511         for (i = split; i < MAX; i++)
512                 r->item[i-split] = l->item[i];
513         if (r->depth) {
514                 for (i = split+1; i <= MAX; i++) {
515                         r->branch[i-split] = l->branch[i];
516                         r->branch[i-split]->parent = r;
517                         r->branch[i-split]->k = i-split;
518                 }
519         }
520         
521         /* Update counts. */
522         l->count = split;
523         r->count = MAX - split;
524         
525         /*
526          * The nodes are now split, but the key isn't inserted yet.
527          *
528          * Insert key into left or right half,
529          * depending on which side it fell on.
530          */
531         if (k <= MIN)
532                 node_insert(*x, *xr, l, k);
533         else
534                 node_insert(*x, *xr, r, k - split);
535         
536         /*
537          * Give l's rightmost branch to r because l's rightmost item
538          * is going up to become the median.
539          */
540         if (r->depth) {
541                 r->branch[0] = l->branch[l->count];
542                 r->branch[0]->parent = r;
543                 r->branch[0]->k = 0;
544         }
545         
546         /*
547          * Take up l's rightmost item to make it the median.
548          * That item's right branch is now r.
549          */
550         *x = l->item[--l->count];
551         *xr = r;
552 }
553
554 /*
555  * Removes item k from node p, shifting successor items back and
556  * decrementing the count.
557  *
558  * Assumes node p has the item k and is a leaf.
559  */
560 static void node_remove_leaf_item(struct btree_node *node, unsigned int k)
561 {
562         unsigned int i;
563         for (i = k+1; i < node->count; i++)
564                 node->item[i-1] = node->item[i];
565         node->count--;
566 }
567
568 static void move_left(struct btree_node *node, unsigned int k);
569 static void move_right(struct btree_node *node, unsigned int k);
570 static void combine(struct btree_node *node, unsigned int k);
571
572 /*
573  * Fixes node->branch[k]'s problem of having one less than MIN items.
574  * May or may not cause node to fall below MIN items, depending on whether
575  * two branches are combined or not.
576  */
577 void node_restore(struct btree_node *node, unsigned int k)
578 {
579         if (k == 0) {
580                 if (node->branch[1]->count > MIN)
581                         move_left(node, 0);
582                 else
583                         combine(node, 0);
584         } else if (k == node->count) {
585                 if (node->branch[k-1]->count > MIN)
586                         move_right(node, k-1);
587                 else
588                         combine(node, k-1);
589         } else if (node->branch[k-1]->count > MIN) {
590                 move_right(node, k-1);
591         } else if (node->branch[k+1]->count > MIN) {
592                 move_left(node, k);
593         } else {
594                 combine(node, k-1);
595         }
596 }
597
598 static void move_left(struct btree_node *node, unsigned int k)
599 {
600         struct btree_node *l = node->branch[k], *r = node->branch[k+1], *mv;
601         unsigned int i;
602         
603         l->item[l->count] = node->item[k];
604         node->item[k] = r->item[0];
605         for (i = 1; i < r->count; i++)
606                 r->item[i-1] = r->item[i];
607         
608         if (r->depth) {
609                 mv = r->branch[0];
610                 l->branch[l->count+1] = mv;
611                 mv->parent = l;
612                 mv->k = l->count+1;
613                 
614                 for (i = 1; i <= r->count; i++) {
615                         r->branch[i-1] = r->branch[i];
616                         r->branch[i-1]->k = i-1;
617                 }
618         }
619         
620         l->count++;
621         r->count--;
622 }
623
624 static void move_right(struct btree_node *node, unsigned int k)
625 {
626         struct btree_node *l = node->branch[k], *r = node->branch[k+1];
627         unsigned int i;
628         
629         for (i = r->count; i--;)
630                 r->item[i+1] = r->item[i];
631         r->item[0] = node->item[k];
632         node->item[k] = l->item[l->count-1];
633         
634         if (r->depth) {
635                 for (i = r->count+1; i--;) {
636                         r->branch[i+1] = r->branch[i];
637                         r->branch[i+1]->k = i+1;
638                 }
639                 r->branch[0] = l->branch[l->count];
640                 r->branch[0]->parent = r;
641                 r->branch[0]->k = 0;
642         }
643         
644         l->count--;
645         r->count++;
646 }
647
648 /* Combine node->branch[k] and node->branch[k+1]. */
649 static void combine(struct btree_node *node, unsigned int k)
650 {
651         struct btree_node *l = node->branch[k], *r = node->branch[k+1], *mv;
652         const void **o = &l->item[l->count];
653         unsigned int i;
654         
655         //append node->item[k] followed by right node's items to left node
656         *o++ = node->item[k];
657         for (i=0; i<r->count; i++)
658                 *o++ = r->item[i];
659         
660         //if applicable, append right node's branches to left node
661         if (r->depth) {
662                 for (i=0; i<=r->count; i++) {
663                         mv = r->branch[i];
664                         l->branch[l->count + i + 1] = mv;
665                         mv->parent = l;
666                         mv->k = l->count + i + 1;
667                 }
668         }
669         
670         //remove k and its right branch from parent node
671         for (i = k+1; i < node->count; i++) {
672                 node->item[i-1] = node->item[i];
673                 node->branch[i] = node->branch[i+1];
674                 node->branch[i]->k = i;
675         }
676         
677         //don't forget to update the left and parent node's counts and to free the right node
678         l->count += r->count + 1;
679         node->count--;
680         free(r);
681 }
682
683 static int node_walk_backward(const struct btree_node *node,
684                                 btree_action_t action, void *ctx)
685 {
686         unsigned int i, count = node->count;
687         
688         if (!node->depth) {
689                 for (i=count; i--;)
690                         if (!action((void*)node->item[i], ctx))
691                                 return 0;
692         } else {
693                 if (!node_walk_backward(node->branch[count], action, ctx))
694                         return 0;
695                 for (i=count; i--;) {
696                         if (!action((void*)node->item[i], ctx))
697                                 return 0;
698                         if (!node_walk_backward(node->branch[i], action, ctx))
699                                 return 0;
700                 }
701         }
702         
703         return 1;
704 }
705
706 static int node_walk_forward(const struct btree_node *node,
707                                 btree_action_t action, void *ctx)
708 {
709         unsigned int i, count = node->count;
710         
711         if (!node->depth) {
712                 for (i=0; i<count; i++)
713                         if (!action((void*)node->item[i], ctx))
714                                 return 0;
715         } else {
716                 for (i=0; i<count; i++) {
717                         if (!node_walk_forward(node->branch[i], action, ctx))
718                                 return 0;
719                         if (!action((void*)node->item[i], ctx))
720                                 return 0;
721                 }
722                 if (!node_walk_forward(node->branch[count], action, ctx))
723                         return 0;
724         }
725         
726         return 1;
727 }