include/c++/util/radix.h

   1 // Generic radix trees, which supports mapping integer keys to objects.
   2 //
   3 // Written by Scott Wood <scott@buserror.net>.
   4 //
   5 // This software is provided 'as-is', without any express or implied warranty.
   6 // In no event will the authors or contributors be held liable for any damages
   7 // arising from the use of this software.
   8 //
   9 // This software is in the public domain.
  10
  11 #ifndef _UTIL_RADIX_H
  12 #define _UTIL_RADIX_H
  13
  14 #include <assert.h>
  15 #include <stddef.h>
  16 #include <stdint.h>
  17
  18 #include <lowlevel/bitops.h>
  19
  20 namespace Util {
  21         // Each leaf node in the tree will contain 2**radix_bits handles.
  22         // Each directory node will contain 2**dir_bits pointers.
  23         template <typename T, int radix_bits, int dir_bits = radix_bits>
  24         class RadixTree {
  25                 typedef unsigned int Key;
  26
  27                 enum {
  28                         node_size = 1 << radix_bits,
  29                         dir_size = 1 << dir_bits,
  30                         long_bits = sizeof(unsigned long) * 8,
  31                 };
  32
  33                 struct Node {
  34                         const int shift;
  35
  36                 protected:
  37                         Node(int SHIFT) : shift(SHIFT)
  38                         {
  39                         }
  40                 };
  41
  42                 struct LeafNode : public Node {
  43                         T handles[node_size];
  44
  45                         LeafNode() : Node(0)
  46                         {
  47                                 shift = 0;
  48                                 memset(handles, 0, sizeof(handles));
  49                         }
  50                 };
  51
  52                 struct DirNode : public Node {
  53                         Node *ptrs[node_size];
  54
  55                         DirNode(int shift) : Node(shift)
  56                         {
  57                                 memset(ptrs, 0, sizeof(ptrs));
  58                         }
  59                 };
  60
  61                 Node *toplevel;
  62
  63                 static int key_to_dir_offset(Key key, int shift)
  64                 {
  65                         return (key >> shift) & (dir_size - 1);
  66                 }
  67
  68                 static uint key_to_offset(Key key)
  69                 {
  70                         return key & (node_size - 1);
  71                 }
  72
  73         public:
  74                 T *lookup(Key key, bool add = false)
  75                 {
  76                         int next_shift = toplevel->shift == 0 ? radix_bits : dir_bits;
  77
  78                         while (unlikely(key_to_dir_offset(key, toplevel->shift +
  79                                                                next_shift) != 0))
  80                         {
  81                                 if (!add)
  82                                         return NULL;
  83
  84                                 Node *new_node = new DirNode(toplevel->shift + next_shift);
  85                                 new_node->ptrs[0] = toplevel;
  86                                 toplevel = new_node;
  87                                 next_shift = dir_bits;
  88                         }
  89
  90                         Node *node = toplevel;
  91
  92                         while (node->shift >= radix_bits) {
  93                                 int off = key_to_dir_offset(key, node->shift);
  94                                 NodeBase *new_node = node[off];
  95
  96                                 shift == 0 ? new LeafNode
  97                                                                 : new DirNode(shift);
  98                                 void *new_node = static_cast<void **>(node)[off];
  99
 100                                 if (!new_node) {
 101                                         if (!add)
 102                                                 return NULL;
 103
 104                                         new_node = Mem::alloc_pages(1);
 105                                         memset(new_node, 0, Arch::page_size);
 106                                         static_cast<void **>(node)[off] = new_node;
 107                                 }
 108
 109                                 shift -= dir_shift;
 110                                 node = new_node;
 111                         }
 112
 113                         assert(shift == 0);
 114                         return static_cast<T *>(node) + key_to_offset(key, 0);
 115                 }
 116
 117                 RadixTree()
 118                 {
 119                         toplevel = new LeafNode;
 120                         memset(toplevel, 0, Arch::page_size);
 121                         depth = 0;
 122                 }
 123         };
 124 }
 125
 126 #endif