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