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

}