package natlab.utils;

import java.util.List;

import nodecases.AbstractNodeCaseHandler;
import ast.ASTNode;

import com.google.common.collect.Lists;

public class NodeFinder {
  public static <T extends ASTNode<?>> List<T> find(ASTNode<?> n, final Class<T> type) {
    final List<T> res = Lists.newLinkedList();
    new AbstractNodeCaseHandler() {
      public void caseASTNode(@SuppressWarnings("rawtypes") ASTNode n) {
        if (type.isInstance(n)) {
          res.add(type.cast(n));
        }
        for (int i = 0; i < n.getNumChild(); i++) {
          caseASTNode(n.getChild(i));
        }
      }
    }.caseASTNode(n);
    return res;
  }

  /**
   * Walks up the tree to find a parent of the specified type.
   * Returns null if no such parent exists.
   */
  public static <T extends ASTNode<?>> T findParent(ASTNode<?> n, Class<T> type) {
    while (n != null && (!type.isInstance(n))) {
      n = n.getParent();
    }
    if (type.isInstance(n)) {
      return type.cast(n);
    }
    return null;
  }

  public static <T extends ASTNode<?>> void apply(final ASTNode<?> n, final Class<T> type,
      final AbstractNodeFunction<T> func) {
    new AbstractNodeCaseHandler() {
      public void caseASTNode(@SuppressWarnings("rawtypes") ASTNode n) {
        if (type.isInstance(n)) {
          func.apply(type.cast(n));
        }
        for (int i = 0; i < n.getNumChild(); i++) {
          caseASTNode(n.getChild(i));
        }
      }
    }.caseASTNode(n);
  }
}
@@ -103,16 +103,99 @@
 TNode;%0A%0A
+import com.google.common.base.Predicate;%0Aimport com.google.common.base.Predicates;%0A
 import c
@@ -226,16 +226,16 @@
 .Lists;%0A
-
 %0Apublic 
@@ -338,24 +338,208 @@
 %3CT%3E type) %7B%0A
+    return find(n, type, Predicates.%3CT%3EalwaysTrue());%0A  %7D%0A%0A  public static %3CT extends ASTNode%3C?%3E%3E List%3CT%3E find(ASTNode%3C?%3E n, final Class%3CT%3E type,%0A      final Predicate%3CT%3E predicate) %7B%0A
     final Li
@@ -730,16 +730,17 @@
     
-res.add(
+T cast = 
 type
@@ -739,34 +739,111 @@
 t = type.cast(n)
-)
 ;
+%0A          if (predicate.apply(cast)) %7B%0A            res.add(cast);%0A          %7D
 %0A        %7D%0A     

Allow NodeFinder.find to take a Predicate