/*
 * Copyright (c) 2002 - 2014 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */

package com.ibm.wala.ipa.callgraph.pruned;

import com.ibm.wala.ipa.callgraph.CGNode;
import com.ibm.wala.ipa.callgraph.CallGraph;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public final class CallGraphPruning {

  public CallGraphPruning(CallGraph cg) {
    this.cg = cg;
  }

  private PruningPolicy pruningPolicy;
  private Set<CGNode> keep;
  private ArrayDeque<CGNode> visited;
  private List<CGNode> marked;
  private int depth;
  private final CallGraph cg;

  private final boolean DEBUG = false;

  /**
   * Searches all nodes in the callgraph that correspond to a method of the application (and not the
   * system library). It includes methods from the system library that transitively may call back
   * into the application.
   *
   * @return Set of relevant callgraph nodes.
   */
  public Set<CGNode> findApplicationNodes() {
    return findApplicationNodes(0);
  }

  /**
   * Searches all nodes in the callgraph that correspond to a method of the application (and not the
   * system library). It includes methods from the system library that transitively may call back
   * into the application. Library methods that do not transitively call back into application
   * methods are cut at the level provided by parameter depth.
   *
   * @param depth The level at which non-returning library methods are cut off.
   * @return Set of relevant callgraph nodes.
   */
  public Set<CGNode> findApplicationNodes(final int depth) {
    return findNodes(depth, ApplicationLoaderPolicy.INSTANCE);
  }

  /**
   * Searches all nodes in the callgraph according to the given pruning policy. It includes all
   * methods which transitively may call methods which comply to the given pruning policy. All other
   * methods are cut at the level provided by parameter 'depth'
   *
   * @param depth the level at which methods which do not comply to the given pruning policy are cut
   *     off.
   * @param policy pruning policy which decides which branches are kept in the call graph
   * @return set of relevant callgraph nodes
   */
  public Set<CGNode> findNodes(final int depth, PruningPolicy policy) {
    if (DEBUG) {
      System.out.println("Running optimization with depth: " + depth);
    }

    this.marked = new ArrayList<>();
    this.keep = new HashSet<>();
    this.visited = new ArrayDeque<>();
    this.depth = depth;
    this.pruningPolicy = policy;

    dfs(cg.getFakeRootNode());

    return keep;
  }

  private void dfs(CGNode root) {

    visited.addLast(root);

    Iterator<CGNode> it = cg.getSuccNodes(root);
    while (it.hasNext()) {
      CGNode next = it.next();
      if (!marked.contains(next)) {
        marked.add(next);
        dfs(next);
      } else {
        if (keep.contains(next)) {
          keep.addAll(visited);
        }
      }
    }

    if (pruningPolicy.check(root)) {
      keep.addAll(visited);
      addDepth(root);
    }
    visited.removeLast();
  }

  private void addDepth(CGNode node) {

    ArrayList<CGNode> A = new ArrayList<>();
    ArrayList<CGNode> B = new ArrayList<>();
    int i = depth;
    A.add(node);
    while (i > 0) {

      for (CGNode n : A) {
        Iterator<CGNode> it = cg.getSuccNodes(n);
        while (it.hasNext()) {
          B.add(it.next());
        }
      }

      if (DEBUG) {
        System.out.println("Tiefe: " + B);
      }

      keep.addAll(B);
      A.clear();
      A.addAll(B);
      B.clear();
      i--;
    }
  }
}
