/* #include "stdafx.h" */

#define OWN_ALLOCATOR

#ifdef OWN_ALLOCATOR

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stddef.h>
/* #include <windows.h> */
#include "memmgrs.h"

#if (defined USE_BLOWFISH)	// To fix link error.
#include <tee_internal_api.h>

#ifndef malloc
#define malloc(x) TEE_Malloc(x, 0)
#endif	// End of malloc

#ifndef free
#define free(x) TEE_Free(x)
#endif	// End of free
#endif	// End of USE_BLOWFISH

/* For free function */
const int NONE_FREE = 0;
const int LEFT_FREE = 1;
const int RIGHT_FREE = 2;
const int BOTH_FREE = 3;

/* Block free/occupied flag; */
const int FREE_BLOCK = 0;
const int OCCUPIED_BLOCK = 1;

const unsigned int CLEAR_MASK = 0xFFFFFFFE;
const unsigned int SET_MASK = 1;

static Memmngr mngr;
static int mngr_initialized = 0;

static __inline Header * pth(void * p) /* pointer to header */
{
	return (Header*)((char*)p - sizeof(Header));
}
static __inline void * htp(Header* h) /* header to pointer */
{
	return ((char*)h + sizeof(Header));
}

static __inline Header * nth(RBNode * n) /* node to header */
{
	return (Header*)((char*)n - sizeof(void*));
}

static __inline RBNode * htn(Header * h)
{
	return (RBNode*)((char*)h + sizeof(void*));
}

static __inline void Header_setbit(Header * h, const int i)
{
	if (!(i % 2)) *((int*)h) =*((int*)h) & CLEAR_MASK;
	else *((int*)h) =*((int*)h) | SET_MASK;
}

static __inline int Header_getbit(Header * h)
{
	if (*((int*)h) % 2) return 1;
	return 0;
}

static __inline int inborder(const void *p)
{
	return (p >= mngr.beg_ptr) && (p < mngr.end_ptr);
}

static void Header_setprev(Header * h, void * p)
{
	int i;

	if (!h) return;
	i = Header_getbit(h);
	h->p = p;
	Header_setbit(h, i);
}

static void * Header_getprev(Header * h)
{
	void * res;
	void ** pr;

	if (!h) return NULL;
	res = h->p;
	pr = &res;
	*((int*)pr) =*((int*)pr) & CLEAR_MASK;
	return res;
}

static __inline Header * Header_next_header(Header * h)
{
	if (!h) return NULL;
	return (Header*)((char*)h + sizeof(Header) + h->size);
}


/*/////////////////////RBTree////////////////////////////////// */

void RBTree_init(RBTree * tree)
{
	if(!tree) return;
	tree->root = NULL;
}

static void RBTree_LeftRotate(RBTree * tree, RBNode * x)
{
	RBNode * y , * temp;

	if(!(tree) || !(x))
		return;
	temp = NULL;
	y = x->right;
	x->right=y->left;
	if (y->left) y->left->parent=x;

	temp = x->parent;
	y->parent=temp;

	if (!temp) tree->root = y;
	else
	{
		if (temp->left == x) temp->left=y;
		else temp->right=y;
	}

	y->left=x;
	x->parent=y;
}

static void RBTree_RightRotate(RBTree * tree, RBNode * x)
{
	RBNode * y , * temp;

	if(!(tree) || !(x))
		return;
	temp = NULL;
	y = x->left;
	x->left=y->right;
	if (y->right) y->right->parent=x;

	temp = x->parent;
	y->parent=temp;

	if (!temp) tree->root = y;
	else if (temp->left == x) temp->left=y;
	else temp->right=y;

	y->right=x;
	x->parent=y;
}

static void RBTree_InsFixup(RBTree * tree, RBNode * z)
{
	RBNode * y;

	if(!(tree) || !(z))
		return;
	while ((z->parent) && (z->parent->color == RED))
	{
		if(z->parent == tree->root) tree->root->color=BLACK;
		else
			if ((z->parent) == (z->parent->parent->left))
		{
			y = z->parent->parent->right;
			if ((y) && (y->color == RED))
			{
				y->color=BLACK;
				z->parent->color=BLACK;
				z->parent->parent->color=RED;
				z = z->parent->parent;
			}
			else
			{
				if (z == z->parent->right)
				{
					z = z->parent;
					RBTree_LeftRotate(tree, z);
				}
				z->parent->color=BLACK;
				z->parent->parent->color=RED;
				RBTree_RightRotate(tree, z->parent->parent);
			}
		}
		else
		{
			y = z->parent->parent->left;
			if ((y) && (y->color == RED))
			{
				y->color=BLACK;
				z->parent->color=BLACK;
				z->parent->parent->color=RED;
				z = z->parent->parent;
			}
			else
			{
				if (z == z->parent->left)
				{
					z = z->parent;
					RBTree_RightRotate(tree, z);
				}
				z->parent->color=BLACK;
				z->parent->parent->color=RED;
				RBTree_LeftRotate(tree, z->parent->parent);
			}
		}
	}
	tree->root->color=BLACK;
}

static void RBTree_RBInsert(RBTree * tree, RBNode * z)
{
	RBNode *x, *y;

	if(!(tree) || !(z))
		return;
	x = tree->root;
	y = NULL;

	z->color=RED;
	z->left=NULL;
	z->right=NULL;
	z->parent=NULL;
	z->next=NULL;

	while (x)
	{
		y = x;
		if ((z->key) == (y->key))
		{
			z->next=y->next;
			y->next=z;
			return; /* goto endf; */
		}
		if ((z->key) < (y->key)) x = y->left;
		else x = y->right;
	}
	z->parent=y;
	if (!y)
	{
		tree->root = z;
		z->color=BLACK;
	}
	else
	{
		if (z->key < y->key) y->left=z;
		else y->right=z;
	}
	RBTree_InsFixup(tree, z);
/* endf: ; */
}

static void RBTree_RBDeleteFixup(RBTree * tree, RBNode *x)
{
	RBNode * w;

	if(!(tree) || !(x))
		return;
	while ((x != tree->root) && (x->color == BLACK))
	{
		if (x == (x->parent->left))
		{
			w = x->parent->right;
			if (w->color == RED)
			{
				w->color=BLACK;
				x->parent->color=RED;
				RBTree_LeftRotate(tree, x->parent);
				w = x->parent->right;
			}
			if (((!(w->left)) || (w->left->color == BLACK))
				&& ((!(w->right)) || (w->right->color == BLACK)))
			{
				w->color=RED;
				x = x->parent;
			}
			else
			{
				if ((!(w->right)) ||(w->right->color == BLACK))
				{
					w->left->color=BLACK;
					w->color=RED;
					RBTree_RightRotate(tree, w);
					w = x->parent->right;
				}
				w->color=x->parent->color;
				x->parent->color=BLACK;
				w->right->color=BLACK;
				RBTree_LeftRotate(tree, x->parent);
				x = tree->root;
			}
		}
		else
		{
			w = x->parent->left;
			if (w->color == RED)
			{
				w->color= BLACK;
				x->parent->color=RED;
				RBTree_RightRotate(tree, x->parent);
				w = x->parent->left;
			}
			if (((!(w->left)) || (w->left->color == BLACK))
				&& ((!(w->right)) || (w->right->color == BLACK)))
			{
				w->color=RED;
				x = x->parent;
			}
			else
			{
				if ((!(w->left)) || (w->left->color == BLACK))
				{
					w->right->color=BLACK;
					w->color=RED;
					RBTree_LeftRotate(tree, w);
					w = x->parent->left;
				}
				w->color=x->parent->color;
				x->parent->color=BLACK;
				w->left->color=BLACK;
				RBTree_RightRotate(tree, x->parent);
				x = tree->root;
			}
		}
	}
	x->color=BLACK;
}

static void RBTree_RBDelete(RBTree * tree, RBNode *z)
{
	RBNode *x, *y;
	Color c;
	int f;

	if(!(tree) || !(z))
		return;
	f = 0;
	if ((!(z->left)) || (!(z->right))) y = z;
	else
	{
		y = z->right;
		while (y->left) y = y->left;
	}
	if ((!(y->left)) && (!(y->right)))
	{
		y->left=&(tree->tempnode);
		y->left->color=BLACK;
		y->left->parent=y;
		f = 1;
	}
	if (y->left) x = y->left;
	else x = y->right;

	x->parent=y->parent;
	if (!(y->parent)) tree->root = x;
	else
	{
		if (y == y->parent->left) y->parent->left=x;
		else y->parent->right=x;
	}
	c = y->color;
	if (z != y)
	{
		if (z->parent)
		{
			if (z == z->parent->left) z->parent->left=y;
			else z->parent->right=y;
		}
		else tree->root = y;

		y->color=z->color;
		y->left=z->left;
		y->right=z->right;
		y->parent=z->parent;
		if (z->left) z->left->parent=y;
		if (z->right) z->right->parent=y;
	}
	if (c == BLACK) RBTree_RBDeleteFixup(tree, x);
	if (f)
	{
		if (x->parent)
			if (x == x->parent->left) x->parent->left=NULL;
			else x->parent->right=NULL;
		else tree->root = NULL;
	}
}

static RBNode * RBTree_find_node(RBTree * tree, size_t size)
{
	RBNode *temp1, *temp2;

	if(!(tree) || (size == 0))
		return NULL;
	temp2 = NULL;
	temp1 = tree->root;
	if (tree->root == NULL) return NULL;
	while(1)
	{
		if (size == temp1->key) break;
		if (size > temp1->key) temp2 = temp1->right;
		else temp2 = temp1->left;
		if (!temp2) break;
		temp1 = temp2;
	}
	while (temp1)
	{
		if (temp1->key >= size) break;
		temp1 = temp1->parent;
	}
	return temp1;
}

static RBNode * RBTree_find_prec(RBTree * tree, size_t size)
{
	RBNode *temp1, *temp2;

	if(!(tree) || (size == 0))
		return NULL;
	temp2 = NULL;
	temp1 = tree->root;
	while(temp1)
	{
		if (temp1->key == size)
		{
			temp2 = temp1;
			break;
		}
		if (size > temp1->key) temp1 = temp1->right;
		else temp1 = temp1->left;
	}
	return temp2;
}

static void RBTree_Delete(RBTree * tree, RBNode * n)
{
	unsigned int s;
	RBNode /* *tn1, */*tn;

	if(!(tree) || !(n))
		return;
	s = n->key;
	if((tn = RBTree_find_prec(tree, s)) == NULL)
		return;
	if ((tn->next) && (tn == n))
	{
		if (tn != tree->root)
		{
			if (tn->parent->left == tn)	tn->parent->left = tn->next;
			else tn->parent->right= tn->next;
		}
		else
		{
			tree->root = tn->next;
		}
		tn->next->left=tn->left;
		if (tn->left) tn->left->parent=tn->next;
		tn->next->right=tn->right;
		if (tn->right) tn->right->parent=tn->next;
		tn->next->parent=tn->parent;
		tn->next->color=tn->color;
		return;
	}
	if ((!(tn->next)) && (tn == n))
	{
		RBTree_RBDelete(tree, n);
		return;
	}
	/* tn1 = tn; */
	while (tn->next != n && tn->next != NULL) tn = tn->next;
	tn->next=n->next;
}

int memmngr_init(void * start, size_t size)
{
	Header * ph;
	RBNode * pn;

	mngr_initialized = 0;
	if (!start || (((ptrdiff_t)start) % MALLOC_ALIGN) || !size || (size < MEMMNGR_INIT_MIN_SIZE) || (size > MEMMNGR_INIT_MAX_SIZE)) 
	{
		return 0;
	}

	mngr.bp = start;
	mngr.end_ptr = (char*)(mngr.bp) + size;
	mngr.beg_ptr = (char*)(mngr.bp);
	ph = (Header*)(mngr.bp);
	pn = htn(ph);
	mngr.fh = ph;
	Header_setprev(ph, NULL);
	Header_setbit(ph, FREE_BLOCK);
	ph->size = size - sizeof(Header);
	pn ->color=(BLACK);
	pn->left=(NULL);
	pn->right=(NULL);
	pn->parent=(NULL);

	pn->next=(NULL);
	RBTree_init(&(mngr.tree));
	RBTree_RBInsert(&(mngr.tree), pn);

	mngr_initialized = 1;

	return 1;
}

void mem_dump()
{
/*
	Header * h;
	int f;
	const char * s;
	void * p;

	h = mngr.fh;
	printf("%s \n", "BEGIN DUMP");
	while( inborder(h))
	{
		f = Header_getbit(h);
		s = ((f) ? "OCCUPIED " : "FREE  \t  ");
		p = htp(h);
		printf(" Address 0x%p \t Size: 0x%x\t %s BLOCK \n", p, (unsigned)(h->size), s);
		h = Header_next_header(h);
	}
	printf("%s \n\n", "END DUMP");
*/
}

void * mem_alloc(size_t size)
{
	RBNode *node;

	unsigned int s;
	Header *ph;
	Header *th, *th1;
	RBNode *n;

	if(!mngr_initialized) return NULL;
	if (size < MALLOC_MIN_ALLOC_SIZE) size = MALLOC_MIN_ALLOC_SIZE;
	if (size % MALLOC_ALIGN) size+= (MALLOC_ALIGN - (size % MALLOC_ALIGN));

	node = RBTree_find_node(&(mngr.tree), size);
	if (!node) return NULL;
	RBTree_Delete(&(mngr.tree), node);
	ph = nth(node);

	/* printf("Allocating memory: %x\n", size); */

	Header_setbit(ph, OCCUPIED_BLOCK);
	if ((ph->size - size) >= (sizeof(RBNode) + sizeof(void*)))
	{
		s = ph->size - size - sizeof(Header);
		ph->size = size;
		/* free header */
		th =  Header_next_header(ph);
		n = htn(th);
		th->size = s;
		Header_setprev(th, htp(ph));
		Header_setbit(th, FREE_BLOCK);
		RBTree_RBInsert(&(mngr.tree), n);
		/* header */
		th1 = Header_next_header(th);
		if (inborder(th1)) Header_setprev(th1, htp(th));

	}
	/* if ((unsigned long)htp(ph) == 0x717b28) printf("Allocating memory: %p, size%x\n", htp(ph), size); */
	return htp(ph);
}

void mem_free(void * addr)
{
	Header *h, *th, *rh;
	RBNode *n;
	unsigned int s;
	ptrdiff_t i;

	if(!mngr_initialized) return;
	if (!inborder(addr))
	{
		return;
	}

	h = pth(addr);
	s = Header_getbit(h);
	i = (ptrdiff_t) addr;
	if (!s || (i % MALLOC_ALIGN))
	{
		return;
	}
	th = pth(Header_getprev(h));
	if (inborder(th) && (Header_getbit(th) == FREE_BLOCK))
	{
		n = htn(th);
		RBTree_Delete(&(mngr.tree), n);
		th->size += sizeof(Header) + h->size;
		h = th;
	}
	th = Header_next_header(h);
	if (inborder(th))
	{
		if (Header_getbit(th) == FREE_BLOCK)
		{
			n = htn(th);
			rh = Header_next_header(th);
			RBTree_Delete(&(mngr.tree), n);
			if (inborder(rh)) Header_setprev(rh, htp(h));
			h->size += sizeof(Header) + th->size;
		}
		else
		{
			Header_setprev(th, htp(h));
		}
	}
	Header_setbit(h, FREE_BLOCK);
	n = htn(h);
	RBTree_RBInsert(&(mngr.tree), n);
}

void * mem_realloc(void * addr, size_t size)
{
	Header * h, * ht;
	void * p;
	size_t n;
	unsigned i;

	if(!mngr_initialized) return NULL;
	if (!addr) return malloc(size);
	h = pth(addr);
	if ((h->size == size) || !(inborder(addr))) return NULL;
	p = malloc(size);
	if (!p) return NULL;
	ht = pth(p);
	n = (ht->size > h->size) ? (h->size / MALLOC_ALIGN) : (ht->size / MALLOC_ALIGN);
	for (i = 0; i < n; i++)	((MALLOC_ALIGN_TYPE * )p)[i] = ((MALLOC_ALIGN_TYPE * )addr)[i];
	free(addr);
	return p;
}

int is_ptr_in_memmngr_range(const void* ptr)
{
	return inborder(ptr);
}
#endif /* OWN_ALLOCATOR */