1 files changed, 59 insertions, 0 deletions
diff --git a/klm/search/vertex_generator.hh b/klm/search/vertex_generator.hh
new file mode 100644
index 00000000..6b98da3e
--- /dev/null
+++ b/klm/search/vertex_generator.hh
@@ -0,0 +1,59 @@
+#ifndef SEARCH_VERTEX_GENERATOR__
+#define SEARCH_VERTEX_GENERATOR__
+
+#include "search/note.hh"
+#include "search/vertex.hh"
+
+#include <boost/unordered_map.hpp>
+
+#include <queue>
+
+namespace lm {
+namespace ngram {
+class ChartState;
+} // namespace ngram
+} // namespace lm
+
+namespace search {
+
+class ContextBase;
+class Final;
+struct PartialEdge;
+
+class VertexGenerator {
+  public:
+    VertexGenerator(ContextBase &context, Vertex &gen);
+
+    void NewHypothesis(const PartialEdge &partial, Note note);
+
+    void FinishedSearch() {
+      root_.under->SortAndSet(context_, NULL);
+    }
+
+    const Vertex &Generating() const { return gen_; }
+
+  private:
+    // Parallel structure to VertexNode.  
+    struct Trie {
+      Trie() : under(NULL) {}
+
+      VertexNode *under;
+      boost::unordered_map<uint64_t, Trie> extend;
+    };
+
+    Trie &FindOrInsert(Trie &node, uint64_t added, const lm::ngram::ChartState &state, unsigned char left, bool left_full, unsigned char right, bool right_full);
+
+    Final *CompleteTransition(Trie &node, const lm::ngram::ChartState &state, Note note, const PartialEdge &partial);
+
+    ContextBase &context_;
+
+    Vertex &gen_;
+
+    Trie root_;
+
+    typedef boost::unordered_map<uint64_t, Final*> Existing;
+    Existing existing_;
+};
+
+} // namespace search
+#endif // SEARCH_VERTEX_GENERATOR__