3 files changed, 115 insertions, 1 deletions

diff --git a/llvm/lib/ProfileData/SampleProf.cpp b/llvm/lib/ProfileData/SampleProf.cpp
index ce0f537f8d5..36c39c12432 100644
--- a/llvm/lib/ProfileData/SampleProf.cpp
+++ b/llvm/lib/ProfileData/SampleProf.cpp
@@ -15,8 +15,11 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/Compression.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LEB128.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/raw_ostream.h"
 #include <string>
@@ -66,6 +69,12 @@ class SampleProfErrorCategoryType : public std::error_category {
       return "Counter overflow";
     case sampleprof_error::ostream_seek_unsupported:
       return "Ostream does not support seek";
+    case sampleprof_error::compress_failed:
+      return "Compress failure";
+    case sampleprof_error::uncompress_failed:
+      return "Uncompress failure";
+    case sampleprof_error::zlib_unavailable:
+      return "Zlib is unavailable";
     }
     llvm_unreachable("A value of sampleprof_error has no message.");
   }
@@ -188,3 +197,75 @@ FunctionSamples::findFunctionSamples(const DILocation *DIL) const {
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void FunctionSamples::dump() const { print(dbgs(), 0); }
 #endif
+
+std::error_code ProfileSymbolList::read(uint64_t CompressSize,
+                                        uint64_t UncompressSize,
+                                        const uint8_t *Data) {
+  const char *ListStart = reinterpret_cast<const char *>(Data);
+  // CompressSize being non-zero means the profile is compressed and
+  // needs to be uncompressed first.
+  if (CompressSize) {
+    if (!llvm::zlib::isAvailable())
+      return sampleprof_error::zlib_unavailable;
+
+    StringRef CompressedStrings(reinterpret_cast<const char *>(Data),
+                                CompressSize);
+    char *Buffer = Allocator.Allocate<char>(UncompressSize);
+    llvm::Error E = zlib::uncompress(CompressedStrings, Buffer, UncompressSize);
+    if (E)
+      return sampleprof_error::uncompress_failed;
+    ListStart = Buffer;
+  }
+
+  uint64_t Size = 0;
+  while (Size < UncompressSize) {
+    StringRef Str(ListStart + Size);
+    add(Str);
+    Size += Str.size() + 1;
+  }
+  return sampleprof_error::success;
+}
+
+std::error_code ProfileSymbolList::write(raw_ostream &OS) {
+  // Sort the symbols before doing compression. It will make the
+  // compression much more effective.
+  std::vector<StringRef> SortedList;
+  SortedList.insert(SortedList.begin(), Syms.begin(), Syms.end());
+  llvm::sort(SortedList);
+
+  std::string UncompressedStrings;
+  for (auto &Sym : SortedList) {
+    UncompressedStrings.append(Sym.str());
+    UncompressedStrings.append(1, '\0');
+  }
+
+  if (ToCompress) {
+    if (!llvm::zlib::isAvailable())
+      return sampleprof_error::zlib_unavailable;
+    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();
+  } else {
+    encodeULEB128(UncompressedStrings.size(), OS);
+    // If profile symbol list is not compressed, we will still save
+    // a compressed size value, but the value of the size is 0.
+    encodeULEB128(0, OS);
+    OS << UncompressedStrings;
+  }
+  return sampleprof_error::success;
+}
+
+void ProfileSymbolList::dump(raw_ostream &OS) const {
+  OS << "======== Dump profile symbol list ========\n";
+  std::vector<StringRef> SortedList;
+  SortedList.insert(SortedList.begin(), Syms.begin(), Syms.end());
+  llvm::sort(SortedList);
+
+  for (auto &Sym : SortedList)
+    OS << Sym << "\n";
+}
diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
index e309d5b270a..f6b33d96261 100644
--- a/llvm/lib/ProfileData/SampleProfReader.cpp
+++ b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -486,12 +486,40 @@ SampleProfileReaderExtBinary::readOneSection(const uint8_t *Start,
         return EC;
     }
     break;
+  case SecProfileSymbolList:
+    if (std::error_code EC = readProfileSymbolList())
+      return EC;
+    break;
   default:
     break;
   }
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileReaderExtBinary::readProfileSymbolList() {
+  auto UncompressSize = readNumber<uint64_t>();
+  if (std::error_code EC = UncompressSize.getError())
+    return EC;
+
+  auto CompressSize = readNumber<uint64_t>();
+  if (std::error_code EC = CompressSize.getError())
+    return EC;
+
+  if (!ProfSymList)
+    ProfSymList = std::make_unique<ProfileSymbolList>();
+
+  if (std::error_code EC =
+          ProfSymList->read(*CompressSize, *UncompressSize, Data))
+    return EC;
+
+  // CompressSize is zero only when ProfileSymbolList is not compressed.
+  if (*CompressSize == 0)
+    Data = Data + *UncompressSize;
+  else
+    Data = Data + *CompressSize;
+  return sampleprof_error::success;
+}
+
 std::error_code SampleProfileReaderExtBinaryBase::read() {
   const uint8_t *BufStart =
       reinterpret_cast<const uint8_t *>(Buffer->getBufferStart());
diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index 47f97de2a3b..068ce5bf959 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -121,7 +121,12 @@ std::error_code SampleProfileWriterExtBinary::writeSections(
 
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
-  addNewSection(SecLBRProfile, SectionStart);
+  SectionStart = addNewSection(SecLBRProfile, SectionStart);
+
+  if (ProfSymList && ProfSymList->size() > 0)
+    if (std::error_code EC = ProfSymList->write(*OutputStream))
+      return EC;
+  addNewSection(SecProfileSymbolList, SectionStart);
 
   return sampleprof_error::success;
 }