/*
 * MemDebug.cpp
 *
 *  Created on: Apr 28, 2010
 *      Author: crueger
 */

// NDEBUG implies NO_MEMDEBUG
#ifdef NDEBUG
# ifndef NO_MEMDEBUG
#   define NO_MEMDEBUG
# endif
#endif

// NO_MEMDEBUG and MEMDEBUG are mutually exclusive, but at least one must be set
#ifdef NO_MEMDEBUG
# ifdef MEMDEBUG
#   undef MEMDEBUG
# endif
#else
# ifndef MEMDEBUG
#   define MEMDEBUG
# endif
#endif

#ifdef MEMDEBUG

#include <iostream>
#include <cstdlib>
#include <cstring>
#include <boost/thread.hpp>

#ifdef __GNUC__
#include <execinfo.h>
#include <cxxabi.h>
#endif

using namespace std;

// we need our own low level mutexex, since we cannot assure the time of construction and destruction
// otherwise
#if defined(unix) || defined(__unix)

#include <pthread.h>
#include <cassert>
#define mutex_t    pthread_mutex_t
#define mutex_init PTHREAD_MUTEX_INITIALIZER
#define mutex_lock(mtx) \
  do{\
    int res = pthread_mutex_lock(&(mtx));\
    assert(!res && "Could not lock mutex!");\
  }while(0)

#define mutex_unlock(mtx) \
  do{\
    int res = pthread_mutex_unlock(&(mtx));\
    assert(!res && "Could not unlock mutex!");\
  }while(0)

#else
# error "No thread structure defined for this plattform..."
#endif

namespace Memory {

  // This struct is added before each memory chunk
  // and contains tracking information. Anything used
  // to track memory cannot use any dynamic memory, so
  // we have to resort to classic C-idioms here.
  // This struct also contains pointers to the next
  // an previous chunks to allow fast traversion of
  // all allocated memory blocks
  struct entry_t {
    // we seperate the tracking info from the rest
    // A checksum will be calculated for this part of
    // the struct, so the information in here should
    // not change during the lifetime of the memory
    struct info_t {
      enum {length = 64};
      char file[length+1];
      int line;
#ifdef __GNUC__  // function tracking only works with GCC
      // function names can get looooong
      enum {length2 = 256};
      char function[length2+1];
#endif
      size_t nbytes;
      bool isUsed;
      void *location;
    } info;
    bool isIgnored;
    char checksum;
    entry_t *prev;
    entry_t *next;
  };


  mutex_t memorylock = mutex_init;

  // start and end of the doubly-linked list
  entry_t *begin=0;
  entry_t *end=0;

  // current amount of allocated memory
  size_t state = 0;
  // maximum amount of allocated memory
  size_t max = 0;
  // number of allocations that have been done so far
  unsigned int allocs = 0;


  // this sets the alignment of the returned memory block
  // malloc guarantees an alignment at the 8 byte border,
  // so we just do the same
  const int alignment = 8;

  // calculates a simple checksum for the info block
  // the checksum is used to find memory corruptions
  inline char calcChecksum(entry_t::info_t *info){
    char *buffer = (char*)info;
    char checksum =0;
    for(size_t i=0;i<sizeof(entry_t::info_t);i++){
      checksum+=buffer[i];
    }
    return checksum;
  }

  // gets the next alignet point which is greater than nbytes
  // this function is only called a fixed number of times, so
  // there is no need to optimize
  inline size_t doAlign(size_t nbytes){
    int nonaligned = nbytes % alignment;
    if(nonaligned) {
     return(nbytes - nonaligned + alignment);
    }
    else{
      return nbytes;
    }
  }

  // Output some state information
  void getState(){
    cout << "Maximum allocated Memory: " << max << " bytes" << endl;
    cout << "Currently allocated Memory: " << state <<" bytes" << endl;
    cout << allocs << " allocated chunks total" << endl;

    // simple traversal of the chunk list
    for(entry_t *pos=begin;pos;pos=pos->next){
      cout << "\nChunk of " << pos->info.nbytes << " bytes" << " still available" << endl;
#ifdef __GNUC__
      cout << "Chunk reserved at: " << pos->info.function
           << " (" << pos->info.file << ":" << pos->info.line  << ")" << endl;
#else
      cout << "Chunk reserved at: " << pos->info.file << ":" << pos->info.line << endl;
#endif
    }
  }

  // Adds an entry to the linked list
  void addEntry(entry_t *entry){
    // check if the entry is already in the list
    if(!entry->isIgnored)
      return;

    mutex_lock(Memory::memorylock);

    entry->next=0;            // the created block is last in the list
    entry->prev=Memory::end;  // the created block is last in the list
    if(!Memory::begin){
      // the list was empty... start a new one
      Memory::begin=entry;
    }
    else {
      // other blocks present... we can add to the last one
      Memory::end->next=entry;
    }
    Memory::end=entry;

    // update some global info
    Memory::state  += entry->info.nbytes;
    if(Memory::state>Memory::max){
        Memory::max = Memory::state;
    }
    ++Memory::allocs;
    // done with the list... it is safe to unlock now
    mutex_unlock(Memory::memorylock);
    entry->isIgnored = false;
  }

  // Deletes an entry from the linked list
  void deleteEntry(entry_t *entry){
    if(entry->isIgnored)
      return;

    mutex_lock(memorylock);
    if(entry->prev){
      entry->prev->next = entry->next;
    }
    else{
      // this node was the beginning of the list
      begin = entry->next;
    }

    if(entry->next){
      entry->next->prev = entry->prev;
    }
    else{
      // this node was the end of the list
      end = entry->prev;
    }
    Memory::state  -= entry->info.nbytes;
    mutex_unlock(memorylock);
    entry->isIgnored = true;

  }

  void _ignore(void *ptr){
    // just deletes the node from the list, but leaves the info intact
    static const size_t entrySpace = Memory::doAlign(sizeof(Memory::entry_t));
    entry_t *entry = (Memory::entry_t*)((char*)ptr-entrySpace);
    deleteEntry(entry);
  }

#ifdef __GNUC__
  // this function let's us find the caller's name
  char* getCaller(){
    // stack looks like this:
    // getCaller();
    // operator new();
    // function_we_are_looking_for(); <-
    const size_t max_depth = 3;
    void* stack_addrs[max_depth];
    size_t stack_depth;
    char **stack_strings=0;
    const char *func_name=0;
    const char *toplevel = "Global scope";
    char *retval=0;

    // get the backtrace, depth three
    stack_depth   = backtrace(stack_addrs,max_depth);
    stack_strings = backtrace_symbols(stack_addrs, stack_depth);
    // used later for demangling
    // reserved here, so we can free it unconditionally
    char *dm_function = static_cast<char*>(malloc(entry_t::info_t::length2));
    if(!dm_function){
      // malloc failed... we are out of luck
      throw std::bad_alloc();
    }

    // see if we found our function name
    if(stack_depth==max_depth){
      // find the mangled function name
      char *begin = stack_strings[max_depth-1];
      // function name starts with a (
      while(*begin && *begin!='(') ++begin;
      char *end=begin;
      while(*end && *end!='+') ++end;

      // see if we found our function name
      if(*begin && *end){
        *begin++ = 0;
        *end = 0;
        // use the C++ demangler

        size_t sz = entry_t::info_t::length2;
        int status;
        char *func_ret = abi::__cxa_demangle(begin, dm_function, &sz, &status);
        if(func_ret){
          // abi might have realloced...
          dm_function = func_ret;
          func_name = dm_function;
        }
        else{
          // demangling failed... get the function name without demangling
          func_name = begin;
        }
      }
      else{
        // function name not found... get the whole line
        func_name = stack_strings[max_depth-1];
      }

    }
    else{
      func_name = toplevel;
    }

    // now we copy the desired function name
    if((retval = static_cast<char*>(malloc(strlen(func_name)+1)))){
      // we know that the string will fit, so strcpy is safe here
      strcpy(retval,func_name);
    }
    else{
      free(stack_strings); // malloc()ed by backtrace_symbols
      free(dm_function);
      // uh-uh ... seems we are out of luck for allocations now
      throw std::bad_alloc();
    }
    free(dm_function);
    free(stack_strings); // malloc()ed by backtrace_symbols
    return retval;
  }
#endif
}

#ifdef __GNUC__

void *operator new(size_t nbytes,const char* file, int line, const char* func) throw(std::bad_alloc) {

  // to avoid allocations of 0 bytes if someone screws up
  // allocation with 0 byte size are undefined behavior, so we are
  // free to handle it this way
  if(!nbytes) {
    nbytes = 1;
  }

  // get the size of the entry, including alignment
  static const size_t entrySpace = Memory::doAlign(sizeof(Memory::entry_t));

  void *res;
  if(!(res=malloc(entrySpace + nbytes))){
    // new must throw, when space is low
    throw std::bad_alloc();
  }

  // build the entry in front of the space
  Memory::entry_t *entry = (Memory::entry_t*) res;
  memset(res,0,entrySpace);
  entry->info.nbytes = nbytes;
  entry->info.isUsed = true;
  strncpy(entry->info.file,file,Memory::entry_t::info_t::length);
  entry->info.file[Memory::entry_t::info_t::length] = '\0';
  entry->info.line=line;
  strncpy(entry->info.function,func,Memory::entry_t::info_t::length2);
  entry->info.function[Memory::entry_t::info_t::length2] = '\0';
  // the space starts behind the info
  entry->info.location = (char*)res + entrySpace;

  // mark the block as not in the list (will be changed by addEntry)
  entry->isIgnored = true;
  Memory::addEntry(entry);

  // get the checksum...
  entry->checksum = Memory::calcChecksum(&entry->info);

  // ok, space is prepared... the user can have it.
  // the rest (constructor, deleting when something is thrown etc)
  // is handled automatically
  return entry->info.location;
}

#else

void *operator new(size_t nbytes,const char* file, int line) throw(std::bad_alloc) {

  // to avoid allocations of 0 bytes if someone screws up
  // allocation with 0 byte size are undefined behavior, so we are
  // free to handle it this way
  if(!nbytes) {
    nbytes = 1;
  }

  // get the size of the entry, including alignment
  static const size_t entrySpace = Memory::doAlign(sizeof(Memory::entry_t));

  void *res;
  if(!(res=malloc(entrySpace + nbytes))){
    // new must throw, when space is low
    throw std::bad_alloc();
  }

  // build the entry in front of the space
  Memory::entry_t *entry = (Memory::entry_t*) res;
  memset(res,0,entrySpace);
  entry->info.nbytes = nbytes;
  entry->info.isUsed = true;
  strncpy(entry->info.file,file,Memory::entry_t::info_t::length);
  entry->info.file[Memory::entry_t::info_t::length] = '\0';
  entry->info.line=line;
  // the space starts behind the info
  entry->info.location = (char*)res + entrySpace;

  // mark the block as not in the list (will be changed by addEntry)
  entry->isIgnored = true;
  Memory::addEntry(entry);

  // get the checksum...
  entry->checksum = Memory::calcChecksum(&entry->info);
  // this will be set to true, when the block is removed from
  // the list for any reason
  entry->isIgnored = false;

  // ok, space is prepared... the user can have it.
  // the rest (constructor, deleting when something is thrown etc)
  // is handled automatically
  return entry->info.location;
}

#endif

void *operator new(size_t nbytes) throw(std::bad_alloc) {
  // Just forward to the other operator, when we do not know from
  // where the allocation came
#ifdef __GNUC__
  // this might throw bad_alloc
  char *caller = Memory::getCaller();
  void* retval = 0;

  // if this throws, we have to clean up the caller anyway
  try{
    retval = operator new(nbytes,"Unknown",0,caller);
  }
  catch(...)
  {
    free(caller); // malloc()ed by Memory::getCaller();
    throw;
  }
  free(caller); // malloc()ed by Memory::getCaller();
  return retval;
#else
  return operator new(nbytes,"Unknown",0);
#endif
}

#ifdef __GNUC__

void *operator new[] (size_t nbytes,const char* file, int line, const char* func) throw(std::bad_alloc) {
  // The difference between new and new[] is just for compiler bookkeeping.
  return operator new(nbytes,file,line,func);
}

#else

void *operator new[] (size_t nbytes,const char* file, int line) throw(std::bad_alloc) {
  // The difference between new and new[] is just for compiler bookkeeping.
  return operator new(nbytes,file,line);
}

#endif

void *operator new[] (size_t nbytes) throw(std::bad_alloc) {
  // Forward again
#ifdef __GNUC__
  // this might throw bad_alloc
    char *caller = Memory::getCaller();
    void *retval=0;

    // if this throws, we have to clean up the caller anyway
    try{
      retval = operator new[] (nbytes,"Unknown",0,caller);
    }
    catch(...)
    {
      free(caller); // malloc()ed by Memory::getCaller();
      throw;
    }
    free(caller); // malloc()ed by Memory::getCaller();
    return retval;
#else
  return operator new[] (nbytes,"Unknown",0);
#endif
}

void operator delete(void *ptr) throw() {
  if(!ptr){
    cerr << "Warning: Deleting NULL pointer" << endl;
    return;
  }

  // get the size for the entry, including alignment
  static const size_t entrySpace = Memory::doAlign(sizeof(Memory::entry_t));

  // get the position for the entry from the pointer the user gave us
  Memory::entry_t *entry = (Memory::entry_t*)((char*)ptr-entrySpace);

  // let's see if the checksum is still matching
  if(Memory::calcChecksum(&entry->info)!=entry->checksum){
    cerr << "Possible memory corruption detected!" << endl;
    cerr << "Trying to recover allocation information..." << endl;
    cerr << "Memory was allocated at " << entry->info.file << ":" << entry->info.line << endl;
    terminate();
  }

  // this will destroy the checksum, so double deletes are caught
  entry->info.isUsed = false;
  Memory::deleteEntry(entry);

  // delete the space reserved by malloc
  free((char*)ptr-entrySpace);
}

// operator that is called when the constructor throws
// do not call manually
void operator delete(void *ptr,const char*, int) throw() {
  operator delete(ptr);
}

void operator delete[](void *ptr){
  // again difference between delete and delete[] is just in compiler bookkeeping
  operator delete(ptr);
}

// and another operator that can be called when a constructor throws
void operator delete[](void *ptr,const char*, int) throw(){
  operator delete(ptr);
}
#endif