1 files changed, 10 insertions, 23 deletions
diff --git a/klm/search/vertex_generator.hh b/klm/search/vertex_generator.hh
index 6b98da3e..60e86112 100644
--- a/klm/search/vertex_generator.hh
+++ b/klm/search/vertex_generator.hh
@@ -1,13 +1,11 @@
 #ifndef SEARCH_VERTEX_GENERATOR__
 #define SEARCH_VERTEX_GENERATOR__
 
-#include "search/note.hh"
+#include "search/edge.hh"
 #include "search/vertex.hh"
 
 #include <boost/unordered_map.hpp>
 
-#include <queue>
-
 namespace lm {
 namespace ngram {
 class ChartState;
@@ -18,40 +16,29 @@ 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);
+    void NewHypothesis(PartialEdge partial) {
+      const lm::ngram::ChartState &state = partial.CompletedState();
+      std::pair<Existing::iterator, bool> ret(existing_.insert(std::make_pair(hash_value(state), partial)));
+      if (!ret.second && ret.first->second < partial) {
+        ret.first->second = partial;
+      }
     }
 
+    void FinishedSearch();
+
     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;
+    typedef boost::unordered_map<uint64_t, PartialEdge> Existing;
     Existing existing_;
 };