1 files changed, 46 insertions, 0 deletions
diff --git a/decoder/ff_ruleshape.h b/decoder/ff_ruleshape.h
index 9f20faf3..488cfd84 100644
--- a/decoder/ff_ruleshape.h
+++ b/decoder/ff_ruleshape.h
@@ -2,6 +2,7 @@
 #define _FF_RULESHAPE_H_
 
 #include <vector>
+#include <map>
 #include "ff.h"
 
 class RuleShapeFeatures : public FeatureFunction {
@@ -28,4 +29,49 @@ class RuleShapeFeatures : public FeatureFunction {
   }
 };
 
+class RuleShapeFeatures2 : public FeatureFunction {
+ public:
+  ~RuleShapeFeatures2();
+  RuleShapeFeatures2(const std::string& param);
+ protected:
+  virtual void TraversalFeaturesImpl(const SentenceMetadata& smeta,
+                                     const HG::Edge& edge,
+                                     const std::vector<const void*>& ant_contexts,
+                                     SparseVector<double>* features,
+                                     SparseVector<double>* estimated_features,
+                                     void* context) const;
+ private:
+  struct Node {
+    int fid_;
+    Node() : fid_() {}
+    std::map<WordID, Node> next_;
+  };
+  mutable Node fidtree_;
+
+  inline WordID MapE(WordID w) const {
+    if (w <= 0) return kNT;
+    unsigned res = 0;
+    if (w < e2class_.size()) res = e2class_[w];
+    if (!res) res = kUNK;
+    return res;
+  }
+
+  inline WordID MapF(WordID w) const {
+    if (w <= 0) return kNT;
+    unsigned res = 0;
+    if (w < f2class_.size()) res = f2class_[w];
+    if (!res) res = kUNK;
+    return res;
+  }
+
+  // prfx_size=0 => use full word classes otherwise truncate to specified length
+  void LoadWordClasses(const std::string& fname, unsigned pfxsize, std::vector<WordID>* pv);
+  const WordID kNT;
+  const WordID kUNK;
+  std::vector<WordID> e2class_;
+  std::vector<WordID> f2class_;
+  bool has_src_;
+  bool has_trg_;
+};
+
 #endif