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, unsigned long index, int options)
  20 {
  21     unsigned long lz;
  22     unsigned long bitmask;
  23     unsigned long i = 0;
  24     struct btrie_node* tree = root->btrie_root;
  25
  26     lz = index ? ICEIL(msbitl - clzl(index), root->order) : 0;
  27     lz = lz * root->order;
  28     bitmask = ((1 << root->order) - 1) << lz;
  29
  30     // Time complexity: O(log_2(log_2(N))) where N is the index to lookup
  31     while (bitmask && tree) {
  32         i = (index & bitmask) >> lz;
  33         struct btrie_node *subtree = 0, *pos, *n;
  34
  35         llist_for_each(pos, n, &tree->children, siblings)
  36         {
  37             if (pos->index == i) {
  38                 subtree = pos;
  39                 break;
  40             }
  41         }
  42
  43         if (!subtree && (options & BTRIE_INSERT)) {
  44             struct btrie_node* new_tree = vzalloc(sizeof(struct btrie_node));
  45             new_tree->index = i;
  46             new_tree->parent = tree;
  47             llist_init_head(&new_tree->children);
  48
  49             llist_append(&tree->children, &new_tree->siblings);
  50             llist_append(&root->btrie_root->nodes, &new_tree->nodes);
  51             tree = new_tree;
  52         } else {
  53             tree = subtree;
  54         }
  55         bitmask = bitmask >> root->order;
  56         lz -= root->order;
  57     }
  58
  59     return tree;
  60 }
  61
  62 void
  63 btrie_init(struct btrie* btrie, unsigned int order)
  64 {
  65     btrie->btrie_root = vzalloc(sizeof(struct btrie_node));
  66     llist_init_head(&btrie->btrie_root->nodes);
  67     llist_init_head(&btrie->btrie_root->children);
  68     btrie->order = order;
  69 }
  70
  71 void*
  72 btrie_get(struct btrie* root, unsigned long index)
  73 {
  74     struct btrie_node* node = __btrie_traversal(root, index, 0);
  75     if (!node) {
  76         return NULL;
  77     }
  78     return node->data;
  79 }
  80
  81 void
  82 btrie_set(struct btrie* root, unsigned long index, void* data)
  83 {
  84     struct btrie_node* node = __btrie_traversal(root, index, BTRIE_INSERT);
  85     node->data = data;
  86 }
  87
  88 void
  89 __btrie_rm_recursive(struct btrie_node* node)
  90 {
  91     struct btrie_node* parent = node->parent;
  92
  93     if (parent) {
  94         llist_delete(&node->siblings);
  95         llist_delete(&node->nodes);
  96         vfree(node);
  97         if (llist_empty(&parent->children) && !parent->data) {
  98             __btrie_rm_recursive(parent);
  99         }
 100     }
 101 }
 102
 103 void*
 104 btrie_remove(struct btrie* root, unsigned long index)
 105 {
 106     struct btrie_node* node = __btrie_traversal(root, index, 0);
 107     if (!node) {
 108         return 0;
 109     }
 110     void* data = node->data;
 111     if (!llist_empty(&node->children)) {
 112         node->data = NULL;
 113     } else {
 114         __btrie_rm_recursive(node);
 115     }
 116     return data;
 117 }
 118
 119 void
 120 btrie_release(struct btrie* tree)
 121 {
 122     struct btrie_node *pos, *n;
 123     llist_for_each(pos, n, &tree->btrie_root->nodes, nodes)
 124     {
 125         vfree(pos);
 126     }
 127
 128     vfree(tree->btrie_root);
 129 }