796 lines · c
1// SPDX-License-Identifier: GPL-2.0-only2/*3 * lib/btree.c - Simple In-memory B+Tree4 *5 * Copyright (c) 2007-2008 Joern Engel <joern@purestorage.com>6 * Bits and pieces stolen from Peter Zijlstra's code, which is7 * Copyright 2007, Red Hat Inc. Peter Zijlstra8 *9 * see http://programming.kicks-ass.net/kernel-patches/vma_lookup/btree.patch10 *11 * A relatively simple B+Tree implementation. I have written it as a learning12 * exercise to understand how B+Trees work. Turned out to be useful as well.13 *14 * B+Trees can be used similar to Linux radix trees (which don't have anything15 * in common with textbook radix trees, beware). Prerequisite for them working16 * well is that access to a random tree node is much faster than a large number17 * of operations within each node.18 *19 * Disks have fulfilled the prerequisite for a long time. More recently DRAM20 * has gained similar properties, as memory access times, when measured in cpu21 * cycles, have increased. Cacheline sizes have increased as well, which also22 * helps B+Trees.23 *24 * Compared to radix trees, B+Trees are more efficient when dealing with a25 * sparsely populated address space. Between 25% and 50% of the memory is26 * occupied with valid pointers. When densely populated, radix trees contain27 * ~98% pointers - hard to beat. Very sparse radix trees contain only ~2%28 * pointers.29 *30 * This particular implementation stores pointers identified by a long value.31 * Storing NULL pointers is illegal, lookup will return NULL when no entry32 * was found.33 *34 * A tricks was used that is not commonly found in textbooks. The lowest35 * values are to the right, not to the left. All used slots within a node36 * are on the left, all unused slots contain NUL values. Most operations37 * simply loop once over all slots and terminate on the first NUL.38 */39 40#include <linux/btree.h>41#include <linux/cache.h>42#include <linux/kernel.h>43#include <linux/slab.h>44#include <linux/module.h>45 46#define NODESIZE MAX(L1_CACHE_BYTES, 128)47 48struct btree_geo {49 int keylen;50 int no_pairs;51 int no_longs;52};53 54struct btree_geo btree_geo32 = {55 .keylen = 1,56 .no_pairs = NODESIZE / sizeof(long) / 2,57 .no_longs = NODESIZE / sizeof(long) / 2,58};59EXPORT_SYMBOL_GPL(btree_geo32);60 61#define LONG_PER_U64 (64 / BITS_PER_LONG)62struct btree_geo btree_geo64 = {63 .keylen = LONG_PER_U64,64 .no_pairs = NODESIZE / sizeof(long) / (1 + LONG_PER_U64),65 .no_longs = LONG_PER_U64 * (NODESIZE / sizeof(long) / (1 + LONG_PER_U64)),66};67EXPORT_SYMBOL_GPL(btree_geo64);68 69struct btree_geo btree_geo128 = {70 .keylen = 2 * LONG_PER_U64,71 .no_pairs = NODESIZE / sizeof(long) / (1 + 2 * LONG_PER_U64),72 .no_longs = 2 * LONG_PER_U64 * (NODESIZE / sizeof(long) / (1 + 2 * LONG_PER_U64)),73};74EXPORT_SYMBOL_GPL(btree_geo128);75 76#define MAX_KEYLEN (2 * LONG_PER_U64)77 78static struct kmem_cache *btree_cachep;79 80void *btree_alloc(gfp_t gfp_mask, void *pool_data)81{82 return kmem_cache_alloc(btree_cachep, gfp_mask);83}84EXPORT_SYMBOL_GPL(btree_alloc);85 86void btree_free(void *element, void *pool_data)87{88 kmem_cache_free(btree_cachep, element);89}90EXPORT_SYMBOL_GPL(btree_free);91 92static unsigned long *btree_node_alloc(struct btree_head *head, gfp_t gfp)93{94 unsigned long *node;95 96 node = mempool_alloc(head->mempool, gfp);97 if (likely(node))98 memset(node, 0, NODESIZE);99 return node;100}101 102static int longcmp(const unsigned long *l1, const unsigned long *l2, size_t n)103{104 size_t i;105 106 for (i = 0; i < n; i++) {107 if (l1[i] < l2[i])108 return -1;109 if (l1[i] > l2[i])110 return 1;111 }112 return 0;113}114 115static unsigned long *longcpy(unsigned long *dest, const unsigned long *src,116 size_t n)117{118 size_t i;119 120 for (i = 0; i < n; i++)121 dest[i] = src[i];122 return dest;123}124 125static unsigned long *longset(unsigned long *s, unsigned long c, size_t n)126{127 size_t i;128 129 for (i = 0; i < n; i++)130 s[i] = c;131 return s;132}133 134static void dec_key(struct btree_geo *geo, unsigned long *key)135{136 unsigned long val;137 int i;138 139 for (i = geo->keylen - 1; i >= 0; i--) {140 val = key[i];141 key[i] = val - 1;142 if (val)143 break;144 }145}146 147static unsigned long *bkey(struct btree_geo *geo, unsigned long *node, int n)148{149 return &node[n * geo->keylen];150}151 152static void *bval(struct btree_geo *geo, unsigned long *node, int n)153{154 return (void *)node[geo->no_longs + n];155}156 157static void setkey(struct btree_geo *geo, unsigned long *node, int n,158 unsigned long *key)159{160 longcpy(bkey(geo, node, n), key, geo->keylen);161}162 163static void setval(struct btree_geo *geo, unsigned long *node, int n,164 void *val)165{166 node[geo->no_longs + n] = (unsigned long) val;167}168 169static void clearpair(struct btree_geo *geo, unsigned long *node, int n)170{171 longset(bkey(geo, node, n), 0, geo->keylen);172 node[geo->no_longs + n] = 0;173}174 175static inline void __btree_init(struct btree_head *head)176{177 head->node = NULL;178 head->height = 0;179}180 181void btree_init_mempool(struct btree_head *head, mempool_t *mempool)182{183 __btree_init(head);184 head->mempool = mempool;185}186EXPORT_SYMBOL_GPL(btree_init_mempool);187 188int btree_init(struct btree_head *head)189{190 __btree_init(head);191 head->mempool = mempool_create(0, btree_alloc, btree_free, NULL);192 if (!head->mempool)193 return -ENOMEM;194 return 0;195}196EXPORT_SYMBOL_GPL(btree_init);197 198void btree_destroy(struct btree_head *head)199{200 mempool_free(head->node, head->mempool);201 mempool_destroy(head->mempool);202 head->mempool = NULL;203}204EXPORT_SYMBOL_GPL(btree_destroy);205 206void *btree_last(struct btree_head *head, struct btree_geo *geo,207 unsigned long *key)208{209 int height = head->height;210 unsigned long *node = head->node;211 212 if (height == 0)213 return NULL;214 215 for ( ; height > 1; height--)216 node = bval(geo, node, 0);217 218 longcpy(key, bkey(geo, node, 0), geo->keylen);219 return bval(geo, node, 0);220}221EXPORT_SYMBOL_GPL(btree_last);222 223static int keycmp(struct btree_geo *geo, unsigned long *node, int pos,224 unsigned long *key)225{226 return longcmp(bkey(geo, node, pos), key, geo->keylen);227}228 229static int keyzero(struct btree_geo *geo, unsigned long *key)230{231 int i;232 233 for (i = 0; i < geo->keylen; i++)234 if (key[i])235 return 0;236 237 return 1;238}239 240static void *btree_lookup_node(struct btree_head *head, struct btree_geo *geo,241 unsigned long *key)242{243 int i, height = head->height;244 unsigned long *node = head->node;245 246 if (height == 0)247 return NULL;248 249 for ( ; height > 1; height--) {250 for (i = 0; i < geo->no_pairs; i++)251 if (keycmp(geo, node, i, key) <= 0)252 break;253 if (i == geo->no_pairs)254 return NULL;255 node = bval(geo, node, i);256 if (!node)257 return NULL;258 }259 return node;260}261 262void *btree_lookup(struct btree_head *head, struct btree_geo *geo,263 unsigned long *key)264{265 int i;266 unsigned long *node;267 268 node = btree_lookup_node(head, geo, key);269 if (!node)270 return NULL;271 272 for (i = 0; i < geo->no_pairs; i++)273 if (keycmp(geo, node, i, key) == 0)274 return bval(geo, node, i);275 return NULL;276}277EXPORT_SYMBOL_GPL(btree_lookup);278 279int btree_update(struct btree_head *head, struct btree_geo *geo,280 unsigned long *key, void *val)281{282 int i;283 unsigned long *node;284 285 node = btree_lookup_node(head, geo, key);286 if (!node)287 return -ENOENT;288 289 for (i = 0; i < geo->no_pairs; i++)290 if (keycmp(geo, node, i, key) == 0) {291 setval(geo, node, i, val);292 return 0;293 }294 return -ENOENT;295}296EXPORT_SYMBOL_GPL(btree_update);297 298/*299 * Usually this function is quite similar to normal lookup. But the key of300 * a parent node may be smaller than the smallest key of all its siblings.301 * In such a case we cannot just return NULL, as we have only proven that no302 * key smaller than __key, but larger than this parent key exists.303 * So we set __key to the parent key and retry. We have to use the smallest304 * such parent key, which is the last parent key we encountered.305 */306void *btree_get_prev(struct btree_head *head, struct btree_geo *geo,307 unsigned long *__key)308{309 int i, height;310 unsigned long *node, *oldnode;311 unsigned long *retry_key = NULL, key[MAX_KEYLEN];312 313 if (keyzero(geo, __key))314 return NULL;315 316 if (head->height == 0)317 return NULL;318 longcpy(key, __key, geo->keylen);319retry:320 dec_key(geo, key);321 322 node = head->node;323 for (height = head->height ; height > 1; height--) {324 for (i = 0; i < geo->no_pairs; i++)325 if (keycmp(geo, node, i, key) <= 0)326 break;327 if (i == geo->no_pairs)328 goto miss;329 oldnode = node;330 node = bval(geo, node, i);331 if (!node)332 goto miss;333 retry_key = bkey(geo, oldnode, i);334 }335 336 if (!node)337 goto miss;338 339 for (i = 0; i < geo->no_pairs; i++) {340 if (keycmp(geo, node, i, key) <= 0) {341 if (bval(geo, node, i)) {342 longcpy(__key, bkey(geo, node, i), geo->keylen);343 return bval(geo, node, i);344 } else345 goto miss;346 }347 }348miss:349 if (retry_key) {350 longcpy(key, retry_key, geo->keylen);351 retry_key = NULL;352 goto retry;353 }354 return NULL;355}356EXPORT_SYMBOL_GPL(btree_get_prev);357 358static int getpos(struct btree_geo *geo, unsigned long *node,359 unsigned long *key)360{361 int i;362 363 for (i = 0; i < geo->no_pairs; i++) {364 if (keycmp(geo, node, i, key) <= 0)365 break;366 }367 return i;368}369 370static int getfill(struct btree_geo *geo, unsigned long *node, int start)371{372 int i;373 374 for (i = start; i < geo->no_pairs; i++)375 if (!bval(geo, node, i))376 break;377 return i;378}379 380/*381 * locate the correct leaf node in the btree382 */383static unsigned long *find_level(struct btree_head *head, struct btree_geo *geo,384 unsigned long *key, int level)385{386 unsigned long *node = head->node;387 int i, height;388 389 for (height = head->height; height > level; height--) {390 for (i = 0; i < geo->no_pairs; i++)391 if (keycmp(geo, node, i, key) <= 0)392 break;393 394 if ((i == geo->no_pairs) || !bval(geo, node, i)) {395 /* right-most key is too large, update it */396 /* FIXME: If the right-most key on higher levels is397 * always zero, this wouldn't be necessary. */398 i--;399 setkey(geo, node, i, key);400 }401 BUG_ON(i < 0);402 node = bval(geo, node, i);403 }404 BUG_ON(!node);405 return node;406}407 408static int btree_grow(struct btree_head *head, struct btree_geo *geo,409 gfp_t gfp)410{411 unsigned long *node;412 int fill;413 414 node = btree_node_alloc(head, gfp);415 if (!node)416 return -ENOMEM;417 if (head->node) {418 fill = getfill(geo, head->node, 0);419 setkey(geo, node, 0, bkey(geo, head->node, fill - 1));420 setval(geo, node, 0, head->node);421 }422 head->node = node;423 head->height++;424 return 0;425}426 427static void btree_shrink(struct btree_head *head, struct btree_geo *geo)428{429 unsigned long *node;430 int fill;431 432 if (head->height <= 1)433 return;434 435 node = head->node;436 fill = getfill(geo, node, 0);437 BUG_ON(fill > 1);438 head->node = bval(geo, node, 0);439 head->height--;440 mempool_free(node, head->mempool);441}442 443static int btree_insert_level(struct btree_head *head, struct btree_geo *geo,444 unsigned long *key, void *val, int level,445 gfp_t gfp)446{447 unsigned long *node;448 int i, pos, fill, err;449 450 BUG_ON(!val);451 if (head->height < level) {452 err = btree_grow(head, geo, gfp);453 if (err)454 return err;455 }456 457retry:458 node = find_level(head, geo, key, level);459 pos = getpos(geo, node, key);460 fill = getfill(geo, node, pos);461 /* two identical keys are not allowed */462 BUG_ON(pos < fill && keycmp(geo, node, pos, key) == 0);463 464 if (fill == geo->no_pairs) {465 /* need to split node */466 unsigned long *new;467 468 new = btree_node_alloc(head, gfp);469 if (!new)470 return -ENOMEM;471 err = btree_insert_level(head, geo,472 bkey(geo, node, fill / 2 - 1),473 new, level + 1, gfp);474 if (err) {475 mempool_free(new, head->mempool);476 return err;477 }478 for (i = 0; i < fill / 2; i++) {479 setkey(geo, new, i, bkey(geo, node, i));480 setval(geo, new, i, bval(geo, node, i));481 setkey(geo, node, i, bkey(geo, node, i + fill / 2));482 setval(geo, node, i, bval(geo, node, i + fill / 2));483 clearpair(geo, node, i + fill / 2);484 }485 if (fill & 1) {486 setkey(geo, node, i, bkey(geo, node, fill - 1));487 setval(geo, node, i, bval(geo, node, fill - 1));488 clearpair(geo, node, fill - 1);489 }490 goto retry;491 }492 BUG_ON(fill >= geo->no_pairs);493 494 /* shift and insert */495 for (i = fill; i > pos; i--) {496 setkey(geo, node, i, bkey(geo, node, i - 1));497 setval(geo, node, i, bval(geo, node, i - 1));498 }499 setkey(geo, node, pos, key);500 setval(geo, node, pos, val);501 502 return 0;503}504 505int btree_insert(struct btree_head *head, struct btree_geo *geo,506 unsigned long *key, void *val, gfp_t gfp)507{508 BUG_ON(!val);509 return btree_insert_level(head, geo, key, val, 1, gfp);510}511EXPORT_SYMBOL_GPL(btree_insert);512 513static void *btree_remove_level(struct btree_head *head, struct btree_geo *geo,514 unsigned long *key, int level);515static void merge(struct btree_head *head, struct btree_geo *geo, int level,516 unsigned long *left, int lfill,517 unsigned long *right, int rfill,518 unsigned long *parent, int lpos)519{520 int i;521 522 for (i = 0; i < rfill; i++) {523 /* Move all keys to the left */524 setkey(geo, left, lfill + i, bkey(geo, right, i));525 setval(geo, left, lfill + i, bval(geo, right, i));526 }527 /* Exchange left and right child in parent */528 setval(geo, parent, lpos, right);529 setval(geo, parent, lpos + 1, left);530 /* Remove left (formerly right) child from parent */531 btree_remove_level(head, geo, bkey(geo, parent, lpos), level + 1);532 mempool_free(right, head->mempool);533}534 535static void rebalance(struct btree_head *head, struct btree_geo *geo,536 unsigned long *key, int level, unsigned long *child, int fill)537{538 unsigned long *parent, *left = NULL, *right = NULL;539 int i, no_left, no_right;540 541 if (fill == 0) {542 /* Because we don't steal entries from a neighbour, this case543 * can happen. Parent node contains a single child, this544 * node, so merging with a sibling never happens.545 */546 btree_remove_level(head, geo, key, level + 1);547 mempool_free(child, head->mempool);548 return;549 }550 551 parent = find_level(head, geo, key, level + 1);552 i = getpos(geo, parent, key);553 BUG_ON(bval(geo, parent, i) != child);554 555 if (i > 0) {556 left = bval(geo, parent, i - 1);557 no_left = getfill(geo, left, 0);558 if (fill + no_left <= geo->no_pairs) {559 merge(head, geo, level,560 left, no_left,561 child, fill,562 parent, i - 1);563 return;564 }565 }566 if (i + 1 < getfill(geo, parent, i)) {567 right = bval(geo, parent, i + 1);568 no_right = getfill(geo, right, 0);569 if (fill + no_right <= geo->no_pairs) {570 merge(head, geo, level,571 child, fill,572 right, no_right,573 parent, i);574 return;575 }576 }577 /*578 * We could also try to steal one entry from the left or right579 * neighbor. By not doing so we changed the invariant from580 * "all nodes are at least half full" to "no two neighboring581 * nodes can be merged". Which means that the average fill of582 * all nodes is still half or better.583 */584}585 586static void *btree_remove_level(struct btree_head *head, struct btree_geo *geo,587 unsigned long *key, int level)588{589 unsigned long *node;590 int i, pos, fill;591 void *ret;592 593 if (level > head->height) {594 /* we recursed all the way up */595 head->height = 0;596 head->node = NULL;597 return NULL;598 }599 600 node = find_level(head, geo, key, level);601 pos = getpos(geo, node, key);602 fill = getfill(geo, node, pos);603 if ((level == 1) && (keycmp(geo, node, pos, key) != 0))604 return NULL;605 ret = bval(geo, node, pos);606 607 /* remove and shift */608 for (i = pos; i < fill - 1; i++) {609 setkey(geo, node, i, bkey(geo, node, i + 1));610 setval(geo, node, i, bval(geo, node, i + 1));611 }612 clearpair(geo, node, fill - 1);613 614 if (fill - 1 < geo->no_pairs / 2) {615 if (level < head->height)616 rebalance(head, geo, key, level, node, fill - 1);617 else if (fill - 1 == 1)618 btree_shrink(head, geo);619 }620 621 return ret;622}623 624void *btree_remove(struct btree_head *head, struct btree_geo *geo,625 unsigned long *key)626{627 if (head->height == 0)628 return NULL;629 630 return btree_remove_level(head, geo, key, 1);631}632EXPORT_SYMBOL_GPL(btree_remove);633 634int btree_merge(struct btree_head *target, struct btree_head *victim,635 struct btree_geo *geo, gfp_t gfp)636{637 unsigned long key[MAX_KEYLEN];638 unsigned long dup[MAX_KEYLEN];639 void *val;640 int err;641 642 BUG_ON(target == victim);643 644 if (!(target->node)) {645 /* target is empty, just copy fields over */646 target->node = victim->node;647 target->height = victim->height;648 __btree_init(victim);649 return 0;650 }651 652 /* TODO: This needs some optimizations. Currently we do three tree653 * walks to remove a single object from the victim.654 */655 for (;;) {656 if (!btree_last(victim, geo, key))657 break;658 val = btree_lookup(victim, geo, key);659 err = btree_insert(target, geo, key, val, gfp);660 if (err)661 return err;662 /* We must make a copy of the key, as the original will get663 * mangled inside btree_remove. */664 longcpy(dup, key, geo->keylen);665 btree_remove(victim, geo, dup);666 }667 return 0;668}669EXPORT_SYMBOL_GPL(btree_merge);670 671static size_t __btree_for_each(struct btree_head *head, struct btree_geo *geo,672 unsigned long *node, unsigned long opaque,673 void (*func)(void *elem, unsigned long opaque,674 unsigned long *key, size_t index,675 void *func2),676 void *func2, int reap, int height, size_t count)677{678 int i;679 unsigned long *child;680 681 for (i = 0; i < geo->no_pairs; i++) {682 child = bval(geo, node, i);683 if (!child)684 break;685 if (height > 1)686 count = __btree_for_each(head, geo, child, opaque,687 func, func2, reap, height - 1, count);688 else689 func(child, opaque, bkey(geo, node, i), count++,690 func2);691 }692 if (reap)693 mempool_free(node, head->mempool);694 return count;695}696 697static void empty(void *elem, unsigned long opaque, unsigned long *key,698 size_t index, void *func2)699{700}701 702void visitorl(void *elem, unsigned long opaque, unsigned long *key,703 size_t index, void *__func)704{705 visitorl_t func = __func;706 707 func(elem, opaque, *key, index);708}709EXPORT_SYMBOL_GPL(visitorl);710 711void visitor32(void *elem, unsigned long opaque, unsigned long *__key,712 size_t index, void *__func)713{714 visitor32_t func = __func;715 u32 *key = (void *)__key;716 717 func(elem, opaque, *key, index);718}719EXPORT_SYMBOL_GPL(visitor32);720 721void visitor64(void *elem, unsigned long opaque, unsigned long *__key,722 size_t index, void *__func)723{724 visitor64_t func = __func;725 u64 *key = (void *)__key;726 727 func(elem, opaque, *key, index);728}729EXPORT_SYMBOL_GPL(visitor64);730 731void visitor128(void *elem, unsigned long opaque, unsigned long *__key,732 size_t index, void *__func)733{734 visitor128_t func = __func;735 u64 *key = (void *)__key;736 737 func(elem, opaque, key[0], key[1], index);738}739EXPORT_SYMBOL_GPL(visitor128);740 741size_t btree_visitor(struct btree_head *head, struct btree_geo *geo,742 unsigned long opaque,743 void (*func)(void *elem, unsigned long opaque,744 unsigned long *key,745 size_t index, void *func2),746 void *func2)747{748 size_t count = 0;749 750 if (!func2)751 func = empty;752 if (head->node)753 count = __btree_for_each(head, geo, head->node, opaque, func,754 func2, 0, head->height, 0);755 return count;756}757EXPORT_SYMBOL_GPL(btree_visitor);758 759size_t btree_grim_visitor(struct btree_head *head, struct btree_geo *geo,760 unsigned long opaque,761 void (*func)(void *elem, unsigned long opaque,762 unsigned long *key,763 size_t index, void *func2),764 void *func2)765{766 size_t count = 0;767 768 if (!func2)769 func = empty;770 if (head->node)771 count = __btree_for_each(head, geo, head->node, opaque, func,772 func2, 1, head->height, 0);773 __btree_init(head);774 return count;775}776EXPORT_SYMBOL_GPL(btree_grim_visitor);777 778static int __init btree_module_init(void)779{780 btree_cachep = kmem_cache_create("btree_node", NODESIZE, 0,781 SLAB_HWCACHE_ALIGN, NULL);782 return 0;783}784 785static void __exit btree_module_exit(void)786{787 kmem_cache_destroy(btree_cachep);788}789 790/* If core code starts using btree, initialization should happen even earlier */791module_init(btree_module_init);792module_exit(btree_module_exit);793 794MODULE_AUTHOR("Joern Engel <joern@logfs.org>");795MODULE_AUTHOR("Johannes Berg <johannes@sipsolutions.net>");796