view contrib/python-zstandard/c-ext/compressiondict.c @ 47072:4c041c71ec01

revlog: introduce an explicit tracking of what the revlog is about

Since the dawn of time, people have been forced to rely to lossy introspection
of the index filename to determine what the purpose and role of the revlog they
encounter is. This is hacky, error prone, inflexible, abstraction-leaky,
<insert-your-own-complaints-here>.

In f63299ee7e4d Raphaël introduced a new attribute to track this information:
`revlog_kind`. However it is initialized in an odd place and various instances
end up not having it set. In addition is only tracking some of the information
we end up having to introspect in various pieces of code.

So we add a new attribute that holds more data and is more strictly enforced.
This work is done in collaboration with Raphaël.

The `revlog_kind` one will be removed/adapted in the next changeset. We expect
to be able to clean up various existing piece of code and to simplify coming
work around the newer revlog format.

Differential Revision: https://phab.mercurial-scm.org/D10352

/**
* 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, &notifications, &dictID, &level, &steps, &threads)) {
		return NULL;
	}

	if (threads < 0) {
		threads = cpu_count();
	}

	memset(&params, 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, &params);
	}
	/* 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);
}