src/ZstdCompressor.cpp

   1 // Copyright (C) 2019-2020 Joel Rosdahl and other contributors
   2 //
   3 // See doc/AUTHORS.adoc for a complete list of contributors.
   4 //
   5 // This program is free software; you can redistribute it and/or modify it
   6 // under the terms of the GNU General Public License as published by the Free
   7 // Software Foundation; either version 3 of the License, or (at your option)
   8 // any later version.
   9 //
  10 // This program is distributed in the hope that it will be useful, but WITHOUT
  11 // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  12 // FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
  13 // more details.
  14 //
  15 // You should have received a copy of the GNU General Public License along with
  16 // this program; if not, write to the Free Software Foundation, Inc., 51
  17 // Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  18
  19 #include "ZstdCompressor.hpp"
  20
  21 #include "Logging.hpp"
  22 #include "assertions.hpp"
  23 #include "exceptions.hpp"
  24
  25 #include <algorithm>
  26
  27 using Logging::log;
  28
  29 ZstdCompressor::ZstdCompressor(FILE* stream, int8_t compression_level)
  30   : m_stream(stream), m_zstd_stream(ZSTD_createCStream())
  31 {
  32   if (compression_level == 0) {
  33     compression_level = default_compression_level;
  34     log("Using default compression level {}", compression_level);
  35   }
  36
  37   // libzstd 1.3.4 and newer support negative levels. However, the query
  38   // function ZSTD_minCLevel did not appear until 1.3.6, so perform detection
  39   // based on version instead.
  40   if (ZSTD_versionNumber() < 10304 && compression_level < 1) {
  41     log(
  42       "Using compression level 1 (minimum level supported by libzstd) instead"
  43       " of {}",
  44       compression_level);
  45     compression_level = 1;
  46   }
  47
  48   m_compression_level = std::min<int>(compression_level, ZSTD_maxCLevel());
  49   if (m_compression_level != compression_level) {
  50     log("Using compression level {} (max libzstd level) instead of {}",
  51         m_compression_level,
  52         compression_level);
  53   }
  54
  55   size_t ret = ZSTD_initCStream(m_zstd_stream, m_compression_level);
  56   if (ZSTD_isError(ret)) {
  57     ZSTD_freeCStream(m_zstd_stream);
  58     throw Error("error initializing zstd compression stream");
  59   }
  60 }
  61
  62 ZstdCompressor::~ZstdCompressor()
  63 {
  64   ZSTD_freeCStream(m_zstd_stream);
  65 }
  66
  67 int8_t
  68 ZstdCompressor::actual_compression_level() const
  69 {
  70   return m_compression_level;
  71 }
  72
  73 void
  74 ZstdCompressor::write(const void* data, size_t count)
  75 {
  76   m_zstd_in.src = data;
  77   m_zstd_in.size = count;
  78   m_zstd_in.pos = 0;
  79
  80   int flush = data ? 0 : 1;
  81
  82   size_t ret;
  83   while (m_zstd_in.pos < m_zstd_in.size) {
  84     uint8_t buffer[READ_BUFFER_SIZE];
  85     m_zstd_out.dst = buffer;
  86     m_zstd_out.size = sizeof(buffer);
  87     m_zstd_out.pos = 0;
  88     ret = ZSTD_compressStream(m_zstd_stream, &m_zstd_out, &m_zstd_in);
  89     ASSERT(!(ZSTD_isError(ret)));
  90     size_t compressed_bytes = m_zstd_out.pos;
  91     if (fwrite(buffer, 1, compressed_bytes, m_stream) != compressed_bytes
  92         || ferror(m_stream)) {
  93       throw Error("failed to write to zstd output stream ");
  94     }
  95   }
  96   ret = flush;
  97   while (ret > 0) {
  98     uint8_t buffer[READ_BUFFER_SIZE];
  99     m_zstd_out.dst = buffer;
 100     m_zstd_out.size = sizeof(buffer);
 101     m_zstd_out.pos = 0;
 102     ret = ZSTD_endStream(m_zstd_stream, &m_zstd_out);
 103     size_t compressed_bytes = m_zstd_out.pos;
 104     if (fwrite(buffer, 1, compressed_bytes, m_stream) != compressed_bytes
 105         || ferror(m_stream)) {
 106       throw Error("failed to write to zstd output stream");
 107     }
 108   }
 109 }
 110
 111 void
 112 ZstdCompressor::finalize()
 113 {
 114   write(nullptr, 0);
 115 }