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 "llvm/ADT/GraphTraits.h"
  19 #include "llvm/Support/DOTGraphTraits.h"
  20
  21 #include <vector>
  22
  23 namespace llvm {
  24
  25 // FIXME:
  26 // - Labels are usually small, maybe it's better to use SmallString
  27 // - Should we use char* during construction?
  28 // - Should we templatize Empty with traits-like interface?
  29
  30 template<class Payload>
  31 class Trie {
  32   friend class GraphTraits<Trie<Payload> >;
  33   friend class DOTGraphTraits<Trie<Payload> >;
  34 public:
  35   class Node {
  36     friend class Trie;
  37     friend class GraphTraits<Trie<Payload> >;
  38
  39     typedef enum {
  40       Same           = -3,
  41       StringIsPrefix = -2,
  42       LabelIsPrefix  = -1,
  43       DontMatch      = 0,
  44       HaveCommonPart
  45     } QueryResult;
  46     typedef std::vector<Node*> NodeVector;
  47     typedef typename std::vector<Node*>::iterator NodeVectorIter;
  48
  49     struct NodeCmp {
  50       bool operator() (Node* N1, Node* N2) {
  51         return (N1->Label[0] < N2->Label[0]);
  52       }
  53       bool operator() (Node* N, char Id) {
  54         return (N->Label[0] < Id);
  55       }
  56     };
  57
  58     std::string Label;
  59     Payload Data;
  60     NodeVector Children;
  61
  62     // Do not implement
  63     Node(const Node&);
  64     Node& operator=(const Node&);
  65
  66     inline void addEdge(Node* N) {
  67       if (Children.empty())
  68         Children.push_back(N);
  69       else {
  70         NodeVectorIter I = std::lower_bound(Children.begin(), Children.end(),
  71                                             N, NodeCmp());
  72         // FIXME: no dups are allowed
  73         Children.insert(I, N);
  74       }
  75     }
  76
  77     inline void setEdge(Node* N) {
  78       char Id = N->Label[0];
  79       NodeVectorIter I = std::lower_bound(Children.begin(), Children.end(),
  80                                           Id, NodeCmp());
  81       assert(I != Children.end() && "Node does not exists!");
  82       *I = N;
  83     }
  84
  85     QueryResult query(const std::string& s) const {
  86       unsigned i, l;
  87       unsigned l1 = s.length();
  88       unsigned l2 = Label.length();
  89
  90       // Find the length of common part
  91       l = std::min(l1, l2);
  92       i = 0;
  93       while ((i < l) && (s[i] == Label[i]))
  94         ++i;
  95
  96       if (i == l) { // One is prefix of another, find who is who
  97         if (l1 == l2)
  98           return Same;
  99         else if (i == l1)
 100           return StringIsPrefix;
 101         else
 102           return LabelIsPrefix;
 103       } else // s and Label have common (possible empty) part, return its length
 104         return (QueryResult)i;
 105     }
 106
 107   public:
 108     inline explicit Node(const Payload& data, const std::string& label = ""):
 109         Label(label), Data(data) { }
 110
 111     inline const Payload& data() const { return Data; }
 112     inline void setData(const Payload& data) { Data = data; }
 113
 114     inline const std::string& label() const { return Label; }
 115
 116 #if 0
 117     inline void dump() {
 118       std::cerr << "Node: " << this << "\n"
 119                 << "Label: " << Label << "\n"
 120                 << "Children:\n";
 121
 122       for (NodeVectorIter I = Children.begin(), E = Children.end(); I != E; ++I)
 123         std::cerr << (*I)->Label << "\n";
 124     }
 125 #endif
 126
 127     inline Node* getEdge(char Id) {
 128       Node* fNode = NULL;
 129       NodeVectorIter I = std::lower_bound(Children.begin(), Children.end(),
 130                                           Id, NodeCmp());
 131       if (I != Children.end() && (*I)->Label[0] == Id)
 132         fNode = *I;
 133
 134       return fNode;
 135     }
 136   };
 137
 138 private:
 139   std::vector<Node*> Nodes;
 140   Payload Empty;
 141
 142   inline Node* addNode(const Payload& data, const std::string label = "") {
 143     Node* N = new Node(data, label);
 144     Nodes.push_back(N);
 145     return N;
 146   }
 147
 148   inline Node* splitEdge(Node* N, char Id, size_t index) {
 149     Node* eNode = N->getEdge(Id);
 150     assert(eNode && "Node doesn't exist");
 151
 152     const std::string &l = eNode->Label;
 153     assert(index > 0 && index < l.length() && "Trying to split too far!");
 154     std::string l1 = l.substr(0, index);
 155     std::string l2 = l.substr(index);
 156
 157     Node* nNode = addNode(Empty, l1);
 158     N->setEdge(nNode);
 159
 160     eNode->Label = l2;
 161     nNode->addEdge(eNode);
 162
 163     return nNode;
 164   }
 165
 166   // Do not implement
 167   Trie(const Trie&);
 168   Trie& operator=(const Trie&);
 169
 170 public:
 171   inline explicit Trie(const Payload& empty):Empty(empty) {
 172     addNode(Empty);
 173   }
 174   inline ~Trie() {
 175     for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
 176       delete Nodes[i];
 177   }
 178
 179   inline Node* getRoot() const { return Nodes[0]; }
 180
 181   bool addString(const std::string& s, const Payload& data);
 182   const Payload& lookup(const std::string& s) const;
 183
 184 };
 185
 186 // Define this out-of-line to dissuade the C++ compiler from inlining it.
 187 template<class Payload>
 188 bool Trie<Payload>::addString(const std::string& s, const Payload& data) {
 189   Node* cNode = getRoot();
 190   Node* tNode = NULL;
 191   std::string s1(s);
 192
 193   while (tNode == NULL) {
 194     char Id = s1[0];
 195     if (Node* nNode = cNode->getEdge(Id)) {
 196       typename Node::QueryResult r = nNode->query(s1);
 197
 198       switch (r) {
 199       case Node::Same:
 200       case Node::StringIsPrefix:
 201         // Currently we don't allow to have two strings in the trie one
 202         // being a prefix of another. This should be fixed.
 203         assert(0 && "FIXME!");
 204         return false;
 205       case Node::DontMatch:
 206         assert(0 && "Impossible!");
 207         return false;
 208       case Node::LabelIsPrefix:
 209         s1 = s1.substr(nNode->label().length());
 210         cNode = nNode;
 211         break;
 212       default:
 213         nNode = splitEdge(cNode, Id, r);
 214         tNode = addNode(data, s1.substr(r));
 215         nNode->addEdge(tNode);
 216       }
 217     } else {
 218       tNode = addNode(data, s1);
 219       cNode->addEdge(tNode);
 220     }
 221   }
 222
 223   return true;
 224 }
 225
 226 template<class Payload>
 227 const Payload& Trie<Payload>::lookup(const std::string& s) const {
 228   Node* cNode = getRoot();
 229   Node* tNode = NULL;
 230   std::string s1(s);
 231
 232   while (tNode == NULL) {
 233     char Id = s1[0];
 234     if (Node* nNode = cNode->getEdge(Id)) {
 235       typename Node::QueryResult r = nNode->query(s1);
 236
 237       switch (r) {
 238       case Node::Same:
 239         tNode = nNode;
 240         break;
 241       case Node::StringIsPrefix:
 242         return Empty;
 243       case Node::DontMatch:
 244         assert(0 && "Impossible!");
 245         return Empty;
 246       case Node::LabelIsPrefix:
 247         s1 = s1.substr(nNode->label().length());
 248         cNode = nNode;
 249         break;
 250       default:
 251         return Empty;
 252       }
 253     } else
 254       return Empty;
 255   }
 256
 257   return tNode->data();
 258 }
 259
 260 template<class Payload>
 261 struct GraphTraits<Trie<Payload> > {
 262   typedef typename Trie<Payload>::Node NodeType;
 263   typedef typename std::vector<NodeType*>::iterator ChildIteratorType;
 264
 265   static inline NodeType *getEntryNode(const Trie<Payload>& T) {
 266     return T.getRoot();
 267   }
 268
 269   static inline ChildIteratorType child_begin(NodeType *N) {
 270     return N->Children.begin();
 271   }
 272   static inline ChildIteratorType child_end(NodeType *N) {
 273     return N->Children.end();
 274   }
 275
 276   typedef typename std::vector<NodeType*>::const_iterator nodes_iterator;
 277
 278   static inline nodes_iterator nodes_begin(const Trie<Payload>& G) {
 279     return G.Nodes.begin();
 280   }
 281   static inline nodes_iterator nodes_end(const Trie<Payload>& G) {
 282     return G.Nodes.end();
 283   }
 284
 285 };
 286
 287 template<class Payload>
 288 struct DOTGraphTraits<Trie<Payload> > : public DefaultDOTGraphTraits {
 289   typedef typename Trie<Payload>::Node NodeType;
 290   typedef typename GraphTraits<Trie<Payload> >::ChildIteratorType EdgeIter;
 291
 292   static std::string getGraphName(const Trie<Payload>& T) {
 293     return "Trie";
 294   }
 295
 296   static std::string getNodeLabel(NodeType* Node, const Trie<Payload>& T) {
 297     if (T.getRoot() == Node)
 298       return "<Root>";
 299     else
 300       return Node->label();
 301   }
 302
 303   static std::string getEdgeSourceLabel(NodeType* Node, EdgeIter I) {
 304     NodeType* N = *I;
 305     return N->label().substr(0, 1);
 306   }
 307
 308   static std::string getNodeAttributes(const NodeType* Node,
 309                                        const Trie<Payload>& T) {
 310     if (Node->data() != T.Empty)
 311       return "color=blue";
 312
 313     return "";
 314   }
 315
 316 };
 317
 318 } // end of llvm namespace
 319
 320 #endif // LLVM_ADT_TRIE_H