#ifndef Py_BUILD_CORE_BUILTIN

# define Py_BUILD_CORE_MODULE 1

#include "Python.h"

#include "_zstdmodule.h"

#include "buffer.h"

#include "internal/pycore_lock.h" // PyMutex_IsLocked

#include <stddef.h> // offsetof()

#include <zstd.h> // ZSTD_*()

typedef struct {

PyObject_HEAD

/* Compression context */

ZSTD_CCtx *cctx;

/* ZstdDict object in use */

PyObject *dict;

/* Last mode, initialized to ZSTD_e_end */

int last_mode;

/* (nbWorker >= 1) ? 1 : 0 */

int use_multithread;

/* Compression level */

int compression_level;

/* Lock to protect the compression context */

PyMutex lock;

} ZstdCompressor;

#define ZstdCompressor_CAST(op) ((ZstdCompressor *)op)

static int

zstd_contentsize_converter(PyObject *size, unsigned long long *p)

{

// None means the user indicates the size is unknown.

if (size == Py_None) {

*p = ZSTD_CONTENTSIZE_UNKNOWN;

}

else {

/* ZSTD_CONTENTSIZE_UNKNOWN is 0ULL - 1

unsigned long long pledged_size = PyLong_AsUnsignedLongLong(size);

/* Here we check for (unsigned long long)-1 as a sign of an error in

if (pledged_size == (unsigned long long)-1 && PyErr_Occurred()) {

*p = ZSTD_CONTENTSIZE_ERROR;

if (PyErr_ExceptionMatches(PyExc_OverflowError)) {

PyErr_Format(PyExc_ValueError,

"size argument should be a positive int less "

"than %llu", ZSTD_CONTENTSIZE_ERROR);

return 0;

}

return 0;

}

if (pledged_size >= ZSTD_CONTENTSIZE_ERROR) {

*p = ZSTD_CONTENTSIZE_ERROR;

PyErr_Format(PyExc_ValueError,

"size argument should be a positive int less "

"than %llu", ZSTD_CONTENTSIZE_ERROR);

return 0;

}

*p = pledged_size;

}

return 1;

}

#include "clinic/compressor.c.h"

static int

_zstd_set_c_level(ZstdCompressor *self, int level)

{

/* Set integer compression level */

int min_level = ZSTD_minCLevel();

int max_level = ZSTD_maxCLevel();

if (level < min_level || level > max_level) {

PyErr_Format(PyExc_ValueError,

"illegal compression level %d; the valid range is [%d, %d]",

level, min_level, max_level);

return -1;

}

/* Save for generating ZSTD_CDICT */

self->compression_level = level;

/* Set compressionLevel to compression context */

size_t zstd_ret = ZSTD_CCtx_setParameter(

self->cctx, ZSTD_c_compressionLevel, level);

/* Check error */

if (ZSTD_isError(zstd_ret)) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

set_zstd_error(mod_state, ERR_SET_C_LEVEL, zstd_ret);

return -1;

}

return 0;

}

static int

_zstd_set_c_parameters(ZstdCompressor *self, PyObject *options)

{

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

if (mod_state == NULL) {

return -1;

}

if (!PyDict_Check(options)) {

PyErr_Format(PyExc_TypeError,

"ZstdCompressor() argument 'options' must be dict, not %T",

options);

return -1;

}

Py_ssize_t pos = 0;

PyObject *key, *value;

while (PyDict_Next(options, &pos, &key, &value)) {

/* Check key type */

if (Py_TYPE(key) == mod_state->DParameter_type) {

PyErr_SetString(PyExc_TypeError,

"compression options dictionary key must not be a "

"DecompressionParameter attribute");

return -1;

}

Py_INCREF(key);

Py_INCREF(value);

int key_v = PyLong_AsInt(key);

Py_DECREF(key);

if (key_v == -1 && PyErr_Occurred()) {

Py_DECREF(value);

return -1;

}

int value_v = PyLong_AsInt(value);

Py_DECREF(value);

if (value_v == -1 && PyErr_Occurred()) {

return -1;

}

if (key_v == ZSTD_c_compressionLevel) {

if (_zstd_set_c_level(self, value_v) < 0) {

return -1;

}

continue;

}

if (key_v == ZSTD_c_nbWorkers) {

/* From the zstd library docs:

if (value_v != 0) {

self->use_multithread = 1;

}

}

/* Set parameter to compression context */

size_t zstd_ret = ZSTD_CCtx_setParameter(self->cctx, key_v, value_v);

/* Check error */

if (ZSTD_isError(zstd_ret)) {

set_parameter_error(1, key_v, value_v);

return -1;

}

}

return 0;

}

static void

capsule_free_cdict(PyObject *capsule)

{

ZSTD_CDict *cdict = PyCapsule_GetPointer(capsule, NULL);

ZSTD_freeCDict(cdict);

}

ZSTD_CDict *

_get_CDict(ZstdDict *self, int compressionLevel)

{

assert(PyMutex_IsLocked(&self->lock));

PyObject *level = NULL;

PyObject *capsule = NULL;

ZSTD_CDict *cdict;

int ret;

/* int level object */

level = PyLong_FromLong(compressionLevel);

if (level == NULL) {

goto error;

}

/* Get PyCapsule object from self->c_dicts */

ret = PyDict_GetItemRef(self->c_dicts, level, &capsule);

if (ret < 0) {

goto error;

}

if (capsule == NULL) {

/* Create ZSTD_CDict instance */

Py_BEGIN_ALLOW_THREADS

cdict = ZSTD_createCDict(self->dict_buffer, self->dict_len,

compressionLevel);

Py_END_ALLOW_THREADS

if (cdict == NULL) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

if (mod_state != NULL) {

PyErr_SetString(mod_state->ZstdError,

"Failed to create a ZSTD_CDict instance from "

"Zstandard dictionary content.");

}

goto error;

}

/* Put ZSTD_CDict instance into PyCapsule object */

capsule = PyCapsule_New(cdict, NULL, capsule_free_cdict);

if (capsule == NULL) {

ZSTD_freeCDict(cdict);

goto error;

}

/* Add PyCapsule object to self->c_dicts */

ret = PyDict_SetItem(self->c_dicts, level, capsule);

if (ret < 0) {

goto error;

}

}

else {

/* ZSTD_CDict instance already exists */

cdict = PyCapsule_GetPointer(capsule, NULL);

}

goto success;

error:

cdict = NULL;

success:

Py_XDECREF(level);

Py_XDECREF(capsule);

return cdict;

}

static int

_zstd_load_impl(ZstdCompressor *self, ZstdDict *zd,

_zstd_state *mod_state, int type)

{

size_t zstd_ret;

if (type == DICT_TYPE_DIGESTED) {

/* Get ZSTD_CDict */

ZSTD_CDict *c_dict = _get_CDict(zd, self->compression_level);

if (c_dict == NULL) {

return -1;

}

/* Reference a prepared dictionary.

zstd_ret = ZSTD_CCtx_refCDict(self->cctx, c_dict);

}

else if (type == DICT_TYPE_UNDIGESTED) {

/* Load a dictionary.

zstd_ret = ZSTD_CCtx_loadDictionary(self->cctx, zd->dict_buffer,

zd->dict_len);

}

else if (type == DICT_TYPE_PREFIX) {

/* Load a prefix */

zstd_ret = ZSTD_CCtx_refPrefix(self->cctx, zd->dict_buffer,

zd->dict_len);

}

else {

Py_UNREACHABLE();

}

/* Check error */

if (ZSTD_isError(zstd_ret)) {

set_zstd_error(mod_state, ERR_LOAD_C_DICT, zstd_ret);

return -1;

}

return 0;

}

static int

_zstd_load_c_dict(ZstdCompressor *self, PyObject *dict)

{

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

/* When compressing, use undigested dictionary by default. */

int type = DICT_TYPE_UNDIGESTED;

ZstdDict *zd = _Py_parse_zstd_dict(mod_state, dict, &type);

if (zd == NULL) {

return -1;

}

int ret;

PyMutex_Lock(&zd->lock);

ret = _zstd_load_impl(self, zd, mod_state, type);

PyMutex_Unlock(&zd->lock);

return ret;

}

static PyObject *

_zstd_ZstdCompressor_new_impl(PyTypeObject *type, PyObject *level,

PyObject *options, PyObject *zstd_dict)

{

ZstdCompressor* self = PyObject_GC_New(ZstdCompressor, type);

if (self == NULL) {

goto error;

}

self->use_multithread = 0;

self->compression_level = ZSTD_CLEVEL_DEFAULT;

self->dict = NULL;

self->lock = (PyMutex){0};

/* Compression context */

self->cctx = ZSTD_createCCtx();

if (self->cctx == NULL) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

if (mod_state != NULL) {

PyErr_SetString(mod_state->ZstdError,

"Unable to create ZSTD_CCtx instance.");

}

goto error;

}

/* Last mode */

self->last_mode = ZSTD_e_end;

if (level != Py_None && options != Py_None) {

PyErr_SetString(PyExc_TypeError,

"Only one of level or options should be used.");

goto error;

}

/* Set compression level */

if (level != Py_None) {

if (!PyLong_Check(level)) {

PyErr_SetString(PyExc_TypeError,

"invalid type for level, expected int");

goto error;

}

int level_v = PyLong_AsInt(level);

if (level_v == -1 && PyErr_Occurred()) {

if (PyErr_ExceptionMatches(PyExc_OverflowError)) {

PyErr_Format(PyExc_ValueError,

"illegal compression level; the valid range is [%d, %d]",

ZSTD_minCLevel(), ZSTD_maxCLevel());

}

goto error;

}

if (_zstd_set_c_level(self, level_v) < 0) {

goto error;

}

}

/* Set options dictionary */

if (options != Py_None) {

if (_zstd_set_c_parameters(self, options) < 0) {

goto error;

}

}

/* Load Zstandard dictionary to compression context */

if (zstd_dict != Py_None) {

if (_zstd_load_c_dict(self, zstd_dict) < 0) {

goto error;

}

Py_INCREF(zstd_dict);

self->dict = zstd_dict;

}

// We can only start GC tracking once self->dict is set.

PyObject_GC_Track(self);

return (PyObject*)self;

error:

Py_XDECREF(self);

return NULL;

}

static void

ZstdCompressor_dealloc(PyObject *ob)

{

ZstdCompressor *self = ZstdCompressor_CAST(ob);

PyObject_GC_UnTrack(self);

/* Free compression context */

if (self->cctx) {

ZSTD_freeCCtx(self->cctx);

}

assert(!PyMutex_IsLocked(&self->lock));

/* Py_XDECREF the dict after free the compression context */

Py_CLEAR(self->dict);

PyTypeObject *tp = Py_TYPE(self);

PyObject_GC_Del(ob);

Py_DECREF(tp);

}

static PyObject *

compress_lock_held(ZstdCompressor *self, Py_buffer *data,

ZSTD_EndDirective end_directive)

{

assert(PyMutex_IsLocked(&self->lock));

ZSTD_inBuffer in;

ZSTD_outBuffer out;

_BlocksOutputBuffer buffer = {.writer = NULL};

size_t zstd_ret;

PyObject *ret;

/* Prepare input & output buffers */

if (data != NULL) {

in.src = data->buf;

in.size = data->len;

in.pos = 0;

}

else {

in.src = &in;

in.size = 0;

in.pos = 0;

}

/* Calculate output buffer's size */

size_t output_buffer_size = ZSTD_compressBound(in.size);

if (output_buffer_size > (size_t) PY_SSIZE_T_MAX) {

PyErr_NoMemory();

goto error;

}

if (_OutputBuffer_InitWithSize(&buffer, &out, -1,

(Py_ssize_t) output_buffer_size) < 0) {

goto error;

}

/* Zstandard stream compress */

while (1) {

Py_BEGIN_ALLOW_THREADS

zstd_ret = ZSTD_compressStream2(self->cctx, &out, &in, end_directive);

Py_END_ALLOW_THREADS

/* Check error */

if (ZSTD_isError(zstd_ret)) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

set_zstd_error(mod_state, ERR_COMPRESS, zstd_ret);

goto error;

}

/* Finished */

if (zstd_ret == 0) {

break;

}

/* Output buffer should be exhausted, grow the buffer. */

assert(out.pos == out.size);

if (out.pos == out.size) {

if (_OutputBuffer_Grow(&buffer, &out) < 0) {

goto error;

}

}

}

/* Return a bytes object */

ret = _OutputBuffer_Finish(&buffer, &out);

if (ret != NULL) {

return ret;

}

error:

_OutputBuffer_OnError(&buffer);

return NULL;

}

#ifndef NDEBUG

static inline int

mt_continue_should_break(ZSTD_inBuffer *in, ZSTD_outBuffer *out)

{

return in->size == in->pos && out->size != out->pos;

}

static PyObject *

compress_mt_continue_lock_held(ZstdCompressor *self, Py_buffer *data)

{

assert(PyMutex_IsLocked(&self->lock));

ZSTD_inBuffer in;

ZSTD_outBuffer out;

_BlocksOutputBuffer buffer = {.writer = NULL};

size_t zstd_ret;

PyObject *ret;

/* Prepare input & output buffers */

in.src = data->buf;

in.size = data->len;

in.pos = 0;

if (_OutputBuffer_InitAndGrow(&buffer, &out, -1) < 0) {

goto error;

}

/* Zstandard stream compress */

while (1) {

Py_BEGIN_ALLOW_THREADS

do {

zstd_ret = ZSTD_compressStream2(self->cctx, &out, &in,

ZSTD_e_continue);

} while (out.pos != out.size

&& in.pos != in.size

&& !ZSTD_isError(zstd_ret));

Py_END_ALLOW_THREADS

/* Check error */

if (ZSTD_isError(zstd_ret)) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

set_zstd_error(mod_state, ERR_COMPRESS, zstd_ret);

goto error;

}

/* Like compress_lock_held(), output as much as possible. */

if (out.pos == out.size) {

if (_OutputBuffer_Grow(&buffer, &out) < 0) {

goto error;

}

}

else if (in.pos == in.size) {

/* Finished */

assert(mt_continue_should_break(&in, &out));

break;

}

}

/* Return a bytes object */

ret = _OutputBuffer_Finish(&buffer, &out);

if (ret != NULL) {

return ret;

}

error:

_OutputBuffer_OnError(&buffer);

return NULL;

}

static PyObject *

_zstd_ZstdCompressor_compress_impl(ZstdCompressor *self, Py_buffer *data,

int mode)

{

PyObject *ret;

/* Check mode value */

if (mode != ZSTD_e_continue &&

mode != ZSTD_e_flush &&

mode != ZSTD_e_end)

{

PyErr_SetString(PyExc_ValueError,

"mode argument wrong value, it should be one of "

"ZstdCompressor.CONTINUE, ZstdCompressor.FLUSH_BLOCK, "

"ZstdCompressor.FLUSH_FRAME.");

return NULL;

}

/* Thread-safe code */

PyMutex_Lock(&self->lock);

/* Compress */

if (self->use_multithread && mode == ZSTD_e_continue) {

ret = compress_mt_continue_lock_held(self, data);

}

else {

ret = compress_lock_held(self, data, mode);

}

if (ret) {

self->last_mode = mode;

}

else {

self->last_mode = ZSTD_e_end;

/* Resetting cctx's session never fail */

ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);

}

PyMutex_Unlock(&self->lock);

return ret;

}

static PyObject *

_zstd_ZstdCompressor_flush_impl(ZstdCompressor *self, int mode)

{

PyObject *ret;

/* Check mode value */

if (mode != ZSTD_e_end && mode != ZSTD_e_flush) {

PyErr_SetString(PyExc_ValueError,

"mode argument wrong value, it should be "

"ZstdCompressor.FLUSH_FRAME or "

"ZstdCompressor.FLUSH_BLOCK.");

return NULL;

}

/* Thread-safe code */

PyMutex_Lock(&self->lock);

ret = compress_lock_held(self, NULL, mode);

if (ret) {

self->last_mode = mode;

}

else {

self->last_mode = ZSTD_e_end;

/* Resetting cctx's session never fail */

ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);

}

PyMutex_Unlock(&self->lock);

return ret;

}

static PyObject *

_zstd_ZstdCompressor_set_pledged_input_size_impl(ZstdCompressor *self,

unsigned long long size)

{

// Error occurred while converting argument, should be unreachable

assert(size != ZSTD_CONTENTSIZE_ERROR);

/* Thread-safe code */

PyMutex_Lock(&self->lock);

/* Check the current mode */

if (self->last_mode != ZSTD_e_end) {

PyErr_SetString(PyExc_ValueError,

"set_pledged_input_size() method must be called "

"when last_mode == FLUSH_FRAME");

PyMutex_Unlock(&self->lock);

return NULL;

}

/* Set pledged content size */

size_t zstd_ret = ZSTD_CCtx_setPledgedSrcSize(self->cctx, size);

PyMutex_Unlock(&self->lock);

if (ZSTD_isError(zstd_ret)) {

_zstd_state* mod_state = PyType_GetModuleState(Py_TYPE(self));

set_zstd_error(mod_state, ERR_SET_PLEDGED_INPUT_SIZE, zstd_ret);

return NULL;

}

Py_RETURN_NONE;

}

static PyMethodDef ZstdCompressor_methods[] = {

_ZSTD_ZSTDCOMPRESSOR_COMPRESS_METHODDEF

_ZSTD_ZSTDCOMPRESSOR_FLUSH_METHODDEF

_ZSTD_ZSTDCOMPRESSOR_SET_PLEDGED_INPUT_SIZE_METHODDEF

{NULL, NULL}

};

PyDoc_STRVAR(ZstdCompressor_last_mode_doc,

"flushed, or a frame ended.");

static PyMemberDef ZstdCompressor_members[] = {

{"last_mode", Py_T_INT, offsetof(ZstdCompressor, last_mode),

Py_READONLY, ZstdCompressor_last_mode_doc},

{NULL}

};

static int

ZstdCompressor_traverse(PyObject *ob, visitproc visit, void *arg)

{

ZstdCompressor *self = ZstdCompressor_CAST(ob);

Py_VISIT(self->dict);

return 0;

}

static int

ZstdCompressor_clear(PyObject *ob)

{

ZstdCompressor *self = ZstdCompressor_CAST(ob);

Py_CLEAR(self->dict);

return 0;

}

static PyType_Slot zstdcompressor_slots[] = {

{Py_tp_new, _zstd_ZstdCompressor_new},

{Py_tp_dealloc, ZstdCompressor_dealloc},

{Py_tp_methods, ZstdCompressor_methods},

{Py_tp_members, ZstdCompressor_members},

{Py_tp_doc, (void *)_zstd_ZstdCompressor_new__doc__},

{Py_tp_traverse, ZstdCompressor_traverse},

{Py_tp_clear, ZstdCompressor_clear},

{0, 0}

};

PyType_Spec zstd_compressor_type_spec = {

.name = "compression.zstd.ZstdCompressor",

.basicsize = sizeof(ZstdCompressor),

// Py_TPFLAGS_IMMUTABLETYPE is not used here as several

// associated constants need to be added to the type.

// PyType_Freeze is called later to set the flag.

.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,

.slots = zstdcompressor_slots,

};