/*
 * Copyright 2002-2007, Haiku Inc.
 * Distributed under the terms of the MIT License.
 */

/* Hoard: A Fast, Scalable, and Memory-Efficient Allocator
 *              for Shared-Memory Multiprocessors
 * Contact author: Emery Berger, http://www.cs.utexas.edu/users/emery
 *
 * Copyright (c) 1998-2000, The University of Texas at Austin.
 *
 * This library is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Library General Public License as
 * published by the Free Software Foundation, http://www.fsf.org.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 */

#include "config.h"
#include "threadheap.h"
#include "processheap.h"
#include "arch-specific.h"

#include <image.h>

#include <errno.h>
#include <string.h>

#include <errno_private.h>
#include <user_thread.h>

#include "tracing_config.h"

using namespace BPrivate;


#if USER_MALLOC_TRACING
#       define KTRACE(format...)        ktrace_printf(format)
#else
#       define KTRACE(format...)        do {} while (false)
#endif


#if HEAP_LEAK_CHECK
static block* sUsedList = NULL;
static hoardLockType sUsedLock = MUTEX_INITIALIZER("");


/*!
        Finds the closest symbol that comes before the given address.
*/
static status_t
get_symbol_for_address(void* address, char *imageBuffer, size_t imageBufferSize,
        char* buffer, size_t bufferSize, int32& offset)
{
        offset = -1;

        image_info info;
        int32 cookie = 0;
        while (get_next_image_info(0, &cookie, &info) == B_OK) {
                if (((addr_t)info.text > (addr_t)address
                                || (addr_t)info.text + info.text_size < (addr_t)address)
                        && ((addr_t)info.data > (addr_t)address
                                || (addr_t)info.data + info.data_size < (addr_t)address))
                        continue;

                char name[256];
                int32 index = 0;
                int32 nameLength = sizeof(name);
                int32 symbolType;
                void* location;
                while (get_nth_image_symbol(info.id, index, name, &nameLength,
                                &symbolType, &location) == B_OK) {
                        if ((addr_t)address >= (addr_t)location) {
                                // see if this is better than what we have
                                int32 newOffset = (addr_t)address - (addr_t)location;

                                if (offset == -1 || offset > newOffset) {
                                        const char* imageName = strrchr(info.name, '/');
                                        if (imageName != NULL)
                                                strlcpy(imageBuffer, imageName + 1, imageBufferSize);
                                        else
                                                strlcpy(imageBuffer, info.name, imageBufferSize);

                                        strlcpy(buffer, name, bufferSize);
                                        offset = newOffset;
                                }
                        }

                        nameLength = sizeof(name);
                        index++;
                }
        }

        return offset != -1 ? B_OK : B_ENTRY_NOT_FOUND;
}


static void
dump_block(block* b)
{
        printf("  %p, %ld bytes: call stack", b + 1, b->getAllocatedSize());

        for (int i = 0; i < HEAP_CALL_STACK_SIZE; i++) {
                if (b->getCallStack(i) != NULL) {
                        char image[256];
                        char name[256];
                        int32 offset;
                        if (get_symbol_for_address(b->getCallStack(i), image, sizeof(image),
                                        name, sizeof(name), offset) != B_OK) {
                                strcpy(name, "???");
                                offset = 0;
                        }

                        printf(": %p (%s:%s+0x%lx)", b->getCallStack(i), image, name, offset);
                }
        }
        putchar('\n');
}


extern "C" void __dump_allocated(void);

extern "C" void
__dump_allocated(void)
{
        hoardLock(sUsedLock);

        puts("allocated:\n");

        block* b = sUsedList;
        while (b != NULL) {
                dump_block(b);

                b = b->getNext();
        }

        hoardUnlock(sUsedLock);
}


static void
add_address(void* address, size_t size)
{
        block *b = (block *)address - 1;

#ifdef __i386__
        // set call stack
        struct stack_frame {
                struct stack_frame*     previous;
                void*                           return_address;
        };

        stack_frame* frame = (stack_frame*)get_stack_frame();

        for (int i = 0; i < HEAP_CALL_STACK_SIZE; i++) {
                if (frame != NULL) {
                        b->setCallStack(i, frame->return_address);
                        frame = frame->previous;
                } else
                        b->setCallStack(i, NULL);
        }

        b->setAllocatedSize(size);
#endif

        hoardLock(sUsedLock);

        b->setNext(sUsedList);
        sUsedList = b;

        hoardUnlock(sUsedLock);
}


static void
remove_address(void* address)
{
        block* b = (block *)address - 1;
        hoardLock(sUsedLock);

        if (sUsedList == b) {
                // we're lucky, it's the first block in the list
                sUsedList = b->getNext();
        } else {
                // search for block in the used list (very slow!)
                block* last = sUsedList;
                while (last != NULL && last->getNext() != b) {
                        last = last->getNext();
                }

                if (last == NULL) {
                        printf("freed block not in used list!\n");
                        dump_block(b);
                } else
                        last->setNext(b->getNext());
        }

        hoardUnlock(sUsedLock);
}

#endif  // HEAP_LEAK_CHECK

#if HEAP_WALL

static void*
set_wall(void* addr, size_t size)
{
        size_t *start = (size_t*)addr;

        start[0] = size;
        memset(start + 1, 0x88, HEAP_WALL_SIZE - sizeof(size_t));
        memset((uint8*)addr + size - HEAP_WALL_SIZE, 0x66, HEAP_WALL_SIZE);

        return (uint8*)addr + HEAP_WALL_SIZE;
}


static void*
check_wall(uint8* buffer)
{
        buffer -= HEAP_WALL_SIZE;
        size_t size = *(size_t*)buffer;

        if (threadHeap::objectSize(buffer) < size)
                debugger("invalid size");

        for (size_t i = 0; i < HEAP_WALL_SIZE; i++) {
                if (i >= sizeof(size_t) && buffer[i] != 0x88) {
                        debug_printf("allocation %p, size %ld front wall clobbered at byte %ld.\n",
                                buffer + HEAP_WALL_SIZE, size - 2 * HEAP_WALL_SIZE, i);
                        debugger("front wall clobbered");
                }
                if (buffer[i + size - HEAP_WALL_SIZE] != 0x66) {
                        debug_printf("allocation %p, size %ld back wall clobbered at byte %ld.\n",
                                buffer + HEAP_WALL_SIZE, size - 2 * HEAP_WALL_SIZE, i);
                        debugger("back wall clobbered");
                }
        }

        return buffer;
}

#endif  // HEAP_WALL

inline static processHeap *
getAllocator(void)
{
        static char *buffer = (char *)hoardSbrk(sizeof(processHeap));
        static processHeap *theAllocator = new (buffer) processHeap();

        return theAllocator;
}


extern "C" void
__heap_before_fork(void)
{
        static processHeap *pHeap = getAllocator();
        for (int i = 0; i < pHeap->getMaxThreadHeaps(); i++)
                pHeap->getHeap(i).lock();

        static_cast<hoardHeap*>(pHeap)->lock();
}

void __init_after_fork(void);

extern "C" void
__heap_after_fork_child(void)
{
        __init_after_fork();
        static processHeap *pHeap = getAllocator();
        for (int i = 0; i < pHeap->getMaxThreadHeaps(); i++)
                pHeap->getHeap(i).initLock();

        pHeap->initLock();
}


extern "C" void
__heap_after_fork_parent(void)
{
        static processHeap *pHeap = getAllocator();
        for (int i = 0; i < pHeap->getMaxThreadHeaps(); i++)
                pHeap->getHeap(i).unlock();

        static_cast<hoardHeap*>(pHeap)->unlock();
}


extern "C" void
__heap_thread_init(void)
{
}


extern "C" void
__heap_thread_exit(void)
{
}


//      #pragma mark - public functions


extern "C" void *
malloc(size_t size)
{
        static processHeap *pHeap = getAllocator();

#if HEAP_WALL
        size += 2 * HEAP_WALL_SIZE;
#endif

        defer_signals();

        void *addr = pHeap->getHeap(pHeap->getHeapIndex()).malloc(size);
        if (addr == NULL) {
                undefer_signals();
                __set_errno(B_NO_MEMORY);
                KTRACE("malloc(%lu) -> NULL", size);
                return NULL;
        }

#if HEAP_LEAK_CHECK
        add_address(addr, size);
#endif

        undefer_signals();

#if HEAP_WALL
        addr = set_wall(addr, size);
#endif

        KTRACE("malloc(%lu) -> %p", size, addr);

        return addr;
}


extern "C" void *
calloc(size_t nelem, size_t elsize)
{
        static processHeap *pHeap = getAllocator();
        size_t size = nelem * elsize;
        void *ptr = NULL;

        if ((nelem > 0) && ((size/nelem) != elsize))
                goto nomem;

#if HEAP_WALL
        size += 2 * HEAP_WALL_SIZE;

        if (nelem == 0 || elsize == 0)
                goto ok;
        if (size < (nelem * size)&& size < (elsize * size))
                goto nomem;

ok:
#endif

        defer_signals();

        ptr = pHeap->getHeap(pHeap->getHeapIndex()).malloc(size);
        if (ptr == NULL) {
                undefer_signals();
        nomem:
                __set_errno(B_NO_MEMORY);
                KTRACE("calloc(%lu, %lu) -> NULL", nelem, elsize);
                return NULL;
        }

#if HEAP_LEAK_CHECK
        add_address(ptr, size);
#endif

        undefer_signals();

#if HEAP_WALL
        ptr = set_wall(ptr, size);
        size -= 2 * HEAP_WALL_SIZE;
#endif

        // Zero out the malloc'd block.
        memset(ptr, 0, size);
        KTRACE("calloc(%lu, %lu) -> %p", nelem, elsize, ptr);
        return ptr;
}


extern "C" void
free(void *ptr)
{
        static processHeap *pHeap = getAllocator();

#if HEAP_WALL
        if (ptr == NULL)
                return;
        KTRACE("free(%p)", ptr);
        ptr = check_wall((uint8*)ptr);
#else
        KTRACE("free(%p)", ptr);
#endif

        defer_signals();

#if HEAP_LEAK_CHECK
        if (ptr != NULL)
                remove_address(ptr);
#endif
        pHeap->free(ptr);

        undefer_signals();
}


extern "C" void *
memalign(size_t alignment, size_t size)
{
        static processHeap *pHeap = getAllocator();

#if HEAP_WALL
        debug_printf("memalign() is not yet supported by the wall code.\n");
        return NULL;
#endif

        defer_signals();

        void *addr = pHeap->getHeap(pHeap->getHeapIndex()).memalign(alignment,
                size);
        if (addr == NULL) {
                undefer_signals();
                __set_errno(B_NO_MEMORY);
                KTRACE("memalign(%lu, %lu) -> NULL", alignment, size);
                return NULL;
        }

#if HEAP_LEAK_CHECK
        add_address(addr, size);
#endif

        undefer_signals();

        KTRACE("memalign(%lu, %lu) -> %p", alignment, size, addr);
        return addr;
}


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


extern "C" int
posix_memalign(void **_pointer, size_t alignment, size_t size)
{
        if ((alignment & (sizeof(void *) - 1)) != 0 || _pointer == NULL)
                return B_BAD_VALUE;

#if HEAP_WALL
        debug_printf("posix_memalign() is not yet supported by the wall code.\n");
        return -1;
#endif
        static processHeap *pHeap = getAllocator();
        defer_signals();
        void *pointer = pHeap->getHeap(pHeap->getHeapIndex()).memalign(alignment,
                size);
        if (pointer == NULL) {
                undefer_signals();
                KTRACE("posix_memalign(%p, %lu, %lu) -> NULL", _pointer, alignment,
                        size);
                return B_NO_MEMORY;
        }

#if HEAP_LEAK_CHECK
        add_address(pointer, size);
#endif

        undefer_signals();

        *_pointer = pointer;
        KTRACE("posix_memalign(%p, %lu, %lu) -> %p", _pointer, alignment, size,
                pointer);
        return 0;
}


extern "C" void *
valloc(size_t size)
{
        return memalign(B_PAGE_SIZE, size);
}


extern "C" void *
realloc(void *ptr, size_t size)
{
        if (ptr == NULL)
                return malloc(size);

        if (size == 0) {
                free(ptr);
                return NULL;
        }

        // If the existing object can hold the new size,
        // just return it.

#if HEAP_WALL
        size += 2 * HEAP_WALL_SIZE;
        ptr = (uint8*)ptr - HEAP_WALL_SIZE;
#endif

        size_t objSize = threadHeap::objectSize(ptr);
        if (objSize >= size) {
#if HEAP_WALL
                check_wall((uint8*)ptr + HEAP_WALL_SIZE);
                ptr = set_wall(ptr, size);
#endif
                KTRACE("realloc(%p, %lu) -> %p", ptr, size, ptr);
                return ptr;
        }

#if HEAP_WALL
        size -= 2 * HEAP_WALL_SIZE;
        objSize -= 2 * HEAP_WALL_SIZE;
        ptr = (uint8*)ptr + HEAP_WALL_SIZE;
#endif

        // Allocate a new block of size sz.
        void *buffer = malloc(size);
        if (buffer == NULL) {
                // Allocation failed, leave old block and return
                __set_errno(B_NO_MEMORY);
                KTRACE("realloc(%p, %lu) -> NULL", ptr, size);
                return NULL;
        }

        // Copy the contents of the original object
        // up to the size of the new block.

        size_t minSize = (objSize < size) ? objSize : size;
        memcpy(buffer, ptr, minSize);

        // Free the old block.
        free(ptr);

        // Return a pointer to the new one.
        KTRACE("realloc(%p, %lu) -> %p", ptr, size, buffer);
        return buffer;
}


extern "C" size_t
malloc_usable_size(void *ptr)
{
        if (ptr == NULL)
                return 0;
        return threadHeap::objectSize(ptr);
}


//      #pragma mark - BeOS specific extensions


struct mstats {
        size_t bytes_total;
        size_t chunks_used;
        size_t bytes_used;
        size_t chunks_free;
        size_t bytes_free;
};


extern "C" struct mstats mstats(void);

extern "C" struct mstats
mstats(void)
{
        // Note, the stats structure is not thread-safe, but it doesn't
        // matter that much either
        processHeap *heap = getAllocator();
        static struct mstats stats;

        int allocated = 0;
        int used = 0;
        int chunks = 0;

        for (int i = 0; i < hoardHeap::SIZE_CLASSES; i++) {
                int classUsed, classAllocated;
                heap->getStats(i, classUsed, classAllocated);

                if (classUsed > 0)
                        chunks++;

                allocated += classAllocated;
                used += classUsed;
        }

        stats.bytes_total = allocated;
        stats.chunks_used = chunks;
        stats.bytes_used = used;
        stats.chunks_free = hoardHeap::SIZE_CLASSES - chunks;
        stats.bytes_free = allocated - used;

        return stats;
}