[llvm] Add API to get decompressed size of a zstd compressed buffer.

llvm/include/llvm/Support/Compression.h

@@ -71,6 +71,8 @@ Error decompress(ArrayRef<uint8_t> Input, uint8_t *Output,
 Error decompress(ArrayRef<uint8_t> Input, SmallVectorImpl<uint8_t> &Output,
                  size_t UncompressedSize);
 
+// Get the size of the decompressed data.
+Error getDecompressedSize(ArrayRef<uint8_t> Input, size_t &UncompressedSize);
 } // End of namespace zstd
 
 enum class Format {

llvm/lib/Support/Compression.cpp

@@ -224,6 +224,20 @@ Error zstd::decompress(ArrayRef<uint8_t> Input,
   return E;
 }
 
+Error zstd::getDecompressedSize(ArrayRef<uint8_t> Input,
+                                size_t &UncompressedSize) {
+
+  unsigned long long Res = ZSTD_getFrameContentSize(Input.data(), Input.size());
+
+  // ZSTD_getFrameContentSize returns unsigned long long, but the size
+  // of uncompressed data should be bounded by size_t.
+  UncompressedSize = static_cast<size_t>(Res);
+
+  return ZSTD_isError(Res) ? make_error<StringError>(ZSTD_getErrorName(Res),
+                                                     inconvertibleErrorCode())
+                           : Error::success();
+}
+
 #else
 bool zstd::isAvailable() { return false; }
 void zstd::compress(ArrayRef<uint8_t> Input,
@@ -240,4 +254,8 @@ Error zstd::decompress(ArrayRef<uint8_t> Input,
                        size_t UncompressedSize) {
   llvm_unreachable("zstd::decompress is unavailable");
 }
+Error zstd::getDecompressedSize(ArrayRef<uint8_t> Input,
+                                size_t &UncompressedSize) {
+  llvm_unreachable("zstd::getDecompressedSize is unavailable");
+}
 #endif

llvm/unittests/Support/CompressionTest.cpp

@@ -73,8 +73,15 @@ static void testZstdCompression(StringRef Input, int Level) {
   SmallVector<uint8_t, 0> Uncompressed;
   zstd::compress(arrayRefFromStringRef(Input), Compressed, Level);
 
+  // Check that getDecompressedSize returns the size of the original buffer.
+  size_t DecompressedSize;
+  Error E = zstd::getDecompressedSize(Compressed, DecompressedSize);
+  EXPECT_FALSE(std::move(E));
+
+  EXPECT_EQ(DecompressedSize, Input.size());
+
   // Check that uncompressed buffer is the same as original.
-  Error E = zstd::decompress(Compressed, Uncompressed, Input.size());
+  E = zstd::decompress(Compressed, Uncompressed, Input.size());
   EXPECT_FALSE(std::move(E));
   EXPECT_EQ(Input, toStringRef(Uncompressed));