contrib/python-zstandard/zstd.c

wireprotov2: advertise set of valid values for requestable fields

changesetdata, manifestdata, and filedata all allow the caller to
specify what data fields to request.

Data fields are extensible and may evolve over time. In order to
prevent clients from making requests for fields that are not
available, the client needs to know what fields are available.

This commit teaches the server to declare a set of "valid values"
for wire protocol command arguments. That set of values is exposed
in the command's capabilities descriptor. The changesetdata,
manifestdata, and filedata commands all declare their set of
available "fields."

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

/**
 * 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.
 */

/* A Python C extension for Zstandard. */

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#elif defined(__APPLE__) || defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)
#include <sys/types.h>
#include <sys/sysctl.h>
#endif

#include "python-zstandard.h"

PyObject *ZstdError;

PyDoc_STRVAR(estimate_decompression_context_size__doc__,
"estimate_decompression_context_size()\n"
"\n"
"Estimate the amount of memory allocated to a decompression context.\n"
);

static PyObject* estimate_decompression_context_size(PyObject* self) {
	return PyLong_FromSize_t(ZSTD_estimateDCtxSize());
}

PyDoc_STRVAR(frame_content_size__doc__,
"frame_content_size(data)\n"
"\n"
"Obtain the decompressed size of a frame."
);

static PyObject* frame_content_size(PyObject* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"source",
		NULL
	};

	Py_buffer source;
	PyObject* result = NULL;
	unsigned long long size;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_content_size",
#else
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_content_size",
#endif
		kwlist, &source)) {
		return NULL;
	}

	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;
	}

	size = ZSTD_getFrameContentSize(source.buf, source.len);

	if (size == ZSTD_CONTENTSIZE_ERROR) {
		PyErr_SetString(ZstdError, "error when determining content size");
	}
	else if (size == ZSTD_CONTENTSIZE_UNKNOWN) {
		result = PyLong_FromLong(-1);
	}
	else {
		result = PyLong_FromUnsignedLongLong(size);
	}

finally:
	PyBuffer_Release(&source);

	return result;
}

PyDoc_STRVAR(frame_header_size__doc__,
"frame_header_size(data)\n"
"\n"
"Obtain the size of a frame header.\n"
);

static PyObject* frame_header_size(PyObject* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"source",
		NULL
	};

	Py_buffer source;
	PyObject* result = NULL;
	size_t zresult;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_header_size",
#else
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_header_size",
#endif
		kwlist, &source)) {
		return NULL;
	}

	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;
	}

	zresult = ZSTD_frameHeaderSize(source.buf, source.len);
	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "could not determine frame header size: %s",
			ZSTD_getErrorName(zresult));
	}
	else {
		result = PyLong_FromSize_t(zresult);
	}

finally:

	PyBuffer_Release(&source);

	return result;
}

PyDoc_STRVAR(get_frame_parameters__doc__,
"get_frame_parameters(data)\n"
"\n"
"Obtains a ``FrameParameters`` instance by parsing data.\n");

PyDoc_STRVAR(train_dictionary__doc__,
"train_dictionary(dict_size, samples, k=None, d=None, steps=None,\n"
"                 threads=None,notifications=0, dict_id=0, level=0)\n"
"\n"
"Train a dictionary from sample data using the COVER algorithm.\n"
"\n"
"A compression dictionary of size ``dict_size`` will be created from the\n"
"iterable of ``samples``. The raw dictionary bytes will be returned.\n"
"\n"
"The COVER algorithm has 2 parameters: ``k`` and ``d``. These control the\n"
"*segment size* and *dmer size*. A reasonable range for ``k`` is\n"
"``[16, 2048+]``. A reasonable range for ``d`` is ``[6, 16]``.\n"
"``d`` must be less than or equal to ``k``.\n"
"\n"
"``steps`` can be specified to control the number of steps through potential\n"
"values of ``k`` and ``d`` to try. ``k`` and ``d`` will only be varied if\n"
"those arguments are not defined. i.e. if ``d`` is ``8``, then only ``k``\n"
"will be varied in this mode.\n"
"\n"
"``threads`` can specify how many threads to use to test various ``k`` and\n"
"``d`` values. ``-1`` will use as many threads as available CPUs. By default,\n"
"a single thread is used.\n"
"\n"
"When ``k`` and ``d`` are not defined, default values are used and the\n"
"algorithm will perform multiple iterations - or steps - to try to find\n"
"ideal parameters. If both ``k`` and ``d`` are specified, then those values\n"
"will be used. ``steps`` or ``threads`` triggers optimization mode to test\n"
"multiple ``k`` and ``d`` variations.\n"
);

static char zstd_doc[] = "Interface to zstandard";

static PyMethodDef zstd_methods[] = {
	{ "estimate_decompression_context_size", (PyCFunction)estimate_decompression_context_size,
	METH_NOARGS, estimate_decompression_context_size__doc__ },
	{ "frame_content_size", (PyCFunction)frame_content_size,
	METH_VARARGS | METH_KEYWORDS, frame_content_size__doc__ },
	{ "frame_header_size", (PyCFunction)frame_header_size,
	METH_VARARGS | METH_KEYWORDS, frame_header_size__doc__ },
	{ "get_frame_parameters", (PyCFunction)get_frame_parameters,
	METH_VARARGS | METH_KEYWORDS, get_frame_parameters__doc__ },
	{ "train_dictionary", (PyCFunction)train_dictionary,
	METH_VARARGS | METH_KEYWORDS, train_dictionary__doc__ },
	{ NULL, NULL }
};

void bufferutil_module_init(PyObject* mod);
void compressobj_module_init(PyObject* mod);
void compressor_module_init(PyObject* mod);
void compressionparams_module_init(PyObject* mod);
void constants_module_init(PyObject* mod);
void compressiondict_module_init(PyObject* mod);
void compressionreader_module_init(PyObject* mod);
void compressionwriter_module_init(PyObject* mod);
void compressoriterator_module_init(PyObject* mod);
void decompressor_module_init(PyObject* mod);
void decompressobj_module_init(PyObject* mod);
void decompressionreader_module_init(PyObject *mod);
void decompressionwriter_module_init(PyObject* mod);
void decompressoriterator_module_init(PyObject* mod);
void frameparams_module_init(PyObject* mod);

void zstd_module_init(PyObject* m) {
	/* python-zstandard relies on unstable zstd C API features. This means
	   that changes in zstd may break expectations in python-zstandard.

	   python-zstandard is distributed with a copy of the zstd sources.
	   python-zstandard is only guaranteed to work with the bundled version
	   of zstd.

	   However, downstream redistributors or packagers may unbundle zstd
	   from python-zstandard. This can result in a mismatch between zstd
	   versions and API semantics. This essentially "voids the warranty"
	   of python-zstandard and may cause undefined behavior.

	   We detect this mismatch here and refuse to load the module if this
	   scenario is detected.
	*/
	if (ZSTD_VERSION_NUMBER != 10304 || ZSTD_versionNumber() != 10304) {
		PyErr_SetString(PyExc_ImportError, "zstd C API mismatch; Python bindings not compiled against expected zstd version");
		return;
	}

	bufferutil_module_init(m);
	compressionparams_module_init(m);
	compressiondict_module_init(m);
	compressobj_module_init(m);
	compressor_module_init(m);
	compressionreader_module_init(m);
	compressionwriter_module_init(m);
	compressoriterator_module_init(m);
	constants_module_init(m);
	decompressor_module_init(m);
	decompressobj_module_init(m);
	decompressionreader_module_init(m);
	decompressionwriter_module_init(m);
	decompressoriterator_module_init(m);
	frameparams_module_init(m);
}

#if defined(__GNUC__) && (__GNUC__ >= 4)
#  define PYTHON_ZSTD_VISIBILITY __attribute__ ((visibility ("default")))
#else
#  define PYTHON_ZSTD_VISIBILITY
#endif

#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef zstd_module = {
	PyModuleDef_HEAD_INIT,
	"zstd",
	zstd_doc,
	-1,
	zstd_methods
};

PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC PyInit_zstd(void) {
	PyObject *m = PyModule_Create(&zstd_module);
	if (m) {
		zstd_module_init(m);
		if (PyErr_Occurred()) {
			Py_DECREF(m);
			m = NULL;
		}
	}
	return m;
}
#else
PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC initzstd(void) {
	PyObject *m = Py_InitModule3("zstd", zstd_methods, zstd_doc);
	if (m) {
		zstd_module_init(m);
	}
}
#endif

/* Attempt to resolve the number of CPUs in the system. */
int cpu_count() {
	int count = 0;

#if defined(_WIN32)
	SYSTEM_INFO si;
	si.dwNumberOfProcessors = 0;
	GetSystemInfo(&si);
	count = si.dwNumberOfProcessors;
#elif defined(__APPLE__)
	int num;
	size_t size = sizeof(int);

	if (0 == sysctlbyname("hw.logicalcpu", &num, &size, NULL, 0)) {
		count = num;
	}
#elif defined(__linux__)
	count = sysconf(_SC_NPROCESSORS_ONLN);
#elif defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)
	int mib[2];
	size_t len = sizeof(count);
	mib[0] = CTL_HW;
	mib[1] = HW_NCPU;
	if (0 != sysctl(mib, 2, &count, &len, NULL, 0)) {
		count = 0;
	}
#elif defined(__hpux)
	count = mpctl(MPC_GETNUMSPUS, NULL, NULL);
#endif

	return count;
}

size_t roundpow2(size_t i) {
	i--;
	i |= i >> 1;
	i |= i >> 2;
	i |= i >> 4;
	i |= i >> 8;
	i |= i >> 16;
	i++;

	return i;
}

/* Safer version of _PyBytes_Resize().
 *
 * _PyBytes_Resize() only works if the refcount is 1. In some scenarios,
 * we can get an object with a refcount > 1, even if it was just created
 * with PyBytes_FromStringAndSize()! That's because (at least) CPython
 * pre-allocates PyBytes instances of size 1 for every possible byte value.
 *
 * If non-0 is returned, obj may or may not be NULL.
 */
int safe_pybytes_resize(PyObject** obj, Py_ssize_t size) {
	PyObject* tmp;

	if ((*obj)->ob_refcnt == 1) {
		return _PyBytes_Resize(obj, size);
	}

	tmp = PyBytes_FromStringAndSize(NULL, size);
	if (!tmp) {
		return -1;
	}

	memcpy(PyBytes_AS_STRING(tmp), PyBytes_AS_STRING(*obj),
		PyBytes_GET_SIZE(*obj));

	Py_DECREF(*obj);
	*obj = tmp;

	return 0;
}