1 files changed, 24 insertions, 25 deletions
diff --git a/klm/search/edge_generator.hh b/klm/search/edge_generator.hh
index 875ccc5e..582c78b7 100644
--- a/klm/search/edge_generator.hh
+++ b/klm/search/edge_generator.hh
@@ -3,11 +3,8 @@
 
 #include "search/edge.hh"
 #include "search/note.hh"
+#include "search/types.hh"
 
-#include <boost/pool/pool.hpp>
-#include <boost/unordered_map.hpp>
-
-#include <functional>
 #include <queue>
 
 namespace lm {
@@ -20,38 +17,40 @@ namespace search {
 
 template <class Model> class Context;
 
-class VertexGenerator;
-
-struct PartialEdgePointerLess : std::binary_function<const PartialEdge *, const PartialEdge *, bool> {
-  bool operator()(const PartialEdge *first, const PartialEdge *second) const {
-    return *first < *second;
-  }
-};
-
 class EdgeGenerator {
   public:
-    EdgeGenerator(PartialEdge &root, unsigned char arity, Note note);
+    EdgeGenerator() {}
 
-    Score TopScore() const {
-      return top_score_;
+    PartialEdge AllocateEdge(Arity arity) {
+      return PartialEdge(partial_edge_pool_, arity);
     }
 
-    Note GetNote() const {
-      return note_;
+    void AddEdge(PartialEdge edge) {
+      generate_.push(edge);
     }
 
-    // Pop.  If there's a complete hypothesis, return it.  Otherwise return NULL.  
-    template <class Model> PartialEdge *Pop(Context<Model> &context, boost::pool<> &partial_edge_pool);
+    bool Empty() const { return generate_.empty(); }
+
+    // Pop.  If there's a complete hypothesis, return it.  Otherwise return an invalid PartialEdge.
+    template <class Model> PartialEdge Pop(Context<Model> &context);
+
+    template <class Model, class Output> void Search(Context<Model> &context, Output &output) {
+      unsigned to_pop = context.PopLimit();
+      while (to_pop > 0 && !generate_.empty()) {
+        PartialEdge got(Pop(context));
+        if (got.Valid()) {
+          output.NewHypothesis(got);
+          --to_pop;
+        }
+      }
+      output.FinishedSearch();
+    }
 
   private:
-    Score top_score_;
-
-    unsigned char arity_;
+    util::Pool partial_edge_pool_;
 
-    typedef std::priority_queue<PartialEdge*, std::vector<PartialEdge*>, PartialEdgePointerLess> Generate;
+    typedef std::priority_queue<PartialEdge> Generate;
     Generate generate_;
-
-    Note note_;
 };
 
 } // namespace search