1 files changed, 17 insertions, 7 deletions
diff --git a/src/lattice.h b/src/lattice.h
index 1177e768..71589b92 100644
--- a/src/lattice.h
+++ b/src/lattice.h
@@ -4,6 +4,14 @@
 #include <string>
 #include <vector>
 #include "wordid.h"
+#include "array2d.h"
+
+class Lattice;
+struct LatticeTools {
+  static bool LooksLikePLF(const std::string &line);
+  static void ConvertTextToLattice(const std::string& text, Lattice* pl);
+  static void ConvertTextOrPLF(const std::string& text_or_plf, Lattice* pl);
+};
 
 struct LatticeArc {
   WordID label;
@@ -14,18 +22,20 @@ struct LatticeArc {
 };
 
 class Lattice : public std::vector<std::vector<LatticeArc> > {
+  friend void LatticeTools::ConvertTextOrPLF(const std::string& text_or_plf, Lattice* pl);
  public:
   Lattice() {}
   explicit Lattice(size_t t, const std::vector<LatticeArc>& v = std::vector<LatticeArc>()) :
    std::vector<std::vector<LatticeArc> >(t, v) {}
- 
-  // TODO add distance functions
-};
+  int Distance(int from, int to) const {
+    if (dist_.empty())
+      return (to - from);
+    return dist_(from, to);
+  }
 
-struct LatticeTools {
-  static bool LooksLikePLF(const std::string &line);
-  static void ConvertTextToLattice(const std::string& text, Lattice* pl);
-  static void ConvertTextOrPLF(const std::string& text_or_plf, Lattice* pl);
+ private:
+  void ComputeDistances();
+  Array2D<int> dist_;
 };
 
 #endif