#include <linux/sched.h>
#include <linux/slab.h>

#include <litmus/litmus.h>
#include <litmus/dgl.h>
#include <litmus/sched_trace.h>

#define MASK_SIZE     (sizeof(unsigned long) * 8)

/* Return number of MASK_SIZE fields needed to store a mask in d */
#define WP(num, word) (num / word + (num % word != 0))
#define MASK_WORDS(d) WP(d->num_resources, MASK_SIZE)

/* Word, bit -> resource id */
#define ri(w, b) (w * MASK_SIZE + b)

 /* For loop, where @i iterates over each set bit in @bit_arr */
#define for_each_resource(bit_arr, d, w, b, i)				\
	for(w = 0; w < MASK_WORDS(d); ++w)				\
		 for(b = find_first_bit(&bit_arr[w],MASK_SIZE), i = ri(w, b);  \
		     b < MASK_SIZE;					       \
		     b = find_next_bit(&bit_arr[w],MASK_SIZE,b+1), i = ri(w, b))

/* Return resource id in dgl @d for resource @r */
#define resource_id(d, r) ((((void*)r) - (void*)((d)->resources))/ sizeof(*r))

/* Return request group of req @r for resource @i */
#define req_group(r, i) (container_of(((void*)r) - sizeof(*r)*(i),	\
				      struct dgl_group_req, requests))

/* Resource id -> word, bit */
static inline void mask_idx(int resource, int *word, int *bit)
{
	*word = resource / MASK_SIZE;
	*bit  = resource % MASK_SIZE;
}


static void print_waiting(struct dgl *dgl, struct dgl_resource *resource)
{
	struct dgl_req *pos;
	struct dgl_group_req *greq;
	unsigned long long last = 0;

	TRACE("List for rid %d\n", resource_id(dgl, resource));
	list_for_each_entry(pos, &resource->waiting, list) {
		greq = pos->greq;
		TRACE("  0x%p with timestamp %llu\n", greq, greq->ts);
		BUG_ON(greq->ts < last);
		last = greq->ts;
	}
}

void dgl_init(struct dgl *dgl, unsigned long num_resources,
	      unsigned long num_replicas)
{
	int i;
	struct dgl_resource *resource;

	dgl->num_replicas  = num_replicas;
	dgl->num_resources = num_resources;

	dgl->resources = kmalloc(sizeof(*dgl->resources) * num_resources,
				 GFP_ATOMIC);
	dgl->acquired  = kmalloc(sizeof(*dgl->acquired) * num_online_cpus(),
				 GFP_ATOMIC);

	for (i = 0; i < num_online_cpus(); ++i)
		dgl->acquired[i] = NULL;

	for (i = 0; i < num_resources; i++) {
		resource = &dgl->resources[i];

		INIT_LIST_HEAD(&resource->waiting);
		resource->free_replicas = dgl->num_replicas;
	}

	dgl->requests = 0;
	dgl->running  = 0;
	dgl->ts = 0;
}

void dgl_free(struct dgl *dgl)
{
	kfree(dgl->resources);
	kfree(dgl->acquired);
}

void dgl_group_req_init(struct dgl *dgl, struct dgl_group_req *greq)
{
	int i;

	greq->requested = kmalloc(sizeof(*greq->requested) * MASK_WORDS(dgl),
				  GFP_ATOMIC);
	greq->waiting   = kmalloc(sizeof(*greq->waiting) * MASK_WORDS(dgl),
				  GFP_ATOMIC);
	greq->requests  = kmalloc(sizeof(*greq->requests) * dgl->num_resources,
				  GFP_ATOMIC);

	BUG_ON(!greq->requested);
	BUG_ON(!greq->waiting);
	BUG_ON(!greq->requests);

	greq->cpu = NO_CPU;
	for (i = 0; i < MASK_WORDS(dgl); ++i) {
		greq->requested[i] = 0;
		greq->waiting[i]   = 0;
	}
}

void dgl_group_req_free(struct dgl_group_req *greq)
{
	kfree(greq->requested);
	kfree(greq->waiting);
	kfree(greq->requests);
}

/**
 * set_req - create request for @replicas of @resource.
 */
void set_req(struct dgl *dgl, struct dgl_group_req *greq,
	     int resource, int replicas)
{
	int word, bit;
	struct dgl_req *req;

	if (replicas > dgl->num_replicas)
		replicas = dgl->num_replicas;

	mask_idx(resource, &word, &bit);
	__set_bit(bit, &greq->requested[word]);

	TRACE("0x%p requesting %d of %d\n", greq, replicas, resource);

	req = &greq->requests[resource];
	req->greq = greq;
	INIT_LIST_HEAD(&req->list);
	req->replicas = replicas;
}

/*
 * Attempt to fulfill request @req for @resource.
 * Return 1 if successful. If the matching group request has acquired all of
 * its needed resources, this will then set that req as dgl->acquired[cpu].
 */
static unsigned long try_acquire(struct dgl *dgl, struct dgl_resource *resource,
				 struct dgl_req *req)
{
	int word, bit, rid, head, empty, room;
	unsigned long waiting;
	struct dgl_group_req *greq;

	rid  = resource_id(dgl, resource);
	greq = req->greq;

	TRACE("0x%p greq\n", greq);

	head  = resource->waiting.next == &req->list;
	empty = list_empty(&resource->waiting);
	room  = resource->free_replicas >= req->replicas;

	if (! (room && (head || empty)) ) {
		TRACE("0x%p cannot acquire %d replicas, %d free\n",
		      greq, req->replicas, resource->free_replicas,
		      room, head, empty);
		return 0;
	}

	resource->free_replicas -= req->replicas;

	TRACE("0x%p acquired %d replicas of rid %d\n",
	      greq, req->replicas, rid);

	mask_idx(rid, &word, &bit);


	TRACE("0x%p, %lu, 0x%p\n", greq->waiting, greq->waiting[word],
	      &greq->waiting[word]);

	clear_bit(bit, &greq->waiting[word]);

        waiting = 0;
	for (word = 0; word < MASK_WORDS(dgl); word++) {
		waiting |= greq->waiting[word];
		if (waiting)
			break;
	}

	if (!waiting) {
		TRACE("0x%p acquired all resources\n", greq);
		BUG_ON(dgl->acquired[greq->cpu]);
		dgl->acquired[greq->cpu] = greq;
		litmus_reschedule(greq->cpu);
		dgl->running++;
	}

	return 1;
}

/**
 * add_group_req - initiate group request.
 */
void add_group_req(struct dgl *dgl, struct dgl_group_req *greq, int cpu)
{
	int b, w, i, succ, all_succ = 1;
	struct dgl_req *req;
	struct dgl_resource *resource;

	greq->cpu = cpu;
	greq->ts = dgl->ts++;

	TRACE("0x%p group request added for CPU %d\n", greq, cpu);
	BUG_ON(dgl->acquired[cpu] == greq);

	++dgl->requests;

	for_each_resource(greq->requested, dgl, w, b, i) {
		__set_bit(b, &greq->waiting[w]);
	}

	for_each_resource(greq->requested, dgl, w, b, i) {
		req = &greq->requests[i];
		resource = &dgl->resources[i];

		succ = try_acquire(dgl, resource, req);
		all_succ &= succ;

		if (!succ) {
			TRACE("0x%p waiting on rid %d\n", greq, i);
			list_add_tail(&req->list, &resource->waiting);
		}
	}

	/* Grant empty requests */
	if (all_succ && !dgl->acquired[cpu]) {
		TRACE("0x%p empty group request acquired cpu %d\n", greq, cpu);
		dgl->acquired[cpu] = greq;
		++dgl->running;
	}

	BUG_ON(dgl->requests && !dgl->running);
}

/**
 * remove_group_req - abandon group request.
 *
 * This will also progress the waiting queues of resources acquired by @greq.
 */
void remove_group_req(struct dgl *dgl, struct dgl_group_req *greq)
{
	int b, w, i;
	struct dgl_req *req, *next;
	struct dgl_resource *resource;

	TRACE("0x%p removing group request for CPU %d\n", greq, greq->cpu);

	--dgl->requests;

	if (dgl->acquired[greq->cpu] == greq) {
		TRACE("0x%p no longer acquired on CPU %d\n", greq, greq->cpu);
		dgl->acquired[greq->cpu] = NULL;
		--dgl->running;
	}

	for_each_resource(greq->requested, dgl, w, b, i) {
		req = &greq->requests[i];
		resource = &dgl->resources[i];

		if (!list_empty(&req->list)) {
			/* Waiting on resource */
			clear_bit(b, &greq->waiting[w]);
			list_del_init(&req->list);
			TRACE("Quitting 0x%p from rid %d\n",
			      req, i);
		} else {
			/* Have resource */
			resource->free_replicas += req->replicas;
			BUG_ON(resource->free_replicas > dgl->num_replicas);
			TRACE("0x%p releasing %d of %d replicas, rid %d\n",
			      greq, req->replicas, resource->free_replicas, i);

			if (!list_empty(&resource->waiting)) {
				/* Give it to the next guy */
				next = list_first_entry(&resource->waiting,
							struct dgl_req,
							list);

				BUG_ON(next->greq->ts < greq->ts);

				if (try_acquire(dgl, resource, next)) {
					list_del_init(&next->list);
					print_waiting(dgl, resource);

				}
			}
		}
	}

	BUG_ON(dgl->requests && !dgl->running);
}