Compression.h

Go to the documentation of this file.

/*
 * Copyright 2013-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <vector>

#include <folly/Optional.h>
#include <folly/Range.h>
#include <folly/compression/Counters.h>
#include <folly/io/IOBuf.h>

namespace folly {
namespace io {

enum class CodecType {
  USER_DEFINED = 0,

  NO_COMPRESSION = 1,

  LZ4 = 2,

  SNAPPY = 3,

  ZLIB = 4,

  LZ4_VARINT_SIZE = 5,

  LZMA2 = 6,
  LZMA2_VARINT_SIZE = 7,

  ZSTD = 8,

  GZIP = 9,

  LZ4_FRAME = 10,

  BZIP2 = 11,

  ZSTD_FAST = 12,

  NUM_CODEC_TYPES = 13,
};

class Codec {
 public:
  virtual ~Codec() {}

  static constexpr uint64_t UNLIMITED_UNCOMPRESSED_LENGTH = uint64_t(-1);
  uint64_t maxUncompressedLength() const;

  CodecType type() const {
    return type_;
  }

  bool needsUncompressedLength() const;

  std::unique_ptr<IOBuf> compress(const folly::IOBuf* data);

  std::string compress(StringPiece data);

  std::unique_ptr<IOBuf> uncompress(
      const IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength = folly::none);

  std::string uncompress(
      StringPiece data,
      folly::Optional<uint64_t> uncompressedLength = folly::none);

  uint64_t maxCompressedLength(uint64_t uncompressedLength) const;

  folly::Optional<uint64_t> getUncompressedLength(
      const folly::IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength = folly::none) const;

 protected:
  Codec(
      CodecType type,
      folly::Optional<int> level = folly::none,
      folly::StringPiece name = {},
      bool counters = true);

 public:
  virtual std::vector<std::string> validPrefixes() const;

  virtual bool canUncompress(
      const folly::IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength = folly::none) const;

 private:
  // default: no limits (save for special value UNKNOWN_UNCOMPRESSED_LENGTH)
  virtual uint64_t doMaxUncompressedLength() const;
  // default: doesn't need uncompressed length
  virtual bool doNeedsUncompressedLength() const;
  virtual std::unique_ptr<IOBuf> doCompress(const folly::IOBuf* data) = 0;
  virtual std::unique_ptr<IOBuf> doUncompress(
      const folly::IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength) = 0;
  // default: an implementation is provided by default to wrap the strings into
  // IOBufs and delegate to the IOBuf methods. This incurs a copy of the output
  // from IOBuf to string. Implementers, at their discretion, can override
  // these methods to avoid the copy.
  virtual std::string doCompressString(StringPiece data);
  virtual std::string doUncompressString(
      StringPiece data,
      folly::Optional<uint64_t> uncompressedLength);

  virtual uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const = 0;
  // default: returns the passed uncompressedLength.
  virtual folly::Optional<uint64_t> doGetUncompressedLength(
      const folly::IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength) const;

  CodecType type_;
  folly::detail::CompressionCounter bytesBeforeCompression_;
  folly::detail::CompressionCounter bytesAfterCompression_;
  folly::detail::CompressionCounter bytesBeforeDecompression_;
  folly::detail::CompressionCounter bytesAfterDecompression_;
  folly::detail::CompressionCounter compressions_;
  folly::detail::CompressionCounter decompressions_;
  folly::detail::CompressionCounter compressionMilliseconds_;
  folly::detail::CompressionCounter decompressionMilliseconds_;
};

class StreamCodec : public Codec {
 public:
  ~StreamCodec() override {}

  bool needsDataLength() const;

  /*****************************************************************************
   * Streaming API
   *****************************************************************************
   * A low-level stateful streaming API.
   * Streaming operations can be started in two ways:
   *   1. From a clean Codec on which no non-const methods have been called.
   *   2. A call to resetStream(), which will reset any codec to a clean state.
   * After a streaming operation has begun, either compressStream() or
   * uncompressStream() must be called until the streaming operation ends.
   * compressStream() ends when it returns true with flushOp END.
   * uncompressStream() ends when it returns true. At this point the codec
   * may be reused by calling resetStream().
   *
   * compress() and uncompress() can be called at any time, but they interrupt
   * any ongoing streaming operations (state is lost and resetStream() must be
   * called before another streaming operation).
   */

  void resetStream(folly::Optional<uint64_t> uncompressedLength = folly::none);

  enum class FlushOp { NONE, FLUSH, END };

  bool compressStream(
      folly::ByteRange& input,
      folly::MutableByteRange& output,
      FlushOp flushOp = StreamCodec::FlushOp::NONE);

  bool uncompressStream(
      folly::ByteRange& input,
      folly::MutableByteRange& output,
      FlushOp flushOp = StreamCodec::FlushOp::NONE);

 protected:
  StreamCodec(
      CodecType type,
      folly::Optional<int> level = folly::none,
      folly::StringPiece name = {},
      bool counters = true)
      : Codec(type, std::move(level), name, counters) {}

  // Returns the uncompressed length last passed to resetStream() or none if it
  // hasn't been called yet.
  folly::Optional<uint64_t> uncompressedLength() const {
    return uncompressedLength_;
  }

 private:
  // default: Implemented using the streaming API.
  std::unique_ptr<IOBuf> doCompress(const folly::IOBuf* data) override;
  std::unique_ptr<IOBuf> doUncompress(
      const folly::IOBuf* data,
      folly::Optional<uint64_t> uncompressedLength) override;

  // default: Returns false
  virtual bool doNeedsDataLength() const;
  virtual void doResetStream() = 0;
  virtual bool doCompressStream(
      folly::ByteRange& input,
      folly::MutableByteRange& output,
      FlushOp flushOp) = 0;
  virtual bool doUncompressStream(
      folly::ByteRange& input,
      folly::MutableByteRange& output,
      FlushOp flushOp) = 0;

  enum class State {
    RESET,
    COMPRESS,
    COMPRESS_FLUSH,
    COMPRESS_END,
    UNCOMPRESS,
    END,
  };
  void assertStateIs(State expected) const;

  State state_{State::RESET};
  ByteRange previousInput_{};
  folly::Optional<uint64_t> uncompressedLength_{};
  bool progressMade_{true};
};

constexpr int COMPRESSION_LEVEL_FASTEST = -1;
constexpr int COMPRESSION_LEVEL_DEFAULT = -2;
constexpr int COMPRESSION_LEVEL_BEST = -3;

std::unique_ptr<Codec> getCodec(
    CodecType type,
    int level = COMPRESSION_LEVEL_DEFAULT);

std::unique_ptr<StreamCodec> getStreamCodec(
    CodecType type,
    int level = COMPRESSION_LEVEL_DEFAULT);

std::unique_ptr<Codec> getAutoUncompressionCodec(
    std::vector<std::unique_ptr<Codec>> customCodecs = {},
    std::unique_ptr<Codec> terminalCodec = {});

bool hasCodec(CodecType type);

bool hasStreamCodec(CodecType type);
} // namespace io
} // namespace folly