/*
 * ControlFlowGraph.java - This file is part of the Jakstab project.
 * Copyright 2007-2015 Johannes Kinder <jk@jakstab.org>
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, see <http://www.gnu.org/licenses/>.
 */
package org.jakstab.cfa;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.jakstab.Program;
import org.jakstab.rtl.statements.BasicBlock;
import org.jakstab.rtl.statements.RTLAssume;
import org.jakstab.rtl.statements.RTLCallReturn;
import org.jakstab.rtl.statements.RTLSkip;
import org.jakstab.rtl.statements.RTLStatement;
import org.jakstab.util.FastSet;
import org.jakstab.util.Logger;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.SetMultimap;

public class ControlFlowGraph {
	
	private static final Logger logger = Logger.getLogger(ControlFlowGraph.class);

	private Location entryPoint;
	private Set<Location> locations;
	private SetMultimap<Location, CFAEdge> outEdges;
	private SetMultimap<Location, CFAEdge> inEdges;

	private Map<Location, BasicBlock> basicBlocks;
	private SetMultimap<Location, CFAEdge> bbOutEdges;
	private SetMultimap<Location, CFAEdge> bbInEdges;
	
	protected ControlFlowGraph() {
		outEdges = HashMultimap.create();
		inEdges = HashMultimap.create();
		locations = new HashSet<Location>();

		basicBlocks = new HashMap<Location, BasicBlock>();
		bbOutEdges = HashMultimap.create();
		bbInEdges = HashMultimap.create();
	}
	
	
	public ControlFlowGraph(Set<CFAEdge> edges) {
		this();
		buildFromEdgeSet(edges);
	}
	
	public BasicBlock getBasicBlock(Location l) {
		return basicBlocks.get(l);
	}
	
	public Set<CFAEdge> getBasicBlockEdges() {
		return Collections.unmodifiableSet(
				new HashSet<CFAEdge>(bbOutEdges.values()));
	}

	public Set<Location> getBasicBlockNodes() {
		return Collections.unmodifiableSet(basicBlocks.keySet());
	}
	
	public Set<CFAEdge> getBasicBlockOutEdges(Location l) {
		return bbOutEdges.get(l);
	}
	
	public Map<Location, BasicBlock> getBasicBlocks() {
		return Collections.unmodifiableMap(basicBlocks);
	}
	
	public CFAEdge getEdgeBetween(Location src, Location tgt) {
		Set<CFAEdge> out = outEdges.get(src);
		if (out != null) for (CFAEdge e : out)
			if (e.getTarget().equals(tgt))
				return e;
		return null;
	}
	
	public Set<CFAEdge> getEdges() {
		return Collections.unmodifiableSet(
				new HashSet<CFAEdge>(outEdges.values()));
	}
	
	public Location getEntryPoint() {
		return entryPoint;
	}
	
	public int getInDegree(Location l) {
		return inEdges.get(l).size();
	}
	
	public Set<CFAEdge> getInEdges(Location l) {
		return Collections.unmodifiableSet(inEdges.get(l));
	}
	
	public Set<Location> getNodes() {
		return Collections.unmodifiableSet(locations);
	}
	
	public int getOutDegree(Location l) {
		return outEdges.get(l).size();
	}
	
	public Set<CFAEdge> getOutEdges(Location l) {
		return Collections.unmodifiableSet(outEdges.get(l));
	}
	
	public Set<Location> getSuccessorLocations(Location l) {
		Set<Location> res = new FastSet<Location>();
		for (CFAEdge e : outEdges.get(l))
			res.add(e.getTarget());
		return res;
	}
	
	public int numEdges() {
		return outEdges.size();
	}
	
	protected final void buildFromEdgeSet(Set<CFAEdge> edges) {
		for (CFAEdge e : edges) {
			addEdge(e);
		}
		
		findEntryPoint();		
		buildBasicBlocks();
		assert valid();
	}

	protected void addEdge(CFAEdge e) {
		
		if (e.getTransformer() instanceof BasicBlock) {
			BasicBlock bb = (BasicBlock)e.getTransformer();
			for (RTLStatement stmt : bb) {
				CFAEdge newEdge = new CFAEdge(stmt.getLabel(), stmt.getNextLabel(), stmt);
				addEdge(newEdge);
			}
			return;
		}
		
		// Check if edge already exists (equality on edges uses not only source & target)
		Set<CFAEdge> existingEdges = outEdges.get(e.getSource());
		for (CFAEdge existing : existingEdges) {
			if (existing.getTarget().equals(e.getTarget()))
				return;
		}
		
		outEdges.put(e.getSource(), e);
		inEdges.put(e.getTarget(), e);
		locations.add(e.getSource());
		locations.add(e.getTarget());
	}
	
	protected boolean isBasicBlockHead(Location l) {
		Set<CFAEdge> in = getInEdges(l);
		
		if (in.size() <= 0) {
			logger.debug("Orphan block at " + l);
			return true;
		}
		
		// If it has in-degree greater than 1, it's a head
		if (in.size() > 1) {
			return true;
		}
		
		// There's only one edge
		CFAEdge e = in.iterator().next();
		Location predLoc = e.getSource();

		// If out-degree of predecessor is greater than 1, this is a head
		if (getOutDegree(predLoc) > 1)
			return true;
		
		Program program = Program.getProgram();
		
		// If inEdge crosses into a stub, make this stub a new block
		if (!program.isStub(predLoc.getAddress()) &&
				program.isStub(l.getAddress())) {
			return true;
		}
		
		return false;		
	}
	
	/**
	 * Audit method for checking validity of the CFG.
	 * 
	 * @return true if the CFG is error free, false otherwise.
	 */
	protected boolean valid() {
		
		int errors = 0;
		
		for (CFAEdge e : outEdges.values()) {
			if (!locations.contains(e.getSource())) {
				logger.error("Locations do not contain " + e.getSource() + " from edge " + e);
				errors++;
			}
			if (!locations.contains(e.getTarget())) {
				logger.error("Locations do not contain " + e.getTarget() + " from edge " + e);
				errors++;
			}
			if (!outEdges.get(e.getSource()).contains(e)) {
				logger.error("Out-edges do not contain out-edge " + e + " with statement " + e.getTransformer() + " and hashcode " + e.hashCode());
				errors++;
			}
			if (!inEdges.get(e.getTarget()).contains(e)) {
				logger.error("In-edges do not contain out-edge " + e + " with statement " + e.getTransformer() + " and hashcode " + e.hashCode());
				errors++;
			}
		}
		
		for (CFAEdge e : inEdges.values()) {
			if (!locations.contains(e.getSource())) {
				logger.error("Locations do not contain " + e.getSource() + " from edge " + e);
				errors++;
			}
			if (!locations.contains(e.getTarget())) {
				logger.error("Locations do not contain " + e.getTarget() + " from edge " + e);
				errors++;
			}
			if (!outEdges.get(e.getSource()).contains(e)) {
				logger.error("Out-edges do not contain in-edge " + e + " with statement " + e.getTransformer() + " and hashcode " + e.hashCode());
				errors++;
			}
			if (!inEdges.get(e.getTarget()).contains(e)) {
				logger.error("In-edges do not contain in-edge " + e + " with statement " + e.getTransformer() + " and hashcode " + e.hashCode());
				errors++;
			}
		}
		
		Set<Location> bbLocations = basicBlocks.keySet();

		for (CFAEdge e : bbOutEdges.values()) {
			if (!bbLocations.contains(e.getSource())) {
				logger.error("Basicblock locations do not contain " + e.getSource() + " from edge " + e);
				errors++;
			}
			if (!bbLocations.contains(e.getTarget())) {
				logger.error("Basicblock locations do not contain " + e.getTarget() + " from edge " + e);
				errors++;
			}
			if (!bbInEdges.get(e.getTarget()).contains(e)) {
				logger.error("BB in-edges do not contain out-edge " + e);
				errors++;
			}
		}
		
		for (CFAEdge e : bbInEdges.values()) {
			if (!bbLocations.contains(e.getSource())) {
				logger.error("Basicblock locations do not contain " + e.getSource() + " from edge " + e);
				errors++;
			}
			if (!bbLocations.contains(e.getTarget())) {
				logger.error("Basicblock locations do not contain " + e.getTarget() + " from edge " + e);
				errors++;
			}
			if (!bbOutEdges.get(e.getSource()).contains(e)) {
				logger.error("BB out-edges do not contain in-edge " + e);
				errors++;
			}
		}
		if (errors != 0) {
			logger.error(errors + " errors in CFG audit.");
			return false;
		} else {
			return true;
		}
	}
	
	private void buildBasicBlocks() {
		
		Set<Location> basicBlockHeads = findBasicBlockHeads();
		
		for (Location l : basicBlockHeads)
			basicBlocks.put(l, new BasicBlock());
		
		for (Map.Entry<Location, BasicBlock> entry : basicBlocks.entrySet()) {
			Location head = entry.getKey();
			BasicBlock bb = entry.getValue();
			
			//bb.add(program.getStatement(head.getLabel()));
			Location l = head;
			Set<CFAEdge> out = outEdges.get(l);
			while (!out.isEmpty()) {
				
				CFAEdge edge = out.iterator().next();				

				// If there is more than one out edge, we'll break out
				if (out.size() > 1) {
					// Normally this is because of an assume - add the Goto to the BB instead of an assume
					if (edge.getTransformer() instanceof RTLAssume) {
						bb.add(((RTLAssume)edge.getTransformer()).getSource());
					} else if (edge.getTransformer() instanceof RTLCallReturn) {
						// Don't show fall-through edges of calls for now
					} else {
						// Multiple edges on a non-assume statement - this can happen with VPC CFGs where
						// the VPC value is modified directly. In this case, put the statement and break
						logger.verbose("Non-assume edge in outgoing set of edges with size > 1: " + edge.getTransformer());
						bb.add((RTLStatement)edge.getTransformer());
					}
					break;
				}

				bb.add((RTLStatement)edge.getTransformer());
				

				l = edge.getTarget();
				if (basicBlocks.containsKey(l))
					break;
				//bb.add(program.getStatement(l.getLabel()));
				out = outEdges.get(l);
			}
			// If there's no statement (because there's an immediate jump), add a skip
			if (bb.isEmpty()) {
				RTLStatement dummy = new RTLSkip();
				dummy.setLabel(l.getLabel());
				if (!out.isEmpty())
					dummy.setNextLabel(out.iterator().next().getTarget().getLabel());
				bb.add(dummy);
			}
			
			for (CFAEdge e : out) {
				RTLStatement edgeStmt;
				if (out.size() > 1 && e.getTransformer() instanceof RTLAssume) {
					edgeStmt = (RTLStatement)e.getTransformer();
				} else {
					edgeStmt = new RTLSkip();
					RTLStatement oldStmt = (RTLStatement)e.getTransformer();
					edgeStmt.setLabel(oldStmt.getLabel());
					edgeStmt.setNextLabel(oldStmt.getNextLabel());
				}
				CFAEdge bbEdge = new CFAEdge(head, e.getTarget(), edgeStmt);
				if (!basicBlocks.containsKey(e.getTarget())) {
					logger.error("Target not in basic block head list? " + bbEdge);
				} else {
					bbOutEdges.put(head, bbEdge);
					//assert basicBlocks.containsKey(e.getTarget()) : "Target not in basic block head list? " + bbEdge;
					bbInEdges.put(e.getTarget(), bbEdge);
				}
			}
		}
		
	}

	private Set<Location> findBasicBlockHeads() {
		Set<Location> result = new HashSet<Location>();
		
		// Find basic block heads
		for (Location l : getNodes()) {
			if (isBasicBlockHead(l))
				result.add(l);
		}
		logger.debug(result.size() + " basic blocks.");
		return result;
	}
	
	private void findEntryPoint() {
		for (Location l : getNodes()) {
			if (getInDegree(l) == 0) {
				assert entryPoint == null : "Graph has multiple entry points: " + entryPoint + " and " + l; 
				entryPoint = l;
			}
		}
		assert entryPoint != null : "No entry point found! First statement in cycle?";
	}
	
}
