/*
  Very fast and accurate memory allocator
  
  reserves memory via standard malloc
  
  Copyright (C) 1998 Valery Shchedrin
*/

#include <stdio.h>
#include <stdlib.h>
/* not needed for malloc() #include <malloc.h> */
#include <string.h>

#include "config.h"

/*#define DEB(x) x */
#define DEB(x)
#define BIG_ALIGN	4   /* big block align */
#define HUNK_ALIGN	4   /* must be pow of 2 */
#define HUNK_THRESHOLD	4   /* must be pow of 2 */
#define HUNK_MAXHUNKS	512
#define HUNK_MAXSIZE	2000
#define ALIGN(size,align) (((size)+(align)-1)&(~((align)-1)))
#define RESERVE_AMOUNT	11000*4 /* reserve 64*4 kb */

#define BT_NONE   0 /* initial value */
#define BT_BUSY   1 /* busy block */
#define BT_HUNKED 2 /* hunked blocks */
#define BT_PSTART 4 /* start of new block */

#define REGISTER register
#define VOLATILE volatile
#define INLINE inline

typedef struct Node_s Node_t;
typedef struct Root_s Root_t;
typedef struct Block_s Block_t;
typedef struct Hunk_s Hunk_t;

struct Node_s {
	int bal;       /* balance */
	Node_t *l, *r; /* left & right subtrees */
};
struct Root_s {
	int offs; /* offset of Node_t in structure */
	Node_t *root; /* pointer to root node */
	int comp_type; /* 1 if comparing <ptrs>, 0 if comparing <free_mem> */
};
struct Block_s {
	char *ptr; /* pointer to data */
	int size;  /* total size of block */
	int btype; /* type of block (BT_xx) */
	Block_t *next; /* next in free_mem (linked list) */
	Node_t free_mem_node, ptrs_node; /* nodes in AVL trees */
};
struct Hunk_s { /* header of hunkblock */
	Block_t b;
	int size, total_hunks, used_hunks;
	Hunk_t *size_next, *free_next;
};

static Root_t free_mem, ptrs;
static Hunk_t *size_h[HUNK_MAXSIZE], *free_h[HUNK_MAXSIZE];
static Block_t tmp_block;

/* balances partially balanced AVL tree */
/* returns 1 if tree height changed 0 otherwise */
static INLINE int balance(Node_t **root)
{
	REGISTER Node_t *p = *root, *pp;
	REGISTER int ht_changed;
	
	if (p->bal < -1) {
		if (p->l->bal == 1) {
			pp = p->l; *root = p->l->r; p->l = (*root)->r;
			(*root)->r = p; pp->r = (*root)->l;
			p = *root; p->l = pp;
			if (p->bal > 0) p->l->bal = -p->bal; else p->l->bal = 0;
			if (p->bal < 0) p->r->bal = -p->bal; else p->r->bal = 0;
			p->bal = 0;
			return 1;
		} else {
			ht_changed = (p->l->bal)?1:0;
			*root = p->l;
			p->l = (*root)->r; (*root)->r = p;
			p->bal = - (++((*root)->bal));
			return ht_changed;
		}
	} else if (p->bal > 1) {
		if (p->r->bal == -1) {
			pp = p->r; *root = p->r->l; p->r = (*root)->l;
			(*root)->l = p; pp->l = (*root)->r;
			p = *root; p->r = pp;
			if (p->bal > 0) p->l->bal = -p->bal; else p->l->bal = 0;
			if (p->bal < 0) p->r->bal = -p->bal; else p->r->bal = 0;
			p->bal = 0;
			return 1;
		} else {
			ht_changed = (p->r->bal)?1:0;
			*root = p->r;
			p->r = (*root)->l; (*root)->l = p;
			p->bal = - (--((*root)->bal));
			return ht_changed;
		}
	}
	return 0;
}

/* INTERNAL insertion/deletion routines */
static Root_t *cur_tree;
static Node_t *fix_node, *new_node;

#define NODE_DATA(t) (((char*)(t))-cur_tree->offs)

/* rather optimal compare function */
#define COMPARE(var, ctp, a, b) \
	if (ctp) { \
		var = ((char*)(a)) - ((char*)(b));\
	} else { \
		var = ((Block_t*)(a))->size - ((Block_t*)(b))->size; \
	}

static int r_delfix(Node_t **root);
static int r_insert(Node_t **root) /* recursive insert */
{
	REGISTER int i; /* height increase */
	
	/* insert prepared `new_node` */
	if (!*root) { *root = new_node; new_node = NULL; return 1; }
	
	COMPARE(i, cur_tree->comp_type, NODE_DATA(new_node), NODE_DATA(*root));
	
	if (i < 0) { /* insert into the left subtree */
		i = -r_insert(&((*root)->l));
		if (new_node != NULL) return 0; /* already there */
	} else if (i > 0) { /* insert into the right subtree */
		i = r_insert(&((*root)->r));
		if (new_node != NULL) return 0; /* already there */
	} else { /* found */
		new_node = *root;
		return 0;
	}
	if (!i) return 0;
	(*root)->bal += i; /* update balance factor */
	return ((*root)->bal)?(1-balance(root)):0;
}
static int r_delete(Node_t **root) /* recursive delete */
{
	REGISTER int i; /* height decrease */
	
	if (!*root) { new_node = NULL; return 0; } /* data not found */

	COMPARE(i, cur_tree->comp_type, NODE_DATA(new_node), NODE_DATA(*root));
	
	if (i < 0) {        i = -r_delete(&((*root)->l));
	} else if (i > 0) { i =  r_delete(&((*root)->r));
	} else {
		new_node = *root;
		if (!new_node->l) {
			*root = new_node->r; return 1;
		} else {
			if (!new_node->r) { *root = new_node->l; return 1; }
			i = r_delfix(&((*root)->r));
			fix_node->l = new_node->l;
			fix_node->r = new_node->r;
			fix_node->bal = new_node->bal;
			*root = fix_node;
		}
	}
	if (!i) return 0;
	(*root)->bal -= i;
	if ((*root)->bal) return balance(root);
	return 1;
}
static int r_delfix(Node_t **root)
{
	REGISTER int i; /* height increase */
	
	if (!(*root)->l) {
		fix_node = *root;
		*root = fix_node->r;
		return 1;
	}
	i = -r_delfix(&((*root)->l));
	if (!i) return 0;
	(*root)->bal -= i;
	if ((*root)->bal) return balance(root);
	return 1;
}
static INLINE void *avl_ins(Root_t *t, void *data) /* insert node */
{
	cur_tree = t; new_node = (Node_t*)(((char*)data)+t->offs);
	memset(new_node, 0, sizeof(Node_t));
	r_insert(&(cur_tree->root));
	if (!new_node) return NULL;
	return (void*)(((char*)new_node) - t->offs);
}
static INLINE int avl_del(Root_t *t, void *data) /* delete node */
{
	cur_tree = t; new_node = (Node_t*)(((char*)data)+t->offs);
	r_delete(&(cur_tree->root));
	if (!new_node) return 0;
	return 1;
}
static INLINE void avl_replace(Root_t *t, void *data) /* replace node */
{
	REGISTER Node_t **p = &t->root;
	REGISTER int cmp;
	while (*p) {
		COMPARE(cmp, t->comp_type, data, ((char*)*p) - t->offs);
		if (cmp < 0) p = &((*p)->l);
		else if (cmp > 0) p = &((*p)->r);
		else {
			memcpy(((char*)data)+t->offs, *p, sizeof(Node_t));
			*p = (Node_t*)(((char*)data)+t->offs);
			return;
		}
	}
}

static Block_t *bl_alloc(size_t size)
{
	REGISTER Node_t *nd;
	REGISTER Block_t *p, *pp;
	
	size = ALIGN(size, BIG_ALIGN);

	for (nd=free_mem.root, pp=NULL;nd;) {
		p = (Block_t*)(((char*)nd) - free_mem.offs);
		if (p->size > size) {
			pp = p; nd = nd->l;
		} else if (p->size < size) nd = nd->r;
		  else { pp = p; break; }
	}
	if (!pp) { /* try to allocate block large enough via malloc */
		pp = (Block_t *)malloc(RESERVE_AMOUNT + size);
		DEB(printf(" reserved new block 0x%x (%d)\n", (unsigned)pp, RESERVE_AMOUNT + size));
		if (!pp) return NULL; /* unable to allocate */
		memset(pp, 0, sizeof(Block_t));
		pp->ptr = (char*)ALIGN((unsigned long)(pp+1), BIG_ALIGN);
		pp->size = ALIGN(RESERVE_AMOUNT + size - (pp->ptr - ((char*)pp)) - BIG_ALIGN + 1, BIG_ALIGN);
		pp->btype |= BT_PSTART;
		avl_ins(&ptrs, pp);
	} else {
		if (pp->next) avl_replace(&free_mem, pp->next);
		  else avl_del(&free_mem, pp);
	}
	
	if (pp->size-size >= BIG_ALIGN+ALIGN(sizeof(Block_t), BIG_ALIGN)) {
		REGISTER Block_t *bt; /* split large block */
		
		bt = (Block_t *)(pp->ptr + size);
		memset(bt, 0, sizeof(Block_t));
		bt->ptr = ((char*)bt) + ALIGN(sizeof(Block_t),BIG_ALIGN);
		bt->size = (pp->ptr + pp->size) - bt->ptr;
		DEB(printf(" splitted block 0x%x inserted (%d)\n", (unsigned)bt, bt->size));
		avl_ins(&ptrs, bt);
		p = (Block_t *)avl_ins(&free_mem, bt);
		if (p) { bt->next = p->next; p->next = bt; }
		else bt->next = NULL;
		pp->size = size;
	}
	
	pp->btype |= BT_BUSY;
	DEB(printf(" allocated block 0x%x (%d) - %d\n", (unsigned)pp, pp->size, pp->btype));
	DEB(ptrs_prn(ptrs.root)); DEB(fr_prn(free_mem.root));
	return pp;
}

static void ptrs_prn(Node_t *p)
{
	Block_t *b;
	
	b = (Block_t*)(((char*)p) - ptrs.offs);
	if (b->ptrs_node.l) ptrs_prn(b->ptrs_node.l);
	printf(" PTR 0x%x size %d btype %d\n", (unsigned)b, b->size, b->btype);
	if (b->ptrs_node.r) ptrs_prn(b->ptrs_node.r);
}

static void fr1_prn(Block_t *p)
{
	printf(" TFR 0x%x size %d btype %d\n", (unsigned)p, p->size, p->btype);
	if (p->next) fr1_prn(p->next);
}

static void fr_prn(Node_t *p)
{
	Block_t *b;
	
	b = (Block_t*)(((char*)p) - free_mem.offs);
	if (b->free_mem_node.l) fr_prn(b->free_mem_node.l);
	fr1_prn(b);
	if (b->free_mem_node.r) fr_prn(b->free_mem_node.r);
}

static void bl_free(Block_t *b)
{
	REGISTER Node_t *nd;
	REGISTER Block_t *p, *pp;
	REGISTER Block_t *bl_bef;

	/* Look for block before & after `b` */
	for (nd=ptrs.root,pp=NULL,bl_bef=NULL;nd;) {
		p = (Block_t*)(((char*)nd) - ptrs.offs);
		if (((char*)p) > ((char*)b)) {
			pp = p; nd = nd->l;
		} else if (((char*)p) < ((char*)b)) {
			bl_bef = p; nd = nd->r;
		} else break;
	}
	/* Do this if `b` has subtrees */
	nd = &b->ptrs_node;
	if (nd->l) {
		for (nd = nd->l; nd->r; nd=nd->r);
		bl_bef = (Block_t*)(((char*)nd) - ptrs.offs);
		nd = &b->ptrs_node;
	}
	if (nd->r) {
		for (nd = nd->r; nd->l; nd=nd->l);
		pp = (Block_t*)(((char*)nd) - ptrs.offs);
	}
	DEB(printf(" freeing block 0x%x (%d) - %d, prev 0x%x (%d), next 0x%x (%d)\n",
		(unsigned)b, b->size, b->btype,
		(unsigned)bl_bef, bl_bef?bl_bef->size:0,
		(unsigned)pp, pp?pp->size:0));
	DEB(ptrs_prn(ptrs.root));
	DEB(fr_prn(free_mem.root));
	/* Combine with block just after */
	if (pp && !(pp->btype & (BT_PSTART|BT_BUSY))) {
		avl_del(&ptrs, pp);
		for (nd=free_mem.root,p=NULL;nd;) {
			p = (Block_t*)(((char*)nd) - free_mem.offs);
			if (p->size > pp->size) nd = nd->l;
			  else if (p->size < pp->size) nd = nd->r;
			  else break;
		}
		if (!nd) {
			printf(" xmalloc failure - bad afterblock\n");
			exit(-1);
		}
		if (p == pp) {
			if (pp->next) avl_replace(&free_mem, pp->next);
			else avl_del(&free_mem, pp);
		} else {
			for (;p->next != pp; p = p->next);
			p->next = pp->next;
		}
		b->size = (pp->ptr + pp->size) - b->ptr;
	}
	/* Combine with free block just before */
	if (bl_bef && !(bl_bef->btype & BT_BUSY) && !(b->btype & BT_PSTART)) {
		pp = bl_bef;
		avl_del(&ptrs, b);
		for (nd=free_mem.root,p=NULL;nd;) {
			p = (Block_t*)(((char*)nd) - free_mem.offs);
			if (p->size > pp->size) nd = nd->l;
			  else if (p->size < pp->size) nd = nd->r;
			  else break;
		}
		if (!nd) {
			printf(" xmalloc failure - bad before block\n");
			exit(-1);
		}
		if (p == pp) {
			if (pp->next) avl_replace(&free_mem, pp->next);
			else avl_del(&free_mem, pp);
		} else {
			for(;p->next != pp; p = p->next);
			p->next = pp->next;
		}
		pp->size = (b->ptr + b->size) - pp->ptr;
		b = pp;
	}
	b->btype &= ~BT_BUSY;
	DEB(printf(" after combination:\n"));
	DEB(ptrs_prn(ptrs.root));
	DEB(fflush(stdout));
	if (b->btype & BT_PSTART) { /* free unused blocks */
		nd = &b->ptrs_node;
		if (nd->r) {
			for (nd = nd->r; nd->l; nd=nd->l);
			pp = (Block_t *)(((char*)nd) - ptrs.offs);
		} else {
			for (nd = ptrs.root,pp=NULL;nd;) {
				p = (Block_t*)(((char*)nd)-ptrs.offs);
				if (((char *)p) > ((char *)b)) {
					pp = p; nd = nd->l;
				} else if (((char *)p) < ((char *)b)) {
					nd = nd->r;
				} else break;
			}
		}
		if (!pp || (pp->btype & BT_PSTART)) {
			avl_del(&ptrs, b);
			DEB(printf(" disposing block 0x%x\n", (unsigned)b));
			free(b);
			return;
		}
	}
	
	/* Put freed element into the tree */
	pp = (Block_t *)avl_ins(&free_mem, b);
	if (pp) { b->next = pp->next; pp->next = b; }
	else b->next = NULL;
}

#define usagemap(h) (((unsigned char *)(h))+sizeof(Hunk_t))

static void *hunk_block_alloc(Hunk_t *h)
{
	REGISTER unsigned long *cpl;
	REGISTER int i,c;
	
	/* Locate free point in usagemap */
	for (cpl=(unsigned long*)usagemap(h);*cpl==0xFFFFFFFF;cpl++);
	i = ((unsigned char *)cpl) - usagemap(h);
	if (*(unsigned short*)cpl != 0xFFFF) {
		if (*(unsigned char*)cpl == 0xFF) {
			c = *(int*)(((unsigned char *)cpl)+1); i++;
		} else  c = *(int*)(unsigned char *)cpl;
	} else {
		i+=2; c = *(((unsigned char *)cpl)+2);
		if (c == 0xFF) { c = *(int*)(((unsigned char *)cpl)+3); i++; }
	}
	i *= 8;
	if ((c & 0xF) == 0xF) { c >>= 4; i+=4; }
	if ((c & 0x3) == 0x3) { c >>= 2; i+=2; }
	if (c & 1) i++;

	usagemap(h)[i/8] |= (1 << (i&7)); /* Set bit */
	
	/* Increment counter and update hashes */
	if (++h->used_hunks == h->total_hunks) free_h[h->size] = h->free_next;

	return h->b.ptr+(i*h->size);
}

static void *hunk_alloc(int size)
{
	REGISTER Hunk_t *p;
	
	if (size >= HUNK_THRESHOLD) size = ALIGN(size, HUNK_ALIGN);
	
	/* Look for already allocated partially free Hunk_t */
	p = free_h[size];
	if (!p) { /* try to allocate new one */
		REGISTER int bound, hunks;
		
		/* Set hunks to allocate */
		if (size_h[size]) {
			hunks = size_h[size]->total_hunks*2;
			if (hunks > HUNK_MAXHUNKS) hunks = HUNK_MAXHUNKS;
		} else {
			bound = sizeof(Hunk_t)+4;
			hunks = (ALIGN(bound,BIG_ALIGN) - bound)/size;
			if (!hunks) hunks = 1;
		}
		
		/* Allocate new Hunk_t */
		bound = ALIGN((hunks+7)/8, 4);
		p = (Hunk_t*)bl_alloc(sizeof(Hunk_t)+bound+size*hunks);
		if (!p) return NULL; /* failed */
		
		/* Initialize new Hunk_t */
		p->total_hunks = hunks; p->used_hunks = 0; p->size = size;
		memset(usagemap(p), 0, bound);
		p->b.btype |= BT_HUNKED;
		p->b.ptr = usagemap(p)+bound;
		p->free_next = free_h[size]; free_h[size] = p;
		p->size_next = size_h[size]; size_h[size] = p;
	}
	return hunk_block_alloc(p);
}

static void hunk_free(Hunk_t *h, char *ptr)
{
	REGISTER int i;
	REGISTER unsigned char *up;
	
	/* Look in usagemap */
	i = (ptr - h->b.ptr)/h->size;
	if (i < 0 || i >= h->total_hunks) return;
	up = &(usagemap(h)[i/8]); i = 1 << (i&7);
	if (!(*up & i)) return;
	*up ^= i;
	
	/* Correct hunk counters */
	if (h->used_hunks == h->total_hunks) {
		if (--h->used_hunks) { /* insert into free_h */
			h->free_next = free_h[h->size];
			free_h[h->size] = h;
		}
	} else {
		if (!--h->used_hunks) { /* delete from free_h */
			REGISTER Hunk_t *p, *pp;
			for (p=free_h[h->size],pp=NULL;p!=h;pp=p,p=p->free_next);
			if (!pp) free_h[h->size] = p->free_next; else pp->free_next = p->free_next;
		}
	}
	
	/* Kill empty Hunk_t */
	if (!h->used_hunks) {
		{ /* delete from size_h */
			REGISTER Hunk_t *p, *pp;
			for (p=size_h[h->size],pp=NULL;p!=h;pp=p,p=p->size_next);
			if (!pp) size_h[h->size] = p->size_next; else pp->size_next = p->size_next;
		}
		h->b.btype &= ~BT_HUNKED;
		h->b.ptr = (char*)ALIGN((unsigned long)((&h->b)+1), BIG_ALIGN);
		bl_free(&h->b);
	}
}

/* PUBLIC xmalloc/xfree functions */

static VOLATILE int initialized = 0;
static int empty_element;
void *xmalloc(size_t size)
{
	if (!initialized) {
		/* initialize trees & hashes */
		free_mem.offs = (((char*)&(tmp_block.free_mem_node))-((char*)&tmp_block));
		ptrs.offs = (((char*)&(tmp_block.ptrs_node))-((char*)&tmp_block));
		free_mem.root = NULL; ptrs.root = NULL;
		free_mem.comp_type = 0; ptrs.comp_type = 1;
		memset(size_h, 0, sizeof(Hunk_t*)*HUNK_MAXSIZE);
		memset(free_h, 0, sizeof(Hunk_t*)*HUNK_MAXSIZE);
		initialized = 1;
	}
	if (size <= 0) return (void*)&empty_element;
	if (size < HUNK_MAXSIZE) return hunk_alloc(size);
	{
		REGISTER Block_t *p;
		
		p = bl_alloc(size);
		ass(p!=NULL);
		return p->ptr;
	}
}
void xfree(void *ptr)
{
	REGISTER Block_t *pp;
	REGISTER Node_t *nd, *bnd;
	
	if (!initialized) return;
	for (nd=ptrs.root, bnd=NULL;nd;) {
		if (((char*)nd) > (char*)ptr) nd = nd->l;
		else { bnd = nd; nd = nd->r; }
	}
	if (!bnd) return;
	pp = (Block_t*)(((char*)bnd)-ptrs.offs);
	if (pp->btype & BT_HUNKED) hunk_free((Hunk_t*)pp, ptr);
	else if (pp->btype & BT_BUSY) {
		if (((char*)ptr) != pp->ptr) return;
		bl_free(pp);
	}
}