/*

* Copyright (c) 2007, John Cronin
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of the copyright holder nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY John Cronin ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL John Cronin BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*/

#include "pmem.h"
#include "multiboot.h"
#include "sys/types.h"
#include "panic.h"

/* The physical memory allocater (V3)

 The allocator uses a system of buddies, which are basically bitmaps representing free
  pages of varying sizes
 e.g. there is a buddy for 4kB pages, 8kB pages, 16kB pages ... 2MB pages, 4MB pages
  (with default setup)
 Initially, only the largest size pages are available, an attempt to allocate a small
  page will not find a match, so the allocator will ask the next buddy up and so on
  until it finds a free page.  This page will then be split by each buddy on the way
  back down the buddy size tree, marking one as free and the other as used, returning
  the used one, until it reaches the buddy at which the original request was made.
 Freeing a page works in reverse
 There is also a function (pmem_allocp) which attempts to allocate a particular address
  This is useful in setting, e.g. memory mapped areas as used
 The special symbol ~0 (bitwise NOT of zero) is used to represent an invalid address
 This particular allocator used 2 bits (a bit-pair) to represent each page.
 This allows 4 states - free, used, used but not swappable and memory mapped device
 */

static buddy_header *bh;

static physaddr_t buddy_alloc(uintptr buddy_no, uint8 bpval, uint8 buddy_loc);
static void buddy_free(uintptr buddy_no, physaddr_t addr);
static physaddr_t buddy_allocp(uintptr buddy_no, physaddr_t addr, uint8 bpval);
static inline uint8 buddy_checkbitpair(uintptr base, uintptr bitpair);
static inline void buddy_setbitpair(uintptr base, uintptr bitpair, uint8 bpval);
static inline uintptr buddy_calcbitpair(uintptr pagesize, physaddr_t p);
static inline uintptr buddy_complementbitpair(uintptr pagesize, physaddr_t p);
static uintptr count_buddies(uintptr buddy_no, uint8 bpval);
static bool check_addr(physaddr_t, uintptr size);

uintptr allocate_buddy(physaddr_t addr, uintptr smallest_page, uintptr largest_page)
{
	/* if addr == ~0, then just calculate the size of the buddies and return
	  total size = buddy_header + size of buddy_def array + size of each bitmap
	 if addr != ~0 then allocate the buddy_header, buddy_def array and zero the buddies */

	physaddr_t base_addr;
	physaddr_t cur_addr;
	buddy_def *bd;
	uintptr cur_ps;
	uintptr i;
	uintptr total_mem;
	uintptr total_size;
	uintptr cur_page;

	if(addr == (physaddr_t)~0)
		base_addr = 0;
	else
		base_addr = addr;
	cur_addr = base_addr;

	if(addr != (physaddr_t)~0)
		bh = (buddy_header *)cur_addr;
	cur_addr += sizeof(buddy_header);

	bd = (buddy_def *)cur_addr;

	for(cur_ps = smallest_page, i = 0; cur_ps <= largest_page; cur_ps *= 2, i++) {
		if(addr != (physaddr_t)~0)
			bd[i].b_ps = cur_ps;
		cur_addr += sizeof(buddy_def);
	}
	if(addr != (physaddr_t)~0) {
		bh->bh_count = i;
		bh->bh_smallest = smallest_page;
		bh->bh_largest = largest_page;
		bh->bh_bd = bd;
	}
	
	/* determine the amount of free memory */
	total_mem = 0;
	if(mbheader->flags & MULTIBOOT_MMAP_AVAIL) {
		for(i = 0; i < mbmmapcount; i++) {
			if(mbmem[i].type == 1) {
				if((mbmem[i].base_addr_union.base_addr + mbmem[i].length_union.length) 
					> total_mem)
					total_mem = mbmem[i].base_addr_union.base_addr + 
					mbmem[i].length_union.length;
			}
		}
	} else if(mbheader->flags & MULTIBOOT_MEM_AVAIL) {
		total_mem = ((uintptr)(mbheader->mem_upper + 1)) << 10;
	} else {
		/* probe memory */
		kpanic("allocate_buddy: probing not yet implemented");
	}

#ifdef DEBUG_PMEM
	printf("allocate_buddy: detected %x bytes of memory\n", total_mem);
#endif

	ALIGN_UP(cur_addr);

	for(cur_ps = smallest_page, i = 0; cur_ps <= largest_page; cur_ps *= 2, i++) {
		uintptr bitmap_size;
		uintptr num_pages;
		
		num_pages = total_mem / cur_ps;
		if(total_mem % cur_ps)
			num_pages ++;

		/* we allocate two bits per page */
		bitmap_size = num_pages >> 2;
		
		if(addr != (physaddr_t)~0) {
			bd[i].b_ps = cur_ps;
			bd[i].b_start = cur_addr;
			bd[i].b_end = cur_addr + bitmap_size;
			zeromem((uint8 *)cur_addr, bitmap_size);
		}

		cur_addr += bitmap_size;
		ALIGN_UP(cur_addr);
	}

	total_size = (uintptr)cur_addr - (uintptr)base_addr;

	/* Set the .buddy section of the elf file to point here, if we are allocating space
	 This ensures we mark the space as allocated */
	if(addr != (physaddr_t)~0) {
		bh->bh_size = total_size;

		for(i = 0; i < mbheader->elf_num; i++) {
			if(!strcmp(&mbstr[mbelf[i].sh_name], ".buddy")) {
				mbelf[i].sh_addr = (Elf32_Addr)base_addr;
				mbelf[i].sh_size = (Elf32_Word)total_size;
				break;
			}
		}

#ifdef DEBUG_PMEM
		printf("allocate_buddy: setting up buddies with default values\n");
#endif

#ifdef DEBUG_PMEM
		printf("allocate_buddy: freeing all total memory\n");
#endif
		/* Free total memory */
		for(i = 0; i < total_mem; i += 0x100000)
			pmem_free(0x100000, (physaddr_t)i);

		/* Mark 1st 1MB as used */
		pmem_allocp(0x100000, 0, PMEM_PAGESTATUS_USEDNOSWAP);

#ifdef DEBUG_PMEM
		printf("allocate_buddy: allocating memory for mmap'ed devices\n");
#endif

		/* Allocate all pages belonging to a memory mapped device */
		for(i = 0; i < mbmmapcount; i++) {
			for(cur_page = (mbmem[i].base_addr_union.base_addr & (uintptr)(~0xfff)); 
				cur_page <
				(mbmem[i].base_addr_union.base_addr + mbmem[i].length_union.length + 4096); 
				cur_page += 4096) {
					if(cur_page >= total_mem)
						break;
					if(mbmem[i].type != 1)
						pmem_allocp(4096, cur_page, PMEM_PAGESTATUS_USEDNOSWAP);
			}
		}

#ifdef DEBUG_PMEM
		printf("allocate_buddy: allocating memory for kernel code and data sections\n");
#endif

		/* Allocate all pages belonging to an elf section */
		for(i = 0; i < mbheader->elf_num; i++) {
			if(mbelf[i].sh_size > 0) {
				for(cur_page = (mbelf[i].sh_addr & (uintptr)(~0xfff)); cur_page < 
					(mbelf[i].sh_addr + mbelf[i].sh_size + 4096); cur_page += 4096)
					pmem_allocp(4096, cur_page, PMEM_PAGESTATUS_USEDNOSWAP);
			}
		}

#ifdef DEBUG_PMEM
		printf("allocate_buddy: allocating memory for currently loaded modules\n");
#endif

		/* Allocate all pages belonging to a module */
		for(i = 0; i < mbheader->mods_count; i++) {
			for(cur_page = (mbmods[i].mod_start & (uintptr)(~0xfff)); cur_page < 
				(mbmods[i].mod_end + 4096); cur_page += 4096)
				pmem_allocp(4096, cur_page, PMEM_PAGESTATUS_USED);
		}

#ifdef DEBUG_PMEM
		pmem_dump_buddies();
#endif
	}

	return total_size;
}

bool check_addr(physaddr_t addr, uintptr size)
{
	physaddr_t end;
	uintptr cur_sect;

	if(addr < 0x100000)
		return FALSE;

	end = addr + (physaddr_t)size;

	/* check if a given address is not included in any elf section */
	for(cur_sect = 0; cur_sect < (uintptr)mbheader->elf_num; cur_sect++) {
		physaddr_t s_addr = (physaddr_t)mbelf[cur_sect].sh_addr;
		uintptr s_size = (uintptr)mbelf[cur_sect].sh_size;
		physaddr_t s_end = s_addr + (physaddr_t)s_size;

		/* only pass a section when either:
		  both start and end are <= start
		  both start and end are >= end */

		if(((s_addr <= addr) && (s_end <= addr)) || ((s_addr >= end) && (s_end >= end))) {
			/* check is given address is not included in any module section */
			uintptr cur_mod;
			for(cur_mod = 0; cur_mod < (uintptr)mbheader->mods_count; cur_mod++) {
				physaddr_t m_addr = (physaddr_t)mbmods[cur_mod].mod_start;
				physaddr_t m_end = (physaddr_t)mbmods[cur_mod].mod_end;

				if(((m_addr <= addr) && (m_end <= addr)) || ((m_addr >= end) && (m_end >= end)))
					continue;
				return FALSE;
			}
			continue;
		}
		return FALSE;
	}
	return TRUE;
}

physaddr_t find_free_address(uintptr size)
{
	/* utility function to find a free memory region of 'size' bytes
	 this is only necessary before the pmem allocator is running
	 it searches the multiboot-provided memory map, elf header sections and
	 module definitions

	 it will only return areas above 1MB as there seem to be some issues with allocating
	  space below 1MB
	 in particular, random bits of memory at the start of the first megabyte get
	  overwritten.  I assume this is due to the IVT or bios data area, but GRUB reports
	  this memory as available.  Hmmm.

	 first, find a free memory region */
	uintptr cur_mmap;

	for(cur_mmap = 0; cur_mmap < mbmmapcount; cur_mmap++) {
		/* type 1 is available RAM, anything else isn't
		 discard if the region is smaller that requested */
		if((mbmem[cur_mmap].type == 1) && ((uintptr)mbmem[cur_mmap].length_union.length 
		>= size)) {
			/* If a free region is to be found anywhere, it is either:
			  a) at the start of the block
			  b) after any section which ends within the block, given that there
			      is sufficient room after the section within the block
			  c) if the block starts below 1 MB, and 0x100000 is included within the
			      block, then, assuming base_addr + length - 0x100000 > size, try this */

			uintptr cur_elf;

			if(check_addr((physaddr_t)mbmem[cur_mmap].base_addr_union.base_addr, size))
				return (physaddr_t)mbmem[cur_mmap].base_addr_union.base_addr;

			if((mbmem[cur_mmap].base_addr_union.base_addr < 0x100000) &&
				((mbmem[cur_mmap].base_addr_union.base_addr + 
				mbmem[cur_mmap].length_union.length) > (size + 0x100000))) {
					if(check_addr(0x100000, size))
						return 0x100000;
			}

			for(cur_elf = 0; cur_elf < (uintptr)mbheader->elf_num; cur_elf++) {
				physaddr_t s_end = (physaddr_t)mbelf[cur_elf].sh_addr +
					(physaddr_t)mbelf[cur_elf].sh_size;
				physaddr_t mmap_end = (physaddr_t)mbmem[cur_mmap].base_addr_union.base_addr +
					(physaddr_t)mbmem[cur_mmap].length_union.length;

				/* does the section end within the block? */
				if((s_end >= mbmem[cur_mmap].base_addr_union.base_addr) && 
					(s_end < mmap_end)) {
						/* is there sufficient free space after the end of the block? */

						if((physaddr_t)size <= (mmap_end - s_end)) {
							/* now see if this will overlap any other section */
							if(check_addr(s_end, size))
								return s_end;
						}
				}
			}
		}
	}

	return (physaddr_t)~0;
}

physaddr_t pmem_alloc(uintptr page_size, uint8 bpval, uint8 buddy_loc)
{
	uintptr i;
	for(i = 0; i < bh->bh_count; i++) {
		if(page_size == bh->bh_bd[i].b_ps)
			return buddy_alloc(i, bpval, buddy_loc);
	}

	return (physaddr_t)~0;
}

physaddr_t pmem_allocp(uintptr page_size, physaddr_t address, uint8 bpval)
{
	uintptr i;
	for(i = 0; i < bh->bh_count; i++) {
		if(page_size == bh->bh_bd[i].b_ps)
			return buddy_allocp(i, address, bpval);
	}

	return (physaddr_t)~0;
}

void pmem_free(uintptr page_size, physaddr_t address)
{
	uintptr i;
	for(i = 0; i < bh->bh_count; i++) {
		if(page_size == bh->bh_bd[i].b_ps)
			buddy_free(i, address);
	}
}

#define CHECK_BITPAIR(a) \
	if(((*b >> ((a) << 1)) & 0x3) == PMEM_PAGESTATUS_FREE) { \
		buddy_setbitpair(cur_byte, (a), bpval); \
		return (((cur_byte - bh->bh_bd[buddy_no].b_start) << 2) + (a)) * \
			bh->bh_bd[buddy_no].b_ps; \
	}

physaddr_t buddy_alloc(uintptr buddy_no, uint8 bpval, uint8 buddy_loc)
{
	/* allocate the first free page from this buddy, otherwise ask a higher buddy */
	uintptr cur_byte;
	uintptr start_byte, end_byte;
	physaddr_t next_buddy_ret;

	if(buddy_no >= bh->bh_count)
		return (physaddr_t)~0;

	/* decide upon starting and ending bytes */

	if(buddy_loc == PMEM_BUDDYLOCATION_HIGH) {
		start_byte = ((0x1000000 / bh->bh_bd[buddy_no].b_ps) >> 2) + 
			bh->bh_bd[buddy_no].b_start;
		end_byte = bh->bh_bd[buddy_no].b_end;
	} else if(buddy_loc == PMEM_BUDDYLOCATION_LOW) {
		start_byte = bh->bh_bd[buddy_no].b_start;
		end_byte = ((0x1000000 / bh->bh_bd[buddy_no].b_ps) >> 2) + 
			bh->bh_bd[buddy_no].b_start;
	} else if(buddy_loc == PMEM_BUDDYLOCATION_ALL) {
		start_byte = bh->bh_bd[buddy_no].b_start;
		end_byte = bh->bh_bd[buddy_no].b_end;
	} else
		return (physaddr_t)~0;

	/* search for a free bit-pair
	  NB that 0 = free, 1,2,3 are not free */
	for(cur_byte = start_byte; cur_byte < end_byte; cur_byte++) {
		uint8 *b = (uint8 *)cur_byte;

		/* check each bit pair */
		CHECK_BITPAIR(0);
		CHECK_BITPAIR(1);
		CHECK_BITPAIR(2);
		CHECK_BITPAIR(3);
	}

	/* Nothing in this buddy, search the next */
	next_buddy_ret = buddy_alloc(buddy_no + 1, bpval, buddy_loc);
	if(next_buddy_ret != (physaddr_t)~0) {
		buddy_setbitpair(bh->bh_bd[buddy_no].b_start, 
			buddy_complementbitpair(bh->bh_bd[buddy_no].b_ps, next_buddy_ret),
			PMEM_PAGESTATUS_FREE);
		buddy_setbitpair(bh->bh_bd[buddy_no].b_start,
			buddy_calcbitpair(bh->bh_bd[buddy_no].b_ps, next_buddy_ret), bpval);
	}
	return next_buddy_ret;
}

physaddr_t buddy_allocp(uintptr buddy_no, physaddr_t addr, uint8 bpval)
{
	/* allocate a specific page from this buddy, or split a higher buddy if available */
	physaddr_t q;
	if(buddy_no >= bh->bh_count)
		return (physaddr_t)~0;

	/* see if the page is free in this buddy */
	if((buddy_checkbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
		bh->bh_bd[buddy_no].b_ps, addr)) & 0x3) == PMEM_PAGESTATUS_FREE) {
			buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
				bh->bh_bd[buddy_no].b_ps, addr), bpval);
			return addr;
	}

	/* else see if its free in a higher buddy */
	q = buddy_allocp(buddy_no + 1, addr, bpval);
	if(q != (physaddr_t)~0) {
		buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_complementbitpair(
			bh->bh_bd[buddy_no].b_ps, addr), PMEM_PAGESTATUS_FREE);
		buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
			bh->bh_bd[buddy_no].b_ps, addr), bpval);
	}
	return q;
}

void buddy_free(uintptr buddy_no, physaddr_t addr)
{
	/* Free a particular page from a buddy, merge with complement if possible and free in
	  a higher buddy */
	if(buddy_no == (bh->bh_count - 1)) {
		/* highest buddy - just free it (TM) */
		buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
			bh->bh_bd[buddy_no].b_ps, addr), PMEM_PAGESTATUS_FREE);
	} else {
		/* see if the complement page is free */
		if((buddy_checkbitpair(bh->bh_bd[buddy_no].b_start, buddy_complementbitpair(
			bh->bh_bd[buddy_no].b_ps, addr)) & 0x3) == PMEM_PAGESTATUS_FREE) {
				/* as the complement is free, we should set both it and this page to be
				  used, and then free the page in a higher buddy. */
				buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
					bh->bh_bd[buddy_no].b_ps, addr), PMEM_PAGESTATUS_USEDNOSWAP);
				buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_complementbitpair(
					bh->bh_bd[buddy_no].b_ps, addr), PMEM_PAGESTATUS_USEDNOSWAP);

				buddy_free(buddy_no + 1, addr);
		} else {
			/* just free it */
			buddy_setbitpair(bh->bh_bd[buddy_no].b_start, buddy_calcbitpair(
				bh->bh_bd[buddy_no].b_ps, addr), PMEM_PAGESTATUS_FREE);
		}
	}
}

inline uint8 buddy_checkbitpair(uintptr base, uintptr bitpair)
{
	/* check whether a particular bitpair is free
	 bitpair = bitpair number (0, 1, 2 or 3)
	 base = base of buddy */

	uintptr byte;
	uint8 *b;

	uintptr actbit = bitpair % 0x4;
	actbit = actbit << 1;
	byte = bitpair >> 2;
	b = (uint8 *)(base + byte);
	return ((*b >> actbit) & 0xff);
}

inline void buddy_setbitpair(uintptr base, uintptr bitpair, uint8 bpval)
{
	/* set the value of a particular bitpair
	 bitpair = bitpair number (0, 1, 2 or 3)
	 base = base of buddy
	 bpval = value to set to

	 first we mask out the current value of the bitpair (but leave the rest of the byte
	  intact), then bitwise OR it with the new value */

	uintptr byte;
	uint8 mask, val, *b;

	uintptr actbit = bitpair % 0x4;
	actbit = actbit << 1;
	byte = bitpair >> 2;
	b = (uint8 *)(base + byte);
	mask = ~(0x3 << actbit);
	val = (bpval & 0x3) << actbit;
	*b &= mask;
	*b |= val;
}

inline uintptr buddy_calcbitpair(uintptr pagesize, physaddr_t p)
{
	/* calculates the particular bitpair which corresponds to this address
	 it is simply p / pagesize */
	return p / pagesize;
}

inline uintptr buddy_complementbitpair(uintptr pagesize, physaddr_t p)
{
	/* calculate the complement of the bitpair to which this page refers */
	uintptr b = buddy_calcbitpair(pagesize, p);
	if(b & 0x1)
		return b - 1;
	else
		return b + 1;
}

uintptr count_buddies(uintptr buddy_no, uint8 bpval)
{
	uintptr cur_byte;
	uintptr total_count = 0;

	for(cur_byte = bh->bh_bd[buddy_no].b_start; cur_byte < bh->bh_bd[buddy_no].b_end;
		cur_byte++) {
			uint8 *b = (uint8 *)cur_byte;

			if(((*b >> 0) & 0x3) == bpval)
				total_count++;
			if(((*b >> 2) & 0x3) == bpval)
				total_count++;
			if(((*b >> 4) & 0x3) == bpval)
				total_count++;
			if(((*b >> 6) & 0x3) == bpval)
				total_count++;
	}

	return total_count;
}

void pmem_dump_buddies()
{
	uintptr i;
	uintptr freemem = 0;

	printf("Dump of buddies:\n");
	printf("Page size       Base            Free            Used            Used, non-swap\n");

	for(i = 0; i < bh->bh_count; i++) {
		printf("%x      %x      %x      %x      %x\n",
			bh->bh_bd[i].b_ps,
			bh->bh_bd[i].b_start,
			count_buddies(i, PMEM_PAGESTATUS_FREE),
			count_buddies(i, PMEM_PAGESTATUS_USED),
			count_buddies(i, PMEM_PAGESTATUS_USEDNOSWAP));
		freemem += (count_buddies(i, PMEM_PAGESTATUS_FREE) * bh->bh_bd[i].b_ps);
	}
	printf("Total free: %x bytes\n", freemem);
}