1 files changed, 8 insertions, 24 deletions

diff --git a/klm/search/vertex_generator.cc b/klm/search/vertex_generator.cc
index 78948c97..d94e6e06 100644
--- a/klm/search/vertex_generator.cc
+++ b/klm/search/vertex_generator.cc
@@ -2,45 +2,30 @@
 
 #include "lm/left.hh"
 #include "search/context.hh"
+#include "search/edge.hh"
 
 #include <stdint.h>
 
 namespace search {
 
-template <class Model> VertexGenerator::VertexGenerator(Context<Model> &context, Vertex &gen) : context_(context), edges_(gen.edges_.size()), partial_edge_pool_(sizeof(PartialEdge), context.PopLimit() * 2) {
-  for (std::size_t i = 0; i < gen.edges_.size(); ++i) {
-    if (edges_[i].Init(*gen.edges_[i], *this))
-      generate_.push(&edges_[i]);
-  }
+VertexGenerator::VertexGenerator(ContextBase &context, Vertex &gen) : context_(context), gen_(gen) {
   gen.root_.InitRoot();
   root_.under = &gen.root_;
-  to_pop_ = context.PopLimit();
-  while (to_pop_ > 0 && !generate_.empty()) {
-    EdgeGenerator *top = generate_.top();
-    generate_.pop();
-    if (top->Pop(context, *this)) {
-      generate_.push(top);
-    }
-  }
-  gen.root_.SortAndSet(context, NULL);
 }
 
-template VertexGenerator::VertexGenerator(Context<lm::ngram::ProbingModel> &context, Vertex &gen);
-template VertexGenerator::VertexGenerator(Context<lm::ngram::RestProbingModel> &context, Vertex &gen);
-
 namespace {
 const uint64_t kCompleteAdd = static_cast<uint64_t>(-1);
 } // namespace
 
-void VertexGenerator::NewHypothesis(const lm::ngram::ChartState &state, const Edge &from, const PartialEdge &partial) {
+void VertexGenerator::NewHypothesis(const PartialEdge &partial, Note note) {
+  const lm::ngram::ChartState &state = partial.CompletedState();
   std::pair<Existing::iterator, bool> got(existing_.insert(std::pair<uint64_t, Final*>(hash_value(state), NULL)));
   if (!got.second) {
     // Found it already.  
     Final &exists = *got.first->second;
     if (exists.Bound() < partial.score) {
-      exists.Reset(partial.score, from, partial.nt[0].End(), partial.nt[1].End());
+      exists.Reset(partial.score, note, partial.nt[0].End(), partial.nt[1].End());
     }
-    --to_pop_;
     return;
   }
   unsigned char left = 0, right = 0;
@@ -67,8 +52,7 @@ void VertexGenerator::NewHypothesis(const lm::ngram::ChartState &state, const Ed
   }
 
   node = &FindOrInsert(*node, kCompleteAdd - state.left.full, state, state.left.length, true, state.right.length, true);
-  got.first->second = CompleteTransition(*node, state, from, partial);
-  --to_pop_;
+  got.first->second = CompleteTransition(*node, state, note, partial);
 }
 
 VertexGenerator::Trie &VertexGenerator::FindOrInsert(VertexGenerator::Trie &node, uint64_t added, const lm::ngram::ChartState &state, unsigned char left, bool left_full, unsigned char right, bool right_full) {
@@ -86,12 +70,12 @@ VertexGenerator::Trie &VertexGenerator::FindOrInsert(VertexGenerator::Trie &node
   return next;
 }
 
-Final *VertexGenerator::CompleteTransition(VertexGenerator::Trie &starter, const lm::ngram::ChartState &state, const Edge &from, const PartialEdge &partial) {
+Final *VertexGenerator::CompleteTransition(VertexGenerator::Trie &starter, const lm::ngram::ChartState &state, Note note, const PartialEdge &partial) {
   VertexNode &node = *starter.under;
   assert(node.State().left.full == state.left.full);
   assert(!node.End());
   Final *final = context_.NewFinal();
-  final->Reset(partial.score, from, partial.nt[0].End(), partial.nt[1].End());
+  final->Reset(partial.score, note, partial.nt[0].End(), partial.nt[1].End());
   node.SetEnd(final);
   return final;
 }