1 files changed, 32 insertions, 21 deletions

diff --git a/rst_parser/arc_factored.h b/rst_parser/arc_factored.h
index 3003a86e..a95f8230 100644
--- a/rst_parser/arc_factored.h
+++ b/rst_parser/arc_factored.h
@@ -5,37 +5,52 @@
 #include <cassert>
 #include <vector>
 #include <utility>
+#include <boost/shared_ptr.hpp>
 #include "array2d.h"
 #include "sparse_vector.h"
 #include "prob.h"
 #include "weights.h"
+#include "wordid.h"
+
+struct TaggedSentence {
+  std::vector<WordID> words;
+  std::vector<WordID> pos;
+};
 
 struct EdgeSubset {
   EdgeSubset() {}
   std::vector<short> roots; // unless multiroot trees are supported, this
                             // will have a single member
-  std::vector<std::pair<short, short> > h_m_pairs; // h,m start at *1*
+  std::vector<std::pair<short, short> > h_m_pairs; // h,m start at 0
 };
 
+struct ArcFeatureFunction;
 class ArcFactoredForest {
  public:
-  explicit ArcFactoredForest(short num_words) :
-      num_words_(num_words),
-      root_edges_(num_words),
-      edges_(num_words, num_words) {
+  ArcFactoredForest() : num_words_() {}
+  explicit ArcFactoredForest(short num_words) {
+    resize(num_words);
+  }
+
+  void resize(unsigned num_words) {
+    num_words_ = num_words;
+    root_edges_.clear();
+    edges_.clear();
+    root_edges_.resize(num_words);
+    edges_.resize(num_words, num_words);
     for (int h = 0; h < num_words; ++h) {
       for (int m = 0; m < num_words; ++m) {
-        edges_(h, m).h = h + 1;
-        edges_(h, m).m = m + 1;
+        edges_(h, m).h = h;
+        edges_(h, m).m = m;
       }
-      root_edges_[h].h = 0;
-      root_edges_[h].m = h + 1;
+      root_edges_[h].h = -1;
+      root_edges_[h].m = h;
     }
   }
 
   // compute the maximum spanning tree based on the current weighting
   // using the O(n^2) CLE algorithm
-  void MaximumEdgeSubset(EdgeSubset* st) const;
+  void MaximumSpanningTree(EdgeSubset* st) const;
 
   // Reweight edges so that edge_prob is the edge's marginals
   // optionally returns log partition
@@ -52,20 +67,16 @@ class ArcFactoredForest {
     prob_t edge_prob;
   };
 
+  // set eges_[*].features
+  void ExtractFeatures(const TaggedSentence& sentence,
+                       const std::vector<boost::shared_ptr<ArcFeatureFunction> >& ffs);
+
   const Edge& operator()(short h, short m) const {
-    assert(m > 0);
-    assert(m <= num_words_);
-    assert(h >= 0);
-    assert(h <= num_words_);
-    return h ? edges_(h - 1, m - 1) : root_edges_[m - 1];
+    return h >= 0 ? edges_(h, m) : root_edges_[m];
   }
 
   Edge& operator()(short h, short m) {
-    assert(m > 0);
-    assert(m <= num_words_);
-    assert(h >= 0);
-    assert(h <= num_words_);
-    return h ? edges_(h - 1, m - 1) : root_edges_[m - 1];
+    return h >= 0 ? edges_(h, m) : root_edges_[m];
   }
 
   float Weight(short h, short m) const {
@@ -87,7 +98,7 @@ class ArcFactoredForest {
   }
 
  private:
-  unsigned num_words_;
+  int num_words_;
   std::vector<Edge> root_edges_;
   Array2D<Edge> edges_;
 };