Add LeafIterator and ConstLeafIterator to MosesTraining::Syntax::Tree

3 files changed, 130 insertions, 3 deletions

diff --git a/phrase-extract/syntax-common/tree-inl.h b/phrase-extract/syntax-common/tree-inl.h
index 9101fc490..811bae2d2 100644
--- a/phrase-extract/syntax-common/tree-inl.h
+++ b/phrase-extract/syntax-common/tree-inl.h
@@ -118,5 +118,92 @@ bool Tree<T>::PreOrderIter<V>::operator!=(const PreOrderIter &rhs) {
   return node_ != rhs.node_;
 }
 
+template<typename T>
+template<typename V>
+class Tree<T>::LeafIter {
+ public:
+  LeafIter();
+  LeafIter(V &);
+
+  V &operator*() { return *node_; }
+  V *operator->() { return node_; }
+
+  LeafIter &operator++();
+  LeafIter operator++(int);
+
+  bool operator==(const LeafIter &);
+  bool operator!=(const LeafIter &);
+
+ private:
+  // Pointer to the current node.
+  V *node_;
+
+  // Stack of indices defining the position of node_ within the child vectors
+  // of its ancestors.
+  std::stack<std::size_t> index_stack_;
+};
+
+template<typename T>
+template<typename V>
+Tree<T>::LeafIter<V>::LeafIter()
+    : node_(0) {
+}
+
+template<typename T>
+template<typename V>
+Tree<T>::LeafIter<V>::LeafIter(V &t)
+    : node_(&t) {
+  // Navigate to the first leaf.
+  while (!node_->IsLeaf()) {
+    index_stack_.push(0);
+    node_ = node_->children()[0];
+  }
+}
+
+template<typename T>
+template<typename V>
+Tree<T>::LeafIter<V> &Tree<T>::LeafIter<V>::operator++() {
+  // Try node's ancestors until either a node is found with a sibling to the
+  // right or we reach the root (in which case the traversal is complete).
+  V *ancestor = node_->parent_;
+  while (ancestor) {
+    std::size_t index = index_stack_.top();
+    index_stack_.pop();
+    if (index+1 < ancestor->children_.size()) {
+      index_stack_.push(index+1);
+      node_ = ancestor->children()[index+1];
+      // Navigate to the first leaf.
+      while (!node_->IsLeaf()) {
+        index_stack_.push(0);
+        node_ = node_->children()[0];
+      }
+      return *this;
+    }
+    ancestor = ancestor->parent_;
+  }
+  node_ = 0;
+  return *this;
+}
+
+template<typename T>
+template<typename V>
+Tree<T>::LeafIter<V> Tree<T>::LeafIter<V>::operator++(int) {
+  LeafIter tmp(*this);
+  ++*this;
+  return tmp;
+}
+
+template<typename T>
+template<typename V>
+bool Tree<T>::LeafIter<V>::operator==(const LeafIter &rhs) {
+  return node_ == rhs.node_;
+}
+
+template<typename T>
+template<typename V>
+bool Tree<T>::LeafIter<V>::operator!=(const LeafIter &rhs) {
+  return node_ != rhs.node_;
+}
+
 }  // namespace Syntax
 }  // namespace MosesTraining
diff --git a/phrase-extract/syntax-common/tree.h b/phrase-extract/syntax-common/tree.h
index e37c2c21f..8cec07a54 100644
--- a/phrase-extract/syntax-common/tree.h
+++ b/phrase-extract/syntax-common/tree.h
@@ -72,7 +72,7 @@ class Tree {
   // iterators.  V is the value type: either Tree<T> or const Tree<T>.
   template<typename V> class PreOrderIter;
   // template<typename V> class PostOrderIter; TODO
-  // template<typename V> class LeafIter; TODO
+  template<typename V> class LeafIter;
 
  public:
   // Pre-order iterators.
@@ -84,8 +84,8 @@ class Tree {
   // typedef PostOrderIter<const Tree<T> > ConstPostOrderIterator; TODO
 
   // Leaf iterators (left-to-right).
-  // typedef LeafIter<Tree<T> > LeafIterator;  TODO
-  // typedef LeafIter<const Tree<T> > ConstLeafIterator; TODO
+  typedef LeafIter<Tree<T> > LeafIterator;
+  typedef LeafIter<const Tree<T> > ConstLeafIterator;
 
  private:
   T value_;
diff --git a/phrase-extract/syntax-common/tree_test.cc b/phrase-extract/syntax-common/tree_test.cc
index 198f52310..8e689f000 100644
--- a/phrase-extract/syntax-common/tree_test.cc
+++ b/phrase-extract/syntax-common/tree_test.cc
@@ -101,6 +101,46 @@ BOOST_AUTO_TEST_CASE(const_pre_order_1) {
   BOOST_REQUIRE(p == end);
 }
 
+// Test Tree<>::LeafIterator with a trivial, single-node tree.
+BOOST_AUTO_TEST_CASE(leaf_1) {
+  boost::scoped_ptr<Tree<int> > root(new Tree<int>(123));
+  Tree<int>::LeafIterator p(*root);
+  BOOST_REQUIRE(p != Tree<int>::LeafIterator());
+  BOOST_REQUIRE(p->value() == 123);
+  ++p;
+  BOOST_REQUIRE(p == Tree<int>::LeafIterator());
+}
+
+// Test Tree<>::LeafIterator on this tree: (1 (2 3) (4) (5 6 (7 8)))
+BOOST_AUTO_TEST_CASE(leaf_2) {
+  boost::scoped_ptr<Tree<int> > root(new Tree<int>(1));
+  root->children().push_back(new Tree<int>(2));
+  root->children()[0]->children().push_back(new Tree<int>(3));
+  root->children().push_back(new Tree<int>(4));
+  root->children().push_back(new Tree<int>(5));
+  root->children()[2]->children().push_back(new Tree<int>(6));
+  root->children()[2]->children().push_back(new Tree<int>(7));
+  root->children()[2]->children()[1]->children().push_back(new Tree<int>(8));
+  root->SetParents();
+
+  Tree<int>::LeafIterator p(*root);
+  Tree<int>::LeafIterator end;
+
+  BOOST_REQUIRE(p != end);
+  BOOST_REQUIRE(p->value() == 3);
+  ++p;
+  BOOST_REQUIRE(p != end);
+  BOOST_REQUIRE(p->value() == 4);
+  ++p;
+  BOOST_REQUIRE(p != end);
+  BOOST_REQUIRE(p->value() == 6);
+  ++p;
+  BOOST_REQUIRE(p != end);
+  BOOST_REQUIRE(p->value() == 8);
+  ++p;
+  BOOST_REQUIRE(p == end);
+}
+
 }  // namespace
 }  // namespace Syntax
 }  // namespace MosesTraining