/*
UC Santa Barbara ArchLab
www.cs.ucsb.edu/~arch
Wavelet-Based Phase Classification Source Code
www.cs.ucsb.edu/~arch/wavelet

Copyright (c) 2006 The Regents of the University of California.
All Rights Reserved.

Permission to use, copy, modify, and distribute this software and its
documentation for educational, research and non-profit purposes,
without fee, and without a written agreement is hereby granted,
provided that the above copyright notice, this paragraph and the
following three paragraphs appear in all copies.

Permission to incorporate this software into commercial products may
be obtained by contacting the University of California. For
information about obtaining such a license contact:
Tim Sherwood <sherwood at cs dot ucsb dot edu>

This software program and documentation are copyrighted by The Regents
of the University of California. The software program and
documentation are supplied "as is", without any accompanying services
from The Regents. The Regents does not warrant that the operation of
the program will be uninterrupted or error-free. The end-user
understands that the program was developed for research purposes and
is advised not to rely exclusively on the program for any reason.

IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO ANY PARTY
FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES,
INCLUDING LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND
ITS DOCUMENTATION, EVEN IF THE UNIVERSITY OF CALIFORNIA HAS BEEN
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. THE UNIVERSITY OF
CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS ON AN "AS IS"
BASIS, AND THE UNIVERSITY OF CALIFORNIA HAS NO OBLIGATIONS TO PROVIDE
MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.

Ted Huffmire, (www.cs.ucsb.edu/~huffmire, huffmire at cs dot ucsb dot edu)
Tim Sherwood (www.cs.ucsb.edu/~sherwood, sherwood at cs dot ucsb dot edu)
14 June 2006

File: pintool.C

Publication:
Ted Huffmire and Tim Sherwood. Wavelet-Based Phase Classification. Proceedings 
of the Fifteenth International Conference on Parallel Architectures and 
Compilation Techniques (PACT'06), Seattle, Washington, September 16-20, 2006.
*/

#include "pin.H"
#include <iostream>
#include <fstream>
#include <iomanip>
#include <linux/unistd.h>
#include <assert.h>

/* ===================================================================== */
/* Global Variables */
/* ===================================================================== */

// The running count of instructions is kept here
UINT64 icount = 0;
std::ofstream TraceFile;
std::ofstream BBFile;

#define STEP 10000
long filesize=0L;
int filenum=0;
int dots=0;
int execbit=0;
char **prevargs;
VOID *bbaddr;

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

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool",
    "o", "~/trace.out", "specify trace file name");
KNOB<BOOL> KnobValues(KNOB_MODE_WRITEONCE, "pintool",
    "values", "1", "Output memory values reads and written");

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

static INT32 Usage()
{
    cerr <<
      "This tool produces a memory address trace.\n";

    cerr << KNOB_BASE::StringKnobSummary();

    cerr << endl;

    return -1;
}


static VOID EmitMem(VOID * ea, INT32 size)
{
    if (!KnobValues)
        return;
    
    switch(size)
    {
      case 0:
        TraceFile << setw(1);
        break;
        
      case 1:
        TraceFile << static_cast<UINT32>(*static_cast<UINT8*>(ea));
        break;
        
      case 2:
        TraceFile << *static_cast<UINT16*>(ea);
        break;
        
      case 4:
        TraceFile << *static_cast<UINT32*>(ea);
        break;
        
      case 8:
        TraceFile << *static_cast<UINT64*>(ea);
        break;
        
      default:
        TraceFile.unsetf(ios::showbase);
        TraceFile << setw(1) << "0x";
        for (INT32 i = 0; i < size; i++)
        {
            TraceFile << static_cast<UINT32>(static_cast<UINT8*>(ea)[i]);
        }
        TraceFile.setf(ios::showbase);
        break;
    }
}

static VOID RecordMem(VOID * ip, CHAR r, VOID * addr, INT32 size, UINT32 tid, CHAR *routine_name, CHAR *section_name)
{
  int totaldots = 0;
  totaldots=icount/STEP;
  for (int i=0; i < totaldots-dots; i++)
    fprintf(stderr, "[%d:%d] %d\n", getpid(), tid, icount);
  fflush(stderr);
  dots=totaldots;

  TraceFile << icount;
  TraceFile << " ";
  TraceFile << ip;
  TraceFile << " ";
  TraceFile << addr;
  TraceFile << " ";
  TraceFile << size;
  TraceFile << "\n";
  TraceFile.flush();
  filesize += 64L;
  if (filesize > 1000000000L) {
    char buf[256];
    TraceFile.close();
    sprintf(buf, "~/%dtrace_%d_.out", getpid(), filenum);
    TraceFile.open(buf);
    TraceFile.setf(ios::showbase);
    filesize = 0L;
    BBFile.close();
    sprintf(buf, "~/%dbb_%d_.out", getpid(), filenum);
    BBFile.open(buf);
    BBFile.setf(ios::showbase);
    filenum++;
  }
}

VOID Exec(CHAR *syscall_name, CHAR *v, VOID *p, VOID *envp, UINT32 tid)
{
  FILE *logfile;
  int i = 0;
  char buf[256];
  sprintf(buf, "~/%dexec.log", getpid());
  fprintf(stderr, "about to open %s\n", buf);
  logfile = fopen(buf, "a");
  fprintf(logfile, "%s %s\n", syscall_name, v);
  if (strcmp(syscall_name, "execve") == 0) {
    char **ptr = (char **)p;
    char **argv;
    argv = (char **)malloc(1024 * sizeof(char *));
    argv[0] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[0], "/cs/arch/huffmire/pin/Bin/pin");
    argv[1] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[1], "-mt");
    argv[2] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[2], "-t");
    argv[3] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[3], "~/trace");
    argv[4] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[4], "--");
    argv[5] = (char *)malloc(256 * sizeof(char));
    strcpy(argv[5], v);
    for (i = 1; ptr && ptr[i] && i < 11; i++) {
      argv[i+5] = (char *)malloc(256 * sizeof(char));
      strcpy(argv[i+5], ptr[i]);
    }
    argv[i+5] = 0;
    fprintf(stderr, "about to call exec(pin):\n");
    fprintf(logfile, "about to call exec(pin):\n");
    for (i = 0; argv[i] != 0; i++) {
      fprintf(stderr, "argv[i] = %s\n", argv[i]);
      fprintf(logfile, "argv[i] = %s\n", argv[i]);
    }
    fflush(logfile);
    fclose(logfile);
    fflush(stderr);
    execve("/cs/arch/huffmire/pin/Bin/pin", argv, (char **)envp);
    fprintf(stderr, "This should never be printed!\n");
    fflush(stderr);
    return;
  }
}

VOID SysBefore(VOID * ip, INT32 num, VOID * arg0, VOID * arg1, VOID * arg2, VOID * arg3, VOID * arg4, VOID * arg5, UINT32 tid)
{
  if (num == SYS_execve) {
    Exec("execve", (CHAR *)arg0, arg1, arg2, tid);
  }
}


VOID Instruction(INS ins, VOID *v)
{

  if (INS_IsSyscall(ins)) {
    INS_InsertCall(ins, IPOINT_BEFORE, AFUNPTR(SysBefore),
		   IARG_INST_PTR, IARG_SYSCALL_NUMBER,
		   IARG_SYSCALL_ARG0, IARG_SYSCALL_ARG1,
		   IARG_SYSCALL_ARG2, IARG_SYSCALL_ARG3,
		   IARG_SYSCALL_ARG4, IARG_SYSCALL_ARG5,
		   IARG_THREAD_ID, IARG_END);
  }
  
  
  if (INS_IsMemoryRead(ins)) {
    RTN r = INS_Rtn(ins);
    SEC section = RTN_Sec(r);
    IMG image = SEC_Img(section);
    INS_InsertPredicatedCall(
			     ins, IPOINT_BEFORE, (AFUNPTR)RecordMem,
			     IARG_INST_PTR,
			     IARG_UINT32, 'R',
			     IARG_MEMORYREAD_EA,
			     IARG_MEMORYREAD_SIZE,
			     IARG_THREAD_ID, 
			     IARG_PTR,
			     RTN_Name(r).c_str(),
			     IARG_PTR,
			     IMG_Name(image).c_str(),
			     IARG_END);
  }
  
  if (INS_HasMemoryRead2(ins)) {
    RTN r = INS_Rtn(ins);
    SEC section = RTN_Sec(r);
    IMG image = SEC_Img(section);
    INS_InsertPredicatedCall(
			     ins, IPOINT_BEFORE, (AFUNPTR)RecordMem,
			     IARG_INST_PTR,
			     IARG_UINT32, 'R',
			     IARG_MEMORYREAD2_EA,
			     IARG_MEMORYREAD_SIZE,
			     IARG_THREAD_ID,
			     IARG_PTR,
			     RTN_Name(r).c_str(),
			     IARG_PTR,
			     IMG_Name(image).c_str(),
			     IARG_END);
  }
}

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

VOID Fini(INT32 code, VOID *v)
{
  TraceFile.flush();
  TraceFile.close();
  BBFile.flush();
  BBFile.close();
}

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

// This function is called before every block
VOID docount(INT32 c, VOID * addr) 
{ 
  icount += c; 
  bbaddr = addr;
  BBFile << addr;
  BBFile << " ";
  BBFile << icount;
  BBFile << "\n";
  BBFile.flush();
}
    
// Pin calls this function every time a new basic block is encountered
// It inserts a call to docount
VOID Trace(TRACE trace, VOID *v)
{
    // Visit every basic block  in the trace
    for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl))
    {
        // Insert a call to docount for every bbl, passing the number of instructions.
        // IPOINT_ANYWHERE allows Pin to schedule the call anywhere in the bbl to obtain best performance.
        
        BBL_InsertCall(bbl, IPOINT_ANYWHERE, (AFUNPTR)docount, IARG_UINT32, BBL_NumIns(bbl),  IARG_ADDRINT, BBL_Address(bbl), IARG_END);
    }
}

int main(int argc, char *argv[])
{
  char buf[256];
  string trace_header = string("#\n"
			       "# Memory Access Trace Generated By Pin\n"
			       "#\n");
  
  
  PIN_InitSymbols();
  PIN_Init(argc, argv);

  if( ParseCommandLine(argc,argv) )
    {
      return Usage();
    }
  
  sprintf(buf, "~/%dtrace.out", getpid());
  TraceFile.open(buf);
  sprintf(buf, "~/%dbb.out", getpid());
  BBFile.open(buf);
  TraceFile.setf(ios::showbase);
  BBFile.setf(ios::showbase);

  prevargs = (char **)malloc(16 * sizeof(char *));
  for (int i = 0; i < 16; i++) {
    prevargs[i] = (char *)malloc(256 * sizeof(char));
    strcpy(prevargs[i], "");
  }

  
  INS_AddInstrumentFunction(Instruction, 0);
  TRACE_AddInstrumentFunction(Trace, 0);
  PIN_AddFiniFunction(Fini, 0);
  
  // Never returns
  
  PIN_StartProgram();
  
  if (0) EmitMem (0,0); 
  
  return 0;
}