git.ozlabs.org Git - ccan/blob - 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 }