#include "heap.h"
#include "vm2.h"
void* operator new(u32 size)
{
	return slab.alloc(size);
}
void operator delete (void *p) {
	slab.free(p);
}

#include "terminal.h"
KernelSlabAllocator slab;
extern AddressSpace kernel_address_space;
static bool DEBUGMM=false;
#define debug if(DEBUGMM)
KernelSlabAllocator::KernelSlabAllocator() {
	AddressSpace &addr=kernel_address_space;
	heap_32= 0xc0800000;
	heap_128=0xc1000000;
	heap_256=0xc1800000;
	heap_512=0xc2000000;
	heap_len=0x00400000;
	head32=head128=head256=head512=0;
	addr.ensure_page_table(heap_32,0);
	addr.ensure_page_table(heap_128,0);
	addr.ensure_page_table(heap_256,0);
	addr.ensure_page_table(heap_512,0);
}
void* KernelSlabAllocator::alloc(u32 size) {
	if(size<32) {
		debug con->printf("Allocating %d bytes for a %d byte object.\n",32,size);
		void* p=alloc_32();
		debug con->printf("Got %X\n",p);
		return p;
	} else if(size<128) {
		debug con->printf("Allocating %d bytes for a %d byte object.\n",128,size);
		return alloc_128();
	} else if(size<256) {
		debug con->printf("Allocating %d bytes for a %d byte object.\n",256,size);
		return alloc_256();
	} else {
		con->printf("Requested size (%d bytes) not servicable\n",size);
		PANIC("NO SUPPORT FOR THIS SIZE!");
		return 0;
	}
}
void KernelSlabAllocator::free(void* ob) {
	u32 addr=(u32)ob;

	if(addr>=heap_32&&addr<heap_32+heap_len) {
		free_32(addr);
		debug con->printf("Freeing 32 byte ob\n");
	} else if(addr>=heap_128&&addr<heap_128+heap_len) {
		debug con->printf("Freeing 128 byte ob\n");
		free_128(addr);
	} else if(addr>=heap_256&&addr<heap_256+heap_len) {
		debug con->printf("Freeing 256 byte ob\n");
		free_256(addr);
	//} else if(addr>=heap_512&&addr<heap_512+heap_len) {
	} else {
		con->printf("failed to free %X\n",ob);
		PANIC("INVALID FREE (doesn't fit any of the buckets\n");
	}
}
void KernelSlabAllocator::get_more_256() {
	if(pages256>=1024)
		PANIC("can't allocate any more space for 128 byte bits...\n");
	u32 start_addr=((pages256<<12) + heap_256);
	u32 end_addr=start_addr+0x1000;
	kernel_address_space.map_anon(start_addr,0,1);
	for(u32 it=start_addr;it!=end_addr;it+=256) {
		free_256(it);
	}
	++pages256;
}
void KernelSlabAllocator::get_more_128() {
	if(pages128>=1024)
		PANIC("can't allocate any more space for 128 byte bits...\n");
	u32 start_addr=((pages128<<12) + heap_128);
	u32 end_addr=start_addr+0x1000;
	kernel_address_space.map_anon(start_addr,0,1);
	for(u32 it=start_addr;it!=end_addr;it+=128) {
		free_128(it);
	}
	++pages128;
}
void KernelSlabAllocator::get_more_32() {
	if(pages32>=1024)
		PANIC("can't allocate any more space for 128 byte bits...\n");
	u32 start_addr=((pages32<<12) + heap_32);
	debug con->printf("Allocating another page for the 32 bit heap: %X\n",start_addr);
	u32 end_addr=start_addr+0x1000;
	kernel_address_space.map_anon(start_addr,0,1);
	for(u32 it=start_addr;it!=end_addr;it+=32) {
		free_32(it);
	}
	++pages32;
}
void KernelSlabAllocator::free_32(u32 addr) {
	if(addr<heap_32||addr>=heap_128)
		PANIC("BAD FREE");
	debug con->printf("free32: addr=%X\thead=%X\n",addr,head32);
	if(head32)
	if((u32)head32<heap_32||(u32)head32>=heap_32+heap_len)
		PANIC("The heap32 head got changed underneath us...");
	u32* p=(u32*)(addr);
	*p=(u32)head32;
	head32=(void*)addr;
}
void KernelSlabAllocator::free_128(u32 addr) {
	u32* p=(u32*)(addr);
	*p=(u32)head128;
	head128=(void*)p;
}
void KernelSlabAllocator::free_256(u32 addr) {
	if(addr<heap_256||addr>=heap_512)
		PANIC("BAD FREE");
	u32* p=(u32*)(addr);
	*p=(u32)head256;
	debug con->printf("free256: addr=%X\thead=%X\n",addr,head256);

	head256=(void*)p;
}
void* KernelSlabAllocator::alloc_256() {
	if(!head256)
		get_more_256();
	u32 *h=(u32*)head256;
	void*ret=head256;
	head256=(void*)*h;
	return ret;
}
void* KernelSlabAllocator::alloc_128() {
	if(!head128)
		get_more_128();
	u32 *h=(u32*)head128;
	void*ret=head128;
	head128=(void*)*h;
	return ret;
}
void* KernelSlabAllocator::alloc_32() {
	if(!head32) {
		debug con->printf("Had to make another page for heap32\n");
		get_more_32();
	}
	u32 *h=(u32*)head32;
	debug con->printf("head32=%X\n",head32);
	if((u32)head32<heap_32||(u32)head32>=heap_32+heap_len)
		PANIC("Head32 is not legal.");
	void*ret=head32;
	head32=(void*)*h;
	return ret;
}