/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) CNRS, INRIA, Universite Bordeaux 1, Telecom SudParis
 * See COPYING in top-level directory.
 */

#ifndef _REENTRANT
#define _REENTRANT
#endif

#include <eztrace-core/eztrace_config.h>
#include <eztrace-lib/eztrace.h>
#include <eztrace-lib/eztrace_module.h>
//#include <eztrace-lib/eztrace_litl_packed.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

/* set to 1 when all the hooks are set.
 * This is useful in order to avoid recursive calls
 */
static int _memory_initialized = 0;

#define CURRENT_MODULE memory
DECLARE_CURRENT_MODULE;

#define MAGIC_PATTERN 0xdeadbeef

/* when a buffer is allocated with eztrace, it has the following pattern:
 * PADDING BLOCK_INFO USER_BLOCK
 * block_info has to be right-aligned because of calloc.
 * the address should be a multiple of 8 bytes in order to avoid bugs
 * (if not aligned, some weird bugs may happen when using -O3)
 */
enum _memory_type {
  MEM_TYPE_MALLOC,
  MEM_TYPE_CUSTOM_MALLOC
};

/* todo: we could add information like:
 * - date of malloc
 * - thread id that allocated the block
 * - NUMA node ?
 */
struct mem_block_info {
  void* u_ptr; /* address of the user block */
  void* p_ptr; /* address of the padding */

  enum _memory_type mem_type;

  size_t total_size; /* size allocated (including this structure) */
  size_t size;       /* size of the buffer (not including this structure) */

  /* WARNING: this must be the last field of the structure */
  uint32_t canary; /* this is used for checking that we malloc'ed the buffer */
} __attribute__((__packed__));

/* size of the padding + mem_info structure */
#define GET_HEADER_SIZE() (sizeof(struct mem_block_info))

#define CANARY_OK(u_ptr) ((*(uint32_t*)((u_ptr) - sizeof(uint32_t))) == MAGIC_PATTERN)

/* converts a pointer to a user_block into a pointer to the block info */
#define USER_PTR_TO_BLOCK_INFO(u_ptr, b_ptr)                     \
  do {                                                           \
    if (!CANARY_OK(u_ptr)) {                                     \
      /* we didn't malloc this buffer */                         \
      (b_ptr) = NULL;                                            \
      break;                                                     \
    }                                                            \
    b_ptr = (void*)(ptr - (void*)sizeof(struct mem_block_info)); \
  } while (0)

/* converts a pointer to a user_block into a pointer to the padding */
#define USER_PTR_TO_PADDING(u_ptr, p_ptr) \
  do {                                    \
    struct mem_block_info* b_ptr;         \
    USER_PTR_TO_BLOCK_INFO(u_ptr, b_ptr); \
    if (!b_ptr) {                         \
      (p_ptr) = NULL;                     \
      break;                              \
    }                                     \
    (p_ptr) = b_ptr->p_ptr;               \
  } while (0)

/* fill a mem_info structure
 * @param p_mem the mem_info* structure to fill
 * @param ptr the address returned by (m,c,re)alloc
 * @param nmemb the number of elements
 * @param block_size the size of 1 element
 */
#define INIT_MEM_INFO(p_mem, ptr, nmemb, block_size)              \
  do {                                                            \
    unsigned int nb_memb_header = GET_HEADER_SIZE() / block_size; \
    if (block_size * nb_memb_header < GET_HEADER_SIZE())          \
      nb_memb_header++;                                           \
    void* u_ptr = ptr + (block_size * nb_memb_header);            \
    p_mem = u_ptr - sizeof(struct mem_block_info);                \
    p_mem->p_ptr = ptr;                                           \
    p_mem->total_size = (nmemb + nb_memb_header) * block_size;    \
    p_mem->size = nmemb * block_size;                             \
    p_mem->u_ptr = u_ptr;                                         \
    p_mem->canary = MAGIC_PATTERN;                                \
  } while (0)

/* todo: also implement mmap and munmap ?
 */
void* (*libcalloc)(size_t nmemb, size_t size) = NULL;
void* (*libmalloc)(size_t size) = NULL;
void (*libfree)(void* ptr) = NULL;
void* (*librealloc)(void* ptr, size_t size) = NULL;

static int malloc_protect_on = 0;

/* Custom malloc function. It is used when libmalloc=NULL (e.g. during startup)
 * This function is not thread-safe and is very likely to be bogus, so use with
 * caution
 */
static void* hand_made_malloc(size_t size) {
  /* allocate a 1MB buffer */
#define POOL_SIZE (1024 * 1024)
  static char mem[POOL_SIZE] = {'\0'};

  /* since this function is only used before we found libmalloc, there's no
   * fancy memory management mechanism (block reuse, etc.)
   */
  static char* next_slot = &mem[0];
  static int total_alloc = 0;

  if (libmalloc)
    /* let's use the real malloc */
    return malloc(size);

  struct mem_block_info* p_block = NULL;
  INIT_MEM_INFO(p_block, next_slot, size, 1);

  p_block->mem_type = MEM_TYPE_CUSTOM_MALLOC;
  total_alloc += size;
  next_slot = next_slot + p_block->total_size;

  return p_block->u_ptr;
}

void* malloc(size_t size) {
  /* if memory_init hasn't been called yet, we need to get libc's malloc
   * address
   */
  if (!libmalloc) {
    if (malloc_protect_on)
      /* protection flag says that malloc is already trying to retrieve the
       * address of malloc.
       * If we call dlsym now, there will be an infinite recursion, so let's
       * allocate memory 'by hand'
       */
      return hand_made_malloc(size);

    /* set the protection flag and retrieve the address of malloc.
     * If dlsym calls malloc, memory will be allocated 'by hand'
     */
    malloc_protect_on = 1;
    INTERCEPT_FUNCTION("malloc", libmalloc);
    /* it is now safe to call libmalloc */
    malloc_protect_on = 0;
  }
  void* return_value = NULL;

  FUNCTION_ENTRY_WITH_ARGS(size);

  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();

    void* pptr = libmalloc(size + GET_HEADER_SIZE());
    struct mem_block_info* p_block = NULL;
    INIT_MEM_INFO(p_block, pptr, size, 1);
    p_block->mem_type = MEM_TYPE_MALLOC;

    //    EZTRACE_EVENT_PACKED_2(EZTRACE_MEMORY_MALLOC, p_block->size, (app_ptr)p_block->u_ptr);

#if 0
    /* for debugging purpose only */
    uint32_t* canary = p_block->u_ptr-sizeof(uint32_t);
    if(*canary != MAGIC_PATTERN) {
      eztrace_warn( "warning: canary = %x instead of %x\n", *canary, MAGIC_PATTERN);
    }
#endif
    return_value=p_block->u_ptr;
    EZTRACE_PROTECT_OFF();
    goto out;
  }

  return_value = libmalloc(size);
 out:
  FUNCTION_EXIT_WITH_ARGS(return_value);
  return return_value;
}

void* realloc(void* ptr, size_t size) {

  /* if ptr is NULL, realloc behaves like malloc */
  if (!ptr)
    return malloc(size);

  /* if size=0 and ptr isn't NULL, realloc behaves like free */
  if (!size && ptr) {
    free(ptr);
    return NULL;
  }

  INTERCEPT_FUNCTION("realloc", librealloc);

  if (!CANARY_OK(ptr)) {
    /* we didn't malloc'ed this buffer */
    return librealloc(ptr, size);
  }

  FUNCTION_ENTRY_WITH_ARGS(ptr, size);
  void* return_value = NULL;
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    struct mem_block_info* p_block;
    USER_PTR_TO_BLOCK_INFO(ptr, p_block);
    size_t old_size = p_block->size;
    size_t header_size = p_block->total_size - p_block->size;

    if (p_block->mem_type != MEM_TYPE_MALLOC) {
      eztrace_warn( "Warning: realloc a ptr that was allocated by hand_made_malloc\n");
    }

    void* pptr = librealloc(p_block->p_ptr, size + header_size);

    if (!p_block) {
      EZTRACE_PROTECT_OFF();
      return NULL;
    }

    INIT_MEM_INFO(p_block, pptr, size + header_size, 1);

    p_block->mem_type = MEM_TYPE_MALLOC;

    //    EZTRACE_EVENT_PACKED_3(EZTRACE_MEMORY_REALLOC, old_size, p_block->size, (app_ptr)p_block->u_ptr);

    return_value = p_block->u_ptr;
    EZTRACE_PROTECT_OFF();
    goto out;
  }
  return_value = librealloc(ptr, size);
 out:
  FUNCTION_EXIT_WITH_ARGS(return_value);
  return return_value;
}

void* calloc(size_t nmemb, size_t size) {
  if (!libcalloc) {
    void* ret = hand_made_malloc(nmemb * size);
    if (ret) {
      memset(ret, 0, nmemb * size);
    }
    return ret;
  }

  void* return_value = NULL;
  
  FUNCTION_ENTRY_WITH_ARGS(nmemb, size);
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    /* compute the number of blocks for header */
    int nb_memb_header = sizeof(struct mem_block_info) / size;
    if (size * nb_memb_header < sizeof(struct mem_block_info))
      nb_memb_header++;

    /* allocate buffer + header */
    void* p_ptr = libcalloc(nmemb + nb_memb_header, size);

    struct mem_block_info* p_block = NULL;
    INIT_MEM_INFO(p_block, p_ptr, nmemb, size);
    p_block->mem_type = MEM_TYPE_MALLOC;

    //    EZTRACE_EVENT_PACKED_2(EZTRACE_MEMORY_MALLOC, p_block->size, (app_ptr)p_block->u_ptr);

#if 0
    /* for debugging purpose only */
    uint32_t* canary = p_block->u_ptr-sizeof(uint32_t);
    if(*canary != MAGIC_PATTERN) {
      eztrace_warn("warning: canary = %x instead of %x\n", *canary, MAGIC_PATTERN);
    }
#endif
  
    return_value =  p_block->u_ptr;

    EZTRACE_PROTECT_OFF();
    goto out;
  }
  return_value = libcalloc(nmemb, size);

 out:
  FUNCTION_EXIT_WITH_ARGS(return_value);
  return return_value;
}

void free(void* ptr) {

  if (!ptr) {
    return;
  }

  if (!libfree) {
    static int plop=0;
    if(plop++ > 10) {
      plop=0;
      return;
    }
    INTERCEPT_FUNCTION("free", libfree);

    if(!libfree) {
      /* dlsym failed, so it is not safe to call libfree. Don't free the memory */
      return;
    }
    plop--;
  }

  /* first, check wether we malloc'ed the buffer */
  if (!CANARY_OK(ptr)) {
    /* we didn't malloc this buffer */
    libfree(ptr);
    return;
  }

  FUNCTION_ENTRY_WITH_ARGS(ptr);

  /* retrieve the block information and free it */
  EZTRACE_PROTECT {
    struct mem_block_info* p_block;
    USER_PTR_TO_BLOCK_INFO(ptr, p_block);

    if (p_block->mem_type == MEM_TYPE_MALLOC) {
      //      EZTRACE_EVENT_PACKED_2(EZTRACE_MEMORY_FREE, p_block->size, (app_ptr)p_block->u_ptr);
      libfree(p_block->p_ptr);
    } else {
      /* the buffer was allocated by hand_made_malloc, there's nothing to free */
    }
  }
  FUNCTION_EXIT;
}

PPTRACE_START_INTERCEPT_FUNCTIONS(memory)
INTERCEPT3("free", libfree)
INTERCEPT3("malloc", libmalloc)
INTERCEPT3("calloc", libcalloc)
INTERCEPT3("realloc", librealloc)
PPTRACE_END_INTERCEPT_FUNCTIONS(memory)

static void init_memory() {
  malloc_protect_on = 1;
  INSTRUMENT_FUNCTIONS(memory);

  if (eztrace_autostart_enabled())
    eztrace_start();

  malloc_protect_on = 0;
  _memory_initialized = 1;
}


static void finalize_memory(void) {
  _memory_initialized = 0;

  eztrace_stop();
}

static void _memory_init(void) __attribute__((constructor));
static void _memory_init(void) {
  malloc_protect_on = 1;
  eztrace_log(dbg_lvl_debug, "eztrace_memory constructor starts\n");
  EZT_REGISTER_MODULE(memory, "Module for dynamic memory (malloc, free, etc.)",
		      init_memory, finalize_memory);
  eztrace_log(dbg_lvl_debug, "eztrace_memory constructor ends\n");
  malloc_protect_on = 0;
}
