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
  24         template <typename T, int radix_bits, int dir_bits = radix_bits>
  25         class RadixTree {
  26                 typedef unsigned int Key;
  27
  28                 enum {
  29                         node_size = 1 << radix_bits,
  30                         dir_size = 1 << dir_bits,
  31                         long_bits = sizeof(unsigned long) * 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 handles[node_size];
  45
  46                         LeafNode() : Node(0)
  47                         {
  48                                 shift = 0;
  49                                 memset(handles, 0, sizeof(handles));
  50                         }
  51                 };
  52
  53                 struct DirNode : public Node {
  54                         Node *ptrs[node_size];
  55
  56                         DirNode(int shift) : Node(shift)
  57                         {
  58                                 memset(ptrs, 0, sizeof(ptrs));
  59                         }
  60                 };
  61
  62                 Node *toplevel;
  63
  64                 static int key_to_dir_offset(Key key, int shift)
  65                 {
  66                         return (key >> shift) & (dir_size - 1);
  67                 }
  68
  69                 static uint key_to_offset(Key key)
  70                 {
  71                         return key & (node_size - 1);
  72                 }
  73
  74         public:
  75                 // The tree depth is contained in each node, and nodes are never
  76                 // removed from the tree, so lockless lookups are possible.
  77                 // Synchronization is required when add == true.
  78
  79                 T *lookup(Key key, bool add = false)
  80                 {
  81                         int next_shift = toplevel->shift == 0 ? radix_bits : dir_bits;
  82
  83                         while (unlikely(key_to_dir_offset(key, toplevel->shift +
  84                                                                next_shift) != 0))
  85                         {
  86                                 if (!add)
  87                                         return NULL;
  88
  89                                 Node *new_node = new DirNode(toplevel->shift + next_shift);
  90                                 new_node->ptrs[0] = toplevel;
  91                                 toplevel = new_node;
  92                                 next_shift = dir_bits;
  93                         }
  94
  95                         Node *node = toplevel;
  96
  97                         while (node->shift > 0) {
  98                                 assert(node->shift >= radix_bits);
  99
 100                                 int off = key_to_dir_offset(key, node->shift);
 101                                 NodeBase *new_node = node[off];
 102
 103                                 if (!new_node) {
 104                                         if (!add)
 105                                                 return NULL;
 106
 107                                         if (node->shift == radix_bits)
 108                                                 new_node = new LeafNode;
 109                                         else
 110                                                 new_node = new DirNode(node->shift - dir_shift);
 111
 112                                         node[off] = new_node;
 113                                 }
 114
 115                                 node = new_node;
 116                         }
 117
 118                         assert(node->shift == 0);
 119                         return node[key_to_offset(key)];
 120                 }
 121
 122                 RadixTree()
 123                 {
 124                         toplevel = new LeafNode;
 125                 }
 126         };
 127 }
 128
 129 #endif