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