/* $Id: kLdrHlpHeap.c 2847 2006-11-01 19:17:21Z bird $ */
/** @file
 *
 * kLdr - The Dynamic Loader, Helper Functions, Heap.
 *
 * Copyright (c) 2006 knut st. osmundsen <bird-kbuild-src@anduin.net>
 *
 *
 * This file is part of kLdr.
 *
 * kLdr is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * kLdr 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with kLdr; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#define KLDRHEAP_STRICT

/*******************************************************************************
*   Header Files                                                               *
*******************************************************************************/
#ifdef __OS2__
# define INCL_BASE
# define INCL_ERRORS
# include <os2.h>
#elif defined(__WIN__)
# include <Windows.h>
#else
# error "port me"
#endif

#include <kLdr.h>
#include "kLdrHlp.h"



/*******************************************************************************
*   Structures and Typedefs                                                    *
*******************************************************************************/
/**
 * A heap block.
 */
typedef struct KLDRHEAPBLOCK
{
    /** Next block in the global list. */
    struct KLDRHEAPBLOCK   *pNext;
    /** Previous block in the global list. */
    struct KLDRHEAPBLOCK   *pPrev;
    /** The size of this block including this header. */
    size_t                  cb;
    /** The flags. */
    size_t                  fFlags;
} KLDRHEAPBLOCK, *PKLDRHEAPBLOCK;

/** Indicates whether the block is free (set) or allocated (clear). */
#define KLDRHEAPBLOCK_FLAG_FREE     ((size_t)1)
/** Valid flag mask. */
#define KLDRHEAPBLOCK_FLAG_MASK     ((size_t)1)

/** Checks if the block is freed. */
#define KLDRHEAPBLOCK_IS_FREE(pB)       ( (pB)->fFlags & KLDRHEAPBLOCK_FLAG_FREE )
/** Check if the block is allocated. */
#define KLDRHEAPBLOCK_IS_ALLOCATED(pB)  !KLDRHEAPBLOCK_IS_FREE(pB)
/** Checks if the two blocks are adjacent.
 * Assumes pB1 < pB2. */
#define KLDRHEAPBLOCK_IS_ADJACENT(pB1, pB2) \
    ( ((uintptr_t)(pB1) + (pB1)->cb) == (uintptr_t)(pB2) )

/** The block alignment. */
#define KLDRHEAPBLOCK_ALIGNMENT     sizeof(KLDRHEAPBLOCK)

/** @def KLDRHEAP_ASSERT
 * Heap assertion. */
/** @def KLDRHEAP_ASSERT_BLOCK
 * Assert that a heap block is valid. */
/** @def KLDRHEAP_ASSERT_FREE
 * Assert that a heap free block is valid. */
#ifdef KLDRHEAP_STRICT
# define KLDRHEAP_ASSERT(expr)      kldrHlpAssert(expr)

# define KLDRHEAP_ASSERT_BLOCK(pHeap, pBlock) \
    do { \
        KLDRHEAP_ASSERT(!((pBlock)->fFlags & ~KLDRHEAPBLOCK_FLAG_MASK)); \
        KLDRHEAP_ASSERT(!((pBlock)->cb & (KLDRHEAPBLOCK_ALIGNMENT - 1))); \
        KLDRHEAP_ASSERT((uintptr_t)(pBlock)->pPrev < (uintptr_t)(pBlock)); \
        KLDRHEAP_ASSERT((uintptr_t)(pBlock)->pNext > (uintptr_t)(pBlock) || !(pBlock)->pNext); \
    } while (0)

# define KLDRHEAP_ASSERT_FREE(pHeap, pFree) \
    do { \
        KLDRHEAP_ASSERT_BLOCK(pHeap, &(pFree)->Core); \
        KLDRHEAP_ASSERT((uintptr_t)(pFree)->pPrev < (uintptr_t)(pFree)); \
        KLDRHEAP_ASSERT((uintptr_t)(pFree)->pNext > (uintptr_t)(pFree) || !(pFree)->pNext); \
    } while (0)

#else
# define KLDRHEAP_ASSERT(expr)          do { } while (0)
# define KLDRHEAP_ASSERT_BLOCK(pH, pB)  do { } while (0)
# define KLDRHEAP_ASSERT_FREE(pH, pF)   do { } while (0)
#endif


/**
 * A free heap block.
 */
typedef struct KLDRHEAPFREE
{
    /** The core bit which we have in common with used blocks. */
    KLDRHEAPBLOCK           Core;
    /** The next free block. */
    struct KLDRHEAPFREE    *pNext;
    /** The previous free block. */
    struct KLDRHEAPFREE    *pPrev;
} KLDRHEAPFREE, *PKLDRHEAPFREE;


/**
 * A heap segment.
 */
typedef struct KLDRHEAPSEG
{
    /** The base address of the segment. */
    void   *pvBase;
    /** The length of the segment (in bytes). */
    size_t  cb;
} KLDRHEAPSEG, *PKLDRHEAPSEG;

/**
 * Bundle of heap segments.
 */
typedef struct KLDRHEAPSEGS
{
    /** Pointer to the next segment bundle. */
    struct KLDRHEAPSEGS    *pNext;
    /** The number of segments used. */
    uint32_t                cSegs;
    /** Array of chunks. */
    KLDRHEAPSEG             aSegs[64];
} KLDRHEAPSEGS, *PKLDRHEAPSEGS;


/**
 * Heap anchor block.
 */
typedef struct KLDRHEAPANCHOR
{
    /** Head of the block list. */
    PKLDRHEAPBLOCK      pHead;
    /** Tail of the block list. */
    PKLDRHEAPBLOCK      pTail;
    /** Head of the free list. */
    PKLDRHEAPFREE       pFreeHead;
    /** Head segment bundle.
     * The order of this list is important, but a bit peculiar.
     * Logically, SegsHead::pNext is the tail pointer. */
    KLDRHEAPSEGS        SegsHead;
} KLDRHEAPANCHOR, *PKLDRHEAPANCHOR;



/*******************************************************************************
*   Global Variables                                                           *
*******************************************************************************/
/** The heap anchor block. */
static KLDRHEAPANCHOR g_Heap;


/*******************************************************************************
*   Internal Functions                                                         *
*******************************************************************************/
static int      kLdrHeapInit(PKLDRHEAPANCHOR pHeap);
static void     kLdrHeapDelete(PKLDRHEAPANCHOR pHeap);
static void *   kLdrHeapAlloc(PKLDRHEAPANCHOR pHeap, size_t cb);
static void     kLdrHeapFree(PKLDRHEAPANCHOR pHeap, void *pv);
static void     kLdrHeapDonate(PKLDRHEAPANCHOR pHeap, void *pv, size_t cb);
static int      kLdrHeapSegAlloc(PKLDRHEAPSEG pSeg, size_t cb);
static void     kLdrHeapSegFree(PKLDRHEAPSEG pSeg);


/**
 * Initializes the kLdr heap.
 *
 * @returns 0 on success, non-zero OS specific status code on failure.
 */
int kldrHlpHeapInit(void)
{
    return kLdrHeapInit(&g_Heap);
}


/**
 * Terminates the kLdr heap.
 */
void kldrHlpHeapTerm(void)
{
    kLdrHeapDelete(&g_Heap);
}


/**
 * Allocates memory from the kLdr heap.
 *
 * @returns Pointer to the allocated heap block. NULL if we can't satify the request.
 * @param   cb      The requested heap block size.
 */
void *kldrHlpAlloc(size_t cb)
{
    return kLdrHeapAlloc(&g_Heap, cb);
}


/**
 * Allocates zero initialized memory from the kLdr heap.
 *
 * @returns Pointer to the allocated heap block. NULL if we can't satify the request.
 * @param   cb      The requested heap block size.
 */
void *kldrHlpAllocZ(size_t cb)
{
    void *pv = kLdrHeapAlloc(&g_Heap, cb);
    if (pv)
        kLdrHlpMemSet(pv, 0, cb);
    return pv;
}


/**
 * Frees memory allocated off the kLdr heap.
 *
 * @param   pv      Pointer to the heap block returned by kldrHlpAlloc().
 */
void kldrHlpFree(void *pv)
{
    kLdrHeapFree(&g_Heap, pv);
}


/**
 * Donates memory to the heap.
 *
 * @param   pv      The address of the memory.
 * @param   cb      The amount of memory.
 */
void kldrHlpHeapDonate(void *pv, size_t cb)
{
    kLdrHeapDonate(&g_Heap, pv, cb);
}



/**
 * Initializes the heap anchor.
 *
 * @returns 0 on success, non-zero on failure.
 * @param   pHeap   The heap anchor to be initialized.
 */
static int kLdrHeapInit(PKLDRHEAPANCHOR pHeap)
{
    pHeap->pHead = NULL;
    pHeap->pTail = NULL;
    pHeap->pFreeHead = NULL;
    pHeap->SegsHead.pNext = NULL;
    pHeap->SegsHead.cSegs = 0;
    return 0;
}


/**
 * Deletes a heap.
 * This will free all resources (memory) associated with the heap.
 *
 * @param   pHeap   The heap to be deleted.
 */
static void kLdrHeapDelete(PKLDRHEAPANCHOR pHeap)
{
    /*
     * Free the segments, LIFO order.
     * The head element is the last to be free, while the
     * head.pNext is really the tail pointer - neat or what?
     */
    while (     pHeap->SegsHead.cSegs
           ||   pHeap->SegsHead.pNext)
    {
        /* find the tail. */
        uint32_t        iSeg;
        PKLDRHEAPSEGS   pSegs = pHeap->SegsHead.pNext;
        if (!pSegs)
            pSegs = &pHeap->SegsHead;
        else
        {
            pHeap->SegsHead.pNext = pSegs->pNext;
            pSegs->pNext = NULL;
        }

        /* free the segments */
        iSeg = pSegs->cSegs;
        while (iSeg-- > 0)
            kLdrHeapSegFree(&pSegs->aSegs[iSeg]);
        pSegs->cSegs = 0;
    }

    /* Zap the anchor. */
    pHeap->pHead = NULL;
    pHeap->pTail = NULL;
    pHeap->pFreeHead = NULL;
    pHeap->SegsHead.pNext = NULL;
    pHeap->SegsHead.cSegs = 0;
}


/**
 * Internal heap block allocator.
 */
static void *   kldrHeapAllocSub(PKLDRHEAPANCHOR pHeap, size_t cb)
{
    /*
     * Find a fitting free block.
     */
    const size_t    cbReq = KLDR_ALIGN_Z(cb + sizeof(KLDRHEAPBLOCK), KLDRHEAPBLOCK_ALIGNMENT);
    PKLDRHEAPFREE   pCur = pHeap->pFreeHead;
    while (pCur)
    {
        if (pCur->Core.cb >= cbReq)
        {
            if (pCur->Core.cb != cbReq)
            {
                /* check and see if there is a better match close by. */
                PKLDRHEAPFREE pCur2 = pCur->pNext;
                unsigned i = 16;
                while (i-- > 0 && pCur2)
                {
                    if (pCur2->Core.cb >= cbReq)
                    {
                        if (pCur2->Core.cb == cbReq)
                        {
                            pCur = pCur2;
                            break;
                        }
                        if (pCur2->Core.cb < pCur->Core.cb)
                            pCur = pCur2;
                    }

                    /* next */
                    KLDRHEAP_ASSERT_FREE(pHeap, pCur2);
                    pCur2 = pCur2->pNext;
                }
            }
            break;
        }

        /* next */
        KLDRHEAP_ASSERT_FREE(pHeap, pCur);
        pCur = pCur->pNext;
    }
    if (!pCur)
        return NULL;
    KLDRHEAP_ASSERT_FREE(pHeap, pCur);

    /*
     * Do we need to split out a block?
     */
    if (pCur->Core.cb - cbReq >= KLDRHEAPBLOCK_ALIGNMENT * 2)
    {
        PKLDRHEAPBLOCK pNew;

        pCur->Core.cb -= cbReq;

        pNew = (PKLDRHEAPBLOCK)((uintptr_t)pCur + pCur->Core.cb);
        pNew->fFlags = 0;
        pNew->cb = cbReq;
        pNew->pNext = pCur->Core.pNext;
        if (pNew->pNext)
            pNew->pNext->pPrev = pNew;
        else
            pHeap->pTail = pNew;
        pNew->pPrev = &pCur->Core;
        pCur->Core.pNext = pNew;

        KLDRHEAP_ASSERT_FREE(pHeap, pCur);
        KLDRHEAP_ASSERT_BLOCK(pHeap, pNew);
        return pNew + 1;
    }

    /*
     * No, just unlink it from the free list and return.
     */
    if (pCur->pNext)
        pCur->pNext->pPrev = pCur->pPrev;
    if (pCur->pPrev)
        pCur->pPrev->pNext = pCur->pNext;
    else
        pHeap->pFreeHead = pCur->pNext;
    pCur->Core.fFlags &= ~KLDRHEAPBLOCK_FLAG_FREE;

    KLDRHEAP_ASSERT_BLOCK(pHeap, &pCur->Core);
    return &pCur->Core + 1;
}


/**
 * Allocate a heap block.
 *
 * @returns Pointer to the allocated heap block on success. On failure NULL is returned.
 * @param   pHeap   The heap.
 * @param   cb      The requested heap block size.
 */
static void *   kLdrHeapAlloc(PKLDRHEAPANCHOR pHeap, size_t cb)
{
    void *pv;

    /* adjust the requested block size. */
    cb = KLDR_ALIGN_Z(cb, KLDRHEAPBLOCK_ALIGNMENT);
    if (!cb)
        cb = KLDRHEAPBLOCK_ALIGNMENT;

    /* try allocate the block. */
    pv = kldrHeapAllocSub(pHeap, cb);
    if (!pv)
    {
        /*
         * Failed, add another segment and try again.
         */
        KLDRHEAPSEG Seg;
        if (kLdrHeapSegAlloc(&Seg, cb + sizeof(KLDRHEAPSEGS) + sizeof(KLDRHEAPBLOCK) * 16))
            return NULL;

        /* donate before insterting the segment, this makes sure we got heap to expand the segment list. */
        kLdrHeapDonate(pHeap, Seg.pvBase, Seg.cb);

        /* insert the segment. */
        if (pHeap->SegsHead.cSegs < sizeof(pHeap->SegsHead.aSegs) / sizeof(pHeap->SegsHead.aSegs[0]))
            pHeap->SegsHead.aSegs[pHeap->SegsHead.cSegs++] = Seg;
        else if (   pHeap->SegsHead.pNext
                 && pHeap->SegsHead.pNext->cSegs < sizeof(pHeap->SegsHead.aSegs) / sizeof(pHeap->SegsHead.aSegs[0]))
            pHeap->SegsHead.pNext->aSegs[pHeap->SegsHead.pNext->cSegs++] = Seg;
        else
        {
            PKLDRHEAPSEGS pSegs = (PKLDRHEAPSEGS)kldrHeapAllocSub(pHeap, sizeof(*pSegs));
            KLDRHEAP_ASSERT(pSegs);
            pSegs->pNext = pHeap->SegsHead.pNext;
            pHeap->SegsHead.pNext = pSegs;
            pSegs->aSegs[0] = Seg;
            pSegs->cSegs = 1;
        }

        /* retry (should succeed) */
        pv = kldrHeapAllocSub(pHeap, cb);
        KLDRHEAP_ASSERT(pv);
    }

    return pv;
}


/**
 * Frees a heap block.
 *
 * @param   pHeap   The heap.
 * @param   pv      The pointer returned by kLdrHeapAlloc().
 */
static void     kLdrHeapFree(PKLDRHEAPANCHOR pHeap, void *pv)
{
    PKLDRHEAPFREE pFree, pLeft, pRight;

    /* ignore NULL pointers. */
    if (!pv)
        return;

    pFree = (PKLDRHEAPFREE)((PKLDRHEAPBLOCK)pv - 1);
    KLDRHEAP_ASSERT_BLOCK(pHeap, &pFree->Core);
    KLDRHEAP_ASSERT(KLDRHEAPBLOCK_IS_ALLOCATED(&pFree->Core));

    /*
     * Merge or link with left node?
     */
    pLeft = (PKLDRHEAPFREE)pFree->Core.pPrev;
    if (    pLeft
        &&  KLDRHEAPBLOCK_IS_FREE(&pLeft->Core)
        &&  KLDRHEAPBLOCK_IS_ADJACENT(&pLeft->Core, &pFree->Core)
       )
    {
        /* merge left */
        pLeft->Core.pNext = pFree->Core.pNext;
        if (pFree->Core.pNext)
            pFree->Core.pNext->pPrev = &pLeft->Core;
        else
            pHeap->pTail = &pLeft->Core;

        pLeft->Core.cb += pFree->Core.cb;
        pFree->Core.fFlags = ~0;
        pFree = pLeft;
    }
    else
    {
        /* link left */
        while (pLeft && !KLDRHEAPBLOCK_IS_FREE(&pLeft->Core))
            pLeft = (PKLDRHEAPFREE)pLeft->Core.pPrev;
        if (pLeft)
        {
            pFree->pPrev = pLeft;
            pFree->pNext = pLeft->pNext;
            if (pLeft->pNext)
                pLeft->pNext->pPrev = pFree;
            pLeft->pNext  = pFree;
        }
        else
        {
            pFree->pPrev = NULL;
            pFree->pNext = pHeap->pFreeHead;
            if (pHeap->pFreeHead)
                pHeap->pFreeHead->pPrev = pFree;
            pHeap->pFreeHead = pFree;
        }
        pFree->Core.fFlags |= KLDRHEAPBLOCK_FLAG_FREE;
    }
    KLDRHEAP_ASSERT_FREE(pHeap, pFree);

    /*
     * Merge right?
     */
    pRight = (PKLDRHEAPFREE)pFree->Core.pNext;
    if (    pRight
        &&  KLDRHEAPBLOCK_IS_FREE(&pRight->Core)
        &&  KLDRHEAPBLOCK_IS_ADJACENT(&pFree->Core, pRight)
       )
    {
        /* unlink pRight from the global list. */
        pFree->Core.pNext = pRight->Core.pNext;
        if (pRight->Core.pNext)
            pRight->Core.pNext->pPrev = &pFree->Core;
        else
            pHeap->pTail = &pFree->Core;

        /* unlink pRight from the free list. */
        pFree->pNext = pRight->pNext;
        if (pRight->pNext)
            pRight->pNext->pPrev = pFree;

        /* update size and invalidate pRight. */
        pFree->Core.cb += pRight->Core.cb;
        pRight->Core.fFlags = ~0;
    }
}


/**
 * Donates memory to the heap.
 *
 * The donated memory is returned to the donator when the heap is deleted.
 *
 * @param pHeap The heap
 * @param pv    The pointer to the donated memory.
 * @param cb    Size of the donated memory.
 */
static void     kLdrHeapDonate(PKLDRHEAPANCHOR pHeap, void *pv, size_t cb)
{
    PKLDRHEAPBLOCK pBlock;

    /*
     * Don't bother with small donations.
     */
    if (cb < KLDRHEAPBLOCK_ALIGNMENT * 4)
        return;

    /*
     * Align the donation on a heap block boundrary.
     */
    if ((uintptr_t)pv & (KLDRHEAPBLOCK_ALIGNMENT - 1))
    {
        cb -= (uintptr_t)pv & 31;
        pv = KLDR_ALIGN_P(pv, KLDRHEAPBLOCK_ALIGNMENT);
    }
    cb &= ~(size_t)(KLDRHEAPBLOCK_ALIGNMENT - 1);

    /*
     * Create an allocated block, link it and free it.
     */
    pBlock = (PKLDRHEAPBLOCK)pv;
    pBlock->pNext = NULL;
    pBlock->pPrev = NULL;
    pBlock->cb = cb;
    pBlock->fFlags = 0;

    /* insert */
    if ((uintptr_t)pBlock < (uintptr_t)pHeap->pHead)
    {
        /* head */
        pBlock->pNext = pHeap->pHead;
        pHeap->pHead->pPrev = pBlock;
        pHeap->pHead = pBlock;
    }
    else if ((uintptr_t)pBlock > (uintptr_t)pHeap->pTail)
    {
        if (pHeap->pTail)
        {
            /* tail */
            pBlock->pPrev = pHeap->pTail;
            pHeap->pTail->pNext = pBlock;
            pHeap->pTail = pBlock;
        }
        else
        {
            /* first */
            pHeap->pHead = pBlock;
            pHeap->pTail = pBlock;
        }
    }
    else
    {
        /* in list (unlikely) */
        PKLDRHEAPBLOCK pPrev = pHeap->pHead;
        PKLDRHEAPBLOCK pCur = pPrev->pNext;
        for (;;)
        {
            KLDRHEAP_ASSERT_BLOCK(pHeap, pCur);
            if ((uintptr_t)pCur > (uintptr_t)pBlock)
                break;
            pPrev = pCur;
            pCur = pCur->pNext;
        }

        pBlock->pNext = pCur;
        pBlock->pPrev = pPrev;
        pPrev->pNext = pBlock;
        pCur->pPrev = pBlock;
    }
    KLDRHEAP_ASSERT_BLOCK(pHeap, pBlock);

    /* free it */
    kLdrHeapFree(pHeap, pBlock + 1);
}



/**
 * Allocates a new segment.
 *
 * @returns 0 on success, non-zero OS status code on failure.
 * @param   pSeg    Where to put the info about the allocated segment.
 * @param   cbMin   The minimum segment size.
 */
static int kLdrHeapSegAlloc(PKLDRHEAPSEG pSeg, size_t cbMin)
{
#ifdef __OS2__
    APIRET rc;

    pSeg->cb = (cbMin + 0xffff) & ~(size_t)0xffff;
    pSeg->pvBase = NULL;
    rc = DosAllocMem(&pSeg->pvBase, pSeg->cb, PAG_COMMIT | PAG_READ | PAG_WRITE | OBJ_ANY);
    if (rc == ERROR_INVALID_PARAMETER)
        rc = DosAllocMem(&pSeg->pvBase, pSeg->cb, PAG_COMMIT | PAG_READ | PAG_WRITE);
    if (rc)
    {
        pSeg->pvBase = NULL;
        pSeg->cb = 0;
        return rc;
    }

#elif defined(__WIN__)
    pSeg->cb = (cbMin + 0xffff) & ~(size_t)0xffff;
    pSeg->pvBase = VirtualAlloc(NULL, pSeg->cb, MEM_COMMIT, PAGE_READWRITE);
    if (!pSeg->pvBase)
    {
        pSeg->cb = 0;
        return GetLastError();
    }

#else
# error "Port me"
#endif

    return 0;
}


/**
 * Frees a segment.
 *
 * @param   pSeg    The segment to be freed.
 */
static void kLdrHeapSegFree(PKLDRHEAPSEG pSeg)
{
#ifdef __OS2__
    APIRET rc = DosFreeMem(pSeg->pvBase);
    KLDRHEAP_ASSERT(!rc); (void)rc;

#elif defined(__WIN__)
    BOOL fRc = VirtualFree(pSeg->pvBase, 0 /*pSeg->cb*/, MEM_RELEASE);
    KLDRHEAP_ASSERT(fRc); (void)fRc;

#else
# error "Port me"
#endif
}