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