//***********************************************************************/
//    Author                    : Garry
//    Original Date             : May,20 2006
//    Module Name               : HEAP.CPP
//    Module Funciton           : 
//                                This module countains Heap's implementation code.
//    Last modified Author      :
//    Last modified Date        :
//    Last modified Content     :
//                                1.
//                                2.
//    Lines number              :
//***********************************************************************/

#ifndef __STDAFX_H__
#include "StdAfx.h"
#endif

//#ifndef __HEAP_H__
//#include "HEAP.H"
//#endif

//
//The implementation of CreateHeap routine.
//This routine does the following:
// 1. Allocate a virtual area according to dwInitSize;
// 3. Create a virtual area node object,to manage the virtual area;
// 2. Create a heap object,and initialize it;
// 4. Insert the virtual area node into heap's virtual area node list;
// 5. Insert the virtual area into heap object's free list;
// 6. Insert the heap object into kernel thread's list;
// 7. If all successful,return the help object's base address.
//
static __HEAP_OBJECT* CreateHeap(DWORD dwInitSize)
{
	__HEAP_OBJECT*              lpHeapObject   = NULL;
	__HEAP_OBJECT*              lpHeapRoot     = NULL;
	__VIRTUAL_AREA_NODE*        lpVirtualArea  = NULL;
	__FREE_BLOCK_HEADER*        lpFreeHeader   = NULL;
	LPVOID                      lpVirtualAddr  = NULL;
	BOOL                        bResult        = FALSE;
	DWORD                       dwFlags        = 0;

	if(dwInitSize > MAX_VIRTUAL_AREA_SIZE)  //Requested size too big.
		return NULL;
	if(dwInitSize < DEFAULT_VIRTUAL_AREA_SIZE)
		dwInitSize = DEFAULT_VIRTUAL_AREA_SIZE;

	//
	//Now,allocate the virtual area.
	//
	lpVirtualAddr = GET_VIRTUAL_AREA(dwInitSize);
	if(NULL == lpVirtualAddr)   //Can not get virtual area.
		goto __TERMINAL;
	lpFreeHeader = (__FREE_BLOCK_HEADER*)lpVirtualAddr;
	lpFreeHeader->dwFlags     = BLOCK_FLAGS_FREE;
	lpFreeHeader->dwBlockSize = dwInitSize - sizeof(__FREE_BLOCK_HEADER); //Caution!!!

	//
	//Now,create a virtual area node object,to manage virtual area.
	//
	lpVirtualArea = (__VIRTUAL_AREA_NODE*)GET_KERNEL_MEMORY(
		sizeof(__VIRTUAL_AREA_NODE));
	if(NULL == lpVirtualArea)  //Can not get memory.
		goto __TERMINAL;
	lpVirtualArea->lpStartAddress = lpVirtualAddr;
	lpVirtualArea->dwAreaSize     = dwInitSize;
	lpVirtualArea->lpNext         = NULL;

	//
	//Now,create a heap object,and initialize it.
	//
	lpHeapObject = (__HEAP_OBJECT*)GET_KERNEL_MEMORY(sizeof(__HEAP_OBJECT));
	if(NULL == lpHeapObject)  //Can not allocate memory.
		goto __TERMINAL;
	lpHeapObject->lpVirtualArea               = lpVirtualArea;  //Virtual area node list.
	lpHeapObject->FreeBlockHeader.dwFlags     |= BLOCK_FLAGS_FREE;
	lpHeapObject->FreeBlockHeader.dwFlags     &= ~BLOCK_FLAGS_USED;
	lpHeapObject->FreeBlockHeader.dwBlockSize = 0;
	lpHeapObject->lpPrev                      = lpHeapObject; //Pointing to itself.
	lpHeapObject->lpNext                      = lpHeapObject; //Pointing to itself.

	__ENTER_CRITICAL_SECTION(NULL,dwFlags);  //Critical section here.
	lpHeapObject->lpKernelThread        = CURRENT_KERNEL_THREAD;
	lpHeapRoot = (__HEAP_OBJECT*)CURRENT_KERNEL_THREAD->lpHeapObject;
	if(NULL == lpHeapRoot)  //Has not any heap yet.
	{
		CURRENT_KERNEL_THREAD->lpHeapObject = (LPVOID)lpHeapObject;
	}
	else  //Has at least one heap object,so insert it into the list.
	{
		lpHeapObject->lpPrev = lpHeapRoot->lpPrev;
		lpHeapObject->lpNext = lpHeapRoot;
		lpHeapObject->lpNext->lpPrev = lpHeapObject;
		lpHeapObject->lpPrev->lpNext = lpHeapObject;
	}
	__LEAVE_CRITICAL_SECTION(NULL,dwFlags);

	//
	//Now,add the virtual area into the heap object's free list.
	//
	lpFreeHeader->lpPrev = &(lpHeapObject->FreeBlockHeader);
	lpFreeHeader->lpNext = &(lpHeapObject->FreeBlockHeader);
	lpHeapObject->FreeBlockHeader.lpPrev      = lpFreeHeader;
	lpHeapObject->FreeBlockHeader.lpNext      = lpFreeHeader;

	bResult = TRUE;    //The whole operation is successful.

__TERMINAL:
	if(!bResult)    //Failed.
	{
		if(lpVirtualAddr)  //Should release it.
			RELEASE_VIRTUAL_AREA(lpVirtualAddr);
		if(lpHeapObject)
			RELEASE_KERNEL_MEMORY((LPVOID)lpHeapObject);
		if(lpVirtualArea)
			RELEASE_KERNEL_MEMORY((LPVOID)lpVirtualArea);
		lpHeapObject = NULL;  //Should return a NULL flags.
	}

	return lpHeapObject;
}

//
//The implementation of DestroyHeap routine.
//This routine does the following:
// 1. Delete the heap object from kernel thread's heap list;
// 2. Release all virtual areas belong to this heap;
// 3. Release the virtual area list;
// 4. Release the heap object itself.
//
static VOID DestroyHeap(__HEAP_OBJECT* lpHeapObject)
{
	__VIRTUAL_AREA_NODE*       lpVirtualArea  = NULL;
	__VIRTUAL_AREA_NODE*       lpVirtualTmp   = NULL;
	LPVOID                     lpVirtualAddr  = NULL;
	DWORD                      dwFlags        = 0L;

	if(NULL == lpHeapObject)  //Parameter check.
		return;

	if(lpHeapObject == lpHeapObject->lpNext)  //Only one heap object in current thread.
	{
		__ENTER_CRITICAL_SECTION(NULL,dwFlags);
		lpHeapObject->lpKernelThread->lpHeapObject = NULL;
		__LEAVE_CRITICAL_SECTION(NULL,dwFlags);
	}
	else  //Delete itself from the kernel thread's heap list.
	{
		__ENTER_CRITICAL_SECTION(NULL,dwFlags);
		if(lpHeapObject->lpKernelThread->lpHeapObject == lpHeapObject)
		{
			lpHeapObject->lpKernelThread->lpHeapObject = (LPVOID)lpHeapObject->lpNext;
		}
		__LEAVE_CRITICAL_SECTION(NULL,dwFlags);

		lpHeapObject->lpPrev->lpNext = lpHeapObject->lpNext;
		lpHeapObject->lpNext->lpPrev = lpHeapObject->lpPrev;
	}

	lpVirtualArea = lpHeapObject->lpVirtualArea;
	while(lpVirtualArea)
	{
		lpVirtualTmp = lpVirtualArea;
		lpVirtualArea = lpVirtualArea->lpNext;
		RELEASE_VIRTUAL_AREA(lpVirtualTmp->lpStartAddress);  //Release the virtual area.
		RELEASE_KERNEL_MEMORY((LPVOID)lpVirtualTmp);
	}

	//
	//Now,should release the heap object itself.
	//
	RELEASE_KERNEL_MEMORY((LPVOID)lpHeapObject);

	return;
}

//
//DestroyAllHeap's implementation.
//This routine destroys all heaps of the current kernel thread.
//
static VOID DestroyAllHeap()
{
	__HEAP_OBJECT*       lpHeapObj1 = NULL;
	__HEAP_OBJECT*       lpHeapObj2 = NULL;
	__HEAP_OBJECT*       lpHeapRoot = NULL;

	lpHeapRoot = (__HEAP_OBJECT*)CURRENT_KERNEL_THREAD->lpHeapObject;
	if(NULL == lpHeapRoot)  //Not any heap object.
		return;
	lpHeapObj1 = lpHeapRoot->lpNext;
	while(lpHeapRoot != lpHeapObj1)
	{
		lpHeapObj2 = lpHeapObj1->lpNext;
		DestroyHeap(lpHeapObj1);
		lpHeapObj1 = lpHeapObj2;
	}
	DestroyHeap(lpHeapRoot);  //Destroy the root heap.
}

//
//The following is a help routine,used to print out the heap information of
//a kernel thread.
//
VOID PrintHeapInfo(__KERNEL_THREAD_OBJECT* lpKernelThread)
{
	__HEAP_OBJECT*   lpHeapObject = NULL;
	__HEAP_OBJECT*   lpHeapTmp    = NULL;
	DWORD            dwFlags      = 0L;
	__VIRTUAL_AREA_NODE*  lpVirtualArea = NULL;

	if(NULL == lpKernelThread)  //Parameter check.
		return;

	__ENTER_CRITICAL_SECTION(NULL,dwFlags);
	lpHeapObject = (__HEAP_OBJECT*)lpKernelThread->lpHeapObject;
	__LEAVE_CRITICAL_SECTION(NULL,dwFlags);
	
	//printf("Start to print out the heap information.\r\n");
	lpHeapTmp = lpHeapObject;
	while(lpHeapTmp)
	{
		//printf("Begin a heap.......\r\n");
		lpVirtualArea = lpHeapTmp->lpVirtualArea;
		while(lpVirtualArea)
		{
			//printf("Start address is : %X\r\n",lpVirtualArea->lpStartAddress);
			lpVirtualArea = lpVirtualArea->lpNext;
		}
		lpHeapTmp = lpHeapTmp->lpNext;
		if(lpHeapObject == lpHeapTmp)  //The same heap object.
			break;
	}
}

//
//The following routine is a help routine,used to dumpout the free block list.
//
VOID DumpFreeList(__HEAP_OBJECT* lpHeapObj)
{
	__FREE_BLOCK_HEADER	*        lpFreeHeader = NULL;

	if(NULL == lpHeapObj)  //Invalid parameter
		return;
	//printf("\r\nBegin to dump out the free list:\r\n");
	lpFreeHeader = lpHeapObj->FreeBlockHeader.lpNext;
	while(lpFreeHeader != &lpHeapObj->FreeBlockHeader)
	{
		//printf("StartAddr: 0x%8X Size: %d\r\n",
		//	(DWORD)lpFreeHeader + sizeof(__FREE_BLOCK_HEADER),
		//	lpFreeHeader->dwBlockSize);
		lpFreeHeader = lpFreeHeader->lpNext;
	}
}

//
//The implementation of HeapAlloc routine.
//This routine does the following actions:
// 1. Check if the heap object given by user belong to current kernel thread,
//    that is,only the thread that heap object belong to can allocate memory
//    from it;
// 2. Check the free list of the heap,try to find a free block can statisfy
//    user's request;
// 3. If can find the block,then split the block in case of the block size is
//    big,and return one to user,or return the whole block to user in case of
//    the block is not too big;
// 4. If can not find the block,then allocate a virtual area according to user's
//    request,and split the virtual area,one part returned to user,insert another
//    part into free list;
// 5. If any failure,returns NULL to indicate failed.
//
static LPVOID HeapAlloc(__HEAP_OBJECT* lpHeapObject,DWORD dwSize)
{
	__VIRTUAL_AREA_NODE*             lpVirtualArea   = NULL;
	__FREE_BLOCK_HEADER*             lpFreeBlock     = NULL;
	__FREE_BLOCK_HEADER*             lpTmpHeader     = NULL;
	LPVOID                           lpResult        = NULL;
	DWORD                            dwFlags         = 0L;
	DWORD                            dwFindSize      = 0L;

	if((NULL == lpHeapObject) || (0 == dwSize)) //Parameter check.
		return lpResult;
	
	__ENTER_CRITICAL_SECTION(NULL,dwFlags);
	if(lpHeapObject->lpKernelThread != CURRENT_KERNEL_THREAD) //Check the heap's owner.
	{
		__LEAVE_CRITICAL_SECTION(NULL,dwFlags);
		return lpResult;
	}
	__LEAVE_CRITICAL_SECTION(NULL,dwFlags);

	if(dwSize < MIN_BLOCK_SIZE)
		dwSize = MIN_BLOCK_SIZE;
	dwFindSize = dwSize + MIN_BLOCK_SIZE + sizeof(__FREE_BLOCK_HEADER);

	//
	//Now,check the free list,try to find a free block.
	//
	lpFreeBlock = lpHeapObject->FreeBlockHeader.lpNext;
	while(lpFreeBlock != &lpHeapObject->FreeBlockHeader)
	{
		if(lpFreeBlock->dwBlockSize >= dwSize)  //Find one.
		{
			if(lpFreeBlock->dwBlockSize >= dwFindSize)  //Should split it into two free blocks.
			{
				lpTmpHeader = (__FREE_BLOCK_HEADER*)((DWORD)lpFreeBlock + dwSize
					+ sizeof(__FREE_BLOCK_HEADER));     //Pointing to second part.
				lpTmpHeader->dwFlags     = BLOCK_FLAGS_FREE;
				lpTmpHeader->dwBlockSize = lpFreeBlock->dwBlockSize - dwSize
					- sizeof(__FREE_BLOCK_HEADER);      //Calculate second part's size.
				//
				//Now,should replace the lpFreeBlock with lpTmpHeader.
				//