/*BEGIN_LEGAL 
Intel Open Source License 

Copyright (c) 2002-2005 Intel Corporation 
All rights reserved. 
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.  Redistributions
in binary form must reproduce the above copyright notice, this list of
conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.  Neither the name of
the Intel Corporation nor the names of its contributors may be used to
endorse or promote products derived from this software without
specific prior written permission.
 
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE INTEL OR
ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
END_LEGAL */
/* ===================================================================== */
/*
  @ORIGINAL_AUTHOR: Robert Muth
*/

/* ===================================================================== */
/*! @file
 *  This file contains a static and dynamic opcode/ISA extension/ISA category mix profiler
 */

#include "pin.H"
#include "instlib.H"
#include "portability.H"
#include <vector>
#include <iostream>
#include <iomanip>
#include <fstream>


using namespace INSTLIB;

/* ===================================================================== */
/* Commandline Switches */
/* ===================================================================== */

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE,         "pintool",
    "o", "mix.out", "specify profile file name");
KNOB<BOOL>   KnobPid(KNOB_MODE_WRITEONCE,                "pintool",
    "i", "0", "append pid to output");
KNOB<BOOL>   KnobProfilePredicated(KNOB_MODE_WRITEONCE,  "pintool",
    "p", "0", "enable accurate profiling for predicated instructions");
KNOB<BOOL>   KnobProfileStaticOnly(KNOB_MODE_WRITEONCE,  "pintool",
    "s", "0", "terminate after collection of static profile for main image");
#ifndef TARGET_WINDOWS
KNOB<BOOL>   KnobProfileDynamicOnly(KNOB_MODE_WRITEONCE, "pintool",
    "d", "0", "Only collect dynamic profile");
#else
KNOB<BOOL>   KnobProfileDynamicOnly(KNOB_MODE_WRITEONCE, "pintool",
    "d", "1", "Only collect dynamic profile");
#endif
KNOB<BOOL>   KnobNoSharedLibs(KNOB_MODE_WRITEONCE,       "pintool",
    "no_shared_libs", "0", "do not instrument shared libraries");

KNOB<BOOL>   KnobOpcodeMix(KNOB_MODE_WRITEONCE,       "pintool",
    "opcode", "0", "Compute opcode mix (default)");
KNOB<BOOL>   KnobExtensionMix(KNOB_MODE_WRITEONCE,       "pintool",
    "extension", "0", "Compute ISA extension mix");
KNOB<BOOL>   KnobCategoryMix(KNOB_MODE_WRITEONCE,       "pintool",
    "category", "0", "Compute ISA category mix");


/* ===================================================================== */

INT32 Usage()
{
    cerr << "This pin tool computes a static and dynamic opcode, "
         << "extension or category mix profile\n\n";
    cerr << KNOB_BASE::StringKnobSummary();
    cerr << endl;
    cerr << "One of -category, -extension or -opcode is required" << endl;
    cerr << endl;
    return -1;
}

/* ===================================================================== */
/* INDEX HELPERS */
/* ===================================================================== */

const UINT32 MAX_INDEX = 4096;       // enough even for IPF
const UINT32 INDEX_SPECIAL =  3000;
const UINT32 MAX_MEM_SIZE = 520;


const UINT32 INDEX_TOTAL =          INDEX_SPECIAL + 0;
const UINT32 INDEX_MEM_ATOMIC =     INDEX_SPECIAL + 1;
const UINT32 INDEX_STACK_READ =     INDEX_SPECIAL + 2;
const UINT32 INDEX_STACK_WRITE =    INDEX_SPECIAL + 3;
const UINT32 INDEX_IPREL_READ =     INDEX_SPECIAL + 4;
const UINT32 INDEX_IPREL_WRITE =    INDEX_SPECIAL + 5;
const UINT32 INDEX_MEM_READ_VARIABLE =  INDEX_SPECIAL + 6;
const UINT32 INDEX_MEM_WRITE_VARIABLE =  INDEX_SPECIAL + 7;
const UINT32 INDEX_MEM_READ_SIZE =  INDEX_SPECIAL + 8;
const UINT32 INDEX_MEM_WRITE_SIZE = INDEX_SPECIAL + 8 + MAX_MEM_SIZE;
const UINT32 INDEX_SPECIAL_END   =  INDEX_SPECIAL + 8 + MAX_MEM_SIZE + MAX_MEM_SIZE;


BOOL IsMemReadIndex(UINT32 i)
{
    return (INDEX_MEM_READ_SIZE <= i && i < INDEX_MEM_READ_SIZE + MAX_MEM_SIZE );
}

BOOL IsMemWriteIndex(UINT32 i)
{
    return (INDEX_MEM_WRITE_SIZE <= i && i < INDEX_MEM_WRITE_SIZE + MAX_MEM_SIZE );
}


/* ===================================================================== */
LOCALFUN UINT32 INS_GetIndex(INS ins)
{
    UINT32 index = 0;
    if (KnobOpcodeMix.Value())
    {
        index = INS_Opcode(ins);
    }
    else if (KnobCategoryMix.Value())
    {
        index = INS_Category(ins);
    }
    else if (KnobExtensionMix.Value())
    {
        index =  INS_Extension(ins);
    }
    return index;
}


LOCALFUN UINT32 INS_GetIndexOffset(INS ins)
{
    UINT32 index = INS_GetIndex(ins);

    if( INS_IsPredicated(ins) )
        return MAX_INDEX + index;
    else
        return index;

}

/* ===================================================================== */

LOCALFUN  UINT32 IndexStringLength(BBL bbl, BOOL memory_acess_profile)
{
    UINT32 count = 0;

    for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins))
    {
        count++;
        if( memory_acess_profile )
        {
            if( INS_IsMemoryRead(ins) ) count++;   // for size

            if( INS_IsStackRead(ins) ) count++;

            if( INS_IsIpRelRead(ins) ) count++;

            
            if( INS_IsMemoryWrite(ins) ) count++; // for size

            if( INS_IsStackWrite(ins) ) count++;

            if( INS_IsIpRelWrite(ins) ) count++;

            
            if( INS_IsAtomicUpdate(ins) ) count++;
        }
    }
    
    return count;
}


/* ===================================================================== */
LOCALFUN UINT32 MemsizeToIndex(UINT32 size, BOOL write)
{
    if( size == VARIABLE_MEMORY_REFERENCE_SIZE )
    {
        return write ?  INDEX_MEM_WRITE_VARIABLE : INDEX_MEM_READ_VARIABLE;
    }
    else
    {
        return (write ? INDEX_MEM_WRITE_SIZE : INDEX_MEM_READ_SIZE ) + size;
    }
}

/* ===================================================================== */
LOCALFUN UINT16 *INS_GenerateIndexString(INS ins, UINT16 *stats, BOOL memory_acess_profile)
{
    *stats++ = INS_GetIndexOffset(ins);

    if( memory_acess_profile )
    {
        if( INS_IsMemoryRead(ins) )  *stats++ = MemsizeToIndex( INS_MemoryReadSize(ins), 0 );
        if( INS_IsMemoryWrite(ins) ) *stats++ = MemsizeToIndex( INS_MemoryWriteSize(ins), 1 );
        
        if( INS_IsAtomicUpdate(ins) ) *stats++ = INDEX_MEM_ATOMIC;
        
        if( INS_IsStackRead(ins) ) *stats++ = INDEX_STACK_READ;
        if( INS_IsStackWrite(ins) ) *stats++ = INDEX_STACK_WRITE;
        
        if( INS_IsIpRelRead(ins) ) *stats++ = INDEX_IPREL_READ;
        if( INS_IsIpRelWrite(ins) ) *stats++ = INDEX_IPREL_WRITE;
    }

    return stats;
}


/* ===================================================================== */

LOCALFUN string IndexToString( UINT32 index )
{
    if( INDEX_SPECIAL <= index  && index < INDEX_SPECIAL_END)
    {
        if( index == INDEX_TOTAL )            return  "*total";
        else if( IsMemReadIndex(index) )      return  "*mem-read-" + decstr( index - INDEX_MEM_READ_SIZE );
        else if( IsMemWriteIndex(index))      return  "*mem-write-" + decstr( index - INDEX_MEM_WRITE_SIZE );
        else if( index == INDEX_MEM_READ_VARIABLE )  return  "*mem-read-variable";
        else if( index == INDEX_MEM_WRITE_VARIABLE ) return  "*mem-write-variable";
        else if( index == INDEX_MEM_ATOMIC )  return  "*mem-atomic";
        else if( index == INDEX_STACK_READ )  return  "*stack-read";
        else if( index == INDEX_STACK_WRITE ) return  "*stack-write";
        else if( index == INDEX_IPREL_READ )  return  "*iprel-read";
        else if( index == INDEX_IPREL_WRITE ) return  "*iprel-write";

        else
        {
            ASSERTX(0);
            return "";
        }
    }
    else if (KnobOpcodeMix.Value())
    {
        return OPCODE_StringShort(index);
    }
    else if (KnobCategoryMix.Value())
    {
        return CATEGORY_StringShort(index);
    }
    else if (KnobExtensionMix.Value())
    {
        return EXTENSION_StringShort(index);
    }
    ASSERTX(0);
    return "";
    
}

/* ===================================================================== */
/* ===================================================================== */
typedef UINT64 COUNTER;


/* zero initialized */

class STATS
{
  public:
    COUNTER unpredicated[MAX_INDEX];
    COUNTER predicated[MAX_INDEX];
    COUNTER predicated_true[MAX_INDEX];

    VOID Clear()
    {
        for ( UINT32 i = 0; i < MAX_INDEX; i++)
        {
            unpredicated[i] = 0;
            predicated[i] = 0;
            predicated_true[i] = 0;
        }
    }
};


STATS GlobalStatsStatic;  // summary stats for static analysis
STATS GlobalStatsDynamic; // summary stats for dyanamic analysis

class BBLSTATS
{
    // each basic block has a counter. We update the individual stats
    // based on the block count. Our first pass sets up the types of stats
    // we need to update for this block. We have one stat per instruction
    // in the block. The _stats array is null terminated.
  public:
    COUNTER _counter;
    const UINT16 * const _stats;

  public:
    BBLSTATS(UINT16 * stats) : _counter(0), _stats(stats) {};

};



LOCALVAR vector<BBLSTATS*> statsList;



/* ===================================================================== */

LOCALVAR INT32 enabled = 0;

VOID activate_counting()
{
    enabled = 1;
}
VOID deactivate_counting()
{
    enabled = 0;
}

VOID emit_stats(); //forward prototype
VOID zero_stats(); //forward prototype

static std::ofstream* out = 0;

LOCALFUN VOID Handler(CONTROL_EVENT ev, VOID *val, CONTEXT * ctxt, VOID *ip, VOID * tid)
{
    switch(ev)
    {
      case CONTROL_START:
        *out << "# Start counting" << endl;
        activate_counting();
        break;
      case CONTROL_STOP:
        *out << "# Stop counting" << endl;
        deactivate_counting();
        break;
      default:
        ASSERTX(false);
    }
}




LOCALFUN VOID HandlerStats(CONTROL_STATS_EVENT ev, VOID *val, CONTEXT* dummy_context, VOID *ip, VOID * tid)
{
    switch(ev)
    {
      case CONTROL_STATS_EMIT:
        *out << "# Emit stats" << endl;
        emit_stats();
        break;
      case CONTROL_STATS_RESET:
        *out << "# Reset stats" << endl;
        zero_stats();
        break;
      default:
        ASSERTX(false);
    }
}

LOCALVAR CONTROL control;
LOCALVAR CONTROL_STATS control_stats;


/* ===================================================================== */

VOID docount(COUNTER * counter)
{
    (*counter) += enabled;
}

/* ===================================================================== */
VOID zero_stats()
{
    GlobalStatsDynamic.Clear();

    for (vector<BBLSTATS*>::iterator bi = statsList.begin(); bi != statsList.end(); bi++)
    {
        BBLSTATS *b = *bi;
        if (b) 
        {
            b->_counter = 0;
        }
    }
}
/* ===================================================================== */

VOID CheckForSpecialMarkers(INS ins, ADDRINT pc, unsigned int instruction_size)
{
    // This checks for single instances of special 3B NOPs.  
    // 0F1FF3 - start
    // 0F1FF4 - stop
    // 0F1FF5 - emit stats
    // 0F1FF6 - zero stats

    // FIXME: if there are collisions with existing instructions, we can
    // change them here.

    //FIXME: Ideally this would be integrated in to the control.H so file
    //so that anything can use it.

    UINT8* pc_ptr = reinterpret_cast<UINT8*>(pc);
    if (pc_ptr[0] == 0x0F &&
        pc_ptr[1] == 0x1F)
    {
        switch(pc_ptr[2])
        {
          case 0xF3: // start
            INS_InsertCall(ins, 
                           IPOINT_BEFORE, 
                           (AFUNPTR)activate_counting,
                           IARG_END);
            break;
          case 0xF4: // stop
            INS_InsertCall(ins, 
                           IPOINT_BEFORE, 
                           (AFUNPTR)deactivate_counting,
                           IARG_END);
            break;
          case 0xF5: // emit
            INS_InsertCall(ins, 
                           IPOINT_BEFORE, 
                           (AFUNPTR)emit_stats,
                           IARG_END);
            break;
          case 0xF6: // zero
            INS_InsertCall(ins, 
                           IPOINT_BEFORE, 
                           (AFUNPTR)zero_stats,
                           IARG_END);
            break;
          default:
            break;
        }
    }
}

/* ===================================================================== */

VOID Trace(TRACE trace, VOID *v)
{
    if ( KnobNoSharedLibs.Value()
         && IMG_Type(SEC_Img(RTN_Sec(TRACE_Rtn(trace)))) == IMG_TYPE_SHAREDLIB)
        return;
    
    const BOOL accurate_handling_of_predicates = KnobProfilePredicated.Value();
    ADDRINT pc = TRACE_Address(trace);
    for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl))
    {
        const INS head = BBL_InsHead(bbl);
        if (! INS_Valid(head)) continue;

        // Summarize the stats for the bbl in a 0 terminated list
        // This is done at instrumentation time
        const UINT32 n = IndexStringLength(bbl, 1);

        // stats is an array of index types. We later multiply it by the
        // dynamic count for a block.
        UINT16 *const stats = new UINT16[ n + 1];
        UINT16 *const stats_end = stats + (n + 1);
        UINT16 *curr = stats;
        
        for (INS ins = head; INS_Valid(ins); ins = INS_Next(ins))
        {
            unsigned int instruction_size = INS_Size(ins);
            CheckForSpecialMarkers(ins, pc, instruction_size);

            // Count the number of times a predicated instruction is actually executed
            // this is expensive and hence disabled by default
            if( INS_IsPredicated(ins) && accurate_handling_of_predicates )
            {
                INS_InsertPredicatedCall(ins,
                                         IPOINT_BEFORE,
                                         AFUNPTR(docount),
                                         IARG_PTR, 
                                         &(GlobalStatsDynamic.predicated_true[INS_GetIndex(ins)]),
                                         IARG_END);    
            }

            curr = INS_GenerateIndexString(ins,curr,1);
            pc = pc + instruction_size;
        }
        
        // stats terminator
        *curr++ = 0;
        
        ASSERTX( curr == stats_end );

        
        // Insert instrumentation to count the number of times the bbl is executed
        BBLSTATS * bblstats = new BBLSTATS(stats);
        INS_InsertCall(head,
                       IPOINT_BEFORE, 
                       AFUNPTR(docount), 
                       IARG_PTR,
                       &(bblstats->_counter), 
                       IARG_END);

        // Remember the counter and stats so we can compute a summary at the end
        statsList.push_back(bblstats);
    }
}

/* ===================================================================== */
VOID DumpStats(ofstream& out, STATS& stats, BOOL predicated_true,  const string& title)
{
    out << "#\n# "
        << title 
        << "\n#\n"
        << "#     ";
    if (KnobOpcodeMix.Value())
        out << "opcode";
    else if (KnobCategoryMix.Value())
        out << "catgry";
    else if (KnobExtensionMix.Value())
        out << "extnsn";
    out<< "       count-unpredicated    count-predicated";

    if( predicated_true )
        out << "    count-predicated-true";

    out << "\n#\n";

    // Compute the "total" bin.
    for ( UINT32 i = 0; i < INDEX_TOTAL; i++)
    {
        stats.unpredicated[INDEX_TOTAL] += stats.unpredicated[i];
        stats.predicated[INDEX_TOTAL] += stats.predicated[i];
        stats.predicated_true[INDEX_TOTAL] += stats.predicated_true[i];
    }
    
    UINT32 limit = MAX_INDEX;
    // only print the special values when doing opcode mix
    if (KnobExtensionMix.Value() || KnobCategoryMix.Value())
    {
        limit = INDEX_TOTAL + 1;
    }
    for ( UINT32 i = 0; i < limit; i++)
    {
        if( stats.unpredicated[i] == 0 &&
            stats.predicated[i] == 0 ) continue;
        
        out << setw(4) << i << " " <<  ljstr(IndexToString(i),15) << " " <<
            setw(16) << stats.unpredicated[i] << " " <<
            setw(16) << stats.predicated[i];
        if( predicated_true ) out << " " << setw(16) << stats.predicated_true[i];
        out << endl;
    }
}



/* ===================================================================== */

VOID emit_stats()
{
    static UINT32 stat_dump_count = 0;

    stat_dump_count++;

    *out << "# EMIT_STATS " << stat_dump_count << endl;

    if (stat_dump_count == 1)
    {
        // static counts -- only dump the static stats the first time around
        DumpStats(*out, GlobalStatsStatic, false, "$static-counts");
        *out << endl;
    }

    // dynamic Counts 
    //statsList.push_back(0); // add terminator marker

    for (vector<BBLSTATS*>::iterator bi = statsList.begin(); bi != statsList.end(); bi++)
    {
        const BBLSTATS *b = (*bi);

        if ( b == 0 ) 
            continue;
        
        for (const UINT16 * stats = b->_stats; *stats; stats++)
        {
            GlobalStatsDynamic.unpredicated[*stats] += b->_counter;
        }
    }

    DumpStats(*out, GlobalStatsDynamic, KnobProfilePredicated, "$dynamic-counts");                
    *out << "# END_STATS" <<  endl;

}

/* ===================================================================== */

VOID Fini(int, VOID * v)
{
    *out << "# FINI: end of program" << endl;
    emit_stats();
    out->close();
}


/* ===================================================================== */

VOID Image(IMG img, VOID * v)
{
    for (SEC sec = IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec))
    {
        for (RTN rtn = SEC_RtnHead(sec); RTN_Valid(rtn); rtn = RTN_Next(rtn))
        {
            // Prepare for processing of RTN, an  RTN is not broken up into BBLs,
            // it is merely a sequence of INSs 
            RTN_Open(rtn);
            
            for (INS ins = RTN_InsHead(rtn); INS_Valid(ins); ins = INS_Next(ins))
            {
                UINT16 array[128];
                UINT16 *end  = INS_GenerateIndexString(ins,array,1);

                if( INS_IsPredicated(ins) )
                {
                    for( UINT16 *start= array; start < end; start++) 
                    {
                        GlobalStatsStatic.predicated[ *start ]++;
                    }
                }
                else
                {
                    for( UINT16 *start= array; start < end; start++) 
                    {
                        GlobalStatsStatic.unpredicated[ *start ]++;
                    }
                }
            }

            // to preserve space, release data associated with RTN after we have processed it
            RTN_Close(rtn);
        }
    }

    if( KnobProfileStaticOnly.Value() )
    {
        Fini(0,0);
        exit(0);
    }
}

    
/* ===================================================================== */

int main(int argc, CHAR *argv[])
{
    PIN_InitSymbols();
    
    if( PIN_Init(argc,argv) )
    {
        return Usage();
    }

    string filename =  KnobOutputFile.Value();
    if( KnobPid )
    {
        filename += "." + decstr( getpid_portable() );
    }
    out = new std::ofstream(filename.c_str());
    
    control.CheckKnobs(Handler, 0);
    control_stats.CheckKnobs(HandlerStats, 0);

    // make sure that exactly one thing-to-count knob is specified.
    UINT32 kinds = 0;
    kinds += KnobOpcodeMix.Value();
    kinds += KnobCategoryMix.Value();
    kinds += KnobExtensionMix.Value();
    if (kinds != 1)
    {
        cerr << "Must have just one of: -category, -opcode or -extension as a pintool option"
             << endl;
        exit(1);
    }
    
    TRACE_AddInstrumentFunction(Trace, 0);

    PIN_AddFiniFunction(Fini, 0);

    if( !KnobProfileDynamicOnly.Value() )
        IMG_AddInstrumentFunction(Image, 0);

    // Never returns

    PIN_StartProgram();
    
    return 0;
}

/* ===================================================================== */
/* eof */
/* ===================================================================== */