Allow decompresion without uncompressedSize
[folly.git] / folly / io / Compression.cpp
1 /*
2  * Copyright 2017 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #include <folly/io/Compression.h>
18
19 #if FOLLY_HAVE_LIBLZ4
20 #include <lz4.h>
21 #include <lz4frame.h>
22 #include <lz4hc.h>
23 #endif
24
25 #include <glog/logging.h>
26
27 #if FOLLY_HAVE_LIBSNAPPY
28 #include <snappy.h>
29 #include <snappy-sinksource.h>
30 #endif
31
32 #if FOLLY_HAVE_LIBZ
33 #include <zlib.h>
34 #endif
35
36 #if FOLLY_HAVE_LIBLZMA
37 #include <lzma.h>
38 #endif
39
40 #if FOLLY_HAVE_LIBZSTD
41 #include <zstd.h>
42 #endif
43
44 #include <folly/Conv.h>
45 #include <folly/Memory.h>
46 #include <folly/Portability.h>
47 #include <folly/ScopeGuard.h>
48 #include <folly/Varint.h>
49 #include <folly/io/Cursor.h>
50
51 namespace folly { namespace io {
52
53 Codec::Codec(CodecType type) : type_(type) { }
54
55 // Ensure consistent behavior in the nullptr case
56 std::unique_ptr<IOBuf> Codec::compress(const IOBuf* data) {
57   uint64_t len = data->computeChainDataLength();
58   if (len == 0) {
59     return IOBuf::create(0);
60   }
61   if (len > maxUncompressedLength()) {
62     throw std::runtime_error("Codec: uncompressed length too large");
63   }
64
65   return doCompress(data);
66 }
67
68 std::string Codec::compress(const StringPiece data) {
69   const uint64_t len = data.size();
70   if (len == 0) {
71     return "";
72   }
73   if (len > maxUncompressedLength()) {
74     throw std::runtime_error("Codec: uncompressed length too large");
75   }
76
77   return doCompressString(data);
78 }
79
80 std::unique_ptr<IOBuf> Codec::uncompress(const IOBuf* data,
81                                          uint64_t uncompressedLength) {
82   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
83     if (needsUncompressedLength()) {
84       throw std::invalid_argument("Codec: uncompressed length required");
85     }
86   } else if (uncompressedLength > maxUncompressedLength()) {
87     throw std::runtime_error("Codec: uncompressed length too large");
88   }
89
90   if (data->empty()) {
91     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
92         uncompressedLength != 0) {
93       throw std::runtime_error("Codec: invalid uncompressed length");
94     }
95     return IOBuf::create(0);
96   }
97
98   return doUncompress(data, uncompressedLength);
99 }
100
101 std::string Codec::uncompress(
102     const StringPiece data,
103     uint64_t uncompressedLength) {
104   if (uncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH) {
105     if (needsUncompressedLength()) {
106       throw std::invalid_argument("Codec: uncompressed length required");
107     }
108   } else if (uncompressedLength > maxUncompressedLength()) {
109     throw std::runtime_error("Codec: uncompressed length too large");
110   }
111
112   if (data.empty()) {
113     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
114         uncompressedLength != 0) {
115       throw std::runtime_error("Codec: invalid uncompressed length");
116     }
117     return "";
118   }
119
120   return doUncompressString(data, uncompressedLength);
121 }
122
123 bool Codec::needsUncompressedLength() const {
124   return doNeedsUncompressedLength();
125 }
126
127 uint64_t Codec::maxUncompressedLength() const {
128   return doMaxUncompressedLength();
129 }
130
131 bool Codec::doNeedsUncompressedLength() const {
132   return false;
133 }
134
135 uint64_t Codec::doMaxUncompressedLength() const {
136   return UNLIMITED_UNCOMPRESSED_LENGTH;
137 }
138
139 std::string Codec::doCompressString(const StringPiece data) {
140   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
141   auto outputBuffer = doCompress(&inputBuffer);
142   std::string output;
143   output.reserve(outputBuffer->computeChainDataLength());
144   for (auto range : *outputBuffer) {
145     output.append(reinterpret_cast<const char*>(range.data()), range.size());
146   }
147   return output;
148 }
149
150 std::string Codec::doUncompressString(
151     const StringPiece data,
152     uint64_t uncompressedLength) {
153   const IOBuf inputBuffer{IOBuf::WRAP_BUFFER, data};
154   auto outputBuffer = doUncompress(&inputBuffer, uncompressedLength);
155   std::string output;
156   output.reserve(outputBuffer->computeChainDataLength());
157   for (auto range : *outputBuffer) {
158     output.append(reinterpret_cast<const char*>(range.data()), range.size());
159   }
160   return output;
161 }
162
163 namespace {
164
165 /**
166  * No compression
167  */
168 class NoCompressionCodec final : public Codec {
169  public:
170   static std::unique_ptr<Codec> create(int level, CodecType type);
171   explicit NoCompressionCodec(int level, CodecType type);
172
173  private:
174   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
175   std::unique_ptr<IOBuf> doUncompress(
176       const IOBuf* data,
177       uint64_t uncompressedLength) override;
178 };
179
180 std::unique_ptr<Codec> NoCompressionCodec::create(int level, CodecType type) {
181   return make_unique<NoCompressionCodec>(level, type);
182 }
183
184 NoCompressionCodec::NoCompressionCodec(int level, CodecType type)
185   : Codec(type) {
186   DCHECK(type == CodecType::NO_COMPRESSION);
187   switch (level) {
188   case COMPRESSION_LEVEL_DEFAULT:
189   case COMPRESSION_LEVEL_FASTEST:
190   case COMPRESSION_LEVEL_BEST:
191     level = 0;
192   }
193   if (level != 0) {
194     throw std::invalid_argument(to<std::string>(
195         "NoCompressionCodec: invalid level ", level));
196   }
197 }
198
199 std::unique_ptr<IOBuf> NoCompressionCodec::doCompress(
200     const IOBuf* data) {
201   return data->clone();
202 }
203
204 std::unique_ptr<IOBuf> NoCompressionCodec::doUncompress(
205     const IOBuf* data,
206     uint64_t uncompressedLength) {
207   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
208       data->computeChainDataLength() != uncompressedLength) {
209     throw std::runtime_error(to<std::string>(
210         "NoCompressionCodec: invalid uncompressed length"));
211   }
212   return data->clone();
213 }
214
215 #if (FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA)
216
217 namespace {
218
219 void encodeVarintToIOBuf(uint64_t val, folly::IOBuf* out) {
220   DCHECK_GE(out->tailroom(), kMaxVarintLength64);
221   out->append(encodeVarint(val, out->writableTail()));
222 }
223
224 inline uint64_t decodeVarintFromCursor(folly::io::Cursor& cursor) {
225   uint64_t val = 0;
226   int8_t b = 0;
227   for (int shift = 0; shift <= 63; shift += 7) {
228     b = cursor.read<int8_t>();
229     val |= static_cast<uint64_t>(b & 0x7f) << shift;
230     if (b >= 0) {
231       break;
232     }
233   }
234   if (b < 0) {
235     throw std::invalid_argument("Invalid varint value. Too big.");
236   }
237   return val;
238 }
239
240 }  // namespace
241
242 #endif  // FOLLY_HAVE_LIBLZ4 || FOLLY_HAVE_LIBLZMA
243
244 #if FOLLY_HAVE_LIBLZ4
245
246 /**
247  * LZ4 compression
248  */
249 class LZ4Codec final : public Codec {
250  public:
251   static std::unique_ptr<Codec> create(int level, CodecType type);
252   explicit LZ4Codec(int level, CodecType type);
253
254  private:
255   bool doNeedsUncompressedLength() const override;
256   uint64_t doMaxUncompressedLength() const override;
257
258   bool encodeSize() const { return type() == CodecType::LZ4_VARINT_SIZE; }
259
260   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
261   std::unique_ptr<IOBuf> doUncompress(
262       const IOBuf* data,
263       uint64_t uncompressedLength) override;
264
265   bool highCompression_;
266 };
267
268 std::unique_ptr<Codec> LZ4Codec::create(int level, CodecType type) {
269   return make_unique<LZ4Codec>(level, type);
270 }
271
272 LZ4Codec::LZ4Codec(int level, CodecType type) : Codec(type) {
273   DCHECK(type == CodecType::LZ4 || type == CodecType::LZ4_VARINT_SIZE);
274
275   switch (level) {
276   case COMPRESSION_LEVEL_FASTEST:
277   case COMPRESSION_LEVEL_DEFAULT:
278     level = 1;
279     break;
280   case COMPRESSION_LEVEL_BEST:
281     level = 2;
282     break;
283   }
284   if (level < 1 || level > 2) {
285     throw std::invalid_argument(to<std::string>(
286         "LZ4Codec: invalid level: ", level));
287   }
288   highCompression_ = (level > 1);
289 }
290
291 bool LZ4Codec::doNeedsUncompressedLength() const {
292   return !encodeSize();
293 }
294
295 // The value comes from lz4.h in lz4-r117, but older versions of lz4 don't
296 // define LZ4_MAX_INPUT_SIZE (even though the max size is the same), so do it
297 // here.
298 #ifndef LZ4_MAX_INPUT_SIZE
299 # define LZ4_MAX_INPUT_SIZE 0x7E000000
300 #endif
301
302 uint64_t LZ4Codec::doMaxUncompressedLength() const {
303   return LZ4_MAX_INPUT_SIZE;
304 }
305
306 std::unique_ptr<IOBuf> LZ4Codec::doCompress(const IOBuf* data) {
307   IOBuf clone;
308   if (data->isChained()) {
309     // LZ4 doesn't support streaming, so we have to coalesce
310     clone = data->cloneCoalescedAsValue();
311     data = &clone;
312   }
313
314   uint32_t extraSize = encodeSize() ? kMaxVarintLength64 : 0;
315   auto out = IOBuf::create(extraSize + LZ4_compressBound(data->length()));
316   if (encodeSize()) {
317     encodeVarintToIOBuf(data->length(), out.get());
318   }
319
320   int n;
321   auto input = reinterpret_cast<const char*>(data->data());
322   auto output = reinterpret_cast<char*>(out->writableTail());
323   const auto inputLength = data->length();
324 #if LZ4_VERSION_NUMBER >= 10700
325   if (highCompression_) {
326     n = LZ4_compress_HC(input, output, inputLength, out->tailroom(), 0);
327   } else {
328     n = LZ4_compress_default(input, output, inputLength, out->tailroom());
329   }
330 #else
331   if (highCompression_) {
332     n = LZ4_compressHC(input, output, inputLength);
333   } else {
334     n = LZ4_compress(input, output, inputLength);
335   }
336 #endif
337
338   CHECK_GE(n, 0);
339   CHECK_LE(n, out->capacity());
340
341   out->append(n);
342   return out;
343 }
344
345 std::unique_ptr<IOBuf> LZ4Codec::doUncompress(
346     const IOBuf* data,
347     uint64_t uncompressedLength) {
348   IOBuf clone;
349   if (data->isChained()) {
350     // LZ4 doesn't support streaming, so we have to coalesce
351     clone = data->cloneCoalescedAsValue();
352     data = &clone;
353   }
354
355   folly::io::Cursor cursor(data);
356   uint64_t actualUncompressedLength;
357   if (encodeSize()) {
358     actualUncompressedLength = decodeVarintFromCursor(cursor);
359     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
360         uncompressedLength != actualUncompressedLength) {
361       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
362     }
363   } else {
364     actualUncompressedLength = uncompressedLength;
365     if (actualUncompressedLength == UNKNOWN_UNCOMPRESSED_LENGTH ||
366         actualUncompressedLength > maxUncompressedLength()) {
367       throw std::runtime_error("LZ4Codec: invalid uncompressed length");
368     }
369   }
370
371   auto sp = StringPiece{cursor.peekBytes()};
372   auto out = IOBuf::create(actualUncompressedLength);
373   int n = LZ4_decompress_safe(
374       sp.data(),
375       reinterpret_cast<char*>(out->writableTail()),
376       sp.size(),
377       actualUncompressedLength);
378
379   if (n < 0 || uint64_t(n) != actualUncompressedLength) {
380     throw std::runtime_error(to<std::string>(
381         "LZ4 decompression returned invalid value ", n));
382   }
383   out->append(actualUncompressedLength);
384   return out;
385 }
386
387 class LZ4FrameCodec final : public Codec {
388  public:
389   static std::unique_ptr<Codec> create(int level, CodecType type);
390   explicit LZ4FrameCodec(int level, CodecType type);
391   ~LZ4FrameCodec();
392
393  private:
394   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
395   std::unique_ptr<IOBuf> doUncompress(
396       const IOBuf* data,
397       uint64_t uncompressedLength) override;
398
399   // Reset the dctx_ if it is dirty or null.
400   void resetDCtx();
401
402   int level_;
403   LZ4F_dctx* dctx_{nullptr};
404   bool dirty_{false};
405 };
406
407 /* static */ std::unique_ptr<Codec> LZ4FrameCodec::create(
408     int level,
409     CodecType type) {
410   return make_unique<LZ4FrameCodec>(level, type);
411 }
412
413 static size_t lz4FrameThrowOnError(size_t code) {
414   if (LZ4F_isError(code)) {
415     throw std::runtime_error(
416         to<std::string>("LZ4Frame error: ", LZ4F_getErrorName(code)));
417   }
418   return code;
419 }
420
421 void LZ4FrameCodec::resetDCtx() {
422   if (dctx_ && !dirty_) {
423     return;
424   }
425   if (dctx_) {
426     LZ4F_freeDecompressionContext(dctx_);
427   }
428   lz4FrameThrowOnError(LZ4F_createDecompressionContext(&dctx_, 100));
429   dirty_ = false;
430 }
431
432 LZ4FrameCodec::LZ4FrameCodec(int level, CodecType type) : Codec(type) {
433   DCHECK(type == CodecType::LZ4_FRAME);
434   switch (level) {
435     case COMPRESSION_LEVEL_FASTEST:
436     case COMPRESSION_LEVEL_DEFAULT:
437       level_ = 0;
438       break;
439     case COMPRESSION_LEVEL_BEST:
440       level_ = 16;
441       break;
442     default:
443       level_ = level;
444       break;
445   }
446 }
447
448 LZ4FrameCodec::~LZ4FrameCodec() {
449   if (dctx_) {
450     LZ4F_freeDecompressionContext(dctx_);
451   }
452 }
453
454 std::unique_ptr<IOBuf> LZ4FrameCodec::doCompress(const IOBuf* data) {
455   // LZ4 Frame compression doesn't support streaming so we have to coalesce
456   IOBuf clone;
457   if (data->isChained()) {
458     clone = data->cloneCoalescedAsValue();
459     data = &clone;
460   }
461   // Set preferences
462   const auto uncompressedLength = data->length();
463   LZ4F_preferences_t prefs{};
464   prefs.compressionLevel = level_;
465   prefs.frameInfo.contentSize = uncompressedLength;
466   // Compress
467   auto buf = IOBuf::create(LZ4F_compressFrameBound(uncompressedLength, &prefs));
468   const size_t written = lz4FrameThrowOnError(LZ4F_compressFrame(
469       buf->writableTail(),
470       buf->tailroom(),
471       data->data(),
472       data->length(),
473       &prefs));
474   buf->append(written);
475   return buf;
476 }
477
478 std::unique_ptr<IOBuf> LZ4FrameCodec::doUncompress(
479     const IOBuf* data,
480     uint64_t uncompressedLength) {
481   // Reset the dctx if any errors have occurred
482   resetDCtx();
483   // Coalesce the data
484   ByteRange in = *data->begin();
485   IOBuf clone;
486   if (data->isChained()) {
487     clone = data->cloneCoalescedAsValue();
488     in = clone.coalesce();
489   }
490   data = nullptr;
491   // Select decompression options
492   LZ4F_decompressOptions_t options;
493   options.stableDst = 1;
494   // Select blockSize and growthSize for the IOBufQueue
495   IOBufQueue queue(IOBufQueue::cacheChainLength());
496   auto blockSize = uint64_t{64} << 10;
497   auto growthSize = uint64_t{4} << 20;
498   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH) {
499     // Allocate uncompressedLength in one chunk (up to 64 MB)
500     const auto allocateSize = std::min(uncompressedLength, uint64_t{64} << 20);
501     queue.preallocate(allocateSize, allocateSize);
502     blockSize = std::min(uncompressedLength, blockSize);
503     growthSize = std::min(uncompressedLength, growthSize);
504   } else {
505     // Reduce growthSize for small data
506     const auto guessUncompressedLen = 4 * std::max(blockSize, in.size());
507     growthSize = std::min(guessUncompressedLen, growthSize);
508   }
509   // Once LZ4_decompress() is called, the dctx_ cannot be reused until it
510   // returns 0
511   dirty_ = true;
512   // Decompress until the frame is over
513   size_t code = 0;
514   do {
515     // Allocate enough space to decompress at least a block
516     void* out;
517     size_t outSize;
518     std::tie(out, outSize) = queue.preallocate(blockSize, growthSize);
519     // Decompress
520     size_t inSize = in.size();
521     code = lz4FrameThrowOnError(
522         LZ4F_decompress(dctx_, out, &outSize, in.data(), &inSize, &options));
523     if (in.empty() && outSize == 0 && code != 0) {
524       // We passed no input, no output was produced, and the frame isn't over
525       // No more forward progress is possible
526       throw std::runtime_error("LZ4Frame error: Incomplete frame");
527     }
528     in.uncheckedAdvance(inSize);
529     queue.postallocate(outSize);
530   } while (code != 0);
531   // At this point the decompression context can be reused
532   dirty_ = false;
533   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
534       queue.chainLength() != uncompressedLength) {
535     throw std::runtime_error("LZ4Frame error: Invalid uncompressedLength");
536   }
537   return queue.move();
538 }
539
540 #endif // FOLLY_HAVE_LIBLZ4
541
542 #if FOLLY_HAVE_LIBSNAPPY
543
544 /**
545  * Snappy compression
546  */
547
548 /**
549  * Implementation of snappy::Source that reads from a IOBuf chain.
550  */
551 class IOBufSnappySource final : public snappy::Source {
552  public:
553   explicit IOBufSnappySource(const IOBuf* data);
554   size_t Available() const override;
555   const char* Peek(size_t* len) override;
556   void Skip(size_t n) override;
557  private:
558   size_t available_;
559   io::Cursor cursor_;
560 };
561
562 IOBufSnappySource::IOBufSnappySource(const IOBuf* data)
563   : available_(data->computeChainDataLength()),
564     cursor_(data) {
565 }
566
567 size_t IOBufSnappySource::Available() const {
568   return available_;
569 }
570
571 const char* IOBufSnappySource::Peek(size_t* len) {
572   auto sp = StringPiece{cursor_.peekBytes()};
573   *len = sp.size();
574   return sp.data();
575 }
576
577 void IOBufSnappySource::Skip(size_t n) {
578   CHECK_LE(n, available_);
579   cursor_.skip(n);
580   available_ -= n;
581 }
582
583 class SnappyCodec final : public Codec {
584  public:
585   static std::unique_ptr<Codec> create(int level, CodecType type);
586   explicit SnappyCodec(int level, CodecType type);
587
588  private:
589   uint64_t doMaxUncompressedLength() const override;
590   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
591   std::unique_ptr<IOBuf> doUncompress(
592       const IOBuf* data,
593       uint64_t uncompressedLength) override;
594 };
595
596 std::unique_ptr<Codec> SnappyCodec::create(int level, CodecType type) {
597   return make_unique<SnappyCodec>(level, type);
598 }
599
600 SnappyCodec::SnappyCodec(int level, CodecType type) : Codec(type) {
601   DCHECK(type == CodecType::SNAPPY);
602   switch (level) {
603   case COMPRESSION_LEVEL_FASTEST:
604   case COMPRESSION_LEVEL_DEFAULT:
605   case COMPRESSION_LEVEL_BEST:
606     level = 1;
607   }
608   if (level != 1) {
609     throw std::invalid_argument(to<std::string>(
610         "SnappyCodec: invalid level: ", level));
611   }
612 }
613
614 uint64_t SnappyCodec::doMaxUncompressedLength() const {
615   // snappy.h uses uint32_t for lengths, so there's that.
616   return std::numeric_limits<uint32_t>::max();
617 }
618
619 std::unique_ptr<IOBuf> SnappyCodec::doCompress(const IOBuf* data) {
620   IOBufSnappySource source(data);
621   auto out =
622     IOBuf::create(snappy::MaxCompressedLength(source.Available()));
623
624   snappy::UncheckedByteArraySink sink(reinterpret_cast<char*>(
625       out->writableTail()));
626
627   size_t n = snappy::Compress(&source, &sink);
628
629   CHECK_LE(n, out->capacity());
630   out->append(n);
631   return out;
632 }
633
634 std::unique_ptr<IOBuf> SnappyCodec::doUncompress(const IOBuf* data,
635                                                  uint64_t uncompressedLength) {
636   uint32_t actualUncompressedLength = 0;
637
638   {
639     IOBufSnappySource source(data);
640     if (!snappy::GetUncompressedLength(&source, &actualUncompressedLength)) {
641       throw std::runtime_error("snappy::GetUncompressedLength failed");
642     }
643     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
644         uncompressedLength != actualUncompressedLength) {
645       throw std::runtime_error("snappy: invalid uncompressed length");
646     }
647   }
648
649   auto out = IOBuf::create(actualUncompressedLength);
650
651   {
652     IOBufSnappySource source(data);
653     if (!snappy::RawUncompress(&source,
654                                reinterpret_cast<char*>(out->writableTail()))) {
655       throw std::runtime_error("snappy::RawUncompress failed");
656     }
657   }
658
659   out->append(actualUncompressedLength);
660   return out;
661 }
662
663 #endif  // FOLLY_HAVE_LIBSNAPPY
664
665 #if FOLLY_HAVE_LIBZ
666 /**
667  * Zlib codec
668  */
669 class ZlibCodec final : public Codec {
670  public:
671   static std::unique_ptr<Codec> create(int level, CodecType type);
672   explicit ZlibCodec(int level, CodecType type);
673
674  private:
675   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
676   std::unique_ptr<IOBuf> doUncompress(
677       const IOBuf* data,
678       uint64_t uncompressedLength) override;
679
680   std::unique_ptr<IOBuf> addOutputBuffer(z_stream* stream, uint32_t length);
681   bool doInflate(z_stream* stream, IOBuf* head, uint32_t bufferLength);
682
683   int level_;
684 };
685
686 std::unique_ptr<Codec> ZlibCodec::create(int level, CodecType type) {
687   return make_unique<ZlibCodec>(level, type);
688 }
689
690 ZlibCodec::ZlibCodec(int level, CodecType type) : Codec(type) {
691   DCHECK(type == CodecType::ZLIB || type == CodecType::GZIP);
692   switch (level) {
693   case COMPRESSION_LEVEL_FASTEST:
694     level = 1;
695     break;
696   case COMPRESSION_LEVEL_DEFAULT:
697     level = Z_DEFAULT_COMPRESSION;
698     break;
699   case COMPRESSION_LEVEL_BEST:
700     level = 9;
701     break;
702   }
703   if (level != Z_DEFAULT_COMPRESSION && (level < 0 || level > 9)) {
704     throw std::invalid_argument(to<std::string>(
705         "ZlibCodec: invalid level: ", level));
706   }
707   level_ = level;
708 }
709
710 std::unique_ptr<IOBuf> ZlibCodec::addOutputBuffer(z_stream* stream,
711                                                   uint32_t length) {
712   CHECK_EQ(stream->avail_out, 0);
713
714   auto buf = IOBuf::create(length);
715   buf->append(length);
716
717   stream->next_out = buf->writableData();
718   stream->avail_out = buf->length();
719
720   return buf;
721 }
722
723 bool ZlibCodec::doInflate(z_stream* stream,
724                           IOBuf* head,
725                           uint32_t bufferLength) {
726   if (stream->avail_out == 0) {
727     head->prependChain(addOutputBuffer(stream, bufferLength));
728   }
729
730   int rc = inflate(stream, Z_NO_FLUSH);
731
732   switch (rc) {
733   case Z_OK:
734     break;
735   case Z_STREAM_END:
736     return true;
737   case Z_BUF_ERROR:
738   case Z_NEED_DICT:
739   case Z_DATA_ERROR:
740   case Z_MEM_ERROR:
741     throw std::runtime_error(to<std::string>(
742         "ZlibCodec: inflate error: ", rc, ": ", stream->msg));
743   default:
744     CHECK(false) << rc << ": " << stream->msg;
745   }
746
747   return false;
748 }
749
750 std::unique_ptr<IOBuf> ZlibCodec::doCompress(const IOBuf* data) {
751   z_stream stream;
752   stream.zalloc = nullptr;
753   stream.zfree = nullptr;
754   stream.opaque = nullptr;
755
756   // Using deflateInit2() to support gzip.  "The windowBits parameter is the
757   // base two logarithm of the maximum window size (...) The default value is
758   // 15 (...) Add 16 to windowBits to write a simple gzip header and trailer
759   // around the compressed data instead of a zlib wrapper. The gzip header
760   // will have no file name, no extra data, no comment, no modification time
761   // (set to zero), no header crc, and the operating system will be set to 255
762   // (unknown)."
763   int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
764   // All other parameters (method, memLevel, strategy) get default values from
765   // the zlib manual.
766   int rc = deflateInit2(&stream,
767                         level_,
768                         Z_DEFLATED,
769                         windowBits,
770                         /* memLevel */ 8,
771                         Z_DEFAULT_STRATEGY);
772   if (rc != Z_OK) {
773     throw std::runtime_error(to<std::string>(
774         "ZlibCodec: deflateInit error: ", rc, ": ", stream.msg));
775   }
776
777   stream.next_in = stream.next_out = nullptr;
778   stream.avail_in = stream.avail_out = 0;
779   stream.total_in = stream.total_out = 0;
780
781   bool success = false;
782
783   SCOPE_EXIT {
784     rc = deflateEnd(&stream);
785     // If we're here because of an exception, it's okay if some data
786     // got dropped.
787     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
788       << rc << ": " << stream.msg;
789   };
790
791   uint64_t uncompressedLength = data->computeChainDataLength();
792   uint64_t maxCompressedLength = deflateBound(&stream, uncompressedLength);
793
794   // Max 64MiB in one go
795   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
796   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
797
798   auto out = addOutputBuffer(
799       &stream,
800       (maxCompressedLength <= maxSingleStepLength ?
801        maxCompressedLength :
802        defaultBufferLength));
803
804   for (auto& range : *data) {
805     uint64_t remaining = range.size();
806     uint64_t written = 0;
807     while (remaining) {
808       uint32_t step = (remaining > maxSingleStepLength ?
809                        maxSingleStepLength : remaining);
810       stream.next_in = const_cast<uint8_t*>(range.data() + written);
811       stream.avail_in = step;
812       remaining -= step;
813       written += step;
814
815       while (stream.avail_in != 0) {
816         if (stream.avail_out == 0) {
817           out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
818         }
819
820         rc = deflate(&stream, Z_NO_FLUSH);
821
822         CHECK_EQ(rc, Z_OK) << stream.msg;
823       }
824     }
825   }
826
827   do {
828     if (stream.avail_out == 0) {
829       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
830     }
831
832     rc = deflate(&stream, Z_FINISH);
833   } while (rc == Z_OK);
834
835   CHECK_EQ(rc, Z_STREAM_END) << stream.msg;
836
837   out->prev()->trimEnd(stream.avail_out);
838
839   success = true;  // we survived
840
841   return out;
842 }
843
844 std::unique_ptr<IOBuf> ZlibCodec::doUncompress(const IOBuf* data,
845                                                uint64_t uncompressedLength) {
846   z_stream stream;
847   stream.zalloc = nullptr;
848   stream.zfree = nullptr;
849   stream.opaque = nullptr;
850
851   // "The windowBits parameter is the base two logarithm of the maximum window
852   // size (...) The default value is 15 (...) add 16 to decode only the gzip
853   // format (the zlib format will return a Z_DATA_ERROR)."
854   int windowBits = 15 + (type() == CodecType::GZIP ? 16 : 0);
855   int rc = inflateInit2(&stream, windowBits);
856   if (rc != Z_OK) {
857     throw std::runtime_error(to<std::string>(
858         "ZlibCodec: inflateInit error: ", rc, ": ", stream.msg));
859   }
860
861   stream.next_in = stream.next_out = nullptr;
862   stream.avail_in = stream.avail_out = 0;
863   stream.total_in = stream.total_out = 0;
864
865   bool success = false;
866
867   SCOPE_EXIT {
868     rc = inflateEnd(&stream);
869     // If we're here because of an exception, it's okay if some data
870     // got dropped.
871     CHECK(rc == Z_OK || (!success && rc == Z_DATA_ERROR))
872       << rc << ": " << stream.msg;
873   };
874
875   // Max 64MiB in one go
876   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
877   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
878
879   auto out = addOutputBuffer(
880       &stream,
881       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
882         uncompressedLength <= maxSingleStepLength) ?
883        uncompressedLength :
884        defaultBufferLength));
885
886   bool streamEnd = false;
887   for (auto& range : *data) {
888     if (range.empty()) {
889       continue;
890     }
891
892     stream.next_in = const_cast<uint8_t*>(range.data());
893     stream.avail_in = range.size();
894
895     while (stream.avail_in != 0) {
896       if (streamEnd) {
897         throw std::runtime_error(to<std::string>(
898             "ZlibCodec: junk after end of data"));
899       }
900
901       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
902     }
903   }
904
905   while (!streamEnd) {
906     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
907   }
908
909   out->prev()->trimEnd(stream.avail_out);
910
911   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
912       uncompressedLength != stream.total_out) {
913     throw std::runtime_error(to<std::string>(
914         "ZlibCodec: invalid uncompressed length"));
915   }
916
917   success = true;  // we survived
918
919   return out;
920 }
921
922 #endif  // FOLLY_HAVE_LIBZ
923
924 #if FOLLY_HAVE_LIBLZMA
925
926 /**
927  * LZMA2 compression
928  */
929 class LZMA2Codec final : public Codec {
930  public:
931   static std::unique_ptr<Codec> create(int level, CodecType type);
932   explicit LZMA2Codec(int level, CodecType type);
933
934  private:
935   bool doNeedsUncompressedLength() const override;
936   uint64_t doMaxUncompressedLength() const override;
937
938   bool encodeSize() const { return type() == CodecType::LZMA2_VARINT_SIZE; }
939
940   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
941   std::unique_ptr<IOBuf> doUncompress(
942       const IOBuf* data,
943       uint64_t uncompressedLength) override;
944
945   std::unique_ptr<IOBuf> addOutputBuffer(lzma_stream* stream, size_t length);
946   bool doInflate(lzma_stream* stream, IOBuf* head, size_t bufferLength);
947
948   int level_;
949 };
950
951 std::unique_ptr<Codec> LZMA2Codec::create(int level, CodecType type) {
952   return make_unique<LZMA2Codec>(level, type);
953 }
954
955 LZMA2Codec::LZMA2Codec(int level, CodecType type) : Codec(type) {
956   DCHECK(type == CodecType::LZMA2 || type == CodecType::LZMA2_VARINT_SIZE);
957   switch (level) {
958   case COMPRESSION_LEVEL_FASTEST:
959     level = 0;
960     break;
961   case COMPRESSION_LEVEL_DEFAULT:
962     level = LZMA_PRESET_DEFAULT;
963     break;
964   case COMPRESSION_LEVEL_BEST:
965     level = 9;
966     break;
967   }
968   if (level < 0 || level > 9) {
969     throw std::invalid_argument(to<std::string>(
970         "LZMA2Codec: invalid level: ", level));
971   }
972   level_ = level;
973 }
974
975 bool LZMA2Codec::doNeedsUncompressedLength() const {
976   return false;
977 }
978
979 uint64_t LZMA2Codec::doMaxUncompressedLength() const {
980   // From lzma/base.h: "Stream is roughly 8 EiB (2^63 bytes)"
981   return uint64_t(1) << 63;
982 }
983
984 std::unique_ptr<IOBuf> LZMA2Codec::addOutputBuffer(
985     lzma_stream* stream,
986     size_t length) {
987
988   CHECK_EQ(stream->avail_out, 0);
989
990   auto buf = IOBuf::create(length);
991   buf->append(length);
992
993   stream->next_out = buf->writableData();
994   stream->avail_out = buf->length();
995
996   return buf;
997 }
998
999 std::unique_ptr<IOBuf> LZMA2Codec::doCompress(const IOBuf* data) {
1000   lzma_ret rc;
1001   lzma_stream stream = LZMA_STREAM_INIT;
1002
1003   rc = lzma_easy_encoder(&stream, level_, LZMA_CHECK_NONE);
1004   if (rc != LZMA_OK) {
1005     throw std::runtime_error(folly::to<std::string>(
1006       "LZMA2Codec: lzma_easy_encoder error: ", rc));
1007   }
1008
1009   SCOPE_EXIT { lzma_end(&stream); };
1010
1011   uint64_t uncompressedLength = data->computeChainDataLength();
1012   uint64_t maxCompressedLength = lzma_stream_buffer_bound(uncompressedLength);
1013
1014   // Max 64MiB in one go
1015   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20;    // 64MiB
1016   constexpr uint32_t defaultBufferLength = uint32_t(4) << 20;     // 4MiB
1017
1018   auto out = addOutputBuffer(
1019     &stream,
1020     (maxCompressedLength <= maxSingleStepLength ?
1021      maxCompressedLength :
1022      defaultBufferLength));
1023
1024   if (encodeSize()) {
1025     auto size = IOBuf::createCombined(kMaxVarintLength64);
1026     encodeVarintToIOBuf(uncompressedLength, size.get());
1027     size->appendChain(std::move(out));
1028     out = std::move(size);
1029   }
1030
1031   for (auto& range : *data) {
1032     if (range.empty()) {
1033       continue;
1034     }
1035
1036     stream.next_in = const_cast<uint8_t*>(range.data());
1037     stream.avail_in = range.size();
1038
1039     while (stream.avail_in != 0) {
1040       if (stream.avail_out == 0) {
1041         out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1042       }
1043
1044       rc = lzma_code(&stream, LZMA_RUN);
1045
1046       if (rc != LZMA_OK) {
1047         throw std::runtime_error(folly::to<std::string>(
1048           "LZMA2Codec: lzma_code error: ", rc));
1049       }
1050     }
1051   }
1052
1053   do {
1054     if (stream.avail_out == 0) {
1055       out->prependChain(addOutputBuffer(&stream, defaultBufferLength));
1056     }
1057
1058     rc = lzma_code(&stream, LZMA_FINISH);
1059   } while (rc == LZMA_OK);
1060
1061   if (rc != LZMA_STREAM_END) {
1062     throw std::runtime_error(folly::to<std::string>(
1063       "LZMA2Codec: lzma_code ended with error: ", rc));
1064   }
1065
1066   out->prev()->trimEnd(stream.avail_out);
1067
1068   return out;
1069 }
1070
1071 bool LZMA2Codec::doInflate(lzma_stream* stream,
1072                           IOBuf* head,
1073                           size_t bufferLength) {
1074   if (stream->avail_out == 0) {
1075     head->prependChain(addOutputBuffer(stream, bufferLength));
1076   }
1077
1078   lzma_ret rc = lzma_code(stream, LZMA_RUN);
1079
1080   switch (rc) {
1081   case LZMA_OK:
1082     break;
1083   case LZMA_STREAM_END:
1084     return true;
1085   default:
1086     throw std::runtime_error(to<std::string>(
1087         "LZMA2Codec: lzma_code error: ", rc));
1088   }
1089
1090   return false;
1091 }
1092
1093 std::unique_ptr<IOBuf> LZMA2Codec::doUncompress(const IOBuf* data,
1094                                                uint64_t uncompressedLength) {
1095   lzma_ret rc;
1096   lzma_stream stream = LZMA_STREAM_INIT;
1097
1098   rc = lzma_auto_decoder(&stream, std::numeric_limits<uint64_t>::max(), 0);
1099   if (rc != LZMA_OK) {
1100     throw std::runtime_error(folly::to<std::string>(
1101       "LZMA2Codec: lzma_auto_decoder error: ", rc));
1102   }
1103
1104   SCOPE_EXIT { lzma_end(&stream); };
1105
1106   // Max 64MiB in one go
1107   constexpr uint32_t maxSingleStepLength = uint32_t(64) << 20; // 64MiB
1108   constexpr uint32_t defaultBufferLength = uint32_t(256) << 10; // 256 KiB
1109
1110   folly::io::Cursor cursor(data);
1111   if (encodeSize()) {
1112     const uint64_t actualUncompressedLength = decodeVarintFromCursor(cursor);
1113     if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1114         uncompressedLength != actualUncompressedLength) {
1115       throw std::runtime_error("LZMA2Codec: invalid uncompressed length");
1116     }
1117     uncompressedLength = actualUncompressedLength;
1118   }
1119
1120   auto out = addOutputBuffer(
1121       &stream,
1122       ((uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1123         uncompressedLength <= maxSingleStepLength)
1124            ? uncompressedLength
1125            : defaultBufferLength));
1126
1127   bool streamEnd = false;
1128   auto buf = cursor.peekBytes();
1129   while (!buf.empty()) {
1130     stream.next_in = const_cast<uint8_t*>(buf.data());
1131     stream.avail_in = buf.size();
1132
1133     while (stream.avail_in != 0) {
1134       if (streamEnd) {
1135         throw std::runtime_error(to<std::string>(
1136             "LZMA2Codec: junk after end of data"));
1137       }
1138
1139       streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1140     }
1141
1142     cursor.skip(buf.size());
1143     buf = cursor.peekBytes();
1144   }
1145
1146   while (!streamEnd) {
1147     streamEnd = doInflate(&stream, out.get(), defaultBufferLength);
1148   }
1149
1150   out->prev()->trimEnd(stream.avail_out);
1151
1152   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH &&
1153       uncompressedLength != stream.total_out) {
1154     throw std::runtime_error(
1155         to<std::string>("LZMA2Codec: invalid uncompressed length"));
1156   }
1157
1158   return out;
1159 }
1160
1161 #endif  // FOLLY_HAVE_LIBLZMA
1162
1163 #ifdef FOLLY_HAVE_LIBZSTD
1164
1165 /**
1166  * ZSTD compression
1167  */
1168 class ZSTDCodec final : public Codec {
1169  public:
1170   static std::unique_ptr<Codec> create(int level, CodecType);
1171   explicit ZSTDCodec(int level, CodecType type);
1172
1173  private:
1174   bool doNeedsUncompressedLength() const override;
1175   std::unique_ptr<IOBuf> doCompress(const IOBuf* data) override;
1176   std::unique_ptr<IOBuf> doUncompress(
1177       const IOBuf* data,
1178       uint64_t uncompressedLength) override;
1179
1180   int level_;
1181 };
1182
1183 std::unique_ptr<Codec> ZSTDCodec::create(int level, CodecType type) {
1184   return make_unique<ZSTDCodec>(level, type);
1185 }
1186
1187 ZSTDCodec::ZSTDCodec(int level, CodecType type) : Codec(type) {
1188   DCHECK(type == CodecType::ZSTD);
1189   switch (level) {
1190     case COMPRESSION_LEVEL_FASTEST:
1191       level = 1;
1192       break;
1193     case COMPRESSION_LEVEL_DEFAULT:
1194       level = 1;
1195       break;
1196     case COMPRESSION_LEVEL_BEST:
1197       level = 19;
1198       break;
1199   }
1200   if (level < 1 || level > ZSTD_maxCLevel()) {
1201     throw std::invalid_argument(
1202         to<std::string>("ZSTD: invalid level: ", level));
1203   }
1204   level_ = level;
1205 }
1206
1207 bool ZSTDCodec::doNeedsUncompressedLength() const {
1208   return false;
1209 }
1210
1211 void zstdThrowIfError(size_t rc) {
1212   if (!ZSTD_isError(rc)) {
1213     return;
1214   }
1215   throw std::runtime_error(
1216       to<std::string>("ZSTD returned an error: ", ZSTD_getErrorName(rc)));
1217 }
1218
1219 std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
1220   // Support earlier versions of the codec (working with a single IOBuf,
1221   // and using ZSTD_decompress which requires ZSTD frame to contain size,
1222   // which isn't populated by streaming API).
1223   if (!data->isChained()) {
1224     auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
1225     const auto rc = ZSTD_compress(
1226         out->writableData(),
1227         out->capacity(),
1228         data->data(),
1229         data->length(),
1230         level_);
1231     zstdThrowIfError(rc);
1232     out->append(rc);
1233     return out;
1234   }
1235
1236   auto zcs = ZSTD_createCStream();
1237   SCOPE_EXIT {
1238     ZSTD_freeCStream(zcs);
1239   };
1240
1241   auto rc = ZSTD_initCStream(zcs, level_);
1242   zstdThrowIfError(rc);
1243
1244   Cursor cursor(data);
1245   auto result = IOBuf::createCombined(ZSTD_compressBound(cursor.totalLength()));
1246
1247   ZSTD_outBuffer out;
1248   out.dst = result->writableTail();
1249   out.size = result->capacity();
1250   out.pos = 0;
1251
1252   for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
1253     ZSTD_inBuffer in;
1254     in.src = buffer.data();
1255     in.size = buffer.size();
1256     for (in.pos = 0; in.pos != in.size;) {
1257       rc = ZSTD_compressStream(zcs, &out, &in);
1258       zstdThrowIfError(rc);
1259     }
1260     cursor.skip(in.size);
1261     buffer = cursor.peekBytes();
1262   }
1263
1264   rc = ZSTD_endStream(zcs, &out);
1265   zstdThrowIfError(rc);
1266   CHECK_EQ(rc, 0);
1267
1268   result->append(out.pos);
1269   return result;
1270 }
1271
1272 static std::unique_ptr<IOBuf> zstdUncompressBuffer(
1273     const IOBuf* data,
1274     uint64_t uncompressedLength) {
1275   // Check preconditions
1276   DCHECK(!data->isChained());
1277   DCHECK(uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH);
1278
1279   auto uncompressed = IOBuf::create(uncompressedLength);
1280   const auto decompressedSize = ZSTD_decompress(
1281       uncompressed->writableTail(),
1282       uncompressed->tailroom(),
1283       data->data(),
1284       data->length());
1285   zstdThrowIfError(decompressedSize);
1286   if (decompressedSize != uncompressedLength) {
1287     throw std::runtime_error("ZSTD: invalid uncompressed length");
1288   }
1289   uncompressed->append(decompressedSize);
1290   return uncompressed;
1291 }
1292
1293 static std::unique_ptr<IOBuf> zstdUncompressStream(
1294     const IOBuf* data,
1295     uint64_t uncompressedLength) {
1296   auto zds = ZSTD_createDStream();
1297   SCOPE_EXIT {
1298     ZSTD_freeDStream(zds);
1299   };
1300
1301   auto rc = ZSTD_initDStream(zds);
1302   zstdThrowIfError(rc);
1303
1304   ZSTD_outBuffer out{};
1305   ZSTD_inBuffer in{};
1306
1307   auto outputSize = ZSTD_DStreamOutSize();
1308   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH) {
1309     outputSize = uncompressedLength;
1310   }
1311
1312   IOBufQueue queue(IOBufQueue::cacheChainLength());
1313
1314   Cursor cursor(data);
1315   for (rc = 0;;) {
1316     if (in.pos == in.size) {
1317       auto buffer = cursor.peekBytes();
1318       in.src = buffer.data();
1319       in.size = buffer.size();
1320       in.pos = 0;
1321       cursor.skip(in.size);
1322       if (rc > 1 && in.size == 0) {
1323         throw std::runtime_error(to<std::string>("ZSTD: incomplete input"));
1324       }
1325     }
1326     if (out.pos == out.size) {
1327       if (out.pos != 0) {
1328         queue.postallocate(out.pos);
1329       }
1330       auto buffer = queue.preallocate(outputSize, outputSize);
1331       out.dst = buffer.first;
1332       out.size = buffer.second;
1333       out.pos = 0;
1334       outputSize = ZSTD_DStreamOutSize();
1335     }
1336     rc = ZSTD_decompressStream(zds, &out, &in);
1337     zstdThrowIfError(rc);
1338     if (rc == 0) {
1339       break;
1340     }
1341   }
1342   if (out.pos != 0) {
1343     queue.postallocate(out.pos);
1344   }
1345   if (in.pos != in.size || !cursor.isAtEnd()) {
1346     throw std::runtime_error("ZSTD: junk after end of data");
1347   }
1348   if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1349       queue.chainLength() != uncompressedLength) {
1350     throw std::runtime_error("ZSTD: invalid uncompressed length");
1351   }
1352
1353   return queue.move();
1354 }
1355
1356 std::unique_ptr<IOBuf> ZSTDCodec::doUncompress(
1357     const IOBuf* data,
1358     uint64_t uncompressedLength) {
1359   {
1360     // Read decompressed size from frame if available in first IOBuf.
1361     const auto decompressedSize =
1362         ZSTD_getDecompressedSize(data->data(), data->length());
1363     if (decompressedSize != 0) {
1364       if (uncompressedLength != Codec::UNKNOWN_UNCOMPRESSED_LENGTH &&
1365           uncompressedLength != decompressedSize) {
1366         throw std::runtime_error("ZSTD: invalid uncompressed length");
1367       }
1368       uncompressedLength = decompressedSize;
1369     }
1370   }
1371   // Faster to decompress using ZSTD_decompress() if we can.
1372   if (uncompressedLength != UNKNOWN_UNCOMPRESSED_LENGTH && !data->isChained()) {
1373     return zstdUncompressBuffer(data, uncompressedLength);
1374   }
1375   // Fall back to slower streaming decompression.
1376   return zstdUncompressStream(data, uncompressedLength);
1377 }
1378
1379 #endif  // FOLLY_HAVE_LIBZSTD
1380
1381 }  // namespace
1382
1383 typedef std::unique_ptr<Codec> (*CodecFactory)(int, CodecType);
1384 static constexpr CodecFactory
1385     codecFactories[static_cast<size_t>(CodecType::NUM_CODEC_TYPES)] = {
1386         nullptr, // USER_DEFINED
1387         NoCompressionCodec::create,
1388
1389 #if FOLLY_HAVE_LIBLZ4
1390         LZ4Codec::create,
1391 #else
1392         nullptr,
1393 #endif
1394
1395 #if FOLLY_HAVE_LIBSNAPPY
1396         SnappyCodec::create,
1397 #else
1398         nullptr,
1399 #endif
1400
1401 #if FOLLY_HAVE_LIBZ
1402         ZlibCodec::create,
1403 #else
1404         nullptr,
1405 #endif
1406
1407 #if FOLLY_HAVE_LIBLZ4
1408         LZ4Codec::create,
1409 #else
1410         nullptr,
1411 #endif
1412
1413 #if FOLLY_HAVE_LIBLZMA
1414         LZMA2Codec::create,
1415         LZMA2Codec::create,
1416 #else
1417         nullptr,
1418         nullptr,
1419 #endif
1420
1421 #if FOLLY_HAVE_LIBZSTD
1422         ZSTDCodec::create,
1423 #else
1424         nullptr,
1425 #endif
1426
1427 #if FOLLY_HAVE_LIBZ
1428         ZlibCodec::create,
1429 #else
1430         nullptr,
1431 #endif
1432
1433 #if FOLLY_HAVE_LIBLZ4
1434         LZ4FrameCodec::create,
1435 #else
1436         nullptr,
1437 #endif
1438 };
1439
1440 bool hasCodec(CodecType type) {
1441   size_t idx = static_cast<size_t>(type);
1442   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1443     throw std::invalid_argument(
1444         to<std::string>("Compression type ", idx, " invalid"));
1445   }
1446   return codecFactories[idx] != nullptr;
1447 }
1448
1449 std::unique_ptr<Codec> getCodec(CodecType type, int level) {
1450   size_t idx = static_cast<size_t>(type);
1451   if (idx >= static_cast<size_t>(CodecType::NUM_CODEC_TYPES)) {
1452     throw std::invalid_argument(
1453         to<std::string>("Compression type ", idx, " invalid"));
1454   }
1455   auto factory = codecFactories[idx];
1456   if (!factory) {
1457     throw std::invalid_argument(to<std::string>(
1458         "Compression type ", idx, " not supported"));
1459   }
1460   auto codec = (*factory)(level, type);
1461   DCHECK_EQ(static_cast<size_t>(codec->type()), idx);
1462   return codec;
1463 }
1464
1465 }}  // namespaces