/*
 * DFS.h
 *
 *  Created on: Nov 5, 2015
 *      Author: yancai
 */

#ifndef LIB_TRANSFORMS_DFIXER_DFS_H_
#define LIB_TRANSFORMS_DFIXER_DFS_H_

#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SCCIterator.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Support/CommandLine.h"

#include "llvm/IR/CFG.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;
using namespace std;
using namespace cl;

enum Color
{
	WHITE, GREY, BLACK
};

// Color marks per vertex (BB).
typedef DenseMap<const BasicBlock *, Color> BBColorMap;
// Collects vertices (BBs) in "finish" order. The first finished vertex is
// first, and so on.
typedef SmallVector<const BasicBlock *, 32> BBVector;

//DFS search
class TopoSorter
{
	public:
		void runToposort(const llvm::Function &F)
		{
			outs() << "Topological sort of " << F.getName() << ":\n";
			// Initialize the color map by marking all the vertices white.
			for (Function::const_iterator I = F.begin(), IE = F.end(); I != IE;
					++I)
			{
				ColorMap[I] = WHITE;
			}

			// The BB graph has a single entry vertex from which the other BBs should
			// be discoverable - the function entry block.
			bool success = recursiveDFSToposort(&F.getEntryBlock());
			if (success)
			{
				// Now we have all the BBs inside SortedBBs in reverse topological order.
				for (BBVector::const_reverse_iterator RI = SortedBBs.rbegin(),
						RE = SortedBBs.rend(); RI != RE; ++RI)
				{
					outs() << " " << (*RI)->getName() << "\n";
				}
			}
			else
			{
				outs() << " Sorting failed\n";
			}
		}
	private:

		BBColorMap ColorMap;
		BBVector SortedBBs;

		// Helper function to recursively run topological sort from a given BB.
		// Returns true if the sort succeeded and false otherwise; topological sort
		// may fail if, for example, the graph is not a DAG (detected a cycle).
		bool recursiveDFSToposort(const BasicBlock *BB)
		{
			ColorMap[BB] = GREY;
			// For demonstration, using the lowest-level APIs here. A BB's successors
			// are determined by looking at its terminator instruction.
			const TerminatorInst *TInst = BB->getTerminator();
			for (unsigned I = 0, NSucc = TInst->getNumSuccessors(); I < NSucc;
					++I)
			{
				BasicBlock *Succ = TInst->getSuccessor(I);
				Color SuccColor = ColorMap[Succ];
				if (SuccColor == WHITE)
				{
					if (!recursiveDFSToposort(Succ))
						return false;
				}
				else if (SuccColor == GREY)
				{
					// This detects a cycle because grey vertices are all ancestors of the
					// currently explored vertex (in other words, they're "on the stack").
					outs() << " Detected cycle: edge from " << BB->getName()
							<< " to " << Succ->getName() << "\n";
					return false;
				}
			}
			// This BB is finished (fully explored), so we can add it to the vector.
			ColorMap[BB] = BLACK;
			SortedBBs.push_back(BB);
			return true;
		}
};


class DeepDFS
{
		bool (*isTargetInst) (llvm::CallInst*);

	public:
		void runToposort(const llvm::Function &F)
		{
			outs() << "Topological sort of " << F.getName() << ":\n";
			// Initialize the color map by marking all the vertices white.
			for (Function::const_iterator I = F.begin(), IE = F.end(); I != IE;
					++I)
			{
				ColorMap[I] = WHITE;
			}

			// The BB graph has a single entry vertex from which the other BBs should
			// be discoverable - the function entry block.
			bool success = recursiveDFSToposort(&F.getEntryBlock());
			if (success)
			{
				// Now we have all the BBs inside SortedBBs in reverse topological order.
				for (BBVector::const_reverse_iterator RI = SortedBBs.rbegin(),
						RE = SortedBBs.rend(); RI != RE; ++RI)
				{
					outs() << " " << (*RI)->getName() << "\n";
				}
			}
			else
			{
				outs() << " Sorting failed\n";
			}
		}
	private:

		BBColorMap ColorMap;
		BBVector SortedBBs;

		// Helper function to recursively run topological sort from a given BB.
		// Returns true if the sort succeeded and false otherwise; topological sort
		// may fail if, for example, the graph is not a DAG (detected a cycle).
		bool recursiveDFSToposort(const BasicBlock *BB)
		{
			ColorMap[BB] = GREY;
			// For demonstration, using the lowest-level APIs here. A BB's successors
			// are determined by looking at its terminator instruction.
			const TerminatorInst *TInst = BB->getTerminator();
			for (unsigned I = 0, NSucc = TInst->getNumSuccessors(); I < NSucc;
					++I)
			{
				BasicBlock *Succ = TInst->getSuccessor(I);
				Color SuccColor = ColorMap[Succ];
				if (SuccColor == WHITE)
				{
					if (!recursiveDFSToposort(Succ))
						return false;
				}
				else if (SuccColor == GREY)
				{
					// This detects a cycle because grey vertices are all ancestors of the
					// currently explored vertex (in other words, they're "on the stack").
					outs() << " Detected cycle: edge from " << BB->getName()
							<< " to " << Succ->getName() << "\n";
					return false;
				}
			}
			// This BB is finished (fully explored), so we can add it to the vector.
			ColorMap[BB] = BLACK;
			SortedBBs.push_back(BB);
			return true;
		}
};



#endif /* LIB_TRANSFORMS_DFIXER_DFS_H_ */
