1 files changed, 29 insertions, 5 deletions
diff --git a/llvm/lib/ProfileData/Coverage/CoverageMappingWriter.cpp b/llvm/lib/ProfileData/Coverage/CoverageMappingWriter.cpp
index d75854a60d1ee..8d3c429c44843 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMappingWriter.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMappingWriter.cpp
@@ -11,9 +11,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/Coverage/CoverageMappingWriter.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
@@ -34,12 +36,34 @@ CoverageFilenamesSectionWriter::CoverageFilenamesSectionWriter(
 #endif
 }
 
-void CoverageFilenamesSectionWriter::write(raw_ostream &OS) {
-  encodeULEB128(Filenames.size(), OS);
-  for (const auto &Filename : Filenames) {
-    encodeULEB128(Filename.size(), OS);
-    OS << Filename;
+void CoverageFilenamesSectionWriter::write(raw_ostream &OS, bool Compress) {
+  std::string FilenamesStr;
+  {
+    raw_string_ostream FilenamesOS{FilenamesStr};
+    for (const auto &Filename : Filenames) {
+      encodeULEB128(Filename.size(), FilenamesOS);
+      FilenamesOS << Filename;
+    }
+  }
+
+  SmallString<128> CompressedStr;
+  bool doCompression =
+      Compress && zlib::isAvailable() && DoInstrProfNameCompression;
+  if (doCompression) {
+    auto E =
+        zlib::compress(FilenamesStr, CompressedStr, zlib::BestSizeCompression);
+    if (E)
+      report_bad_alloc_error("Failed to zlib compress coverage data");
   }
+
+  // ::= <num-filenames>
+  //     <uncompressed-len>
+  //     <compressed-len-or-zero>
+  //     (<compressed-filenames> | <uncompressed-filenames>)
+  encodeULEB128(Filenames.size(), OS);
+  encodeULEB128(FilenamesStr.size(), OS);
+  encodeULEB128(doCompression ? CompressedStr.size() : 0U, OS);
+  OS << (doCompression ? StringRef(CompressedStr) : StringRef(FilenamesStr));
 }
 
 namespace {