/*
//
// suballoc.c
//
// Memory suballocation
//
// Copyright (c) 1995-96 Jim Nelson.  Permission to distribute
// granted by the author.  No warranties are made on the fitness of this
// source code.
//
*/

#include "htp.h"

#ifdef USE_SUBALLOC

/*
// performance counters
*/
#if DEBUG
uint freePoolHits = 0;
uint totalAllocations = 0;
#endif

/*
// pattern used to clear memory on allocation and free
*/

#define SUBALLOC_ALLOC_BYTE                 (0xA5)
#define SUBALLOC_FREE_BYTE                  (0x5A)

/*
// each memory allocation has a small header, which is used by
// this module alone
*/

#define SUBALLOC_MAX_HEADER_FILE_SIZE               (16)

typedef struct tagSUBALLOC_HEADER
{
    struct tagSUBALLOC_HEADER *next;
    struct tagSUBALLOC_HEADER *prev;
    char            file[SUBALLOC_MAX_HEADER_FILE_SIZE];
    uint            line;
    uint            allocSize;
    uint            userSize;
    DWORD           signature;
} SUBALLOC_HEADER;

#define SUBALLOC_SIGNATURE                  (0xFEA55AEF)
#define SUBALLOC_END_SIGNATURE              (0xA5EFFE5A)

/*
// the memory pool is kept on a doubly-linked list
*/

typedef struct tagSUBALLOC_POOL
{
    SUBALLOC_HEADER *head;
    SUBALLOC_HEADER *tail;
    uint            totalMemorySize;
} SUBALLOC_POOL;

SUBALLOC_POOL freePool;
SUBALLOC_POOL reservedPool;

/*
// memory pool functions
*/

void InitializePool(SUBALLOC_POOL *pool)
{
    assert(pool != NULL);

    pool->head = NULL;
    pool->tail = NULL;
    pool->totalMemorySize = 0;
}

SUBALLOC_HEADER *PointerToHeader(void *ptr)
{
    assert(ptr != NULL);
    return (SUBALLOC_HEADER *) (((BYTE *) ptr) - sizeof(SUBALLOC_HEADER));
}

void *GetUserBuffer(SUBALLOC_HEADER *header)
{
    /* cant assert good header, this function might be used during its */
    /* construction */
    assert(header != NULL);
    return ((BYTE *) header) + sizeof(SUBALLOC_HEADER);
}

BOOL IsHeaderOkay(SUBALLOC_HEADER *header)
{
    if(header == NULL)
    {
        DebugMsg("IsHeaderOkay: NULL header\n");
        return FALSE;
    }

    /* check head signature */
    if(header->signature != SUBALLOC_SIGNATURE)
    {
        DebugMsg("IsHeaderOkay: bad start signature, file %s line %u\n",
            header->file, header->line);
        return FALSE;
    }

    return TRUE;
}

void SetEndSignature(SUBALLOC_HEADER *header)
{
    BYTE *startSignature;
    static DWORD endSignature = SUBALLOC_END_SIGNATURE;

    assert(header != NULL);

    /* find the first byte of the area beyond the user-allocated buffer */
    startSignature = (BYTE *) header + sizeof(SUBALLOC_HEADER)
        + header->userSize;

    /* since some machines don't like misaligned accesses, copy the signature */
    /* in byte by byte */
    memcpy(startSignature, &endSignature, sizeof(DWORD));
}

BOOL IsEndSignatureOkay(SUBALLOC_HEADER *header)
{
    BYTE *startSignature;
    static DWORD endSignature = SUBALLOC_END_SIGNATURE;

    assert(header != NULL);

    /* find the first byte beyond the user-allocated buffer */
    startSignature = (BYTE *) header + sizeof(SUBALLOC_HEADER)
        + header->userSize;

    /* since some machines don't like misaligned accessed, compare the */
    /* signature byte by byte */
    return (memcmp(startSignature, &endSignature, sizeof(DWORD)) == 0)
        ? TRUE : FALSE;
}

SUBALLOC_HEADER *CreatePoolElement(uint size)
{
    SUBALLOC_HEADER *header;
    uint allocSize;

#if SUBALLOC_MINALLOCSIZE
    allocSize = (size < SUBALLOC_MINALLOCSIZE) ? SUBALLOC_MINALLOCSIZE : size;
#else
    allocSize = size;
#endif

    /* allocate header, end signature, and user buffer */
    if((header = malloc(allocSize + sizeof(SUBALLOC_HEADER) + sizeof(DWORD))) == NULL)
    {
        return NULL;
    }

#if DEBUG
    header->next = (void *) 0x12345678;
    header->prev = (void *) 0x12345678;
#endif

    /* set up as much of the header as possible */
    header->allocSize = allocSize;
    header->userSize = size;
    header->signature = SUBALLOC_SIGNATURE;

#if SUBALLOC_CLEARMEM
    memset(GetUserBuffer(header), SUBALLOC_ALLOC_BYTE, allocSize);
#endif

    /* set the end signature */
    SetEndSignature(header);

    assert(IsHeaderOkay(header) == TRUE);
    assert(IsEndSignatureOkay(header) == TRUE);

    return header;
}

void FreePoolElement(SUBALLOC_HEADER *header)
{
    assert(IsHeaderOkay(header) == TRUE);

#if SUBALLOC_CLEARMEM
    memset(GetUserBuffer(header), SUBALLOC_FREE_BYTE, header->allocSize);
#endif

    free(header);
}

SUBALLOC_HEADER *ResizePoolElement(SUBALLOC_HEADER *header, uint newSize)
{
    uint allocSize;
    SUBALLOC_HEADER *newHeader;

    assert(IsHeaderOkay(header) == TRUE);
    assert(newSize != 0);

#if SUBALLOC_MINALLOCSIZE
    allocSize = (newSize < SUBALLOC_MINALLOCSIZE) ? SUBALLOC_MINALLOCSIZE
        : newSize;
#else
    allocSize = newSize;
#endif

    /* it COULD already be this size */
    if(header->allocSize >= allocSize)
    {
        header->userSize = newSize;
        SetEndSignature(header);

        assert(IsHeaderOkay(header) == TRUE);
        assert(IsEndSignatureOkay(header) == TRUE);

        return header;
    }

    newHeader = realloc(header, newSize + sizeof(SUBALLOC_HEADER) + sizeof(DWORD));
    if(newHeader == NULL)
    {
        /* couldnt resize block */
        return NULL;
    }

#if DEBUG
    newHeader->prev = (void *) 0x12345678;
    newHeader->next = (void *) 0x12345678;
#endif

    /* fill in the new header */
    newHeader->userSize = newSize;
    newHeader->allocSize = allocSize;
    newHeader->signature = SUBALLOC_SIGNATURE;

    /* set the end signature */
    SetEndSignature(newHeader);

    assert(IsHeaderOkay(newHeader) == TRUE);
    assert(IsEndSignatureOkay(newHeader) == TRUE);

    return newHeader;
}

void AddPoolElement(SUBALLOC_POOL *pool, SUBALLOC_HEADER *header)
{
    assert(pool != NULL);
    assert(IsHeaderOkay(header) == TRUE);

    if(pool->head != NULL)
    {
        header->prev = pool->tail;
        pool->tail->next = header;
    }
    else
    {
        pool->head = header;
        header->prev = NULL;
    }

    header->next = NULL;
    pool->tail = header;

    pool->totalMemorySize += header->allocSize;
}

SUBALLOC_HEADER *RemoveFirstElement(SUBALLOC_POOL *pool)
{
    SUBALLOC_HEADER *header;

    assert(pool != NULL);

    if(pool->head == NULL)
    {
        return NULL;
    }

    header = pool->head;
    if((pool->head = header->next) != NULL)
    {
        pool->head->prev = NULL;
    }

    assert(pool->totalMemorySize >= header->allocSize);
    pool->totalMemorySize -= header->allocSize;

    assert(IsHeaderOkay(header) == TRUE);

    return header;
}

void RemoveElement(SUBALLOC_POOL *pool, SUBALLOC_HEADER *header)
{
    assert(pool != NULL);
    assert(IsHeaderOkay(header) == TRUE);

    if(pool->head == header)
    {
        pool->head = header->next;
        if(pool->head != NULL)
        {
            pool->head->prev = NULL;
        }
    }
    else if(pool->tail == header)
    {
        pool->tail = header->prev;
        if(pool->tail != NULL)
        {
            pool->tail->next = NULL;
        }
    }
    else
    {
        header->prev->next = header->next;
        header->next->prev = header->prev;
    }

    assert(pool->totalMemorySize >= header->allocSize);
    pool->totalMemorySize -= header->allocSize;
}

SUBALLOC_HEADER *RemovePoolElementBySize(SUBALLOC_POOL *pool, uint size)
{
    SUBALLOC_HEADER *curr;

#if !SUBALLOC_FIRSTFIT
    SUBALLOC_HEADER *bestfit;
    uint bestfitDiff;
    uint currDiff;
#endif

    assert(pool != NULL);
    assert(size != 0);

#if !SUBALLOC_FIRSTFIT
    bestfit = NULL;
    bestfitDiff = UINT_MAX;
#endif

    curr = pool->head;
    while(curr != NULL)
    {
        assert(IsHeaderOkay(curr) == TRUE);

        if(curr->allocSize < size)
        {
            /* too small */
            curr = curr->next;
            continue;
        }

#if SUBALLOC_FIRSTFIT
        /* found one, unlink from the list */
        RemoveElement(pool, curr);

#if SUBALLOC_DEBLOG
        DebugMsg("suballoc: first fit found block of %u bytes for alloc of %u bytes\n",
            curr->allocSize, size);
#endif

        return curr;
#else

#if SUBALLOC_MINALLOCSIZE
        /* if the block is the minimum allocation size, then it is the */
        /* best fit, by definition */
        if(curr->allocSize == SUBALLOC_MINALLOCSIZE)
        {
            bestfit = curr;
            break;
        }
#endif

        currDiff = curr->allocSize - size;
        if(currDiff == 0)
        {
            /* this is as good as it gets */
            bestfit = curr;
            break;
        }

        if(currDiff < bestfitDiff)
        {
            bestfitDiff = currDiff;
            bestfit = curr;
        }

        curr = curr->next;
#endif
    }

#if !SUBALLOC_FIRSTFIT
    if(bestfit != NULL)
    {
        RemoveElement(pool, bestfit);

#if SUBALLOC_DEBLOG
        DebugMsg("suballoc: best fit found block of %u bytes for alloc of %u bytes\n",
            bestfit->allocSize, size);
#endif

        return bestfit;
    }
#endif

    /* nothing was found */
    return NULL;
}

BOOL RemovePoolElementByAddr(SUBALLOC_POOL *pool, SUBALLOC_HEADER *header)
{
    SUBALLOC_HEADER *curr;

    assert(pool != NULL);
    assert(IsHeaderOkay(header) == TRUE);

    curr = pool->head;
    while(curr != NULL)
    {
        assert(IsHeaderOkay(curr) == TRUE);

        if(curr == header)
        {
            RemoveElement(pool, curr);
            return TRUE;
        }

        curr = curr->next;
    }

    return FALSE;
}

void RemoveLargestPoolElement(SUBALLOC_POOL *pool)
{
    SUBALLOC_HEADER *curr;
    SUBALLOC_HEADER *largest;

    assert(pool != NULL);

    largest = NULL;

    /* walk the list, looking for the largest allocated block */
    for(curr = pool->head; curr != NULL; curr = curr->next)
    {
        assert(IsHeaderOkay(curr) == TRUE);

        if(largest == NULL)
        {
            largest = curr;
            continue;
        }

        if(curr->allocSize > largest->allocSize)
        {
            largest = curr;
            continue;
        }
    }

    /* if one was found, remove it */
    if(largest != NULL)
    {
#if SUBALLOC_DEBLOG
        DebugMsg("removing largest buffer of %u bytes from pool of %u bytes\n",
            largest->allocSize, pool->totalMemorySize);
#endif

        if(largest->prev != NULL)
        {
            largest->prev->next = largest->next;
        }

        if(largest->next != NULL)
        {
            largest->next->prev = largest->prev;
        }

        /* removed, so destroy it and account for it in the pool */
        assert(pool->totalMemorySize >= largest->allocSize);
        pool->totalMemorySize -= largest->allocSize;
        FreePoolElement(largest);
    }
}

/*
// public functions
*/

void InitializeMemory(void)
{
    InitializePool(&freePool);
    InitializePool(&reservedPool);
}

void TerminateMemory(void)
{
    SUBALLOC_HEADER *header;

    if(reservedPool.head != NULL)
    {
#if SUBALLOC_DEBLOG
        DebugMsg("UNFREED MEMORY: %u bytes held in reserved pool\n",
            reservedPool.totalMemorySize);
#endif

#if SUBALLOC_WARNING
        printf("suballoc: %u bytes unfreed\n", reservedPool.totalMemorySize);
#endif
    }

#if SUBALLOC_DEBLOG
    DebugMsg("suballoc: %u bytes held in free pool\n", freePool.totalMemorySize);
#endif

    while((header = RemoveFirstElement(&reservedPool)) != NULL)
    {
        assert(IsHeaderOkay(header) == TRUE);
        assert(IsEndSignatureOkay(header) == TRUE);

#if SUBALLOC_DEBLOG
        DebugMsg("Removing reserved alloc for %u bytes from %s line %u\n", header->userSize,
            header->file, header->line);
#endif

        FreePoolElement(header);
    }

    while((header = RemoveFirstElement(&freePool)) != NULL)
    {
        assert(IsHeaderOkay(header) == TRUE);
        /* dont check end signature, it is not valid in free pool */

#if SUBALLOC_DEBLOG
        DebugMsg("Removing free alloc for %u bytes from %s line %u\n",
            header->userSize, header->file, header->line);
#endif
        FreePoolElement(header);
    }
}

void *_AllocMemory(uint size, const char *file, uint line)
{
    SUBALLOC_HEADER *header;

#if SUBALLOC_DEBLOG
    DebugMsg("Allocating %u bytes in file %s line %u\n", size, file, line);
#endif

#if DEBUG
    totalAllocations++;
#endif

    if(size == 0)
    {
#if SUBALLOC_DEBLOG
        DebugMsg("Allocation for 0 bytes in file %s line %u\n", file, line);
#endif
        return NULL;
    }

    /* if buffer of sufficient size already in pool, use that one */
    if((header = RemovePoolElementBySize(&freePool, size)) == NULL)
    {
        /* need to create a new pool item */
        if((header = CreatePoolElement(size)) == NULL)
        {
            /* out of memory! */
#if SUBALLOC_DEBLOG
            DebugMsg("System heap out of memory! file %s line %u\n", file, line);
#endif
            return NULL;
        }
    }
    else
    {
#if DEBUG
        freePoolHits++;
#endif
    }

    assert(IsHeaderOkay(header) == TRUE);

    /* replace the headers file and line with this new information */
    if(file != NULL)
    {
        StringCopy(header->file, file, SUBALLOC_MAX_HEADER_FILE_SIZE);
    }
    else
    {
        header->file[0] = NUL;
    }
    header->line = line;

    header->userSize = size;

    /* set the end signature */
    SetEndSignature(header);

    /* add the new suballoc to the reserved pool */
    AddPoolElement(&reservedPool, header);

    assert(IsHeaderOkay(header) == TRUE);
    assert(IsEndSignatureOkay(header) == TRUE);

    return GetUserBuffer(header);
}

void _FreeMemory(void *ptr, const char *file, uint line)
{
    SUBALLOC_HEADER *header;

#if !SUBALLOC_DEBLOG
    UNREF_PARAM(file);
    UNREF_PARAM(line);
#endif

    if(ptr == NULL)
    {
#if SUBALLOC_DEBLOG
        DebugMsg("Tried to free NULL pointer in file %s line %u\n", file, line);
#endif
        return;
    }

    header = (SUBALLOC_HEADER *) (((BYTE *) ptr) - sizeof(SUBALLOC_HEADER));
    assert(IsHeaderOkay(header) == TRUE);
    assert(IsEndSignatureOkay(header) == TRUE);

#if SUBALLOC_DEBLOG
    DebugMsg("Freeing %u bytes (%u alloced) in file %s line %u\n",
        header->userSize, header->allocSize, file, line);
#endif

    if(RemovePoolElementByAddr(&reservedPool, header) == FALSE)
    {
        /* doh! */
        HtpMsg(MSG_ERROR, NULL, "Fatal error: block of memory freed not in reserved pool, exiting ...\n");
        exit(1);
    }

#if SUBALLOC_CLEARMEM
    memset(GetUserBuffer(header), SUBALLOC_FREE_BYTE, header->allocSize);
#endif

    /* move it back into the free pool */
    AddPoolElement(&freePool, header);

#if SUBALLOC_MAXFREEPOOLSIZE
    /* check to see that the free pool size hasnt grown too big */
    while(freePool.totalMemorySize > SUBALLOC_MAXFREEPOOLSIZE)
    {
        RemoveLargestPoolElement(&freePool);
    }
#endif
}

void *_ResizeMemory(void *ptr, uint newSize, const char *file, uint line)
{
    SUBALLOC_HEADER *header;
    SUBALLOC_HEADER *newHeader;

#if SUBALLOC_DEBLOG
    DebugMsg("Reallocing to new size of %u bytes in file %s line %u\n",
        newSize, file, line);
#endif

    if((newSize == 0) && (ptr != NULL))
    {
        /* treat as a free */
        FreeMemory(ptr);
        return NULL;
    }

    if((ptr == NULL) && (newSize != 0))
    {
        /* treat as an alloc */
        return AllocMemory(newSize);
    }

    if((ptr == NULL) && (newSize == 0))
    {
        /* treat as a dumbo programmer */
        return NULL;
    }

    header = PointerToHeader(ptr);
    assert(IsHeaderOkay(header) == TRUE);
    assert(IsEndSignatureOkay(header) == TRUE);

    if(RemovePoolElementByAddr(&reservedPool, header) == FALSE)
    {
        /* oof */
        HtpMsg(MSG_ERROR, NULL, "Fatal error: block of memory freed not in reserved pool, exiting ...\n");
        exit(1);
    }

    /* resize memory buffer */
    newHeader = ResizePoolElement(header, newSize);
    if(newHeader == NULL)
    {
        /* failed to resize */
        AddPoolElement(&reservedPool, header);
        return NULL;
    }

    /* succeeded, fill in rest of header */
    if(file != NULL)
    {
        StringCopy(newHeader->file, file, SUBALLOC_MAX_HEADER_FILE_SIZE);
    }
    else
    {
        newHeader->file[0] = NUL;
    }
    newHeader->line = line;

    /* add to reserve pool and return to user */
    AddPoolElement(&reservedPool, newHeader);

    assert(IsHeaderOkay(newHeader) == TRUE);
    assert(IsEndSignatureOkay(newHeader) == TRUE);

    return GetUserBuffer(newHeader);
}

uint MemorySize(void *ptr)
{
    SUBALLOC_HEADER *header;

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

    header = PointerToHeader(ptr);
    assert(IsHeaderOkay(header) == TRUE);

    if(IsHeaderOkay(header) == FALSE)
    {
        return 0;
    }

    return header->userSize;
}

#endif
