#include "thread_local_malloc.h"

#include <cassert>
#include <malloc/malloc.h>
#include <sys/mman.h>
#include <unistd.h>
#include <pthread/pthread.h>
#include <thread>


static pthread_key_t threadLocalZoneKey = 0;


/**
 * This macro contains all the data that can be used to
 * generate the callbacks of a malloc_zone_t structure.
 *
 * This macro will call the macro functions A,B,C,D with the following parameters:
 *
 * A(x, y)
 *   x := the return type
 *   y := the malloc_zone_t callback name
 *
 * B(x, y)
 *   x := an argument type
 *   y := an argument name
 *
 * C(x, y)
 *   x := a trailing argument type
 *   y := a trailing argument name
 *   The difference between B and C is that C signals that
 *   this pair is the last one in the list of arguments.
 *   Thus e.g. the trailing comma can be ommited.
 *
 * D(x)
 *   x := the malloc_zone_t callback name (again)
 *
 * E(x)
 *   x := an argument name
 *
 * F(x)
 *   x := a trailing argument name
 *
 * B and C can be used to generate the arguments of a function.
 * E and F can be used to generate the matching parameters for a function.
 */
#define MAP(A, B, C, D, E, F)                                                                                                                                                                                           \
	A(size_t, size)             B(malloc_zone_t*, zone) C(const void*, ptr)                                                            D(size)               E(zone) F(ptr)                                             \
	A(void*, malloc)            B(malloc_zone_t*, zone) C(size_t, size)                                                                D(malloc)             E(zone) F(size)                                            \
	A(void*, calloc)            B(malloc_zone_t*, zone) B(size_t, num_items)   C(size_t, size)                                         D(calloc)             E(zone) E(num_items)   F(size)                             \
	A(void*, valloc)            B(malloc_zone_t*, zone) C(size_t, size)                                                                D(valloc)             E(zone) F(size)                                            \
	A(void, free)               B(malloc_zone_t*, zone) C(void*, ptr)                                                                  D(free)               E(zone) F(ptr)                                             \
	A(void*, realloc)           B(malloc_zone_t*, zone) B(void*, ptr)          C(size_t, size)                                         D(realloc)            E(zone) E(ptr)         F(size)                             \
	A(void, destroy)            C(malloc_zone_t*, zone)                                                                                D(destroy)            F(zone)                                                    \
	A(unsigned, batch_malloc)   B(malloc_zone_t*, zone) B(size_t, size)        B(void**, results)           C(unsigned, num_requested) D(batch_malloc)       E(zone) E(size)        E(results)         F(num_requested) \
	A(void, batch_free)         B(malloc_zone_t*, zone) B(void**, to_be_freed) C(unsigned, num_to_be_freed)                            D(batch_free)         E(zone) E(to_be_freed) F(num_to_be_freed)                  \
	A(void*, memalign)          B(malloc_zone_t*, zone) B(size_t, alignment)   C(size_t, size)                                         D(memalign)           E(zone) E(alignment)   F(size)                             \
	A(void, free_definite_size) B(malloc_zone_t*, zone) B(void*, ptr)          C(size_t, size)                                         D(free_definite_size) E(zone) E(ptr)         F(size)                             \
	A(size_t, pressure_relief)  B(malloc_zone_t*, zone) C(size_t, goal)                                                                D(pressure_relief)    E(zone) F(goal)                                            \

#define MAP_NIL(...)
#define MAP_B(X, Y) X Y,
#define MAP_C(X, Y) X Y
#define MAP_E(X) X,
#define MAP_F(X) X


/*
 * This generates the necessary function pointers to store the original callback functions.
 * E.g.: void*(*orig_malloc)(malloc_zone_t* zone, size_t size);
 */
#define A(X, Y) static X(*orig_##Y)(
#define D(X) );

MAP(A, MAP_B, MAP_C, D, MAP_NIL, MAP_NIL)

#undef A
#undef D


/*
 * This generates the patched functions.
 * E.g.: void* tl_malloc(malloc_zone_t* zone, size_t size) { ... }
 */
#define A(X, Y) static X tl_##Y(
#define D(X) ) {                                       \
	void* z = pthread_getspecific(threadLocalZoneKey); \
	if (z) zone = (malloc_zone_t*)z;                   \
	return orig_##X(
#define F(X) X); }

MAP(A, MAP_B, MAP_C, D, MAP_E, F)

#undef A
#undef D
#undef F


void tl_malloc_init() {
	pthread_key_create(&threadLocalZoneKey, [](void* data) {
		malloc_destroy_zone((malloc_zone_t*)data);
	});

	malloc_zone_t* zone = malloc_default_zone();

	vm_address_t pageStart = reinterpret_cast<vm_address_t>(zone) & static_cast<vm_size_t>(~(getpagesize() - 1));
	vm_size_t len = reinterpret_cast<vm_address_t>(zone) - pageStart + sizeof(malloc_zone_t);
	mprotect(reinterpret_cast<void*>(pageStart), len, PROT_READ | PROT_WRITE);

	/*
	 * Save the original callbacks in the function pointers.
	 */
#define A(X, Y) orig_##Y = zone->Y;
MAP(A, MAP_NIL, MAP_NIL, MAP_NIL, MAP_NIL, MAP_NIL)
#undef A

	/*
	 * Overwrite using the patched functions.
	 */
#define A(X, Y) zone->Y = tl_##Y;
MAP(A, MAP_NIL, MAP_NIL, MAP_NIL, MAP_NIL, MAP_NIL)
#undef A

	mprotect(reinterpret_cast<void*>(pageStart), len, PROT_READ);
}

void tl_malloc_register() {
	assert(threadLocalZoneKey != 0);

	void* zone = pthread_getspecific(threadLocalZoneKey);

	if (!zone) {
		zone = malloc_create_zone(0, 0);
		pthread_setspecific(threadLocalZoneKey, zone);
	}
}