/** @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.28 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2004-07-15 23:52:34 $
 *
 *  Copyright 2004 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 <stdbool.h>
#include <stdint.h>

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

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

static int 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;

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

/**
 *  Default memExhausted callback.
 */
static void *memExhaustedCallbackDefault(size_t size)
{
    return NULL;
}

static psMemExhaustedCallback memExhaustedCallback = memExhaustedCallbackDefault;

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

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

    return old;
}

static void memProblemCallbackDefault(const psMemBlock *ptr,
                                      const char *file, int lineno)
{
    if (ptr->refCounter < 1) {
        psError(__func__,
                "Block %ld allocated at %s:%d freed more than once at %s:%d\n",
                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);
    }
}
static psMemProblemCallback memProblemCallback = memProblemCallbackDefault;

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;
}

/*
 * 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;
}

/*
 * The default callbacks, and the routines to change them
 */
static psMemAllocateCallback memAllocateCallback = memAllocateCallbackDefault;
static psMemFreeCallback memFreeCallback = memFreeCallbackDefault;

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;
}

/*
 * 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 int 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(funcName,"Memory Corruption: NULL memory block found.");
        return 1;
    }

    if (m->refCounter == 0) {
        // using an unreferenced block of memory, are you?
        psError(__func__,"Memory Corruption: memory block %ld was freed but still used.",
                m->id);
        return 1;
    }

    if (m->startblock != P_PS_MEMMAGIC || m->endblock != P_PS_MEMMAGIC) {
        psError(funcName, "Memory Corruption: memory block %ld is corrupted (buffer underflow)",
                m->id);
        return 1;
    }
    if (*(void**)((int8_t*)(m+1)+m->userMemorySize) != P_PS_MEMMAGIC) {
        psError(funcName, "Memory Corruption: memory block %ld is corrupted (buffer overflow)",
                m->id);
        return 1;
    }

    return 0;
}

int psMemCheckCorruption(bool abort_on_error)
{
    int nbad = 0;                       // number of bad blocks

    // 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) {
        if (checkMemBlock(iter, __func__)) {
            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;
}

void *p_psAlloc(size_t size, const char *file, int lineno)
{
    psMemBlock *ptr = malloc(sizeof(psMemBlock) + size + sizeof(void*));

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

    // increment the memory id safely.
    pthread_mutex_lock(&memBlockListMutex);
    *(psMemoryId*)&ptr->id = ++memid;
    pthread_mutex_unlock(&memBlockListMutex);

    ptr->file = file;
    ptr->freeFcn = NULL;
    *(unsigned int*)&ptr->lineno = lineno;
    ptr->startblock = P_PS_MEMMAGIC;
    ptr->endblock = P_PS_MEMMAGIC;
    ptr->userMemorySize = size;
    *(void**)((int8_t*)(ptr+1)+size) = P_PS_MEMMAGIC;
    ptr->previousBlock = NULL;

    pthread_mutex_init(&ptr->refCounterMutex, 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
}

void *p_psRealloc(void *vptr, size_t size, const char *file, int lineno)
{
    if (vptr == NULL) {
        return p_psAlloc(size, file, lineno);
    } else {
        psMemBlock *ptr = ((psMemBlock *)vptr) - 1;
        bool isBlockLast = false;

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

        pthread_mutex_lock(&memBlockListMutex);

        isBlockLast = (ptr == lastMemBlockAllocated);

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

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

        ptr->userMemorySize = size;
        *(void**)((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(void *vptr, const char *file, int lineno)
{
    (void)p_psMemDecrRefCounter(vptr,file,lineno);   // this handles the free, if required.
}

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

    pthread_mutex_lock(&memBlockListMutex);

    for (psMemBlock* iter = topBlock; iter != NULL; iter=iter->nextBlock) {
        if ( (psMemGetRefCounter(iter+1) > 0) && (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, iter->lineno, 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 ( (psMemGetRefCounter(iter+1) > 0) && (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(void *vptr)
{
    psMemBlock *ptr;
    unsigned int 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
void* p_psMemIncrRefCounter(void *vptr, const char *file, int 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
void* p_psMemDecrRefCounter(void *vptr,const char *file, int lineno)
{
    if (vptr == NULL) {
        return NULL;
    }

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

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

    pthread_mutex_lock(&ptr->refCounterMutex);

    if (ptr->refCounter > 1) {
        /// XXX - Probably should have another mutex here.
        ptr->refCounter--;          // multiple references, just decrement the count.
        pthread_mutex_unlock(&ptr->refCounterMutex);

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

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

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

        #ifdef PS_MEM_DEBUG
        // for debug mode, we should keep the memBlock around and mark as not referenced.
        ptr->refCounter = 0;
        (void)p_psRealloc(vptr,0,file,lineno);
        #else

        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);

        pthread_mutex_destroy(&ptr->refCounterMutex);

        free(ptr);
        #endif

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

    return vptr;
}

void p_psMemSetDeallocator(void* 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 p_psMemGetDeallocator(void* vptr)
{
    if (vptr == NULL) {
        return NULL;
    }

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

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

    return ptr->freeFcn;
}
