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

llvm/include/llvm/Support/Compression.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/Error.h"
 
 namespace llvm {
 template <typename T> class SmallVectorImpl;
@@ -74,6 +75,8 @@ LLVM_ABI Error decompress(ArrayRef<uint8_t> Input,
                           SmallVectorImpl<uint8_t> &Output,
                           size_t UncompressedSize);
 
+// Get the size of the decompressed data.
+Expected<uint64_t> getDecompressedSize(ArrayRef<uint8_t> Input);
 } // End of namespace zstd
 
 enum class Format {

llvm/lib/Support/Compression.cpp

@@ -225,6 +225,17 @@ Error zstd::decompress(ArrayRef<uint8_t> Input,
   return E;
 }
 
+Expected<uint64_t> zstd::getDecompressedSize(ArrayRef<uint8_t> Input) {
+
+  unsigned long long Res = ZSTD_getFrameContentSize(Input.data(), Input.size());
+
+  if (ZSTD_isError(Res))
+    return make_error<StringError>(ZSTD_getErrorName(Res),
+                                                     inconvertibleErrorCode());
+
+  return Res;
+}
+
 #else
 bool zstd::isAvailable() { return false; }
 void zstd::compress(ArrayRef<uint8_t> Input,
@@ -241,4 +252,7 @@ Error zstd::decompress(ArrayRef<uint8_t> Input,
                        size_t UncompressedSize) {
   llvm_unreachable("zstd::decompress is unavailable");
 }
+Expected<uint64_t> zstd::getDecompressedSize(ArrayRef<uint8_t> Input) {
+  llvm_unreachable("zstd::getDecompressedSize is unavailable");
+}
 #endif

llvm/unittests/Support/CompressionTest.cpp

@@ -73,6 +73,11 @@ 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.
+  Expected<uint64_t> DecompressedSize = zstd::getDecompressedSize(Compressed);
+  EXPECT_FALSE(!DecompressedSize);
+  EXPECT_EQ(DecompressedSize.get(), Input.size());
+
   // Check that uncompressed buffer is the same as original.
   Error E = zstd::decompress(Compressed, Uncompressed, Input.size());
   EXPECT_FALSE(std::move(E));