1 files changed, 84 insertions, 18 deletions

diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index 068ce5bf959..03446367665 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -21,6 +21,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/ErrorOr.h"
@@ -72,21 +73,58 @@ SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
   return sampleprof_error::success;
 }
 
+SecHdrTableEntry &
+SampleProfileWriterExtBinaryBase::getEntryInLayout(SecType Type) {
+  auto SecIt = std::find_if(
+      SectionLayout.begin(), SectionLayout.end(),
+      [=](const auto &Entry) -> bool { return Entry.Type == Type; });
+  return *SecIt;
+}
+
 /// Return the current position and prepare to use it as the start
 /// position of a section.
-uint64_t SampleProfileWriterExtBinaryBase::markSectionStart() {
-  return OutputStream->tell();
+uint64_t SampleProfileWriterExtBinaryBase::markSectionStart(SecType Type) {
+  uint64_t SectionStart = OutputStream->tell();
+  auto &Entry = getEntryInLayout(Type);
+  // Use LocalBuf as a temporary output for writting data.
+  if (hasSecFlag(Entry, SecFlagCompress))
+    LocalBufStream.swap(OutputStream);
+  return SectionStart;
+}
+
+std::error_code SampleProfileWriterExtBinaryBase::compressAndOutput() {
+  if (!llvm::zlib::isAvailable())
+    return sampleprof_error::zlib_unavailable;
+  std::string &UncompressedStrings =
+      static_cast<raw_string_ostream *>(LocalBufStream.get())->str();
+  if (UncompressedStrings.size() == 0)
+    return sampleprof_error::success;
+  auto &OS = *OutputStream;
+  SmallString<128> CompressedStrings;
+  llvm::Error E = zlib::compress(UncompressedStrings, CompressedStrings,
+                                 zlib::BestSizeCompression);
+  if (E)
+    return sampleprof_error::compress_failed;
+  encodeULEB128(UncompressedStrings.size(), OS);
+  encodeULEB128(CompressedStrings.size(), OS);
+  OS << CompressedStrings.str();
+  UncompressedStrings.clear();
+  return sampleprof_error::success;
 }
 
-/// Add a new section into section header table. Return the position
-/// of SectionEnd.
-uint64_t
-SampleProfileWriterExtBinaryBase::addNewSection(SecType Sec,
+/// Add a new section into section header table.
+std::error_code
+SampleProfileWriterExtBinaryBase::addNewSection(SecType Type,
                                                 uint64_t SectionStart) {
-  uint64_t SectionEnd = OutputStream->tell();
-  SecHdrTable.push_back(
-      {Sec, 0, SectionStart - FileStart, SectionEnd - SectionStart});
-  return SectionEnd;
+  auto Entry = getEntryInLayout(Type);
+  if (hasSecFlag(Entry, SecFlagCompress)) {
+    LocalBufStream.swap(OutputStream);
+    if (std::error_code EC = compressAndOutput())
+      return EC;
+  }
+  SecHdrTable.push_back({Type, Entry.Flags, SectionStart - FileStart,
+                         OutputStream->tell() - SectionStart});
+  return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::write(
@@ -94,6 +132,8 @@ std::error_code SampleProfileWriterExtBinaryBase::write(
   if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
+  std::string LocalBuf;
+  LocalBufStream = std::make_unique<raw_string_ostream>(LocalBuf);
   if (std::error_code EC = writeSections(ProfileMap))
     return EC;
 
@@ -105,28 +145,38 @@ std::error_code SampleProfileWriterExtBinaryBase::write(
 
 std::error_code SampleProfileWriterExtBinary::writeSections(
     const StringMap<FunctionSamples> &ProfileMap) {
-  uint64_t SectionStart = markSectionStart();
+  uint64_t SectionStart = markSectionStart(SecProfSummary);
   computeSummary(ProfileMap);
   if (auto EC = writeSummary())
     return EC;
-  SectionStart = addNewSection(SecProfSummary, SectionStart);
+  if (std::error_code EC = addNewSection(SecProfSummary, SectionStart))
+    return EC;
 
   // Generate the name table for all the functions referenced in the profile.
+  SectionStart = markSectionStart(SecNameTable);
   for (const auto &I : ProfileMap) {
     addName(I.first());
     addNames(I.second);
   }
   writeNameTable();
-  SectionStart = addNewSection(SecNameTable, SectionStart);
+  if (std::error_code EC = addNewSection(SecNameTable, SectionStart))
+    return EC;
 
+  SectionStart = markSectionStart(SecLBRProfile);
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
-  SectionStart = addNewSection(SecLBRProfile, SectionStart);
+  if (std::error_code EC = addNewSection(SecLBRProfile, SectionStart))
+    return EC;
+
+  if (ProfSymList && ProfSymList->toCompress())
+    setToCompressSection(SecProfileSymbolList);
 
+  SectionStart = markSectionStart(SecProfileSymbolList);
   if (ProfSymList && ProfSymList->size() > 0)
     if (std::error_code EC = ProfSymList->write(*OutputStream))
       return EC;
-  addNewSection(SecProfileSymbolList, SectionStart);
+  if (std::error_code EC = addNewSection(SecProfileSymbolList, SectionStart))
+    return EC;
 
   return sampleprof_error::success;
 }
@@ -308,6 +358,23 @@ std::error_code SampleProfileWriterBinary::writeHeader(
   return sampleprof_error::success;
 }
 
+void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
+  for (auto &Entry : SectionLayout)
+    addSecFlags(Entry, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::setToCompressSection(SecType Type) {
+  addSectionFlags(Type, SecFlagCompress);
+}
+
+void SampleProfileWriterExtBinaryBase::addSectionFlags(SecType Type,
+                                                       SecFlags Flags) {
+  for (auto &Entry : SectionLayout) {
+    if (Entry.Type == Type)
+      addSecFlags(Entry, Flags);
+  }
+}
+
 void SampleProfileWriterExtBinaryBase::allocSecHdrTable() {
   support::endian::Writer Writer(*OutputStream, support::little);
 
@@ -342,9 +409,9 @@ std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
   // to adjust the order in SecHdrTable to be consistent with
   // SectionLayout when we write SecHdrTable to the memory.
   for (uint32_t i = 0; i < SectionLayout.size(); i++) {
-    uint32_t idx = IndexMap[static_cast<uint32_t>(SectionLayout[i])];
+    uint32_t idx = IndexMap[static_cast<uint32_t>(SectionLayout[i].Type)];
     Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Type));
-    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Flag));
+    Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Flags));
     Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Offset));
     Writer.write(static_cast<uint64_t>(SecHdrTable[idx].Size));
   }
@@ -362,7 +429,6 @@ std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
   FileStart = OS.tell();
   writeMagicIdent(Format);
 
-  initSectionLayout();
   allocSecHdrTable();
   return sampleprof_error::success;
 }