include/llvm/ADT/Trie.h

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