Branch data Line data Source code
1 : : /*
2 : : _BlocksOutputBuffer is used to maintain an output buffer
3 : : that has unpredictable size. Suitable for compression/decompression
4 : : API (bz2/lzma/zlib) that has stream->next_out and stream->avail_out:
5 : :
6 : : stream->next_out: point to the next output position.
7 : : stream->avail_out: the number of available bytes left in the buffer.
8 : :
9 : : It maintains a list of bytes object, so there is no overhead of resizing
10 : : the buffer.
11 : :
12 : : Usage:
13 : :
14 : : 1, Initialize the struct instance like this:
15 : : _BlocksOutputBuffer buffer = {.list = NULL};
16 : : Set .list to NULL for _BlocksOutputBuffer_OnError()
17 : :
18 : : 2, Initialize the buffer use one of these functions:
19 : : _BlocksOutputBuffer_InitAndGrow()
20 : : _BlocksOutputBuffer_InitWithSize()
21 : :
22 : : 3, If (avail_out == 0), grow the buffer:
23 : : _BlocksOutputBuffer_Grow()
24 : :
25 : : 4, Get the current outputted data size:
26 : : _BlocksOutputBuffer_GetDataSize()
27 : :
28 : : 5, Finish the buffer, and return a bytes object:
29 : : _BlocksOutputBuffer_Finish()
30 : :
31 : : 6, Clean up the buffer when an error occurred:
32 : : _BlocksOutputBuffer_OnError()
33 : : */
34 : :
35 : : #ifndef Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
36 : : #define Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
37 : : #ifdef __cplusplus
38 : : extern "C" {
39 : : #endif
40 : :
41 : : #include "Python.h"
42 : :
43 : : typedef struct {
44 : : // List of bytes objects
45 : : PyObject *list;
46 : : // Number of whole allocated size
47 : : Py_ssize_t allocated;
48 : : // Max length of the buffer, negative number means unlimited length.
49 : : Py_ssize_t max_length;
50 : : } _BlocksOutputBuffer;
51 : :
52 : : static const char unable_allocate_msg[] = "Unable to allocate output buffer.";
53 : :
54 : : /* In 32-bit build, the max block size should <= INT32_MAX. */
55 : : #define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)
56 : :
57 : : /* Block size sequence */
58 : : #define KB (1024)
59 : : #define MB (1024*1024)
60 : : static const Py_ssize_t BUFFER_BLOCK_SIZE[] =
61 : : { 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,
62 : : 32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,
63 : : OUTPUT_BUFFER_MAX_BLOCK_SIZE };
64 : : #undef KB
65 : : #undef MB
66 : :
67 : : /* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole
68 : : allocated size growth step is:
69 : : 1 32 KB +32 KB
70 : : 2 96 KB +64 KB
71 : : 3 352 KB +256 KB
72 : : 4 1.34 MB +1 MB
73 : : 5 5.34 MB +4 MB
74 : : 6 13.34 MB +8 MB
75 : : 7 29.34 MB +16 MB
76 : : 8 45.34 MB +16 MB
77 : : 9 77.34 MB +32 MB
78 : : 10 109.34 MB +32 MB
79 : : 11 141.34 MB +32 MB
80 : : 12 173.34 MB +32 MB
81 : : 13 237.34 MB +64 MB
82 : : 14 301.34 MB +64 MB
83 : : 15 429.34 MB +128 MB
84 : : 16 557.34 MB +128 MB
85 : : 17 813.34 MB +256 MB
86 : : 18 1069.34 MB +256 MB
87 : : 19 1325.34 MB +256 MB
88 : : 20 1581.34 MB +256 MB
89 : : 21 1837.34 MB +256 MB
90 : : 22 2093.34 MB +256 MB
91 : : ...
92 : : */
93 : :
94 : : /* Initialize the buffer, and grow the buffer.
95 : :
96 : : max_length: Max length of the buffer, -1 for unlimited length.
97 : :
98 : : On success, return allocated size (>=0)
99 : : On failure, return -1
100 : : */
101 : : static inline Py_ssize_t
102 : 144186 : _BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,
103 : : const Py_ssize_t max_length,
104 : : void **next_out)
105 : : {
106 : : PyObject *b;
107 : : Py_ssize_t block_size;
108 : :
109 : : // ensure .list was set to NULL
110 : : assert(buffer->list == NULL);
111 : :
112 : : // get block size
113 [ + + + + ]: 144186 : if (0 <= max_length && max_length < BUFFER_BLOCK_SIZE[0]) {
114 : 31258 : block_size = max_length;
115 : : } else {
116 : 112928 : block_size = BUFFER_BLOCK_SIZE[0];
117 : : }
118 : :
119 : : // the first block
120 : 144186 : b = PyBytes_FromStringAndSize(NULL, block_size);
121 [ - + ]: 144186 : if (b == NULL) {
122 : 0 : return -1;
123 : : }
124 : :
125 : : // create the list
126 : 144186 : buffer->list = PyList_New(1);
127 [ - + ]: 144186 : if (buffer->list == NULL) {
128 : 0 : Py_DECREF(b);
129 : 0 : return -1;
130 : : }
131 : 144186 : PyList_SET_ITEM(buffer->list, 0, b);
132 : :
133 : : // set variables
134 : 144186 : buffer->allocated = block_size;
135 : 144186 : buffer->max_length = max_length;
136 : :
137 : 144186 : *next_out = PyBytes_AS_STRING(b);
138 : 144186 : return block_size;
139 : : }
140 : :
141 : : /* Initialize the buffer, with an initial size.
142 : :
143 : : Check block size limit in the outer wrapper function. For example, some libs
144 : : accept UINT32_MAX as the maximum block size, then init_size should <= it.
145 : :
146 : : On success, return allocated size (>=0)
147 : : On failure, return -1
148 : : */
149 : : static inline Py_ssize_t
150 : 6611 : _BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,
151 : : const Py_ssize_t init_size,
152 : : void **next_out)
153 : : {
154 : : PyObject *b;
155 : :
156 : : // ensure .list was set to NULL
157 : : assert(buffer->list == NULL);
158 : :
159 : : // the first block
160 : 6611 : b = PyBytes_FromStringAndSize(NULL, init_size);
161 [ - + ]: 6611 : if (b == NULL) {
162 : 0 : PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
163 : 0 : return -1;
164 : : }
165 : :
166 : : // create the list
167 : 6611 : buffer->list = PyList_New(1);
168 [ - + ]: 6611 : if (buffer->list == NULL) {
169 : 0 : Py_DECREF(b);
170 : 0 : return -1;
171 : : }
172 : 6611 : PyList_SET_ITEM(buffer->list, 0, b);
173 : :
174 : : // set variables
175 : 6611 : buffer->allocated = init_size;
176 : 6611 : buffer->max_length = -1;
177 : :
178 : 6611 : *next_out = PyBytes_AS_STRING(b);
179 : 6611 : return init_size;
180 : : }
181 : :
182 : : /* Grow the buffer. The avail_out must be 0, please check it before calling.
183 : :
184 : : On success, return allocated size (>=0)
185 : : On failure, return -1
186 : : */
187 : : static inline Py_ssize_t
188 : 1849 : _BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,
189 : : void **next_out,
190 : : const Py_ssize_t avail_out)
191 : : {
192 : : PyObject *b;
193 : 1849 : const Py_ssize_t list_len = Py_SIZE(buffer->list);
194 : : Py_ssize_t block_size;
195 : :
196 : : // ensure no gaps in the data
197 [ - + ]: 1849 : if (avail_out != 0) {
198 : 0 : PyErr_SetString(PyExc_SystemError,
199 : : "avail_out is non-zero in _BlocksOutputBuffer_Grow().");
200 : 0 : return -1;
201 : : }
202 : :
203 : : // get block size
204 [ + - ]: 1849 : if (list_len < (Py_ssize_t) Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE)) {
205 : 1849 : block_size = BUFFER_BLOCK_SIZE[list_len];
206 : : } else {
207 : 0 : block_size = BUFFER_BLOCK_SIZE[Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE) - 1];
208 : : }
209 : :
210 : : // check max_length
211 [ + + ]: 1849 : if (buffer->max_length >= 0) {
212 : : // if (rest == 0), should not grow the buffer.
213 : 721 : Py_ssize_t rest = buffer->max_length - buffer->allocated;
214 : : assert(rest > 0);
215 : :
216 : : // block_size of the last block
217 [ + + ]: 721 : if (block_size > rest) {
218 : 394 : block_size = rest;
219 : : }
220 : : }
221 : :
222 : : // check buffer->allocated overflow
223 [ - + ]: 1849 : if (block_size > PY_SSIZE_T_MAX - buffer->allocated) {
224 : 0 : PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
225 : 0 : return -1;
226 : : }
227 : :
228 : : // create the block
229 : 1849 : b = PyBytes_FromStringAndSize(NULL, block_size);
230 [ - + ]: 1849 : if (b == NULL) {
231 : 0 : PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
232 : 0 : return -1;
233 : : }
234 [ - + ]: 1849 : if (PyList_Append(buffer->list, b) < 0) {
235 : 0 : Py_DECREF(b);
236 : 0 : return -1;
237 : : }
238 : 1849 : Py_DECREF(b);
239 : :
240 : : // set variables
241 : 1849 : buffer->allocated += block_size;
242 : :
243 : 1849 : *next_out = PyBytes_AS_STRING(b);
244 : 1849 : return block_size;
245 : : }
246 : :
247 : : /* Return the current outputted data size. */
248 : : static inline Py_ssize_t
249 : 30446 : _BlocksOutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer,
250 : : const Py_ssize_t avail_out)
251 : : {
252 : 30446 : return buffer->allocated - avail_out;
253 : : }
254 : :
255 : : /* Finish the buffer.
256 : :
257 : : Return a bytes object on success
258 : : Return NULL on failure
259 : : */
260 : : static inline PyObject *
261 : 150037 : _BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,
262 : : const Py_ssize_t avail_out)
263 : : {
264 : : PyObject *result, *block;
265 : 150037 : const Py_ssize_t list_len = Py_SIZE(buffer->list);
266 : :
267 : : // fast path for single block
268 [ + + + + : 150037 : if ((list_len == 1 && avail_out == 0) ||
+ + ]
269 [ - + ]: 1242 : (list_len == 2 && Py_SIZE(PyList_GET_ITEM(buffer->list, 1)) == avail_out))
270 : : {
271 : 29181 : block = PyList_GET_ITEM(buffer->list, 0);
272 : 29181 : Py_INCREF(block);
273 : :
274 [ + - ]: 29181 : Py_CLEAR(buffer->list);
275 : 29181 : return block;
276 : : }
277 : :
278 : : // final bytes object
279 : 120856 : result = PyBytes_FromStringAndSize(NULL, buffer->allocated - avail_out);
280 [ - + ]: 120856 : if (result == NULL) {
281 : 0 : PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
282 : 0 : return NULL;
283 : : }
284 : :
285 : : // memory copy
286 [ + - ]: 120856 : if (list_len > 0) {
287 : 120856 : char *posi = PyBytes_AS_STRING(result);
288 : :
289 : : // blocks except the last one
290 : 120856 : Py_ssize_t i = 0;
291 [ + + ]: 122705 : for (; i < list_len-1; i++) {
292 : 1849 : block = PyList_GET_ITEM(buffer->list, i);
293 : 1849 : memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block));
294 : 1849 : posi += Py_SIZE(block);
295 : : }
296 : : // the last block
297 : 120856 : block = PyList_GET_ITEM(buffer->list, i);
298 : 120856 : memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block) - avail_out);
299 : : } else {
300 : : assert(Py_SIZE(result) == 0);
301 : : }
302 : :
303 [ + - ]: 120856 : Py_CLEAR(buffer->list);
304 : 120856 : return result;
305 : : }
306 : :
307 : : /* Clean up the buffer when an error occurred. */
308 : : static inline void
309 : 760 : _BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)
310 : : {
311 [ + - ]: 760 : Py_CLEAR(buffer->list);
312 : 760 : }
313 : :
314 : : #ifdef __cplusplus
315 : : }
316 : : #endif
317 : : #endif /* Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H */
|
