JezK
Edit File: pycore_blocks_output_buffer.h
/* _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 */