

#include <iostream>
#include <fstream>
#include <string.h>
#include "pin.H"
#include <list>

using namespace std;

ofstream OutFile;

bool main_flag;

bool mflag;

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

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

list<sprof> mylist;

sprof func1;

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

list<char*> call_stack;
list<UINT64> ins_stack;

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

VOID docount() 
{ 

if(main_flag == true)
{ 
	icount++;
}

}


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

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();
 
 mflag = 0;
 
 list<sprof> ::iterator it;

 for ( it=mylist.begin(); it!=mylist.end() ; ++it)
   { 
     if ( (strcmp( it->name, call_stack.back() ) == 0)  &&  ( it->level ==  call_stack.size() ) )
	{
	  it->instance++;
	  it->count += (icount-ins_stack.back());
          mflag = 1;	
	}
   }
   
   if (mflag == 0)
   {
	func1.name =  call_stack.back();
	func1.instance = 1;
	func1.count = (icount-ins_stack.back());	
	func1.level = call_stack.size();
	mylist.push_back(func1);
	    
   }

   
call_stack.pop_back(); 
ins_stack.pop_back(); 

PIN_UnlockClient();
}

}
    
///////////////////////////////////////////////////////////////////////////////////////////////////////////
//--------------------------------------- Analysis Functions --------------------------------------------//
///////////////////////////////////////////////////////////////////////////////////////////////////////////


// Pin calls this function every time a new instruction is encountered
VOID Instruction(INS ins, VOID *v)
{
 // Insert a call to docount before every conditional Branch instruction, no arguments are passed
 INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)docount, IARG_END);
 
}


//----------------------------------------------------------------------------------------------------------
// Pin calls this function every time a new routine is encountered
VOID Check_Main(RTN rtn, VOID *v)
{
    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 * x = new RTN(rtn);
	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", "sprofile.out", "specify output file name");


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

// This function is called when the application exits
VOID Fini(INT32 code, VOID *v)
{

    list<sprof> ::iterator it;
	
    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 routines
    //RTN_AddInstrumentFunction(function_Profiler,0); 	

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

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