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 enum btrie_traverse_mode {
  17     BTRIE_FIND,
  18     BTRIE_FIND_PARENT,
  19     BTRIE_INSERT
  20 };
  21 struct btrie_locator
  22 {
  23     enum btrie_traverse_mode mode;
  24     int shifts;
  25     unsigned long index;
  26 };
  27
  28 static inline void
  29 __btrie_locator(struct btrie_locator* loc,
  30                 unsigned long index, enum btrie_traverse_mode mode)
  31 {
  32     loc->mode = mode;
  33     loc->index = index;
  34 }
  35
  36 static inline void
  37 __btrie_addchild(struct btrie* root,
  38                  struct btrie_node* parent, struct btrie_node* child)
  39 {
  40     if (unlikely(!parent->children)) {
  41         parent->children = vcalloc(sizeof(parent), 1 << root->order);
  42     }
  43
  44     parent->children_cnt++;
  45     parent->children[child->index] = child;
  46 }
  47
  48 static inline struct btrie_node*
  49 __btrie_get_child(struct btrie_node* node, int index)
  50 {
  51     if (unlikely(!node->children)) {
  52         return NULL;
  53     }
  54
  55     return node->children[index];
  56 }
  57
  58 static inline void
  59 __btrie_free_node(struct btrie_node* node)
  60 {
  61     vfree_safe(node->children);
  62     vfree(node);
  63 }
  64
  65 static inline bool
  66 __full_node(struct btrie* root, struct btrie_node* node)
  67 {
  68     return node->children_cnt == 1 << root->order;
  69 }
  70
  71 static inline struct btrie_node*
  72 __btrie_create(struct btrie* root, struct btrie_node* parent, int index)
  73 {
  74     struct btrie_node* node;
  75
  76     node = vzalloc(sizeof(struct btrie_node));
  77     node->index = index;
  78     node->parent = parent;
  79
  80     if (likely(parent)) {
  81         llist_append(&root->btrie_root->nodes, &node->nodes);
  82         __btrie_addchild(root, parent, node);
  83     }
  84
  85     return node;
  86 }
  87
  88 static struct btrie_node*
  89 __btrie_traversal(struct btrie* root, struct btrie_locator* locator)
  90 {
  91     unsigned long lz, index;
  92     unsigned long bitmask;
  93     unsigned long i = 0;
  94     struct btrie_node *subtree, *tree = root->btrie_root;
  95
  96     index = locator->index;
  97     lz = index ? (msbitl - clzl(index)) / root->order : 0;
  98     lz = lz * root->order;
  99     bitmask = ((1 << root->order) - 1) << lz;
 100
 101     // Time complexity: O(log_2(log_2(N))) where N is the index to lookup
 102     while (bitmask && tree) {
 103         i = (index & bitmask) >> lz;
 104         subtree = __btrie_get_child(tree, i);
 105
 106         if (!subtree && (locator->mode != BTRIE_FIND))
 107         {
 108             tree = __btrie_create(root, tree, i);
 109         }
 110         else {
 111             tree = subtree;
 112         }
 113
 114         if (lz == root->order && locator->mode == BTRIE_FIND_PARENT)
 115         {
 116             break;
 117         }
 118
 119         bitmask = bitmask >> root->order;
 120         lz -= root->order;
 121     }
 122
 123     locator->shifts = lz;
 124     return tree;
 125 }
 126
 127 static unsigned long
 128 __btrie_alloc_slot(struct btrie* root, struct btrie_node **slot,
 129                    unsigned long start, unsigned long end)
 130 {
 131     unsigned int index, od;
 132     unsigned long result, mask, result_next;
 133     struct btrie_node *tree, *subtree;
 134     struct btrie_locator locator;
 135
 136     result = 0;
 137     od   = root->order;
 138     mask = (1 << od) - 1;
 139     tree = root->btrie_root;
 140
 141     if (start) {
 142         __btrie_locator(&locator, start, BTRIE_FIND_PARENT);
 143         tree = __btrie_traversal(root, &locator);
 144         result = start;
 145     }
 146
 147     while (tree && result < end)
 148     {
 149         index = result & mask;
 150         result_next = result << od;
 151
 152         subtree = __btrie_get_child(tree, index);
 153
 154         if (subtree) {
 155             if (result >= start && !subtree->data) {
 156                 tree = subtree;
 157                 goto done;
 158             }
 159         }
 160         else {
 161             subtree = __btrie_create(root, tree, index);
 162
 163             if (result >= start)
 164                 goto done;
 165
 166             goto descend;
 167         }
 168
 169         if (!__full_node(root, tree) && index != mask) {
 170             result++;
 171             continue;
 172         }
 173
 174 descend:
 175         if (result_next >= end) {
 176             tree = tree->parent;
 177             result >>= od;
 178             result++;
 179             continue;
 180         }
 181
 182         tree   = subtree;
 183         result = result_next;
 184     }
 185
 186     *slot = NULL;
 187     return -1;
 188
 189 done:
 190     *slot = subtree;
 191     return result;
 192 }
 193
 194 void
 195 btrie_init(struct btrie* btrie, unsigned int order)
 196 {
 197     btrie->btrie_root = __btrie_create(btrie, NULL, 0);
 198     btrie->order = order;
 199
 200     llist_init_head(&btrie->btrie_root->nodes);
 201 }
 202
 203 void*
 204 btrie_get(struct btrie* root, unsigned long index)
 205 {
 206     struct btrie_node* node;
 207     struct btrie_locator locator;
 208
 209     __btrie_locator(&locator, index, BTRIE_FIND);
 210     node = __btrie_traversal(root, &locator);
 211     if (!node) {
 212         return NULL;
 213     }
 214
 215     return node->data;
 216 }
 217
 218 void
 219 btrie_set(struct btrie* root, unsigned long index, void* data)
 220 {
 221     struct btrie_node* node;
 222     struct btrie_locator locator;
 223
 224     __btrie_locator(&locator, index, BTRIE_INSERT);
 225     node = __btrie_traversal(root, &locator);
 226     node->data = data;
 227 }
 228
 229 void
 230 __btrie_rm_recursive(struct btrie_node* node)
 231 {
 232     struct btrie_node* parent = node->parent;
 233
 234     if (!parent) {
 235         return;
 236     }
 237
 238     parent->children[node->index] = NULL;
 239     parent->children_cnt--;
 240     __btrie_free_node(node);
 241
 242     if (!parent->children_cnt && !parent->data) {
 243         __btrie_rm_recursive(parent);
 244     }
 245 }
 246
 247 unsigned long
 248 btrie_map(struct btrie* root,
 249           unsigned long start, unsigned long end, void* data)
 250 {
 251     struct btrie_node* node = NULL;
 252     unsigned long alloced;
 253
 254     alloced = __btrie_alloc_slot(root, &node, start, end);
 255
 256     if (!node)
 257         return -1;
 258
 259     node->data = data;
 260     return alloced;
 261 }
 262
 263 void*
 264 btrie_remove(struct btrie* root, unsigned long index)
 265 {
 266     void* data;
 267     struct btrie_node* node;
 268     struct btrie_locator locator;
 269
 270     __btrie_locator(&locator, index, BTRIE_FIND);
 271     node = __btrie_traversal(root, &locator);
 272     if (!node) {
 273         return 0;
 274     }
 275
 276     data = node->data;
 277     if (!node->children_cnt) {
 278         node->data = NULL;
 279     } else {
 280         __btrie_rm_recursive(node);
 281     }
 282
 283     return data;
 284 }
 285
 286 void
 287 btrie_release(struct btrie* tree)
 288 {
 289     struct btrie_node *pos, *n;
 290     llist_for_each(pos, n, &tree->btrie_root->nodes, nodes)
 291     {
 292         __btrie_free_node(pos);
 293     }
 294
 295     __btrie_free_node(tree->btrie_root);
 296 }