#include <linux/module.h>
#include <linux/types.h>
#include <linux/kernel.h>
#include <linux/mm.h>
#include <linux/highmem.h>
#include <linux/pagemap.h>
#include <linux/errno.h>
#include <linux/string.h>
#include <linux/delay.h>
#include <linux/fb.h>
#include <linux/fs.h>
#include <linux/file.h>
#include <linux/ioctl.h>
#include <linux/init.h>
#include <linux/pci.h>
#include <asm/io.h>
#include <linux/uaccess.h>

#ifdef CONFIG_MTRR
#include <asm/mtrr.h>
#endif

//#define DEBUG

#ifdef DEBUG
#define dprintk(fmt, args...)		\
	printk(fmt, ##args)
#define debug_show_mtrr()	show_mtrr()
#else
//#define pr_debug(format, ...) do {} while (0)
#define	dprintk(fmt, ...)	do {} while (0)
#define	debug_show_mtrr()	do {} while (0)
#endif

MODULE_LICENSE("GPL");

#define NR_PAGES_ORDER	1

static u64 __do_write(struct page *page, int page_nr)
{
	void *addr = page_address(page);
	u64 total = 0;
	int i;
	
	for (i = 0; i < PAGE_SIZE * page_nr; i++) {
		((char *)addr)[i] = (char)0x99;
	}

	return total;
}

static u64 __do_read(struct page *page, int page_nr)
{
	void *addr = page_address(page);
	u64 total = 0;
	int i;
	
	for (i = 0; i < PAGE_SIZE * page_nr; i++) {
		total += ((char *)addr)[i];
		if (((char *)addr)[i] != (char)0x99) {
			printk("read failed %x expected 0x99.\n", ((char *)addr)[i]);
			break;
		}
	}

	return total;
}

static void do_rw(struct page *page, int page_nr, const char *cachetype)
{
	ktime_t start, end;
	u64 total;
	
	start = ktime_get();
	__do_write(page, page_nr);
	total = __do_read(page, page_nr);
	end = ktime_get();
	printk("Running in %s %lld ns.\n", cachetype,
	       ktime_to_ns(ktime_sub(end, start)));	
}

static int add_region_to_mtrr(struct page *page, int page_nr, int type,
			      const char *cachetype)
{
	unsigned long pfn;
	int ret;
	
	pfn = page_to_pfn(page);
	
	ret = mtrr_add(pfn << PAGE_SHIFT, PAGE_SIZE * page_nr, type, 1);
	if (ret < 0)
		printk("mtrr_add %s failed, return %d.\n", cachetype, ret);

	dprintk("%s add region %lx nrpage %d.\n", cachetype, pfn << PAGE_SHIFT, page_nr);	
	return ret;
}

static int del_region_from_mtrr(int entry, struct page *page, int page_nr, const char *cachetype)
{
	unsigned long pfn;
	int ret;
	
	pfn = page_to_pfn(page);
	
	ret = mtrr_del(entry, pfn << PAGE_SHIFT, PAGE_SIZE * page_nr);
	if (ret < 0)
		printk("mtrr_del %s failed, return %d.\n", cachetype, ret);
	dprintk("%s del region %lx nrpage %d.\n", cachetype, pfn << PAGE_SHIFT, page_nr);	
	return ret;
}

static int memory_rw(void)
{
	struct page *page;
	int page_order, page_nr, ret;
	
	page_order = NR_PAGES_ORDER;
	page_nr = 1 << page_order;
	
	page = alloc_pages(GFP_KERNEL, page_order);
	if (!page)
		return -ENOMEM;
	
	ret = add_region_to_mtrr(page, page_nr, MTRR_TYPE_UNCACHABLE,
				 "MTRR_TYPE_UNCACHABLE");
	if (ret < 0)
		goto exit;

	do_rw(page, page_nr, "preload MTRR_TYPE_UNCACHABLE");
	do_rw(page, page_nr, "MTRR_TYPE_UNCACHABLE");
	
	ret = del_region_from_mtrr(ret, page, page_nr, "MTRR_TYPE_UNCACHABLE");
	if (ret < 0)
		goto exit;
	
	ret = add_region_to_mtrr(page, page_nr, MTRR_TYPE_WRBACK,
				 "MTRR_TYPE_WRBACK");
	if (ret < 0)
		goto exit;
	
	do_rw(page, page_nr, "preload MTRR_TYPE_WRBACK");
	do_rw(page, page_nr, "MTRR_TYPE_WRBACK");

	ret = del_region_from_mtrr(ret, page, page_nr, "MTRR_TYPE_WRBACK");
	if (ret < 0)
		goto exit;
		
	ret = 0;
exit:
	__free_pages(page, page_order);
	return ret;
		
}

static int __init my_init(void)
{
	printk("__INIT__.\n");

	dprintk("######### INIT MTRR...\n");

	return memory_rw();
}

static void __exit my_exit(void)
{
	dprintk("###### EXIT MTRR...\n");
	return;
}

module_init(my_init)
module_exit(my_exit)