package natlab.utils;

import java.util.Iterator;
import java.util.Stack;

import ast.ASTNode;

import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.AbstractSequentialIterator;
import com.google.common.collect.FluentIterable;

/**
 * A utility that finds nodes of a certain type in the AST, and returns them as a lazy iterable.
 * This class is intended to be used together with Guava's Iterable utilities; see also
 * AstPredicates and AstFunctions for utilities related to this. For example, the
 * following snippet searches <tt>tree</tt> for all functions that aren't nested, returns
 * a list of their names:
 * 
 * <pre>
 *  NodeFinder.find(ast.Function.class, tree))
 *    .filter(Predicates.not(AstPredicates.nestedFunction()))
 *    .transform(AstFunctions.functionToName())
 *    .toImmutableList();
 * </pre>
 */
public class NodeFinder {
  /**
   * Returns an iterable of all the descendants of <tt>tree</tt>. This corresponds to a
   * depth-first traversal of the subtree rooted at <tt>tree</tt>.
   */
  public static Iterable<ASTNode<?>> allDescendantsOf(final ASTNode<?> tree) {
    Preconditions.checkNotNull(tree);
    return new Iterable<ASTNode<?>>() {
      @Override public Iterator<ASTNode<?>> iterator() {
        final Stack<ASTNode<?>> toVisit = new Stack<ASTNode<?>>();
        toVisit.push(tree);
        return new AbstractIterator<ASTNode<?>>() {
          @Override protected ASTNode<?> computeNext() {
            if (toVisit.empty()) {
              return endOfData();
            }
            ASTNode<?> next = toVisit.pop();
            for (ASTNode<?> child : next) {
              toVisit.push(child);
            }
            return next;
          }
        };
      }
    };
  }
  
  /**
   * Returns an iterable of the ancestors of <tt>tree</tt>. 
   */
  public static Iterable<ASTNode<?>> allAncestorsOf(final ASTNode<?> tree) {
    Preconditions.checkNotNull(tree);
    return new Iterable<ASTNode<?>>() {
      @Override public Iterator<ASTNode<?>> iterator() {
        return new AbstractSequentialIterator<ASTNode<?>>(tree.getParent()) {
          @Override protected ASTNode<?> computeNext(ASTNode<?> previous) {
            return previous.getParent();
          }
        };
      }
    };
  }

  /**
   * Returns a lazy iterable of the nodes of this finder's type in <tt>tree</tt>.
   */
  public static <T> FluentIterable<T> find(final Class<T> clazz, final ASTNode<?> tree) {
    return FluentIterable.from(allDescendantsOf(tree)).filter(clazz);
  }

  /**
   * Walks up the tree to find a parent of the specified type.
   * Returns null if no such parent exists.
   */
  public static <T> T findParent(Class<T> clazz, ASTNode<?> node) {
    return FluentIterable.from(allAncestorsOf(node)).filter(clazz).first().orNull();
  }

  /**
   * Applies <tt>func</tt> to each node of type <tt>type</tt> in <tt>n</tt>.
   */
  public static <T> void apply(final Class<T> type, final ASTNode<?> n,
      final AbstractNodeFunction<T> func) {
    for (T node : NodeFinder.find(type, n)) {
      func.apply(node);
    }
  }
}
@@ -60,21 +60,20 @@
 va.util.
-Stack
+List
 ;%0A%0Aimpor
@@ -137,59 +137,8 @@
 ns;%0A
-import com.google.common.collect.AbstractIterator;%0A
 impo
@@ -242,16 +242,104 @@
 terable;
+%0Aimport com.google.common.collect.Lists;%0Aimport com.google.common.collect.TreeTraverser;
 %0A%0A/**%0A *
@@ -961,173 +961,102 @@
 %7B%0A  
-/**%0A   * Returns an iterable of all the descendants of %3Ctt%3Etree%3C/tt%3E. This corresponds to a%0A   * depth-first traversal of the subtree rooted at %3Ctt%3Etree%3C/tt%3E.%0A   */%0A
+private static class AstNodeTreeTraverser extends TreeTraverser%3CASTNode%3C?%3E%3E %7B%0A    @Override%0A  
   pu
@@ -1052,39 +1052,32 @@
 ride%0A    public 
-static 
 Iterable%3CASTNode
@@ -1085,31 +1085,17 @@
 ?%3E%3E 
-allDescendantsOf(final 
+children(
 ASTN
@@ -1097,36 +1097,36 @@
 (ASTNode%3C?%3E 
-tree
+root
 ) %7B%0A    
 Precondition
@@ -1117,477 +1117,257 @@
     
-Preconditions.checkNotNull(tree);%0A    return new Iterable%3CASTNode%3C?%3E%3E() %7B%0A      @Override public Iterator%3CASTNode%3C?%3E%3E iterator() %7B%0A        final Stack%3CASTNode%3C?%3E%3E toVisit = new Stack%3CASTNode%3C?%3E%3E();%0A        toVisit.push(tree);%0A        return new AbstractIterator%3CASTNode%3C?%3E%3E() %7B%0A          @Override protected ASTNode%3C?%3E computeNext() %7B%0A            if (toVisit.empty()) %7B%0A              return endOfData();%0A            %7D%0A            ASTNode%3C?%3E next = toVisit.pop();%0A      
+  // ASTNode is iterable, but it's Iterable%3CT extends ASTNode%3E (without the %3C?%3E).%0A      // If it was Iterable%3CT extends ASTNode%3C?%3E%3E, we could just write%0A      // return Iterables.concat(root);%0A      List%3CASTNode%3C?%3E%3E children = Lists.newArrayList();%0A
     
@@ -1392,19 +1392,19 @@
 child : 
-nex
+roo
 t) %7B%0A   
@@ -1412,26 +1412,20 @@
     
-      toVisit.push
+children.add
 (chi
@@ -1439,77 +1439,373 @@
     
+%7D%0A
       
-%7D%0A            return next;%0A          %7D%0A        %7D;%0A      %7D%0A    %7D
+return children;%0A    %7D%0A  %7D%0A%0A  /**%0A   * Returns an iterable of all the descendants of %3Ctt%3Etree%3C/tt%3E. This corresponds to a%0A   * depth-first traversal of the subtree rooted at %3Ctt%3Etree%3C/tt%3E.%0A   */%0A  public static Iterable%3CASTNode%3C?%3E%3E allDescendantsOf(ASTNode%3C?%3E tree) %7B%0A    return new AstNodeTreeTraverser().preOrderTraversal(Preconditions.checkNotNull(tree))
 ;%0A  

Simplify NodeFinder using Guava's TreeTraverser.