From e630b7fa11de0fedcf72dabc920d1f42da11f8dc Mon Sep 17 00:00:00 2001 From: Matheus Date: Sun, 12 Jul 2026 13:28:13 -0300 Subject: [PATCH] lib: trie implementation --- kernel-src/include/trie.h | 35 +++++ kernel-src/lib/trie.c | 305 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 340 insertions(+) create mode 100644 kernel-src/include/trie.h create mode 100644 kernel-src/lib/trie.c diff --git a/kernel-src/include/trie.h b/kernel-src/include/trie.h new file mode 100644 index 00000000..cf322b0b --- /dev/null +++ b/kernel-src/include/trie.h @@ -0,0 +1,35 @@ +#ifndef _TRIE_H +#define _TRIE_H + +#include + +#define TRIE_BITS 6 +#define TRIE_NODE_COUNT (1 << TRIE_BITS) + +typedef struct trie_node trie_node_t; +struct trie_node { + uint8_t count; + union { + trie_node_t *nodes[TRIE_NODE_COUNT]; + void *items[TRIE_NODE_COUNT]; + }; +}; + +typedef struct { + uint8_t height; + union { + void *root_item; + trie_node_t *root; + }; +} trie_t; + +typedef void (trie_iterate_fn_t)(void *item); + +void trie_init(trie_t *trie); +int trie_insert(trie_t *trie, uint64_t key, void *value); +int trie_lookup(trie_t *trie, uint64_t key, void **ret); +int trie_remove(trie_t *trie, uint64_t key); +void trie_truncate(trie_t *trie, uint64_t max_key); +void trie_iterate(trie_t *trie, uint64_t min_key, uint64_t max_key, trie_iterate_fn_t); + +#endif diff --git a/kernel-src/lib/trie.c b/kernel-src/lib/trie.c new file mode 100644 index 00000000..33fcf568 --- /dev/null +++ b/kernel-src/lib/trie.c @@ -0,0 +1,305 @@ +#include +#include +#include +#include +#include +#include +#include + +void trie_init(trie_t *trie) { + trie->height = 0; + trie->root = NULL; +} + +static inline uint8_t key_height(uint64_t key) { + if (key == 0) + return 0; + + return (64 - __builtin_clzll(key) + TRIE_BITS - 1) / TRIE_BITS; +} + +static inline uint8_t get_offset(uint64_t key, uint8_t height) { + return (key >> (height * TRIE_BITS)) & (TRIE_NODE_COUNT - 1); +} + +static void trie_trim(trie_t *trie) { + if (trie->root->count == 0) { + free(trie->root); + trie->root = NULL; + trie->height = 0; + } else while (trie->height && trie->root->count == 1 && trie->root->nodes[0]) { + --trie->height; + void *old_root = trie->root; + trie->root = trie->root->nodes[0]; + free(old_root); + } +} + +static int trie_insert_internal(trie_node_t *trie_node, uint64_t key, void *value, uint8_t height) { + uint8_t offset = get_offset(key, height); + if (height == 0) { + if (trie_node->items[offset]) + return EEXIST; + + trie_node->items[offset] = value; + ++trie_node->count; + return 0; + } + + if (trie_node->nodes[offset] == NULL) { + trie_node->nodes[offset] = alloc(sizeof(trie_node_t)); + if (trie_node->nodes[offset] == NULL) + return ENOMEM; + + memset(trie_node->nodes[offset], 0, sizeof(trie_node_t)); + + trie_node->nodes[offset]->count = 0; + ++trie_node->count; + } + + int error = trie_insert_internal(trie_node->nodes[offset], key, value, height - 1); + if (error && trie_node->nodes[offset]->count == 0) { + free(trie_node->nodes[offset]); + trie_node->nodes[offset] = NULL; + --trie_node->count; + } + + return error; +} + +static bool expand_trie(trie_t *trie, uint8_t difference) { + if (difference == 0) + return true; + + trie_node_t *new = alloc(sizeof(trie_node_t)); + if (new == NULL) + return false; + + memset(new, 0, sizeof(trie_node_t)); + + if (trie->root) { + new->nodes[0] = trie->root; + new->count = 1; + } + trie->root = new; + ++trie->height; + + if (!expand_trie(trie, difference - 1)) { + trie->root = new->nodes[0]; + --trie->height; + free(new); + return false; + } + + return true; +} + +static void undo_expand(trie_t *trie, uint8_t difference) { + if (difference == 0) + return; + + trie_node_t *trie_node = trie->root; + trie->root = trie_node->nodes[0]; + free(trie_node); + --trie->height; + + undo_expand(trie, difference - 1); +} + +int trie_insert(trie_t *trie, uint64_t key, void *value) { + __assert(value); + + uint8_t height = key_height(key); + + if (height == 0 && trie->height == 0) { + if (trie->root_item) + return EEXIST; + + trie->root_item = value; + return 0; + } + + uint8_t difference = height > trie->height ? height - trie->height : 0; + if (!expand_trie(trie, difference)) + return ENOMEM; + + int error = trie_insert_internal(trie->root, key, value, trie->height - 1); + if (error) + undo_expand(trie, difference); + + return error; +} + +static int trie_lookup_internal(trie_node_t *trie_node, uint64_t key, uint8_t height, void **ret) { + if (trie_node == NULL) + return ENOENT; + + if (height == 0) { + *ret = trie_node->items[get_offset(key, height)]; + return *ret ? 0 : ENOENT; + } + + return trie_lookup_internal(trie_node->nodes[get_offset(key, height)], key, height - 1, ret); +} + +int trie_lookup(trie_t *trie, uint64_t key, void **ret) { + uint8_t height = key_height(key); + if (height > trie->height) + return ENOENT; + + if (trie->height == 0) { + *ret = trie->root_item; + return *ret ? 0 : ENOENT; + } + + return trie_lookup_internal(trie->root, key, trie->height - 1, ret); +} + +static int trie_remove_internal(trie_node_t *trie_node, uint64_t key, uint8_t height) { + if (trie_node == NULL) + return ENOENT; + + if (height == 0) { + uint8_t offset = get_offset(key, height); + if (trie_node->items[offset] == NULL) + return ENOENT; + + trie_node->items[offset] = NULL; + --trie_node->count; + return 0; + } + + int error = trie_remove_internal(trie_node->nodes[get_offset(key, height)], key, height - 1); + if (error) + return error; + + if (trie_node->nodes[get_offset(key, height)]->count == 0) { + free(trie_node->nodes[get_offset(key, height)]); + trie_node->nodes[get_offset(key, height)] = NULL; + --trie_node->count; + } + + return 0; +} + +int trie_remove(trie_t *trie, uint64_t key) { + uint8_t height = key_height(key); + if (height > trie->height) + return ENOENT; + + if (trie->height == 0) { + if (trie->root_item == NULL) + return ENOENT; + + trie->root_item = NULL; + return 0; + } + + int error = trie_remove_internal(trie->root, key, trie->height - 1); + if (error) + return error; + + trie_trim(trie); + + return 0; +} + +static void trie_recursively_free(trie_node_t *trie_node, uint8_t height) { + if (trie_node == NULL) + return; + + if (height) { + for (uint8_t i = 0; i < TRIE_NODE_COUNT; ++i) + trie_recursively_free(trie_node->nodes[i], height - 1); + } + + free(trie_node); +} + +static void trie_truncate_internal(trie_node_t *trie_node, uint64_t max_key, uint8_t height) { + uint8_t offset = get_offset(max_key, height); + + if (height) { + trie_node_t *child_node = trie_node->nodes[offset]; + if (child_node) { + trie_truncate_internal(child_node, max_key, height - 1); + + if (child_node->count == 0) { + free(child_node); + trie_node->nodes[offset] = NULL; + --trie_node->count; + } + } + } + + for (uint8_t i = offset + (height == 0 ? 0 : 1); i < TRIE_NODE_COUNT; ++i) { + if (trie_node->nodes[i] == NULL) + continue; + + if (height) + trie_recursively_free(trie_node->nodes[i], height - 1); + + trie_node->nodes[i] = NULL; + --trie_node->count; + } +} + +void trie_truncate(trie_t *trie, uint64_t max_key) { + if (key_height(max_key) > trie->height) + return; + + if (max_key == 0) { + if (trie->height) + trie_recursively_free(trie->root, trie->height - 1); + + trie->root_item = NULL; + trie->height = 0; + return; + } + + if (max_key == 1) { + void *item; + if (trie_lookup(trie, 0, &item)) + item = NULL; + + if (trie->height) + trie_recursively_free(trie->root, trie->height - 1); + + trie->root_item = item; + trie->height = 0; + return; + } + + trie_truncate_internal(trie->root, max_key, trie->height - 1); + + trie_trim(trie); +} + +static void trie_iterate_internal(trie_node_t *trie_node, uint64_t min_key, uint64_t max_key, uint8_t height, trie_iterate_fn_t fn, bool edge_low, bool edge_high) { + if (trie_node == NULL) + return; + + for (uint8_t i = 0; i < TRIE_NODE_COUNT; ++i) { + if ((edge_low && i < get_offset(min_key, height)) || (edge_high && i > get_offset(max_key, height))) + continue; + + if (height) + trie_iterate_internal(trie_node->nodes[i], min_key, max_key, height - 1, fn, edge_low && i == get_offset(min_key, height), edge_high && i == get_offset(max_key, height)); + else if (trie_node->items[i]) + fn(trie_node->items[i]); + } +} + +void trie_iterate(trie_t *trie, uint64_t min_key, uint64_t max_key, trie_iterate_fn_t fn) { + if (key_height(min_key) > trie->height || min_key > max_key) + return; + + if (trie->height == 0) { + if (min_key == 0 && trie->root_item) + fn(trie->root_item); + + return; + } + + trie_iterate_internal(trie->root, min_key, max_key, trie->height - 1, fn, true, key_height(max_key) <= trie->height); +}