// vmm.c
// Author: Josh Holtrop
// Date: 09/30/03
// Rewritten from scratch: 12/23/03
// Modified: 12/30/03


// This is the initialization procedure for the Virtual Memory Manager
//  It sets up the final page directory/page table setup and maps video memory, if present
void	vmm_init()
{
	dword *pageTables = (dword *)0xC0104000;	//this is the location of the page directory
	pageTables[0x3FF] = 0x104000|0x03;		//the last page directory entry points to the page directory itself
	if (videoMode)	//we are in a graphical mode
	{
		dword vidPages = video_mode.XResolution * video_mode.YResolution * (video_mode.BitsPerPixel >> 3);
		if (vidPages % 4096)
			vidPages = (vidPages >> 12) + 1;
		else
			vidPages = (vidPages >> 12);
		vmm_mapn(0xF0000000, video_mode.PhysBasePtr, vidPages);
	}
	dword firstHeapEntryBlock = (dword)mm_palloc();
	vmm_map1((dword)firstHeapEntry, firstHeapEntryBlock);
	HeapEntryBlock *heb = (HeapEntryBlock *)firstHeapEntry;
	vmm_heb_init(heb);
	heb->entry[0].base = (dword)firstHeapEntry;	//start of kernel's heap memory
	heb->entry[0].size = 4096;			//the first HeapEntryBlock is 1 page long
	heb->entry[0].attributes = VMM_HE_HEB;		//this is a HeapEntryBlock
	heb->entry[1].base = (dword)firstHeapEntry+4096;	//this is the start of the rest of the kernel's heap memory
	heb->entry[1].size = 0x10000000-4096;		//the rest of the kernel's heap memory: 256mb - 4kb
	heb->entry[1].attributes = VMM_HE_HOLE;		//this is a hold - an unmapped section of heap memory
}


// This function initialzes a Heap Entry Block to unused entries linked together
void	vmm_heb_init(HeapEntryBlock *heb)
{
	int a;
	for (a = 0; a < 256; a++)
	{
		heb->entry[a].base = 0;
		heb->entry[a].size = 0;
		heb->entry[a].attributes = VMM_HE_UNUSED;
		heb->entry[a].link = (dword)&(heb->entry[a+1]);
	}
	heb->entry[255].link = 0;
}


// This function maps a virtual address to a physical address using the page directory / page table
void	vmm_map1(dword virt, dword physical)
{
	dword pde = virt >> 22;
	dword pte = (virt & 0x003FF000) >> 12;
	dword *pageTables = (dword *)0xC0104000;	//this is the location of the page directory
	if (!(pageTables[pde] & 0x01))		//the page directory entry does not exist, we must allocate a page for it
	{
		dword *newpagetable = mm_palloc();
		pageTables[pde] = ((dword)newpagetable) | 0x03;
		invlpg(virt);
		dword *newpteptr = (dword *)(0xFFC00000 | (pde << 12));	//points to first dword of newly allocated page table
		int a;
		for (a = 0; a < 1024; a++)
		{
			newpteptr[a] = 0;
		}
	}
	dword *pteptr = (dword *)(0xFFC00000 | (pde << 12) | (pte << 2));
	*pteptr = physical | 0x03;
	invlpg(virt);
}


// This function maps a variable number of pages in a row
void	vmm_mapn(dword virt, dword physical, dword n)
{
	for (; n > 0; n--)
	{
		vmm_map1(virt, physical);
		virt += 4096;
		physical += 4096;
	}
}


// This function removes the virtual address's entry in the page directory / page table
void	vmm_unmap1(dword virt)
{
	dword *pteptr = (dword *)(0xFFC00000 | ((virt & 0xFFC00000) >> 10) | ((virt & 0x003FF000) >> 10));
	*pteptr = 0;
 	invlpg(virt);
}


// This function removes multiple pages' entries
void	vmm_unmapn(dword virt, dword n)
{
	for (; n > 0; n--)
	{
		vmm_unmap1(virt);
		virt += 4096;
	}
}


// Virtual Memory allocator function
void	*malloc(dword bytes)
{
 	if (bytes % VMM_MALLOC_GRANULARITY)
 		bytes = bytes + VMM_MALLOC_GRANULARITY - (bytes % VMM_MALLOC_GRANULARITY);
	void *attempt = vmm_getFreeChunk(bytes);
	if (attempt)
		return attempt;
	if(vmm_moreCore(bytes))
		return 0;			//we could not get any more heap memory
	return vmm_getFreeChunk(bytes);
}


// This function returns a pointer if a free chunk of memory exists
void	*vmm_getFreeChunk(dword bytes)
{
	HeapEntry *he = firstHeapEntry;
	for (;;)
	{
		if (he->attributes == VMM_HE_FREE)
		{
			if (he->size > bytes)
			{
				HeapEntry *nhe = vmm_nextHeapEntry();
				nhe->base = he->base;
				nhe->size = bytes;
				nhe->attributes = VMM_HE_USED;
				he->base += bytes;
				he->size -= bytes;
				return (void *)(nhe->base);
			}
			if (he->size == bytes)
			{
				he->attributes = VMM_HE_USED;
				return (void *)(he->base);
			}
		}
		he = (HeapEntry *)he->link;
		if (!he)
			break;
	}
	return 0;
}


// This function coalesces any two adjacent heap entries into one entry
void	vmm_coalesceHeapEntry(HeapEntry *he)
{
	if (!(he->attributes == VMM_HE_FREE || he->attributes == VMM_HE_HOLE))
		return;
	if (he->size == 0)
	{
		he->attributes = VMM_HE_UNUSED;
		return;
	}
	HeapEntry *hec = firstHeapEntry;
	for (;;)
	{
		if (hec->attributes == he->attributes)
		{
			if ((hec->base + hec->size) == he->base)	//hec ends where he begins
			{
				he->base = hec->base;
				he->size += hec->size;
				hec->attributes = VMM_HE_UNUSED;
			}
			if ((he->base + he->size) == hec->base)		//he ends where hec begins
			{
				he->size += hec->size;
				hec->attributes = VMM_HE_UNUSED;
			}
		}
		hec = (HeapEntry *)hec->link;
		if (!hec)
			break;
	}
}


// This function retrieves more physical memory for the heap
int	vmm_moreCore(dword bytes)
{
	dword pages = (bytes >> 12) + 1;
	bytes = pages << 12;
	HeapEntry *he = vmm_getFirstHoleHeapEntry(bytes);
	dword virt = he->base;
	for (; pages > 0; pages--)
	{
		dword phys = (dword)mm_palloc();
		if (!phys)
			return 1;
		vmm_map1(virt, phys);
		virt += 4096;
	}
	if (he->size == bytes)
	{
		he->attributes = VMM_HE_FREE;
		vmm_coalesceHeapEntry(he);
	}
	else
	{
		HeapEntry *nhe = vmm_nextHeapEntry();
		nhe->base = he->base;
		nhe->size = bytes;
		nhe->attributes = VMM_HE_FREE;
		he->base += bytes;
		he->size -= bytes;
		vmm_coalesceHeapEntry(nhe);
	}
	return 0;
}


// This function returns the next available heap entry, creates more entries if we are running low
HeapEntry *vmm_nextHeapEntry()
{
	if (vmm_heapEntriesLeft() < 10)
		vmm_addHeapEntryBlock();
	return vmm_getFirstUnusedHeapEntry();
}


// This function creates a new block (page) of heap entries (256)
void	vmm_addHeapEntryBlock()
{
	HeapEntry *he = vmm_getFirstHoleHeapEntry(4096);
	HeapEntry *newBlock = vmm_getFirstUnusedHeapEntry();
	dword heb = (dword)mm_palloc();
	vmm_map1(he->base, heb);
	vmm_heb_init((HeapEntryBlock *)he->base);
	HeapEntry *lhe = vmm_getLastHeapEntry();
	if (he->size == 4096)
	{
		he->attributes = VMM_HE_HEB;
		lhe->link = (dword)he->base;
	}
	else
	{
		newBlock->base = he->base;
		newBlock->size = 4096;
		newBlock->attributes = VMM_HE_HEB;
		he->base += 4096;
		he->size -= 4096;
		lhe->link = (dword)newBlock->base;
	}
	return;
}


// This function returns the last heap entry in the linked list, useful for setting
//  its link field to point to the first entry of a newly allocated list
HeapEntry *vmm_getLastHeapEntry()
{
	HeapEntry *he = firstHeapEntry;
	for (;;)
	{
		if (he->link == 0)
			return he;
		he = (HeapEntry *)he->link;
	}
}


// This function returns the first heap entry corresponding to a memory "hole"
HeapEntry *vmm_getFirstHoleHeapEntry(dword minBytes)
{
	HeapEntry *he = firstHeapEntry;
	for (;;)
	{
		if ((he->attributes == VMM_HE_HOLE) && (he->size >= minBytes))
			return he;
		he = (HeapEntry *)he->link;
		if (!he)
			break;
	}
	return 0;
}


// This function returns the first heap entry that is not being used
HeapEntry *vmm_getFirstUnusedHeapEntry()
{
	HeapEntry *he = firstHeapEntry;
	for (;;)
	{
		if (he->attributes == VMM_HE_UNUSED)
			return he;
		he = (HeapEntry *)he->link;
		if (!he)
			break;
	}
	return 0;
}


// This function returns the number of heap entries available for use
dword	vmm_heapEntriesLeft()
{
 	HeapEntry *he = firstHeapEntry;
 	dword entries = 0;
 	for (;;)
 	{
 		if (he->attributes == VMM_HE_UNUSED)
 			entries++;
 		he = (HeapEntry *)he->link;
 		if (!he)
 			break;
 	}
	return entries;
}


// This function "frees" an area of memory previously allocated with malloc()
int	free(void *ptr)
{
	HeapEntry *he = vmm_getHeapEntryByBase((dword)ptr);
	if (!he)
		return 1;	//a heap entry starting at the given address was not found
	he->attributes = VMM_HE_FREE;
	vmm_coalesceHeapEntry(he);
	return 0;
}


// This function scans the heap entry linked list for an entry that begins at the given address
HeapEntry *vmm_getHeapEntryByBase(dword base)
{
	HeapEntry *he = firstHeapEntry;
	for (;;)
	{
		if (he->base == base)
			return he;
		he = (HeapEntry *)he->link;
		if (!he)
			break;
	}
	return 0;
}