ChartKBestExtractor: fix memory leak, clean-up code

1 files changed, 33 insertions, 30 deletions

diff --git a/moses/ChartKBestExtractor.h b/moses/ChartKBestExtractor.h
index 66430ec1e..05b016d50 100644
--- a/moses/ChartKBestExtractor.h
+++ b/moses/ChartKBestExtractor.h
@@ -24,6 +24,7 @@
 #include "ScoreComponentCollection.h"
 
 #include <boost/unordered_set.hpp>
+#include <boost/weak_ptr.hpp>
 
 #include <queue>
 #include <vector>
@@ -53,51 +54,30 @@ public:
 
     UnweightedHyperarc edge;
     std::vector<std::size_t> backPointers;
+    std::vector<boost::shared_ptr<Derivation> > subderivations;
     ScoreComponentCollection scoreBreakdown;
     float score;
   };
 
   struct DerivationOrderer {
-    bool operator()(const boost::shared_ptr<Derivation> &d1,
-                    const boost::shared_ptr<Derivation> &d2) const {
-      return d1->score < d2->score;
-    }
-  };
-
-  struct DerivationHasher {
-    std::size_t operator()(const boost::shared_ptr<Derivation> &d) const {
-      std::size_t seed = 0;
-      boost::hash_combine(seed, d->edge.head);
-      boost::hash_combine(seed, d->edge.tail);
-      boost::hash_combine(seed, d->backPointers);
-      return seed;
-    }
-  };
-
-  struct DerivationEqualityPred {
-    bool operator()(const boost::shared_ptr<Derivation> &d1,
-                    const boost::shared_ptr<Derivation> &d2) const {
-      return d1->edge.head == d2->edge.head &&
-             d1->edge.tail == d2->edge.tail &&
-             d1->backPointers == d2->backPointers;
+    bool operator()(const boost::weak_ptr<Derivation> &d1,
+                    const boost::weak_ptr<Derivation> &d2) const {
+      boost::shared_ptr<Derivation> s1(d1);
+      boost::shared_ptr<Derivation> s2(d2);
+      return s1->score < s2->score;
     }
   };
 
   struct Vertex {
-    typedef std::priority_queue<boost::shared_ptr<Derivation>,
-                                std::vector<boost::shared_ptr<Derivation> >,
+    typedef std::priority_queue<boost::weak_ptr<Derivation>,
+                                std::vector<boost::weak_ptr<Derivation> >,
                                 DerivationOrderer> DerivationQueue;
 
-    typedef boost::unordered_set<boost::shared_ptr<Derivation>,
-                                 DerivationHasher,
-                                 DerivationEqualityPred> DerivationSet;
-
     Vertex(const ChartHypothesis &h) : hypothesis(h), visited(false) {}
 
     const ChartHypothesis &hypothesis;
-    std::vector<boost::shared_ptr<Derivation> > kBestList;
+    std::vector<boost::weak_ptr<Derivation> > kBestList;
     DerivationQueue candidates;
-    DerivationSet seen;
     bool visited;
   };
 
@@ -114,6 +94,28 @@ private:
   typedef boost::unordered_map<const ChartHypothesis *,
                                boost::shared_ptr<Vertex> > VertexMap;
 
+  struct DerivationHasher {
+    std::size_t operator()(const boost::shared_ptr<Derivation> &d) const {
+      std::size_t seed = 0;
+      boost::hash_combine(seed, d->edge.head);
+      boost::hash_combine(seed, d->edge.tail);
+      boost::hash_combine(seed, d->backPointers);
+      return seed;
+    }
+  };
+
+  struct DerivationEqualityPred {
+    bool operator()(const boost::shared_ptr<Derivation> &d1,
+                    const boost::shared_ptr<Derivation> &d2) const {
+      return d1->edge.head == d2->edge.head &&
+             d1->edge.tail == d2->edge.tail &&
+             d1->backPointers == d2->backPointers;
+    }
+  };
+
+  typedef boost::unordered_set<boost::shared_ptr<Derivation>, DerivationHasher,
+                               DerivationEqualityPred> DerivationSet;
+
   UnweightedHyperarc CreateEdge(const ChartHypothesis &);
   boost::shared_ptr<Vertex> FindOrCreateVertex(const ChartHypothesis &);
   void GetCandidates(Vertex &, std::size_t);
@@ -121,6 +123,7 @@ private:
   void LazyNext(Vertex &, const Derivation &, std::size_t);
 
   VertexMap m_vertexMap;
+  DerivationSet m_derivations;
 };
 
 }  // namespace Moses