diff --git a/clang/lib/Serialization/ASTReader.cpp b/clang/lib/Serialization/ASTReader.cpp --- a/clang/lib/Serialization/ASTReader.cpp +++ b/clang/lib/Serialization/ASTReader.cpp @@ -1452,19 +1452,24 @@ unsigned RecCode = MaybeRecCode.get(); if (RecCode == SM_SLOC_BUFFER_BLOB_COMPRESSED) { - if (!llvm::compression::zlib::isAvailable()) { - Error("zlib is not available"); + // Inspect the first two bytes to differentiate zlib (\x1f\x8b) and zstd. + const llvm::compression::Format F = + Blob.size() >= 2 && memcmp(Blob.data(), "\x1f\x8b", 2) == 0 + ? llvm::compression::Format::Zlib + : llvm::compression::Format::Zstd; + if (const char *Reason = llvm::compression::getReasonIfUnsupported(F)) { + Error(Reason); return nullptr; } - SmallVector Uncompressed; - if (llvm::Error E = llvm::compression::zlib::decompress( - llvm::arrayRefFromStringRef(Blob), Uncompressed, Record[0])) { + SmallVector Decompressed; + if (llvm::Error E = llvm::compression::decompress( + F, llvm::arrayRefFromStringRef(Blob), Decompressed, Record[0])) { Error("could not decompress embedded file contents: " + llvm::toString(std::move(E))); return nullptr; } return llvm::MemoryBuffer::getMemBufferCopy( - llvm::toStringRef(Uncompressed), Name); + llvm::toStringRef(Decompressed), Name); } else if (RecCode == SM_SLOC_BUFFER_BLOB) { return llvm::MemoryBuffer::getMemBuffer(Blob.drop_back(1), Name, true); } else { diff --git a/clang/lib/Serialization/ASTWriter.cpp b/clang/lib/Serialization/ASTWriter.cpp --- a/clang/lib/Serialization/ASTWriter.cpp +++ b/clang/lib/Serialization/ASTWriter.cpp @@ -1979,6 +1979,14 @@ // Compress the buffer if possible. We expect that almost all PCM // consumers will not want its contents. SmallVector CompressedBuffer; + if (llvm::compression::zstd::isAvailable()) { + llvm::compression::zstd::compress( + llvm::arrayRefFromStringRef(Blob.drop_back(1)), CompressedBuffer, 9); + RecordDataType Record[] = {SM_SLOC_BUFFER_BLOB_COMPRESSED, Blob.size() - 1}; + Stream.EmitRecordWithBlob(SLocBufferBlobCompressedAbbrv, Record, + llvm::toStringRef(CompressedBuffer)); + return; + } if (llvm::compression::zlib::isAvailable()) { llvm::compression::zlib::compress( llvm::arrayRefFromStringRef(Blob.drop_back(1)), CompressedBuffer); diff --git a/clang/test/Modules/embed-files-compressed.cpp b/clang/test/Modules/embed-files-compressed.cpp --- a/clang/test/Modules/embed-files-compressed.cpp +++ b/clang/test/Modules/embed-files-compressed.cpp @@ -1,4 +1,4 @@ -// REQUIRES: zlib +// REQUIRES: zlib || zstd // REQUIRES: shell // // RUN: rm -rf %t