#include <litmus/way_tracker.h>
#include <asm/io.h>
#include <litmus/color.h>
#include <linux/sched.h>
#include <linux/module.h>
#include <linux/bitops.h>
#include <linux/slab.h>
#include <linux/list.h>

#include <litmus/litmus.h>
#include <litmus/color.h>

static raw_spinlock_t lock;

static unsigned long *ways;

static int peek_next_way(unsigned int color)
{
	int ret;

	BUG_ON(color > color_cache_info.nr_colors);

	ret = find_first_bit(&ways[color], sizeof(unsigned long)*8);
	if (ret == sizeof(unsigned long)*8) {
		return -1;
	} else {
		return ret;
	}
}

static int take_next_way(unsigned int color)
{
	int idx;

	raw_spin_lock(&lock);
	idx = peek_next_way(color);
	if (idx != -1) {
		clear_bit(idx, &ways[color]);
		TRACE("Took, now %lu free of color %d\n", hweight_long(ways[color]), color);
	} else {
		/* Seriously bad. */
		BUG();
	}
	raw_spin_unlock(&lock);
	return idx;
}

static void release_way(unsigned int color, int way)
{
	BUG_ON(way < 0 || way > color_cache_info.ways);
	raw_spin_lock(&lock);
	__set_bit(way, &ways[color]);
	TRACE("Released, now %lu free of color %d\n", hweight_long(ways[color]), color);
	raw_spin_unlock(&lock);
}

void color_page_info_take_ways(struct task_struct *t)
{
	struct color_page_info *cur;

	list_for_each_entry(cur, &tsk_rt(t)->color_page_info_list, list) {
		cur->way = take_next_way(cur->color);
		TRACE_TASK(t, "assign color %d way %d for info %p\n",
				cur->color, cur->way, cur);
	}
}

void color_page_info_release_ways(struct task_struct *t)
{
	struct color_page_info *cur;

	list_for_each_entry(cur, &tsk_rt(t)->color_page_info_list, list) {
		release_way(cur->color, cur->way);
		TRACE_TASK(t, "release color %d way %d for info %p\n",
				cur->color, cur->way, cur);
	}
}

void reset_way_tracker(void)
{
	int color, unused;

	for (color = 0; color < color_cache_info.nr_colors; ++color) {
		unused = sizeof(unsigned long)*8 - color_cache_info.ways;
		ways[color] = ULONG_MAX >> unused;
	}

}

static int __init init_way_tracker(void)
{
	int size;

	BUG_ON(!color_cache_info.nr_colors);
	BUG_ON(color_cache_info.ways > sizeof(unsigned long)*8);

	printk(KERN_INFO "Starting way tracker\n");

	size = sizeof(unsigned long) * color_cache_info.nr_colors;
	ways = kmalloc(size, GFP_ATOMIC);

	reset_way_tracker();

	raw_spin_lock_init(&lock);

	return 0;
}

module_init(init_way_tracker);