/* 
 * Copyright 2011 Steffen Vogel, Chair for Operating Systems,
 *                               RWTH Aachen University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * This file is part of MetalSVM.
 */

#include <metalsvm/vma.h>
#include <metalsvm/stdlib.h>
#include <metalsvm/stdio.h>
#include <metalsvm/tasks_types.h>
#include <metalsvm/spinlock.h>
#include <metalsvm/errno.h>

/*
 * Kernel space VMA list and lock
 * 
 * For bootstrapping we initialize the VMA list with one empty VMA
 * (start == end) and expand this VMA by calls to vma_alloc()
 */
static vma_t vma_boot = { VMA_KERN_MAX, VMA_KERN_MAX, VMA_HEAP };
static vma_t* vma_list = &vma_boot;
static spinlock_t vma_lock = SPINLOCK_INIT;

size_t vma_alloc(size_t size, uint32_t flags)
{
	task_t* task = per_core(current_task);
	spinlock_t* lock;
	vma_t** list;
	size_t ret = 0;

	kprintf("vma_alloc(0x%lx, 0x%x)\n", size, flags);

	size_t base, limit; // boundaries for search
	size_t start, end;

	if (BUILTIN_EXPECT(!size, 0))
		return 0;

	if (flags & VMA_USER) {
		base = VMA_KERN_MAX;
		limit = VMA_USER_MAX;
		list = &task->vma_list;
		lock = &task->vma_lock;
	}
	else {
		base = 0;
		limit = VMA_KERN_MAX;
		list = &vma_list;
		lock = &vma_lock;
	}

	spinlock_lock(lock);

	// "last" fit search for free memory area
	vma_t* pred = *list;  // vma before current gap
	vma_t* succ = NULL; // vma after current gap
	do {
		start = (pred) ? pred->end : base;
		end = (succ) ? succ->start : limit;

		if (end > start && end - start > size)
			break; // we found a gap

		succ = pred;
		pred = (pred) ? pred->prev : NULL;
	} while (pred || succ);

	if (BUILTIN_EXPECT(end > limit || end < start || end - start < size, 0)) {
		spinlock_unlock(lock);
		return 0;
	}

	// resize existing vma
	if (succ && succ->flags == flags) {
		succ->start -= size;
		ret = succ->start;
	}
	// insert new vma
	else {
		vma_t* new = kmalloc(sizeof(vma_t));
		if (BUILTIN_EXPECT(!new, 0))
			return 0;

		new->start = end-size;
		new->end = end;
		new->flags = flags;
		new->next = succ;
		new->prev = pred;

		if (pred)
			pred->next = new;
		if (succ)
			succ->prev = new;
		else
			*list = new;

		ret = new->start;
	}

	spinlock_unlock(lock);
	return ret;
}

int vma_free(size_t start, size_t end)
{
	task_t* task = per_core(current_task);
	spinlock_t* lock;
	vma_t* vma;
	vma_t** list;

	if (BUILTIN_EXPECT(start >= end, 0))
		return -EINVAL;

	if (end <= VMA_KERN_MAX) {
		lock = &vma_lock;
		list = &vma_list;
	}
	else if (start >= VMA_KERN_MAX) {
		lock = &task->vma_lock;
		list = &task->vma_list;
	}
	else
		return -EINVAL;

	if (BUILTIN_EXPECT(!*list, 0))
		return -EINVAL;

	spinlock_lock(lock);

	// search vma
	vma = *list;
	while (vma) {
		if (start >= vma->start && end <= vma->end) break;
		vma = vma->prev;
	}

	if (BUILTIN_EXPECT(!vma, 0)) {
		spinlock_unlock(lock);
		return -EINVAL;
	}

	// free/resize vma
	if (start == vma->start && end == vma->end) {
		if (vma == *list)
			*list = vma->next; // update list head
		if (vma->prev)
			vma->prev->next = vma->next;
		if (vma->next)
			vma->next->prev = vma->prev;
		kfree(vma);
	}
	else if (start == vma->start)
		vma->start = end;
	else if (end == vma->end)
		vma->end = start;
	else {
		vma_t* new = kmalloc(sizeof(vma_t));
		if (BUILTIN_EXPECT(!new, 0)) {
			spinlock_unlock(lock);
			return -ENOMEM;
		}

		new->start = end;
		vma->end = start;

		new->end = vma->end;
		new->next = vma->next;
		new->prev = vma;
		vma->next = new;
	}

	spinlock_unlock(lock);

	return 0;
}

int vma_add(size_t start, size_t end, uint32_t flags)
{
	task_t* task = per_core(current_task);
	spinlock_t* lock;
	vma_t** list;

	kprintf("vma_add(0x%lx, 0x%lx, 0x%x)\n", start, end, flags);

	if (BUILTIN_EXPECT(start >= end, 0))
		return -EINVAL;

	if (flags & VMA_USER) {
		list = &task->vma_list;
		lock = &task->vma_lock;

		// check if address is in userspace
		if (BUILTIN_EXPECT(start < VMA_KERN_MAX, 0))
			return -EINVAL;
	}
	else {
		list = &vma_list;
		lock = &vma_lock;

		// check if address is in kernelspace
		if (BUILTIN_EXPECT(end > VMA_KERN_MAX, 0))
			return -EINVAL;
	}

	spinlock_lock(lock);

	// search gap
	vma_t* pred = *list;
	vma_t* succ = NULL;
	while (pred) {
		if ((!pred || pred->end <= start) &&
		    (!succ || succ->start >= end))
			break;

		succ = pred;
		pred = pred->prev;
	}

	// resize existing vma
	if (pred && pred->end == start && pred->flags == flags)
		pred->end = end;
	else if (succ && succ->start == end && succ->flags == flags)
		succ->start = start;
	// insert new vma
	else {
		vma_t* new = kmalloc(sizeof(vma_t));
		if (BUILTIN_EXPECT(!new, 0))
			return 0;

		new->start = start;
		new->end = end;
		new->flags = flags;
		new->next = succ;
		new->prev = pred;

		if (pred)
			pred->next = new;
		if (succ)
			succ->prev = new;
		else
			*list = new;
	}

	spinlock_unlock(lock);

	return 0;
}

int copy_vma_list(task_t* task)
{
	task_t* parent_task = per_core(current_task);

	spinlock_init(&task->vma_lock);
	spinlock_lock(&parent_task->vma_lock);
	spinlock_lock(&task->vma_lock);

	int ret = 0;
	vma_t* last = NULL;
	vma_t* parent = parent_task->vma_list;

	while (parent) {
		vma_t *new = kmalloc(sizeof(vma_t));
		if (BUILTIN_EXPECT(!new, 0)) {
			ret = -ENOMEM;
			goto out;
		}

		new->start = parent->start;
		new->end = parent->end;
		new->flags = parent->flags;
		new->prev = last;

		if (last)
			last->next = new;
		else
			task->vma_list = new;

		last = new;
		parent = parent->next;
	}

out:
	spinlock_unlock(&task->vma_lock);
	spinlock_unlock(&parent_task->vma_lock);
	return ret;
}

int drop_vma_list()
{
	task_t* task = per_core(current_task);

	spinlock_lock(&task->vma_lock);

	while(task->vma_list)
		pfree((void*) task->vma_list->start, task->vma_list->end - task->vma_list->start);

	spinlock_unlock(&task->vma_lock);

	return 0;
}

void vma_dump()
{
	void print_vma(vma_t *vma) {
		while (vma) {
			kprintf("0x%lx - 0x%lx: size=%x, flags=%c%c%c\n", vma->start, vma->end, vma->end - vma->start,
				(vma->flags & VMA_READ) ? 'r' : '-',
				(vma->flags & VMA_WRITE) ? 'w' : '-',
				(vma->flags & VMA_EXECUTE) ? 'x' : '-');
			vma = vma->prev;
		}
	}

	task_t* task = per_core(current_task);

	kputs("Kernelspace VMAs:\n");
	spinlock_lock(&vma_lock);
	print_vma(vma_list);
	spinlock_unlock(&vma_lock);

	kputs("Userspace VMAs:\n");
	spinlock_lock(&task->vma_lock);
	print_vma(task->vma_list);
	spinlock_unlock(&task->vma_lock);
}