1 files changed, 204 insertions, 0 deletions
diff --git a/Modules/_zstd/_zstdmodule.h b/Modules/_zstd/_zstdmodule.h
new file mode 100644
index 00000000000..d50f1489e6f
--- /dev/null
+++ b/Modules/_zstd/_zstdmodule.h
@@ -0,0 +1,204 @@
+#pragma once
+/*
+Low level interface to Meta's zstd library for use in the compression.zstd
+Python module.
+*/
+
+/* Declarations shared between different parts of the _zstd module*/
+
+#include "Python.h"
+
+#include "zstd.h"
+#include "zdict.h"
+
+
+// if you update the minimum version, you should update the compile
+// check in configure.ac
+#define PYTHON_MINIMUM_SUPPORTED_ZSTD_VERSION 10405
+
+#if ZSTD_VERSION_NUMBER < PYTHON_MINIMUM_SUPPORTED_ZSTD_VERSION
+    #error "_zstd module requires zstd v1.4.5+"
+#endif
+
+/* Forward declaration of module state */
+typedef struct _zstd_state _zstd_state;
+
+/* Forward reference of module def */
+extern PyModuleDef _zstdmodule;
+
+/* For clinic type calculations */
+static inline _zstd_state *
+get_zstd_state_from_type(PyTypeObject *type) {
+    PyObject *module = PyType_GetModuleByDef(type, &_zstdmodule);
+    if (module == NULL) {
+        return NULL;
+    }
+    void *state = PyModule_GetState(module);
+    assert(state != NULL);
+    return (_zstd_state *)state;
+}
+
+extern PyType_Spec zstddict_type_spec;
+extern PyType_Spec zstdcompressor_type_spec;
+extern PyType_Spec ZstdDecompressor_type_spec;
+
+struct _zstd_state {
+    PyObject *empty_bytes;
+    PyObject *empty_readonly_memoryview;
+    PyObject *str_read;
+    PyObject *str_readinto;
+    PyObject *str_write;
+    PyObject *str_flush;
+
+    PyTypeObject *ZstdDict_type;
+    PyTypeObject *ZstdCompressor_type;
+    PyTypeObject *ZstdDecompressor_type;
+    PyObject *ZstdError;
+
+    PyTypeObject *CParameter_type;
+    PyTypeObject *DParameter_type;
+};
+
+typedef struct {
+    PyObject_HEAD
+
+    /* Reusable compress/decompress dictionary, they are created once and
+       can be shared by multiple threads concurrently, since its usage is
+       read-only.
+       c_dicts is a dict, int(compressionLevel):PyCapsule(ZSTD_CDict*) */
+    ZSTD_DDict *d_dict;
+    PyObject *c_dicts;
+
+    /* Content of the dictionary, bytes object. */
+    PyObject *dict_content;
+    /* Dictionary id */
+    uint32_t dict_id;
+
+    /* __init__ has been called, 0 or 1. */
+    int inited;
+} ZstdDict;
+
+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;
+
+    /* __init__ has been called, 0 or 1. */
+    int inited;
+} ZstdCompressor;
+
+typedef struct {
+    PyObject_HEAD
+
+    /* Decompression context */
+    ZSTD_DCtx *dctx;
+
+    /* ZstdDict object in use */
+    PyObject *dict;
+
+    /* Unconsumed input data */
+    char *input_buffer;
+    size_t input_buffer_size;
+    size_t in_begin, in_end;
+
+    /* Unused data */
+    PyObject *unused_data;
+
+    /* 0 if decompressor has (or may has) unconsumed input data, 0 or 1. */
+    char needs_input;
+
+    /* For decompress(), 0 or 1.
+       1 when both input and output streams are at a frame edge, means a
+       frame is completely decoded and fully flushed, or the decompressor
+       just be initialized. */
+    char at_frame_edge;
+
+    /* For ZstdDecompressor, 0 or 1.
+       1 means the end of the first frame has been reached. */
+    char eof;
+
+    /* Used for fast reset above three variables */
+    char _unused_char_for_align;
+
+    /* __init__ has been called, 0 or 1. */
+    int inited;
+} ZstdDecompressor;
+
+typedef enum {
+    TYPE_DECOMPRESSOR,          // <D>, ZstdDecompressor class
+    TYPE_ENDLESS_DECOMPRESSOR,  // <E>, decompress() function
+} decompress_type;
+
+typedef enum {
+    ERR_DECOMPRESS,
+    ERR_COMPRESS,
+    ERR_SET_PLEDGED_INPUT_SIZE,
+
+    ERR_LOAD_D_DICT,
+    ERR_LOAD_C_DICT,
+
+    ERR_GET_C_BOUNDS,
+    ERR_GET_D_BOUNDS,
+    ERR_SET_C_LEVEL,
+
+    ERR_TRAIN_DICT,
+    ERR_FINALIZE_DICT
+} error_type;
+
+typedef enum {
+    DICT_TYPE_DIGESTED = 0,
+    DICT_TYPE_UNDIGESTED = 1,
+    DICT_TYPE_PREFIX = 2
+} dictionary_type;
+
+static inline int
+mt_continue_should_break(ZSTD_inBuffer *in, ZSTD_outBuffer *out) {
+    return in->size == in->pos && out->size != out->pos;
+}
+
+/* Format error message and set ZstdError. */
+extern void
+set_zstd_error(const _zstd_state* const state,
+               const error_type type, size_t zstd_ret);
+
+extern void
+set_parameter_error(const _zstd_state* const state, int is_compress,
+                    int key_v, int value_v);
+
+static const char init_twice_msg[] = "__init__ method is called twice.";
+
+extern int
+_PyZstd_load_c_dict(ZstdCompressor *self, PyObject *dict);
+
+extern int
+_PyZstd_load_d_dict(ZstdDecompressor *self, PyObject *dict);
+
+extern int
+_PyZstd_set_c_parameters(ZstdCompressor *self, PyObject *level_or_options,
+                         const char *arg_name, const char *arg_type);
+
+extern int
+_PyZstd_set_d_parameters(ZstdDecompressor *self, PyObject *options);
+
+extern PyObject *
+decompress_impl(ZstdDecompressor *self, ZSTD_inBuffer *in,
+                Py_ssize_t max_length,
+                Py_ssize_t initial_size,
+                decompress_type type);
+
+extern PyObject *
+compress_impl(ZstdCompressor *self, Py_buffer *data,
+              ZSTD_EndDirective end_directive);