mirror of https://github.com/python/cpython.git
322 lines
8.6 KiB
C
322 lines
8.6 KiB
C
/*
|
|
_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"
|
|
|
|
#ifndef Py_BUILD_CORE
|
|
# error "this header requires Py_BUILD_CORE define"
|
|
#endif
|
|
|
|
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 */
|