/* Leaker.h - a simple program to track C memory allocations
** see left404.com/2011/03/19/leaker for more information.
**
** Copyright (C) 3/19/2011  Dara Hazeghi
**
** This software may be freely modified and redistributed.  It has no warranty.
** Special thanks to Marshall Thomas for many of the ideas contained therein.
*/

#ifndef _LEAKER_H
#define _LEAKER_H

#include <stdlib.h>
#include <string.h>
#include <assert.h>

#define START_SIZE	128 /* initial table size */

typedef struct { void *ptr; const char *filename; int line; } LEAK_T;
typedef struct { LEAK_T *table; int size; int count; } LEAKTABLE_T;

/* global allocation table */
static LEAKTABLE_T _leaker = { NULL, 0, 0 };

static void _Leaker_Init(void);
static void _Leaker_Add(void *ptr, const char *filename, int line);
static void *_Leaker_Remove(void *ptr);
static int _Leaker_Find(void *ptr);
void *_Leaker_Alloc(int size, const char *filename, int line);
void *_Leaker_Realloc(void *ptr, int size, const char *filename, int line);
void _Leaker_Dealloc(void *ptr, const char *filename, int line);
void _Leaker_Report(void);

/* initialize allocation table */
static void _Leaker_Init(void)
{
	_leaker.table = (LEAK_T *)malloc(sizeof(LEAK_T)*START_SIZE);

	assert(_leaker.table);

	_leaker.size = START_SIZE;

	atexit(_Leaker_Report);
}

/* insert new entry (address + line) in allocation table */
static void _Leaker_Add(void *ptr, const char *filename, int line)
{
	if(!_leaker.table) _Leaker_Init();

	if(_leaker.size == _leaker.count)
	{
		_leaker.table = (LEAK_T *)realloc(_leaker.table, sizeof(LEAK_T)*_leaker.size*2);
		assert(_leaker.table);

		_leaker.size *= 2;
	}

	_leaker.table[_leaker.count].ptr = ptr;
	_leaker.table[_leaker.count].filename = filename;
	_leaker.table[_leaker.count].line = line;

	_leaker.count++;
}

/* remove and return existing address from allocation table */
static void *_Leaker_Remove(void *ptr)
{
	int found = _Leaker_Find(ptr);

	if(found != -1)
	{
		_leaker.table[found] = _leaker.table[--_leaker.count];
		return ptr;
	}

	else
		return NULL;
}

/* return the index of given address in the table, or -1 if it is not found */
static int _Leaker_Find(void *ptr)
{
	int i;

	for(i = 0; i < _leaker.count; i++)
	{
		if(_leaker.table[i].ptr == ptr) return i;
	}
	return -1;
}

/* custom allocation function overriding malloc and calloc */
void *_Leaker_Alloc(int size, const char *filename, int line)
{
	void *ptr = malloc(size);

	assert(ptr);

	memset(ptr, 0, size);

	_Leaker_Add(ptr, filename, line);
	return ptr;
}

/* custom reallocation function overriding realloc */
void *_Leaker_Realloc(void *ptr, int size, const char *filename, int line)
{
	int found;

	if(!ptr || !_leaker.table || (found = _Leaker_Find(ptr)) == -1)
	{
		if(!_leaker.table || ptr)
			fprintf(stderr, "%s:%d could not reallocate %P, address not in table.\n",
				filename, line, ptr);
		return _Leaker_Alloc(size, filename, line);
	}

	else
	{
		ptr = realloc(ptr, size);

		assert(ptr);

		_leaker.table[found].ptr = ptr;
		return ptr;
	}	
}

/* custom deallocation function overriding free */
void _Leaker_Dealloc(void *ptr, const char *filename, int line)
{
	if(!_leaker.table || !_Leaker_Remove(ptr))
	{
		fprintf(stderr, "%s:%d could not deallocate %P, address not in table.\n",
			filename, line, ptr);
	}
	else
		free(ptr);
}

/* function to report on memory leaks, deallocate remaining memory */
void _Leaker_Report(void)
{
	fprintf(stderr, "Leaker report:\n");

	if(_leaker.count == 0)
		fprintf(stderr, "No leaks!\n\n");

	else
	{
		int i;
		fprintf(stderr, "Leaks found!\n");

		for(i = 0; i < _leaker.count; i++)
		{
			fprintf(stderr, "%s:%d memory allocated at address %P was never freed.\n",
				_leaker.table[i].filename, _leaker.table[i].line, _leaker.table[i].ptr);
			free(_leaker.table[i].ptr);
		}
		fprintf(stderr, "\n");
	}
	if(_leaker.table) free(_leaker.table);
}

/* preprocessor magic to override built-in allocation functions with our own */
#define malloc(size)		_Leaker_Alloc(size, __FILE__, __LINE__)
#define calloc(n, size)		_Leaker_Alloc(n*size, __FILE__, __LINE__)
#define free(ptr)			_Leaker_Dealloc(ptr, __FILE__, __LINE__)
#define realloc(ptr, size)	_Leaker_Realloc(ptr, size, __FILE__, __LINE__);

#endif