//////////////////////////////////////////////////////////////////////
// Heap.cc
//

// Copyright 1999 by Steven McDougall
// This code is freeware
// It may be copied/modified/redistributed without restriction

#include "Heap.h"
#include <memory.h>
#include <assert.h>
#include <stdio.h>

// Allocated blocks are aligned on kAlign byte boundaries
// Set kAlign to whatever your architecture requires
// Must be a power of 2
const int kAlign     = 4;         

// Allocated blocks are surrounded by kGuard extra bytes,
// to help protect against/detect overwrites
// Must be a multiple of kAlign (possibly 0)
const int kGuard     = kAlign;      

const int kMagic     = 0x53574d20;  // 'SWM '
const int kAlignMask = kAlign-1;


enum ePattern 
{ 
    kpUninit = 0xcc, // fill pattern for uninitialized storage
    kpFreed  = 0xcd  // fill pattern for free'd storage
};


struct Block
{
    unsigned	 magic;     // to detect corruption
    Block	*prev;      // pointer to previous block
    Block	*next;      // pointer to next     block
    unsigned	 size;      // size of this block, 
                            // not including the Block struct
    unsigned	 userSize;  // number of bytes in this block allocated to user
    unsigned	 fFree:1;   // set if block is free
};


static char	 gacGuard[kGuard]; // buffer for checking guard bytes
static Block	*gpStart;          // current search position
static int       gHeapSize;        // used by CheckBlock() for validation
static int       gBlockSize;       // sizeof Block struct, rounded up to kAlign


static unsigned  Roundup   (unsigned n);
static unsigned  RoundDelta(char *pc);
static char     *Alloc     (unsigned userSize, unsigned size);
static void      Split     (Block *pOldBlock, unsigned size);
static void      Join      (Block *pBlock1);
static void      DumpBlock (Block *pBlock);
static void      CheckBlock(Block *pBlock);


// Initialize the heap
void Init(char *pc,     // need not be aligned
          unsigned size)
{
    unsigned nDelta = RoundDelta(pc);
    pc   += nDelta;
    size -= nDelta;

    gBlockSize = Roundup(sizeof(Block));

    Block *pBlock	= reinterpret_cast<Block *>(pc);
    pBlock->magic 	= kMagic;
    pBlock->prev  	= pBlock;
    pBlock->next  	= pBlock;
    pBlock->size  	= size - gBlockSize;
    pBlock->userSize 	= 0;
    pBlock->fFree	= 1;
	
    memset(gacGuard, kpUninit, kGuard);
    gpStart    = pBlock;
    gHeapSize  = size;
}


// Allocate a block of userSize bytes
// The block will be aligned on a kAlign byte boundary
char *Malloc(unsigned userSize)
{
    unsigned size = kGuard + Roundup(userSize) + kGuard;

    Block *pStart = gpStart;
	
    do
    {
	CheckBlock(gpStart);
	if (gpStart->fFree && size <= gpStart->size)
	    return Alloc(userSize, size);
	gpStart = gpStart->next;	
    } while (gpStart != pStart);
	
    return 0;
}


// Round n up to a kAlign byte boundary
unsigned Roundup(unsigned n)
{
    return (n+kAlign-1) & ~kAlignMask;
}


// Return the delta necessary to round pc up to a kAlign byte boundary
unsigned RoundDelta(char *pc)
{
    unsigned n = reinterpret_cast<unsigned>(pc);
    return ((n+kAlign-1) & ~kAlignMask) - n;
}


// Allocate the block at the current search position
char *Alloc(unsigned userSize, unsigned size)
{
    unsigned leftover = gpStart->size - size;
    if (leftover >= gBlockSize + kGuard + kAlign + kGuard)
	Split(gpStart, size);

    gpStart->userSize = userSize;
    gpStart->fFree    = 0;

    char *pc = reinterpret_cast<char *>(gpStart);
    memset(pc+gBlockSize, kpUninit, gpStart->size);
       
    gpStart = gpStart->next;
    return pc + gBlockSize + kGuard;
}


// Split a block in two.
void Split(Block *pOldBlock, unsigned size)
{
    char  *pcOldBlock = reinterpret_cast<char  *>(pOldBlock);
    char  *pcNewBlock = pcOldBlock + gBlockSize + size;
    Block *pNewBlock  = reinterpret_cast<Block *>(pcNewBlock);

    pNewBlock->magic 	= kMagic;
    pNewBlock->prev  	= pOldBlock;
    pNewBlock->next  	= pOldBlock->next;
    pNewBlock->size  	= pOldBlock->size - size - gBlockSize;
    pNewBlock->userSize	= 0;
    pNewBlock->fFree	= 1;
	
    pOldBlock->size = size;
    pOldBlock->next->prev = pNewBlock;
    pOldBlock->next       = pNewBlock;
}


// Free an allocated block
void Free(char *pc)
{
    Block *pBlock = reinterpret_cast<Block *>(pc - kGuard - gBlockSize);
    CheckBlock(pBlock);
    assert(pBlock->fFree==0);
    assert(!memcmp(pc-kGuard	      , gacGuard, kGuard));
    assert(!memcmp(pc+pBlock->userSize, gacGuard, kGuard));

    pBlock->fFree    = 1;
    pBlock->userSize = 0;
    pc = reinterpret_cast<char *>(pBlock);
    memset(pc+gBlockSize, kpFreed, pBlock->size);

    if (pBlock->next->fFree) Join(pBlock);
    if (pBlock->prev->fFree) Join(pBlock->prev);
}


// Join two adjacent blocks
void Join(Block *pBlock1)
{
    Block *pBlock2 = pBlock1->next;

    if (pBlock2<=pBlock1) // don't merge the tail to the head
	return;

    if (gpStart==pBlock2)
	gpStart = pBlock1;

    pBlock1->size = pBlock1->size + gBlockSize + pBlock2->size;
    pBlock1->next = pBlock2->next;
    pBlock2->next->prev = pBlock1;
	
    char *pc = reinterpret_cast<char *>(pBlock1);
    memset(pc+gBlockSize, kpFreed, pBlock1->size);
}


// Dump out all the Block structs in the heap
void Dump()
{
    printf("\n--------------------------------\n");

    Block *pStart = gpStart;

    do
    {
	DumpBlock(pStart);
	pStart = pStart->next;
    } while (pStart != gpStart);
	
}


// Dump a single Block struct
void DumpBlock(Block *pBlock)
{
    printf("Block    %p\n"  ,  pBlock);
    printf("magic    %.4s\n", &pBlock->magic);
    printf("prev     %p\n"  ,  pBlock->prev);
    printf("next     %p\n"  ,  pBlock->next);
    printf("size     %d\n"  ,  pBlock->size);
    printf("userSize %d\n"  ,  pBlock->userSize);
    printf("free     %d\n"  ,  pBlock->fFree);
    printf("\n");
}


// Print summary statistics on heap usage
void Stats()
{
    int nUserBytes  = 0;
    int nFreeBytes  = 0;
    int nOverhead   = 0;
    int nUserBlocks = 0;
    int nFreeBlocks = 0;

    Block *pStart = gpStart;

    do
    {
	if (pStart->fFree)
	{
	    nFreeBlocks++;
	    nFreeBytes += pStart->size - kGuard - kGuard;
	    nOverhead  += gBlockSize + kGuard + kGuard;
	}
	else
	{
	    nUserBlocks++;
	    nUserBytes += pStart->userSize;
	    nOverhead  += gBlockSize + pStart->size - pStart->userSize;
	}

	pStart = pStart->next;
    } while (pStart != gpStart);

    int nTotalBytes = nFreeBytes + nUserBytes + nOverhead;
    
    printf("\n--------------------------------\n");
    printf("%d free blocks, %d user blocks\n", nFreeBlocks, nUserBlocks);
    printf("%d free + %d user + %d overhead = %d total bytes\n",
	   nFreeBytes, nUserBytes, nOverhead, nTotalBytes);
}


// Do an integrity check on the entire heap
void Check()
{
    Block *pStart = gpStart;

    do
    {
	CheckBlock(pStart);
	pStart = pStart->next;
    } while (pStart != gpStart);
}


// Do an integrity check on a single block in the heap
void CheckBlock(Block *pBlock)
{
    assert(pBlock->magic==kMagic);
    assert(pBlock->next->prev==pBlock);
    assert(pBlock->prev->next==pBlock);

    char *pcBlock = reinterpret_cast<char *>(pBlock);
    char *pcNext  = reinterpret_cast<char *>(pBlock->next);

    char *pcEnd   = pcBlock+gBlockSize+pBlock->size;
    
    if (pcBlock<pcNext)
    {
	assert(pcEnd==pcNext);
    }
    else
    {
	assert(pcEnd-pcNext==gHeapSize);
    }
}