/*

_BlocksOutputBuffer is used to maintain an output buffer

that has unpredictable size. Suitable for compression/decompression

API (bz2/lzma/zlib) that has stream->next_out and stream->avail_out:

stream->next_out: point to the next output position.

stream->avail_out: the number of available bytes left in the buffer.

It maintains a list of bytes object, so there is no overhead of resizing

the buffer.

Usage:

1, Initialize the struct instance like this:

_BlocksOutputBuffer buffer = {.list = NULL};

Set .list to NULL for _BlocksOutputBuffer_OnError()

2, Initialize the buffer use one of these functions:

_BlocksOutputBuffer_InitAndGrow()

_BlocksOutputBuffer_InitWithSize()

3, If (avail_out == 0), grow the buffer:

_BlocksOutputBuffer_Grow()

4, Get the current outputted data size:

_BlocksOutputBuffer_GetDataSize()

5, Finish the buffer, and return a bytes object:

_BlocksOutputBuffer_Finish()

6, Clean up the buffer when an error occurred:

_BlocksOutputBuffer_OnError()

*/

#ifndef Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H

#define Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H

#ifdef __cplusplus

extern "C" {

#endif

#include "Python.h"

typedef struct {

// List of bytes objects

PyObject *list;

// Number of whole allocated size

Py_ssize_t allocated;

// Max length of the buffer, negative number means unlimited length.

Py_ssize_t max_length;

} _BlocksOutputBuffer;

static const char unable_allocate_msg[] = "Unable to allocate output buffer.";

/* In 32-bit build, the max block size should <= INT32_MAX. */

#define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)

/* Block size sequence */

#define KB (1024)

#define MB (1024*1024)

static const Py_ssize_t BUFFER_BLOCK_SIZE[] =

{ 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,

32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,

OUTPUT_BUFFER_MAX_BLOCK_SIZE };

#undef KB

#undef MB

/* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole

allocated size growth step is:

1 32 KB +32 KB

2 96 KB +64 KB

3 352 KB +256 KB

4 1.34 MB +1 MB

5 5.34 MB +4 MB

6 13.34 MB +8 MB

7 29.34 MB +16 MB

8 45.34 MB +16 MB

9 77.34 MB +32 MB

10 109.34 MB +32 MB

11 141.34 MB +32 MB

12 173.34 MB +32 MB

13 237.34 MB +64 MB

14 301.34 MB +64 MB

15 429.34 MB +128 MB

16 557.34 MB +128 MB

17 813.34 MB +256 MB

18 1069.34 MB +256 MB

19 1325.34 MB +256 MB

20 1581.34 MB +256 MB

21 1837.34 MB +256 MB

22 2093.34 MB +256 MB

...

*/

/* Initialize the buffer, and grow the buffer.

max_length: Max length of the buffer, -1 for unlimited length.

On success, return allocated size (>=0)

On failure, return -1

*/

static inline Py_ssize_t

_BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,

const Py_ssize_t max_length,

void **next_out)

{

PyObject *b;

Py_ssize_t block_size;

// ensure .list was set to NULL

assert(buffer->list == NULL);

// get block size

if (0 <= max_length && max_length < BUFFER_BLOCK_SIZE[0]) {

block_size = max_length;

} else {

block_size = BUFFER_BLOCK_SIZE[0];

}

// the first block

b = PyBytes_FromStringAndSize(NULL, block_size);

if (b == NULL) {

return -1;

}

// create the list

buffer->list = PyList_New(1);

if (buffer->list == NULL) {

Py_DECREF(b);

return -1;

}

PyList_SET_ITEM(buffer->list, 0, b);

// set variables

buffer->allocated = block_size;

buffer->max_length = max_length;

*next_out = PyBytes_AS_STRING(b);

return block_size;

}

/* Initialize the buffer, with an initial size.

Check block size limit in the outer wrapper function. For example, some libs

accept UINT32_MAX as the maximum block size, then init_size should <= it.

On success, return allocated size (>=0)

On failure, return -1

*/

static inline Py_ssize_t

_BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,

const Py_ssize_t init_size,

void **next_out)

{

PyObject *b;

// ensure .list was set to NULL

assert(buffer->list == NULL);

// the first block

b = PyBytes_FromStringAndSize(NULL, init_size);

if (b == NULL) {

PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);

return -1;

}

// create the list

buffer->list = PyList_New(1);

if (buffer->list == NULL) {

Py_DECREF(b);

return -1;

}

PyList_SET_ITEM(buffer->list, 0, b);

// set variables

buffer->allocated = init_size;

buffer->max_length = -1;

*next_out = PyBytes_AS_STRING(b);

return init_size;

}

/* Grow the buffer. The avail_out must be 0, please check it before calling.

On success, return allocated size (>=0)

On failure, return -1

*/

static inline Py_ssize_t

_BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,

void **next_out,

const Py_ssize_t avail_out)

{

PyObject *b;

const Py_ssize_t list_len = Py_SIZE(buffer->list);

Py_ssize_t block_size;

// ensure no gaps in the data

if (avail_out != 0) {

PyErr_SetString(PyExc_SystemError,

"avail_out is non-zero in _BlocksOutputBuffer_Grow().");

return -1;

}

// get block size

if (list_len < (Py_ssize_t) Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE)) {

block_size = BUFFER_BLOCK_SIZE[list_len];

} else {

block_size = BUFFER_BLOCK_SIZE[Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE) - 1];

}

// check max_length

if (buffer->max_length >= 0) {

// if (rest == 0), should not grow the buffer.

Py_ssize_t rest = buffer->max_length - buffer->allocated;

assert(rest > 0);

// block_size of the last block

if (block_size > rest) {

block_size = rest;

}

}

// check buffer->allocated overflow

if (block_size > PY_SSIZE_T_MAX - buffer->allocated) {

PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);

return -1;

}

// create the block

b = PyBytes_FromStringAndSize(NULL, block_size);

if (b == NULL) {

PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);

return -1;

}

if (PyList_Append(buffer->list, b) < 0) {

Py_DECREF(b);

return -1;

}

Py_DECREF(b);

// set variables

buffer->allocated += block_size;

*next_out = PyBytes_AS_STRING(b);

return block_size;

}

/* Return the current outputted data size. */

static inline Py_ssize_t

_BlocksOutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer,

const Py_ssize_t avail_out)

{

return buffer->allocated - avail_out;

}

/* Finish the buffer.

Return a bytes object on success

Return NULL on failure

*/

static inline PyObject *

_BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,

const Py_ssize_t avail_out)

{

PyObject *result, *block;

const Py_ssize_t list_len = Py_SIZE(buffer->list);

// fast path for single block

if ((list_len == 1 && avail_out == 0) ||

(list_len == 2 && Py_SIZE(PyList_GET_ITEM(buffer->list, 1)) == avail_out))

{

block = PyList_GET_ITEM(buffer->list, 0);

Py_INCREF(block);

Py_CLEAR(buffer->list);

return block;

}

// final bytes object

result = PyBytes_FromStringAndSize(NULL, buffer->allocated - avail_out);

if (result == NULL) {

PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);

return NULL;

}

// memory copy

if (list_len > 0) {

char *posi = PyBytes_AS_STRING(result);

// blocks except the last one

Py_ssize_t i = 0;

for (; i < list_len-1; i++) {

block = PyList_GET_ITEM(buffer->list, i);

memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block));

posi += Py_SIZE(block);

}

// the last block

block = PyList_GET_ITEM(buffer->list, i);

memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block) - avail_out);

} else {

assert(Py_SIZE(result) == 0);

}

Py_CLEAR(buffer->list);

return result;

}

/* Clean up the buffer when an error occurred. */

static inline void

_BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)

{

Py_CLEAR(buffer->list);

}

#ifdef __cplusplus

}

#endif

#endif /* Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H */