Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-30 10:18:05

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