#include "pool.h"

typedef struct _FreeList FreeList;
typedef struct _PoolHeader PoolHeader;
typedef struct _Pool Pool;
typedef struct _FreeListElement FreeListElement;

struct _FreeList
{
  FreeListElement *head;
  gulong count; /* number of elements on the free list */
  GMutex *lock; /* The free list needs to be threadsafe. We should be able  */
                /* to use the fast atomic implementation from GStreamer,    */
		/* provided we can make it count.                           */
};

struct _FreeListElement
{
  FreeListElement *next;
};

/* each pool is 4k, and allocated with 4k alignment */
#define POOL_SIZE 4096
#define POOL_HEADER_SIZE 32
#define POOL_SIZE_SHIFT 12
#define POOL(mem) ((PoolHeader*)(((gulong)mem) & ~(1 << POOL_SIZE_SHIFT - 1)))

struct _PoolHeader
{
  PoolHeader *next;      /* Pools owned by each thread are kept in a doubly-linked list. */
  PoolHeader *previous;  /* Pools on the global free list just use the next pointer.     */

  FreeList free_list;    /* Free list for this pool */

  GThread *owner;        /* Each pool knows which thread it belongs to.                  */
                         /* This information is used to ensure that we only manipulate   */
                         /* the per-thread pool lists from the thread itself.            */

  gboolean full;         /* The full flag indicates whether the pool is on the list      */
                         /* of full pools. We don't simply rely on the free list count   */
                         /* for that, since it may happen that a pool becomes partially  */
                         /* full without being moved off the full list (if the memory    */
                         /* is freed from a different thread).                           */
};

struct _MemChunk 
{
  FreeList free_pools;   /* Global list of free pools */

  GPrivate *pools;       /* Partially filled pools for each thread. */
  GPrivate *full_pools;  /* Full pools for each thread. Note that these lists are only   */
                          * manipulated by the thread that owns them                     */
  gint atom_size;
  gint atoms_per_pool;
};

static inline void 
free_list_init (FreeList *free_list)
{
  free_list->head = NULL;
  free_list->count = 0;
  free_list->lock = g_mutex_new ();
}

static inline void 
free_list_destroy (FreeList *free_list)
{
  g_mutex_free (free_list->lock);
}

static inline void
free_list_push (FreeList *free_list,
		gpointer  mem)
{
  FreeListElement *elem;

  g_mutex_lock (free_list->lock);
  elem = mem;
  elem->next = (gpointer)free_list->head;
  free_list->head = elem;
  free_list->count++;
  g_mutex_unlock (free_list->lock);
}

static inline gpointer 
free_list_pop (FreeList *free_list)
{
  FreeListElement *head;
  
  g_mutex_lock (free_list->lock);
  head = (gpointer)free_list->head;
  
  if (head)
    {
      free_list->head = head->next;
      free_list->count--;
    }

  g_mutex_unlock (free_list->lock);

  return head;
}

/* Find out if the pool is empty. This is the reason 
 * why we need the free lists to count their members.
 */
static inline gboolean
pool_is_empty (MemChunk   *mem_chunk,
	       PoolHeader *pool)
{
  return pool->free_list.count == mem_chunk->atoms_per_pool;
}

/* Create a new pool. We use posix_memalign to ensure
 * that we get a properly aligned chunk of memory. That
 * is necessary to ensure that we can find the pool of 
 * chunk by just looking at the address.
 *
 * The chunk is initialized by pushing all atoms onto
 * the free list.
 *
 * FIXME: should we use mmap ?
 */
static PoolHeader *
create_new_pool (gint atom_size)
{
  PoolHeader *header;
  int i;

  posix_memalign (&header, POOL_SIZE, POOL_SIZE);

  header->next = NULL;
  header->previous = NULL;
  header->owner = NULL;
  header->full = FALSE;
  
  free_list_init (&header->free_list);

  for (i = POOL_HEADER_SIZE; i < POOL_SIZE; i += atom_size)
    free_list_push (&header->free_list, header + i);

  return header;
}

/* Obtain a new pool, either from the global
 * list of free pools, or by creating a new one.
 */
static inline PoolHeader *
pool_alloc (MemChunk *mem_chunk)
{
  PoolHeader *pool;
  
  pool = free_list_pop (&mem_chunk->free_pools);
  
  if (!pool)
    pool = create_new_pool (mem_chunk->atom_size);

  return pool;
}

/* Return a pool to the global list of free pools.
 * 
 * FIXME: could free pools here if the global
 *        list grows large.
 */
static inline void
pool_free (MemChunk   *mem_chunk, 
	   PoolHeader *pool)
{
  free_list_push (&mem_chunk->free_pools, pool);
}

/* When a thread terminates, all pools owned by 
 * it are marked as not belonging to any thread.
 * We can't simply move them to the global free list
 * or free them here, since chunks from the pools
 * may have been passed to other threads and could
 * still be in use.
 * The pools will eventually be returned to the
 * global free list, when they become empty, 
 * see mem_chunk_free().
 */
static void
disown_pools (gpointer data)
{
  PoolHeader *pool, *next;

  pool = data;
  while (pool)
    {
      next = pool->next;

      pool->previous = NULL;
      pool->next = NULL;
      pool->owner = NULL;
      
      pool = next;
    }
}

/* Creates a new memchunk for the given atom size.
 * We don't allow to specify the pool size, since
 * the design of this allocator relies on fixed-size
 * pools.
 */
MemChunk *
mem_chunk_new (gint atom_size)
{
  MemChunk *mem_chunk;

  mem_chunk = g_new (MemChunk, 1);

  free_list_init (&mem_chunk->free_pools);
  
  mem_chunk->pools = g_private_new (disown_pools);
  mem_chunk->full_pools = g_private_new (disown_pools);

  mem_chunk->atom_size = atom_size;
  mem_chunk->atoms_per_pool = (POOL_SIZE - sizeof (PoolHeader)) / atom_size;

  return mem_chunk;
};

/* Allocate a chunk. We first try to find a usable pool
 * on the list of pools owned by this thread, and if 
 * that fails, obtain a new pool for this thread.
 *
 * If the pool becomes full due to this allocation, mark
 * it as full and move it from the list of usable pools 
 * to the list of full pools.
 */  
gpointer 
mem_chunk_alloc (MemChunk *mem_chunk)
{
  PoolHeader *pool;
  gpointer mem;

 again:
  pool = g_private_get (mem_chunk->pools);

  if (!pool)
    {
      pool = pool_alloc (mem_chunk);

      pool->owner = g_thread_self ();
      pool->previous = NULL;
      pool->next = NULL;

      g_private_set (mem_chunk->pools, pool);
    }
  
  mem = free_list_pop (&pool->free_list);

  if (!mem)
    {
      pool->full = TRUE;

      if (pool->next)
	pool->next->previous = NULL;
      g_private_set (mem_chunk->pools, pool->next);

      pool->previous = NULL;
      pool->next = g_private_get (mem_chunk->full_pools);
      pool->next->previous = pool;
      g_private_set (mem_chunk->full_pools, pool);

      goto again;
    }

  return mem;
}

/* Free a chunk. Since free lists are threadsafe, we can just push 
 * the chunk onto the free list of its pool, even if the pool belongs
 * to a different thread.  
 * 
 * If the pool was full before and if it is ours, move it from the
 * list of full pools to the list of usable pools. If it is not ours, 
 * just let it sit in the list of full pools. 
 *
 * FIXME: the idea is that it should be rare for memory to be freed
 * in a different thread, thus the pool will eventually be moved 
 * off the full list, but this will fail if all chunks on the pool
 * are being passed to other threads...
 *
 * If the pool has become empty and it is ours, give it back. We also
 * handle leftover pools of dead threads here. If such an unowned
 * pool becomes empty, we also give it back. If the pool has become
 * empty but is owned by somebody else, just leave it in that threads
 * list of usable pools.
 */			  
void
mem_chunk_free (MemChunk *mem_chunk, 
		gpointer  mem)
{
  PoolHeader *pool;

  pool = POOL (mem);

  free_list_push (&pool->free_list, mem); 

  if (pool->full)
    {
      /* if this is our pool, move it  */
      if (pool->owner == g_thread_self ())
	{
	  if (pool->next)
	    pool->next->previous = pool->previous;
	  if (pool->previous)
	    pool->previous->next = pool->next;
	  else
	    g_private_set (mem_chunk->full_pools, pool->next);

	  pool->full = FALSE;
	  pool->previous = NULL;
	  pool->next = g_private_get (mem_chunk->pools);
	  g_private_set (mem_chunk->pools, pool);
	}

      if (!pool->owner)
	pool->full = FALSE;
    }

  if (pool_is_empty (mem_chunk, pool))
    {
      /* only free the pool if it is unowned or if it is ours */
      if (pool->owner == g_thread_self ())
	{
	  if (pool->next)
	    pool->next->previous = pool->previous;
	  if (pool->previous)
	    {
	      pool->previous->next = pool->next;
	    }
	  else
	    g_private_set (mem_chunk->full_pools, pool->next);

	  pool->previous = NULL;
	  pool->next = NULL;
	  pool->owner = NULL;
	}
      
      if (!pool->owner)
	pool_free (mem_chunk, pool);
    }
}
