#define _GNU_SOURCE

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

#define CACHE_LINE 64

static char tmpbuff[4096];
static unsigned long tmppos = 0;
static unsigned long tmpallocs = 0;

/* =========================================================
 * interception points
 */
static void* (*real_malloc)  (size_t)         = NULL;
static void  (*real_free)    (void*)          = NULL;
static void* (*real_calloc)  (size_t, size_t) = NULL;
static void* (*real_realloc) (void*, size_t)  = NULL;

static void init(void) {
	real_malloc  = dlsym(RTLD_NEXT, "malloc");
	real_free    = dlsym(RTLD_NEXT, "free");
	real_calloc  = dlsym(RTLD_NEXT, "calloc");
	real_realloc = dlsym(RTLD_NEXT, "realloc");

	if (!real_malloc || !real_free || !real_calloc || !real_realloc) {
		fprintf(stderr, "Error in `dlsym`: %s\n", dlerror());
	}
}

static inline size_t alignUp(size_t size, size_t alignment)
{
	if (alignment > 0) {
		assert(((alignment - 1) & alignment) == 0);
		size_t mask = (alignment - 1);
		return (size + mask) & ~mask;
	}
	return size;
}

void *malloc(size_t size)
{
	static int initializing = 0;
	size = alignUp(size, 64);

	if (real_malloc == NULL)
	{
		if (!initializing)
		{
			initializing = 1;
			init();
			initializing = 0;
		}
		else
		{
			if (tmppos + size < sizeof(tmpbuff))
			{
				void *retptr = tmpbuff + tmppos;
				tmppos += size;
				++tmpallocs;
				return retptr;
			}
			else
			{
				fprintf(stderr, "%d in %d allocs\n", tmppos, tmpallocs);
				fprintf(stderr, "jcheck: too much memory requested during initialisation - increase tmpbuff size\n");
				exit(1);
			}
		}
	}

	return real_malloc(size);
}

void free(void *ptr)
{
	// something wrong if we call free before one of the allocators!
	if (real_malloc == NULL)
		init();

	if (!(ptr >= (void*) tmpbuff && ptr <= (void*)(tmpbuff + tmppos)))
	{
		real_free(ptr);
	}
}

void* calloc(size_t nmemb, size_t size)
{
	// Overflows are not detected anymore
	size_t total = alignUp(nmemb * size, 64);
	if (real_calloc == NULL)
	{
		void *ptr = malloc(total);
		if (ptr)
			memset(ptr, 0, total);
		return ptr;
	}

	// Is this always correct ?
	return real_calloc(1, total);
}

void* realloc(void *ptr, size_t size)
{
	size = alignUp(size, 64);
	if (real_realloc == NULL)
	{
		void *nptr = malloc(size);
		if (nptr && ptr)
		{
			memmove(nptr, ptr, size);
			free(ptr);
		}
		return nptr;
	}

	return real_realloc(ptr, size);
}