#include "Python.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

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.  Its

       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 arena.  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 *)PyMem_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 = (char *)_Py_ALIGN_UP(b->ab_mem, ALIGNMENT) -

            (char *)(b->ab_mem);

    return b;

}

static void

block_free(block *b) {

    while (b) {

        block *next = b->ab_next;

        PyMem_Free(b);

        b = next;

    }

}

static void *

block_alloc(block *b, size_t size)

{

    void *p;

    assert(b);

    size = _Py_SIZE_ROUND_UP(size, ALIGNMENT);

    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 *)PyMem_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) {

        PyMem_Free((void *)arena);

        return (PyArena*)PyErr_NoMemory();

    }

    arena->a_objects = PyList_New(0);

    if (!arena->a_objects) {

        block_free(arena->a_head);

        PyMem_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)

{

    assert(arena);

#if defined(Py_DEBUG)

    /*

    fprintf(stderr,

        "alloc=%zu size=%zu blocks=%zu block_size=%zu big=%zu objects=%zu\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);

    */

    Py_DECREF(arena->a_objects);

    PyMem_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;

}