bufferutil.c 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792
  1. /**
  2. * Copyright (c) 2017-present, Gregory Szorc
  3. * All rights reserved.
  4. *
  5. * This software may be modified and distributed under the terms
  6. * of the BSD license. See the LICENSE file for details.
  7. */
  8. #include "python-zstandard.h"
  9. extern PyObject* ZstdError;
  10. PyDoc_STRVAR(BufferWithSegments__doc__,
  11. "BufferWithSegments - A memory buffer holding known sub-segments.\n"
  12. "\n"
  13. "This type represents a contiguous chunk of memory containing N discrete\n"
  14. "items within sub-segments of that memory.\n"
  15. "\n"
  16. "Segments within the buffer are stored as an array of\n"
  17. "``(offset, length)`` pairs, where each element is an unsigned 64-bit\n"
  18. "integer using the host/native bit order representation.\n"
  19. "\n"
  20. "The type exists to facilitate operations against N>1 items without the\n"
  21. "overhead of Python object creation and management.\n"
  22. );
  23. static void BufferWithSegments_dealloc(ZstdBufferWithSegments* self) {
  24. /* Backing memory is either canonically owned by a Py_buffer or by us. */
  25. if (self->parent.buf) {
  26. PyBuffer_Release(&self->parent);
  27. }
  28. else if (self->useFree) {
  29. free(self->data);
  30. }
  31. else {
  32. PyMem_Free(self->data);
  33. }
  34. self->data = NULL;
  35. if (self->useFree) {
  36. free(self->segments);
  37. }
  38. else {
  39. PyMem_Free(self->segments);
  40. }
  41. self->segments = NULL;
  42. PyObject_Del(self);
  43. }
  44. static int BufferWithSegments_init(ZstdBufferWithSegments* self, PyObject* args, PyObject* kwargs) {
  45. static char* kwlist[] = {
  46. "data",
  47. "segments",
  48. NULL
  49. };
  50. Py_buffer segments;
  51. Py_ssize_t segmentCount;
  52. Py_ssize_t i;
  53. memset(&self->parent, 0, sizeof(self->parent));
  54. #if PY_MAJOR_VERSION >= 3
  55. if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*y*:BufferWithSegments",
  56. #else
  57. if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*s*:BufferWithSegments",
  58. #endif
  59. kwlist, &self->parent, &segments)) {
  60. return -1;
  61. }
  62. if (!PyBuffer_IsContiguous(&self->parent, 'C') || self->parent.ndim > 1) {
  63. PyErr_SetString(PyExc_ValueError, "data buffer should be contiguous and have a single dimension");
  64. goto except;
  65. }
  66. if (!PyBuffer_IsContiguous(&segments, 'C') || segments.ndim > 1) {
  67. PyErr_SetString(PyExc_ValueError, "segments buffer should be contiguous and have a single dimension");
  68. goto except;
  69. }
  70. if (segments.len % sizeof(BufferSegment)) {
  71. PyErr_Format(PyExc_ValueError, "segments array size is not a multiple of %zu",
  72. sizeof(BufferSegment));
  73. goto except;
  74. }
  75. segmentCount = segments.len / sizeof(BufferSegment);
  76. /* Validate segments data, as blindly trusting it could lead to arbitrary
  77. memory access. */
  78. for (i = 0; i < segmentCount; i++) {
  79. BufferSegment* segment = &((BufferSegment*)(segments.buf))[i];
  80. if (segment->offset + segment->length > (unsigned long long)self->parent.len) {
  81. PyErr_SetString(PyExc_ValueError, "offset within segments array references memory outside buffer");
  82. goto except;
  83. return -1;
  84. }
  85. }
  86. /* Make a copy of the segments data. It is cheap to do so and is a guard
  87. against caller changing offsets, which has security implications. */
  88. self->segments = PyMem_Malloc(segments.len);
  89. if (!self->segments) {
  90. PyErr_NoMemory();
  91. goto except;
  92. }
  93. memcpy(self->segments, segments.buf, segments.len);
  94. PyBuffer_Release(&segments);
  95. self->data = self->parent.buf;
  96. self->dataSize = self->parent.len;
  97. self->segmentCount = segmentCount;
  98. return 0;
  99. except:
  100. PyBuffer_Release(&self->parent);
  101. PyBuffer_Release(&segments);
  102. return -1;
  103. }
  104. /**
  105. * Construct a BufferWithSegments from existing memory and offsets.
  106. *
  107. * Ownership of the backing memory and BufferSegments will be transferred to
  108. * the created object and freed when the BufferWithSegments is destroyed.
  109. */
  110. ZstdBufferWithSegments* BufferWithSegments_FromMemory(void* data, unsigned long long dataSize,
  111. BufferSegment* segments, Py_ssize_t segmentsSize) {
  112. ZstdBufferWithSegments* result = NULL;
  113. Py_ssize_t i;
  114. if (NULL == data) {
  115. PyErr_SetString(PyExc_ValueError, "data is NULL");
  116. return NULL;
  117. }
  118. if (NULL == segments) {
  119. PyErr_SetString(PyExc_ValueError, "segments is NULL");
  120. return NULL;
  121. }
  122. for (i = 0; i < segmentsSize; i++) {
  123. BufferSegment* segment = &segments[i];
  124. if (segment->offset + segment->length > dataSize) {
  125. PyErr_SetString(PyExc_ValueError, "offset in segments overflows buffer size");
  126. return NULL;
  127. }
  128. }
  129. result = PyObject_New(ZstdBufferWithSegments, &ZstdBufferWithSegmentsType);
  130. if (NULL == result) {
  131. return NULL;
  132. }
  133. result->useFree = 0;
  134. memset(&result->parent, 0, sizeof(result->parent));
  135. result->data = data;
  136. result->dataSize = dataSize;
  137. result->segments = segments;
  138. result->segmentCount = segmentsSize;
  139. return result;
  140. }
  141. static Py_ssize_t BufferWithSegments_length(ZstdBufferWithSegments* self) {
  142. return self->segmentCount;
  143. }
  144. static ZstdBufferSegment* BufferWithSegments_item(ZstdBufferWithSegments* self, Py_ssize_t i) {
  145. ZstdBufferSegment* result = NULL;
  146. if (i < 0) {
  147. PyErr_SetString(PyExc_IndexError, "offset must be non-negative");
  148. return NULL;
  149. }
  150. if (i >= self->segmentCount) {
  151. PyErr_Format(PyExc_IndexError, "offset must be less than %zd", self->segmentCount);
  152. return NULL;
  153. }
  154. if (self->segments[i].length > PY_SSIZE_T_MAX) {
  155. PyErr_Format(PyExc_ValueError,
  156. "item at offset %zd is too large for this platform", i);
  157. return NULL;
  158. }
  159. result = (ZstdBufferSegment*)PyObject_CallObject((PyObject*)&ZstdBufferSegmentType, NULL);
  160. if (NULL == result) {
  161. return NULL;
  162. }
  163. result->parent = (PyObject*)self;
  164. Py_INCREF(self);
  165. result->data = (char*)self->data + self->segments[i].offset;
  166. result->dataSize = (Py_ssize_t)self->segments[i].length;
  167. result->offset = self->segments[i].offset;
  168. return result;
  169. }
  170. #if PY_MAJOR_VERSION >= 3
  171. static int BufferWithSegments_getbuffer(ZstdBufferWithSegments* self, Py_buffer* view, int flags) {
  172. if (self->dataSize > PY_SSIZE_T_MAX) {
  173. view->obj = NULL;
  174. PyErr_SetString(PyExc_BufferError, "buffer is too large for this platform");
  175. return -1;
  176. }
  177. return PyBuffer_FillInfo(view, (PyObject*)self, self->data, (Py_ssize_t)self->dataSize, 1, flags);
  178. }
  179. #else
  180. static Py_ssize_t BufferWithSegments_getreadbuffer(ZstdBufferWithSegments* self, Py_ssize_t segment, void **ptrptr) {
  181. if (segment != 0) {
  182. PyErr_SetString(PyExc_ValueError, "segment number must be 0");
  183. return -1;
  184. }
  185. if (self->dataSize > PY_SSIZE_T_MAX) {
  186. PyErr_SetString(PyExc_ValueError, "buffer is too large for this platform");
  187. return -1;
  188. }
  189. *ptrptr = self->data;
  190. return (Py_ssize_t)self->dataSize;
  191. }
  192. static Py_ssize_t BufferWithSegments_getsegcount(ZstdBufferWithSegments* self, Py_ssize_t* len) {
  193. if (len) {
  194. *len = 1;
  195. }
  196. return 1;
  197. }
  198. #endif
  199. PyDoc_STRVAR(BufferWithSegments_tobytes__doc__,
  200. "Obtain a bytes instance for this buffer.\n"
  201. );
  202. static PyObject* BufferWithSegments_tobytes(ZstdBufferWithSegments* self) {
  203. if (self->dataSize > PY_SSIZE_T_MAX) {
  204. PyErr_SetString(PyExc_ValueError, "buffer is too large for this platform");
  205. return NULL;
  206. }
  207. return PyBytes_FromStringAndSize(self->data, (Py_ssize_t)self->dataSize);
  208. }
  209. PyDoc_STRVAR(BufferWithSegments_segments__doc__,
  210. "Obtain a BufferSegments describing segments in this sintance.\n"
  211. );
  212. static ZstdBufferSegments* BufferWithSegments_segments(ZstdBufferWithSegments* self) {
  213. ZstdBufferSegments* result = (ZstdBufferSegments*)PyObject_CallObject((PyObject*)&ZstdBufferSegmentsType, NULL);
  214. if (NULL == result) {
  215. return NULL;
  216. }
  217. result->parent = (PyObject*)self;
  218. Py_INCREF(self);
  219. result->segments = self->segments;
  220. result->segmentCount = self->segmentCount;
  221. return result;
  222. }
  223. static PySequenceMethods BufferWithSegments_sq = {
  224. (lenfunc)BufferWithSegments_length, /* sq_length */
  225. 0, /* sq_concat */
  226. 0, /* sq_repeat */
  227. (ssizeargfunc)BufferWithSegments_item, /* sq_item */
  228. 0, /* sq_ass_item */
  229. 0, /* sq_contains */
  230. 0, /* sq_inplace_concat */
  231. 0 /* sq_inplace_repeat */
  232. };
  233. static PyBufferProcs BufferWithSegments_as_buffer = {
  234. #if PY_MAJOR_VERSION >= 3
  235. (getbufferproc)BufferWithSegments_getbuffer, /* bf_getbuffer */
  236. 0 /* bf_releasebuffer */
  237. #else
  238. (readbufferproc)BufferWithSegments_getreadbuffer, /* bf_getreadbuffer */
  239. 0, /* bf_getwritebuffer */
  240. (segcountproc)BufferWithSegments_getsegcount, /* bf_getsegcount */
  241. 0 /* bf_getcharbuffer */
  242. #endif
  243. };
  244. static PyMethodDef BufferWithSegments_methods[] = {
  245. { "segments", (PyCFunction)BufferWithSegments_segments,
  246. METH_NOARGS, BufferWithSegments_segments__doc__ },
  247. { "tobytes", (PyCFunction)BufferWithSegments_tobytes,
  248. METH_NOARGS, BufferWithSegments_tobytes__doc__ },
  249. { NULL, NULL }
  250. };
  251. static PyMemberDef BufferWithSegments_members[] = {
  252. { "size", T_ULONGLONG, offsetof(ZstdBufferWithSegments, dataSize),
  253. READONLY, "total size of the buffer in bytes" },
  254. { NULL }
  255. };
  256. PyTypeObject ZstdBufferWithSegmentsType = {
  257. PyVarObject_HEAD_INIT(NULL, 0)
  258. "zstd.BufferWithSegments", /* tp_name */
  259. sizeof(ZstdBufferWithSegments),/* tp_basicsize */
  260. 0, /* tp_itemsize */
  261. (destructor)BufferWithSegments_dealloc, /* tp_dealloc */
  262. 0, /* tp_print */
  263. 0, /* tp_getattr */
  264. 0, /* tp_setattr */
  265. 0, /* tp_compare */
  266. 0, /* tp_repr */
  267. 0, /* tp_as_number */
  268. &BufferWithSegments_sq, /* tp_as_sequence */
  269. 0, /* tp_as_mapping */
  270. 0, /* tp_hash */
  271. 0, /* tp_call */
  272. 0, /* tp_str */
  273. 0, /* tp_getattro */
  274. 0, /* tp_setattro */
  275. &BufferWithSegments_as_buffer, /* tp_as_buffer */
  276. Py_TPFLAGS_DEFAULT, /* tp_flags */
  277. BufferWithSegments__doc__, /* tp_doc */
  278. 0, /* tp_traverse */
  279. 0, /* tp_clear */
  280. 0, /* tp_richcompare */
  281. 0, /* tp_weaklistoffset */
  282. 0, /* tp_iter */
  283. 0, /* tp_iternext */
  284. BufferWithSegments_methods, /* tp_methods */
  285. BufferWithSegments_members, /* tp_members */
  286. 0, /* tp_getset */
  287. 0, /* tp_base */
  288. 0, /* tp_dict */
  289. 0, /* tp_descr_get */
  290. 0, /* tp_descr_set */
  291. 0, /* tp_dictoffset */
  292. (initproc)BufferWithSegments_init, /* tp_init */
  293. 0, /* tp_alloc */
  294. PyType_GenericNew, /* tp_new */
  295. };
  296. PyDoc_STRVAR(BufferSegments__doc__,
  297. "BufferSegments - Represents segments/offsets within a BufferWithSegments\n"
  298. );
  299. static void BufferSegments_dealloc(ZstdBufferSegments* self) {
  300. Py_CLEAR(self->parent);
  301. PyObject_Del(self);
  302. }
  303. #if PY_MAJOR_VERSION >= 3
  304. static int BufferSegments_getbuffer(ZstdBufferSegments* self, Py_buffer* view, int flags) {
  305. return PyBuffer_FillInfo(view, (PyObject*)self,
  306. (void*)self->segments, self->segmentCount * sizeof(BufferSegment),
  307. 1, flags);
  308. }
  309. #else
  310. static Py_ssize_t BufferSegments_getreadbuffer(ZstdBufferSegments* self, Py_ssize_t segment, void **ptrptr) {
  311. if (segment != 0) {
  312. PyErr_SetString(PyExc_ValueError, "segment number must be 0");
  313. return -1;
  314. }
  315. *ptrptr = (void*)self->segments;
  316. return self->segmentCount * sizeof(BufferSegment);
  317. }
  318. static Py_ssize_t BufferSegments_getsegcount(ZstdBufferSegments* self, Py_ssize_t* len) {
  319. if (len) {
  320. *len = 1;
  321. }
  322. return 1;
  323. }
  324. #endif
  325. static PyBufferProcs BufferSegments_as_buffer = {
  326. #if PY_MAJOR_VERSION >= 3
  327. (getbufferproc)BufferSegments_getbuffer,
  328. 0
  329. #else
  330. (readbufferproc)BufferSegments_getreadbuffer,
  331. 0,
  332. (segcountproc)BufferSegments_getsegcount,
  333. 0
  334. #endif
  335. };
  336. PyTypeObject ZstdBufferSegmentsType = {
  337. PyVarObject_HEAD_INIT(NULL, 0)
  338. "zstd.BufferSegments", /* tp_name */
  339. sizeof(ZstdBufferSegments),/* tp_basicsize */
  340. 0, /* tp_itemsize */
  341. (destructor)BufferSegments_dealloc, /* tp_dealloc */
  342. 0, /* tp_print */
  343. 0, /* tp_getattr */
  344. 0, /* tp_setattr */
  345. 0, /* tp_compare */
  346. 0, /* tp_repr */
  347. 0, /* tp_as_number */
  348. 0, /* tp_as_sequence */
  349. 0, /* tp_as_mapping */
  350. 0, /* tp_hash */
  351. 0, /* tp_call */
  352. 0, /* tp_str */
  353. 0, /* tp_getattro */
  354. 0, /* tp_setattro */
  355. &BufferSegments_as_buffer, /* tp_as_buffer */
  356. Py_TPFLAGS_DEFAULT, /* tp_flags */
  357. BufferSegments__doc__, /* tp_doc */
  358. 0, /* tp_traverse */
  359. 0, /* tp_clear */
  360. 0, /* tp_richcompare */
  361. 0, /* tp_weaklistoffset */
  362. 0, /* tp_iter */
  363. 0, /* tp_iternext */
  364. 0, /* tp_methods */
  365. 0, /* tp_members */
  366. 0, /* tp_getset */
  367. 0, /* tp_base */
  368. 0, /* tp_dict */
  369. 0, /* tp_descr_get */
  370. 0, /* tp_descr_set */
  371. 0, /* tp_dictoffset */
  372. 0, /* tp_init */
  373. 0, /* tp_alloc */
  374. PyType_GenericNew, /* tp_new */
  375. };
  376. PyDoc_STRVAR(BufferSegment__doc__,
  377. "BufferSegment - Represents a segment within a BufferWithSegments\n"
  378. );
  379. static void BufferSegment_dealloc(ZstdBufferSegment* self) {
  380. Py_CLEAR(self->parent);
  381. PyObject_Del(self);
  382. }
  383. static Py_ssize_t BufferSegment_length(ZstdBufferSegment* self) {
  384. return self->dataSize;
  385. }
  386. #if PY_MAJOR_VERSION >= 3
  387. static int BufferSegment_getbuffer(ZstdBufferSegment* self, Py_buffer* view, int flags) {
  388. return PyBuffer_FillInfo(view, (PyObject*)self,
  389. self->data, self->dataSize, 1, flags);
  390. }
  391. #else
  392. static Py_ssize_t BufferSegment_getreadbuffer(ZstdBufferSegment* self, Py_ssize_t segment, void **ptrptr) {
  393. if (segment != 0) {
  394. PyErr_SetString(PyExc_ValueError, "segment number must be 0");
  395. return -1;
  396. }
  397. *ptrptr = self->data;
  398. return self->dataSize;
  399. }
  400. static Py_ssize_t BufferSegment_getsegcount(ZstdBufferSegment* self, Py_ssize_t* len) {
  401. if (len) {
  402. *len = 1;
  403. }
  404. return 1;
  405. }
  406. #endif
  407. PyDoc_STRVAR(BufferSegment_tobytes__doc__,
  408. "Obtain a bytes instance for this segment.\n"
  409. );
  410. static PyObject* BufferSegment_tobytes(ZstdBufferSegment* self) {
  411. return PyBytes_FromStringAndSize(self->data, self->dataSize);
  412. }
  413. static PySequenceMethods BufferSegment_sq = {
  414. (lenfunc)BufferSegment_length, /* sq_length */
  415. 0, /* sq_concat */
  416. 0, /* sq_repeat */
  417. 0, /* sq_item */
  418. 0, /* sq_ass_item */
  419. 0, /* sq_contains */
  420. 0, /* sq_inplace_concat */
  421. 0 /* sq_inplace_repeat */
  422. };
  423. static PyBufferProcs BufferSegment_as_buffer = {
  424. #if PY_MAJOR_VERSION >= 3
  425. (getbufferproc)BufferSegment_getbuffer,
  426. 0
  427. #else
  428. (readbufferproc)BufferSegment_getreadbuffer,
  429. 0,
  430. (segcountproc)BufferSegment_getsegcount,
  431. 0
  432. #endif
  433. };
  434. static PyMethodDef BufferSegment_methods[] = {
  435. { "tobytes", (PyCFunction)BufferSegment_tobytes,
  436. METH_NOARGS, BufferSegment_tobytes__doc__ },
  437. { NULL, NULL }
  438. };
  439. static PyMemberDef BufferSegment_members[] = {
  440. { "offset", T_ULONGLONG, offsetof(ZstdBufferSegment, offset), READONLY,
  441. "offset of segment within parent buffer" },
  442. { NULL }
  443. };
  444. PyTypeObject ZstdBufferSegmentType = {
  445. PyVarObject_HEAD_INIT(NULL, 0)
  446. "zstd.BufferSegment", /* tp_name */
  447. sizeof(ZstdBufferSegment),/* tp_basicsize */
  448. 0, /* tp_itemsize */
  449. (destructor)BufferSegment_dealloc, /* tp_dealloc */
  450. 0, /* tp_print */
  451. 0, /* tp_getattr */
  452. 0, /* tp_setattr */
  453. 0, /* tp_compare */
  454. 0, /* tp_repr */
  455. 0, /* tp_as_number */
  456. &BufferSegment_sq, /* tp_as_sequence */
  457. 0, /* tp_as_mapping */
  458. 0, /* tp_hash */
  459. 0, /* tp_call */
  460. 0, /* tp_str */
  461. 0, /* tp_getattro */
  462. 0, /* tp_setattro */
  463. &BufferSegment_as_buffer, /* tp_as_buffer */
  464. Py_TPFLAGS_DEFAULT, /* tp_flags */
  465. BufferSegment__doc__, /* tp_doc */
  466. 0, /* tp_traverse */
  467. 0, /* tp_clear */
  468. 0, /* tp_richcompare */
  469. 0, /* tp_weaklistoffset */
  470. 0, /* tp_iter */
  471. 0, /* tp_iternext */
  472. BufferSegment_methods, /* tp_methods */
  473. BufferSegment_members, /* tp_members */
  474. 0, /* tp_getset */
  475. 0, /* tp_base */
  476. 0, /* tp_dict */
  477. 0, /* tp_descr_get */
  478. 0, /* tp_descr_set */
  479. 0, /* tp_dictoffset */
  480. 0, /* tp_init */
  481. 0, /* tp_alloc */
  482. PyType_GenericNew, /* tp_new */
  483. };
  484. PyDoc_STRVAR(BufferWithSegmentsCollection__doc__,
  485. "Represents a collection of BufferWithSegments.\n"
  486. );
  487. static void BufferWithSegmentsCollection_dealloc(ZstdBufferWithSegmentsCollection* self) {
  488. Py_ssize_t i;
  489. if (self->firstElements) {
  490. PyMem_Free(self->firstElements);
  491. self->firstElements = NULL;
  492. }
  493. if (self->buffers) {
  494. for (i = 0; i < self->bufferCount; i++) {
  495. Py_CLEAR(self->buffers[i]);
  496. }
  497. PyMem_Free(self->buffers);
  498. self->buffers = NULL;
  499. }
  500. PyObject_Del(self);
  501. }
  502. static int BufferWithSegmentsCollection_init(ZstdBufferWithSegmentsCollection* self, PyObject* args) {
  503. Py_ssize_t size;
  504. Py_ssize_t i;
  505. Py_ssize_t offset = 0;
  506. size = PyTuple_Size(args);
  507. if (-1 == size) {
  508. return -1;
  509. }
  510. if (0 == size) {
  511. PyErr_SetString(PyExc_ValueError, "must pass at least 1 argument");
  512. return -1;
  513. }
  514. for (i = 0; i < size; i++) {
  515. PyObject* item = PyTuple_GET_ITEM(args, i);
  516. if (!PyObject_TypeCheck(item, &ZstdBufferWithSegmentsType)) {
  517. PyErr_SetString(PyExc_TypeError, "arguments must be BufferWithSegments instances");
  518. return -1;
  519. }
  520. if (0 == ((ZstdBufferWithSegments*)item)->segmentCount ||
  521. 0 == ((ZstdBufferWithSegments*)item)->dataSize) {
  522. PyErr_SetString(PyExc_ValueError, "ZstdBufferWithSegments cannot be empty");
  523. return -1;
  524. }
  525. }
  526. self->buffers = PyMem_Malloc(size * sizeof(ZstdBufferWithSegments*));
  527. if (NULL == self->buffers) {
  528. PyErr_NoMemory();
  529. return -1;
  530. }
  531. self->firstElements = PyMem_Malloc(size * sizeof(Py_ssize_t));
  532. if (NULL == self->firstElements) {
  533. PyMem_Free(self->buffers);
  534. self->buffers = NULL;
  535. PyErr_NoMemory();
  536. return -1;
  537. }
  538. self->bufferCount = size;
  539. for (i = 0; i < size; i++) {
  540. ZstdBufferWithSegments* item = (ZstdBufferWithSegments*)PyTuple_GET_ITEM(args, i);
  541. self->buffers[i] = item;
  542. Py_INCREF(item);
  543. if (i > 0) {
  544. self->firstElements[i - 1] = offset;
  545. }
  546. offset += item->segmentCount;
  547. }
  548. self->firstElements[size - 1] = offset;
  549. return 0;
  550. }
  551. static PyObject* BufferWithSegmentsCollection_size(ZstdBufferWithSegmentsCollection* self) {
  552. Py_ssize_t i;
  553. Py_ssize_t j;
  554. unsigned long long size = 0;
  555. for (i = 0; i < self->bufferCount; i++) {
  556. for (j = 0; j < self->buffers[i]->segmentCount; j++) {
  557. size += self->buffers[i]->segments[j].length;
  558. }
  559. }
  560. return PyLong_FromUnsignedLongLong(size);
  561. }
  562. Py_ssize_t BufferWithSegmentsCollection_length(ZstdBufferWithSegmentsCollection* self) {
  563. return self->firstElements[self->bufferCount - 1];
  564. }
  565. static ZstdBufferSegment* BufferWithSegmentsCollection_item(ZstdBufferWithSegmentsCollection* self, Py_ssize_t i) {
  566. Py_ssize_t bufferOffset;
  567. if (i < 0) {
  568. PyErr_SetString(PyExc_IndexError, "offset must be non-negative");
  569. return NULL;
  570. }
  571. if (i >= BufferWithSegmentsCollection_length(self)) {
  572. PyErr_Format(PyExc_IndexError, "offset must be less than %zd",
  573. BufferWithSegmentsCollection_length(self));
  574. return NULL;
  575. }
  576. for (bufferOffset = 0; bufferOffset < self->bufferCount; bufferOffset++) {
  577. Py_ssize_t offset = 0;
  578. if (i < self->firstElements[bufferOffset]) {
  579. if (bufferOffset > 0) {
  580. offset = self->firstElements[bufferOffset - 1];
  581. }
  582. return BufferWithSegments_item(self->buffers[bufferOffset], i - offset);
  583. }
  584. }
  585. PyErr_SetString(ZstdError, "error resolving segment; this should not happen");
  586. return NULL;
  587. }
  588. static PySequenceMethods BufferWithSegmentsCollection_sq = {
  589. (lenfunc)BufferWithSegmentsCollection_length, /* sq_length */
  590. 0, /* sq_concat */
  591. 0, /* sq_repeat */
  592. (ssizeargfunc)BufferWithSegmentsCollection_item, /* sq_item */
  593. 0, /* sq_ass_item */
  594. 0, /* sq_contains */
  595. 0, /* sq_inplace_concat */
  596. 0 /* sq_inplace_repeat */
  597. };
  598. static PyMethodDef BufferWithSegmentsCollection_methods[] = {
  599. { "size", (PyCFunction)BufferWithSegmentsCollection_size,
  600. METH_NOARGS, PyDoc_STR("total size in bytes of all segments") },
  601. { NULL, NULL }
  602. };
  603. PyTypeObject ZstdBufferWithSegmentsCollectionType = {
  604. PyVarObject_HEAD_INIT(NULL, 0)
  605. "zstd.BufferWithSegmentsCollection", /* tp_name */
  606. sizeof(ZstdBufferWithSegmentsCollection),/* tp_basicsize */
  607. 0, /* tp_itemsize */
  608. (destructor)BufferWithSegmentsCollection_dealloc, /* tp_dealloc */
  609. 0, /* tp_print */
  610. 0, /* tp_getattr */
  611. 0, /* tp_setattr */
  612. 0, /* tp_compare */
  613. 0, /* tp_repr */
  614. 0, /* tp_as_number */
  615. &BufferWithSegmentsCollection_sq, /* tp_as_sequence */
  616. 0, /* tp_as_mapping */
  617. 0, /* tp_hash */
  618. 0, /* tp_call */
  619. 0, /* tp_str */
  620. 0, /* tp_getattro */
  621. 0, /* tp_setattro */
  622. 0, /* tp_as_buffer */
  623. Py_TPFLAGS_DEFAULT, /* tp_flags */
  624. BufferWithSegmentsCollection__doc__, /* tp_doc */
  625. 0, /* tp_traverse */
  626. 0, /* tp_clear */
  627. 0, /* tp_richcompare */
  628. 0, /* tp_weaklistoffset */
  629. /* TODO implement iterator for performance. */
  630. 0, /* tp_iter */
  631. 0, /* tp_iternext */
  632. BufferWithSegmentsCollection_methods, /* tp_methods */
  633. 0, /* tp_members */
  634. 0, /* tp_getset */
  635. 0, /* tp_base */
  636. 0, /* tp_dict */
  637. 0, /* tp_descr_get */
  638. 0, /* tp_descr_set */
  639. 0, /* tp_dictoffset */
  640. (initproc)BufferWithSegmentsCollection_init, /* tp_init */
  641. 0, /* tp_alloc */
  642. PyType_GenericNew, /* tp_new */
  643. };
  644. void bufferutil_module_init(PyObject* mod) {
  645. Py_TYPE(&ZstdBufferWithSegmentsType) = &PyType_Type;
  646. if (PyType_Ready(&ZstdBufferWithSegmentsType) < 0) {
  647. return;
  648. }
  649. Py_INCREF(&ZstdBufferWithSegmentsType);
  650. PyModule_AddObject(mod, "BufferWithSegments", (PyObject*)&ZstdBufferWithSegmentsType);
  651. Py_TYPE(&ZstdBufferSegmentsType) = &PyType_Type;
  652. if (PyType_Ready(&ZstdBufferSegmentsType) < 0) {
  653. return;
  654. }
  655. Py_INCREF(&ZstdBufferSegmentsType);
  656. PyModule_AddObject(mod, "BufferSegments", (PyObject*)&ZstdBufferSegmentsType);
  657. Py_TYPE(&ZstdBufferSegmentType) = &PyType_Type;
  658. if (PyType_Ready(&ZstdBufferSegmentType) < 0) {
  659. return;
  660. }
  661. Py_INCREF(&ZstdBufferSegmentType);
  662. PyModule_AddObject(mod, "BufferSegment", (PyObject*)&ZstdBufferSegmentType);
  663. Py_TYPE(&ZstdBufferWithSegmentsCollectionType) = &PyType_Type;
  664. if (PyType_Ready(&ZstdBufferWithSegmentsCollectionType) < 0) {
  665. return;
  666. }
  667. Py_INCREF(&ZstdBufferWithSegmentsCollectionType);
  668. PyModule_AddObject(mod, "BufferWithSegmentsCollection", (PyObject*)&ZstdBufferWithSegmentsCollectionType);
  669. }