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, u32_t index, int options)
  20 {
  21     index = index >> root->truncated;
  22     u32_t lz = index ? ROUNDDOWN(31 - clz(index), BTRIE_BITS) : 0;
  23     u32_t bitmask = ((1 << BTRIE_BITS) - 1) << lz;
  24     u32_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 (bitmask && 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, u32_t trunc_bits)
  61 {
  62     btrie->btrie_root = vzalloc(sizeof(struct btrie_node));
  63     llist_init_head(&btrie->btrie_root->nodes);
  64     llist_init_head(&btrie->btrie_root->children);
  65     btrie->truncated = trunc_bits;
  66 }
  67
  68 void*
  69 btrie_get(struct btrie* root, u32_t index)
  70 {
  71     struct btrie_node* node = __btrie_traversal(root, index, 0);
  72     if (!node) {
  73         return NULL;
  74     }
  75     return node->data;
  76 }
  77
  78 void
  79 btrie_set(struct btrie* root, u32_t index, void* data)
  80 {
  81     struct btrie_node* node = __btrie_traversal(root, index, BTRIE_INSERT);
  82     node->data = data;
  83 }
  84
  85 void
  86 __btrie_rm_recursive(struct btrie_node* node)
  87 {
  88     struct btrie_node* parent = node->parent;
  89
  90     if (parent) {
  91         llist_delete(&node->siblings);
  92         llist_delete(&node->nodes);
  93         vfree(node);
  94         if (llist_empty(&parent->children) && !parent->data) {
  95             __btrie_rm_recursive(parent);
  96         }
  97     }
  98 }
  99
 100 void*
 101 btrie_remove(struct btrie* root, u32_t index)
 102 {
 103     struct btrie_node* node = __btrie_traversal(root, index, 0);
 104     if (!node) {
 105         return 0;
 106     }
 107     void* data = node->data;
 108     if (!llist_empty(&node->children)) {
 109         node->data = NULL;
 110     } else {
 111         __btrie_rm_recursive(node);
 112     }
 113     return data;
 114 }
 115
 116 void
 117 btrie_release(struct btrie* tree)
 118 {
 119     struct btrie_node *pos, *n;
 120     llist_for_each(pos, n, &tree->btrie_root->nodes, nodes)
 121     {
 122         vfree(pos);
 123     }
 124
 125     vfree(tree->btrie_root);
 126 }