
#include <iostream>
#include <fstream>
#include <string.h>
#include <list>

#include "pin.H"

using namespace std;

ofstream OutFile;

bool main_flag;

// The running count of instructions is kept here
// make it static to help the compiler optimize icount
static UINT64 icount = 0;

struct sprof 
{
	UINT64 instance;
	UINT64 count;
	UINT32 level;
	char* name;	
};

list<sprof> mylist;

sprof func1;

///////////////////////////

list<char*> call_stack;

list<UINT64> ins_stack;

///////////////////////////////////////////////////////////////////////////////////////////////////////////
//--------------------------------------- RUN Time Functions ------------------------------//
///////////////////////////////////////////////////////////////////////////////////////////////////////////

void PIN_FAST_ANALYSIS_CALL docount(INT32 c) 
{ 
	icount += c; 
}

//--------------------------------------------------------------------------------------------------------//

VOID set_flag( UINT32 flag)
{
	main_flag = flag;
}

//--------------------------------------------------------------------------------------------------------//

VOID fun_prof_before(VOID* name) 
{ 
	if(main_flag == true)
		{
			call_stack.push_back((char*)name);
			ins_stack.push_back(icount);
		}
}

//-------------------------------------------------------------------------------------------------------//

VOID fun_prof_after( )
{
	if(main_flag == true)
		{
			PIN_LockClient();
			
			func1.count = (icount-ins_stack.back());	
			func1.level = call_stack.size();
			func1.name = call_stack.back();
			
			mylist.push_back(func1);
					
			call_stack.pop_back(); 
			ins_stack.pop_back(); 

			PIN_UnlockClient();
		}
}
    
///////////////////////////////////////////////////////////////////////////////////////////////////////////
//--------------------------------------- Analysis Functions ---------------------------------//
///////////////////////////////////////////////////////////////////////////////////////////////////////////

// Pin Callback
void Trace(TRACE trace, void *v)
{
	
for( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
{
	BBL_InsertCall(bbl, IPOINT_ANYWHERE, (AFUNPTR)docount, IARG_FAST_ANALYSIS_CALL, IARG_UINT32, BBL_NumIns(bbl),IARG_END);
}

}

//----------------------------------------------------------------------------------------------------------
// Pin calls this function every time a new routine is encountered
VOID Check_Main( RTN rtn, VOID *v)
{	
	
	if ( RTN_Valid(rtn) )
		{
			
			RTN_Open(rtn);
				
			// Insert a call to docount before every instruction, no arguments are passed
			if( strcmp(RTN_Name(rtn).c_str(),"main") == 0 )
				{  
					RTN_InsertCall(rtn,IPOINT_BEFORE,(AFUNPTR)set_flag,IARG_UINT32,1,IARG_END); 
					RTN_InsertCall(rtn,IPOINT_AFTER,(AFUNPTR)set_flag,IARG_UINT32,0,IARG_END);
				}
			else
				{
					RTN_InsertCall(rtn,IPOINT_BEFORE,(AFUNPTR)fun_prof_before,IARG_PTR, RTN_Name(rtn).c_str(), IARG_END);
					RTN_InsertCall(rtn,IPOINT_AFTER,(AFUNPTR)fun_prof_after, IARG_END);
				}
				       	
			RTN_Close(rtn);
		}

}

//----------------------------------------------------------------------------------------------------------

//////////////////////////////////////////////////////////////////////////////////////////////////////////


KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool","o", "sprofile3.out", "specify output file name");

bool compareLevel( const sprof& one,  const sprof& two )
{
	if( one.level > two.level)  
	{
		return true;
	}
	else 
	{	
		return false;
	}
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////

VOID check_iterations( )
{
	list<sprof> :: iterator it1,it2;
	
	for ( it1=mylist.begin(); it1!=mylist.end() ;  ++it1 )
	{		
		for ( it2=next(it1,1); it2!=mylist.end() ;  ++it2)
			{
				if( (strcmp(it1->name, it2->name) == 0)  && (it1->level == it2->level) )
				{
					it1-> instance++;
					it1->count += it2->count;
					it2 = mylist.erase(it2);
					--it2;
				}					
			}
	}
}

// This function is called when the application exits
VOID Fini(INT32 code, VOID *v)
{
	list<sprof> ::iterator it;
	
	for ( it=mylist.begin(); it!=mylist.end() ; ++it )
	{
		it-> instance = 1;
	}
	
	check_iterations();
	
	
	
	mylist.sort(compareLevel);
	
	OutFile << "Level" << "\t" << " Instants " << "\t" << "Instructions" << "\t\t"    << "Function"  << endl;
	
	OutFile << endl;
	
	for ( it=mylist.begin(); it!=mylist.end() ; ++it )
	{
		OutFile << it->level   << "\t"  << it->instance << "\t" "\t" << it->count  << "\t\t\t"   << it-> name << endl;
	}
    
	// Write to a file since cout and cerr maybe closed by the application
	OutFile.setf(ios::showbase);
	OutFile << endl;
	OutFile << endl;
	OutFile << icount << "\t" "\t" "\t"  << " Total Instruction Count "  << endl;
	OutFile.close();
}

/* ===================================================================== */
/* Print Help Message                                                    */
/* ===================================================================== */

INT32 Usage()
{
    cerr << "This tool counts the number of dynamic instructions executed" << endl;
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

/* ===================================================================== */
/* 										Main			                                                                 */
/* ===================================================================== */
/*   argc, argv are the entire command line: pin -t <toolname> -- ...    */
/* ===================================================================== */

int main(int argc, char * argv[])
{
    PIN_InitSymbols();

    // Initialize pin
    if (PIN_Init(argc, argv)) return Usage();

    OutFile.open(KnobOutputFile.Value().c_str());
	
    // Register Instruction to be called to instrument routines
    RTN_AddInstrumentFunction(Check_Main,0);  

    // Register Instruction to be called to instrument instructions
    TRACE_AddInstrumentFunction(Trace, 0);

    // Register Fini to be called when the application exits
    PIN_AddFiniFunction(Fini, 0);
    
    // Start the program, never returns
    PIN_StartProgram();
    
    return 0;
}
