3 files changed, 94 insertions, 21 deletions

diff --git a/llvm/lib/ProfileData/SampleProf.cpp b/llvm/lib/ProfileData/SampleProf.cpp
index 027fa81790b..c6960ba4bd9 100644
--- a/llvm/lib/ProfileData/SampleProf.cpp
+++ b/llvm/lib/ProfileData/SampleProf.cpp
@@ -38,6 +38,10 @@ class SampleProfErrorCategoryType : public std::error_category {
       return "Malformed profile data";
     case sampleprof_error::unrecognized_format:
       return "Unrecognized profile encoding format";
+    case sampleprof_error::unsupported_writing_format:
+      return "Profile encoding format unsupported for writing operations";
+    case sampleprof_error::truncated_name_table:
+      return "Truncated function name table";
     case sampleprof_error::not_implemented:
       return "Unimplemented feature";
     }
diff --git a/llvm/lib/ProfileData/SampleProfReader.cpp b/llvm/lib/ProfileData/SampleProfReader.cpp
index c620d4ca84a..a058274e8ee 100644
--- a/llvm/lib/ProfileData/SampleProfReader.cpp
+++ b/llvm/lib/ProfileData/SampleProfReader.cpp
@@ -378,6 +378,16 @@ ErrorOr<StringRef> SampleProfileReaderBinary::readString() {
   return Str;
 }
 
+ErrorOr<StringRef> SampleProfileReaderBinary::readStringFromTable() {
+  std::error_code EC;
+  auto Idx = readNumber<unsigned>();
+  if (std::error_code EC = Idx.getError())
+    return EC;
+  if (*Idx >= NameTable.size())
+    return sampleprof_error::truncated_name_table;
+  return NameTable[*Idx];
+}
+
 std::error_code
 SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
   auto Val = readNumber<unsigned>();
@@ -413,7 +423,7 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
       return EC;
 
     for (unsigned J = 0; J < *NumCalls; ++J) {
-      auto CalledFunction(readString());
+      auto CalledFunction(readStringFromTable());
       if (std::error_code EC = CalledFunction.getError())
         return EC;
 
@@ -442,7 +452,7 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
     if (std::error_code EC = Discriminator.getError())
       return EC;
 
-    auto FName(readString());
+    auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
@@ -457,7 +467,7 @@ SampleProfileReaderBinary::readProfile(FunctionSamples &FProfile) {
 
 std::error_code SampleProfileReaderBinary::read() {
   while (!at_eof()) {
-    auto FName(readString());
+    auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
@@ -489,6 +499,18 @@ std::error_code SampleProfileReaderBinary::readHeader() {
   else if (*Version != SPVersion())
     return sampleprof_error::unsupported_version;
 
+  // Read the name table.
+  auto Size = readNumber<size_t>();
+  if (std::error_code EC = Size.getError())
+    return EC;
+  NameTable.reserve(*Size);
+  for (size_t I = 0; I < *Size; ++I) {
+    auto Name(readString());
+    if (std::error_code EC = Name.getError())
+      return EC;
+    NameTable.push_back(*Name);
+  }
+
   return sampleprof_error::success;
 }
 
diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index 6be884b13e0..d3001aac3f6 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -30,7 +30,8 @@ using namespace llvm::sampleprof;
 using namespace llvm;
 
 /// \brief Write samples to a text file.
-bool SampleProfileWriterText::write(StringRef FName, const FunctionSamples &S) {
+std::error_code SampleProfileWriterText::write(StringRef FName,
+                                               const FunctionSamples &S) {
   OS << FName << ":" << S.getTotalSamples();
   if (Indent == 0)
     OS << ":" << S.getHeadSamples();
@@ -61,31 +62,74 @@ bool SampleProfileWriterText::write(StringRef FName, const FunctionSamples &S) {
       OS << Loc.LineOffset << ": ";
     else
       OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
-    write(Loc.CalleeName, CalleeSamples);
+    if (std::error_code EC = write(Loc.CalleeName, CalleeSamples))
+      return EC;
   }
   Indent -= 1;
 
-  return true;
+  return sampleprof_error::success;
 }
 
-SampleProfileWriterBinary::SampleProfileWriterBinary(StringRef F,
-                                                     std::error_code &EC)
-    : SampleProfileWriter(F, EC, sys::fs::F_None) {
-  if (EC)
-    return;
+std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
+  const auto &ret = NameTable.find(FName);
+  if (ret == NameTable.end())
+    return sampleprof_error::truncated_name_table;
+  encodeULEB128(ret->second, OS);
+  return sampleprof_error::success;
+}
+
+void SampleProfileWriterBinary::addName(StringRef FName) {
+  auto NextIdx = NameTable.size();
+  NameTable.insert(std::make_pair(FName, NextIdx));
+}
+
+void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
+  // Add all the names in indirect call targets.
+  for (const auto &I : S.getBodySamples()) {
+    const SampleRecord &Sample = I.second;
+    for (const auto &J : Sample.getCallTargets())
+      addName(J.first());
+  }
+
+  // Recursively add all the names for inlined callsites.
+  for (const auto &J : S.getCallsiteSamples()) {
+    CallsiteLocation Loc = J.first;
+    const FunctionSamples &CalleeSamples = J.second;
+    addName(Loc.CalleeName);
+    addNames(CalleeSamples);
+  }
+}
 
-  // Write the file header.
+std::error_code SampleProfileWriterBinary::writeHeader(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  // Write file magic identifier.
   encodeULEB128(SPMagic(), OS);
   encodeULEB128(SPVersion(), OS);
+
+  // Generate the name table for all the functions referenced in the profile.
+  for (const auto &I : ProfileMap) {
+    addName(I.first());
+    addNames(I.second);
+  }
+
+  // Write out the name table.
+  encodeULEB128(NameTable.size(), OS);
+  for (auto N : NameTable) {
+    OS << N.first;
+    encodeULEB128(0, OS);
+  }
+
+  return sampleprof_error::success;
 }
 
 /// \brief Write samples to a binary file.
 ///
 /// \returns true if the samples were written successfully, false otherwise.
-bool SampleProfileWriterBinary::write(StringRef FName,
-                                      const FunctionSamples &S) {
-  OS << FName;
-  encodeULEB128(0, OS);
+std::error_code SampleProfileWriterBinary::write(StringRef FName,
+                                                 const FunctionSamples &S) {
+  if (std::error_code EC = writeNameIdx(FName))
+    return EC;
+
   encodeULEB128(S.getTotalSamples(), OS);
   encodeULEB128(S.getHeadSamples(), OS);
   encodeULEB128(S.getBodySamples().size(), OS);
@@ -99,10 +143,10 @@ bool SampleProfileWriterBinary::write(StringRef FName,
     encodeULEB128(Sample.getSamples(), OS);
     encodeULEB128(Sample.getCallTargets().size(), OS);
     for (const auto &J : Sample.getCallTargets()) {
-      std::string Callee = J.first();
+      StringRef Callee = J.first();
       unsigned CalleeSamples = J.second;
-      OS << Callee;
-      encodeULEB128(0, OS);
+      if (std::error_code EC = writeNameIdx(Callee))
+        return EC;
       encodeULEB128(CalleeSamples, OS);
     }
   }
@@ -114,10 +158,11 @@ bool SampleProfileWriterBinary::write(StringRef FName,
     const FunctionSamples &CalleeSamples = J.second;
     encodeULEB128(Loc.LineOffset, OS);
     encodeULEB128(Loc.Discriminator, OS);
-    write(Loc.CalleeName, CalleeSamples);
+    if (std::error_code EC = write(Loc.CalleeName, CalleeSamples))
+      return EC;
   }
 
-  return true;
+  return sampleprof_error::success;
 }
 
 /// \brief Create a sample profile writer based on the specified format.
@@ -138,6 +183,8 @@ SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) {
     Writer.reset(new SampleProfileWriterBinary(Filename, EC));
   else if (Format == SPF_Text)
     Writer.reset(new SampleProfileWriterText(Filename, EC));
+  else if (Format == SPF_GCC)
+    EC = sampleprof_error::unsupported_writing_format;
   else
     EC = sampleprof_error::unrecognized_format;