Expose size estimation helpers for Linux (#4064)

diff --git a/contrib/linux-kernel/linux_zstd.h b/contrib/linux-kernel/linux_zstd.h
index f109d49f43f80c141880ff60c949483da04822fb..a7ce080dcb08b185c0c86d5b0292e7937229efcc 100644 (file)
--- a/contrib/linux-kernel/linux_zstd.h
+++ b/contrib/linux-kernel/linux_zstd.h
@@ -153,6 +153,20 @@ typedef ZSTD_CCtx zstd_cctx;
  */
 size_t zstd_cctx_workspace_bound(const zstd_compression_parameters *parameters);
 
+/**
+ * zstd_cctx_workspace_bound_with_ext_seq_prod() - max memory needed to
+ * initialize a zstd_cctx when using the block-level external sequence
+ * producer API.
+ * @parameters: The compression parameters to be used.
+ *
+ * If multiple compression parameters might be used, the caller must call
+ * this function for each set of parameters and use the maximum size.
+ *
+ * Return:      A lower bound on the size of the workspace that is passed to
+ *              zstd_init_cctx().
+ */
+size_t zstd_cctx_workspace_bound_with_ext_seq_prod(const zstd_compression_parameters *parameters);
+
 /**
  * zstd_init_cctx() - initialize a zstd compression context
  * @workspace:      The workspace to emplace the context into. It must outlive
@@ -257,6 +271,16 @@ typedef ZSTD_CStream zstd_cstream;
  */
 size_t zstd_cstream_workspace_bound(const zstd_compression_parameters *cparams);
 
+/**
+ * zstd_cstream_workspace_bound_with_ext_seq_prod() - memory needed to initialize
+ * a zstd_cstream when using the block-level external sequence producer API.
+ * @cparams: The compression parameters to be used for compression.
+ *
+ * Return:   A lower bound on the size of the workspace that is passed to
+ *           zstd_init_cstream().
+ */
+size_t zstd_cstream_workspace_bound_with_ext_seq_prod(const zstd_compression_parameters *cparams);
+
 /**
  * zstd_init_cstream() - initialize a zstd streaming compression context
  * @parameters        The zstd parameters to use for compression.
@@ -416,6 +440,18 @@ size_t zstd_decompress_stream(zstd_dstream *dstream, zstd_out_buffer *output,
  */
 size_t zstd_find_frame_compressed_size(const void *src, size_t src_size);
 
+/**
+ * zstd_register_sequence_producer() - exposes the zstd library function
+ * ZSTD_registerSequenceProducer(). This is used for the block-level external
+ * sequence producer API. See upstream zstd.h for detailed documentation.
+ */
+typedef ZSTD_sequenceProducer_F zstd_sequence_producer_f;
+void zstd_register_sequence_producer(
+  zstd_cctx *cctx,
+  void* sequence_producer_state,
+  zstd_sequence_producer_f sequence_producer
+);
+
 /**
  * struct zstd_frame_params - zstd frame parameters stored in the frame header
  * @frameContentSize: The frame content size, or ZSTD_CONTENTSIZE_UNKNOWN if not

diff --git a/contrib/linux-kernel/zstd_compress_module.c b/contrib/linux-kernel/zstd_compress_module.c
index 8ecf43226af2f32cee0c59d06ab9f06511e16701..039db55f9f302dbdbf3626c1fd7a3534e5001f1f 100644 (file)
--- a/contrib/linux-kernel/zstd_compress_module.c
+++ b/contrib/linux-kernel/zstd_compress_module.c
@@ -16,6 +16,7 @@
 
 #include "common/zstd_deps.h"
 #include "common/zstd_internal.h"
+#include "compress/zstd_compress_internal.h"
 
 #define ZSTD_FORWARD_IF_ERR(ret)            \
        do {                                \
@@ -85,6 +86,52 @@ size_t zstd_cctx_workspace_bound(const zstd_compression_parameters *cparams)
 }
 EXPORT_SYMBOL(zstd_cctx_workspace_bound);
 
+// Used by zstd_cctx_workspace_bound_with_ext_seq_prod()
+static size_t dummy_external_sequence_producer(
+       void *sequenceProducerState,
+       ZSTD_Sequence *outSeqs, size_t outSeqsCapacity,
+       const void *src, size_t srcSize,
+       const void *dict, size_t dictSize,
+       int compressionLevel,
+       size_t windowSize)
+{
+       (void)sequenceProducerState;
+       (void)outSeqs; (void)outSeqsCapacity;
+       (void)src; (void)srcSize;
+       (void)dict; (void)dictSize;
+       (void)compressionLevel;
+       (void)windowSize;
+       return ZSTD_SEQUENCE_PRODUCER_ERROR;
+}
+
+static void init_cctx_params_from_compress_params(
+       ZSTD_CCtx_params *cctx_params,
+       const zstd_compression_parameters *compress_params)
+{
+       ZSTD_parameters zstd_params;
+       memset(&zstd_params, 0, sizeof(zstd_params));
+       zstd_params.cParams = *compress_params;
+       ZSTD_CCtxParams_init_advanced(cctx_params, zstd_params);
+}
+
+size_t zstd_cctx_workspace_bound_with_ext_seq_prod(const zstd_compression_parameters *compress_params)
+{
+       ZSTD_CCtx_params cctx_params;
+       init_cctx_params_from_compress_params(&cctx_params, compress_params);
+       ZSTD_CCtxParams_registerSequenceProducer(&cctx_params, NULL, dummy_external_sequence_producer);
+       return ZSTD_estimateCCtxSize_usingCCtxParams(&cctx_params);
+}
+EXPORT_SYMBOL(zstd_cctx_workspace_bound_with_ext_seq_prod);
+
+size_t zstd_cstream_workspace_bound_with_ext_seq_prod(const zstd_compression_parameters *compress_params)
+{
+       ZSTD_CCtx_params cctx_params;
+       init_cctx_params_from_compress_params(&cctx_params, compress_params);
+       ZSTD_CCtxParams_registerSequenceProducer(&cctx_params, NULL, dummy_external_sequence_producer);
+       return ZSTD_estimateCStreamSize_usingCCtxParams(&cctx_params);
+}
+EXPORT_SYMBOL(zstd_cstream_workspace_bound_with_ext_seq_prod);
+
 zstd_cctx *zstd_init_cctx(void *workspace, size_t workspace_size)
 {
        if (workspace == NULL)
@@ -160,5 +207,14 @@ size_t zstd_end_stream(zstd_cstream *cstream, zstd_out_buffer *output)
 }
 EXPORT_SYMBOL(zstd_end_stream);
 
+void zstd_register_sequence_producer(
+  zstd_cctx *cctx,
+  void* sequence_producer_state,
+  zstd_sequence_producer_f sequence_producer
+) {
+       ZSTD_registerSequenceProducer(cctx, sequence_producer_state, sequence_producer);
+}
+EXPORT_SYMBOL(zstd_register_sequence_producer);
+
 MODULE_LICENSE("Dual BSD/GPL");
 MODULE_DESCRIPTION("Zstd Compressor");