#ifndef LITMUS_BINARY_HEAP_H
#define LITMUS_BINARY_HEAP_H

/**
 * Simple binary heap with add, arbitrary delete, delete_root, and top
 * operations.
 *
 * Style meant to conform with list.h.
 *
 * Motivation: Linux's prio_heap.h is of fixed size. Litmus's binomial
 * heap may be overkill (and perhaps not general enough) for some applications.
 *
 * Note: In order to make node swaps fast, a node inserted with a data pointer
 * may not always hold said data pointer. This is similar to the binomial heap
 * implementation. This does make node deletion tricky since we have to
 * (1) locate the node that holds the data pointer to delete, and (2) the
 * node that was originally inserted with said data pointer. These have to be
 * coalesced into a single node before removal (see usage of
 * __binheap_safe_swap()). We have to track node references to accomplish this.
 */

struct binheap_node {
	void	*data;
	struct binheap_node *parent;
	struct binheap_node *left;
	struct binheap_node *right;

	/* pointer to binheap_node that holds *data for which this binheap_node
	 * was originally inserted.  (*data "owns" this node)
	 */
	struct binheap_node *ref;
	struct binheap_node **ref_ptr;
};

/**
 * Signature of compator function.  Assumed 'less-than' (min-heap).
 * Pass in 'greater-than' for max-heap.
 *
 * TODO: Consider macro-based implementation that allows comparator to be
 * inlined (similar to Linux red/black tree) for greater efficiency.
 */
typedef int (*binheap_order_t)(struct binheap_node *a,
							   struct binheap_node *b);


struct binheap_handle {
	struct binheap_node *root;

	/* pointer to node to take next inserted child */
	struct binheap_node *next;

	/* pointer to last node in complete binary tree */
	struct binheap_node *last;

	/* comparator function pointer */
	binheap_order_t compare;
};


/**
 * binheap_entry - get the struct for this heap node.
 *  Only valid when called upon heap nodes other than the root handle.
 * @ptr:	the heap node.
 * @type:	the type of struct pointed to by binheap_node::data.
 * @member:	unused.
 */
#define binheap_entry(ptr, type, member) \
((type *)((ptr)->data))

/**
 * binheap_node_container - get the struct that contains this node.
 *  Only valid when called upon heap nodes other than the root handle.
 * @ptr:	the heap node.
 * @type:	the type of struct the node is embedded in.
 * @member:	the name of the binheap_struct within the (type) struct.
 */
#define binheap_node_container(ptr, type, member) \
container_of((ptr), type, member)

/**
 * binheap_top_entry - get the struct for the node at the top of the heap.
 *  Only valid when called upon the heap handle node.
 * @ptr:    the special heap-handle node.
 * @type:   the type of the struct the head is embedded in.
 * @member:	the name of the binheap_struct within the (type) struct.
 */
#define binheap_top_entry(ptr, type, member) \
binheap_entry((ptr)->root, type, member)

/**
 * binheap_delete_root - remove the root element from the heap.
 * @handle:	 handle to the heap.
 * @type:    the type of the struct the head is embedded in.
 * @member:	 the name of the binheap_struct within the (type) struct.
 */
#define binheap_delete_root(handle, type, member) \
__binheap_delete_root((handle), &((type *)((handle)->root->data))->member)

/**
 * binheap_delete - remove an arbitrary element from the heap.
 * @to_delete:  pointer to node to be removed.
 * @handle:	 handle to the heap.
 */
#define binheap_delete(to_delete, handle) \
__binheap_delete((to_delete), (handle))

/**
 * binheap_add - insert an element to the heap
 * new_node: node to add.
 * @handle:	 handle to the heap.
 * @type:    the type of the struct the head is embedded in.
 * @member:	 the name of the binheap_struct within the (type) struct.
 */
#define binheap_add(new_node, handle, type, member) \
__binheap_add((new_node), (handle), container_of((new_node), type, member))

/**
 * binheap_decrease - re-eval the position of a node (based upon its
 * original data pointer).
 * @handle: handle to the heap.
 * @orig_node: node that was associated with the data pointer
 *             (whose value has changed) when said pointer was
 *             added to the heap.
 */
#define binheap_decrease(orig_node, handle) \
__binheap_decrease((orig_node), (handle))

#define BINHEAP_NODE_INIT() { NULL, NULL, NULL, NULL , NULL, NULL}

#define BINHEAP_NODE(name) \
	struct binheap_node name = BINHEAP_NODE_INIT()


static inline void INIT_BINHEAP_NODE(struct binheap_node *n)
{
	n->data = NULL;
	n->parent = NULL;
	n->left = NULL;
	n->right = NULL;
	n->ref = NULL;
	n->ref_ptr = NULL;
}

static inline void INIT_BINHEAP_HANDLE(
	struct binheap_handle *handle,
	binheap_order_t compare)
{
	handle->root = NULL;
	handle->next = NULL;
	handle->last = NULL;
	handle->compare = compare;
}

/* Returns true (1) if binheap is empty. */
static inline int binheap_empty(struct binheap_handle *handle)
{
	return(handle->root == NULL);
}


/* Update the node reference pointers.  Same logic as Litmus binomial heap. */
static inline void __update_ref(struct binheap_node *parent,
								struct binheap_node *child)
{
	*(parent->ref_ptr) = child;
	*(child->ref_ptr) = parent;

	swap(parent->ref_ptr, child->ref_ptr);
}

/* Swaps data between two nodes. */
static inline void __binheap_swap(struct binheap_node *parent,
								  struct binheap_node *child)
{
	swap(parent->data, child->data);
	__update_ref(parent, child);
}

/* Swaps memory and data between two nodes. Actual nodes swap instead of
 * just data.  Needed when we delete nodes from the heap.
 */
static inline void __binheap_swap_safe(struct binheap_handle *handle,
									   struct binheap_node *a,
									   struct binheap_node *b)
{
	swap(a->data, b->data);
	__update_ref(a, b);

	if((a->parent != NULL) && (a->parent == b->parent)) {
		/* special case: shared parent */
		swap(a->parent->left, a->parent->right);
	}
	else {
		/* Update pointers to swap parents. */

		if(a->parent) {
			if(a == a->parent->left) {
				a->parent->left = b;
			}
			else {
				a->parent->right = b;
			}
		}

		if(b->parent) {
			if(b == b->parent->left) {
				b->parent->left = a;
			}
			else {
				b->parent->right = a;
			}
		}

		swap(a->parent, b->parent);
	}

	/* swap children */

	if(a->left) {
		a->left->parent = b;

		if(a->right) {
			a->right->parent = b;
		}
	}

	if(b->left) {
		b->left->parent = a;

		if(b->right) {
			b->right->parent = a;
		}
	}

	swap(a->left, b->left);
	swap(a->right, b->right);


	/* update next/last/root pointers */

	if(a == handle->next) {
		handle->next = b;
	}
	else if(b == handle->next) {
		handle->next = a;
	}

	if(a == handle->last) {
		handle->last = b;
	}
	else if(b == handle->last) {
		handle->last = a;
	}

	if(a == handle->root) {
		handle->root = b;
	}
	else if(b == handle->root) {
		handle->root = a;
	}
}


/**
 * Update the pointer to the last node in the complete binary tree.
 * Called internally after the root node has been deleted.
 */
static inline void __binheap_update_last(struct binheap_handle *handle)
{
	struct binheap_node *temp = handle->last;

	/* find a "bend" in the tree. */
	while(temp->parent && (temp == temp->parent->left)) {
		temp = temp->parent;
	}

	/* step over to sibling if we're not at root */
	if(temp->parent != NULL) {
		temp = temp->parent->left;
	}

	/* now travel right as far as possible. */
	while(temp->right != NULL) {
		temp = temp->right;
	}

	/* take one step to the left if we're not at the bottom-most level. */
	if(temp->left != NULL) {
		temp = temp->left;
	}

	//BUG_ON(!(temp->left == NULL && temp->right == NULL));

	handle->last = temp;
}

/**
 * Update the pointer to the node that will take the next inserted node.
 * Called internally after a node has been inserted.
 */
static inline void __binheap_update_next(struct binheap_handle *handle)
{
	struct binheap_node *temp = handle->next;

	/* find a "bend" in the tree. */
	while(temp->parent && (temp == temp->parent->right)) {
		temp = temp->parent;
	}

	/* step over to sibling if we're not at root */
	if(temp->parent != NULL) {
		temp = temp->parent->right;
	}

	/* now travel left as far as possible. */
	while(temp->left != NULL) {
		temp = temp->left;
	}

	handle->next = temp;
}



/* bubble node up towards root */
static inline void __binheap_bubble_up(
	struct binheap_handle *handle,
	struct binheap_node *node)
{
	/* Note: NULL data pointers are used internally for arbitrary delete */
	while((node->parent != NULL) &&
			((node->data == NULL) /* let null data bubble to the top */ ||
			 handle->compare(node, node->parent))) {
		__binheap_swap(node->parent, node);
		node = node->parent;
	}
}


/* bubble node down, swapping with min-child */
static inline void __binheap_bubble_down(struct binheap_handle *handle)
{
	struct binheap_node *node = handle->root;

	while(node->left != NULL) {
		if(node->right && handle->compare(node->right, node->left)) {
			if(handle->compare(node->right, node)) {
				__binheap_swap(node, node->right);
				node = node->right;
			}
			else {
				break;
			}
		}
		else {
			if(handle->compare(node->left, node)) {
				__binheap_swap(node, node->left);
				node = node->left;
			}
			else {
				break;
			}
		}
	}
}



static inline void __binheap_add(struct binheap_node *new_node,
								 struct binheap_handle *handle,
								 void *data)
{
	/* NULL data pointers are used internally */
	if(!data) {
		WARN_ON(!data);
		return;
	}

	new_node->data = data;
	new_node->ref = new_node;
	new_node->ref_ptr = &(new_node->ref);

	if(!binheap_empty(handle)) {
		/* insert left side first */
		if(handle->next->left == NULL) {
			handle->next->left = new_node;
			new_node->parent = handle->next;
			new_node->left = NULL;
			new_node->right = NULL;

			handle->last = new_node;

			__binheap_bubble_up(handle, new_node);
		}
		else {
			/* left occupied. insert right. */
			handle->next->right = new_node;
			new_node->parent = handle->next;
			new_node->left = NULL;
			new_node->right = NULL;

			handle->last = new_node;

			__binheap_update_next(handle);
			__binheap_bubble_up(handle, new_node);
		}
	}
	else {
		/* first node in heap */

		new_node->parent = NULL;
		new_node->left = NULL;
		new_node->right = NULL;

		handle->root = new_node;
		handle->next = new_node;
		handle->last = new_node;
	}
}



/**
 * Removes the root node from the heap. The node is removed after coalescing
 * the binheap_node with its original data pointer at the root of the tree.
 *
 * The 'last' node in the tree is then swapped up to the root and bubbled
 * down.
 */
static inline void __binheap_delete_root(struct binheap_handle *handle,
										 struct binheap_node *container)
{
	struct binheap_node *root = handle->root;

	if(root != container) {
		/* coalesce */
		__binheap_swap_safe(handle, root, container);
		root = container;
	}

	if(handle->last != root) {
		/* swap 'last' node up to root and bubble it down. */

		struct binheap_node *to_move = handle->last;

		if(to_move->parent != root) {
			handle->next = to_move->parent;

			if(handle->next->right == to_move) {
				/* disconnect from parent */
				to_move->parent->right = NULL;
				handle->last = handle->next->left;
			}
			else {
				/* find new 'last' before we disconnect */
				__binheap_update_last(handle);

				/* disconnect from parent */
				to_move->parent->left = NULL;
			}
		}
		else {
			/* 'last' is direct child of root */

			handle->next = to_move;

			if(to_move == to_move->parent->right) {
				to_move->parent->right = NULL;
				handle->last = to_move->parent->left;
			}
			else {
				to_move->parent->left = NULL;
				handle->last = to_move;
			}
		}
		to_move->parent = NULL;

		/* reconnect as root.  We can't just swap data ptrs since root node
		 * may be freed after this function returns.
		 */
		to_move->left = root->left;
		to_move->right = root->right;
		if(to_move->left != NULL) {
			to_move->left->parent = to_move;
		}
		if(to_move->right != NULL) {
			to_move->right->parent = to_move;
		}

		handle->root = to_move;

		/* bubble down */
		__binheap_bubble_down(handle);
	}
	else {
		/* removing last node in tree */
		handle->root = NULL;
		handle->next = NULL;
		handle->last = NULL;
	}
}


/**
 * Delete an arbitrary node.  Bubble node to delete up to the root,
 * and then delete to root.
 */
static inline void __binheap_delete(
								struct binheap_node *node_to_delete,
								struct binheap_handle *handle)
{
	struct binheap_node *target = node_to_delete->ref;
	void *temp_data = target->data;

	/* temporarily set data to null to allow node to bubble up to the top. */
	target->data = NULL;

	__binheap_bubble_up(handle, target);
	__binheap_delete_root(handle, node_to_delete);

	node_to_delete->data = temp_data;  /* restore node data pointer */
}

/**
 * Bubble up a node whose pointer has decreased in value.
 */
static inline void __binheap_decrease(
				struct binheap_node *orig_node,
				struct binheap_handle *handle)
{
	struct binheap_node *target = orig_node->ref;
	__binheap_bubble_up(handle, target);
}

#endif