lunaix-os/kernel/ds/btrie.c

   1 /**
   2  * @file btrie.c
   3  * @author Lunaixsky
   4  * @brief Bitwise trie tree implementation for sparse array.
   5  * @version 0.1
   6  * @date 2022-08-01
   7  *
   8  * @copyright Copyright (c) 2022
   9  *
  10  */
  11
  12 #include <lunaix/ds/btrie.h>
  13 #include <lunaix/mm/valloc.h>
  14 #include <lunaix/spike.h>
  15
  16 #define BTRIE_INSERT 1
  17
  18 struct btrie_node*
  19 __btrie_traversal(struct btrie* root, uint32_t index, int options)
  20 {
  21     index = index >> root->truncated;
  22     uint32_t lz = index ? ROUNDDOWN(31 - __builtin_clz(index), BTRIE_BITS) : 0;
  23     uint32_t bitmask = ((1 << BTRIE_BITS) - 1) << lz;
  24     uint32_t i = 0;
  25     struct btrie_node* tree = root->btrie_root;
  26
  27     // Time complexity: O(log_2(log_2(N))) where N is the index to lookup
  28     while (lz && tree) {
  29         i = (index & bitmask) >> lz;
  30         struct btrie_node *subtree = 0, *pos, *n;
  31
  32         llist_for_each(pos, n, &tree->children, siblings)
  33         {
  34             if (pos->index == i) {
  35                 subtree = pos;
  36                 break;
  37             }
  38         }
  39
  40         if (!subtree && (options & BTRIE_INSERT)) {
  41             struct btrie_node* new_tree = vzalloc(sizeof(struct btrie_node));
  42             new_tree->index = i;
  43             new_tree->parent = tree;
  44             llist_init_head(&new_tree->children);
  45
  46             llist_append(&tree->children, &new_tree->siblings);
  47             llist_append(&root->btrie_root->nodes, &new_tree->nodes);
  48             tree = new_tree;
  49         } else {
  50             tree = subtree;
  51         }
  52         bitmask = bitmask >> BTRIE_BITS;
  53         lz -= BTRIE_BITS;
  54     }
  55
  56     return tree;
  57 }
  58
  59 void
  60 btrie_init(struct btrie* btrie, uint32_t trunc_bits)
  61 {
  62     btrie->btrie_root = vzalloc(sizeof(struct btrie_node));
  63     llist_init_head(&btrie->btrie_root->nodes);
  64     btrie->truncated = trunc_bits;
  65 }
  66
  67 void*
  68 btrie_get(struct btrie* root, uint32_t index)
  69 {
  70     struct btrie_node* node = __btrie_traversal(root, index, 0);
  71     if (!node) {
  72         return NULL;
  73     }
  74     return node->data;
  75 }
  76
  77 void
  78 btrie_set(struct btrie* root, uint32_t index, void* data)
  79 {
  80     struct btrie_node* node = __btrie_traversal(root, index, BTRIE_INSERT);
  81     node->data = data;
  82 }
  83
  84 void
  85 __btrie_remove(struct btrie_node* node)
  86 {
  87     struct btrie_node* parent = node->parent;
  88     if (parent) {
  89         llist_delete(&node->siblings);
  90         llist_delete(&node->nodes);
  91         vfree(node);
  92         if (llist_empty(&parent->children)) {
  93             __btrie_remove(parent);
  94         }
  95     }
  96 }
  97
  98 void*
  99 btrie_remove(struct btrie* root, uint32_t index)
 100 {
 101     struct btrie_node* node = __btrie_traversal(root, index, 0);
 102     if (!node) {
 103         return 0;
 104     }
 105     void* data = node->data;
 106     __btrie_remove(node);
 107     return data;
 108 }
 109
 110 void
 111 btrie_release(struct btrie* tree)
 112 {
 113     struct btrie_node *pos, *n;
 114     llist_for_each(pos, n, &tree->btrie_root->nodes, nodes)
 115     {
 116         vfree(pos);
 117     }
 118
 119     vfree(tree->btrie_root);
 120 }