view contrib/python-zstandard/zstd.c @ 48282:0d6a099bba58

dirstate: group return logic and clarify each function in flagfunc

As suggested by spectral during review of my fix of the missing return, this
makes it more obvious that all cases are covered and improves readability of
the functions with new names and a small docstring.

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

/**
 * 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 compressionchunker_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 != 10404 || ZSTD_versionNumber() != 10404) {
		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);
	compressionchunker_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;
}