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 <stddef.h>
  15 #include <stdint.h>
  16 #include <lowlevel/bitops.h>
  17
  18 namespace Util {
  19         // Each leaf node in the tree will contain 2**radix_bits nodes.
  20         // Each directory node will contain 2**dir_bits pointers.
  21         // Key must be an unsigned integer.
  22
  23         template <typename T, typename Key, int radix_bits,
  24                   int dir_bits = radix_bits>
  25         class RadixTree {
  26                 enum {
  27                         node_size = 1 << radix_bits,
  28                         dir_size = 1 << dir_bits,
  29                         key_bits = sizeof(Key) * 8,
  30                 };
  31
  32                 struct Node {
  33                         const int shift;
  34
  35                 protected:
  36                         Node(int SHIFT) : shift(SHIFT)
  37                         {
  38                         }
  39                 };
  40
  41                 struct LeafNode : public Node {
  42                         T nodes[node_size];
  43
  44                         LeafNode() : Node(0)
  45                         {
  46                                 memset(nodes, 0, sizeof(nodes));
  47                         }
  48                 };
  49
  50                 struct DirNode : public Node {
  51                         Node *ptrs[dir_size];
  52
  53                         DirNode(int shift) : Node(shift)
  54                         {
  55                                 memset(ptrs, 0, sizeof(ptrs));
  56                         }
  57                 };
  58
  59                 Node *toplevel;
  60
  61                 static int key_to_dir_offset(Key key, int shift)
  62                 {
  63                         return (key >> shift) & (dir_size - 1);
  64                 }
  65
  66                 static unsigned int key_to_offset(Key key)
  67                 {
  68                         return key & (node_size - 1);
  69                 }
  70
  71         public:
  72                 // The tree depth is contained in each node, and nodes are never
  73                 // removed from the tree, so lockless lookups are possible.
  74                 // Synchronization is required when add == true.
  75
  76                 T *lookup(Key key, bool add = false)
  77                 {
  78                         Node *node = toplevel;
  79                         int next_shift = node->shift ? node->shift + dir_bits : radix_bits;
  80
  81                         while (unlikely(next_shift < key_bits && (key >> next_shift)))
  82                         {
  83                                 if (!add)
  84                                         return NULL;
  85
  86                                 DirNode *dn = new DirNode(next_shift);
  87                                 dn->ptrs[0] = node;
  88                                 toplevel = dn;
  89                                 next_shift += dir_bits;
  90                         }
  91
  92                         node = toplevel;
  93
  94                         while (node->shift > 0) {
  95                                 int off = key_to_dir_offset(key, node->shift);
  96                                 DirNode *dn = static_cast<DirNode *>(node);
  97                                 node = dn->ptrs[off];
  98
  99                                 if (!node) {
 100                                         if (!add)
 101                                                 return NULL;
 102
 103                                         if (dn->shift == radix_bits)
 104                                                 node = new LeafNode;
 105                                         else
 106                                                 node = new DirNode(dn->shift == radix_bits ?
 107                                                                    0 : dn->shift - dir_bits);
 108
 109                                         dn->ptrs[off] = node;
 110                                 }
 111                         }
 112
 113                         LeafNode *ln = static_cast<LeafNode *>(node);
 114                         return &ln->nodes[key_to_offset(key)];
 115                 }
 116
 117                 RadixTree()
 118                 {
 119                         // OPT: start with a NULL toplevel.  Besides avoiding wasting
 120                         // memory for empty trees, it avoids the need for a stack of
 121                         // trees down the ptr[0] path if there are no small keys.
 122
 123                         toplevel = new LeafNode;
 124                 }
 125         };
 126
 127 }
 128
 129 #endif