CallGraphHelper.java

package jqian.sootex.util.callgraph;

import soot.*;
import soot.jimple.toolkits.callgraph.*;
import soot.toolkits.graph.*;
import soot.util.queue.*;
import java.util.*;


/**
 */
@SuppressWarnings({ "unchecked", "rawtypes" })
public class CallGraphHelper{ 
	
	public static ReachableMethods getReachableMethod(CallGraph cg, Collection entries){
		ReachableMethods rm;
		Scene scene = Scene.v();			
		if(cg==scene.getCallGraph() && entries.equals(scene.getEntryPoints())){
			rm = scene.getReachableMethods();
		}
		else{
			rm = new ReachableMethods(cg,entries);
			rm.update();			
		}
		
		return rm;
	}

	static class CGNodeFilter extends CallGraphNodeFilter{
		ReachableMethods _rm;
		
		public CGNodeFilter(ReachableMethods rm){
			this._rm = rm;
		}
		
		public boolean isIngored(MethodOrMethodContext m){
			return !_rm.contains(m);
		}
	}
	
	public static CallGraphNodeFilter getCallGraphNodeFilter(ReachableMethods rm){
		CGNodeFilter nodeFilter = new CGNodeFilter(rm);	
		return nodeFilter;
	}
	
	public DirectedGraph<MethodOrMethodContext> toDirectedGraph(CallGraph cg,Collection<MethodOrMethodContext> entries){
		ReachableMethods rm = getReachableMethod(cg, entries);
		CGNodeFilter nodeFilter = new CGNodeFilter(rm);		
		DirectedCallGraph dcg = new DirectedCallGraph(cg, entries, nodeFilter, new CallGraphEdgeFilter());
		return dcg;
	}
		
	private static class SubCallGraph extends CallGraph {
		Set<MethodOrMethodContext> _nodes;
		CallGraph _cg;
		CallGraphEdgeFilter _filter;

		public SubCallGraph(CallGraph cg, Set<MethodOrMethodContext> nodes, CallGraphEdgeFilter filter) {
			_nodes = nodes;
			_cg = cg;
			this._filter = filter;
		}

		public Iterator<Edge> edgesInto(MethodOrMethodContext m) {
			List<Edge> edgesIn = new LinkedList<Edge>();
			for (Iterator<Edge> edges = _cg.edgesInto(m); edges.hasNext();) {
				Edge e = edges.next();
				if (_nodes.contains(e.src()) && !_filter.isIgnored(e)) {
					edgesIn.add(e);
				}
			}
			
			return edgesIn.iterator();
		}

		public Iterator<Edge> edgesOutOf(MethodOrMethodContext m) {
			List<Edge> edgesOut = new LinkedList<Edge>();
			for (Iterator<Edge> edges = _cg.edgesOutOf(m); edges.hasNext();) {
				Edge e = edges.next();
				if (_nodes.contains(e.tgt()) && !_filter.isIgnored(e)) {
					edgesOut.add(e);
				}
			}
			return edgesOut.iterator();
		}

		public Iterator<Edge> edgesOutOf(Unit u) {
			throw new RuntimeException("Not implemented");
		}

		public QueueReader<Edge> listener() {
			throw new RuntimeException("Not implemented");
		}

		public boolean removeEdge(Edge e) {
			throw new RuntimeException("Not implemented");
		}

		// Returns the number of edges in the call graph.
		public int size() {
			throw new RuntimeException("Not implemented");
		}

		public Iterator<MethodOrMethodContext> sourceMethods() {
			return _nodes.iterator();
		}

		public String toString() {
			throw new RuntimeException("Not implemented");
		}
	}
	
	public static CallGraph getSubCallGraph(CallGraph cg,Set nodes,CallGraphEdgeFilter filter){
		return new SubCallGraph(cg,nodes,filter);
	} 
	
	public static CallGraph getSubCallGraph(CallGraph cg,Set nodes){
		return new SubCallGraph(cg,nodes,new CallGraphEdgeFilter());
	} 
	
	public static boolean hasDirectRecurion(CallGraph cg,SootMethod m){
		for(Iterator<Edge> it = cg.edgesOutOf(m);it.hasNext();){
			Edge e = it.next();
			SootMethod tgt = e.tgt();
			if(tgt==m){
				return true;
			}
		}	
		
		return false;
	}
	
	/** Make sure the call information of method <code>m</code> occur in the call graph. */
	public static void assureCallGraph(CallGraph cg, SootMethod m){
		
	}
	
	
    private static boolean isLibMethod(SootMethod m){
    	SootClass cls = m.getDeclaringClass();
    	String className = cls.getName();

    	int size = JAVA_LIB_PACKAGES.length;
    	for(int i=0;i<size;i++){
    		String pkg = JAVA_LIB_PACKAGES[i];
    		if(className.startsWith(pkg)){
            	return true;
            }
    	}		
          
    	return false;      
    }
    
    private static String[] JAVA_LIB_PACKAGES = {
    	"sun.", "sunw.", "com.sun.", "java.", "javax."
    };
    
    
    public static Collection<SootMethod> widthFirstTraverse(CallGraph cg, SootMethod entry, int depth){
		Queue<SootMethod> queue = new LinkedList<SootMethod>();
		queue.add(entry);	    
	    Set<SootMethod> processed = new HashSet<SootMethod>();
	    
	    for(int i=0; i<depth && !queue.isEmpty(); i++){
	    	// process elements generated by last round 
	    	int size = queue.size();
	    	
	    	for(int k=0;k<size;k++){
	    		SootMethod m = (SootMethod)queue.poll(); 
	    		
	    		if(!processed.add(m)){
	    			continue;
	    		}
		    	
		    	for(Iterator<Edge> it = cg.edgesOutOf(m);it.hasNext();){
		    		Edge e = it.next();
		    		
		    		SootMethod tgt = e.tgt();		    		
		    		if(!processed.contains(tgt)){
			           	queue.offer(tgt);
		    		}		    		 
		    	}
	    	}
	    }
      
        return processed; 
	}
    
	/**
	 * Tailor call graph. All methods called in a library method more than depth 
	 * <code>libDepth</code> will be tailored off.
	 */
	public static CallGraph tailorLibMethods(CallGraph cg, Collection<SootMethod> entries, int libDepth){
		libDepth--;
		CallGraph newCg = new CallGraph();	
		Stack<SootMethod> stack = new Stack<SootMethod>();
	    Map<SootMethod,Integer> m2depth = new HashMap<SootMethod,Integer>();
	    Set<SootMethod> processed = new HashSet<SootMethod>();
	    
	    stack.addAll(entries);
	    
	    while(!stack.isEmpty()){
	    	SootMethod m = (SootMethod)stack.pop();
	    	
	    	boolean firstTime = processed.add(m);
	    	Integer depth = m2depth.get(m);
	    	if(depth==null){
	    		depth = 0;
	    		m2depth.put(m, 0);
	    	}
	    	depth = depth + 1;
	    	
	    	for(Iterator<Edge> it = cg.edgesOutOf(m);it.hasNext();){
	    		Edge e = it.next();	    
	    		SootMethod tgt = e.tgt();
	    		
	    		if(firstTime){
	    			newCg.addEdge(e); 
	    		}
	    		
	    		if(isLibMethod(tgt)){	    			
	    			if(depth <= libDepth){
	    				Integer lastDepth = m2depth.get(tgt);
	    				
	    				// non passed, or this time depth is smaller
	    				if(lastDepth==null || depth<lastDepth){
	    					m2depth.put(tgt, depth);
		    				stack.push(tgt);
	    				}
	    			}
	    		}
	    		else{
	    			if(!processed.contains(tgt)){
	    				m2depth.put(tgt, 0);
	    				stack.push(tgt);	    				
	    			}
	    		}
	    	}
	    }
      
        return newCg; 
	}
 }