From e4443b735dd1c64853d863da06b62a10bb8c73f2 Mon Sep 17 00:00:00 2001 From: Lunaixsky Date: Tue, 10 Dec 2024 02:43:53 +0000 Subject: [PATCH] rewrite the btrie key allocation algorithm for better uniformity --- lunaix-os/kernel/ds/btrie.c | 244 ++++++++++++++++++++++-------------- 1 file changed, 150 insertions(+), 94 deletions(-) diff --git a/lunaix-os/kernel/ds/btrie.c b/lunaix-os/kernel/ds/btrie.c index 8cba604..37bf803 100644 --- a/lunaix-os/kernel/ds/btrie.c +++ b/lunaix-os/kernel/ds/btrie.c @@ -13,25 +13,13 @@ #include #include -enum btrie_traverse_mode { +enum btrie_query_mode { BTRIE_FIND, BTRIE_FIND_PARENT, BTRIE_INSERT }; -struct btrie_locator -{ - enum btrie_traverse_mode mode; - int shifts; - unsigned long index; -}; -static inline void -__btrie_locator(struct btrie_locator* loc, - unsigned long index, enum btrie_traverse_mode mode) -{ - loc->mode = mode; - loc->index = index; -} +#define remove_visit_flag(node) ((struct btrie_node*)(__ptr(node) & ~1)) static inline void __btrie_addchild(struct btrie* root, @@ -48,11 +36,42 @@ __btrie_addchild(struct btrie* root, static inline struct btrie_node* __btrie_get_child(struct btrie_node* node, int index) { - if (unlikely(!node->children)) { + if (unlikely(!node || !node->children)) { return NULL; } - return node->children[index]; + return remove_visit_flag(node->children[index]); +} + +static inline bool +__btrie_child_visited(struct btrie_node* node, int index) +{ + if (unlikely(!node || !node->children)) { + return false; + } + + return !!(__ptr(node->children[index]) & 1); +} + +static inline void +__btrie_visit_child(struct btrie_node* node, int index) +{ + if (unlikely(!node || !node->children)) { + return; + } + + node->children[index] + = (struct btrie_node*)(__ptr(node->children[index]) | 1); +} + +static inline void +__btrie_unvisit_child(struct btrie_node* node, int index) +{ + if (unlikely(!node || !node->children)) { + return; + } + + node->children[index] = remove_visit_flag(node->children[index]); } static inline void @@ -86,109 +105,152 @@ __btrie_create(struct btrie* root, struct btrie_node* parent, int index) } static struct btrie_node* -__btrie_traversal(struct btrie* root, struct btrie_locator* locator) +__btrie_traversal(struct btrie* root, + unsigned long key, enum btrie_query_mode mode) { - unsigned long lz, index; - unsigned long bitmask; - unsigned long i = 0; + unsigned long lz, bitmask, i = 0; struct btrie_node *subtree, *tree = root->btrie_root; - - index = locator->index; - lz = index ? (msbitl - clzl(index)) / root->order : 0; + + lz = !key ? 0 : (msbitl - clzl(key)) / root->order; lz = lz * root->order; bitmask = ((1 << root->order) - 1) << lz; - // Time complexity: O(log_2(log_2(N))) where N is the index to lookup + // Time complexity: O(log_2(log_2(N))) where N is the key to lookup while (bitmask && tree) { - i = (index & bitmask) >> lz; + i = (key & bitmask) >> lz; subtree = __btrie_get_child(tree, i); - if (!subtree && (locator->mode != BTRIE_FIND)) - { + if (!lz && mode == BTRIE_FIND_PARENT) { + break; + } + + if (!subtree && (mode != BTRIE_FIND)) { tree = __btrie_create(root, tree, i); } else { tree = subtree; } - if (lz == root->order && locator->mode == BTRIE_FIND_PARENT) - { - break; - } - bitmask = bitmask >> root->order; lz -= root->order; } - locator->shifts = lz; return tree; } -static unsigned long -__btrie_alloc_slot(struct btrie* root, struct btrie_node **slot, - unsigned long start, unsigned long end) +#define check_partial(loc, start, end, mask) \ + (((long)(start) - (long)mask - 1) <= (long)loc && loc < end) + +static inline struct btrie_node* +__get_immediate_node(struct btrie* tree, struct btrie_node *root, + unsigned long start, unsigned long end, unsigned long loc) { - unsigned int index, od; - unsigned long result, mask, result_next; - struct btrie_node *tree, *subtree; - struct btrie_locator locator; - - result = 0; - od = root->order; - mask = (1 << od) - 1; - tree = root->btrie_root; - - if (start) { - __btrie_locator(&locator, start, BTRIE_FIND_PARENT); - tree = __btrie_traversal(root, &locator); - result = start; - } - - while (tree && result < end) - { - index = result & mask; - result_next = result << od; - - subtree = __btrie_get_child(tree, index); - - if (subtree) { - if (result >= start && !subtree->data) { - tree = subtree; - goto done; - } + unsigned int index; + unsigned long mask; + struct btrie_node *node; + + mask = (1UL << tree->order) - 1; + index = loc & mask; + for (int i = 0; i <= mask; ++i, index = (index + 1) & mask) + { + loc = (loc & ~mask) + index; + node = __btrie_get_child(root, index); + + if (loc < start || loc >= end) { + continue; } - else { - subtree = __btrie_create(root, tree, index); - if (result >= start) - goto done; + if (!node) { + return __btrie_create(tree, root, index); + } - goto descend; + if (!node->data) { + return node; } + } + + return NULL; +} + +static struct btrie_node* +__btrie_find_free(struct btrie* tree, struct btrie_node *root, + unsigned long start, unsigned long end, unsigned long loc) +{ + unsigned long mask, next_loc; + struct btrie_node *node, *found; + + if (!root) return NULL; - if (!__full_node(root, tree) && index != mask) { - result++; + mask = (1UL << tree->order) - 1; + + __btrie_visit_child(root->parent, root->index); + + if (!__full_node(tree, root)) { + found = __get_immediate_node(tree, root, start, end, loc); + if (found) goto done; + } + + for (unsigned int i = 0; i <= mask; i++) + { + next_loc = ((loc & ~mask) + i) << tree->order; + + if (!next_loc || !check_partial(next_loc, start, end, mask)) { continue; } -descend: - if (result_next >= end) { - tree = tree->parent; - result >>= od; - result++; + if (__btrie_child_visited(root, i)) { continue; } - tree = subtree; - result = result_next; + node = __btrie_get_child(root, i); + node = node ?: __btrie_create(tree, root, i); + + found = __btrie_find_free(tree, node, start, end, next_loc); + + if (found) { + goto done; + } } - *slot = NULL; - return -1; + loc >>= tree->order; + found = __btrie_find_free(tree, root->parent, start, end, loc + 1); done: - *slot = subtree; - return result; + __btrie_unvisit_child(root->parent, root->index); + return found; +} + +static unsigned long +__btrie_alloc_slot(struct btrie* tree, struct btrie_node **slot, + unsigned long start, unsigned long end) +{ + unsigned int od; + unsigned long result, mask; + struct btrie_node *found, *node; + + od = 0; + mask = (1 << od) - 1; + found = tree->btrie_root; + result = 0; + + if (!start && !__btrie_get_child(found, 0)) { + *slot = __btrie_create(tree, found, 0); + return 0; + } + + found = __btrie_traversal(tree, start, BTRIE_FIND_PARENT); + found = __btrie_find_free(tree, found, start, end, start); + + node = found; + while (node) + { + result |= node->index << od; + od += tree->order; + node = node->parent; + } + + *slot = found; + return found ? result : -1UL; } void @@ -204,10 +266,8 @@ void* btrie_get(struct btrie* root, unsigned long index) { struct btrie_node* node; - struct btrie_locator locator; - - __btrie_locator(&locator, index, BTRIE_FIND); - node = __btrie_traversal(root, &locator); + + node = __btrie_traversal(root, index, BTRIE_FIND); if (!node) { return NULL; } @@ -219,10 +279,8 @@ void btrie_set(struct btrie* root, unsigned long index, void* data) { struct btrie_node* node; - struct btrie_locator locator; - - __btrie_locator(&locator, index, BTRIE_INSERT); - node = __btrie_traversal(root, &locator); + + node = __btrie_traversal(root, index, BTRIE_INSERT); node->data = data; } @@ -265,10 +323,8 @@ btrie_remove(struct btrie* root, unsigned long index) { void* data; struct btrie_node* node; - struct btrie_locator locator; - - __btrie_locator(&locator, index, BTRIE_FIND); - node = __btrie_traversal(root, &locator); + + node = __btrie_traversal(root, index, BTRIE_FIND); if (!node) { return 0; } -- 2.27.0