// Renamed from Map to Tree to avoid conflict with libmach.

/*
Copyright (c) 2003-2007 Russ Cox, Tom Bergan, Austin Clements,
	Massachusetts Institute of Technology
Portions Copyright (c) 2009 The Go Authors. All rights reserved.

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

// Mutable map structure, but still based on
// Okasaki, Red Black Trees in a Functional Setting, JFP 1999,
// which is a lot easier than the traditional red-black
// and plenty fast enough for me.  (Also I could copy
// and edit fmap.c.)

#include <u.h>
#include <libc.h>
#include "tree.h"

enum
{
	Red = 0,
	Black = 1
};


// Red-black trees are binary trees with this property:
//	1. No red node has a red parent.
//	2. Every path from the root to a leaf contains the
//		same number of black nodes.

static TreeNode*
rwTreeNode(TreeNode *p, int color, TreeNode *left, void *key, void *value, TreeNode *right)
{
	if(p == nil)
		p = malloc(sizeof *p);
	if(p == nil)
		sysfatal("out of memory");
	p->color = color;
	p->left = left;
	p->key = key;
	p->value = value;
	p->right = right;
	return p;
}

static TreeNode*
balance(TreeNode *m0)
{
	void *xk, *xv, *yk, *yv, *zk, *zv;
	TreeNode *a, *b, *c, *d;
	TreeNode *m1, *m2;
	int color;
	TreeNode *left, *right;
	void *key, *value;

	color = m0->color;
	left = m0->left;
	key = m0->key;
	value = m0->value;
	right = m0->right;

	// Okasaki notation: (T is mkTreeNode, B is Black, R is Red, x, y, z are key-value.
	//
	// balance B (T R (T R a x b) y c) z d
	// balance B (T R a x (T R b y c)) z d
	// balance B a x (T R (T R b y c) z d)
	// balance B a x (T R b y (T R c z d))
	//
	//     = T R (T B a x b) y (T B c z d)

	if(color == Black){
		if(left && left->color == Red){
			if(left->left && left->left->color == Red){
				a = left->left->left;
				xk = left->left->key;
				xv = left->left->value;
				b = left->left->right;
				yk = left->key;
				yv = left->value;
				c = left->right;
				zk = key;
				zv = value;
				d = right;
				m1 = left;
				m2 = left->left;
				goto hard;
			}else if(left->right && left->right->color == Red){
				a = left->left;
				xk = left->key;
				xv = left->value;
				b = left->right->left;
				yk = left->right->key;
				yv = left->right->value;
				c = left->right->right;
				zk = key;
				zv = value;
				d = right;
				m1 = left;
				m2 = left->right;
				goto hard;
			}
		}else if(right && right->color == Red){
			if(right->left && right->left->color == Red){
				a = left;
				xk = key;
				xv = value;
				b = right->left->left;
				yk = right->left->key;
				yv = right->left->value;
				c = right->left->right;
				zk = right->key;
				zv = right->value;
				d = right->right;
				m1 = right;
				m2 = right->left;
				goto hard;
			}else if(right->right && right->right->color == Red){
				a = left;
				xk = key;
				xv = value;
				b = right->left;
				yk = right->key;
				yv = right->value;
				c = right->right->left;
				zk = right->right->key;
				zv = right->right->value;
				d = right->right->right;
				m1 = right;
				m2 = right->right;
				goto hard;
			}
		}
	}
	return rwTreeNode(m0, color, left, key, value, right);

hard:
	return rwTreeNode(m0, Red, rwTreeNode(m1, Black, a, xk, xv, b),
		yk, yv, rwTreeNode(m2, Black, c, zk, zv, d));
}

static TreeNode*
ins0(TreeNode *p, void *k, void *v, TreeNode *rw)
{
	if(p == nil)
		return rwTreeNode(rw, Red, nil, k, v, nil);
	if(p->key == k){
		if(rw)
			return rwTreeNode(rw, p->color, p->left, k, v, p->right);
		p->value = v;
		return p;
	}
	if(p->key < k)
		p->left = ins0(p->left, k, v, rw);
	else
		p->right = ins0(p->right, k, v, rw);
	return balance(p);
}

static TreeNode*
ins1(Tree *m, TreeNode *p, void *k, void *v, TreeNode *rw)
{
	int i;

	if(p == nil)
		return rwTreeNode(rw, Red, nil, k, v, nil);
	i = m->cmp(p->key, k);
	if(i == 0){
		if(rw)
			return rwTreeNode(rw, p->color, p->left, k, v, p->right);
		p->value = v;
		return p;
	}
	if(i < 0)
		p->left = ins1(m, p->left, k, v, rw);
	else
		p->right = ins1(m, p->right, k, v, rw);
	return balance(p);
}

void
treeputelem(Tree *m, void *key, void *val, TreeNode *rw)
{
	if(m->cmp)
		m->root = ins1(m, m->root, key, val, rw);
	else
		m->root = ins0(m->root, key, val, rw);
}

void
treeput(Tree *m, void *key, void *val)
{
	treeputelem(m, key, val, nil);
}

void*
treeget(Tree *m, void *key)
{
	int i;
	TreeNode *p;

	p = m->root;
	if(m->cmp){
		for(;;){
			if(p == nil)
				return nil;
			i = m->cmp(p->key, key);
			if(i < 0)
				p = p->left;
			else if(i > 0)
				p = p->right;
			else
				return p->value;
		}
	}else{
		for(;;){
			if(p == nil)
				return nil;
			if(p->key == key)
				return p->value;
			if(p->key < key)
				p = p->left;
			else
				p = p->right;
		}
	}
}