/*
 * Copyright 2015, Michael Lotz <mmlr@mlotz.ch>.
 * Distributed under the terms of the MIT License.
 */


#include "malloc_debug_api.h"

#include <malloc.h>
#include <string.h>

#include <stdio.h>
#include <stdlib.h>


static heap_implementation* sCurrentHeap = NULL;


// #pragma mark - Heap Debug API


extern "C" status_t
heap_debug_start_wall_checking(int msInterval)
{
	if (sCurrentHeap->start_wall_checking != NULL)
		return sCurrentHeap->start_wall_checking(msInterval);

	return B_NOT_SUPPORTED;
}


extern "C" status_t
heap_debug_stop_wall_checking()
{
	if (sCurrentHeap->stop_wall_checking != NULL)
		return sCurrentHeap->stop_wall_checking();

	return B_NOT_SUPPORTED;
}


extern "C" void
heap_debug_set_paranoid_validation(bool enabled)
{
	if (sCurrentHeap->set_paranoid_validation != NULL)
		sCurrentHeap->set_paranoid_validation(enabled);
}


extern "C" void
heap_debug_set_memory_reuse(bool enabled)
{
	if (sCurrentHeap->set_memory_reuse != NULL)
		sCurrentHeap->set_memory_reuse(enabled);
}


extern "C" void
heap_debug_set_debugger_calls(bool enabled)
{
	if (sCurrentHeap->set_debugger_calls != NULL)
		sCurrentHeap->set_debugger_calls(enabled);
}


extern "C" void
heap_debug_set_default_alignment(size_t defaultAlignment)
{
	if (sCurrentHeap->set_default_alignment != NULL)
		sCurrentHeap->set_default_alignment(defaultAlignment);
}


extern "C" void
heap_debug_validate_heaps()
{
	if (sCurrentHeap->validate_heaps != NULL)
		sCurrentHeap->validate_heaps();
}


extern "C" void
heap_debug_validate_walls()
{
	if (sCurrentHeap->validate_walls != NULL)
		sCurrentHeap->validate_walls();
}


extern "C" void
heap_debug_dump_allocations(bool statsOnly, thread_id thread)
{
	if (sCurrentHeap->dump_allocations != NULL)
		sCurrentHeap->dump_allocations(statsOnly, thread);
}


extern "C" void
heap_debug_dump_heaps(bool dumpAreas, bool dumpBins)
{
	if (sCurrentHeap->dump_heaps != NULL)
		sCurrentHeap->dump_heaps(dumpAreas, dumpBins);
}


extern "C" void *
heap_debug_malloc_with_guard_page(size_t size)
{
	if (sCurrentHeap->malloc_with_guard_page != NULL)
		return sCurrentHeap->malloc_with_guard_page(size);

	return NULL;
}


extern "C" status_t
heap_debug_get_allocation_info(void *address, size_t *size,
	thread_id *thread)
{
	if (sCurrentHeap->get_allocation_info != NULL)
		return sCurrentHeap->get_allocation_info(address, size, thread);

	return B_NOT_SUPPORTED;
}


extern "C" status_t
heap_debug_set_dump_allocations_on_exit(bool enabled)
{
	if (sCurrentHeap->set_dump_allocations_on_exit != NULL)
		return sCurrentHeap->set_dump_allocations_on_exit(enabled);

	return B_NOT_SUPPORTED;
}


extern "C" status_t
heap_debug_set_stack_trace_depth(size_t stackTraceDepth)
{
	if (sCurrentHeap->set_stack_trace_depth != NULL)
		return sCurrentHeap->set_stack_trace_depth(stackTraceDepth);

	return B_NOT_SUPPORTED;
}


// #pragma mark - Init


extern "C" status_t
__init_heap(void)
{
	const char *mode = getenv("MALLOC_DEBUG");
	if (mode == NULL || strchr(mode, 'g') == NULL)
		sCurrentHeap = &__mallocDebugHeap;
	else
		sCurrentHeap = &__mallocGuardedHeap;

	status_t result = sCurrentHeap->init();
	if (result != B_OK)
		return result;

	if (mode != NULL) {
		if (strchr(mode, 'p') != NULL)
			heap_debug_set_paranoid_validation(true);
		if (strchr(mode, 'r') != NULL)
			heap_debug_set_memory_reuse(false);
		if (strchr(mode, 'e') != NULL)
			heap_debug_set_dump_allocations_on_exit(true);

		size_t defaultAlignment = 0;
		const char *argument = strchr(mode, 'a');
		if (argument != NULL
			&& sscanf(argument, "a%" B_SCNuSIZE, &defaultAlignment) == 1) {
			heap_debug_set_default_alignment(defaultAlignment);
		}

		size_t stackTraceDepth = 0;
		argument = strchr(mode, 's');
		if (argument != NULL
			&& sscanf(argument, "s%" B_SCNuSIZE, &stackTraceDepth) == 1) {
			heap_debug_set_stack_trace_depth(stackTraceDepth);
		}

		int wallCheckInterval = 0;
		argument = strchr(mode, 'w');
		if (argument != NULL
			&& sscanf(argument, "w%d", &wallCheckInterval) == 1) {
			heap_debug_start_wall_checking(wallCheckInterval);
		}
	}

	return B_OK;
}


extern "C" void
__heap_terminate_after()
{
	if (sCurrentHeap->terminate_after != NULL)
		sCurrentHeap->terminate_after();
}


extern "C" void
__heap_before_fork(void)
{
}


extern "C" void
__heap_after_fork_child(void)
{
}


extern "C" void
__heap_after_fork_parent(void)
{
}


extern "C" void
__heap_thread_init(void)
{
}


extern "C" void
__heap_thread_exit(void)
{
}


// #pragma mark - Public API


extern "C" void*
memalign(size_t alignment, size_t size)
{
	return sCurrentHeap->memalign(alignment, size);
}


extern "C" void*
aligned_alloc(size_t alignment, size_t size)
{
	if ((size % alignment) != 0)
		return NULL;

	return sCurrentHeap->memalign(alignment, size);
}


extern "C" void*
malloc(size_t size)
{
	return sCurrentHeap->malloc(size);
}


extern "C" void
free(void* address)
{
	sCurrentHeap->free(address);
}


extern "C" void*
realloc(void* address, size_t newSize)
{
	return sCurrentHeap->realloc(address, newSize);
}


extern "C" void*
calloc(size_t numElements, size_t size)
{
	if (sCurrentHeap->calloc != NULL)
		return sCurrentHeap->calloc(numElements, size);

	void* address = malloc(numElements * size);
	if (address != NULL)
		memset(address, 0, numElements * size);

	return address;
}


extern "C" void*
valloc(size_t size)
{
	if (sCurrentHeap->valloc != NULL)
		return sCurrentHeap->valloc(size);

	return memalign(B_PAGE_SIZE, size);
}


extern "C" int
posix_memalign(void **pointer, size_t alignment, size_t size)
{
	if (sCurrentHeap->posix_memalign != NULL)
		return sCurrentHeap->posix_memalign(pointer, alignment, size);

	if (!is_valid_alignment(alignment))
		return EINVAL;

	*pointer = memalign(alignment, size);
	if (*pointer == NULL)
		return ENOMEM;

	return 0;
}


extern "C" size_t
malloc_usable_size(void *ptr)
{
	size_t size;
	thread_id thread;

	if (ptr == NULL)
		return 0;

	status_t res = sCurrentHeap->get_allocation_info(ptr, &size, &thread);

	if (res != B_OK)
		return 0;

	return size;
}