/* 
 * Copyright 2014 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.
 */

/**
 * @author Steffen Vogel <steffen.vogel@rwth-aachen.de>
 */

#include <metalsvm/stdlib.h>
#include <metalsvm/stdio.h>
#include <metalsvm/stdarg.h>
#include <metalsvm/memory.h>
#include <metalsvm/time.h>
#include <metalsvm/tasks.h>
#include <metalsvm/vma.h>
#include <metalsvm/malloc.h>

#include <asm/page.h>
#include <asm/irqflags.h>
#include <asm/processor.h>
#include <asm/pmc.h>

#define ITERATIONS	1000
#define PAGE_COUNT	40
#define SIZE		(PAGE_COUNT*PAGE_SIZE)
#define VIRT_FROM_ADDR	0x50000000 // Userspace
#define VIRT_TO_ADDR	0x30000000 // Kernelspace

extern atomic_int32_t total_page;
extern atomic_int32_t total_allocated_pages;
extern atomic_int32_t total_available_pages;

/** @brief Simple helper to format our test results */
static void test(size_t expr, char *fmt, ...)
{
	void _putchar(int c, void *arg) { kputchar(c); } // for kvprintf

	static int c = 1;

	va_list ap;
	va_start(ap, fmt);

	kprintf("%s #%u:\t", (expr) ? "PASSED" : "FAILED", c++);
	kvprintf(fmt, _putchar, NULL, 10, ap);
	kputs("\n");

	va_end(ap);

	if (!expr)
		abort();
}

/** @brief Linear feedback shift register PRNG */
static uint16_t rand()
{
	static uint16_t lfsr = 0xACE1u;
	static uint16_t bit;

	bit  = ((lfsr >> 0) ^ (lfsr >> 2) ^ (lfsr >> 3) ^ (lfsr >> 5) ) & 1;
	return lfsr = (lfsr >> 1) | (bit << 15);
}

/** @brief BSD sum algorithm ('sum' Unix command) and used by QEmu */
uint16_t checksum(size_t start, size_t end)
{
	size_t addr;
	uint16_t sum;

	for(addr = start, sum = 0; addr < end; addr++) {
		uint8_t val = *((uint8_t *) addr);
		sum = (sum >> 1) | (sum << 15);
		sum += val;
	}

	return sum;
}

static int paging_stage2(void *arg);

/** @brief Test of the paging subsystem
 *
 * We will map a single physical memory region to two virtual regions.
 * When writing to the first one, we should be able to read the same contents
 * from the second one.
 */
static void paging(void)
{
	size_t c, sum;
	size_t *p1, *p2;
	size_t virt_from, virt_to;
	size_t phys;
	size_t t;
	int ret;
	int flags;

	// disable irqs to prevent context switches for rdtsc measurement
	flags = irq_nested_disable();

	// show original page maps
	t = rdtsc();
	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);
	kprintf("delta_t = %lu\n", rdtsc() - t);

	t = rdtsc();
	page_stats(1); // reset accessed and dirty bits
	kprintf("delta_t = %lu\n", rdtsc() - t);

	irq_nested_enable(flags);

	kprintf("bookkeeping pages:\n");
	kprintf("  - total:\t%lu\n", atomic_int32_read(&total_pages));
	kprintf("  - alloc:\t%lu\n", atomic_int32_read(&total_allocated_pages));
	kprintf("  - avail:\t%lu\n", atomic_int32_read(&total_available_pages));

	// allocate physical page frames
	phys = get_pages(PAGE_COUNT);
	test(phys, "get_pages(%lu) = %#lx", PAGE_COUNT, phys);

	// create first mapping
	virt_from = map_region(VIRT_FROM_ADDR, phys, PAGE_COUNT, MAP_USER_SPACE);
	test(virt_from, "map_region(%#lx, %#lx, %lu, %#x) = %#lx", VIRT_FROM_ADDR, phys, PAGE_COUNT, MAP_USER_SPACE, virt_from);

	// check address translation
	phys = virt_to_phys(virt_from);
	test(phys, "virt_to_phys(%#lx) = %#lx", virt_from, phys);

	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	// test set_page_flags()
	ret = set_page_flags(virt_from, PAGE_COUNT, MAP_CODE);
	test(!ret, "set_page_flags(%#lx, %u, %x)", virt_from, PAGE_COUNT, MAP_USER_SPACE|MAP_CODE); // now executable

	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	// write test data
	p1 = (size_t *) virt_from;
	for (c = 0; c < SIZE/sizeof(size_t); c++) {
		p1[c] = c;
	}

	// create second mapping pointing to the same page frames
	virt_to = map_region(VIRT_TO_ADDR, phys, PAGE_COUNT, 0);
	test(virt_to, "map_region(%#lx, %#lx, %lu, %#x) = %#lx", VIRT_TO_ADDR, phys, PAGE_COUNT, 0, virt_to);

	// check address translation
	phys = virt_to_phys(virt_to);
	test(phys, "virt_to_phys(%#lx) = %#lx", virt_to, phys);

	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	// check if both mapped areas are equal
	p2 = (size_t *) virt_to;
	for (c = 0; c < SIZE/sizeof(size_t); c++) {
		if (p1[c] != p2[c])
			test(0, "data mismatch: *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is equal");

	// try to remap without MAP_REMAP
	virt_to = map_region(VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, 0);
	test(!virt_to, "map_region(%#lx, %#lx, %lu, %#x) = %#lx (without MAP_REMAP flag)", VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, 0, virt_to);

	// try to remap with MAP_REMAP
	virt_to = map_region(VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, MAP_REMAP);
	test(virt_to, "map_region(%#lx, %#lx, %lu, %#x) = %#lx (with MAP_REMAP flag)", VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, MAP_REMAP, virt_to);

	// check if data is not equal anymore (we remapped with +PAGE_SIZE offset)
	p1 = (size_t *) (virt_from + PAGE_SIZE);
	for (c = 0; c < (SIZE-PAGE_SIZE)/sizeof(size_t); c++) {
		if (p1[c] != p2[c])
			test(0, "data mismatch at *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is equal");

	// try to remap with MAP_REMAP
	virt_to = map_region(VIRT_TO_ADDR, phys, PAGE_COUNT, MAP_REMAP);
	test(virt_to, "map_region(%#lx, %#lx, %lu, %#x) = %#lx (with MAP_REMAP flag)", VIRT_TO_ADDR, phys, PAGE_COUNT, MAP_REMAP, virt_to);

	// test unmapping
	ret = unmap_region(VIRT_FROM_ADDR, PAGE_COUNT);
	test(!ret, "unmap_region(%#lx, %lu) = %u", VIRT_FROM_ADDR, PAGE_COUNT, ret);

	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	// calc checksum
	sum = checksum(virt_to, virt_to + SIZE);
	test(sum == 23196, "checksum(%p, %p) = %lu", virt_to, virt_to + SIZE, sum);

	size_t cr3 = read_cr3();
	kprintf("cr3 old = %#lx\n", cr3);

	create_kernel_task(0, paging_stage2, &sum, NORMAL_PRIO);
	wait(&ret);
	test(!ret, "paging stage 2 returned with code = %i", ret);
}

static int paging_stage2(void *arg)
{
	size_t old, new;

	kprintf("PAGING: entering stage 2...\n");

	size_t cr3 = read_cr3();
	kprintf("cr3 new = %#lx\n", cr3);

	old = *((size_t *) arg);
	kprintf("old sum: %lu\n", old);

	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	new = checksum(VIRT_TO_ADDR, VIRT_TO_ADDR + SIZE);
	test(old == new, "checksum(%p, %p) = %lu", VIRT_TO_ADDR, VIRT_TO_ADDR + SIZE, new);

	return 0;
}

/** @brief Test of the VMA allocator */
static void vma(void)
{
	int ret;
	vma_dump();

	// vma_alloc
	size_t a1 = vma_alloc(SIZE, VMA_HEAP);
	test(a1, "vma_alloc(%#x, %#x) = %#lx", SIZE, VMA_HEAP, a1);

	size_t a2 = vma_alloc(SIZE, VMA_HEAP|VMA_USER);
	test(a2 != 0, "vma_alloc(%#x, %#x) = %#lx", SIZE, VMA_HEAP|VMA_USER, a2);
	vma_dump();

	// vma_free
	ret = vma_free(a1, a1+SIZE);
	test(ret >= 0, "vma_free(%#lx, %#lx) = %i", a1, a1+SIZE, ret);

	ret = vma_free(a2, a2+SIZE);
	test(ret >= 0, "vma_free(%#lx, %#lx) = %i", a2, a2+SIZE, ret);
	vma_dump();

	// vma_add
	ret = vma_add(VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(%#lx, %#lx, %#x) = %u", VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, VMA_HEAP|VMA_USER, ret);

	ret = vma_add(VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(%#lx, %#lx, %#x) = %u", VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, VMA_HEAP|VMA_USER, ret);

	ret = vma_add(VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(%#lx, %#lx, %#x) = %u", VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, VMA_HEAP|VMA_USER, ret);
	vma_dump();

	// vma_free
	ret = vma_free(VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR);
	test(ret >= 0, "vma_free(%#lx, %#lx) = %u", VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, ret);

	ret = vma_free(VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE);
	test(ret >= 0, "vma_free(%#lx, %#lx) = %u", VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, ret);

	ret = vma_free(VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE);
	test(ret >= 0, "vma_free(%#lx, %#lx) = %u", VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, ret);
	vma_dump();
}

/** @brief Test of the kernel malloc allocator */
static void malloc(void)
{
	int i;
	int* p[20];
	int* a;

	// kmalloc() test
	buddy_dump();
	a = kmalloc(SIZE);
	test(a != NULL, "kmalloc(%lu) = %p", SIZE, a);
	buddy_dump();

	// simple write/read test
	for (i=0; i<SIZE/sizeof(int); i++)
		a[i] = i;

	for (i=0; i<SIZE/sizeof(int); i++) {
		if (a[i] != i)
			test(0, "data mismatch: *(%p) != %lu", &a[i], i);
	}
	test(1, "data is equal");

	// kfree() test
	kfree(a);
	test(1, "kfree(%p)", a);
	buddy_dump();

	// some random malloc/free patterns to stress the buddy system
	for (i=0; i<20; i++) {
		uint16_t sz = rand();
		p[i] = kmalloc(sz);
		test(p[i] != NULL, "kmalloc(%u) = %p", sz, p[i]);
	}
	buddy_dump();

	for (i=0; i<20; i++) {
		kfree(p[i]);
		test(1, "kfree(%p)", p[i]);
	}
	buddy_dump();
}

/** @brief A memory benchmark for page table walks and TLB misses */
int bench(void)
{
	// init hardware performance counters
	struct pmc_caps* cap = pmc_init();
	if (cap->version == 0x21) { // QEmu returns garbage
		kputs("QEMU does not support PMCs.. skipping benchmark!\n");
		return -1;
	}

	kprintf("PMC architecural version: %u\n", cap->version);
	kprintf("There are %u general purpose PMCs (%u bit wide) available\n", cap->gp_count, cap->gp_width);
	kprintf("There are %u fixed function PMCs (%u bit wide) available\n", cap->ff_count, cap->ff_width);

	// setup PMCs
	pmc_stop_all();
	pmc_gp_config(0, PMC_EVT_PAGE_WALK_CLKS, PMC_EVTSEL_OS | PMC_EVTSEL_EN, 0, 0);
	pmc_gp_config(1, PMC_EVT_PAGE_WALK_COUNT, PMC_EVTSEL_OS | PMC_EVTSEL_EN, 0, 0);

	size_t phyaddr = get_page();
	size_t viraddr;
	size_t pages = 512*511;
	size_t virbase = 2*KERNEL_SPACE;
	kprintf("virbase %#llx KERNEL_SPACE %#llx\n", virbase, KERNEL_SPACE);
	for (viraddr = virbase; viraddr < virbase+pages*PAGE_SIZE; viraddr += PAGE_SIZE) {
		kprintf("map at %#llx\n", viraddr);
		size_t ret = map_region(viraddr, phyaddr, 1, MAP_KERNEL_SPACE);
		if (ret != viraddr) {
			kprintf("map failed at %#llx\n", viraddr);
			break;
		}
	}

	int i;
	for (i=0; i < ITERATIONS; i++) {
		tlb_flush();

		pmc_reset_all();
		pmc_start_all();

		for (viraddr = virbase; viraddr < virbase+pages*PAGE_SIZE; viraddr += PAGE_SIZE) {
			char * p = (char *) viraddr;
			(*p)++;
		}

		pmc_stop_all();

		uint64_t clks = pmc_gp_read(0);
		uint64_t count = pmc_gp_read(1);

		kprintf("%llu\n", 1000000 * clks / count);
	}

	return 0;
}

int smp(void* arg)
{
	kprintf("Hello from Core %d\n", smp_id());
	page_dump(PG_XD | PG_GLOBAL | PG_USER | PG_RW);

	return 33;
}

/** @brief This is a simple procedure to test memory management subsystem */
int memory(void* arg)
{
	int ret;
	tid_t id;

#if 0
	size_t t0, t1, t2, t3;
	size_t pages;
	for (pages = 1; pages < (1 << 25); pages++) {
		t0 = rdtsc();
		size_t ret = map_region((1 << 28), 0x1000, pages, MAP_KERNEL_SPACE);
		t1 = rdtsc();

		if (!ret)
			break;

		t2 = rdtsc();
		ret = unmap_region((1 << 28), pages);
		t3 = rdtsc();

		kprintf("%llu\t%llu\t%llu\n", pages, t1-t0, t3-t2);
	}

	kprintf("======== USER: malloc test...\n");
		char* argv[] = {"/bin/memtest", "25", "10"};

		ret = create_user_task(&id, argv[0], argv);
		test(!ret, "calling %s %s %s with id = %i, ret = %i", argv[0], argv[1], argv[2], id, ret);
		wait(&ret);
		test(!ret, "userspace task returned with code = %d", ret);
	return 0;

	kprintf("======== PAGING: test started...\n");
	paging();

	kprintf("======== VMA: test started...\n");
	vma();

	kprintf("======== MALLOC: test started...\n");
	malloc();

	kprintf("======== USER: test fork...\n");
	char* argv2[] = {"/bin/fork", NULL};
	ret = create_user_task(&id, argv2[0], argv2);
	test(!ret, "calling %s with id = %i, ret = %i", argv2[0], id, ret);
	wait(&ret);
	test(!ret, "userspace task returned with code = %d", ret);
#endif
	kprintf("======== BENCH: memory and TLB benchmark started...\n");
	bench();

	kprintf("======== SMP: test multicore...\n");
	ret = create_kernel_task_on_core(&id, smp, NULL, NORMAL_PRIO, 1);
	wait(&ret);
	test(!ret, "smp task returned with code = %d", ret);

	kprintf("======== All tests finished successfull...\n");

	return 0;
}