mirror of
https://github.com/Mathewnd/Astral.git
synced 2026-10-27 12:48:22 +00:00
lib: trie implementation
This commit is contained in:
Matheus
@@ -0,0 +1,35 @@
|
||||
#ifndef _TRIE_H
|
||||
#define _TRIE_H
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#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
|
||||
@@ -0,0 +1,305 @@
|
||||
#include <trie.h>
|
||||
#include <errno.h>
|
||||
#include <stddef.h>
|
||||
#include <stdbool.h>
|
||||
#include <string.h>
|
||||
#include <kernel/alloc.h>
|
||||
#include <logging.h>
|
||||
|
||||
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);
|
||||
}
|
||||
Reference in New Issue
Block a user