/*
 * Copyright 2003-2013, Axel Dörfler, axeld@pinc-software.de.
 * Copyright 2005-2013, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include <boot/heap.h>

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

#include <algorithm>

#include <boot/platform.h>
#include <util/OpenHashTable.h>


#if KDEBUG
#define DEBUG_ALLOCATIONS
	// if defined, allocated memory is filled with 0xcc and freed memory with 0xdeadbeef
#define DEBUG_MAX_HEAP_USAGE
	// if defined, the maximum heap usage is determined and printed before
	// entering the kernel
#define DEBUG_VALIDATE_HEAP_ON_FREE
	// if defined, the heap integrity is checked on every free
	// (must have DEBUG_MAX_HEAP_USAGE defined also)
#endif

#include <util/SimpleAllocator.h>


//#define TRACE_HEAP
#ifdef TRACE_HEAP
#	define TRACE(format...)	dprintf(format)
#else
#	define TRACE(format...)	do { } while (false)
#endif


const static size_t kAlignment = 8;
	// all memory chunks will be a multiple of this

const static size_t kDefaultHeapSize = (1024 + 512) * 1024;
	// default initial heap size, unless overridden by platform loader
const static size_t kLargeAllocationThreshold = 128 * 1024;
	// allocations of this size or larger are allocated separately


struct LargeAllocation {
	LargeAllocation()
	{
		fAddress = NULL;
		fSize = 0;
	}

	status_t Allocate(size_t size)
	{
		ssize_t actualSize = platform_allocate_heap_region(size, &fAddress);
		if (actualSize < 0)
			return actualSize;

		fSize = actualSize;
		return B_OK;
	}

	void Free()
	{
		platform_free_heap_region(fAddress, fSize);
	}

	void* Address() const
	{
		return fAddress;
	}

	size_t Size() const
	{
		return fSize;
	}

	LargeAllocation*& HashNext()
	{
		return fHashNext;
	}

private:
	void*				fAddress;
	size_t				fSize;
	LargeAllocation*	fHashNext;
};


struct LargeAllocationHashDefinition {
	typedef void*			KeyType;
	typedef	LargeAllocation	ValueType;

	size_t HashKey(void* key) const
	{
		return size_t(key) / kAlignment;
	}

	size_t Hash(LargeAllocation* value) const
	{
		return HashKey(value->Address());
	}

	bool Compare(void* key, LargeAllocation* value) const
	{
		return value->Address() == key;
	}

	LargeAllocation*& GetLink(LargeAllocation* value) const
	{
		return value->HashNext();
	}
};


typedef BOpenHashTable<LargeAllocationHashDefinition> LargeAllocationHashTable;


static void* sHeapBase;
static void* sHeapEnd;
static SimpleAllocator<kAlignment> sAllocator;

static LargeAllocationHashTable sLargeAllocations;


static void*
malloc_large(size_t size)
{
	LargeAllocation* allocation = new(std::nothrow) LargeAllocation;
	if (allocation == NULL) {
		dprintf("malloc_large(): Out of memory!\n");
		return NULL;
	}

	if (allocation->Allocate(size) != B_OK) {
		dprintf("malloc_large(): Out of memory!\n");
		delete allocation;
		return NULL;
	}

	sLargeAllocations.InsertUnchecked(allocation);
	TRACE("malloc_large(%lu) -> %p\n", size, allocation->Address());
	return allocation->Address();
}


static void
free_large(void* address)
{
	LargeAllocation* allocation = sLargeAllocations.Lookup(address);
	if (allocation == NULL)
		panic("free_large(%p): unknown allocation!\n", address);

	sLargeAllocations.RemoveUnchecked(allocation);
	allocation->Free();
	delete allocation;
}


//	#pragma mark -


void
heap_release()
{
	heap_print_statistics();

	LargeAllocation* allocation = sLargeAllocations.Clear(true);
	while (allocation != NULL) {
		LargeAllocation* next = allocation->HashNext();
		allocation->Free();
		allocation = next;
	}

	platform_free_heap_region(sHeapBase, (addr_t)sHeapEnd - (addr_t)sHeapBase);

	sHeapBase = sHeapEnd = NULL;
	memset((void*)&sAllocator, 0, sizeof(sAllocator));
	memset((void*)&sLargeAllocations, 0, sizeof(sLargeAllocations));
}


void
heap_print_statistics()
{
#ifdef DEBUG_MAX_HEAP_USAGE
	dprintf("maximum boot loader heap usage: %" B_PRIu32 ", currently used: %" B_PRIu32 "\n",
		sAllocator.MaxHeapUsage(), sAllocator.MaxHeapSize() - sAllocator.Available());
#endif
}


status_t
heap_init(stage2_args* args)
{
	if (args->heap_size == 0)
		args->heap_size = kDefaultHeapSize;

	void* base;
	ssize_t size = platform_allocate_heap_region(args->heap_size, &base);
	if (size < 0)
		return B_ERROR;

	sHeapBase = base;
	sHeapEnd = (void*)((addr_t)base + size);

	sAllocator.AddChunk(sHeapBase, size);

	if (sLargeAllocations.Init(64) != B_OK)
		return B_NO_MEMORY;

	return B_OK;
}


uint32
heap_available()
{
	return sAllocator.Available();
}


void*
malloc(size_t size)
{
	if (sHeapBase == NULL)
		return NULL;

	if (size >= kLargeAllocationThreshold)
		return malloc_large(size);

	void* allocated = sAllocator.Allocate(size);
	if (allocated == NULL) {
		if (size == 0)
			return allocated;

		if (size > sAllocator.Available()) {
			dprintf("malloc(): Out of memory allocating a block of %ld bytes, "
				"only %" B_PRId32 " left!\n", size, sAllocator.Available());
			return NULL;
		}

		dprintf("malloc(): Out of memory allocating a block of %ld bytes!\n", size);
		return NULL;
	}

	TRACE("malloc(%lu) -> %p\n", size, allocated);
	return allocated;
}


void*
realloc(void* oldBuffer, size_t newSize)
{
	if (newSize == 0) {
		TRACE("realloc(%p, %lu) -> NULL\n", oldBuffer, newSize);
		free(oldBuffer);
		return NULL;
	}

	size_t oldSize = 0;
	if (oldBuffer != NULL) {
		if (oldBuffer >= sHeapBase && oldBuffer < sHeapEnd) {
			oldSize = sAllocator.UsableSize(oldBuffer);
		} else {
			LargeAllocation* allocation = sLargeAllocations.Lookup(oldBuffer);
			if (allocation == NULL) {
				panic("realloc(%p, %zu): unknown allocation!\n", oldBuffer,
					newSize);
			}

			oldSize = allocation->Size();
		}

		// Check if the old buffer still fits, and if it makes sense to keep it.
		if (oldSize >= newSize
			&& (oldSize < 128 || newSize > oldSize / 3)) {
			TRACE("realloc(%p, %lu) old buffer is large enough\n",
				oldBuffer, newSize);
			return oldBuffer;
		}
	}

	void* newBuffer = malloc(newSize);
	if (newBuffer == NULL)
		return NULL;

	if (oldBuffer != NULL) {
		memcpy(newBuffer, oldBuffer, std::min(oldSize, newSize));
		free(oldBuffer);
	}

	TRACE("realloc(%p, %lu) -> %p\n", oldBuffer, newSize, newBuffer);
	return newBuffer;
}


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

	return address;
}


void
free(void* allocated)
{
	if (allocated == NULL)
		return;

	TRACE("free(%p)\n", allocated);

	if (allocated < sHeapBase || allocated >= sHeapEnd) {
		free_large(allocated);
		return;
	}

	sAllocator.Free(allocated);
}