package natlab.toolkits.utils;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import natlab.refactoring.AbstractNodeFunction;
import nodecases.AbstractNodeCaseHandler;
import ast.ASTNode;
import ast.EmptyStmt;

public class NodeFinder {
  public static <T extends ASTNode> List<T> find(ASTNode n, final Class<T> type) {
    final List<T> res = new LinkedList<T>();
    new AbstractNodeCaseHandler() {
      public void caseASTNode(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(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);
  }

  public static Map<ASTNode, Integer> nodeCount(ASTNode t) {
    final Map<ASTNode, Integer> res = new HashMap<ASTNode, Integer>();
    new AbstractNodeCaseHandler() {
      int count = 0;
      public void caseASTNode(ASTNode n) {
        int before = count;
        if (n instanceof EmptyStmt) {
          res.put(n, 0);
          return;
        }
        count++;
        for (int i = 0; i < n.getNumChild(); i++) {
          caseASTNode(n.getChild(i));
        }
        res.put(n, count - before);
      }
    }.caseASTNode(t);
    return res;
  }
}
@@ -29,63 +29,8 @@
 s;%0A%0A
-import java.util.HashMap;%0Aimport java.util.LinkedList;%0A
 impo
@@ -204,16 +204,96 @@
 yStmt;%0A%0A
+import com.google.common.collect.Lists;%0Aimport com.google.common.collect.Maps;%0A%0A
 public c
@@ -344,16 +344,19 @@
  ASTNode
+%3C?%3E
 %3E List%3CT
@@ -365,24 +365,27 @@
 find(ASTNode
+%3C?%3E
  n, final Cl
@@ -423,20 +423,25 @@
 %3E res = 
+Lists.
 new
- 
 LinkedLi
@@ -442,19 +442,16 @@
 nkedList
-%3CT%3E
 ();%0A    
@@ -504,32 +504,62 @@
 oid caseASTNode(
+@SuppressWarnings(%22rawtypes%22) 
 ASTNode n) %7B%0A   
@@ -936,16 +936,19 @@
  ASTNode
+%3C?%3E
 %3E T find
@@ -957,24 +957,27 @@
 rent(ASTNode
+%3C?%3E
  n, Class%3CT%3E
@@ -1184,16 +1184,19 @@
  ASTNode
+%3C?%3E
 %3E void a
@@ -1213,16 +1213,19 @@
  ASTNode
+%3C?%3E
  n, fina
@@ -1343,32 +1343,62 @@
 oid caseASTNode(
+@SuppressWarnings(%22rawtypes%22) 
 ASTNode n) %7B%0A   
@@ -1628,32 +1628,35 @@
 atic Map%3CASTNode
+%3C?%3E
 , Integer%3E nodeC
@@ -1667,16 +1667,19 @@
 (ASTNode
+%3C?%3E
  t) %7B%0A  
@@ -1689,32 +1689,35 @@
 inal Map%3CASTNode
+%3C?%3E
 , Integer%3E res =
@@ -1721,37 +1721,23 @@
 s = 
+Maps.
 new
- 
 HashMap
-%3CASTNode, Integer%3E
 ();%0A
@@ -1823,16 +1823,46 @@
 ASTNode(
+@SuppressWarnings(%22rawtypes%22) 
 ASTNode 

Fix warnings in a few places