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