/* bbv - generates Simpoint Basic Block Vectors               */
/* loosely based on the cachegrind plugin                     */

/*   This program is free software; you can redistribute it and/or
     modify it under the terms of the GNU General Public License as
     published by the Free Software Foundation; either version 2 of the
     License, or (at your option) any later version.
  
     This program 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 for more details.
  
     You should have received a copy of the GNU General Public License
     along with this program; if not, write to the Free Software
     Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
     02111-1307, USA.
  
     The GNU General Public License is contained in the file COPYING.
*/

/* original code by Vince Weaver, < vince _at_ csl.cornell.edu > */
/* pcfile code by Oriol Prat, < oriol.prat _at _ bsc.es >        */

#include "pub_tool_basics.h"
#include "pub_tool_tooliface.h"
#include "pub_tool_options.h"    /* command line options */

#include "pub_tool_vki.h"        /* vki_stat */
#include "pub_tool_libcbase.h"   /* VG_(strlen) */
#include "pub_tool_libcfile.h"   /* VG_(write) */
#include "pub_tool_libcprint.h"  /* VG_(printf) */
#include "pub_tool_libcassert.h" /* VG_(exit) */
#include "pub_tool_mallocfree.h" /* plain_free */
#include "pub_tool_machine.h"    /* VG_(fnptr_to_fnentry) */
#include "pub_tool_debuginfo.h"  /* VG_(get_fnname) */

#include "pub_tool_oset.h"       /* ordered set stuff */

#define DEFAULT_GRAIN_SIZE 100000000  /* 100 million by default */

/* instruction special cases */
#define REP_INSTRUCTION   0x1
#define FLDCW_INSTRUCTION 0x2

static Int interval_size=DEFAULT_GRAIN_SIZE;
static ULong dyn_instr=0,total_instr=0,global_rep_count=0,
       unique_rep_count=0,fldcw_count=0;
static Int interval_num=1;
static Bool new_interval=True;
static UChar buf[100];

static UChar bbFileName[256];
static UChar pcFileName[256];

static Int bbtrace_fd;
static Int block_num=1;

static Bool instr_count_only=False;

struct BB_info {
   Addr       BB_addr;           /* key, must be first */
   Int        n_instrs;
   Int        block_num;
   Int        bb_iterations;
   Int        bb_inst_counter;
   Int        bb_proc_dyn_instr;
   Bool       is_entry;
   UChar      fn_name[20];
};

static OSet* instr_info_table;
static Int  distinct_instrs=0;
static Int rep_count=0;

static void dumpPcFile(void) {
   struct BB_info   *bb_elem;
   Int       pctrace_fd;
   SysRes    sres;
   
   sres = VG_(open)(pcFileName, VKI_O_CREAT|VKI_O_TRUNC|VKI_O_WRONLY,
		              VKI_S_IRUSR|VKI_S_IWUSR|VKI_S_IRGRP|VKI_S_IWGRP);
   if (sres.isError) {
      VG_(printf)("error: opening %s\n", pcFileName);
      VG_(exit)(1);
   } else {
      VG_(printf)("Opened %s...\n", pcFileName);
      pctrace_fd = sres.res;
   }

   VG_(OSetGen_ResetIter)(instr_info_table);
   while ( (bb_elem = VG_(OSetGen_Next)(instr_info_table)) ) {
         VG_(write)(pctrace_fd,"F",1);
         VG_(sprintf)( buf,":%d:%x:%s\n", 
                       bb_elem->block_num, 
                       bb_elem->BB_addr,
                       bb_elem->fn_name);
         VG_(write)(pctrace_fd, (void*)buf, VG_(strlen)(buf));      
   }
   
   VG_(close)(pctrace_fd);
}


static VG_REGPARM(2) void updateBBV(struct BB_info *bbInfo,
				    Int insn_type) {
				    
    struct BB_info *bb_elem;
    Int n_instrs=1;
   
    tl_assert(bbInfo);

      /* Handle the rep prefix on the x86 architecture */
    if (insn_type&REP_INSTRUCTION) {
       rep_count++;
       n_instrs=0;
    }
   
       /* we finished rep but didn't clear out count */
    if (rep_count && (!(insn_type&REP_INSTRUCTION))) {
       n_instrs=2;
       global_rep_count+=rep_count;
       unique_rep_count++;
       rep_count=0;
    }

       /* count fldcw instructions */
    if (insn_type&FLDCW_INSTRUCTION) {
       fldcw_count++;  
    }

    bbInfo->bb_inst_counter+=n_instrs;
    bbInfo->bb_iterations++;

    total_instr+=n_instrs;
    dyn_instr +=n_instrs;

    if ( dyn_instr > interval_size) {
       new_interval = True;
	
       if (!instr_count_only) {
          /* put in an entry to the out.bb file */

          VG_(write)(bbtrace_fd,"T",1);
       
          VG_(OSetGen_ResetIter)(instr_info_table);
          while ( (bb_elem = VG_(OSetGen_Next)(instr_info_table)) ) {
	     if ( bb_elem->bb_inst_counter != 0 ) {
       	        VG_(sprintf)( buf,":%d:%d   ", 
			 bb_elem->block_num, 
			 bb_elem->bb_inst_counter);
	        VG_(write)(bbtrace_fd, (void*)buf, VG_(strlen)(buf));
	        bb_elem->bb_inst_counter = 0;
	     }
          }

          VG_(write)(bbtrace_fd,"\n",1);
       }
       
       dyn_instr -= interval_size;
       interval_num++;
    }
}

static Int get_inst(Int len,Addr addr) {
   
    int result=0;
   
#if defined(VGP_x86_linux) || defined(VGP_amd64_linux)

    unsigned char *first_byte,*second_byte;
    
    if (len>1) {
       
       first_byte=(unsigned char *)addr;
       second_byte=first_byte+1;
	  
	  /* rep */
       if ((*first_byte==0xf3) || (*first_byte==0xf2)) {
	  result|=REP_INSTRUCTION;
       }

          /* fldcw */
          /* opcode is 0xd9/5, ie 1101 1001 oo10 1mmm */
       if ((*first_byte==0xd9) &&
	   (*second_byte<0xb0) &&  /* need this case of fldz, etc, count */
	   ( (*second_byte & 0x38) == 0x28)) {
          result|=FLDCW_INSTRUCTION;
/*	  VG_(printf)("Found fldcw at %x\n",addr); */
       }
    }
#endif
    return result;
}


   /* sbIn = super block to translate */
   /* layout = guest layout           */
   /* gWordTy = size of guest word    */
   /* hWordTy = size of host word     */

static IRSB* bbv_instrument ( VgCallbackClosure* closure,
			     IRSB* sbIn, VexGuestLayout* layout,
			     VexGuestExtents* vge,
			     IRType gWordTy, IRType hWordTy ) {
   
    Int      current_instr,n_instrs=1;
    IRSB     *sbOut;
    IRStmt   *st;
    struct BB_info  *bbInfo;
    Addr64   origAddr,ourAddr;
    IRDirty  *di;
    IRExpr   **argv, *arg1, *arg2;
    Int      argc,inst_opcode;

       /* We don't handle a host/guest word size mismatch */
    if (gWordTy != hWordTy) {
       VG_(tool_panic)("host/guest word size mismatch");
    }
   
       /* Set up SB */
    sbOut = deepCopyIRSBExceptStmts(sbIn);
  
       /* Copy verbatim any IR preamble preceding the first IMark */
    current_instr = 0;
    while ( (current_instr < sbIn->stmts_used) && 
	    (sbIn->stmts[current_instr]->tag!=Ist_IMark)) {
       addStmtToIRSB( sbOut, sbIn->stmts[current_instr] );
       current_instr++;
    }
       
       /* Get the first statement */
    tl_assert(sbIn->stmts_used > 0);
    st = sbIn->stmts[current_instr];
   
       /* double check we are at a Mark statement */
    tl_assert(Ist_IMark == st->tag);
        
    origAddr=st->Ist.IMark.addr;


       /* Get the BB_info */
    bbInfo = VG_(OSetGen_Lookup)(instr_info_table, &origAddr);

    if (bbInfo==NULL) {
       bbInfo=VG_(OSetGen_AllocNode)(instr_info_table, sizeof(struct BB_info));

      /* BB never translated before (at this address, at least;  could have */
      /* been unloaded and then reloaded elsewhere in memory)               */
      bbInfo->BB_addr = origAddr;
      bbInfo->n_instrs = n_instrs;
      bbInfo->block_num=block_num;
      VG_(get_fnname)(origAddr,bbInfo->fn_name,20);
      bbInfo->is_entry=VG_(get_fnname_if_entry)(origAddr, bbInfo->fn_name,20);
       
      VG_(OSetGen_Insert)( instr_info_table, bbInfo );

      distinct_instrs++;
      block_num++;
   }
   

   
       /* Iterate through the basic block, putting the original   */
       /* instructions in place, plus putting a call to updateBBV */
       /* for each original instruction                           */

       /* This is less efficient than only instrumenting the BB   */
       /* But it gives proper results given the fact that         */
       /* valgrind uses superblocks (not basic blocks) by default */


    while(current_instr < sbIn->stmts_used) {
       st=sbIn->stmts[current_instr];

       if (st->tag == Ist_IMark) {
	     
	  ourAddr = st->Ist.IMark.addr;

	  inst_opcode=get_inst(st->Ist.IMark.len,ourAddr);

	     /* Create call to our instrumentation function to insert */
	     /* It has two arguments, the original address of the basic block */
	     /* and an argument spefcifying if it is a *rep* instruction      */
	  argc=2;
	  arg1= mkIRExpr_HWord( (HWord)bbInfo);
          arg2= mkIRExpr_HWord(inst_opcode);
          argv= mkIRExprVec_2(arg1,arg2);
          di= unsafeIRDirty_0_N( argc, "updateBBV", 
				 VG_(fnptr_to_fnentry)( &updateBBV ), argv);
	  
	     /* Insert our call */
          addStmtToIRSB( sbOut,  IRStmt_Dirty(di));	  
       }
       
	  /* Insert the original instruction */
       addStmtToIRSB( sbOut, st );
            
       current_instr++;
   }
   
   return sbOut;
}


/*--------------------------------------------------------------------*/
/*--- Setup                                                        ---*/
/*--------------------------------------------------------------------*/

static void bbv_post_clo_init(void) {
   
   SysRes  sres;
   
   sres = VG_(open)(bbFileName, VKI_O_CREAT|VKI_O_TRUNC|VKI_O_WRONLY,
		              VKI_S_IRUSR|VKI_S_IWUSR|VKI_S_IRGRP|VKI_S_IWGRP);
   if (sres.isError) {
      VG_(printf)("error: opening %s\n",bbFileName);
      VG_(exit)(1);
   } else {
      VG_(printf)("Opened %s...\n",bbFileName);
      bbtrace_fd=sres.res;
   }         
}



static void bbv_thread_called ( ThreadId tid, ULong nDisp ) {
   
   static int old_thread=0;
   
   if (tid!=old_thread) {
      VG_(printf)("Entered thread %d\n",tid);
      old_thread=tid;
   }
}


   /* Parse the command line options */
static Bool bbv_process_cmd_line_option(Char* arg) {
   
        /* 9 is length of "--bbfile=" */
    if (VG_CLO_STREQN(9, arg, "--bbfile=")) {
       VG_(sprintf)(bbFileName,"%s",&arg[9]);
    }
    else if (VG_CLO_STREQN(9, arg, "--pcfile=")) {
       VG_(sprintf)(pcFileName,"%s",&arg[9]);
    }
    else if (VG_CLO_STREQN(16,arg,"--interval_size=")) {
       interval_size=(int)VG_(atoll)(&arg[16]);
    }
    else if (VG_CLO_STREQN(18,arg,"--instr_count_only")) {
       instr_count_only=True;
    }
    else {
       return False;
    }
    return True;
}


static void bbv_print_usage(void) {
    VG_(printf) ("   --bbfile=<file>       filename to put basic block info into\n");
    VG_(printf) ("   --pcfile=<file>       filename to put basic block addresses and function names\n");
    VG_(printf) ("   --interval_size=<num> interval size\n");
}

static void bbv_print_debug_usage(void) {
      VG_(printf)("    (none)\n");
}



static void bbv_fini(Int exitcode) {

   if (pcFileName[0] !=0) {
      dumpPcFile();
   }

   VG_(sprintf)(buf,"\n\n# Total periods: %ld\n# Total instructions: %lld\n"
		        "# Total reps: %lld\n# Unique reps: %lld\n"
		        "# Total fldcw instructions: %lld\n\n",
		(int)(total_instr/(ULong)interval_size),total_instr,
		global_rep_count,unique_rep_count,fldcw_count);
   VG_(write)(2, (void*)buf,VG_(strlen)(buf));
   VG_(write)(bbtrace_fd,(void*)buf,VG_(strlen)(buf));   
   VG_(close)(bbtrace_fd);
}


static void bbv_pre_clo_init(void) {
   
   VG_(details_name)            ("EXP-BBV");
   VG_(details_version)         (NULL);
   VG_(details_description)     ("calculates Simpoint Basic Block Vectors");
   VG_(details_copyright_author)(
      "Copyright (C) 2006-2008 Vince Weaver");
   VG_(details_bug_reports_to)  ("vince _at_ csl.cornell.edu");

    /* set up default output file */
   VG_(sprintf)(bbFileName,"out.bb");
   
   VG_(basic_tool_funcs)          (bbv_post_clo_init,
                                   bbv_instrument,
                                   bbv_fini);
   
   VG_(needs_command_line_options)(bbv_process_cmd_line_option,
                                   bbv_print_usage,
                                   bbv_print_debug_usage);


   VG_(track_start_client_code)( bbv_thread_called );
         
   instr_info_table = VG_(OSetGen_Create)(/*keyOff*/0,
					  NULL,
					  VG_(malloc), VG_(free));
     
}

VG_DETERMINE_INTERFACE_VERSION(bbv_pre_clo_init)

/*--------------------------------------------------------------------*/
/*--- end                                                          ---*/
/*--------------------------------------------------------------------*/
