/*
 * Copyright 2009-2010, Ingo Weinhold, ingo_weinhold@gmx.de
 * Distributed under the terms of the MIT License.
 */


#include <debug_heap.h>

#include <new>

#include <util/AutoLock.h>
#include <vm/vm.h>


#define INITIAL_DEBUG_HEAP_SIZE B_PAGE_SIZE

static char sInitialHeap[INITIAL_DEBUG_HEAP_SIZE] __attribute__ ((aligned (8)));
static void* sHeapBase = sInitialHeap;
static size_t sHeapSize = INITIAL_DEBUG_HEAP_SIZE;

const kdebug_alloc_t kdebug_alloc = {};


struct allocation_header {
        uint32  size : 31;      // size in allocation_header units
        bool    free : 1;
        uint32  previous;
};

struct free_entry : allocation_header {
        uint32  previous_free;
        uint32  next_free;
};


struct DebugAllocPool {
        void Init(void* heap, size_t heapSize)
        {
                fParent = NULL;
                fChild = NULL;

                uint32 size = heapSize / 8;
                fBase = (allocation_header*)heap - 1;
                fEnd = size + 1;
                fFirstFree = 0;
                fLastFree = 0;

                // add free entry spanning the whole area
                fBase[1].size = size - 1;
                fBase[1].previous = 0;
                _InsertFreeEntry(1);
        }

        DebugAllocPool* CreateChildPool()
        {
                // do we already have a child pool?
                if (fChild != NULL)
                        return NULL;

                // create the pool object
                DebugAllocPool* pool
                        = (DebugAllocPool*)Allocate(sizeof(DebugAllocPool));
                if (pool == NULL)
                        return NULL;

                // do we have enough free space?
                if (fLastFree == 0 || fBase[fLastFree].size < 2) {
                        Free(pool);
                        return NULL;
                }

                allocation_header* header = &fBase[fLastFree];
                _RemoveFreeEntry(fLastFree);

                pool->Init(header + 1, header->size * 8);
                pool->fParent = this;

                return fChild = pool;
        }

        void Destroy()
        {
                if (fParent != NULL) {
                        fParent->fChild = NULL;
                        fParent->Free(fBase + 1);
                }
        }

        DebugAllocPool* Parent() const
        {
                return fParent;
        }

        void* Allocate(size_t size)
        {
                size = (size + 7) / 8;
                uint32 index = fFirstFree;
                while (index != 0 && fBase[index].size < size)
                        index = ((free_entry*)&fBase[index])->next_free;

                if (index == 0)
                        return NULL;

                _RemoveFreeEntry(index);

                // if the entry is big enough, we split it
                if (fBase[index].size - size >= 2) {
                        uint32 next = index + 1 + size;
                        uint32 nextNext = index + 1 + fBase[index].size;
                        fBase[next].size = fBase[index].size - size - 1;
                        fBase[next].previous = index;
                        fBase[index].size = size;
                        _InsertFreeEntry(next);

                        if (nextNext < fEnd)
                                fBase[nextNext].previous = next;
                }

                return &fBase[index + 1];
        }

        void Free(void* address)
        {
                // check address
                if (((addr_t)address & 7) != 0 || address <= fBase + 1
                        || address >= fBase + fEnd) {
                        kprintf("DebugAllocPool::Free(%p): bad address\n", address);
                        return;
                }

                // get header
                allocation_header* header = (allocation_header*)address - 1;
                uint32 index = header - fBase;
                if (header->free) {
                        kprintf("DebugAllocPool::Free(%p): double free\n", address);
                        return;
                }

                uint32 next = index + 1 + header->size;

                // join with previous, if possible
                if (index > 1 && fBase[header->previous].free) {
                        uint32 previous = header->previous;
                        _RemoveFreeEntry(previous);

                        fBase[previous].size += 1 + header->size;
                        fBase[next].previous = previous;

                        index = previous;
                        header = fBase + index;
                }

                // join with next, if possible
                if (next < fEnd && fBase[next].free) {
                        _RemoveFreeEntry(next);

                        header->size += 1 + fBase[next].size;

                        uint32 nextNext = index + 1 + header->size;
                        if (nextNext < fEnd)
                                fBase[nextNext].previous = index;
                }

                _InsertFreeEntry(index);
        }

private:
        void _InsertFreeEntry(uint32 index)
        {
                // find the insertion point -- list is sorted by ascending size
                uint32 size = fBase[index].size;
                uint32 next = fFirstFree;
                while (next != 0 && size > fBase[next].size)
                        next = ((free_entry*)&fBase[next])->next_free;

                // insert
                uint32 previous;
                if (next != 0) {
                        previous = ((free_entry*)&fBase[next])->previous_free;
                        ((free_entry*)&fBase[next])->previous_free = index;
                } else {
                        previous = fLastFree;
                        fLastFree = index;
                }

                if (previous != 0)
                        ((free_entry*)&fBase[previous])->next_free = index;
                else
                        fFirstFree = index;

                ((free_entry*)&fBase[index])->previous_free = previous;
                ((free_entry*)&fBase[index])->next_free = next;

                fBase[index].free = true;
        }

        void _RemoveFreeEntry(uint32 index)
        {
                uint32 previous = ((free_entry*)&fBase[index])->previous_free;
                uint32 next = ((free_entry*)&fBase[index])->next_free;

                if (previous != 0)
                        ((free_entry*)&fBase[previous])->next_free = next;
                else
                        fFirstFree = next;

                if (next != 0)
                        ((free_entry*)&fBase[next])->previous_free = previous;
                else
                        fLastFree = previous;

                fBase[index].free = false;
        }

private:
        DebugAllocPool*         fParent;
        DebugAllocPool*         fChild;
        allocation_header*      fBase;          // actually base - 1, so that index 0 is
                                                                        // invalid
        uint32                          fEnd;
        uint32                          fFirstFree;
        uint32                          fLastFree;
};


static DebugAllocPool* sCurrentPool;
static DebugAllocPool sInitialPool;


debug_alloc_pool*
create_debug_alloc_pool()
{
        if (sCurrentPool == NULL) {
                sInitialPool.Init(sHeapBase, sHeapSize);
                sCurrentPool = &sInitialPool;
                return sCurrentPool;
        }

        DebugAllocPool* pool = sCurrentPool->CreateChildPool();
        if (pool == NULL)
                return NULL;

        sCurrentPool = pool;
        return sCurrentPool;
}


void
delete_debug_alloc_pool(debug_alloc_pool* pool)
{
        if (pool == NULL || sCurrentPool == NULL)
                return;

        // find the pool in the hierarchy
        DebugAllocPool* otherPool = sCurrentPool;
        while (otherPool != NULL && otherPool != pool)
                otherPool = otherPool->Parent();

        if (otherPool == NULL)
                return;

        // destroy the pool
        sCurrentPool = pool->Parent();
        pool->Destroy();

        if (pool != &sInitialPool)
                debug_free(pool);
}


void*
debug_malloc(size_t size)
{
        if (sCurrentPool == NULL)
                return NULL;

        return sCurrentPool->Allocate(size);
}


void*
debug_calloc(size_t num, size_t size)
{
        size_t allocationSize = num * size;
        void* allocation = debug_malloc(allocationSize);
        if (allocation == NULL)
                return NULL;

        memset(allocation, 0, allocationSize);
        return allocation;
}


void
debug_free(void* address)
{
        if (address != NULL && sCurrentPool != NULL)
                sCurrentPool->Free(address);
}


void
debug_heap_init()
{
        // create the heap area
        void* base;
        virtual_address_restrictions virtualRestrictions = {};
        virtualRestrictions.address_specification = B_ANY_KERNEL_ADDRESS;
        physical_address_restrictions physicalRestrictions = {};
        area_id area = create_area_etc(B_SYSTEM_TEAM, "kdebug heap", KDEBUG_HEAP,
                B_FULL_LOCK, B_KERNEL_READ_AREA | B_KERNEL_WRITE_AREA,
                CREATE_AREA_DONT_WAIT, 0, &virtualRestrictions, &physicalRestrictions,
                (void**)&base);
        if (area < 0)
                return;

        // switch from the small static buffer to the area
        InterruptsLocker locker;
        sHeapBase = base;
        sHeapSize = KDEBUG_HEAP;
}