Add an ith leaf relation to tregex

Author: AngledLuffa

src/edu/stanford/nlp/trees/tregex/Relation.java

@@ -131,7 +131,7 @@ static Relation getRelation(String s,
   /**
    * Static factory method for relations requiring an argument, including
    * HAS_ITH_CHILD, ITH_CHILD_OF, UNBROKEN_CATEGORY_DOMINATES,
-   * UNBROKEN_CATEGORY_DOMINATED_BY.
+   * UNBROKEN_CATEGORY_DOMINATED_BY, ANCESTOR_OF_ITH_LEAF.
    *
    * @param s The String representation of the relation
    * @param arg The argument to the relation, as a string; could be a node
@@ -168,6 +168,9 @@ static Relation getRelation(String s, String arg,
       case ",+":
         r = new UnbrokenCategoryFollows(arg, basicCatFunction);
         break;
+      case "<<<":
+        r = new AncestorOfIthLeaf(Integer.parseInt(arg));
+        break;
       default:
         throw new ParseException("Unrecognized compound relation " + s + ' '
             + arg);
@@ -362,6 +365,89 @@ void advance() {
     }
   };
 
+  /**
+   * Looks for the ith leaf of the current node
+   */
+  private static class AncestorOfIthLeaf extends Relation {
+
+    private static final long serialVersionUID = -6495191354526L;
+
+    private final int leafNum;
+
+    AncestorOfIthLeaf(int i) {
+      super("<<<" + String.valueOf(i));
+      if (i == 0) {
+        throw new IllegalArgumentException("Error -- no such thing as zeroth leaf!");
+      }
+      leafNum = i;
+    }
+
+    @Override
+    boolean satisfies(Tree t1, Tree t2, Tree root, final TregexMatcher matcher) {
+      if (t1 == t2)
+        return false;
+      if (!t2.isLeaf())
+        return false;
+      // this is kind of lazy
+      // if it somehow became a performance limitation,
+      // a recursive search would be faster
+      List<Tree> leaves = t1.getLeaves();
+      if (leaves.size() < Math.abs(leafNum))
+        return false;
+
+      final int index;
+      if (leafNum > 0) {
+        index = leafNum - 1;
+      } else {
+        // eg, leafNum == -1 means we check leaves.size() - 1
+        index = leaves.size() + leafNum;
+      }
+      return leaves.get(index) == t2;
+    }
+
+    @Override
+    Iterator<Tree> searchNodeIterator(final Tree t,
+                                      final TregexMatcher matcher) {
+      return new SearchNodeIterator() {
+        @Override
+        void initialize() {
+          List<Tree> leaves = t.getLeaves();
+          if (leaves.size() >= Math.abs(leafNum)) {
+            final int index;
+            if (leafNum > 0) {
+              index = leafNum - 1;
+            } else {
+              index = leafNum + leaves.size();
+            }
+            next = leaves.get(index);
+          }
+        }
+      };
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof AncestorOfIthLeaf)) {
+        return false;
+      }
+
+      final AncestorOfIthLeaf other = (AncestorOfIthLeaf) o;
+      if (leafNum != other.leafNum) {
+        return false;
+      }
+
+      return true;
+    }
+
+    @Override
+    public int hashCode() {
+      return leafNum + 20;
+    }
+  };
+
   private static final Relation DOMINATES = new Relation("<<") {
 
     private static final long serialVersionUID = -2580199434621268260L;

src/edu/stanford/nlp/trees/tregex/TregexPattern.java

@@ -114,6 +114,8 @@
  * <tr><td>A &gt;i B <td>A is the ith child of B (i &gt; 0)
  * <tr><td>A &lt;-i B <td>B is the ith-to-last child of A (i &gt; 0)
  * <tr><td>A &gt;-i B <td>A is the ith-to-last child of B (i &gt; 0)
+ * <tr><td>A &lt;&lt;&lt;i B <td>B is the ith leaf of A
+ * <tr><td>A &lt;&lt;&lt;-i B <td>B is the ith-to-last leaf of A
  * <tr><td>A &lt;: B <td>B is the only child of A
  * <tr><td>A &gt;: B <td>A is the only child of B
  * <tr><td>A &lt;&lt;: B <td>A dominates B via an unbroken chain (length &gt; 0) of unary local trees.

test/src/edu/stanford/nlp/trees/tregex/TregexTest.java

@@ -1081,6 +1081,18 @@ public void testAncestorOfLeaf() {
     runTest("A <<< (b . c=foo) <<< =foo", "(ROOT (A (B b) (C c)))", "(A (B b) (C c))");
   }
 
+  public void testAncestorOfIthLeaf() {
+    runTest("A <<<1 b", "(ROOT (A (B b)))", "(A (B b))");
+    runTest("A <<<2 b", "(ROOT (A (B b)))");
+    runTest("A <<<-1 b", "(ROOT (A (B b)))", "(A (B b))");
+    runTest("A <<<1 b", "(ROOT (A (B z) (C b)))");
+    runTest("A <<<2 b", "(ROOT (A (B z) (C b)))", "(A (B z) (C b))");
+    runTest("A <<<-1 b", "(ROOT (A (B z) (C b)))", "(A (B z) (C b))");
+    runTest("A <<<-2 b", "(ROOT (A (B z) (C b)))");
+    runTest("A <<<-1 z", "(ROOT (A (B z) (C b)))");
+    runTest("A <<<-2 z", "(ROOT (A (B z) (C b)))", "(A (B z) (C b))");
+  }
+
   public void testHeadOfPhrase() {
     runTest("NP <# NNS", "(NP (NN work) (NNS practices))", "(NP (NN work) (NNS practices))");
     runTest("NP <# NN", "(NP (NN work) (NNS practices))"); // should have no results