#include "heap.basic.h"

#include "paging.h"

#include <lib/mutex.h>
#include <core/sys.h>

void heap_create(struct heap *heap, size_t start, size_t datasize, size_t maxdatasize) {
	uint32_t i;
	if (start & 0x0FFF) start = (start & 0xFFFFF000) + 0x1000;
	
	// Allocate frames
	for (i = start; i < start + datasize; i += 0x1000) {
		page_map(pagedir_getPage(kernel_pagedir, i, 1), frame_alloc(), 0, 0);
	}

	heap->start_addr = start;
	heap->end_addr = start + datasize;
	heap->max_end = start + maxdatasize;

	struct heap_header *first = (struct heap_header*)start;
	first->magic = HH_FREE_MAGIC;
	first->prev = first->next = first;

	heap->first = first;
	heap->first_free = first;
	
	heap->mutex = 0;
}

// ***** UTILITY FUNCTIONS *****

#define BLK_SIZE(blk) (blk->next == blk ? ((size_t)heap->end_addr - (size_t)blk) : ((size_t)blk->next - (size_t)blk))
#define BLK_FREE(blk) (blk->magic == HH_FREE_MAGIC)

void blk_split(struct heap *heap, struct heap_header *hdr, size_t size) {
	ASSERT (size >= MIN_BLK_SIZE);
	ASSERT (BLK_SIZE(hdr) > size);
	ASSERT (BLK_SIZE(hdr) - size >= MIN_BLK_SIZE);
	ASSERT (BLK_FREE(hdr));

	struct heap_header *hdr_nxt = (struct heap_header*)((size_t)hdr + size);
	hdr_nxt->magic = HH_FREE_MAGIC;
	hdr_nxt->prev = hdr;
	if (hdr->next == hdr) {
		hdr_nxt->next = hdr_nxt;
	} else {
		hdr_nxt->next = hdr->next;
		hdr->next->prev = hdr_nxt;
	}
	hdr->next = hdr_nxt;
}

void blk_combine_next(struct heap *heap, struct heap_header *hdr) {
	ASSERT (hdr != hdr->next);	// make sure we are not the last block
	ASSERT (hdr->magic == HH_FREE_MAGIC);
	ASSERT (hdr->next->magic == HH_FREE_MAGIC);

	struct heap_header *hdr_nxt = hdr->next;
	if (hdr_nxt->next == hdr_nxt) {
		hdr->next = hdr;
	} else {
		hdr->next = hdr_nxt->next;
		hdr_nxt->next->prev = hdr;
	}
}

// ***** Allocation/Freeing functions *****

void *heap_alloc(struct heap *heap, size_t sz) {
	/* Algorithm :
	1 Use heap->first_free to find first effectively free block, update heap->first_free
	2 Find first free AND BIG ENOUGH block
		- If no block can be found, return 0 (no heap expand yet)
	3 If that block is too big, split it in two
	4 Mark that block as used and return it's address + sizeof(heap_header) */
	mutex_lock(&heap->mutex);

	sz += sizeof(struct heap_header);
	if (sz < MIN_BLK_SIZE) sz = MIN_BLK_SIZE;
	if (sz % 4 != 0) sz = sz + 4 - (sz % 4); 

	// 1
	struct heap_header *it = heap->first_free;
	while (!BLK_FREE(it) && it != it->next) it = it->next;
	if (!BLK_FREE(it) && it == it->next) {
		mutex_unlock(&heap->mutex);
		return 0;
	}
	heap->first_free = it;

	// 2
	while (BLK_SIZE(it) < sz && it != it->next) it = it->next;
	if (BLK_SIZE(it) < sz) {		// no block could be found
		mutex_unlock(&heap->mutex);
		return 0;
	}

	// 3
	if (BLK_SIZE(it) >= sz + MIN_BLK_SIZE) blk_split(heap, it, sz);

	// 4
	void *ret = (void*)((size_t)it + sizeof(struct heap_header));
	it->magic = HH_ALLOC_MAGIC;

	mutex_unlock(&heap->mutex);
	return ret;
}

void heap_free(struct heap *heap, void *ptr) {
	/* Algorithm :
	1 Get block pointer to the block we are freeing
	2 Mark that block as free
	3 If block->next is free, combine_next(block)
	4 If block->prev is free, block = block->prev; combine_next(block)
	5 If block < heap->first_free; heap->first_free = block
	*/
	mutex_lock(&heap->mutex);

	// 1
	struct heap_header *hdr = (struct heap_header*)((size_t)ptr - sizeof(struct heap_header));
	if (hdr->magic != HH_ALLOC_MAGIC) return;
	// 2
	hdr->magic = HH_FREE_MAGIC;
	// 3
	if (hdr->next != hdr && hdr->next->magic == HH_FREE_MAGIC) blk_combine_next(heap, hdr);
	// 4
	if (hdr->prev != hdr && hdr->prev->magic == HH_FREE_MAGIC) {
		hdr = hdr->prev;
		blk_combine_next(heap, hdr);
	}
	// 5
	if (hdr < heap->first_free) heap->first_free = hdr;

	mutex_unlock(&heap->mutex);
}