/** @file  psMemory.c
*
*  @brief Contains the definitions for the memory management system
*
*  psMemory.h has additional information and documentation of the routines found in this file.
*
*  @author Robert DeSonia, MHPCC
*  @author Robert Lupton, Princeton University
*
*  @version $Revision: 1.52 $ $Name: not supported by cvs2svn $
*  @date $Date: 2005-06-08 23:40:45 $
*
*  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
*/

#define PS_ALLOW_MALLOC                    // we're allowed to call malloc()

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

#include "psMemory.h"
#include "psError.h"
#include "psAbort.h"
#include "psLogMsg.h"

#include "psSysUtilsErrors.h"

#define P_PS_MEMMAGIC (psPtr )0xdeadbeef   // Magic number in psMemBlock header

#define P_PS_LARGE_BLOCK_SIZE 65536        // size where under, we try to recycle

static psS32 checkMemBlock(const psMemBlock* m, const char *funcName);
static psMemBlock* lastMemBlockAllocated = NULL;
static pthread_mutex_t memBlockListMutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_mutex_t memIdMutex = PTHREAD_MUTEX_INITIALIZER;

static pthread_mutex_t recycleMemBlockListMutex = PTHREAD_MUTEX_INITIALIZER;

static psS32 recycleBins = 13;
static psS32 recycleBinSize[14] = {
                                      8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, P_PS_LARGE_BLOCK_SIZE
                                  };

// N.B. recycleBinSize should be terminated by P_PS_LARGE_BLOCK_SIZE (simplifies search loops)
static psMemBlock* recycleMemBlockList[13] = {
            NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL
        };

#ifdef PS_MEM_DEBUG
static psMemBlock* deadBlockList;       // a place to put dead memBlocks in debug mode.
#endif // #ifdef PS_MEM_DEBUG

/**
 * Unique ID for allocated blocks
 */
static psMemoryId memid = 0;

/**
 *  Default memExhausted callback.
 */
static psPtr memExhaustedCallbackDefault(size_t size)
{
    psPtr ptr = NULL;

    pthread_mutex_lock(&recycleMemBlockListMutex);
    psS32 level = recycleBins - 1;

    while (level >= 0 && ptr == NULL) {
        while (recycleMemBlockList[level] != NULL && ptr == NULL) {
            psMemBlock* old = recycleMemBlockList[level];

            recycleMemBlockList[level] = recycleMemBlockList[level]->nextBlock;
            free(old);
            ptr = malloc(size);
        }
        level--;
    }
    pthread_mutex_unlock(&recycleMemBlockListMutex);

    return ptr;
}

/*
 * Default callback for both allocate and free. Note that the
 * value of p_psMemAllocateID/p_psMemFreeID is incremented
 * by the return value (so returning 0 means that the callback
 * isn't resignalled)
 */
static psMemoryId memAllocateCallbackDefault(const psMemBlock* ptr)
{
    static psMemoryId incr = 0; // "p_psMemAllocateID += incr"

    return incr;
}

static psMemoryId memFreeCallbackDefault(const psMemBlock* ptr)
{
    static psMemoryId incr = 0; // "p_psMemFreeID += incr"

    return incr;
}

static void memProblemCallbackDefault(const psMemBlock* ptr, const char *file, psS32 lineno)
{
    if (ptr->refCounter < 1) {
        psError(PS_ERR_MEMORY_CORRUPTION, false,
                PS_ERRORTEXT_psMemory_MULTIPLE_FREE,
                ptr->id, ptr->file, ptr->lineno, file, lineno);
    }

    if (lineno > 0) {
        psAbort(__func__, "Detected a problem in the memory system at %s:%d", file, lineno);
    }
}

/*
 * Routines to check the consistency of the allocated and/or free memory arena
 *
 * N.b. If the block wasn't allocated by psAlloc, it will appear corrupted
 */
static psS32 checkMemBlock(const psMemBlock* m, const char *funcName)
{
    // n.b. since this is called by psMemCheckCorruption while the memblock list is mutex locked,
    // we shouldn't call such things as p_psAlloc/p_psFree here.

    if (m == NULL) {
        psError(PS_ERR_MEMORY_CORRUPTION, true,
                PS_ERRORTEXT_psMemory_NULL_BLOCK);
        return 1;
    }

    if (m->refCounter == 0) {
        // using an unreferenced block of memory, are you?
        psError(PS_ERR_MEMORY_CORRUPTION, true,
                PS_ERRORTEXT_psMemory_DEREF_BLOCK_USE,
                m->id);
        return 1;
    }

    if (m->startblock != P_PS_MEMMAGIC || m->endblock != P_PS_MEMMAGIC) {
        psError(PS_ERR_MEMORY_CORRUPTION, true,
                PS_ERRORTEXT_psMemory_UNDERFLOW,
                m->id);
        return 1;
    }
    if (*(psPtr *)((int8_t *) (m + 1) + m->userMemorySize) != P_PS_MEMMAGIC) {
        psError(PS_ERR_MEMORY_CORRUPTION, true,
                PS_ERRORTEXT_psMemory_OVERFLOW,
                m->id);
        return 1;
    }

    return 0;
}

/*
 * The default callbacks
 */
static psMemAllocateCallback memAllocateCallback = memAllocateCallbackDefault;
static psMemFreeCallback memFreeCallback = memFreeCallbackDefault;
static psMemProblemCallback memProblemCallback = memProblemCallbackDefault;
static psMemExhaustedCallback memExhaustedCallback = memExhaustedCallbackDefault;

psMemExhaustedCallback psMemExhaustedCallbackSet(psMemExhaustedCallback func)
{
    psMemExhaustedCallback old = memExhaustedCallback;

    if (func != NULL) {
        memExhaustedCallback = func;
    } else {
        memExhaustedCallback = memExhaustedCallbackDefault;
    }

    return old;
}

psMemProblemCallback psMemProblemCallbackSet(psMemProblemCallback func)
{
    psMemProblemCallback old = memProblemCallback;

    if (func != NULL) {
        memProblemCallback = func;
    } else {
        memProblemCallback = memProblemCallbackDefault;
    }

    return old;
}

/*
 * And now the I-want-to-be-informed callbacks
 *
 * Call the callbacks when these IDs are allocated/freed
 */
psMemoryId p_psMemAllocateID = 0;       // notify user this block is allocated
psMemoryId p_psMemFreeID = 0;   // notify user this block is freed

psMemoryId psMemAllocateCallbackSetID(psMemoryId id)
{
    psMemoryId old = p_psMemAllocateID;

    p_psMemAllocateID = id;

    return old;
}

psMemoryId psMemFreeCallbackSetID(psMemoryId id)
{
    psMemoryId old = p_psMemFreeID;

    p_psMemFreeID = id;

    return old;
}

psMemAllocateCallback psMemAllocateCallbackSet(psMemAllocateCallback func)
{
    psMemFreeCallback old = memAllocateCallback;

    if (func != NULL) {
        memAllocateCallback = func;
    } else {
        memAllocateCallback = memAllocateCallbackDefault;
    }

    return old;
}

psMemFreeCallback psMemFreeCallbackSet(psMemFreeCallback func)
{
    psMemFreeCallback old = memFreeCallback;

    if (func != NULL) {
        memFreeCallback = func;
    } else {
        memFreeCallback = memFreeCallbackDefault;
    }

    return old;
}

/*
 * Return memory ID counter for next block to be allocated
 */
psMemoryId psMemGetId(void)
{
    psMemoryId id;

    pthread_mutex_lock(&memIdMutex);
    id = memid + 1;
    pthread_mutex_unlock(&memIdMutex);

    return id;
}

psS32 psMemCheckCorruption(psBool abort_on_error)
{
    psS32 nbad = 0;               // number of bad blocks
    psBool failure = false;

    // get exclusive access to the memBlock list to avoid it changing on us while we use it.
    //    pthread_mutex_lock(&memBlockListMutex);

    for (psMemBlock* iter = lastMemBlockAllocated; iter != NULL; iter = iter->nextBlock) {
        pthread_mutex_unlock(&memBlockListMutex);
        failure = checkMemBlock(iter, __func__);
        pthread_mutex_lock(&memBlockListMutex);
        if ( failure ) {
            nbad++;

            memProblemCallback(iter, __func__, __LINE__);

            if (abort_on_error) {
                // release the lock on the memblock list
                pthread_mutex_unlock(&memBlockListMutex);
                psAbort(__func__, "Detected memory corruption");
                return nbad;
            }
        }
    }

    // release the lock on the memblock list
    pthread_mutex_unlock(&memBlockListMutex);
    return nbad;
}

psPtr p_psAlloc(size_t size, const char *file, psS32 lineno)
{

    psMemBlock* ptr = NULL;

    size = (size < recycleBinSize[0]) ? recycleBinSize[0] : size; // set the minimum size to allocate

    // memory is of the size I want to bother recycling?
    if (size < P_PS_LARGE_BLOCK_SIZE) {
        // find the bin we need.
        psS32 level = 0;

        while (size > recycleBinSize[level]) {
            level++;
        }
        // Are we in one of the bins
        if (level < recycleBins) {

            size = recycleBinSize[level];  // round-up size to next sized bin.

            pthread_mutex_lock(&recycleMemBlockListMutex);

            if (recycleMemBlockList[level] != NULL) {
                ptr = recycleMemBlockList[level];
                recycleMemBlockList[level] = ptr->nextBlock;
                if (recycleMemBlockList[level] != NULL) {
                    recycleMemBlockList[level]->previousBlock = NULL;
                }
                size = ptr->userMemorySize;
            }

            pthread_mutex_unlock(&recycleMemBlockListMutex);
        }
    }

    if (ptr == NULL) {
        ptr = malloc(sizeof(psMemBlock) + size + sizeof(psPtr ));

        if (ptr == NULL) {
            ptr = memExhaustedCallback(size);
            if (ptr == NULL) {
                psAbort(__func__, "Failed to allocate %u bytes at %s:%d", size, file, lineno);
            }
        }

        *(psPtr*)&ptr->startblock = P_PS_MEMMAGIC;
        *(psPtr*)&ptr->endblock = P_PS_MEMMAGIC;
        ptr->userMemorySize = size;
        pthread_mutex_init(&ptr->refCounterMutex, NULL);
    }
    // increment the memory id safely.
    pthread_mutex_lock(&memBlockListMutex);
    *(psMemoryId* ) & ptr->id = ++memid;
    pthread_mutex_unlock(&memBlockListMutex);

    ptr->file = file;
    ptr->freeFcn = NULL;
    ptr->persistent = false;
    *(psU32 *)&ptr->lineno = lineno;
    *(psPtr *)((int8_t *) (ptr + 1) + size) = P_PS_MEMMAGIC;
    ptr->previousBlock = NULL;

    ptr->refCounter = 1;                   // one user so far

    // need exclusive access of the memory block list now...
    pthread_mutex_lock(&memBlockListMutex);

    // insert the new block to the front of the memBlock linked-list
    ptr->nextBlock = lastMemBlockAllocated;
    if (ptr->nextBlock != NULL) {
        ptr->nextBlock->previousBlock = ptr;
    }
    lastMemBlockAllocated = ptr;

    pthread_mutex_unlock(&memBlockListMutex);

    // Did the user ask to be informed about this allocation?
    if (ptr->id == p_psMemAllocateID) {
        p_psMemAllocateID += memAllocateCallback(ptr);
    }
    // And return the user the memory that they allocated
    return ptr + 1;                        // user memory
}

psPtr p_psRealloc(psPtr vptr, size_t size, const char *file, psS32 lineno)
{
    size = (size < recycleBinSize[0]) ? recycleBinSize[0] : size; // set the minimum size to allocate

    if (vptr == NULL) {
        return p_psAlloc(size, file, lineno);
    } else {
        psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;
        psBool isBlockLast = false;

        if (checkMemBlock(ptr, __func__) != 0) {
            memProblemCallback(ptr, file, lineno);
            psAbort(file, "Realloc detected a memory corruption (id %lld @ %s:%d).",
                    ptr->id, ptr->file, ptr->lineno);
        }

        pthread_mutex_lock(&memBlockListMutex);

        isBlockLast = (ptr == lastMemBlockAllocated);

        ptr = (psMemBlock* ) realloc(ptr, sizeof(psMemBlock) + size + sizeof(psPtr ));

        if (ptr == NULL) {
            ptr = memExhaustedCallback(size);
            if(ptr == NULL) {
                psAbort(__func__, "Failed to reallocate %ld bytes at %s:%d", size, file, lineno);
            }
        }

        ptr->userMemorySize = size;
        *(psPtr *)((int8_t *) (ptr + 1) + size) = P_PS_MEMMAGIC;

        if (isBlockLast) {
            lastMemBlockAllocated = ptr;
        }
        // the block location may have changed, so fix the linked list addresses.
        if (ptr->nextBlock != NULL) {
            ptr->nextBlock->previousBlock = ptr;
        }
        if (ptr->previousBlock != NULL) {
            ptr->previousBlock->nextBlock = ptr;
        }

        pthread_mutex_unlock(&memBlockListMutex);

        // Did the user ask to be informed about this allocation?
        if (ptr->id == p_psMemAllocateID) {
            p_psMemAllocateID += memAllocateCallback(ptr);
        }

        return ptr + 1;                    // usr memory
    }
}

void p_psFree(psPtr vptr, const char *file, psS32 lineno)
{
    if (vptr == NULL) {
        return;
    }
    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;
    if (ptr->refCounter < 1) {
        psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

        psAbort(__func__,PS_ERRORTEXT_psMemory_MULTIPLE_FREE,
                ptr->id, ptr->file, ptr->lineno, file, lineno);
    }

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, file, lineno);
        psAbort(__func__,"Memory Corruption Detected.");
    }

    (void)p_psMemDecrRefCounter(vptr, file, lineno);    // this handles the free, if required.
}

/*
 * Check for memory leaks.
 */
psS32 psMemCheckLeaks(psMemoryId id0, psMemBlock* ** arr, FILE * fd, psBool persistence)
{
    psS32 nleak = 0;
    psS32 j = 0;
    psMemBlock* topBlock = lastMemBlockAllocated;

    pthread_mutex_lock(&memBlockListMutex);

    for (psMemBlock* iter = topBlock; iter != NULL; iter = iter->nextBlock) {
        if ( (iter->refCounter > 0) &&
                ( (persistence) || (!persistence && !iter->persistent) ) &&
                (iter->id >= id0)) {

            nleak++;

            if (fd != NULL) {
                if (nleak == 1) {
                    fprintf(fd, "   %20s:line ID\n", "file");
                }

                fprintf(fd, "   %20s:%-4d %ld\n", iter->file, (int)iter->lineno, (long)iter->id);
            }
        }
    }

    pthread_mutex_unlock(&memBlockListMutex);

    if (nleak == 0 || arr == NULL) {
        return nleak;
    }

    *arr = p_psAlloc(nleak * sizeof(psMemBlock), __FILE__, __LINE__);
    pthread_mutex_lock(&memBlockListMutex);

    for (psMemBlock* iter = topBlock; iter != NULL; iter = iter->nextBlock) {
        if ( (iter->refCounter > 0) &&
                ( (persistence) || (!persistence && !iter->persistent) ) &&
                (iter->id >= id0)) {

            (*arr)[j++] = iter;
            if (j == nleak) {              // found them all
                break;
            }
        }
    }

    pthread_mutex_unlock(&memBlockListMutex);

    return nleak;
}

/*
 * Reference counting APIs
 */

// return refCounter
psReferenceCount psMemGetRefCounter(psPtr vptr)
{
    psMemBlock* ptr;
    psU32 refCount;

    if (vptr == NULL) {
        return 0;
    }

    ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, __func__, __LINE__);
    }

    pthread_mutex_lock(&ptr->refCounterMutex);
    refCount = ptr->refCounter;
    pthread_mutex_unlock(&ptr->refCounterMutex);

    return refCount;
}

// increment and return refCounter
psPtr p_psMemIncrRefCounter(psPtr vptr, const char *file, psS32 lineno)
{
    psMemBlock* ptr;

    if (vptr == NULL) {
        return vptr;
    }

    ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__)) {
        memProblemCallback(ptr, file, lineno);
    }

    pthread_mutex_lock(&ptr->refCounterMutex);
    ptr->refCounter++;
    pthread_mutex_unlock(&ptr->refCounterMutex);

    return vptr;
}

// decrement and return refCounter
psPtr p_psMemDecrRefCounter(psPtr vptr, const char *file, psS32 lineno)
{
    if (vptr == NULL) {
        return NULL;
    }

    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, file, lineno);
        return NULL;
    }

    // Did the user ask to be informed about this deallocation?
    if (ptr->id == p_psMemFreeID) {
        p_psMemFreeID += memFreeCallback(ptr);
    }

    pthread_mutex_lock(&ptr->refCounterMutex);

    if (ptr->refCounter > 1) {
        ptr->refCounter--;                 // multiple references, just decrement the count.
        pthread_mutex_unlock(&ptr->refCounterMutex);

    } else {
        pthread_mutex_unlock(&ptr->refCounterMutex);

        if (ptr->freeFcn != NULL) {
            ptr->freeFcn(vptr);
        }

        pthread_mutex_lock(&memBlockListMutex);

        // cut the memBlock out of the memBlock list
        if (ptr->nextBlock != NULL) {
            ptr->nextBlock->previousBlock = ptr->previousBlock;
        }
        if (ptr->previousBlock != NULL) {
            ptr->previousBlock->nextBlock = ptr->nextBlock;
        }
        if (lastMemBlockAllocated == ptr) {
            lastMemBlockAllocated = ptr->nextBlock;
        }

        pthread_mutex_unlock(&memBlockListMutex);

        // do we need to recycle?
        if (ptr->userMemorySize < P_PS_LARGE_BLOCK_SIZE) {

            psS32 level = 1;

            while (ptr->userMemorySize >= recycleBinSize[level]) {
                level++;
            }
            level--;

            ptr->refCounter = 0;
            ptr->previousBlock = NULL;

            pthread_mutex_lock(&recycleMemBlockListMutex);
            ptr->nextBlock = recycleMemBlockList[level];
            if (recycleMemBlockList[level] != NULL) {
                recycleMemBlockList[level]->previousBlock = ptr;
            }
            recycleMemBlockList[level] = ptr;
            pthread_mutex_unlock(&recycleMemBlockListMutex);

        } else {
            // memory is larger than I want to recycle.
            #ifdef PS_MEM_DEBUG
            (void)p_psRealloc(vptr, 0, file, lineno);
            ptr->previousBlock = NULL;
            ptr->nextBlock = deadBlockList;
            if (deadBlockList != NULL) {
                deadBlockList->previous = ptr;
            }
            deadBlockList = ptr;
            #else // #ifdef PS_MEM_DEBUG

            pthread_mutex_destroy(&ptr->refCounterMutex);
            free(ptr);
            #endif // #else - #ifdef PS_MEM_DEBUG

        }

        vptr = NULL;                       // since we freed it, make sure we return NULL.
    }

    return vptr;
}

void psMemSetDeallocator(psPtr vptr, psFreeFcn freeFcn)
{
    if (vptr == NULL) {
        return;
    }

    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, __func__, __LINE__);
    }

    ptr->freeFcn = freeFcn;

}
psFreeFcn psMemGetDeallocator(psPtr vptr)
{
    if (vptr == NULL) {
        return NULL;
    }

    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, __func__, __LINE__);
    }

    return ptr->freeFcn;
}

psBool p_psMemGetPersistent(psPtr vptr)
{
    if (vptr == NULL) {
        return NULL;
    }

    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, __func__, __LINE__);
    }

    return ptr->persistent;
}

void p_psMemSetPersistent(psPtr vptr,psBool value)
{
    if (vptr == NULL) {
        return;
    }

    psMemBlock* ptr = ((psMemBlock* ) vptr) - 1;

    if (checkMemBlock(ptr, __func__) != 0) {
        memProblemCallback(ptr, __func__, __LINE__);
    }

    ptr->persistent = value;
}
