HepMC3/bxzstr/zstd_stream_wrapper.hpp

0001 /* This Source Code Form is subject to the terms of the Mozilla Public
0002  * License, v. 2.0. If a copy of the MPL was not distributed with this
0003  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
0004  *
0005  * This file is a part of bxzstr (https://github.com/tmaklin/bxzstr)
0006  * Written by Tommi Mäklin (tommi@maklin.fi) */
0007
0008 #if defined(BXZSTR_ZSTD_SUPPORT) && (BXZSTR_ZSTD_SUPPORT) == 1
0009
0010 #ifndef BXZSTR_ZSTD_STREAM_WRAPPER_HPP
0011 #define BXZSTR_ZSTD_STREAM_WRAPPER_HPP
0012
0013 #include <zstd.h>
0014
0015 #ifndef ZSTD_CLEVEL_DEFAULT
0016 #define ZSTD_CLEVEL_DEFAULT 5
0017 #endif
0018 #include <string>
0019 #include <exception>
0020
0021 #include "stream_wrapper.hpp"
0022
0023 namespace bxz {
0024 /// Exception class thrown by failed zstd operations.
0025 class zstdException : public std::exception {
0026   public:
0027     zstdException(const size_t err) : msg("zstd error: ") {
0028     this->msg += "[" + std::to_string(err) + "]: ";
0029         this->msg += ZSTD_getErrorName(err);
0030     }
0031     zstdException(const std::string _msg) : msg(_msg) {}
0032
0033     const char * what() const noexcept { return this->msg.c_str(); }
0034
0035   private:
0036     std::string msg;
0037
0038 }; // class zstdException
0039
0040 namespace detail {
0041 class zstd_stream_wrapper : public stream_wrapper {
0042   public:
0043     zstd_stream_wrapper(const bool _isInput = true,
0044             const int level = ZSTD_CLEVEL_DEFAULT, const int = 0)
0045         : isInput(_isInput) {
0046     if (this->isInput) {
0047         this->dctx = ZSTD_createDCtx();
0048         if (this->dctx == NULL) throw zstdException("ZSTD_createDCtx() failed!");
0049     } else {
0050         this->cctx = ZSTD_createCCtx();
0051         if (this->cctx == NULL) throw zstdException("ZSTD_createCCtx() failed!");
0052         this->ret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, level);
0053     }
0054     if (ZSTD_isError(this->ret)) throw zstdException(this->ret);
0055
0056     }
0057
0058     ~zstd_stream_wrapper() {
0059     if (this->isInput) {
0060         ZSTD_freeDCtx(this->dctx);
0061     } else {
0062         ZSTD_freeCCtx(this->cctx);
0063     }
0064     }
0065
0066     int decompress(const int = 0) override {
0067     this->update_inbuffer();
0068     this->update_outbuffer();
0069
0070     this->ret = ZSTD_decompressStream(this->dctx, &output, &input);
0071     if (ZSTD_isError(this->ret)) throw zstdException(this->ret);
0072
0073     this->update_stream_state();
0074
0075     return (int)ret;
0076     }
0077
0078     int compress(const int endStream) override {
0079     this->update_inbuffer();
0080     this->update_outbuffer();
0081
0082     if (endStream) {
0083         this->ret = ZSTD_endStream(this->cctx, &output);
0084         if (ZSTD_isError(this->ret)) throw zstdException(this->ret);
0085     } else {
0086         this->ret = ZSTD_compressStream2(this->cctx, &output, &input, ZSTD_e_continue);
0087         if (ZSTD_isError(this->ret)) throw zstdException(this->ret);
0088
0089         this->ret = (input.pos == input.size);
0090     }
0091
0092     this->update_stream_state();
0093
0094     return (int)ret;
0095     }
0096
0097     bool stream_end() const override { return this->ret == 0; }
0098     bool done() const override { return this->stream_end(); }
0099
0100     const unsigned char* next_in() const override { return static_cast<unsigned char*>(this->buffIn); }
0101     long avail_in() const override { return this->buffInSize; }
0102     unsigned char* next_out() const override { return static_cast<unsigned char*>(this->buffOut); }
0103     long avail_out() const override { return this->buffOutSize; }
0104
0105     void set_next_in(const unsigned char* in) override { this->buffIn = (void*)in; }
0106     void set_avail_in(long in) override { this->buffInSize = (size_t)in; }
0107     void set_next_out(const unsigned char* in) override { this->buffOut = (void*)in; }
0108     void set_avail_out(long in) override { this->buffOutSize = (size_t)in; }
0109
0110   private:
0111     bool isInput;
0112     size_t ret;
0113
0114     size_t buffInSize;
0115     void* buffIn;
0116     size_t buffOutSize;
0117     void* buffOut;
0118
0119     ZSTD_DCtx* dctx;
0120     ZSTD_CCtx* cctx;
0121
0122     ZSTD_inBuffer input;
0123     ZSTD_outBuffer output;
0124
0125     void update_inbuffer() { this->input = { this->buffIn, this->buffInSize, 0 }; }
0126     void update_outbuffer() { this->output =  { this->buffOut, this->buffOutSize, 0 }; }
0127     void update_stream_state() {
0128     this->set_next_out(this->next_out() + this->output.pos);
0129     this->set_avail_out(this->avail_out() - this->output.pos);
0130     this->set_next_in(this->next_in() + this->input.pos);
0131     this->set_avail_in(this->avail_in() - this->input.pos);
0132     }
0133
0134 }; // class zstd_stream_wrapper
0135 } // namespace detail
0136 } // namespace bxz
0137
0138 #endif
0139 #endif