include/llvm/ADT/Trie.h

   1 //===- llvm/ADT/Trie.h ---- Generic trie structure --------------*- C++ -*-===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file was developed by Anton Korobeynikov and is distributed under
   6 // the University of Illinois Open Source License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This class defines a generic trie structure. The trie structure
  11 // is immutable after creation, but the payload contained within it is not.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #ifndef LLVM_ADT_TRIE_H
  16 #define LLVM_ADT_TRIE_H
  17
  18 #include <map>
  19 #include <vector>
  20
  21 namespace llvm {
  22
  23 // FIXME:
  24 // - Labels are usually small, maybe it's better to use SmallString
  25 // - Something efficient for child storage
  26 // - Should we use char* during construction?
  27 // - Should we templatize Empty with traits-like interface?
  28 // - GraphTraits interface
  29
  30 template<class Payload>
  31 class Trie {
  32   class Node {
  33     friend class Trie;
  34
  35     typedef enum {
  36       Same           = -3,
  37       StringIsPrefix = -2,
  38       LabelIsPrefix  = -1,
  39       DontMatch      = 0,
  40       HaveCommonPart
  41     } QueryResult;
  42
  43     std::string Label;
  44     Payload Data;
  45     std::map<char, Node*> Children;
  46   public:
  47     inline explicit Node(const Payload& data, const std::string& label = ""):
  48         Label(label), Data(data) { }
  49
  50     inline Node(const Node& n) {
  51       Data = n.Data;
  52       Children = n.Children;
  53       Label = n.Label;
  54     }
  55     inline Node& operator=(const Node& n) {
  56       if (&n != this) {
  57         Data = n.Data;
  58         Children = n.Children;
  59         Label = n.Label;
  60       }
  61
  62       return *this;
  63     }
  64
  65     inline bool isLeaf() const { return Children.empty(); }
  66
  67     inline const Payload& getData() const { return Data; }
  68     inline void setData(const Payload& data) { Data = data; }
  69
  70     inline void setLabel(const std::string& label) { Label = label; }
  71     inline const std::string& getLabel() const { return Label; }
  72
  73     inline bool addEdge(Node* N) {
  74       const std::string& Label = N->getLabel();
  75       return Children.insert(std::make_pair(Label[0], N)).second;
  76     }
  77
  78     QueryResult query(const std::string& s) const {
  79       unsigned i, l;
  80       unsigned l1 = s.length();
  81       unsigned l2 = Label.length();
  82
  83       // Find the length of common part
  84       l = std::min(l1, l2);
  85       i = 0;
  86       while ((i < l) && (s[i] == Label[i]))
  87         ++i;
  88
  89       if (i == l) { // One is prefix of another, find who is who
  90         if (l1 == l2)
  91           return Same;
  92         else if (i == l1)
  93           return StringIsPrefix;
  94         else
  95           return LabelIsPrefix;
  96       } else // s and Label have common (possible empty) part, return its length
  97         return (QueryResult)i;
  98     }
  99   };
 100
 101   std::vector<Node*> Nodes;
 102   Payload Empty;
 103
 104   inline Node* addNode(const Payload& data, const std::string label = "") {
 105     Node* N = new Node(data, label);
 106     Nodes.push_back(N);
 107     return N;
 108   }
 109
 110   inline Node* splitEdge(Node* N, char Id, size_t index) {
 111     assert(N->Children.count(Id) && "Node doesn't exist");
 112
 113     Node* eNode = N->Children[Id];
 114
 115     const std::string &l = eNode->Label;
 116     assert(index > 0 && index < l.length() && "Trying to split too far!");
 117     std::string l1 = l.substr(0, index);
 118     std::string l2 = l.substr(index);
 119
 120     eNode->Label = l2;
 121
 122     Node* nNode = addNode(Empty, l1);
 123     nNode->addEdge(eNode);
 124
 125     N->Children[Id] = nNode;
 126
 127     return nNode;
 128   }
 129
 130 public:
 131   inline explicit Trie(const Payload& empty):Empty(empty) {
 132     addNode(Empty);
 133   }
 134   inline ~Trie() {
 135     for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
 136       delete Nodes[i];
 137   }
 138
 139   inline Node* getRoot() const { return Nodes[0]; }
 140
 141   bool addString(const std::string& s, const Payload& data) {
 142     Node* cNode = getRoot();
 143     Node* tNode = NULL;
 144     std::string s1(s);
 145
 146     while (tNode == NULL) {
 147       char Id = s1[0];
 148       if (cNode->Children.count(Id)) {
 149         Node* nNode = cNode->Children[Id];
 150         typename Node::QueryResult r = nNode->query(s1);
 151
 152         switch (r) {
 153         case Node::Same:
 154         case Node::StringIsPrefix:
 155           // Currently we don't allow to have two strings in the trie one
 156           // being a prefix of another. This should be fixed.
 157           assert(0 && "FIXME!");
 158           return false;
 159         case Node::DontMatch:
 160           assert(0 && "Impossible!");
 161           return false;
 162         case Node::LabelIsPrefix:
 163           s1 = s1.substr(nNode->getLabel().length());
 164           cNode = nNode;
 165           break;
 166         default:
 167          nNode = splitEdge(cNode, Id, r);
 168          tNode = addNode(data, s1.substr(r));
 169          nNode->addEdge(tNode);
 170         }
 171       } else {
 172         tNode = addNode(data, s1);
 173         cNode->addEdge(tNode);
 174       }
 175     }
 176
 177     return true;
 178   }
 179
 180   const Payload& lookup(const std::string& s) const {
 181     Node* cNode = getRoot();
 182     Node* tNode = NULL;
 183     std::string s1(s);
 184
 185     while (tNode == NULL) {
 186       if (cNode->Children.count(s1[0])) {
 187         Node* nNode = cNode->Children[s1[0]];
 188         typename Node::QueryResult r = nNode->query(s1);
 189
 190         switch (r) {
 191         case Node::Same:
 192           tNode = nNode;
 193           break;
 194         case Node::StringIsPrefix:
 195           return Empty;
 196         case Node::DontMatch:
 197           assert(0 && "Impossible!");
 198           return Empty;
 199         case Node::LabelIsPrefix:
 200           s1 = s1.substr(nNode->getLabel().length());
 201           cNode = nNode;
 202           break;
 203         default:
 204           return Empty;
 205         }
 206       } else
 207         return Empty;
 208     }
 209
 210     return tNode->getData();
 211   }
 212
 213 };
 214
 215 }
 216
 217 #endif // LLVM_ADT_TRIE_H