#include <unistd.h>
#include <sys/mman.h>
#include <stdio.h>
#include <stdlib.h>


#define assert(x) { write(2, x, strlen(x)); exit(1); }
#define SMART_REALLOC 1

struct memblk
{
	int mb_magic;
	struct memblk *mb_next;
	struct memblk *mb_prev;
	size_t mb_len;
	struct poolblk *mb_pool;
};


struct poolblk
{
	struct poolblk *pb_next;
	struct memblk *pb_freelist;
};

static struct poolblk *poollist;

/* CHUNKSIZE should be 2^n - 1 and should be grater than sizeof(struct memblk) */
#define CHUNKSIZE 31
#define POOLSIZE 4*1024

#define MEM_MAGIC_FREE 0x46524545
#define MEM_MAGIC_USED 0x55534544
#ifdef __uclinux__
#define MMAP_FLAGS MAP_SHARED | MAP_ANONYMOUS
#else
#define MMAP_FLAGS MAP_PRIVATE | MAP_ANONYMOUS
#endif

static void *
   real_malloc (size_t len)
{
	void *result = mmap ((void *) 0, len, PROT_READ | PROT_WRITE,
						 MMAP_FLAGS, 0, 0);
	if (result == (void *) -1)
		return 0;

	return result;


}

/* unlink from the freelist */
static 
   mb_unlink (struct memblk *mb)
{
	if (mb->mb_prev)
		mb->mb_prev->mb_next = mb->mb_next;
	else
		mb->mb_pool->pb_freelist = mb->mb_next;
	if (mb->mb_next)
		mb->mb_next->mb_prev = mb->mb_prev;
	mb->mb_magic = MEM_MAGIC_USED;
}

/* relink onto the end of the freelist */
static 
   mb_link (struct memblk *mb)
{
	struct memblk *lastmb;

	lastmb = mb->mb_pool->pb_freelist;
	if (lastmb)
	{
		while (lastmb->mb_next)
			lastmb = lastmb->mb_next;
		lastmb->mb_next = mb;
	}
	else
	{
		mb->mb_pool->pb_freelist = mb;
	}
	mb->mb_prev = lastmb;
	mb->mb_next = 0;
	mb->mb_magic = MEM_MAGIC_FREE;
}


static char *
   alloc_memblk (struct memblk *mb, size_t len, int isfree)
{
	size_t newlen;
	struct memblk *newmb;

	newlen = mb->mb_len - len;
	if (newlen > CHUNKSIZE)
	{
		newmb = (struct memblk *) ((char *) mb + len);

		*newmb = *mb;
		newmb->mb_len = newlen;
		mb->mb_len = len;

		if (isfree) /* we're suballocating from the block on free list */
		{
			/* re-link the freelist */
			if (mb->mb_prev)
				mb->mb_prev->mb_next = newmb;
			else 
				mb->mb_pool->pb_freelist = newmb;

			if (mb->mb_next)
				mb->mb_next->mb_prev = newmb;

			mb->mb_magic = MEM_MAGIC_USED;			

		}
		else
		{
			/* we're suballocating from already allocated block (in realloc) */
			mb_link(newmb);
		}
	}
	else if (isfree)
	{
		mb_unlink (mb);
	}
	mb->mb_prev = mb->mb_next = 0;
	return ((char *) ++mb);
}

#ifdef MEMORY_DEBUG
show_pools ()
{
	struct poolblk *pb;
	struct memblk *mb;

	printf ("====\n");
	for (pb = poollist; pb; pb = pb->pb_next)
	{
		printf ("Pool %x\n", pb);
		for (mb = pb->pb_freelist; mb; mb = mb->mb_next)
		{
			printf ("  Block: at %x %d bytes\n", mb, mb->mb_len);
		}
	}
	printf ("~~~~\n");
}
#endif

static struct memblk *
   m_search_pool (size_t len)
{
	struct poolblk *pb;
	struct memblk *mb;

	for (pb = poollist; pb; pb = pb->pb_next)
	{
		for (mb = pb->pb_freelist; mb; mb = mb->mb_next)
			if (mb->mb_len >= len)
			{
#ifdef MEMORY_DEBUG
				printf ("found free segment at %x\n", mb);
#endif
				return (mb);
			}
	}
	return ((struct memblk *) 0);
}

static 
   defrag_heap ()
{
	struct poolblk *pb;
	struct memblk *mb, *mb2, *mbend;

#ifdef MEMORY_DEBUG
	printf ("Before defrag: ");
	show_pools ();
#endif

	for (pb = poollist; pb; pb = pb->pb_next)
	{
mb_again:
		for (mb = pb->pb_freelist; mb; mb = mb->mb_next)
		{
			mbend = (struct memblk *) ((char *) mb + mb->mb_len);
			for (mb2 = pb->pb_freelist; mb2; mb2 = mb2->mb_next)
			{
				if (mb2 == mbend)
				{
		  /* unlink mb2 */
					mb_unlink (mb2);
					mb->mb_len += mb2->mb_len;
					goto mb_again;
				}
			}
		}
	}

#ifdef MEMORY_DEBUG
	printf ("After defrag: ");
	show_pools ();
#endif
}


void *
   malloc (size_t len)
{
	struct poolblk *pb, *newpb;
	struct memblk *mb;
	size_t lenadj, poolsz;
	char *end;

	lenadj = (len + sizeof (struct memblk) + CHUNKSIZE) & ~CHUNKSIZE;
#ifdef MEMORY_DEBUG
	printf ("searching for segment of %d bytes\n", lenadj);
#endif

	if (poollist)
	{
		mb = m_search_pool (lenadj);
		if (mb)
			return (alloc_memblk (mb, lenadj, 1));
#ifdef MEMORY_DEBUG
		printf ("cleaning the heap\n");
#endif
		defrag_heap ();
		mb = m_search_pool (lenadj);
		if (mb)
			return (alloc_memblk (mb, lenadj, 1));
	}
#ifdef MEMORY_DEBUG
	printf ("allocating new pool\n");
#endif
	if (lenadj + sizeof (struct poolblk) > POOLSIZE)
		poolsz = (lenadj + sizeof (struct poolblk) + CHUNKSIZE) & ~CHUNKSIZE;
	else
		poolsz = POOLSIZE;
	newpb = (struct poolblk *) real_malloc (poolsz);
	if (!newpb)
		return ((char *) 0);
	if (poollist)
	{
		for (pb = poollist; pb->pb_next; pb = pb->pb_next)
			;
		pb->pb_next = newpb;
	}
	else
	{
		poollist = newpb;
	}
	newpb->pb_next = 0;
	end = (char *) newpb + poolsz;
	mb = newpb->pb_freelist = (struct memblk *) (newpb + 1);
	mb->mb_next = 0;
	mb->mb_prev = 0;
	mb->mb_pool = newpb;
	mb->mb_len = end - (char *) mb;
	mb->mb_magic = MEM_MAGIC_FREE;
#ifdef MEMORY_DEBUG
	printf ("newpb=%x\n", newpb);
	printf ("newpb->pb_freelist=%x\n", newpb->pb_freelist);
	printf ("newpb->pb_freelist->mb_len=%x\n", newpb->pb_freelist->mb_len);
	show_pools ();
#endif
	mb = m_search_pool (lenadj);
	if (mb)
		return (alloc_memblk (mb, lenadj, 1));
#ifdef MEMORY_DEBUG
	printf ("huh? m_search_pool failed?");
#endif
	return ((char *) 0);
}

void 
   free (void *p)
{
	struct memblk *mb;

	mb = (struct memblk *) p;
	--mb;
	if (mb->mb_magic != MEM_MAGIC_USED) {
		assert("attempted to free block with corrupt MAGIC");
		kill(getpid(), 11);
	}
	mb_link (mb);
}

void * realloc(void* p, size_t len)
{
	size_t lenadj;
	struct memblk *mb;
	void* newmem;
	
	if (!len)
	{
		if (p) free(p);
		return 0;
	}



	if (p)
	{
		mb = (struct memblk *) p;
		--mb;
		if (mb->mb_magic != MEM_MAGIC_USED) {
			assert("attempted to reelloc block with corrupt MAGIC");
			kill(getpid(), 11);
		}

#ifdef SMART_REALLOC	
		lenadj = (len + sizeof (struct memblk) + CHUNKSIZE) & ~CHUNKSIZE;
		if (lenadj < mb->mb_len)
		{
			alloc_memblk(mb, lenadj, 0);
			return (char*) (mb + 1);
		}
#endif
	}
	

	newmem = malloc(len);
	if (!newmem)
	{
		if (p) free(p);
		return (0);
	}

	if (p)
	{
		size_t clen = mb->mb_len - sizeof(struct memblk);
		if (len < clen)
			clen = len;
		memcpy(newmem, p, clen);
		free(p);
	}
	return newmem;
}

	
#ifdef TEST

int freecnt = 256;
void* ttab[256];

int findfreeslot()
{
	int i, x;
	
	if (!freecnt)
		return 0;

	i = rand() % 256;
	while(ttab[i])
	{
		i = (i+1) % 256;
	}
	return i+1;
}


int findnonfree(int i)
{
	while(!ttab[i])
	{
		i = (i+1) % 256;
	}
	return i;
}

	
	
main(int argc, char* argv[])
{

	int nextfree,nonfree, i;
	size_t sz;
	void* ptr;
	
	while(nextfree = findfreeslot())
	{
		int op = rand() % 10;

		if (op >= 7)
			op = 2;
		else if (op > 3)
			op = 1;
		else
			op = 0;
		
		if (freecnt == 256)
			op = 0;

		

		nextfree--;

		switch(op)
		{
			case 0: /* allocate */
				sz = rand() % (16*1024);
				ttab[nextfree] = ptr = malloc(sz);
				printf("allocated  [%d]=%08x, %d\n",  nextfree, ptr, sz);				
				memset(ptr, 'M', sz);
				freecnt--;
				break;
			case 1:
				nonfree = findnonfree(nextfree);
				sz = rand() % (16*1024);
				ttab[nonfree] = ptr = realloc(ttab[nonfree], sz);
				printf("reallocated  [%d]=%08x, %d\n",  nonfree, ptr, sz);								
				if (sz && (ptr != 0))
					memset(ptr, 'R', sz);
				else
					freecnt++;
				break;
			case 2:
				nonfree = findnonfree(nextfree);
				printf("freeing [%d]=%08x\n", nonfree, ttab[nonfree]);
				free(ttab[nonfree]);
				ttab[nonfree] = 0;
				freecnt++;
				break;
		}
	}

	printf("freeing all\n");
	for(i = 0; i < 256; i++)
	{
		free(ttab[i]);
		freecnt++;
	}

}


#endif