diff --git a/llvm/include/llvm/Support/Compression.h b/llvm/include/llvm/Support/Compression.h index 246ccbd6f6dcf..64dadf96feb43 100644 --- a/llvm/include/llvm/Support/Compression.h +++ b/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 class SmallVectorImpl; @@ -74,6 +75,8 @@ LLVM_ABI Error decompress(ArrayRef Input, SmallVectorImpl &Output, size_t UncompressedSize); +// Get the size of the decompressed data. +Expected getDecompressedSize(ArrayRef Input); } // End of namespace zstd enum class Format { diff --git a/llvm/lib/Support/Compression.cpp b/llvm/lib/Support/Compression.cpp index 3979ca6acaf74..ffa073d0be932 100644 --- a/llvm/lib/Support/Compression.cpp +++ b/llvm/lib/Support/Compression.cpp @@ -225,6 +225,17 @@ Error zstd::decompress(ArrayRef Input, return E; } +Expected zstd::getDecompressedSize(ArrayRef Input) { + + unsigned long long Res = ZSTD_getFrameContentSize(Input.data(), Input.size()); + + if (ZSTD_isError(Res)) + return make_error(ZSTD_getErrorName(Res), + inconvertibleErrorCode()); + + return Res; +} + #else bool zstd::isAvailable() { return false; } void zstd::compress(ArrayRef Input, @@ -241,4 +252,7 @@ Error zstd::decompress(ArrayRef Input, size_t UncompressedSize) { llvm_unreachable("zstd::decompress is unavailable"); } +Expected zstd::getDecompressedSize(ArrayRef Input) { + llvm_unreachable("zstd::getDecompressedSize is unavailable"); +} #endif diff --git a/llvm/unittests/Support/CompressionTest.cpp b/llvm/unittests/Support/CompressionTest.cpp index 5d326cafbe3a1..4ad1d3ba4fec6 100644 --- a/llvm/unittests/Support/CompressionTest.cpp +++ b/llvm/unittests/Support/CompressionTest.cpp @@ -73,6 +73,11 @@ static void testZstdCompression(StringRef Input, int Level) { SmallVector Uncompressed; zstd::compress(arrayRefFromStringRef(Input), Compressed, Level); + // Check that getDecompressedSize returns the size of the original buffer. + Expected 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));