/*
* Copyright (c) 2007, Insomniac Games
* 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 <organization> 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 INSOMNIAC GAMES ``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 INSOMNIAC GAMES 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.
*/
////////////////////////////////////////////////////////////////////////////////////////////////
//
//  win32_penter.cpp
//
//  written by: Rob Wyatt
//
//  File to hook a function being called and to 'adjust' the return address so you get called
//  on the way out as well. In these entry and exit calls timings are taken and recorded which
//  are made into detailed profile output by an external tool.
//
////////////////////////////////////////////////////////////////////////////////////////////////
#include "igg/igprofile.h"

#define MAX_SAMPLE  1024*1024

struct entry
{
  unsigned int   eip_func;
  unsigned int   eip_caller;
  __int64        time;
  __int64        child_time;
  int            stack_depth;
  int            stack_ptr;
};

struct ProfileHeader
{
  double m_cpu_speed;
  DWORD  m_num_samples;
  char   m_module_name[512];
};

static int g_count = 0;
static entry* g_trace_ptr = 0;
static bool g_trace_active = false;
static unsigned int g_call_stack[4096];
static int g_stack_idx = 0;

////////////////////////////////////////////////////////////////////////////////////////////////
static void FlushProfileSamples(const char* fname)
{
  FILE* dump_file = fopen(fname,"wb");

  ProfileHeader header;
  header.m_num_samples = g_count;
  GetModuleFileName(0,header.m_module_name,512);

  header.m_cpu_speed = IGG::g_ProfileCon.m_cpu_speed;
  // write the module filename so we can find the symbols later
  fwrite(&header,sizeof(header),1,dump_file);

  // write the samples
  fwrite(g_trace_ptr,sizeof(entry),g_count,dump_file);
  fclose(dump_file);
}

////////////////////////////////////////////////////////////////////////////////////////////////
//
//  _pleave()
//
//  written by: Rob Wyatt
//
//
//
////////////////////////////////////////////////////////////////////////////////////////////////
extern "C" void static __declspec(naked) _cdecl _pleave( void ) 
{
  _asm
  {
    sub esp,4
    push eax
    push ebx
    push ecx
    push edx
    push ebp
    push edi
    push esi

    mov ecx,g_stack_idx
    dec ecx
    mov g_stack_idx,ecx
    lea eax,[g_call_stack+ecx*8]
    
    mov ebx,[eax+0]
    mov edi,[eax+4]
    mov [edi+4],ebx             // store the calling functions address
    mov [esp+28],ebx

    // get the new time stamp and do a 64bit subtract from the source time
    rdtsc
    sub edx,[edi+12]
    sbb eax,[edi+8]
    mov [edi+8],eax
    mov [edi+12],edx
    mov [edi+24],ecx             // call index

    cmp ecx,1
    jl  noparent

    dec ecx
    lea ebx,[g_call_stack+ecx*8]
    mov esi,[ebx+4]
    add [esi+16],eax
    adc [esi+20],edx
noparent:
    pop esi
    pop edi
    pop ebp
    pop edx
    pop ecx
    pop ebx
    pop eax
    ret
  }
}


////////////////////////////////////////////////////////////////////////////////////////////////
//
//  _penter()
//
//  written by: Rob Wyatt
//
//
//
////////////////////////////////////////////////////////////////////////////////////////////////
extern "C" void __declspec(naked) _cdecl _penter( void ) 
{
  _asm 
  {
    cmp BYTE PTR [g_trace_active],0  // if we are not active then terminate as quickly as possible
    jne  start_profile
    ret

start_profile:
    // store all the registers
    push eax
    push ebx
    push ecx
    push edx
    push ebp
    push edi
    push esi

    mov edi,g_count
    mov eax,g_trace_ptr                      // load the pointer our data
    mov esi,edi
    shl edi,5    
    lea edi,[eax+edi]                        // destination address of where to store
    inc esi
    mov g_count,esi                           // update the sample counter
    
    lea ebx,[esp+28]                          // get the stack address of the return address
    mov eax,[ebx]                             // get the return address
    mov [edi+0],eax                           // store it
    mov DWORD PTR [edi+4],0                   // store the stack pointer
    rdtsc                                     // 64bit time stamp in EDX:EAX
    mov [edi+8],eax                           // store low 32 bits
    mov [edi+12],edx                          // store high 32 bits
    mov DWORD PTR[edi+16],0
    mov DWORD PTR [edi+20],0                            // zero child time
    mov DWORD PTR [edi+24],0xffffffff                   // zero child time    
    mov [edi+28],esp                          // store stack pointer
    xor ebx,ebx
    mov eax,1
    cmp esi,MAX_SAMPLE                        // if we just stored the last sample the disable tracing
    cmovge eax,ebx
    mov DWORD PTR g_trace_active,eax

    // fiddle the return address
    mov eax,g_stack_idx
    lea eax,[g_call_stack+eax*8]
    inc g_stack_idx
    mov ebx,[esp+32]
    mov [eax+0],ebx
    mov [eax+4],edi
    // send the function return back to us
    mov eax,_pleave
    mov DWORD PTR [esp+32],eax

    // restore registers and return
    pop esi
    pop edi
    pop ebp
    pop edx
    pop ecx
    pop ebx
    pop eax
    ret
  }
}


namespace IGG
{
  ////////////////////////////////////////////////////////////////////////////////////////////////
  void StartFunctionProfiler()
  {
    if (g_trace_active)
      return;

    g_trace_ptr = new entry[MAX_SAMPLE];
    g_count = 0;    
    g_trace_active = true;
    g_stack_idx = 0;
  }

  ////////////////////////////////////////////////////////////////////////////////////////////////
  void StopFunctionProfiler(char* fname)
  {
    g_trace_active = false;
    FlushProfileSamples(fname);
    delete[] g_trace_ptr;
    g_trace_ptr = 0;
  }

  ////////////////////////////////////////////////////////////////////////////////////////////////
  bool IsFunctionProfilerActive()
  {
    return g_trace_active;
  }
}