1 files changed, 149 insertions, 35 deletions

diff --git a/contrib/llvm/lib/ProfileData/SampleProfWriter.cpp b/contrib/llvm/lib/ProfileData/SampleProfWriter.cpp
index c95267a..51feee5 100644
--- a/contrib/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/contrib/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -30,16 +30,27 @@ using namespace llvm::sampleprof;
 using namespace llvm;
 
 /// \brief Write samples to a text file.
-bool SampleProfileWriterText::write(StringRef FName, const FunctionSamples &S) {
-  if (S.empty())
-    return true;
-
-  OS << FName << ":" << S.getTotalSamples() << ":" << S.getHeadSamples()
-     << "\n";
-
-  for (const auto &I : S.getBodySamples()) {
-    LineLocation Loc = I.first;
-    const SampleRecord &Sample = I.second;
+///
+/// Note: it may be tempting to implement this in terms of
+/// FunctionSamples::print().  Please don't.  The dump functionality is intended
+/// for debugging and has no specified form.
+///
+/// The format used here is more structured and deliberate because
+/// it needs to be parsed by the SampleProfileReaderText class.
+std::error_code SampleProfileWriterText::write(StringRef FName,
+                                               const FunctionSamples &S) {
+  auto &OS = *OutputStream;
+
+  OS << FName << ":" << S.getTotalSamples();
+  if (Indent == 0)
+    OS << ":" << S.getHeadSamples();
+  OS << "\n";
+
+  SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
+  for (const auto &I : SortedSamples.get()) {
+    LineLocation Loc = I->first;
+    const SampleRecord &Sample = I->second;
+    OS.indent(Indent + 1);
     if (Loc.Discriminator == 0)
       OS << Loc.LineOffset << ": ";
     else
@@ -52,32 +63,89 @@ bool SampleProfileWriterText::write(StringRef FName, const FunctionSamples &S) {
     OS << "\n";
   }
 
-  return true;
+  SampleSorter<CallsiteLocation, FunctionSamples> SortedCallsiteSamples(
+      S.getCallsiteSamples());
+  Indent += 1;
+  for (const auto &I : SortedCallsiteSamples.get()) {
+    CallsiteLocation Loc = I->first;
+    const FunctionSamples &CalleeSamples = I->second;
+    OS.indent(Indent);
+    if (Loc.Discriminator == 0)
+      OS << Loc.LineOffset << ": ";
+    else
+      OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
+    if (std::error_code EC = write(Loc.CalleeName, CalleeSamples))
+      return EC;
+  }
+  Indent -= 1;
+
+  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, *OutputStream);
+  return sampleprof_error::success;
+}
 
-  // Write the file header.
+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);
+  }
+}
+
+std::error_code SampleProfileWriterBinary::writeHeader(
+    const StringMap<FunctionSamples> &ProfileMap) {
+  auto &OS = *OutputStream;
+
+  // 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) {
-  if (S.empty())
-    return true;
+std::error_code SampleProfileWriterBinary::writeBody(StringRef FName,
+                                                     const FunctionSamples &S) {
+  auto &OS = *OutputStream;
+
+  if (std::error_code EC = writeNameIdx(FName))
+    return EC;
 
-  OS << FName;
-  encodeULEB128(0, OS);
   encodeULEB128(S.getTotalSamples(), OS);
-  encodeULEB128(S.getHeadSamples(), OS);
+
+  // Emit all the body samples.
   encodeULEB128(S.getBodySamples().size(), OS);
   for (const auto &I : S.getBodySamples()) {
     LineLocation Loc = I.first;
@@ -87,18 +155,38 @@ 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();
-      unsigned CalleeSamples = J.second;
-      OS << Callee;
-      encodeULEB128(0, OS);
+      StringRef Callee = J.first();
+      uint64_t CalleeSamples = J.second;
+      if (std::error_code EC = writeNameIdx(Callee))
+        return EC;
       encodeULEB128(CalleeSamples, OS);
     }
   }
 
-  return true;
+  // Recursively emit all the callsite samples.
+  encodeULEB128(S.getCallsiteSamples().size(), OS);
+  for (const auto &J : S.getCallsiteSamples()) {
+    CallsiteLocation Loc = J.first;
+    const FunctionSamples &CalleeSamples = J.second;
+    encodeULEB128(Loc.LineOffset, OS);
+    encodeULEB128(Loc.Discriminator, OS);
+    if (std::error_code EC = writeBody(Loc.CalleeName, CalleeSamples))
+      return EC;
+  }
+
+  return sampleprof_error::success;
 }
 
-/// \brief Create a sample profile writer based on the specified format.
+/// \brief Write samples of a top-level function to a binary file.
+///
+/// \returns true if the samples were written successfully, false otherwise.
+std::error_code SampleProfileWriterBinary::write(StringRef FName,
+                                                 const FunctionSamples &S) {
+  encodeULEB128(S.getHeadSamples(), *OutputStream);
+  return writeBody(FName, S);
+}
+
+/// \brief Create a sample profile file writer based on the specified format.
 ///
 /// \param Filename The file to create.
 ///
@@ -110,12 +198,38 @@ bool SampleProfileWriterBinary::write(StringRef FName,
 ErrorOr<std::unique_ptr<SampleProfileWriter>>
 SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) {
   std::error_code EC;
+  std::unique_ptr<raw_ostream> OS;
+  if (Format == SPF_Binary)
+    OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::F_None));
+  else
+    OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::F_Text));
+  if (EC)
+    return EC;
+
+  return create(OS, Format);
+}
+
+/// \brief Create a sample profile stream writer based on the specified format.
+///
+/// \param OS The output stream to store the profile data to.
+///
+/// \param Writer The writer to instantiate according to the specified format.
+///
+/// \param Format Encoding format for the profile file.
+///
+/// \returns an error code indicating the status of the created writer.
+ErrorOr<std::unique_ptr<SampleProfileWriter>>
+SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
+                            SampleProfileFormat Format) {
+  std::error_code EC;
   std::unique_ptr<SampleProfileWriter> Writer;
 
   if (Format == SPF_Binary)
-    Writer.reset(new SampleProfileWriterBinary(Filename, EC));
+    Writer.reset(new SampleProfileWriterBinary(OS));
   else if (Format == SPF_Text)
-    Writer.reset(new SampleProfileWriterText(Filename, EC));
+    Writer.reset(new SampleProfileWriterText(OS));
+  else if (Format == SPF_GCC)
+    EC = sampleprof_error::unsupported_writing_format;
   else
     EC = sampleprof_error::unrecognized_format;