/*
 * Copyright 2009, Axel Dörfler, axeld@pinc-software.de.
 * Distributed under the terms of the MIT License.
 */


/*!	A simple allocator that works directly on an area, based on the boot
	loader's heap. See there for more information about its inner workings.
*/


#include <RealtimeAlloc.h>

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

#include <OS.h>

#include <locks.h>
#include <kernel/util/DoublyLinkedList.h>


//#define TRACE_RTM
#ifdef TRACE_RTM
#	define TRACE(x...) printf(x);
#else
#	define TRACE(x...) ;
#endif


class FreeChunk {
public:
			void				SetTo(size_t size, FreeChunk* next);

			uint32				Size() const;
			uint32				CompleteSize() const { return fSize; }

			FreeChunk*			Next() const { return fNext; }
			void				SetNext(FreeChunk* next) { fNext = next; }

			FreeChunk*			Split(uint32 splitSize);
			bool				IsTouching(FreeChunk* link);
			FreeChunk*			Join(FreeChunk* link);
			void				Remove(rtm_pool* pool,
									FreeChunk* previous = NULL);
			void				Enqueue(rtm_pool* pool);

			void*				AllocatedAddress() const;
	static	FreeChunk*			SetToAllocated(void* allocated);
	static	addr_t				NextOffset() { return sizeof(size_t); }

private:
			size_t				fSize;
			FreeChunk*			fNext;
};


struct rtm_pool : DoublyLinkedListLinkImpl<rtm_pool> {
	area_id		area;
	void*		heap_base;
	size_t		max_size;
	size_t		available;
	FreeChunk	free_anchor;
	mutex		lock;

	bool Contains(void* buffer) const;
	void Free(void* buffer);
};

typedef DoublyLinkedList<rtm_pool> PoolList;


const static uint32 kAlignment = 256;
	// all memory chunks will be a multiple of this

static mutex sPoolsLock = MUTEX_INITIALIZER("rtm pools");
static PoolList sPools;


void
FreeChunk::SetTo(size_t size, FreeChunk* next)
{
	fSize = size;
	fNext = next;
}


/*!	Returns the amount of bytes that can be allocated
	in this chunk.
*/
uint32
FreeChunk::Size() const
{
	return fSize - FreeChunk::NextOffset();
}


/*!	Splits the upper half at the requested location
	and returns it.
*/
FreeChunk*
FreeChunk::Split(uint32 splitSize)
{
	splitSize = (splitSize - 1 + kAlignment) & ~(kAlignment - 1);

	FreeChunk* chunk
		= (FreeChunk*)((uint8*)this + FreeChunk::NextOffset() + splitSize);
	chunk->fSize = fSize - splitSize - FreeChunk::NextOffset();
	chunk->fNext = fNext;

	fSize = splitSize + FreeChunk::NextOffset();

	return chunk;
}


/*!	Checks if the specified chunk touches this chunk, so
	that they could be joined.
*/
bool
FreeChunk::IsTouching(FreeChunk* chunk)
{
	return chunk
		&& (((uint8*)this + fSize == (uint8*)chunk)
			|| (uint8*)chunk + chunk->fSize == (uint8*)this);
}


/*!	Joins the chunk to this chunk and returns the pointer
	to the new chunk - which will either be one of the
	two chunks.
	Note, the chunks must be joinable, or else this method
	doesn't work correctly. Use FreeChunk::IsTouching()
	to check if this method can be applied.
*/
FreeChunk*
FreeChunk::Join(FreeChunk* chunk)
{
	if (chunk < this) {
		chunk->fSize += fSize;
		chunk->fNext = fNext;

		return chunk;
	}

	fSize += chunk->fSize;
	fNext = chunk->fNext;

	return this;
}


void
FreeChunk::Remove(rtm_pool* pool, FreeChunk* previous)
{
	if (previous == NULL) {
		// find the previous chunk in the list
		FreeChunk* chunk = pool->free_anchor.fNext;

		while (chunk != NULL && chunk != this) {
			previous = chunk;
			chunk = chunk->fNext;
		}

		if (chunk == NULL)
			return;
	}

	previous->fNext = fNext;
	fNext = NULL;
}


void
FreeChunk::Enqueue(rtm_pool* pool)
{
	FreeChunk* chunk = pool->free_anchor.fNext;
	FreeChunk* last = &pool->free_anchor;
	while (chunk && chunk->Size() < fSize) {
		last = chunk;
		chunk = chunk->fNext;
	}

	fNext = chunk;
	last->fNext = this;
}


void*
FreeChunk::AllocatedAddress() const
{
	return (void*)&fNext;
}


FreeChunk*
FreeChunk::SetToAllocated(void* allocated)
{
	return (FreeChunk*)((addr_t)allocated - FreeChunk::NextOffset());
}


// #pragma mark - rtm_pool


bool
rtm_pool::Contains(void* buffer) const
{
	return (addr_t)heap_base <= (addr_t)buffer
		&& (addr_t)heap_base - 1 + max_size >= (addr_t)buffer;
}


void
rtm_pool::Free(void* allocated)
{
	FreeChunk* freedChunk = FreeChunk::SetToAllocated(allocated);
	available += freedChunk->CompleteSize();

	// try to join the new free chunk with an existing one
	// it may be joined with up to two chunks

	FreeChunk* chunk = free_anchor.Next();
	FreeChunk* last = &free_anchor;
	int32 joinCount = 0;

	while (chunk) {
		if (chunk->IsTouching(freedChunk)) {
			// almost "insert" it into the list before joining
			// because the next pointer is inherited by the chunk
			freedChunk->SetNext(chunk->Next());
			freedChunk = chunk->Join(freedChunk);

			// remove the joined chunk from the list
			last->SetNext(freedChunk->Next());
			chunk = last;

			if (++joinCount == 2)
				break;
		}

		last = chunk;
		chunk = chunk->Next();
	}

	// enqueue the link at the right position; the
	// free link queue is ordered by size

	freedChunk->Enqueue(this);
}


// #pragma mark -


static rtm_pool*
pool_for(void* buffer)
{
	MutexLocker _(&sPoolsLock);

	PoolList::Iterator iterator = sPools.GetIterator();
	while (rtm_pool* pool = iterator.Next()) {
		if (pool->Contains(buffer))
			return pool;
	}

	return NULL;
}


// #pragma mark - public API


status_t
rtm_create_pool(rtm_pool** _pool, size_t totalSize, const char* name)
{
	rtm_pool* pool = (rtm_pool*)malloc(sizeof(rtm_pool));
	if (pool == NULL)
		return B_NO_MEMORY;

	if (name != NULL)
		mutex_init_etc(&pool->lock, name, MUTEX_FLAG_CLONE_NAME);
	else
		mutex_init(&pool->lock, "realtime pool");

	// Allocate enough space for at least one allocation over \a totalSize
	pool->max_size = (totalSize + sizeof(FreeChunk) - 1 + B_PAGE_SIZE)
		& ~(B_PAGE_SIZE - 1);

	area_id area = create_area(name, &pool->heap_base, B_ANY_ADDRESS,
		pool->max_size, B_LAZY_LOCK, B_READ_AREA | B_WRITE_AREA);
	if (area < 0) {
		mutex_destroy(&pool->lock);
		free(pool);
		return area;
	}

	pool->area = area;
	pool->available = pool->max_size - FreeChunk::NextOffset();

	// declare the whole heap as one chunk, and add it
	// to the free list

	FreeChunk* chunk = (FreeChunk*)pool->heap_base;
	chunk->SetTo(pool->max_size, NULL);

	pool->free_anchor.SetTo(0, chunk);

	*_pool = pool;

	MutexLocker _(&sPoolsLock);
	sPools.Add(pool);
	return B_OK;
}


status_t
rtm_delete_pool(rtm_pool* pool)
{
	if (pool == NULL)
		return B_BAD_VALUE;

	mutex_lock(&pool->lock);

	{
		MutexLocker _(&sPoolsLock);
		sPools.Remove(pool);
	}

	delete_area(pool->area);
	mutex_destroy(&pool->lock);
	free(pool);

	return B_OK;
}


void*
rtm_alloc(rtm_pool* pool, size_t size)
{
	if (pool == NULL)
		return malloc(size);

	if (pool->heap_base == NULL || size == 0)
		return NULL;

	MutexLocker _(&pool->lock);

	// align the size requirement to a kAlignment bytes boundary
	size = (size - 1 + kAlignment) & ~(size_t)(kAlignment - 1);

	if (size > pool->available) {
		TRACE("malloc(): Out of memory!\n");
		return NULL;
	}

	FreeChunk* chunk = pool->free_anchor.Next();
	FreeChunk* last = &pool->free_anchor;
	while (chunk && chunk->Size() < size) {
		last = chunk;
		chunk = chunk->Next();
	}

	if (chunk == NULL) {
		// could not find a free chunk as large as needed
		TRACE("malloc(): Out of memory!\n");
		return NULL;
	}

	if (chunk->Size() > size + sizeof(FreeChunk) + kAlignment) {
		// if this chunk is bigger than the requested size,
		// we split it to form two chunks (with a minimal
		// size of kAlignment allocatable bytes).

		FreeChunk* freeChunk = chunk->Split(size);
		last->SetNext(freeChunk);

		// re-enqueue the free chunk at the correct position
		freeChunk->Remove(pool, last);
		freeChunk->Enqueue(pool);
	} else {
		// remove the chunk from the free list

		last->SetNext(chunk->Next());
	}

	pool->available -= size + sizeof(size_t);

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


status_t
rtm_free(void* allocated)
{
	if (allocated == NULL)
		return B_OK;

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

	// find pool
	rtm_pool* pool = pool_for(allocated);
	if (pool == NULL) {
		free(allocated);
		return B_OK;
	}

	MutexLocker _(&pool->lock);
	pool->Free(allocated);
	return B_OK;
}


status_t
rtm_realloc(void** _buffer, size_t newSize)
{
	if (_buffer == NULL)
		return B_BAD_VALUE;

	TRACE("rtm_realloc(%p, %lu)\n", *_buffer, newSize);

	void* oldBuffer = *_buffer;

	// find pool
	rtm_pool* pool = pool_for(oldBuffer);
	if (pool == NULL) {
		void* buffer = realloc(oldBuffer, newSize);
		if (buffer != NULL) {
			*_buffer = buffer;
			return B_OK;
		}
		return B_NO_MEMORY;
	}

	MutexLocker _(&pool->lock);

	if (newSize == 0) {
		TRACE("realloc(%p, %lu) -> NULL\n", oldBuffer, newSize);
		pool->Free(oldBuffer);
		*_buffer = NULL;
		return B_OK;
	}

	size_t copySize = newSize;
	if (oldBuffer != NULL) {
		FreeChunk* oldChunk = FreeChunk::SetToAllocated(oldBuffer);

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

		if (copySize > oldChunk->Size())
			copySize = oldChunk->Size();
	}

	void* newBuffer = rtm_alloc(pool, newSize);
	if (newBuffer == NULL)
		return B_NO_MEMORY;

	if (oldBuffer != NULL) {
		memcpy(newBuffer, oldBuffer, copySize);
		pool->Free(oldBuffer);
	}

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


status_t
rtm_size_for(void* buffer)
{
	if (buffer == NULL)
		return 0;

	FreeChunk* chunk = FreeChunk::SetToAllocated(buffer);
	// TODO: we currently always return the actual chunk size, not the allocated
	// one
	return chunk->Size();
}


status_t
rtm_phys_size_for(void* buffer)
{
	if (buffer == NULL)
		return 0;

	FreeChunk* chunk = FreeChunk::SetToAllocated(buffer);
	return chunk->Size();
}


size_t
rtm_available(rtm_pool* pool)
{
	if (pool == NULL) {
		// whatever - might want to use system_info instead
		return 1024 * 1024;
	}

	return pool->available;
}


rtm_pool*
rtm_default_pool()
{
	// We always return NULL - the default pool will just use malloc()/free()
	return NULL;
}


#if 0
extern "C" {

// undocumented symbols that BeOS exports
status_t rtm_create_pool_etc(rtm_pool ** out_pool, size_t total_size, const char * name, int32 param4, int32 param5, ...);
void rtm_get_pool(rtm_pool *pool,void *data,int32 param3,int32 param4, ...);

}
#endif