/* 
 * 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/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/processor.h>

#define PAGE_COUNT	10
#define SIZE		(PAGE_COUNT*PAGE_SIZE)
#define VIRT_FROM_ADDR	0x100000000000
#define VIRT_TO_ADDR	0x200000000000

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)
{
	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);

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

	page_dump(0, -1L);

	return 0;
}

/** @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;

	// show original page maps
	t = rdtsc();
	page_dump(0, -1L);
	kprintf("delta_t = %lu\n", rdtsc() - t);

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

	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) = 0x%lx", PAGE_COUNT, phys);

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

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

	// 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, MAP_USER_SPACE);
	test(virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx", VIRT_TO_ADDR, phys, PAGE_COUNT, 0, virt_to);

	// show pagings infos again
	page_dump(0, -1L);
	page_stats(0, -1L, 0);

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

	// 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, MAP_USER_SPACE);
	test(!virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%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|MAP_USER_SPACE);
	test(virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%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)
	p2 = (size_t *) virt_to;
	for (c = 0; c < SIZE/sizeof(size_t); c++) {
		if (p1[c] == p2[c])
			test(0, "data match at *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is unequal");

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

	page_dump(0, -1L);

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

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

	create_kernel_task(0, paging_stage2, &sum, NORMAL_PRIO);
	sleep(5);
}

/** @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(0x%x, 0x%x) = 0x%lx", SIZE, VMA_HEAP, a1);

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

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

	ret = vma_free(a2, a2+SIZE);
	test(ret >= 0, "vma_free(0x%lx, 0x%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(0x%lx, 0x%lx, 0x%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(0x%lx, 0x%lx, 0x%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(0x%lx, 0x%lx, 0x%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(0x%lx, 0x%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(0x%lx, 0x%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(0x%lx, 0x%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_config(0, PMC_EVT_PAGE_WALK_CLKS, PMC_EVTSEL_OS | PMC_EVTSEL_EN, 0, 0);
	pmc_config(1, PMC_EVT_PAGE_WALK_COUNT, PMC_EVTSEL_OS | PMC_EVTSEL_EN, 0, 0);

	// allocate space for results
	uint64_t *data = kmalloc(ITERATIONS * sizeof(uint64_t));
	if (!data)
		return -1;

	// clear caches
	tlb_flush();
	flush_cache();

	int i;
	for (i=0; i < ITERATIONS; i++) {
		pmc_write(0, 0);
		pmc_write(1, 0);

		pmc_start_all();

#if 0
		int i = 100;
		while (i--) {
			tlb_flush();
			page_stats(0);
		}
#else
		//flush_cache();
		//tlb_flush();
		page_stats(0);
#endif

		pmc_stop_all();

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

		/*kprintf("Number of Page table walks: %lu\n", count);
		kprintf("Page table walks clock cycles: %lu\n", clks);
		kprintf("Cycles per table walk: %lu.%u\n", clks / count, (1000 * clks / count) % 1000 );*/

		data[i] = 1000000 * clks / count;
	}

	// dump results
	for (i=0; i<ITERATIONS; i++)
		kprintf("%u\t%lu\n", i, data[i]);

	return 0;
}

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

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

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

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

	kprintf("======== USER: test userspace...\n");
	char* argv[] = {"/bin/fork", NULL};
	ret = create_user_task(&id, argv[0], argv);
	test(!ret, "calling %s with id = %i, ret = %i", argv[0], id, ret);
	wait(&ret);
	test(!ret, "userspace task returned with code = %d", ret);

	kprintf("======== BENCH: memory and TLB benchmark started...\n");
	bench();

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

	return 0;
}