Mercurial > hg
view contrib/python-zstandard/c-ext/compressiondict.c @ 46444:6c778d20c8c2
rust: replace ToString impls with Display
ToString is automatically implementing for everything that implements
Display, and Display can avoid allocating intermediate strings.
Differential Revision: https://phab.mercurial-scm.org/D9904
author | Simon Sapin <simon.sapin@octobus.net> |
---|---|
date | Thu, 28 Jan 2021 19:21:57 +0100 |
parents | e92ca942ddca |
children |
line wrap: on
line source
/** * Copyright (c) 2016-present, Gregory Szorc * All rights reserved. * * This software may be modified and distributed under the terms * of the BSD license. See the LICENSE file for details. */ #include "python-zstandard.h" extern PyObject* ZstdError; ZstdCompressionDict* train_dictionary(PyObject* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "dict_size", "samples", "k", "d", "notifications", "dict_id", "level", "steps", "threads", NULL }; size_t capacity; PyObject* samples; unsigned k = 0; unsigned d = 0; unsigned notifications = 0; unsigned dictID = 0; int level = 0; unsigned steps = 0; int threads = 0; ZDICT_cover_params_t params; Py_ssize_t samplesLen; Py_ssize_t i; size_t samplesSize = 0; void* sampleBuffer = NULL; size_t* sampleSizes = NULL; void* sampleOffset; Py_ssize_t sampleSize; void* dict = NULL; size_t zresult; ZstdCompressionDict* result = NULL; if (!PyArg_ParseTupleAndKeywords(args, kwargs, "nO!|IIIIiIi:train_dictionary", kwlist, &capacity, &PyList_Type, &samples, &k, &d, ¬ifications, &dictID, &level, &steps, &threads)) { return NULL; } if (threads < 0) { threads = cpu_count(); } memset(¶ms, 0, sizeof(params)); params.k = k; params.d = d; params.steps = steps; params.nbThreads = threads; params.zParams.notificationLevel = notifications; params.zParams.dictID = dictID; params.zParams.compressionLevel = level; /* Figure out total size of input samples. */ samplesLen = PyList_Size(samples); for (i = 0; i < samplesLen; i++) { PyObject* sampleItem = PyList_GET_ITEM(samples, i); if (!PyBytes_Check(sampleItem)) { PyErr_SetString(PyExc_ValueError, "samples must be bytes"); return NULL; } samplesSize += PyBytes_GET_SIZE(sampleItem); } sampleBuffer = PyMem_Malloc(samplesSize); if (!sampleBuffer) { PyErr_NoMemory(); goto finally; } sampleSizes = PyMem_Malloc(samplesLen * sizeof(size_t)); if (!sampleSizes) { PyErr_NoMemory(); goto finally; } sampleOffset = sampleBuffer; for (i = 0; i < samplesLen; i++) { PyObject* sampleItem = PyList_GET_ITEM(samples, i); sampleSize = PyBytes_GET_SIZE(sampleItem); sampleSizes[i] = sampleSize; memcpy(sampleOffset, PyBytes_AS_STRING(sampleItem), sampleSize); sampleOffset = (char*)sampleOffset + sampleSize; } dict = PyMem_Malloc(capacity); if (!dict) { PyErr_NoMemory(); goto finally; } Py_BEGIN_ALLOW_THREADS /* No parameters uses the default function, which will use default params and call ZDICT_optimizeTrainFromBuffer_cover under the hood. */ if (!params.k && !params.d && !params.zParams.compressionLevel && !params.zParams.notificationLevel && !params.zParams.dictID) { zresult = ZDICT_trainFromBuffer(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen); } /* Use optimize mode if user controlled steps or threads explicitly. */ else if (params.steps || params.nbThreads) { zresult = ZDICT_optimizeTrainFromBuffer_cover(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen, ¶ms); } /* Non-optimize mode with explicit control. */ else { zresult = ZDICT_trainFromBuffer_cover(dict, capacity, sampleBuffer, sampleSizes, (unsigned)samplesLen, params); } Py_END_ALLOW_THREADS if (ZDICT_isError(zresult)) { PyMem_Free(dict); PyErr_Format(ZstdError, "cannot train dict: %s", ZDICT_getErrorName(zresult)); goto finally; } result = PyObject_New(ZstdCompressionDict, &ZstdCompressionDictType); if (!result) { PyMem_Free(dict); goto finally; } result->dictData = dict; result->dictSize = zresult; result->dictType = ZSTD_dct_fullDict; result->d = params.d; result->k = params.k; result->cdict = NULL; result->ddict = NULL; finally: PyMem_Free(sampleBuffer); PyMem_Free(sampleSizes); return result; } int ensure_ddict(ZstdCompressionDict* dict) { if (dict->ddict) { return 0; } Py_BEGIN_ALLOW_THREADS dict->ddict = ZSTD_createDDict_advanced(dict->dictData, dict->dictSize, ZSTD_dlm_byRef, dict->dictType, ZSTD_defaultCMem); Py_END_ALLOW_THREADS if (!dict->ddict) { PyErr_SetString(ZstdError, "could not create decompression dict"); return 1; } return 0; } PyDoc_STRVAR(ZstdCompressionDict__doc__, "ZstdCompressionDict(data) - Represents a computed compression dictionary\n" "\n" "This type holds the results of a computed Zstandard compression dictionary.\n" "Instances are obtained by calling ``train_dictionary()`` or by passing\n" "bytes obtained from another source into the constructor.\n" ); static int ZstdCompressionDict_init(ZstdCompressionDict* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "data", "dict_type", NULL }; int result = -1; Py_buffer source; unsigned dictType = ZSTD_dct_auto; self->dictData = NULL; self->dictSize = 0; self->cdict = NULL; self->ddict = NULL; #if PY_MAJOR_VERSION >= 3 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|I:ZstdCompressionDict", #else if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*|I:ZstdCompressionDict", #endif kwlist, &source, &dictType)) { return -1; } if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) { PyErr_SetString(PyExc_ValueError, "data buffer should be contiguous and have at most one dimension"); goto finally; } if (dictType != ZSTD_dct_auto && dictType != ZSTD_dct_rawContent && dictType != ZSTD_dct_fullDict) { PyErr_Format(PyExc_ValueError, "invalid dictionary load mode: %d; must use DICT_TYPE_* constants", dictType); goto finally; } self->dictType = dictType; self->dictData = PyMem_Malloc(source.len); if (!self->dictData) { PyErr_NoMemory(); goto finally; } memcpy(self->dictData, source.buf, source.len); self->dictSize = source.len; result = 0; finally: PyBuffer_Release(&source); return result; } static void ZstdCompressionDict_dealloc(ZstdCompressionDict* self) { if (self->cdict) { ZSTD_freeCDict(self->cdict); self->cdict = NULL; } if (self->ddict) { ZSTD_freeDDict(self->ddict); self->ddict = NULL; } if (self->dictData) { PyMem_Free(self->dictData); self->dictData = NULL; } PyObject_Del(self); } PyDoc_STRVAR(ZstdCompressionDict_precompute_compress__doc__, "Precompute a dictionary so it can be used by multiple compressors.\n" ); static PyObject* ZstdCompressionDict_precompute_compress(ZstdCompressionDict* self, PyObject* args, PyObject* kwargs) { static char* kwlist[] = { "level", "compression_params", NULL }; int level = 0; ZstdCompressionParametersObject* compressionParams = NULL; ZSTD_compressionParameters cParams; size_t zresult; if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!:precompute_compress", kwlist, &level, &ZstdCompressionParametersType, &compressionParams)) { return NULL; } if (level && compressionParams) { PyErr_SetString(PyExc_ValueError, "must only specify one of level or compression_params"); return NULL; } if (!level && !compressionParams) { PyErr_SetString(PyExc_ValueError, "must specify one of level or compression_params"); return NULL; } if (self->cdict) { zresult = ZSTD_freeCDict(self->cdict); self->cdict = NULL; if (ZSTD_isError(zresult)) { PyErr_Format(ZstdError, "unable to free CDict: %s", ZSTD_getErrorName(zresult)); return NULL; } } if (level) { cParams = ZSTD_getCParams(level, 0, self->dictSize); } else { if (to_cparams(compressionParams, &cParams)) { return NULL; } } assert(!self->cdict); self->cdict = ZSTD_createCDict_advanced(self->dictData, self->dictSize, ZSTD_dlm_byRef, self->dictType, cParams, ZSTD_defaultCMem); if (!self->cdict) { PyErr_SetString(ZstdError, "unable to precompute dictionary"); return NULL; } Py_RETURN_NONE; } static PyObject* ZstdCompressionDict_dict_id(ZstdCompressionDict* self) { unsigned dictID = ZDICT_getDictID(self->dictData, self->dictSize); return PyLong_FromLong(dictID); } static PyObject* ZstdCompressionDict_as_bytes(ZstdCompressionDict* self) { return PyBytes_FromStringAndSize(self->dictData, self->dictSize); } static PyMethodDef ZstdCompressionDict_methods[] = { { "dict_id", (PyCFunction)ZstdCompressionDict_dict_id, METH_NOARGS, PyDoc_STR("dict_id() -- obtain the numeric dictionary ID") }, { "as_bytes", (PyCFunction)ZstdCompressionDict_as_bytes, METH_NOARGS, PyDoc_STR("as_bytes() -- obtain the raw bytes constituting the dictionary data") }, { "precompute_compress", (PyCFunction)ZstdCompressionDict_precompute_compress, METH_VARARGS | METH_KEYWORDS, ZstdCompressionDict_precompute_compress__doc__ }, { NULL, NULL } }; static PyMemberDef ZstdCompressionDict_members[] = { { "k", T_UINT, offsetof(ZstdCompressionDict, k), READONLY, "segment size" }, { "d", T_UINT, offsetof(ZstdCompressionDict, d), READONLY, "dmer size" }, { NULL } }; static Py_ssize_t ZstdCompressionDict_length(ZstdCompressionDict* self) { return self->dictSize; } static PySequenceMethods ZstdCompressionDict_sq = { (lenfunc)ZstdCompressionDict_length, /* sq_length */ 0, /* sq_concat */ 0, /* sq_repeat */ 0, /* sq_item */ 0, /* sq_ass_item */ 0, /* sq_contains */ 0, /* sq_inplace_concat */ 0 /* sq_inplace_repeat */ }; PyTypeObject ZstdCompressionDictType = { PyVarObject_HEAD_INIT(NULL, 0) "zstd.ZstdCompressionDict", /* tp_name */ sizeof(ZstdCompressionDict), /* tp_basicsize */ 0, /* tp_itemsize */ (destructor)ZstdCompressionDict_dealloc, /* tp_dealloc */ 0, /* tp_print */ 0, /* tp_getattr */ 0, /* tp_setattr */ 0, /* tp_compare */ 0, /* tp_repr */ 0, /* tp_as_number */ &ZstdCompressionDict_sq, /* tp_as_sequence */ 0, /* tp_as_mapping */ 0, /* tp_hash */ 0, /* tp_call */ 0, /* tp_str */ 0, /* tp_getattro */ 0, /* tp_setattro */ 0, /* tp_as_buffer */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */ ZstdCompressionDict__doc__, /* tp_doc */ 0, /* tp_traverse */ 0, /* tp_clear */ 0, /* tp_richcompare */ 0, /* tp_weaklistoffset */ 0, /* tp_iter */ 0, /* tp_iternext */ ZstdCompressionDict_methods, /* tp_methods */ ZstdCompressionDict_members, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ 0, /* tp_dict */ 0, /* tp_descr_get */ 0, /* tp_descr_set */ 0, /* tp_dictoffset */ (initproc)ZstdCompressionDict_init, /* tp_init */ 0, /* tp_alloc */ PyType_GenericNew, /* tp_new */ }; void compressiondict_module_init(PyObject* mod) { Py_SET_TYPE(&ZstdCompressionDictType, &PyType_Type); if (PyType_Ready(&ZstdCompressionDictType) < 0) { return; } Py_INCREF((PyObject*)&ZstdCompressionDictType); PyModule_AddObject(mod, "ZstdCompressionDict", (PyObject*)&ZstdCompressionDictType); }