/*
 * This allocator takes blocks from a coarser allocator (p->alloc) and
 * uses them as arenas.
 * 
 * An arena is split into a sequence of blocks of variable size.  The
 * blocks begin with a Bhdr that denotes the length (including the Bhdr)
 * of the block.  An arena begins with an Arena header block (Arena,
 * ARENA_MAGIC) and ends with a Bhdr block with magic ARENATAIL_MAGIC and
 * size 0.  Intermediate blocks are either allocated or free.  At the end
 * of each intermediate block is a Btail, which contains information
 * about where the block starts.  This is useful for walking backwards.
 * 
 * Free blocks (Free*) have a magic value of FREE_MAGIC in their Bhdr
 * headers.  They are kept in a binary tree (p->freeroot) traversible by
 * walking ->left and ->right.  Each node of the binary tree is a pointer
 * to a circular doubly-linked list (next, prev) of blocks of identical
 * size.  Blocks are added to this ``tree of lists'' by pooladd(), and
 * removed by pooldel().
 * 
 * When freed, adjacent blocks are coalesced to create larger blocks when
 * possible.
 * 
 * Allocated blocks (Alloc*) have one of two magic values: KEMPT_MAGIC or
 * UNKEMPT_MAGIC.  When blocks are released from the pool, they have
 * magic value UNKEMPT_MAGIC.  Once the block has been trimmed by kemb()
 * and the amount of user-requested data has been recorded in the
 * datasize field of the tail, the magic value is changed to KEMPT_MAGIC.
 * All blocks returned to callers should be of type KEMPT_MAGIC, as
 * should all blocks passed to us by callers.  The amount of data the user
 * asked us for can be found by subtracting the short in tail->datasize 
 * from header->size.  Further, the up to at most four bytes between the
 * end of the user-requested data block and the actual Btail structure are
 * marked with a magic value, which is checked to detect user overflow.
 * 
 * The arenas returned by p->alloc are kept in a doubly-linked list
 * (p->arenalist) running through the arena headers, sorted by descending
 * base address (prev, next).  When a new arena is allocated, we attempt
 * to merge it with its two neighbors via p->merge.
 */

#include <u.h>
#include <libc.h>
#include <pool.h>

typedef struct Alloc	Alloc;
typedef struct Arena	Arena;
typedef struct Bhdr	Bhdr;
typedef struct Btail	Btail;
typedef struct Free	Free;

struct Bhdr {
	ulong	magic;
	ulong	size;
};
enum {
	NOT_MAGIC = 0xdeadfa11,
};
#define B2NB(b) ((Bhdr*)((uchar*)(b)+(b)->size))

#define SHORT(x) (((x)[0] << 8) | (x)[1])
#define PSHORT(p, x) \
	(((uchar*)(p))[0] = ((x)>>8)&0xFF, \
	((uchar*)(p))[1] = (x)&0xFF)

enum {
	TAIL_MAGIC0 = 0xBE,
	TAIL_MAGIC1 = 0xEF
};
struct Btail {
	uchar	magic0;
	uchar	datasize[2];
	uchar	magic1;
	ulong	size;	/* same as Bhdr->size */
};
#define B2T(b)	((Btail*)((uchar*)(b)+(b)->size-sizeof(Btail)))
#define B2PT(b) ((Btail*)((uchar*)(b)-sizeof(Btail)))
#define T2HDR(t) ((Bhdr*)((uchar*)(t)+sizeof(Btail)-(t)->size))
struct Free {
			Bhdr;
	Free*	left;
	Free*	right;
	Free*	next;
	Free*	prev;
};
enum {
	FREE_MAGIC = 0xBA5EBA11,
};

/*
 * the point of the notused fields is to make 8c differentiate
 * between Bhdr and Allocblk, and between Kempt and Unkempt.
 */
struct Alloc {
			Bhdr;
};
enum {
	KEMPT_MAGIC = 0x0A110C09,
	UNKEMPT_MAGIC = 0xCAB00D1E+1,
};

struct Arena {
			Bhdr;
	Arena*	aup;
	Arena*	down;
	ulong	asize;
	ulong	pad;	/* to a multiple of 8 bytes */
};
enum {
	ARENA_MAGIC = 0xC0A1E5CE+1,
	ARENATAIL_MAGIC = 0xEC5E1A0C+1,
};
#define A2TB(a)	((Bhdr*)((uchar*)(a)+(a)->asize-sizeof(Bhdr)))
#define A2B(a)	B2NB(a)

enum {
	MINBLOCKSIZE = sizeof(Free)+sizeof(Btail)
};

static uchar datamagic[] = { 0xFE, 0xF1, 0xF0, 0xFA };

#define	Poison	(void*)0xCafeBabe

#define _B2D(a)	((void*)((uchar*)a+sizeof(Bhdr)))
#define _D2B(v)	((Alloc*)((uchar*)v-sizeof(Bhdr)))

// static void*	_B2D(void*);
// static void*	_D2B(void*);
static void*	B2D(Pool*, Alloc*);
static Alloc*	D2B(Pool*, void*);
static Arena*	arenamerge(Pool*, Arena*, Arena*);
static void		blockcheck(Pool*, Bhdr*);
static Alloc*	blockmerge(Pool*, Bhdr*, Bhdr*);
static Alloc*	blocksetdsize(Pool*, Alloc*, ulong);
static Bhdr*	blocksetsize(Bhdr*, ulong);
static ulong	bsize2asize(Pool*, ulong);
static ulong	dsize2bsize(Pool*, ulong);
static ulong	getdsize(Alloc*);
static Alloc*	kemb(Pool*, Alloc*, ulong);
static Free*	listadd(Free*, Free*);
static Free*	listdelete(Free*, Free*);
static void		logstack(Pool*);
static Free**	ltreewalk(Free**, ulong);
static void		memmark(void*, int, ulong);
static Free*	pooladd(Pool*, Alloc*);
static void*	poolallocl(Pool*, ulong);
static void		poolcheckl(Pool*);
static void		poolcheckarena(Pool*, Arena*);
static int		poolcompactl(Pool*);
static Alloc*	pooldel(Pool*, Free*);
static void		pooldumpl(Pool*);
static void		pooldumparena(Pool*, Arena*);
static void		poolfreel(Pool*, void*);
static void		poolnewarena(Pool*, ulong);
static void*	poolreallocl(Pool*, void*, ulong);
static Free*	treedelete(Free*, Free*);
static Free*	treeinsert(Free*, Free*);
static Free*	treelookup(Free*, ulong);
static Free*	treelookupgt(Free*, ulong);

/*
 * Debugging
 * 
 * Antagonism causes blocks to always be filled with garbage if their
 * contents are undefined.  This tickles both programs and the library.
 * It's a linear time hit but not so noticeable during nondegenerate use.
 * It would be worth leaving in except that it negates the benefits of the
 * kernel's demand-paging.  The tail magic and end-of-data magic 
 * provide most of the user-visible benefit that antagonism does anyway.
 *
 * Paranoia causes the library to recheck the entire pool on each lock
 * or unlock.  A failed check on unlock means we tripped over ourselves,
 * while a failed check on lock tends to implicate the user.  Paranoia has
 * the potential to slow things down a fair amount for pools with large
 * numbers of allocated blocks.  It completely negates all benefits won
 * by the binary tree.  Turning on paranoia in the kernel makes it painfully
 * slow.
 *
 * Verbosity induces the dumping of the pool via p->print at each lock operation.
 * By default, only one line is logged for each alloc, free, and realloc.
 */

/* the if(!x);else avoids ``dangling else'' problems */
#define antagonism	if(!(p->flags & POOL_ANTAGONISM));else
#define paranoia	if(!(p->flags & POOL_PARANOIA));else
#define verbosity	if(!(p->flags & POOL_VERBOSITY));else

#define DPRINT	if(!(p->flags & POOL_DEBUGGING));else p->print
#define LOG		if(!(p->flags & POOL_LOGGING));else p->print

/*
 * Tree walking
 */

/* ltreewalk: return address of pointer to node of size == size */
static Free**
ltreewalk(Free **t, ulong size)
{
	Free **ot;
//	Free **hist[40];
//	int nhist;

//	nhist = 0;

	ot = t;

	assert(t != nil /* ltreewalk */);

	for(;;) {
//		hist[nhist++] = t;
		if(*t == nil)
			return t;

		assert((*t)->magic == FREE_MAGIC);

		if(size == (*t)->size)
			return t;
		if(size < (*t)->size)
			t = &(*t)->left;
		else
			t = &(*t)->right;
	}
	return nil;	/* not reached */
}

/* treelookup: find node in tree with size == size */
static Free*
treelookup(Free *t, ulong size)
{
	return *ltreewalk(&t, size);
}

/* treeinsert: insert node into tree */
static Free*
treeinsert(Free *tree, Free *node)
{
	Free **loc, *repl;

	assert(node != nil /* treeinsert */);

	loc = ltreewalk(&tree, node->size);
	if(*loc == nil) {
		node->left = nil;
		node->right = nil;
	} else {	/* replace existing node */
		repl = *loc;
		node->left = repl->left;
		node->right = repl->right;
	}
	*loc = node;
	return tree;
}

/* treedelete: remove node from tree */
static Free*
treedelete(Free *tree, Free *node)
{
	Free **loc, **lsucc, *succ;

	assert(node != nil /* treedelete */);

	loc = ltreewalk(&tree, node->size);
	assert(*loc == node);

	if(node->left == nil)
		*loc = node->right;
	else if(node->right == nil)
		*loc = node->left;
	else {
		/* have two children, use inorder successor as replacement */
		for(lsucc = &node->right; (*lsucc)->left; lsucc = &(*lsucc)->left)
			;
		succ = *lsucc;
		*lsucc = succ->right;
		succ->left = node->left;
		succ->right = node->right;
		*loc = succ;
	}

	node->left = node->right = Poison;
	return tree;
}

/* treelookupgt: find smallest node in tree with size >= size */
static Free*
treelookupgt(Free *t, ulong size)
{
	Free *lastgood;	/* last node we saw that was big enough */

	lastgood = nil;
	for(;;) {
		if(t == nil)
			return lastgood;
		if(size == t->size)
			return t;
		if(size < t->size) {
			lastgood = t;
			t = t->left;
		} else {
			t = t->right;
		}
	}
	return nil;	/* not reached */
}

/* 
 * List maintenance
 */

/* listadd: add a node to a doubled linked list */
static Free*
listadd(Free *list, Free *node)
{
	if(list == nil) {
		node->next = node;
		node->prev = node;
		return node;
	}

	node->prev = list->prev;
	node->next = list;

	node->prev->next = node;
	node->next->prev = node;

	return list;
}

/* listdelete: remove node from a doubly linked list */
static Free*
listdelete(Free *list, Free *node)
{
	if(node->next == node) {	/* singular list */
		node->prev = node->next = Poison;
		return nil;
	}

	node->next->prev = node->prev;
	node->prev->next = node->next;

	if(list == node)
		list = node->next;

	node->prev = node->next = Poison;
	return list;
}

/*
 * Pool maintenance
 */

/* pooladd: add anode to the free pool */
static Free*
pooladd(Pool *p, Alloc *anode)
{
	Free *lst, *olst;
	Free *node;
	Free **parent;

	antagonism {
		memmark(_B2D(anode), 0xF7, anode->size-sizeof(Bhdr)-sizeof(Btail));
	}

	node = (Free*)anode;
	node->magic = FREE_MAGIC;
	parent = ltreewalk(&p->freeroot, node->size);
	olst = *parent;
	lst = listadd(olst, node);
	if(olst != lst)	/* need to update tree */
		*parent = treeinsert(*parent, lst);
	p->curfree += node->size;
	return node;
}

/* pooldel: remove node from the free pool */
static Alloc*
pooldel(Pool *p, Free *node)
{
	Free *lst, *olst;
	Free **parent;

	parent = ltreewalk(&p->freeroot, node->size);
	olst = *parent;
	assert(olst != nil /* pooldel */);

	lst = listdelete(olst, node);
	if(lst == nil)
		*parent = treedelete(*parent, olst);
	else if(lst != olst)
		*parent = treeinsert(*parent, lst);

	node->left = node->right = Poison;
	p->curfree -= node->size;

	antagonism {
		memmark(_B2D(node), 0xF9, node->size-sizeof(Bhdr)-sizeof(Btail));
	}

	node->magic = UNKEMPT_MAGIC;
	return (Alloc*)node;
}

/*
 * Block maintenance 
 */
/* block allocation */
static ulong
dsize2bsize(Pool *p, ulong sz)
{
	sz += sizeof(Bhdr)+sizeof(Btail);
	if(sz < p->minblock)
		sz = p->minblock;
	sz = (sz+p->quantum-1)&~(p->quantum-1);
	return sz;
}

static ulong
bsize2asize(Pool *p, ulong sz)
{
	sz += sizeof(Arena)+sizeof(Btail);
	if(sz < p->minarena)
		sz = p->minarena;
	sz = (sz+p->quantum)&~(p->quantum-1);
	return sz;
}

/* blockmerge: merge a and b, known to be adjacent */
/* both are removed from pool if necessary. */
static Alloc*
blockmerge(Pool *pool, Bhdr *a, Bhdr *b)
{
	Btail *t;

	assert(B2NB(a) == b);

	if(a->magic == FREE_MAGIC)
		pooldel(pool, (Free*)a);
	if(b->magic == FREE_MAGIC)
		pooldel(pool, (Free*)b);

	t = B2T(a);
	t->size = (ulong)Poison;
	t->magic0 = NOT_MAGIC;
	t->magic1 = NOT_MAGIC;
	PSHORT(t->datasize, NOT_MAGIC);

	a->size += b->size;
	t = B2T(a);
	t->size = a->size;
	PSHORT(t->datasize, 0xFFFF);

	b->size = NOT_MAGIC;
	b->magic = NOT_MAGIC;

	a->magic = UNKEMPT_MAGIC;
	return (Alloc*)a;
}

/* blocksetsize: set the total size of a block, fixing tail pointers */
static Bhdr*
blocksetsize(Bhdr *b, ulong bsize)
{
	Btail *t;

	assert(b->magic != FREE_MAGIC /* blocksetsize */);

	b->size = bsize;
	t = B2T(b);
	t->size = b->size;
	t->magic0 = TAIL_MAGIC0;
	t->magic1 = TAIL_MAGIC1;
	return b;
}

/* getdsize: return the requested data size for an allocated block */
static ulong
getdsize(Alloc *b)
{
	Btail *t;
	t = B2T(b);
	return b->size - SHORT(t->datasize);
}

/* blocksetdsize: set the user data size of a block */
static Alloc*
blocksetdsize(Pool *p, Alloc *b, ulong dsize)
{
	Btail *t;
	uchar *q, *eq;

	assert(b->size >= dsize2bsize(p, dsize));
	assert(b->size - dsize < 0x10000);

	t = B2T(b);
	PSHORT(t->datasize, b->size - dsize);

	q=(uchar*)_B2D(b)+dsize;
	eq = (uchar*)t;
	if(eq > q+4)
		eq = q+4;
	for(; q<eq; q++)
		*q = datamagic[((ulong)q)%nelem(datamagic)];

	return b;
}

/* kemb: trim a block down to what is needed to hold dsize bytes of user data */
static Alloc*
kemb(Pool *p, Alloc *b, ulong dsize)
{
	ulong extra, bsize;
	Alloc *frag;

	bsize = dsize2bsize(p, dsize);
	extra = b->size - bsize;
	if(b->size - dsize >= 0x10000 ||
	  (extra >= bsize>>2 && extra >= MINBLOCKSIZE && extra >= p->minblock)) {
		blocksetsize(b, bsize);
		frag = (Alloc*) B2NB(b);

		antagonism {
			memmark(frag, 0xF1, extra);
		}

		frag->magic = UNKEMPT_MAGIC;
		blocksetsize(frag, extra);
		pooladd(p, frag);
	}

	b->magic = KEMPT_MAGIC;
	blocksetdsize(p, b, dsize);
	return b;
}

/*
 * Arena maintenance
 */

/* arenasetsize: set arena size, updating tail */
static void
arenasetsize(Arena *a, ulong asize)
{
	Bhdr *atail;

	a->asize = asize;
	atail = A2TB(a);
	atail->magic = ARENATAIL_MAGIC;
	atail->size = 0;
}

/* poolnewarena: allocate new arena */
static void
poolnewarena(Pool *p, ulong asize)
{
	Arena *a;
	Arena *ap, *lastap;
	Alloc *b;

	LOG(p, "newarena %lud\n", asize);
	if(p->cursize+asize > p->maxsize) {
		if(poolcompactl(p) == 0)
			LOG(p, "pool too big: %lud+%lud > %lud\n",
				p->cursize, asize, p->maxsize);
		return;
	}

	if((a = p->alloc(asize)) == nil) {
		/* assume errstr set by p->alloc */
		return;
	}

	p->cursize += asize;

	/* arena hdr */
	a->magic = ARENA_MAGIC;
	blocksetsize(a, sizeof(Arena));
	arenasetsize(a, asize);
	blockcheck(p, a);

	/* create one large block in arena */
	b = (Alloc*)A2B(a);
	b->magic = UNKEMPT_MAGIC;
	blocksetsize(b, (uchar*)A2TB(a)-(uchar*)b);
	blockcheck(p, b);
	pooladd(p, b);
	blockcheck(p, b);

	/* sort arena into descending sorted arena list */
	for(lastap=nil, ap=p->arenalist; ap > a; lastap=ap, ap=ap->down)
		;

	if(a->down = ap)	/* assign = */
		a->down->aup = a;

	if(a->aup = lastap)	/* assign = */
		a->aup->down = a;
	else
		p->arenalist = a;

	/* merge with surrounding arenas if possible */
	/* must do a with up before down with a (think about it) */
	if(a->aup)
		arenamerge(p, a, a->aup);
	if(a->down)
		arenamerge(p, a->down, a);
}

/* blockresize: grow a block to encompass space past its end, possibly by */
/* trimming it into two different blocks. */
static void
blockgrow(Pool *p, Bhdr *b, ulong nsize)
{
	if(b->magic == FREE_MAGIC) {
		Alloc *a;
		Bhdr *bnxt;
		a = pooldel(p, (Free*)b);
		blockcheck(p, a);
		blocksetsize(a, nsize);
		blockcheck(p, a);
		bnxt = B2NB(a);
		if(bnxt->magic == FREE_MAGIC)
			a = blockmerge(p, a, bnxt);
		blockcheck(p, a);
		pooladd(p, a);
	} else {
		Alloc *a;
		ulong dsize;

		a = (Alloc*)b;
		dsize = getdsize(a);
		blocksetsize(a, nsize);
		kemb(p, a, dsize);
	}
}

/* arenamerge: attempt to coalesce to arenas that might be adjacent */
static Arena*
arenamerge(Pool *p, Arena *bot, Arena *top)
{
	Bhdr *bbot, *btop;
	Btail *t;

	blockcheck(p, bot);
	blockcheck(p, top);
	assert(bot->aup == top && top > bot);

	if(p->merge == nil || p->merge(bot, top) == 0)
		return nil;

	/* remove top from list */
	if(bot->aup = top->aup)	/* assign = */
		bot->aup->down = bot;
	else
		p->arenalist = bot;
	
	/* save ptrs to last block in bot, first block in top */
	t = B2PT(A2TB(bot));
	bbot = T2HDR(t);
	btop = A2B(top);
	blockcheck(p, bbot);
	blockcheck(p, btop);

	/* grow bottom arena to encompass top */
	arenasetsize(bot, top->asize + ((uchar*)top - (uchar*)bot));

	/* grow bottom block to encompass space between arenas */
	blockgrow(p, bbot, (uchar*)btop-(uchar*)bbot);
	blockcheck(p, bbot);
	return bot;
}

/* dumpblock: print block's vital stats */
static void
dumpblock(Pool *p, Bhdr *b)
{
	ulong *dp;
	ulong dsize;
	uchar *cp;

	dp = (ulong*)b;
	p->print(p, "pool %s block %p\nhdr %.8lux %.8