compression.cc

// Copyright (c) 2022-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).

#include "util/compression.h"

#ifdef BZIP2
#include <bzlib.h>
#endif  // BZIP2

#include <limits>

#ifdef LZ4
#include <lz4.h>
#include <lz4hc.h>
#if LZ4_VERSION_NUMBER < 10700  // < r129
#error "LZ4 support requires version >= 1.7.0 (lz4-devel)"
#endif  // LZ4_VERSION_NUMBER < 10700
#endif  // LZ4

#ifdef SNAPPY
#include <snappy-sinksource.h>
#include <snappy.h>
#endif  // SNAPPY

#ifdef ZLIB
#include <zlib.h>
#endif  // ZLIB

#include "options/options_helper.h"
#include "port/likely.h"
#include "rocksdb/convenience.h"
#include "rocksdb/utilities/object_registry.h"
#include "test_util/sync_point.h"
#include "util/cast_util.h"
#include "util/string_util.h"

namespace ROCKSDB_NAMESPACE {

// WART: does not match OptionsHelper::compression_type_string_map
std::string CompressionTypeToString(CompressionType compression_type) {
  switch (compression_type) {
    case kNoCompression:
      return "NoCompression";
    case kSnappyCompression:
      return "Snappy";
    case kZlibCompression:
      return "Zlib";
    case kBZip2Compression:
      return "BZip2";
    case kLZ4Compression:
      return "LZ4";
    case kLZ4HCCompression:
      return "LZ4HC";
    case kXpressCompression:
      return "Xpress";
    case kZSTD:
      return "ZSTD";
    case kDisableCompressionOption:
      return "DisableOption";
    default: {
      bool is_custom = compression_type >= kFirstCustomCompression &&
                       compression_type <= kLastCustomCompression;
      unsigned char c = lossless_cast<unsigned char>(compression_type);
      return (is_custom ? "Custom" : "Reserved") +
             ToBaseCharsString<16>(2, c, /*uppercase=*/true);
    }
  }
}

// WART: does not match OptionsHelper::compression_type_string_map
CompressionType CompressionTypeFromString(std::string compression_type_str) {
  if (!compression_type_str.empty()) {
    switch (compression_type_str[0]) {
      case 'N':
        if (compression_type_str == "NoCompression") {
          return kNoCompression;
        }
        break;
      case 'S':
        if (compression_type_str == "Snappy") {
          return kSnappyCompression;
        }
        break;
      case 'Z':
        if (compression_type_str == "ZSTD") {
          return kZSTD;
        }
        if (compression_type_str == "Zlib") {
          return kZlibCompression;
        }
        break;
      case 'B':
        if (compression_type_str == "BZip2") {
          return kBZip2Compression;
        }
        break;
      case 'L':
        if (compression_type_str == "LZ4") {
          return kLZ4Compression;
        }
        if (compression_type_str == "LZ4HC") {
          return kLZ4HCCompression;
        }
        break;
      case 'X':
        if (compression_type_str == "Xpress") {
          return kXpressCompression;
        }
        break;
      default:;
    }
  }
  // unrecognized
  return kDisableCompressionOption;
}

std::string CompressionOptionsToString(
    const CompressionOptions& compression_options) {
  std::string result;
  result.reserve(512);
  result.append("window_bits=")
      .append(std::to_string(compression_options.window_bits))
      .append("; ");
  result.append("level=")
      .append(std::to_string(compression_options.level))
      .append("; ");
  result.append("strategy=")
      .append(std::to_string(compression_options.strategy))
      .append("; ");
  result.append("max_dict_bytes=")
      .append(std::to_string(compression_options.max_dict_bytes))
      .append("; ");
  result.append("zstd_max_train_bytes=")
      .append(std::to_string(compression_options.zstd_max_train_bytes))
      .append("; ");
  // NOTE: parallel_threads is skipped because it doesn't really affect the file
  // contents written, arguably doesn't belong in CompressionOptions
  result.append("enabled=")
      .append(std::to_string(compression_options.enabled))
      .append("; ");
  result.append("max_dict_buffer_bytes=")
      .append(std::to_string(compression_options.max_dict_buffer_bytes))
      .append("; ");
  result.append("use_zstd_dict_trainer=")
      .append(std::to_string(compression_options.use_zstd_dict_trainer))
      .append("; ");
  result.append("max_compressed_bytes_per_kb=")
      .append(std::to_string(compression_options.max_compressed_bytes_per_kb))
      .append("; ");
  result.append("checksum=")
      .append(std::to_string(compression_options.checksum))
      .append("; ");
  return result;
}

std::unique_ptr<StreamingCompress> StreamingCompress::Create(
    CompressionType compression_type, const CompressionOptions& opts,
    uint32_t compress_format_version, size_t max_output_len) {
  switch (compression_type) {
    case kZSTD: {
      if (!ZSTD_Streaming_Supported()) {
        return nullptr;
      }
      return std::make_unique<ZSTDStreamingCompress>(
          opts, compress_format_version, max_output_len);
    }
    default:
      return nullptr;
  }
}

std::unique_ptr<StreamingUncompress> StreamingUncompress::Create(
    CompressionType compression_type, uint32_t compress_format_version,
    size_t max_output_len) {
  switch (compression_type) {
    case kZSTD: {
      if (!ZSTD_Streaming_Supported()) {
        return nullptr;
      }
      return std::make_unique<ZSTDStreamingUncompress>(compress_format_version,
                                                       max_output_len);
    }
    default:
      return nullptr;
  }
}

int ZSTDStreamingCompress::Compress(const char* input, size_t input_size,
                                    char* output, size_t* output_pos) {
  assert(input != nullptr && output != nullptr && output_pos != nullptr);
  *output_pos = 0;
  // Don't need to compress an empty input
  if (input_size == 0) {
    return 0;
  }
#ifndef ZSTD
  (void)input;
  (void)input_size;
  (void)output;
  return -1;
#else
  if (input_buffer_.src == nullptr || input_buffer_.src != input) {
    // New input
    // Catch errors where the previous input was not fully decompressed.
    assert(input_buffer_.pos == input_buffer_.size);
    input_buffer_ = {input, input_size, /*pos=*/0};
  } else if (input_buffer_.src == input) {
    // Same input, not fully compressed.
  }
  ZSTD_outBuffer output_buffer = {output, max_output_len_, /*pos=*/0};
  const size_t remaining =
      ZSTD_compressStream2(cctx_, &output_buffer, &input_buffer_, ZSTD_e_end);
  if (ZSTD_isError(remaining)) {
    // Failure
    Reset();
    return -1;
  }
  // Success
  *output_pos = output_buffer.pos;
  return (int)remaining;
#endif
}

void ZSTDStreamingCompress::Reset() {
#ifdef ZSTD
  ZSTD_CCtx_reset(cctx_, ZSTD_ResetDirective::ZSTD_reset_session_only);
  input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
#endif
}

int ZSTDStreamingUncompress::Uncompress(const char* input, size_t input_size,
                                        char* output, size_t* output_pos) {
  assert(output != nullptr && output_pos != nullptr);
  *output_pos = 0;
  // Don't need to uncompress an empty input
  if (input_size == 0) {
    return 0;
  }
#ifdef ZSTD
  if (input) {
    // New input
    input_buffer_ = {input, input_size, /*pos=*/0};
  }
  ZSTD_outBuffer output_buffer = {output, max_output_len_, /*pos=*/0};
  size_t ret = ZSTD_decompressStream(dctx_, &output_buffer, &input_buffer_);
  if (ZSTD_isError(ret)) {
    Reset();
    return -1;
  }
  *output_pos = output_buffer.pos;
  return (int)(input_buffer_.size - input_buffer_.pos);
#else
  (void)input;
  (void)input_size;
  (void)output;
  return -1;
#endif
}

void ZSTDStreamingUncompress::Reset() {
#ifdef ZSTD
  ZSTD_DCtx_reset(dctx_, ZSTD_ResetDirective::ZSTD_reset_session_only);
  input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
#endif
}

void DecompressorDict::Populate(Decompressor& from_decompressor, Slice dict) {
  if (UNLIKELY(dict.empty())) {
    dict_str_ = {};
    dict_allocation_ = {};
    // Appropriately reject bad files with empty dictionary block.
    // It is longstanding not to write an empty dictionary block:
    // https://github.com/facebook/rocksdb/blame/10.2.fb/table/block_based/block_based_table_builder.cc#L1841
    decompressor_ = std::make_unique<FailureDecompressor>(
        Status::Corruption("Decompression dictionary is empty"));
  } else {
    Status s = from_decompressor.MaybeCloneForDict(dict, &decompressor_);
    if (decompressor_ == nullptr) {
      dict_str_ = {};
      dict_allocation_ = {};
      assert(!s.ok());
      decompressor_ = std::make_unique<FailureDecompressor>(std::move(s));
    } else {
      assert(s.ok());
      assert(decompressor_->GetSerializedDict() == dict);
    }
  }

  memory_usage_ = sizeof(struct DecompressorDict);
  memory_usage_ += dict_str_.size();
  if (dict_allocation_) {
    auto allocator = dict_allocation_.get_deleter().allocator;
    if (allocator) {
      memory_usage_ +=
          allocator->UsableSize(dict_allocation_.get(), GetRawDict().size());
    } else {
      memory_usage_ += GetRawDict().size();
    }
  }
  memory_usage_ += decompressor_->ApproximateOwnedMemoryUsage();
}

// ZSTD dictionary training implementations
std::string ZSTD_TrainDictionary(const std::string& samples,
                                 const std::vector<size_t>& sample_lens,
                                 size_t max_dict_bytes) {
#ifdef ZSTD
  assert(samples.empty() == sample_lens.empty());
  if (samples.empty()) {
    return "";
  }
  std::string dict_data(max_dict_bytes, '\0');
  size_t dict_len = ZDICT_trainFromBuffer(
      &dict_data[0], max_dict_bytes, &samples[0], &sample_lens[0],
      static_cast<unsigned>(sample_lens.size()));
  if (ZDICT_isError(dict_len)) {
    return "";
  }
  assert(dict_len <= max_dict_bytes);
  dict_data.resize(dict_len);
  return dict_data;
#else
  assert(false);
  (void)samples;
  (void)sample_lens;
  (void)max_dict_bytes;
  return "";
#endif  // ZSTD
}

std::string ZSTD_TrainDictionary(const std::string& samples,
                                 size_t sample_len_shift,
                                 size_t max_dict_bytes) {
#ifdef ZSTD
  // skips potential partial sample at the end of "samples"
  size_t num_samples = samples.size() >> sample_len_shift;
  std::vector<size_t> sample_lens(num_samples, size_t(1) << sample_len_shift);
  return ZSTD_TrainDictionary(samples, sample_lens, max_dict_bytes);
#else
  assert(false);
  (void)samples;
  (void)sample_len_shift;
  (void)max_dict_bytes;
  return "";
#endif  // ZSTD
}

std::string ZSTD_FinalizeDictionary(const std::string& samples,
                                    const std::vector<size_t>& sample_lens,
                                    size_t max_dict_bytes, int level) {
#ifdef ROCKSDB_ZDICT_FINALIZE
  assert(samples.empty() == sample_lens.empty());
  if (samples.empty()) {
    return "";
  }
  level = SanitizeZSTDCompressionLevel(level);
  std::string dict_data(max_dict_bytes, '\0');
  size_t dict_len = ZDICT_finalizeDictionary(
      dict_data.data(), max_dict_bytes, samples.data(),
      std::min(static_cast<size_t>(samples.size()), max_dict_bytes),
      samples.data(), sample_lens.data(),
      static_cast<unsigned>(sample_lens.size()),
      {level, 0 /* notificationLevel */, 0 /* dictID */});
  if (ZDICT_isError(dict_len)) {
    return "";
  } else {
    assert(dict_len <= max_dict_bytes);
    dict_data.resize(dict_len);
    return dict_data;
  }
#else
  assert(false);
  (void)samples;
  (void)sample_lens;
  (void)max_dict_bytes;
  (void)level;
  return "";
#endif  // ROCKSDB_ZDICT_FINALIZE
}

// ***********************************************************************
// BEGIN built-in implementation of customization interface
// ***********************************************************************
Status Decompressor::ExtractUncompressedSize(Args& args) {
  // Default implementation:
  //
  // Standard format for prepending uncompressed size to the compressed
  // payload. (RocksDB compress_format_version=2 except Snappy)
  //
  // This is historically a varint32, but it is preliminarily generalized
  // to varint64, in case that is supported on the write side for some
  // algorithms.
  if (LIKELY(GetVarint64(&args.compressed_data, &args.uncompressed_size))) {
    if (LIKELY(args.uncompressed_size <= SIZE_MAX)) {
      return Status::OK();
    } else {
      return Status::MemoryLimit("Uncompressed size too large for platform");
    }
  } else {
    return Status::Corruption("Unable to extract uncompressed size");
  }
}

const Slice& Decompressor::GetSerializedDict() const {
  // Default: empty slice => no dictionary
  static Slice kEmptySlice;
  return kEmptySlice;
}

namespace {

class CompressorBase : public Compressor {
 public:
  explicit CompressorBase(const CompressionOptions& opts) : opts_(opts) {}

  uint32_t GetRecommendedParallelThreads() const override {
    return opts_.parallel_threads;
  }

 protected:
  CompressionOptions opts_;
};

class CompressorWithSimpleDictBase : public CompressorBase {
 public:
  explicit CompressorWithSimpleDictBase(const CompressionOptions& opts,
                                        std::string&& dict_data = {})
      : CompressorBase(opts), dict_data_(std::move(dict_data)) {}

  DictConfig GetDictGuidance(CacheEntryRole /*block_type*/) const override {
    if (opts_.max_dict_bytes == 0) {
      return DictDisabled{};
    }
    return DictSampling{opts_.max_dict_bytes};
  }

  // NOTE: empty dict is equivalent to no dict
  Slice GetSerializedDict() const override { return dict_data_; }

  std::unique_ptr<Compressor> Clone() const override {
    return CloneForDict(std::string{dict_data_});
  }

  std::unique_ptr<Compressor> MaybeCloneSpecialized(
      CacheEntryRole /*block_type*/,
      DictConfigArgs&& dict_config) const final override {
    if (auto* samples = std::get_if<DictSamples>(&dict_config)) {
      assert(samples->Verify());
      if (samples->empty()) {
        return nullptr;
      }
      return CloneForDict(std::move(samples->sample_data));
    } else if (auto* predef = std::get_if<DictPreDefined>(&dict_config)) {
      if (predef->dict_data.empty()) {
        return nullptr;
      }
      return CloneForDict(std::move(predef->dict_data));
    } else {
      assert(std::holds_alternative<DictDisabled>(dict_config));
      return nullptr;
    }
  }

  virtual std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const = 0;

 protected:
  const std::string dict_data_;
};

// NOTE: the legacy behavior is to pretend to use dictionary compression when
// enabled, including storing a dictionary block, but to ignore it. That is
// matched here.
class BuiltinSnappyCompressorV2 final : public CompressorWithSimpleDictBase {
 public:
  using CompressorWithSimpleDictBase::CompressorWithSimpleDictBase;

  const char* Name() const override { return "BuiltinSnappyCompressorV2"; }

  CompressionType GetPreferredCompressionType() const override {
    return kSnappyCompression;
  }

  // Snappy is more often SLOWER with parallel compression than faster.
  uint32_t GetRecommendedParallelThreads() const override { return 1; }

  std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const override {
    return std::make_unique<BuiltinSnappyCompressorV2>(opts_,
                                                       std::move(dict_data));
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea*) override {
#ifdef SNAPPY
    struct MySink : public snappy::Sink {
      MySink(char* output, size_t output_size)
          : output_(output), output_size_(output_size) {}

      char* output_;
      size_t output_size_;
      size_t pos_ = 0;

      void Append(const char* data, size_t n) override {
        if (pos_ + n <= output_size_) {
          std::memcpy(output_ + pos_, data, n);
          pos_ += n;
        } else {
          // Virtual abort
          pos_ = output_size_ + 1;
        }
      }

      char* GetAppendBuffer(size_t length, char* scratch) override {
        if (pos_ + length <= output_size_) {
          return output_ + pos_;
        }
        return scratch;
      }
    };
    MySink sink{compressed_output, *compressed_output_size};
    snappy::ByteArraySource source{uncompressed_data.data(),
                                   uncompressed_data.size()};

    size_t outlen = snappy::Compress(&source, &sink);
    if (outlen > 0 && sink.pos_ <= sink.output_size_) {
      // Compression kept/successful
      assert(outlen == sink.pos_);
      *compressed_output_size = outlen;
      *out_compression_type = kSnappyCompression;
      return Status::OK();
    }
    // Compression rejected
    *compressed_output_size = 1;
#else
    (void)uncompressed_data;
    (void)compressed_output;
    // Compression bypassed (not supported)
    *compressed_output_size = 0;
#endif
    *out_compression_type = kNoCompression;
    return Status::OK();
  }

  std::shared_ptr<Decompressor> GetOptimizedDecompressor() const override;
};

[[maybe_unused]]
std::pair<char*, size_t> StartCompressBlockV2(Slice uncompressed_data,
                                              char* compressed_output,
                                              size_t compressed_output_size) {
  if (  // Can't compress more than 4GB
      uncompressed_data.size() > std::numeric_limits<uint32_t>::max() ||
      // Need enough output space for encoding uncompressed size
      compressed_output_size <= 5) {
    // Compression bypassed
    return {nullptr, 0};
  }
  // Standard format for prepending uncompressed size to the compressed
  // data in compress_format_version=2
  char* alg_output = EncodeVarint32(
      compressed_output, static_cast<uint32_t>(uncompressed_data.size()));
  size_t alg_max_output_size =
      compressed_output_size - (alg_output - compressed_output);
  return {alg_output, alg_max_output_size};
}

class BuiltinZlibCompressorV2 final : public CompressorWithSimpleDictBase {
 public:
  using CompressorWithSimpleDictBase::CompressorWithSimpleDictBase;

  const char* Name() const override { return "BuiltinZlibCompressorV2"; }

  CompressionType GetPreferredCompressionType() const override {
    return kZlibCompression;
  }

  std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const override {
    return std::make_unique<BuiltinZlibCompressorV2>(opts_,
                                                     std::move(dict_data));
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea*) override {
#ifdef ZLIB
    auto [alg_output, alg_max_output_size] = StartCompressBlockV2(
        uncompressed_data, compressed_output, *compressed_output_size);
    if (alg_max_output_size == 0) {
      // Compression bypassed
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    // The memLevel parameter specifies how much memory should be allocated for
    // the internal compression state.
    // memLevel=1 uses minimum memory but is slow and reduces compression ratio.
    // memLevel=9 uses maximum memory for optimal speed.
    // The default value is 8. See zconf.h for more details.
    static const int memLevel = 8;
    int level = opts_.level;
    if (level == CompressionOptions::kDefaultCompressionLevel) {
      level = Z_DEFAULT_COMPRESSION;
    }

    z_stream stream;
    memset(&stream, 0, sizeof(z_stream));

    // Initialize the zlib stream
    int st = deflateInit2(&stream, level, Z_DEFLATED, opts_.window_bits,
                          memLevel, opts_.strategy);
    if (st != Z_OK) {
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    // Set dictionary if available
    if (!dict_data_.empty()) {
      st = deflateSetDictionary(
          &stream, reinterpret_cast<const Bytef*>(dict_data_.data()),
          static_cast<unsigned int>(dict_data_.size()));
      if (st != Z_OK) {
        deflateEnd(&stream);
        *compressed_output_size = 0;
        *out_compression_type = kNoCompression;
        return Status::OK();
      }
    }

    // Set up input
    stream.next_in = (Bytef*)uncompressed_data.data();
    stream.avail_in = static_cast<unsigned int>(uncompressed_data.size());

    // Set up output
    stream.next_out = reinterpret_cast<Bytef*>(alg_output);
    stream.avail_out = static_cast<unsigned int>(alg_max_output_size);

    // Compress
    st = deflate(&stream, Z_FINISH);
    size_t outlen = alg_max_output_size - stream.avail_out;
    deflateEnd(&stream);

    if (st == Z_STREAM_END) {
      // Compression kept/successful
      *compressed_output_size =
          outlen + /*header size*/ (alg_output - compressed_output);
      *out_compression_type = kZlibCompression;
      return Status::OK();
    }
    // Compression failed or rejected
    *compressed_output_size = 1;
#else
    (void)uncompressed_data;
    (void)compressed_output;
    // Compression bypassed (not supported)
    *compressed_output_size = 0;
#endif
    *out_compression_type = kNoCompression;
    return Status::OK();
  }
};

class BuiltinBZip2CompressorV2 final : public CompressorWithSimpleDictBase {
 public:
  using CompressorWithSimpleDictBase::CompressorWithSimpleDictBase;

  const char* Name() const override { return "BuiltinBZip2CompressorV2"; }

  CompressionType GetPreferredCompressionType() const override {
    return kBZip2Compression;
  }

  std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const override {
    return std::make_unique<BuiltinBZip2CompressorV2>(opts_,
                                                      std::move(dict_data));
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea*) override {
#ifdef BZIP2
    auto [alg_output, alg_max_output_size] = StartCompressBlockV2(
        uncompressed_data, compressed_output, *compressed_output_size);
    if (alg_max_output_size == 0) {
      // Compression bypassed
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    // BZip2 doesn't actually use the dictionary, but we store it for
    // compatibility similar to BuiltinSnappyCompressorV2

    // Initialize the bzip2 stream
    bz_stream stream;
    memset(&stream, 0, sizeof(bz_stream));

    // Block size 1 is 100K.
    // 0 is for silent.
    // 30 is the default workFactor
    int st = BZ2_bzCompressInit(&stream, 1, 0, 30);
    if (st != BZ_OK) {
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    // Set up input
    stream.next_in = const_cast<char*>(uncompressed_data.data());
    stream.avail_in = static_cast<unsigned int>(uncompressed_data.size());

    // Set up output
    stream.next_out = alg_output;
    stream.avail_out = static_cast<unsigned int>(alg_max_output_size);

    // Compress
    st = BZ2_bzCompress(&stream, BZ_FINISH);
    size_t outlen = alg_max_output_size - stream.avail_out;
    BZ2_bzCompressEnd(&stream);

    // Check for success
    if (st == BZ_STREAM_END) {
      // Compression kept/successful
      *compressed_output_size = outlen + (alg_output - compressed_output);
      *out_compression_type = kBZip2Compression;
      return Status::OK();
    }
    // Compression failed or rejected
    *compressed_output_size = 1;
#else
    (void)uncompressed_data;
    (void)compressed_output;
    // Compression bypassed (not supported)
    *compressed_output_size = 0;
#endif
    *out_compression_type = kNoCompression;
    return Status::OK();
  }
};

class BuiltinLZ4CompressorV2WithDict : public CompressorWithSimpleDictBase {
 public:
  using CompressorWithSimpleDictBase::CompressorWithSimpleDictBase;

  const char* Name() const override { return "BuiltinLZ4CompressorV2"; }

  CompressionType GetPreferredCompressionType() const override {
    return kLZ4Compression;
  }

  // LZ4 (accelerated, not LZ4HC) is more often SLOWER with parallel
  // compression than faster.
  uint32_t GetRecommendedParallelThreads() const override { return 1; }

  std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const override {
    return std::make_unique<BuiltinLZ4CompressorV2WithDict>(
        opts_, std::move(dict_data));
  }

  ManagedWorkingArea ObtainWorkingArea() override {
#ifdef LZ4
    return {reinterpret_cast<WorkingArea*>(LZ4_createStream()), this};
#else
    return {};
#endif
  }
  void ReleaseWorkingArea(WorkingArea* wa) override {
    if (wa) {
#ifdef LZ4
      LZ4_freeStream(reinterpret_cast<LZ4_stream_t*>(wa));
#endif
    }
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea* wa) override {
#ifdef LZ4
    auto [alg_output, alg_max_output_size] = StartCompressBlockV2(
        uncompressed_data, compressed_output, *compressed_output_size);
    if (alg_max_output_size == 0) {
      // Compression bypassed
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    ManagedWorkingArea tmp_wa;
    LZ4_stream_t* stream;
    if (wa != nullptr && wa->owner() == this) {
      stream = reinterpret_cast<LZ4_stream_t*>(wa->get());
#if LZ4_VERSION_NUMBER >= 10900  // >= version 1.9.0
      LZ4_resetStream_fast(stream);
#else
      LZ4_resetStream(stream);
#endif
    } else {
      tmp_wa = ObtainWorkingArea();
      stream = reinterpret_cast<LZ4_stream_t*>(tmp_wa.get());
    }
    if (!dict_data_.empty()) {
      // TODO: more optimization possible here?
      LZ4_loadDict(stream, dict_data_.data(),
                   static_cast<int>(dict_data_.size()));
    }
    int acceleration = LZ4AccelerationFromLevel(opts_.level);
    auto outlen = LZ4_compress_fast_continue(
        stream, uncompressed_data.data(), alg_output,
        static_cast<int>(uncompressed_data.size()),
        static_cast<int>(alg_max_output_size), acceleration);
    if (outlen > 0) {
      // Compression kept/successful
      size_t output_size = static_cast<size_t>(
          outlen + /*header size*/ (alg_output - compressed_output));
      assert(output_size <= *compressed_output_size);
      *compressed_output_size = output_size;
      *out_compression_type = kLZ4Compression;
      return Status::OK();
    }
    // Compression rejected
    *compressed_output_size = 1;
#else
    (void)uncompressed_data;
    (void)compressed_output;
    (void)wa;
    // Compression bypassed (not supported)
    *compressed_output_size = 0;
#endif
    *out_compression_type = kNoCompression;
    return Status::OK();
  }

 protected:
  // Translates a non-positive `level` to an LZ4 "acceleration" value (=
  // -level), clamped to the maximum effective acceleration. The clamp avoids
  // signed overflow when negating INT_MIN and reflects that lz4 internally caps
  // acceleration at LZ4_ACCELERATION_MAX (currently 65537; not exposed by
  // lz4.h).
  [[maybe_unused]] static int LZ4AccelerationFromLevel(int level) {
    assert(level <= 0 || level == CompressionOptions::kDefaultCompressionLevel);
    constexpr int kLZ4MaxAcceleration = 65537;  // == LZ4_ACCELERATION_MAX
    if (level <= -kLZ4MaxAcceleration) {
      return kLZ4MaxAcceleration;
    }
    if (level >= 0) {
      return 1;
    }
    return -level;
  }
};

class BuiltinLZ4CompressorV2NoDict final
    : public BuiltinLZ4CompressorV2WithDict {
 public:
  BuiltinLZ4CompressorV2NoDict(const CompressionOptions& opts)
      : BuiltinLZ4CompressorV2WithDict(opts, /*dict_data=*/{}) {}

  std::unique_ptr<Compressor> Clone() const override {
    return std::make_unique<BuiltinLZ4CompressorV2NoDict>(opts_);
  }

  ManagedWorkingArea ObtainWorkingArea() override {
    // Using an LZ4_stream_t between compressions and resetting with
    // LZ4_resetStream_fast is actually slower than using a fresh LZ4_stream_t
    // each time, or not involving a stream at all. Similarly, using an extState
    // does not seem to offer a performance boost, perhaps a small regression.
    return {};
  }

  void ReleaseWorkingArea(WorkingArea* wa) override {
    // Should not be called
    (void)wa;
    assert(wa == nullptr);
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea* wa) override {
#ifdef LZ4
    (void)wa;
    auto [alg_output, alg_max_output_size] = StartCompressBlockV2(
        uncompressed_data, compressed_output, *compressed_output_size);
    if (alg_max_output_size == 0) {
      // Compression bypassed
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }
    int acceleration = LZ4AccelerationFromLevel(opts_.level);
    auto outlen =
        LZ4_compress_fast(uncompressed_data.data(), alg_output,
                          static_cast<int>(uncompressed_data.size()),
                          static_cast<int>(alg_max_output_size), acceleration);
    if (outlen > 0) {
      // Compression kept/successful
      size_t output_size = static_cast<size_t>(
          outlen + /*header size*/ (alg_output - compressed_output));
      assert(output_size <= *compressed_output_size);
      *compressed_output_size = output_size;
      *out_compression_type = kLZ4Compression;
      return Status::OK();
    }
    // Compression rejected
    *compressed_output_size = 1;
#else
    (void)uncompressed_data;
    (void)compressed_output;
    (void)wa;
    // Compression bypassed (not supported)
    *compressed_output_size = 0;
#endif
    *out_compression_type = kNoCompression;
    return Status::OK();
  }
};

class BuiltinLZ4HCCompressorV2 final : public CompressorWithSimpleDictBase {
 public:
  using CompressorWithSimpleDictBase::CompressorWithSimpleDictBase;

  const char* Name() const override { return "BuiltinLZ4HCCompressorV2"; }

  CompressionType GetPreferredCompressionType() const override {
    return kLZ4HCCompression;
  }

  std::unique_ptr<Compressor> CloneForDict(
      std::string&& dict_data) const override {
    return std::make_unique<BuiltinLZ4HCCompressorV2>(opts_,
                                                      std::move(dict_data));
  }

  ManagedWorkingArea ObtainWorkingArea() override {
#ifdef LZ4
    return {reinterpret_cast<WorkingArea*>(LZ4_createStreamHC()), this};
#else
    return {};
#endif
  }
  void ReleaseWorkingArea(WorkingArea* wa) override {
    if (wa) {
#ifdef LZ4
      LZ4_freeStreamHC(reinterpret_cast<LZ4_streamHC_t*>(wa));
#endif
    }
  }

  Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                       size_t* compressed_output_size,
                       CompressionType* out_compression_type,
                       ManagedWorkingArea* wa) override {
#ifdef LZ4
    auto [alg_output, alg_max_output_size] = StartCompressBlockV2(
        uncompressed_data, compressed_output, *compressed_output_size);
    if (alg_max_output_size == 0) {
      // Compression bypassed
      *compressed_output_size = 0;
      *out_compression_type = kNoCompression;
      return Status::OK();
    }

    int level = opts_.level;
    if (level == CompressionOptions::kDefaultCompressionLevel) {
      level = 0;  // lz4hc.h says any value < 1 will be sanitized to default
    } else if (level > LZ4HC_CLEVEL_MAX) {
      // Map large positive levels to the maximum effective level. lz4hc clamps
      // internally too, but make it explicit for clarity/determinism.
      level = LZ4HC_CLEVEL_MAX;
    }

    ManagedWorkingArea tmp_wa;
    LZ4_streamHC_t* stream;
    if (wa != nullptr && wa->owner() == this) {
      stream = reinterpret_cast<LZ4_streamHC_t*>(wa->get());
    } else {
      tmp_wa = ObtainWorkingArea();
      stream = reinterpret_cast<LZ4_streamHC_t*>(tmp_wa.get());
    }
#if LZ4_VERSION_NUMBER >= 10900  // >= version 1.9.0
    LZ4_resetStreamHC_fast(stream, level);
#else
    LZ4_resetStreamHC(stream, level);
#endif
    if (dict_data_.size() > 0) {
      // TODO: more optimization possible here?
      LZ4_loadDictHC(stream, dict_data_.data(),
                     static_cast<int>(dict_data_.size()));