/*
 * fda.c - Free Debug Allocator
 * Copyright (C) 1997 Thomas Helvey <tomh@inxpress.net>
 *
 * 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, 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */
#include <sys/types.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "fda.h"
#include <assert.h>

static void lowmem_fn(void)
{
  fprintf(stderr, "Warning: low memory!\n");
}

static void nomem_fn(void)
{
  fprintf(stderr, "Error: out of memory!\n");
  abort();
}

static void (*lowMemFn)(void) = lowmem_fn;
static void (*noMemFn)(void)  = nomem_fn;

void set_lowmem_handler(void (*fn)(void))
{
  lowMemFn = fn;
}

void set_nomem_handler(void (*fn)(void))
{
  noMemFn = fn;
}

#if defined(_DEBUG) || defined(DEBUG)

#define ULONGSIZE sizeof(unsigned long)
#define base_ptr(p) ((char*)(p) - ULONGSIZE)
#define base_size(s) ((((s) + 3 + 2 * ULONGSIZE) / ULONGSIZE) * ULONGSIZE)

struct _location {
  struct _location* next;
  struct _location* prev;
  const char* file;
  int         line;
  int         count;
};

struct _blkhdr {
  struct _blkhdr* next;     /* Next block in list */
  unsigned char*  buf;      /* Allocated buffer */
  size_t          size;     /* Size of allocated buffer */
  int             ref;      /* Buffer referenced flag */
  struct _location* location;  /* Where the allocation took place */
}; 

typedef struct _blkhdr   BlkHdr;
typedef struct _location Location;

static const unsigned long hell = 0xdeadbeef;
static size_t  byteCount = 0;  /* count of currently allocated bytes */
static size_t  blockCount = 0; /* count of allocated blocks */
static BlkHdr* bhTab[FDA_HASH_TABLE_SIZE];
static Location* locationList = 0;


/*
 * fda_get_byte_count - returns the client memory allocated in bytes
 */
size_t fda_get_byte_count(void)
{
  return byteCount;
}

/*
 * fda_get_block_count - returns the number of blocks allocated
 */
size_t fda_get_block_count(void)
{
  return blockCount;
}

/*
 * find_location - finds a location on the list, this
 * only compares pointers so it should only be used with
 * ANSI __FILE__ and __LINE__ macros.
 */
static Location* find_location(const char* file, int line)
{
  Location* location = locationList;
  for ( ; location; location = location->next) {
    if (file == location->file && line == location->line)
      return location;
  }
  return 0;
}

/*
 * add_location - adds a allocation location to the list
 * returns a pointer to the new location
 */
static Location* add_location(const char* file, int line)
{
  Location* location;
  assert(0 != file);
  if ((location = (Location*) malloc(sizeof(Location))) != 0) {
    location->next = locationList;
    location->prev = 0;
    location->file = file;
    location->line = line;
    location->count = 0;
    if (location->next)
      location->next->prev = location;
    locationList = location;
  }
  return location;
}

/*
 * free_location - frees a file/line info location
 */
static void free_location(Location* location)
{
  assert(0 != location && 0 == location->count);
  if (location->next != 0)
    location->next->prev = location->prev;
  if (location->prev != 0)
    location->prev->next = location->next;
  else
    locationList = location->next;
  free(location);
}

/*
 * hash_ptr - simple pointer hash function
 */
static unsigned long hash_ptr(const void* p)
{
  return ((unsigned long)p >> 2) % FDA_HASH_TABLE_SIZE;
}

/*
 * find_blk_exhaustive - find a block by scanning the
 * entire hash table. This function finds blocks that do not
 * start at the pointer returned from Malloc.
 */
static BlkHdr* find_blk_exhaustive(const unsigned char* p)
{
  int i;
  BlkHdr* bh;

  for (i = 0; i < FDA_HASH_TABLE_SIZE; ++i) {
    for (bh = bhTab[i]; bh; bh = bh->next) {
      if (bh->buf < p && p < (bh->buf + ULONGSIZE + bh->size))
        return bh;
    }
  }
  return 0;
}

/*
 * find_blk - return the block struct associated with the
 * pointer p.
 */
static BlkHdr* find_blk(const unsigned char* p)
{
  BlkHdr* bh = bhTab[hash_ptr(p)];
  for ( ; bh; bh = bh->next) {
    if (p == bh->buf)
      return bh;
  }
  return find_blk_exhaustive(p);
}

/*
 * make_blk - create a block header and add it to the hash table
 */
static int make_blk(unsigned char* p, size_t size, Location* loc)
{
  BlkHdr* bh;

  assert(0 != p && 0 < size && 0 != loc);
  if ((bh = (BlkHdr*)malloc(sizeof(BlkHdr))) != 0) {
    unsigned long h = hash_ptr(p);
    bh->ref = 0;
    bh->buf = p;
    bh->size = size;
    bh->location = loc;
    bh->next = bhTab[h];
    bhTab[h] = bh;
    ++bh->location->count;
    byteCount += size;
    ++blockCount;
  }
  return (0 != bh);
}

/*
 * free_blk - remove a block header and free it
 */
static void free_blk(const unsigned char* p)
{
  BlkHdr* bh_prev = 0;
  BlkHdr* bh;
  unsigned long h = hash_ptr(p);

  for (bh = bhTab[h]; bh; bh = bh->next) {
    if (p == bh->buf) {
      if (0 == bh_prev)
        bhTab[h] = bh->next;
      else
        bh_prev->next = bh->next;
      break;
    }
    bh_prev = bh;
  }
  /* if bh is NULL p was not allocated here */
  assert(0 != bh);
  assert(bh->location->count > 0);
  if (--bh->location->count == 0)
    free_location(bh->location);
  byteCount -= bh->size;
  --blockCount;
  SHRED_MEM(bh, sizeof(BlkHdr));
  free(bh);
}

/*
 * update_blk - update block info, rehash if pointers are different,
 * update location info if needed
 */
static void update_blk(unsigned char* p, unsigned char* np, size_t size, const char* file, int line)
{
  BlkHdr* bh;
  if (p != np) {
    BlkHdr* bh_prev = 0;
    unsigned long h = hash_ptr(p);

    /* remove the old entry from the hash table */
    for (bh = bhTab[h]; bh; bh = bh->next) {
      if (p == bh->buf) {
        if (0 == bh_prev)
          bhTab[h] = bh->next;
        else
          bh_prev->next = bh->next;
        /* put it back in the hash table at hash(np) */
        h = hash_ptr(np);
        bh->next = bhTab[h];
        bhTab[h] = bh;
        break;
      }
      bh_prev = bh;
    }
  }
  else
    bh = find_blk(p);
  /* invalid ptr? */
  assert(0 != bh);
  byteCount -= bh->size;
  byteCount += size;
  bh->buf = np;
  bh->size = size;
  /* update location info */
  if (bh->location->file != file || bh->location->line != line) {
    if (--bh->location->count == 0)
      free_location(bh->location);
    if ((bh->location = find_location(file, line)) == 0) {
      if ((bh->location = add_location(file, line)) == 0)
        noMemFn();
    }
    assert(0 != bh->location);
    ++bh->location->count;
  }
}

/*
 * fda_sizeof - returns the size of block of memory pointed to by p
 */
size_t fda_sizeof(const void* p)
{
  BlkHdr* bh = find_blk(base_ptr(p));
  assert(0 != bh && (unsigned char*)base_ptr(p) == bh->buf);
  return bh->size;
}

/*
 * fda_clear_refs - clear referenced markers on all blocks
 */
void fda_clear_refs(void)
{
  int i;
  BlkHdr* bh;

  for (i = 0; i < FDA_HASH_TABLE_SIZE; ++i) {
    for (bh = bhTab[i]; bh; bh = bh->next)
      bh->ref = 0;
  }
}

/*
 * fda_set_ref - mark block as referenced
 */
void fda_set_ref(const void* p)
{
  BlkHdr* bh = find_blk(base_ptr(p));
  assert(0 != bh);
  bh->ref = 1;
}

/*
 * fda_assert_refs - scan for all blocks and check for null
 * ptrs and unreferenced (lost) blocks
 */
void fda_assert_refs(void)
{
  int i;
  BlkHdr* bh;

  for (i = 0; i < FDA_HASH_TABLE_SIZE; ++i) {
    for (bh = bhTab[i]; bh; bh = bh->next) {
      assert(0 != bh->buf && 0 < bh->size);
      assert(1 == bh->ref);
    }
  }
}

/*
 * valid_ptr - returns true if p points to allocated memory and
 * has at least size available
 */
int valid_ptr(const void* p, size_t size)
{
  BlkHdr* bh;
  assert(0 != p && 0 < size);
  bh = find_blk(base_ptr(p));
  /* check that there are at least size bytes available from p */
  assert(((const unsigned char*)p + size) <= (bh->buf + ULONGSIZE + bh->size));
  return 1;
}

/*
 * fda_enum_locations - calls enumfn to list file, line, and count
 * info for allocations, returns the number of locations found
 */
int fda_enum_locations(void (*enumfn)(const char*, int, int))
{
  int count = 0;
  Location* location;
  assert(0 != enumfn);
  for (location = locationList; location; location = location->next) {
    (*enumfn)(location->file, location->line, location->count);
    ++count;
  }
  return count;
}

/*
 * fda_enum_leaks - scan hash table for leaks and call enumfn to
 * report them.
 */
int fda_enum_leaks(void (*enumfn)(const char*, int, size_t, void*))
{
  int count = 0;
  BlkHdr* bh;
  int i;
  for (i = 0; i < FDA_HASH_TABLE_SIZE; ++i) {
    for (bh = bhTab[i]; bh; bh = bh->next) {
      if (0 == bh->ref) {
        (*enumfn)(bh->location->file, bh->location->line, bh->size, bh->buf);
        ++count;
      }
    }
  }
  return count;
}

/*
 * fda_malloc - allocate size chunk of memory and create debug
 * records for it.
 */
void* fda_malloc(size_t size, const char* file, int line)
{
  void* p;
  size_t blk_size;
  Location* location;

  assert(0 < size && 0 != file);
  /* 
   * Make sure that there is enough room for prefix/postfix 
   * and we get an aligned buffer 
   */
  blk_size = base_size(size);

  if ((p = malloc(blk_size)) == 0) {
    lowMemFn();
    if ((p = malloc(blk_size)) == 0)
      noMemFn();
  }
  /* don't allow malloc to fail */
  assert(0 != p);
  /* shred the memory and set bounds markers */
  SHRED_MEM(p, blk_size);
  *((unsigned long*)p) = hell;
  *((unsigned long*)((char*)p + blk_size - ULONGSIZE)) = hell;

  /* find the location or create a new one */
  if ((location = find_location(file, line)) == 0) {
    if ((location = add_location(file, line)) == 0) {
      free(p);
      noMemFn();
    }
  }
  /* don't allow noMemFn to return */
  assert(0 != location);
  if (!make_blk(p, size, location)) {
    free(p);
    p = 0;
    noMemFn();
  }
  /* don't allow noMemFn to return */
  assert(0 != p);
  return (void*)((char*)p + ULONGSIZE);
}

/*
 * fda_free - check chunk of memory for overruns and free it
 */
void fda_free(void* p)
{
  if (p) {
    size_t size;
    void* ptr = base_ptr(p);
    BlkHdr* bh = find_blk(ptr);
    /* ptr invalid? */
    assert(0 != bh && ptr == bh->buf);
    size = base_size(bh->size);
    /* buffer underflow? */
    assert(hell == *((unsigned long*) ptr));
    /* 
     * buffer overflow?
     * Note: it's possible to have up to 3 bytes of unchecked space 
     * between size and hell
     */
    assert(hell == *((unsigned long*)((char*)ptr + size - ULONGSIZE)));
    SHRED_MEM(ptr, size);
    free_blk(ptr);
  }
}

/*
 * fda_realloc - resize a buffer, force reallocation if new size is
 * larger than old size
 */
void* fda_realloc(void* p, size_t size, const char* file, int line)
{
  void* np;
  size_t old_size;
  size_t blk_size;
  /* don't allow malloc or free through realloc */
  assert(0 != p && 0 < size);
  old_size = fda_sizeof(p);
  
  if (size < old_size)
    SHRED_MEM((char*)p + size, old_size - size);
  else if (size > old_size) {
    void* t = fda_malloc(size, __FILE__, __LINE__);
    memcpy(t, p, old_size);
    fda_free(p);
    p = t;
  }
  blk_size = base_size(size);
  if ((np = realloc(base_ptr(p), blk_size)) == 0) {
    lowMemFn();
    if ((np = realloc(base_ptr(p), blk_size)) == 0)
      noMemFn();
  }
  /* don't allow noMemFn to return */
  assert(0 != np);
  update_blk(base_ptr(p), np, size, file, line);
  /* shred tail */
  if (size > old_size)
    SHRED_MEM((char*)np + ULONGSIZE + old_size, blk_size - (old_size + ULONGSIZE));
  /* Note: it's possible to have up to 3 bytes between size and hell */
  *((unsigned long*)((char*)np + blk_size - ULONGSIZE)) = hell;
  return p;
}

/*
 * fda_calloc - allocate 0 initialized buffer nelems * size length
 */
void* fda_calloc(size_t nelems, size_t size, const char* file, int line)
{
  void* p;
  assert(0 < nelems && 0 < size && 0 != file);
  size *= nelems;
  p = fda_malloc(size, file, line);
  memset(p, 0, size);
  return p;
}

/*
 * fda_dupstring - duplicates a string returns newly allocated string
 */
char* fda_dupstring(const char* src, const char* file, int line)
{
  char* p;
  assert(0 != src);
  p = (char*) fda_malloc(strlen(src) + 1, file, line);
  strcpy(p, src);
  return p;
}

#else /* !defined(_DEBUG) && !defined(DEBUG) */

void* Malloc(size_t size)
{
  void* p = malloc(size);
  if (0 == p) {
    lowMemFn();
    if ((p = malloc(size)) == 0)
      noMemFn();
  }
  return p;
}

void* Realloc(void* p, size_t size)
{
  void* np = realloc(p, size);
  if (0 == np) {
    lowMemFn();
    if ((np = realloc(p, size)) == 0)
      noMemFn();
  }
  return np;
}

void* Calloc(size_t nelems, size_t size)
{
  void* np = calloc(nelems, size);
  if (0 == np) {
    lowMemFn();
    if ((np = calloc(nelems, size)) == 0)
      noMemFn();
  }
  return np;
}

char* DupString(const char* src)
{
  char* p = Malloc(strlen(src) + 1);
  strcpy(p, src);
  return p;
}

#endif /* !defined(_DEBUG) && !defined(DEBUG) */