aboutsummaryrefslogtreecommitdiffstats
path: root/include/llvm/ADT/Trie.h
blob: 30a9f56e1b165c8866a811d416a5bf9a45aabbfb (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
//===- llvm/ADT/Trie.h ---- Generic trie structure --------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file was developed by Anton Korobeynikov and is distributed under
// the University of Illinois Open Source License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This class defines a generic trie structure. The trie structure
// is immutable after creation, but the payload contained within it is not.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_ADT_TRIE_H
#define LLVM_ADT_TRIE_H

#include <map>
#include <vector>

namespace llvm {

// FIXME:
// - Labels are usually small, maybe it's better to use SmallString
// - Something efficient for child storage
// - Should we use char* during construction?
// - GraphTraits interface
// - Eliminate Edge class, which is ok for debugging, but not for end code

template<class Payload>
class Trie {
  class Edge;
  class Node;

  class Edge {
    std::string Label;
    Node *Parent, *Child;

  public:
    typedef enum {
      Same           = -3,
      StringIsPrefix = -2,
      LabelIsPrefix  = -1,
      DontMatch      = 0,
      HaveCommonPart
    } QueryResult;

    inline explicit Edge(std::string label = "",
                         Node* parent = NULL, Node* child = NULL):
      Label(label), Parent(parent), Child(child) { }

    inline void setParent(Node* parent) { Parent = parent; }
    inline Node* getParent() const { return Parent; }
    inline void setChild(Node* child) { Child = child; }
    inline Node* getChild() const { return Child; }
    inline void setLabel(const std::string& label) { Label = label; }
    inline const std::string& getLabel() const { return Label; }

    QueryResult query(const std::string& string) const {
      unsigned i, l;
      unsigned l1 = string.length();
      unsigned l2 = Label.length();

      // Find the length of common part
      l = std::min(l1, l2);
      i = 0;
      while ((i < l) && (string[i] == Label[i]))
        ++i;

      if (i == l) { // One is prefix of another, find who is who
        if (l1 == l2)
          return Same;
        else if (i == l1)
          return StringIsPrefix;
        else
          return LabelIsPrefix;
      } else // String and Label just have common part, return its length
        return (QueryResult)i;
    }
  };

  class Node {
    friend class Trie;

    std::map<char, Edge> Edges;
    Payload Data;
  public:
    inline explicit Node(const Payload& data):Data(data) { }
    inline Node(const Node& n) {
      Data = n.Data;
      Edges = n.Edges;
    }
    inline Node& operator=(const Node& n) {
      if (&n != this) {
        Data = n.Data;
        Edges = n.Edges;
      }

      return *this;
    }

    inline bool isLeaf() const { return Edges.empty(); }

    inline const Payload& getData() const { return Data; }
    inline void setData(const Payload& data) { Data = data; }

    inline Edge* addEdge(const std::string& Label) {
      if (!Edges.insert(std::make_pair(Label[0],
                                       Edge(Label, this))).second) {
        assert(0 && "Edge already exists!");
        return NULL;
      } else
        return &Edges[Label[0]];
    }
  };

  std::vector<Node*> Nodes;
  Payload Empty;

  inline Node* addNode(const Payload& data) {
    Node* N = new Node(data);
    Nodes.push_back(N);
    return N;
  }

  inline Node* splitEdge(Edge& cEdge, size_t index) {
    const std::string& l = cEdge.getLabel();
    assert(index < l.length() && "Trying to split too far!");

    std::string l1 = l.substr(0, index);
    std::string l2 = l.substr(index);

    Node* nNode = addNode(Empty);
    Edge* nEdge = nNode->addEdge(l2);
    nEdge->setChild(cEdge.getChild());
    cEdge.setChild(nNode);
    cEdge.setLabel(l1);

    return nNode;
  }

public:
  inline explicit Trie(const Payload& empty):Empty(empty) {
    addNode(Empty);
  }
  inline ~Trie() {
    for (unsigned i = 0, e = Nodes.size(); i != e; ++i)
      delete Nodes[i];
  }

  inline Node* getRoot() const { return Nodes[0]; }

  bool addString(const std::string& s, const Payload& data) {
    Node* cNode = getRoot();
    Edge* nEdge = NULL;
    std::string s1(s);

    while (nEdge == NULL) {
      if (cNode->Edges.count(s1[0])) {
        Edge& cEdge = cNode->Edges[s1[0]];
        typename Edge::QueryResult r = cEdge.query(s1);

        switch (r) {
        case Edge::Same:
        case Edge::StringIsPrefix:
        case Edge::DontMatch:
          assert(0 && "Impossible!");
          return false;
        case Edge::LabelIsPrefix:
          s1 = s1.substr(cEdge.getLabel().length());
          cNode = cEdge.getChild();
          break;
        default:
          nEdge = splitEdge(cEdge, r)->addEdge(s1.substr(r));
        }
      } else
        nEdge = cNode->addEdge(s1);
    }

    Node* tNode = addNode(data);
    nEdge->setChild(tNode);

    return true;
  }

  const Payload& lookup(const std::string& s) const {
    Node* cNode = getRoot();
    Node* tNode = NULL;
    std::string s1(s);

    while (tNode == NULL) {
      if (cNode->Edges.count(s1[0])) {
        Edge& cEdge = cNode->Edges[s1[0]];
        typename Edge::QueryResult r = cEdge.query(s1);

        switch (r) {
        case Edge::Same:
          tNode = cEdge.getChild();
          break;
        case Edge::StringIsPrefix:
          return Empty;
        case Edge::DontMatch:
          assert(0 && "Impossible!");
          return Empty;
        case Edge::LabelIsPrefix:
          s1 = s1.substr(cEdge.getLabel().length());
          cNode = cEdge.getChild();
          break;
        default:
          return Empty;
        }
      } else
        return Empty;
    }

    return tNode->getData();
  }

};

}

#endif // LLVM_ADT_TRIE_H