Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-11-19 09:50:45

0001 /*
0002    _BlocksOutputBuffer is used to maintain an output buffer
0003    that has unpredictable size. Suitable for compression/decompression
0004    API (bz2/lzma/zlib) that has stream->next_out and stream->avail_out:
0005 
0006         stream->next_out:  point to the next output position.
0007         stream->avail_out: the number of available bytes left in the buffer.
0008 
0009    It maintains a list of bytes object, so there is no overhead of resizing
0010    the buffer.
0011 
0012    Usage:
0013 
0014    1, Initialize the struct instance like this:
0015         _BlocksOutputBuffer buffer = {.list = NULL};
0016       Set .list to NULL for _BlocksOutputBuffer_OnError()
0017 
0018    2, Initialize the buffer use one of these functions:
0019         _BlocksOutputBuffer_InitAndGrow()
0020         _BlocksOutputBuffer_InitWithSize()
0021 
0022    3, If (avail_out == 0), grow the buffer:
0023         _BlocksOutputBuffer_Grow()
0024 
0025    4, Get the current outputted data size:
0026         _BlocksOutputBuffer_GetDataSize()
0027 
0028    5, Finish the buffer, and return a bytes object:
0029         _BlocksOutputBuffer_Finish()
0030 
0031    6, Clean up the buffer when an error occurred:
0032         _BlocksOutputBuffer_OnError()
0033 */
0034 
0035 #ifndef Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
0036 #define Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
0037 #ifdef __cplusplus
0038 extern "C" {
0039 #endif
0040 
0041 #include "Python.h"
0042 
0043 #ifndef Py_BUILD_CORE
0044 #  error "this header requires Py_BUILD_CORE define"
0045 #endif
0046 
0047 typedef struct {
0048     // List of bytes objects
0049     PyObject *list;
0050     // Number of whole allocated size
0051     Py_ssize_t allocated;
0052     // Max length of the buffer, negative number means unlimited length.
0053     Py_ssize_t max_length;
0054 } _BlocksOutputBuffer;
0055 
0056 static const char unable_allocate_msg[] = "Unable to allocate output buffer.";
0057 
0058 /* In 32-bit build, the max block size should <= INT32_MAX. */
0059 #define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)
0060 
0061 /* Block size sequence */
0062 #define KB (1024)
0063 #define MB (1024*1024)
0064 static const Py_ssize_t BUFFER_BLOCK_SIZE[] =
0065     { 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,
0066       32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,
0067       OUTPUT_BUFFER_MAX_BLOCK_SIZE };
0068 #undef KB
0069 #undef MB
0070 
0071 /* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole
0072    allocated size growth step is:
0073     1   32 KB       +32 KB
0074     2   96 KB       +64 KB
0075     3   352 KB      +256 KB
0076     4   1.34 MB     +1 MB
0077     5   5.34 MB     +4 MB
0078     6   13.34 MB    +8 MB
0079     7   29.34 MB    +16 MB
0080     8   45.34 MB    +16 MB
0081     9   77.34 MB    +32 MB
0082     10  109.34 MB   +32 MB
0083     11  141.34 MB   +32 MB
0084     12  173.34 MB   +32 MB
0085     13  237.34 MB   +64 MB
0086     14  301.34 MB   +64 MB
0087     15  429.34 MB   +128 MB
0088     16  557.34 MB   +128 MB
0089     17  813.34 MB   +256 MB
0090     18  1069.34 MB  +256 MB
0091     19  1325.34 MB  +256 MB
0092     20  1581.34 MB  +256 MB
0093     21  1837.34 MB  +256 MB
0094     22  2093.34 MB  +256 MB
0095     ...
0096 */
0097 
0098 /* Initialize the buffer, and grow the buffer.
0099 
0100    max_length: Max length of the buffer, -1 for unlimited length.
0101 
0102    On success, return allocated size (>=0)
0103    On failure, return -1
0104 */
0105 static inline Py_ssize_t
0106 _BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,
0107                                 const Py_ssize_t max_length,
0108                                 void **next_out)
0109 {
0110     PyObject *b;
0111     Py_ssize_t block_size;
0112 
0113     // ensure .list was set to NULL
0114     assert(buffer->list == NULL);
0115 
0116     // get block size
0117     if (0 <= max_length && max_length < BUFFER_BLOCK_SIZE[0]) {
0118         block_size = max_length;
0119     } else {
0120         block_size = BUFFER_BLOCK_SIZE[0];
0121     }
0122 
0123     // the first block
0124     b = PyBytes_FromStringAndSize(NULL, block_size);
0125     if (b == NULL) {
0126         return -1;
0127     }
0128 
0129     // create the list
0130     buffer->list = PyList_New(1);
0131     if (buffer->list == NULL) {
0132         Py_DECREF(b);
0133         return -1;
0134     }
0135     PyList_SET_ITEM(buffer->list, 0, b);
0136 
0137     // set variables
0138     buffer->allocated = block_size;
0139     buffer->max_length = max_length;
0140 
0141     *next_out = PyBytes_AS_STRING(b);
0142     return block_size;
0143 }
0144 
0145 /* Initialize the buffer, with an initial size.
0146 
0147    Check block size limit in the outer wrapper function. For example, some libs
0148    accept UINT32_MAX as the maximum block size, then init_size should <= it.
0149 
0150    On success, return allocated size (>=0)
0151    On failure, return -1
0152 */
0153 static inline Py_ssize_t
0154 _BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,
0155                                  const Py_ssize_t init_size,
0156                                  void **next_out)
0157 {
0158     PyObject *b;
0159 
0160     // ensure .list was set to NULL
0161     assert(buffer->list == NULL);
0162 
0163     // the first block
0164     b = PyBytes_FromStringAndSize(NULL, init_size);
0165     if (b == NULL) {
0166         PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
0167         return -1;
0168     }
0169 
0170     // create the list
0171     buffer->list = PyList_New(1);
0172     if (buffer->list == NULL) {
0173         Py_DECREF(b);
0174         return -1;
0175     }
0176     PyList_SET_ITEM(buffer->list, 0, b);
0177 
0178     // set variables
0179     buffer->allocated = init_size;
0180     buffer->max_length = -1;
0181 
0182     *next_out = PyBytes_AS_STRING(b);
0183     return init_size;
0184 }
0185 
0186 /* Grow the buffer. The avail_out must be 0, please check it before calling.
0187 
0188    On success, return allocated size (>=0)
0189    On failure, return -1
0190 */
0191 static inline Py_ssize_t
0192 _BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,
0193                          void **next_out,
0194                          const Py_ssize_t avail_out)
0195 {
0196     PyObject *b;
0197     const Py_ssize_t list_len = Py_SIZE(buffer->list);
0198     Py_ssize_t block_size;
0199 
0200     // ensure no gaps in the data
0201     if (avail_out != 0) {
0202         PyErr_SetString(PyExc_SystemError,
0203                         "avail_out is non-zero in _BlocksOutputBuffer_Grow().");
0204         return -1;
0205     }
0206 
0207     // get block size
0208     if (list_len < (Py_ssize_t) Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE)) {
0209         block_size = BUFFER_BLOCK_SIZE[list_len];
0210     } else {
0211         block_size = BUFFER_BLOCK_SIZE[Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE) - 1];
0212     }
0213 
0214     // check max_length
0215     if (buffer->max_length >= 0) {
0216         // if (rest == 0), should not grow the buffer.
0217         Py_ssize_t rest = buffer->max_length - buffer->allocated;
0218         assert(rest > 0);
0219 
0220         // block_size of the last block
0221         if (block_size > rest) {
0222             block_size = rest;
0223         }
0224     }
0225 
0226     // check buffer->allocated overflow
0227     if (block_size > PY_SSIZE_T_MAX - buffer->allocated) {
0228         PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
0229         return -1;
0230     }
0231 
0232     // create the block
0233     b = PyBytes_FromStringAndSize(NULL, block_size);
0234     if (b == NULL) {
0235         PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
0236         return -1;
0237     }
0238     if (PyList_Append(buffer->list, b) < 0) {
0239         Py_DECREF(b);
0240         return -1;
0241     }
0242     Py_DECREF(b);
0243 
0244     // set variables
0245     buffer->allocated += block_size;
0246 
0247     *next_out = PyBytes_AS_STRING(b);
0248     return block_size;
0249 }
0250 
0251 /* Return the current outputted data size. */
0252 static inline Py_ssize_t
0253 _BlocksOutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer,
0254                                 const Py_ssize_t avail_out)
0255 {
0256     return buffer->allocated - avail_out;
0257 }
0258 
0259 /* Finish the buffer.
0260 
0261    Return a bytes object on success
0262    Return NULL on failure
0263 */
0264 static inline PyObject *
0265 _BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,
0266                            const Py_ssize_t avail_out)
0267 {
0268     PyObject *result, *block;
0269     const Py_ssize_t list_len = Py_SIZE(buffer->list);
0270 
0271     // fast path for single block
0272     if ((list_len == 1 && avail_out == 0) ||
0273         (list_len == 2 && Py_SIZE(PyList_GET_ITEM(buffer->list, 1)) == avail_out))
0274     {
0275         block = PyList_GET_ITEM(buffer->list, 0);
0276         Py_INCREF(block);
0277 
0278         Py_CLEAR(buffer->list);
0279         return block;
0280     }
0281 
0282     // final bytes object
0283     result = PyBytes_FromStringAndSize(NULL, buffer->allocated - avail_out);
0284     if (result == NULL) {
0285         PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
0286         return NULL;
0287     }
0288 
0289     // memory copy
0290     if (list_len > 0) {
0291         char *posi = PyBytes_AS_STRING(result);
0292 
0293         // blocks except the last one
0294         Py_ssize_t i = 0;
0295         for (; i < list_len-1; i++) {
0296             block = PyList_GET_ITEM(buffer->list, i);
0297             memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block));
0298             posi += Py_SIZE(block);
0299         }
0300         // the last block
0301         block = PyList_GET_ITEM(buffer->list, i);
0302         memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block) - avail_out);
0303     } else {
0304         assert(Py_SIZE(result) == 0);
0305     }
0306 
0307     Py_CLEAR(buffer->list);
0308     return result;
0309 }
0310 
0311 /* Clean up the buffer when an error occurred. */
0312 static inline void
0313 _BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)
0314 {
0315     Py_CLEAR(buffer->list);
0316 }
0317 
0318 #ifdef __cplusplus
0319 }
0320 #endif
0321 #endif /* Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H */