#include "Python.h"
#include "pyarena.h"

/* A simple arena block structure.

  Measurements with standard library modules suggest the average
  allocation is about 20 bytes and that most compiles use a single
  block.

  TODO(jhylton): Think about a realloc API, maybe just for the last
  allocation?
*/

#define DEFAULT_BLOCK_SIZE 8192
#define ALIGNMENT               8
#define ALIGNMENT_MASK          (ALIGNMENT - 1)
#define ROUNDUP(x)              (((x) + ALIGNMENT_MASK) & ~ALIGNMENT_MASK)

typedef struct _block {
       /* Total number of bytes owned by this block available to pass out.
        * Read-only after initialization.  The first such byte starts at
        * ab_mem.
        */
       size_t ab_size;

       /* Total number of bytes already passed out.  The next byte available
        * to pass out starts at ab_mem + ab_offset.
        */
       size_t ab_offset;

       /* An arena maintains a singly-linked, NULL-terminated list of
        * all blocks owned by the arena.  These are linked via the
        * ab_next member.
        */
       struct _block *ab_next;

       /* Pointer to the first allocatable byte owned by this block.  Read-
        * only after initialization.
        */
       void *ab_mem;
} block;

/* The arena manages two kinds of memory, blocks of raw memory
  and a list of PyObject* pointers.  PyObjects are decrefed
  when the arena is freed.
*/

struct _arena {
       /* Pointer to the first block allocated for the arena, never NULL.
          It is used only to find the first block when the arena is
          being freed.
        */
       block *a_head;

       /* Pointer to the block currently used for allocation.  It's
          ab_next field should be NULL.  If it is not-null after a
          call to block_alloc(), it means a new block has been allocated
          and a_cur should be reset to point it.
        */
       block *a_cur;

       /* A Python list object containing references to all the PyObject
          pointers associated with this area.  They will be DECREFed
          when the arena is freed.
       */
       PyObject *a_objects;

#if defined(Py_DEBUG)
       /* Debug output */
       size_t total_allocs;
       size_t total_size;
       size_t total_blocks;
       size_t total_block_size;
       size_t total_big_blocks;
#endif
};

static block *
block_new(size_t size)
{
       /* Allocate header and block as one unit.
          ab_mem points just past header. */
       block *b = (block *)malloc(sizeof(block) + size);
       if (!b)
               return NULL;
       b->ab_size = size;
       b->ab_mem = (void *)(b + 1);
       b->ab_next = NULL;
       b->ab_offset = ROUNDUP((Py_uintptr_t)(b->ab_mem)) -
         (Py_uintptr_t)(b->ab_mem);
       return b;
}

static void
block_free(block *b) {
       while (b) {
               block *next = b->ab_next;
               free(b);
               b = next;
       }
}

static void *
block_alloc(block *b, size_t size)
{
       void *p;
       assert(b);
       size = ROUNDUP(size);
       if (b->ab_offset + size > b->ab_size) {
               /* If we need to allocate more memory than will fit in
                  the default block, allocate a one-off block that is
                  exactly the right size. */
               /* TODO(jhylton): Think about space waste at end of block */
               block *newbl = block_new(
                               size < DEFAULT_BLOCK_SIZE ?
                               DEFAULT_BLOCK_SIZE : size);
               if (!newbl)
                       return NULL;
               assert(!b->ab_next);
               b->ab_next = newbl;
               b = newbl;
       }

       assert(b->ab_offset + size <= b->ab_size);
       p = (void *)(((char *)b->ab_mem) + b->ab_offset);
       b->ab_offset += size;
       return p;
}

PyArena *
PyArena_New()
{
       PyArena* arena = (PyArena *)malloc(sizeof(PyArena));
       if (!arena)
               return (PyArena*)PyErr_NoMemory();

       arena->a_head = block_new(DEFAULT_BLOCK_SIZE);
       arena->a_cur = arena->a_head;
       if (!arena->a_head) {
               free((void *)arena);
               return (PyArena*)PyErr_NoMemory();
       }
       arena->a_objects = PyList_New(0);
       if (!arena->a_objects) {
               block_free(arena->a_head);
               free((void *)arena);
               return (PyArena*)PyErr_NoMemory();
       }
#if defined(Py_DEBUG)
       arena->total_allocs = 0;
       arena->total_size = 0;
       arena->total_blocks = 1;
       arena->total_block_size = DEFAULT_BLOCK_SIZE;
       arena->total_big_blocks = 0;
#endif
       return arena;
}

void
PyArena_Free(PyArena *arena)
{
       int r;
       assert(arena);
#if defined(Py_DEBUG)
       /*
       fprintf(stderr,
               "alloc=%d size=%d blocks=%d block_size=%d big=%d objects=%d\n",
               arena->total_allocs, arena->total_size, arena->total_blocks,
               arena->total_block_size, arena->total_big_blocks,
               PyList_Size(arena->a_objects));
       */
#endif
       block_free(arena->a_head);
       /* This property normally holds, except when the code being compiled
          is sys.getobjects(0), in which case there will be two references.
       assert(arena->a_objects->ob_refcnt == 1);
       */

       /* Clear all the elements from the list.  This is necessary
          to guarantee that they will be DECREFed. */
       r = PyList_SetSlice(arena->a_objects,
                           0, PyList_GET_SIZE(arena->a_objects), NULL);
       assert(r == 0);
       assert(PyList_GET_SIZE(arena->a_objects) == 0);
       Py_DECREF(arena->a_objects);
       free(arena);
}

void *
PyArena_Malloc(PyArena *arena, size_t size)
{
       void *p = block_alloc(arena->a_cur, size);
       if (!p)
               return PyErr_NoMemory();
#if defined(Py_DEBUG)
       arena->total_allocs++;
       arena->total_size += size;
#endif
       /* Reset cur if we allocated a new block. */
       if (arena->a_cur->ab_next) {
               arena->a_cur = arena->a_cur->ab_next;
#if defined(Py_DEBUG)
               arena->total_blocks++;
               arena->total_block_size += arena->a_cur->ab_size;
               if (arena->a_cur->ab_size > DEFAULT_BLOCK_SIZE)
                       ++arena->total_big_blocks;
#endif
       }
       return p;
}

int
PyArena_AddPyObject(PyArena *arena, PyObject *obj)
{
       int r = PyList_Append(arena->a_objects, obj);
       if (r >= 0) {
               Py_DECREF(obj);
       }
       return r;
}