/* rbTree - rbTreeRed-rbTreeBlack Tree - a type of binary tree which 
 * automatically keeps relatively balanced during
 * inserts and deletions.
 *   original author: Shane Saunders
 *   adapted into local conventions: Jim Kent
 */

#include "common.h"
#include "localmem.h"
#include "rbTree.h"



static struct rbTreeNode *restructure(struct rbTree *t, int tos, 
	struct rbTreeNode *x, struct rbTreeNode *y, struct rbTreeNode *z)
/* General restructuring function - checks for all
 * restructuring cases.  Call when insert has messed up tree.
 * Sadly delete has to do even more work. */
{
struct rbTreeNode *parent, *midNode;

if(y == x->left) 
    {
    if(z == y->left) 
        {  /* in-order:  z, y, x */
	midNode = y;
	y->left = z;
	x->left = y->right;
	y->right = x;
	}
    else 
        {  /* in-order:  y, z, x */
	midNode = z;
	y->right = z->left;
	z->left = y;
	x->left = z->right;
	z->right = x;
	}
    }
else 
    {
    if(z == y->left) 
	{  /* in-order:  x, z, y */
	midNode = z;
	x->right = z->left;
	z->left = x;
	y->left = z->right;
	z->right = y;
	}
    else 
	{  /* in-order:  x, y, z */
	midNode = y;
	x->right = y->left;
	y->left = x;
	y->right = z;
	}
    }
if(tos != 0) 
    {
    parent = t->stack[tos-1];
    if(x == parent->left) 
	parent->left = midNode;
    else 
	parent->right = midNode;
    }
else 
    t->root = midNode;

return midNode;
}

struct rbTree *rbTreeNewDetailed(int (*compare)(void *, void *), struct lm *lm, 
	struct rbTreeNode *stack[128])
/* Allocate rbTree on an existing local memory & stack.  This is for cases
 * where you want a lot of trees, and don't want the overhead for each one. 
 * Note, to clean these up, just do freez(&rbTree) rather than rbFreeTree(&rbTree). */
{
struct rbTree *t;
AllocVar(t);
t->root = NULL;
t->compare = compare;
t->lm = lm;
t->stack = stack;	
t->n = 0;
return t;
}

struct rbTree *rbTreeNew(int (*compare)(void *, void *))
/* rbTreeNew() - Allocates space for a red-black tree and returns a pointer
 * to it.  The function compare compares they keys of two items, and returns a
 * negative, zero, or positive integer depending on whether the first item is
 * less than, equal to, or greater than the second.  */
{
/* The stack keeps us from having to keep explicit
 * parent, grandparent, greatgrandparent variables.
 * It needs to be big enough for the maximum depth
 * of tree.  Since the whole point of rb trees is
 * that they are self-balancing, this is not all
 * that deep, just 2*log2(N).  Therefore a stack of
 * 128 is good for up to 2^64 items in stack, which
 * should keep us for the next couple of decades... */
struct lm *lm = lmInit(0);
struct rbTreeNode **stack = lmAlloc(lm, 128 * sizeof(stack[0]));	
return rbTreeNewDetailed(compare, lm, stack);
}


void rbTreeFree(struct rbTree **pTree)
/* rbTreeFree() - Frees space used by the red-black tree pointed to by t. */
{
struct rbTree *tree = *pTree;
if (tree != NULL)
    {
    lmCleanup(&tree->lm);
    freez(pTree);
    }
}

void rbTreeFreeList(struct rbTree **pList)
/* Free up a list of rbTrees. */
{
struct rbTree *tree, *next;
for (tree = *pList; tree != NULL; tree = next)
    {
    next = tree->next;
    rbTreeFree(&tree);
    }
}

void *rbTreeAdd(struct rbTree *t, void *item)
/* rbTreeAdd() - Inserts an item into the red-black tree pointed to by t,
 * according the the value its key.  The key of an item in the red-black
 * tree must be unique among items in the tree.  If an item with the same key
 * already exists in the tree, a pointer to that item is returned.  Otherwise,
 * NULL is returned, indicating insertion was successful.
 */
{
struct rbTreeNode *x, *p, *q, *m, **attachX;
int (* compare)(void *, void *);
int cmpResult;
rbTreeColor col;
struct rbTreeNode **stack = NULL;
int tos;

tos = 0;    
if((p = t->root) != NULL) 
    {
    compare = t->compare;
    stack = t->stack;

    /* Repeatedly explore either the left branch or the right branch
     * depending on the value of the key, until an empty branch is chosen.
     */
    for(;;) 
        {
	stack[tos++] = p;
	cmpResult = compare(item, p->item);
	if(cmpResult < 0) 
	    {
	    p = p->left;
	    if(!p) 
	        {
		p = stack[--tos];
		attachX = &p->left;
		break;
		}
	    }
	else if(cmpResult > 0) 
	    {
	    p = p->right;
	    if(!p) 
	        {
		p = stack[--tos];
		attachX = &p->right;
		break;
		}
	    }
	else 
	    {
	    return p->item;
	    }
	}
    col = rbTreeRed;
    }
else 
    {
    attachX = &t->root;
    col = rbTreeBlack;
    }

/* Allocate new node and place it in tree. */
if ((x = t->freeList) != NULL)
    t->freeList = x->right;
else
    lmAllocVar(t->lm, x);
x->left = x->right = NULL;
x->item = item;
x->color = col;
*attachX = x;
t->n++;

/* Restructuring or recolouring will be needed if node x and its parent, p,
 * are both red.
 */
if(tos > 0) 
    {
    while(p->color == rbTreeRed) 
	{  /* Double red problem. */

	/* Obtain a pointer to p's parent, m, and sibling, q. */
	m = stack[--tos];
	q = p == m->left ? m->right : m->left;
	
	/* Determine whether restructuring or recolouring is needed. */
	if(!q || q->color == rbTreeBlack) 
	    {
	    /* Sibling is black.  ==>  Perform restructuring. */
	    
	    /* Restructure according to the left to right order, of nodes
	     * m, p, and x.
	     */
	    m = restructure(t, tos, m, p, x);
	    m->color = rbTreeBlack;
	    m->left->color = m->right->color = rbTreeRed;

	    /* Restructuring eliminates the double red problem. */
	    break;
	    }
	/* else just need to flip color */
	
	/* Sibling is also red.  ==>  Perform recolouring. */
	p->color = rbTreeBlack;
	q->color = rbTreeBlack;

	if(tos == 0) break;  /* The root node always remains black. */
	    
	m->color = rbTreeRed;

	/* Continue, checking colouring higher up. */
	x = m;
	p = stack[--tos];
	}
    }

return NULL;
}


void *rbTreeFind(struct rbTree *t, void *item)
/* rbTreeFind() - Find an item in the red-black tree with the same key as the
 * item pointed to by `item'.  Returns a pointer to the item found, or NULL
 * if no item was found.
 */
{
struct rbTreeNode *p, *nextP;
int (*compare)(void *, void *) = t->compare;
int cmpResult;
    
/* Repeatedly explore either the left or right branch, depending on the
 * value of the key, until the correct item is found.  */
for (p = t->root; p != NULL; p = nextP)
    {
    cmpResult = compare(item, p->item);
    if(cmpResult < 0) 
	nextP = p->left;
    else if(cmpResult > 0) 
	nextP = p->right;
    else 
	return p->item;
    }
return NULL;
}


void *rbTreeRemove(struct rbTree *t, void *item)
/* rbTreeRemove() - Delete an item in the red-black tree with the same key as
 * the item pointed to by `item'.  Returns a pointer to the deleted item,
 * and NULL if no item was found.
 */
{
struct rbTreeNode *p, *r, *x, *y, *z, *b, *newY;
struct rbTreeNode *m;
rbTreeColor removeCol;
void *returnItem;
int (* compare)(void *, void *);
int cmpResult;
struct rbTreeNode **stack;
int i, tos;


/* Attempt to locate the item to be deleted. */
if((p = t->root)) 
    {
    compare = t->compare;
    stack = t->stack;
    tos = 0;
    
    for(;;) 
	{
	stack[tos++] = p;
	cmpResult = compare(item, p->item);
	if(cmpResult < 0) 
	    p = p->left;
	else if(cmpResult > 0) 
	    p = p->right;
	else 
	    /* Item found. */
	    break;
	if(!p) return NULL;
	}
    }
else 
    return NULL;

/* p points to the node to be deleted, and is currently on the top of the
 * stack.
 */
if(!p->left) 
    {
    tos--;  /* Adjust tos to remove p. */
    /* Right child replaces p. */
    if(tos == 0) 
	{
	r = t->root = p->right;
	x = y = NULL;
	}
    else 
	{
	x = stack[--tos];
	if(p == x->left) 
	    {
	    r = x->left = p->right;
	    y = x->right;
	    }
	else 
	    {
	    r = x->right = p->right;
	    y = x->left;
	    }
	}
    removeCol = p->color;
    }
else if(!p->right) 
    {
    tos--;  /* Adjust tos to remove p. */
    /* Left child replaces p. */
    if(tos == 0) 
	{
	r = t->root = p->left;
	x = y = NULL;
	}
    else 
	{
	x = stack[--tos];
	if(p == x->left) 
	    {
	    r = x->left = p->left;
	    y = x->right;
	    }
	else 
	    {
	    r = x->right = p->left;
	    y = x->left;
	    }
	}
    removeCol = p->color;
    }
else 
    {
    /* Save p's stack position. */
    i = tos-1;
    
    /* Minimum child, m, in the right subtree replaces p. */
    m = p->right;
    do 
	{
	stack[tos++] = m;
	m = m->left;
	} while(m);
    m = stack[--tos];

    /* Update either the left or right child pointers of p's parent. */
    if(i == 0) 
	{
	t->root = m;
	}
    else 
	{
	x = stack[i-1];  /* p's parent. */
	if(p == x->left) 
	    {
	    x->left = m;
	    }
	else 
	    {
	    x->right = m;
	    }
	}
    
    /* Update the tree part m is removed from, and assign m the child
     * pointers of p (only if m is not the right child of p).
     */
    stack[i] = m;  /* Node m replaces node p on the stack. */
    x = stack[--tos];
    r = m->right;
    if(tos != i) 
	{  /* x is equal to the parent of m. */
	y = x->right;
	x->left = r;
	m->right = p->right;
	}
    else 
	{ /* m was the right child of p, and x is equal to m. */
	y = p->left;
	}
    m->left = p->left;

    /* We treat node m as the node which has been removed. */
    removeCol = m->color;
    m->color = p->color;
    }

/* Get return value and reuse the space used by node p. */
returnItem = p->item;
p->right = t->freeList;
t->freeList = p;

t->n--;

/* The pointers x, y, and r point to nodes which may be involved in
 * restructuring and recolouring.
 *  x - the parent of the removed node.
 *  y - the sibling of the removed node.
 *  r - the node which replaced the removed node.
 * From the above code, the next entry off the stack will be the parent of
 * node x.
 */

/* The number of black nodes on paths to all external nodes (NULL child
 * pointers) must remain the same for all paths.  Restructuring or
 * recolouring of nodes may be necessary to enforce this.
 */
if(removeCol == rbTreeBlack) 
    {
    /* Removal of a black node requires some adjustment. */
    
    if(!r || r->color == rbTreeBlack) 
	{
	/* A black node replaced the deleted black node.  Note that
	 * external nodes (NULL child pointers) are always black, so
	 * if r is NULL it is treated as a black node.
	 */

	/* This causes a double-black problem, since node r would need to
	 * be coloured double-black in order for the black color on
	 * paths through r to remain the same as for other paths.
	 */

	/* If r is the root node, the double-black color is not necessary
	 * to maintain the color balance.  Otherwise, some adjustment of
	 * nearby nodes is needed in order to eliminate the double-black
	 * problem.  NOTE:  x points to the parent of r.
	 */
	if(x) for(;;) 
	    {

	    /* There are three adjustment cases:
	     *  1.  r's sibling, y, is black and has a red child, z.
	     *  2.  r's sibling, y, is black and has two black children.
	     *  3.  r's sibling, y, is red.
	     */
	    if(y->color == rbTreeBlack) 
		{

		/* Note the conditional evaluation for assigning z. */
		if(((z = y->left) && z->color == rbTreeRed) ||
		   ((z = y->right) && z->color == rbTreeRed)) 
		       {		    
		    /* Case 1:  perform a restructuring of nodes x, y, and
		     * z.
		     */
		    
		    b = restructure(t, tos, x, y, z);
		    b->color = x->color;
		    b->left->color = b->right->color = rbTreeBlack;
		    
		    break;
		    }
		else 
		    {
		    /* Case 2:  recolour node y red. */
		    
		    y->color = rbTreeRed;
		    
		    if(x->color == rbTreeRed) 
			{
			x->color = rbTreeBlack;
			break;
			}
		    /* else */

		    if(tos == 0) break;  /* Root level reached. */
		    /* else */
		    
		    r = x;
		    x = stack[--tos];  /* x <- parent of x. */
		    y = x->left == r ? x->right : x->left;
		    }
		}
	    else 
		{
		/* Case 3:  Restructure nodes x, y, and z, where:
		 *  - If node y is the left child of x, then z is the left
		 *    child of y.  Otherwise z is the right child of y.
		 */
		if(x->left == y) 
		    {
		    newY = y->right;
		    z = y->left;
		    }
		else 
		    {
		    newY = y->left;
		    z = y->right;
		    }
		
		restructure(t, tos, x, y, z);
		y->color = rbTreeBlack;
		x->color = rbTreeRed;

		/* Since x has moved down a place in the tree, and y is the
		 * new the parent of x, the stack must be adjusted so that
		 * the parent of x is correctly identified in the next call
		 * to restructure().
		 */
		stack[tos++] = y;

		/* After restructuring, node r has a black sibling, newY,
		 * so either case 1 or case 2 applies.  If case 2 applies
		 * the double-black problem does not reappear.
		 */
		y = newY;
		
		/* Note the conditional evaluation for assigning z. */
		if(((z = y->left) && z->color == rbTreeRed) ||
		   ((z = y->right) && z->color == rbTreeRed)) 
		   {		    
		    /* Case 1:  perform a restructuring of nodes x, y, and
		     * z.
		     */
		    
		    b = restructure(t, tos, x, y, z);
		    b->color = rbTreeRed;  /* Since node x was red. */
		    b->left->color = b->right->color = rbTreeBlack;
		    }
		else 
		    {
		    /* Case 2:  recolour node y red. */

		    /* Note that node y is black and node x is red. */
		    
		    y->color = rbTreeRed;
		    x->color = rbTreeBlack;
		    }

		break;
		}
	    }
	}
    else 
	{
	/* A red node replaced the deleted black node. */

	/* In this case we can simply color the red node black. */
	r->color = rbTreeBlack;
	}
    }
return returnItem;
}

/* Some variables to help recursively dump tree. */
static int dumpLevel;	/* Indentation level. */
static FILE *dumpFile;  /* Output file */
static void (*dumpIt)(void *item, FILE *f);  /* Item dumper. */

static void rTreeDump(struct rbTreeNode *n)
/* Recursively dump. */
{
if (n == NULL)
    return;
spaceOut(dumpFile, ++dumpLevel * 3);
fprintf(dumpFile, "%c ", (n->color ==  rbTreeRed ? 'r' : 'b'));
dumpIt(n->item, dumpFile);
fprintf(dumpFile, "\n");
rTreeDump(n->left);
rTreeDump(n->right);
--dumpLevel;
}

void rbTreeDump(struct rbTree *tree, FILE *f, 
	void (*dumpItem)(void *item, FILE *f))
/* Dump out rb tree to file, mostly for debugging. */
{
dumpFile = f;
dumpLevel = 0;
dumpIt = dumpItem;
fprintf(f, "rbTreeDump\n");
rTreeDump(tree->root);
}



/* Variables to help recursively traverse tree. */
static void (*doIt)(void *item);
static void *minIt, *maxIt;
static int (*compareIt)(void *, void *);

static void rTreeTraverseRange(struct rbTreeNode *n)
/* Recursively traverse tree in range applying doIt. */
{
if (n != NULL)
   {
   int minCmp = compareIt(n->item, minIt);
   int maxCmp = compareIt(n->item, maxIt);
   if (minCmp >= 0)
       rTreeTraverseRange(n->left);
   if (minCmp >= 0 && maxCmp <= 0)
       doIt(n->item);
   if (maxCmp <= 0)
       rTreeTraverseRange(n->right);
   }
}

static void rTreeTraverse(struct rbTreeNode *n)
/* Recursively traverse full tree applying doIt. */
{
if (n != NULL)
    {
    rTreeTraverse(n->left);
    doIt(n->item);
    rTreeTraverse(n->right);
    }
}


void rbTreeTraverseRange(struct rbTree *tree, void *minItem, void *maxItem,
	void (*doItem)(void *item))
/* Apply doItem function to all items in tree such that
 * minItem <= item <= maxItem */
{
doIt = doItem;
minIt = minItem;
maxIt = maxItem;
compareIt = tree->compare;
rTreeTraverseRange(tree->root);
}

void rbTreeTraverse(struct rbTree *tree, void (*doItem)(void *item))
/* Apply doItem function to all items in tree */
{
doIt = doItem;
rTreeTraverse(tree->root);
}

struct rTreeContext
/* Context for traversing a tree when you want to be fully thread safe and reentrant. */
    {
    void *context;	/* Some context carried from user and passed to doIt. */
    void (*doItem)(void *item, void *context);
    };

static void rTreeTraverseWithContext(struct rbTreeNode *n, struct rTreeContext *context)
/* Traverse tree with a little context so don't need little static variables that
 * prevent reentrancy of callback functions. */
{
if (n != NULL)
    {
    rTreeTraverseWithContext(n->left, context);
    context->doItem(n->item, context->context);
    rTreeTraverseWithContext(n->right, context);
    }
}

void rbTreeTraverseWithContext(struct rbTree *tree, 
	void (*doItem)(void *item, void *context), void *context)
/* Traverse tree calling doItem on every item with context pointer passed through to doItem.
 * This often avoids having to declare global or static variables for the doItem callback to use. */
{
struct rTreeContext ctx;
ctx.context = context;
ctx.doItem = doItem;
rTreeTraverseWithContext(tree->root, &ctx);
}

struct slRef *itList;  /* List of items that rbTreeItemsInRange returns. */

static void addRef(void *item)
/* Add item it itList. */
{
refAdd(&itList, item);
}

struct slRef *rbTreeItemsInRange(struct rbTree *tree, void *minItem, void *maxItem)
/* Return a sorted list of references to items in tree between range.
 * slFreeList this list when done. */
{
itList = NULL;
rbTreeTraverseRange(tree, minItem, maxItem, addRef);
slReverse(&itList);
return itList;
}

static void addRefWithContext(void *item, void *context)
/* Add item it itList. */
{
struct slRef **pList = context;
refAdd(pList, item);
}


struct slRef *rbTreeItems(struct rbTree *tree)
/* Return sorted list of items.  slFreeList this when done.*/
{
struct slRef *list = NULL;
rbTreeTraverseWithContext(tree, addRefWithContext, &list);
slReverse(&list);
return list;
}

int rbTreeCmpString(void *a, void *b)
/* Set up rbTree so as to work on strings. */
{
return strcmp(a, b);
}

int rbTreeCmpWord(void *a, void *b)	
/* Set up rbTree so as to work on case-insensitive strings. */
{
return differentWord(a,b);
}