[ProfileData] Add unit test infrastructure for sample profile reader/writer

[slingn]

Summary

Adds support for in-memory round-trip of sample profile data along with basic
round trip unit tests. This will also make it easier to include unit tests for
future changes to sample profiling.

[davidxl]

include/llvm/ProfileData/SampleProfReader.h

270

///

include/llvm/ProfileData/SampleProfWriter.h

32

Is this constructor redundant now?

[silvas]

unittests/ProfileData/SampleProfTest.cpp

31

Can you just hold this by value?

[slingn]

Updated for davidxl and silvas comments.

include/llvm/ProfileData/SampleProfWriter.h

32

Yes - we should be able to just use the stream interface.

unittests/ProfileData/SampleProfTest.cpp

31

Yes...that implies that SampleProfileWriter::create() should take a raw_ostream & rather than unique_ptr<raw_ostream> in that case. And that means that a caller of SampleProfileWriter::create() owns the output stream rather than the writer.

[davidxl]

include/llvm/ProfileData/SampleProfWriter.h

34–36

Should it be protected?

66

why changing this interface? It is not like the stream based create method is used multiple different places.

[slingn]

include/llvm/ProfileData/SampleProfWriter.h

66

OK. Having a reference to the ostream (and caller owning it) would be slightly cleaner but precludes us having a factory method that takes a file name since somebody has to own the ostream.

I'll put that back and switch back to unique_ptr<raw_ostream>.

[slingn]

Updated for davidxl's comments.

Reverted to unique_ptr<raw_ostream> interface and put back SampleProfWriter file path factory function.

[davidxl]

LGTM

include/llvm/ProfileData/SampleProfReader.h

void reportError(int64_t LineNumber, Twine Msg) const {
     Ctx.diagnose(DiagnosticInfoSampleProfile(Buffer->getBufferIdentifier(),
                                              LineNumber, Msg));
   }
 
   /// \brief Create a sample profile reader appropriate to the file format.
   static ErrorOr<std::unique_ptr<SampleProfileReader>>
   create(StringRef Filename, LLVMContext &C);
 
+  // \brief Create a sample profile reader from the supplied memory buffer.
+  static ErrorOr<std::unique_ptr<SampleProfileReader>>
+  create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C);
+
 protected:
   /// \brief Map every function to its associated profile.
   ///
   /// The profile of every function executed at runtime is collected
   /// in the structure FunctionSamples. This maps function objects
   /// to their corresponding profiles.
   StringMap<FunctionSamples> Profiles;
 

include/llvm/ProfileData/SampleProfWriter.h

 
 namespace sampleprof {
 
 enum SampleProfileFormat { SPF_None = 0, SPF_Text, SPF_Binary, SPF_GCC };
 
 /// \brief Sample-based profile writer. Base class.
 class SampleProfileWriter {
 public:
   SampleProfileWriter(StringRef Filename, std::error_code &EC,
                       sys::fs::OpenFlags Flags)
-      : OS(Filename, EC, Flags) {}
+      : OutputStream(new raw_fd_ostream(Filename, EC, Flags)) {}
+  SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
+      : OutputStream(std::move(OS)) {}
   virtual ~SampleProfileWriter() {}
 
   /// Write sample profiles in \p S for function \p FName.
   ///
   /// \returns status code of the file update operation.
   virtual std::error_code write(StringRef FName, const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
   /// \returns status code of the file update operation.
   std::error_code write(const StringMap<FunctionSamples> &ProfileMap) {
     if (std::error_code EC = writeHeader(ProfileMap))
       return EC;
 
     for (const auto &I : ProfileMap) {
       StringRef FName = I.first();
       const FunctionSamples &Profile = I.second;
       if (std::error_code EC = write(FName, Profile))
         return EC;
     }
     return sampleprof_error::success;
   }
 
+  raw_ostream &getOutputStream() { return *OutputStream; }
+
   /// Profile writer factory.
   ///
-  /// Create a new writer based on the value of \p Format.
+  /// Create a new file writer based on the value of \p Format.
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(StringRef Filename, SampleProfileFormat Format);
 
+  /// Create a new stream writer based on the value of \p Format.
+  /// For testing.
+  static ErrorOr<std::unique_ptr<SampleProfileWriter>>
+  create(std::unique_ptr<raw_ostream> &OS, SampleProfileFormat Format);
+
 protected:
   /// \brief Write a file header for the profile file.
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) = 0;
 
   /// \brief Output stream where to emit the profile to.
-  raw_fd_ostream OS;
+  std::unique_ptr<raw_ostream> OutputStream;
 };
 
 /// \brief Sample-based profile writer (text format).
 class SampleProfileWriterText : public SampleProfileWriter {
 public:
   SampleProfileWriterText(StringRef F, std::error_code &EC)
       : SampleProfileWriter(F, EC, sys::fs::F_Text), Indent(0) {}
+  SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriter(OS), Indent(0) {}
 
   std::error_code write(StringRef FName, const FunctionSamples &S) override;
 
 protected:
   std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override {
     return sampleprof_error::success;
   }
 
 private:
   /// Indent level to use when writing.
   ///
   /// This is used when printing inlined callees.
   unsigned Indent;
 };
 
 /// \brief Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
   SampleProfileWriterBinary(StringRef F, std::error_code &EC)
       : SampleProfileWriter(F, EC, sys::fs::F_None), NameTable() {}
+  SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriter(OS), NameTable() {}
 
   std::error_code write(StringRef F, const FunctionSamples &S) override;
 
 protected:
   std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeNameIdx(StringRef FName);
   std::error_code writeBody(StringRef FName, const FunctionSamples &S);

lib/ProfileData/SampleProfReader.cpp

 /// \param C The LLVM context to use to emit diagnostics.
 ///
 /// \returns an error code indicating the status of the created reader.
 ErrorOr<std::unique_ptr<SampleProfileReader>>
 SampleProfileReader::create(StringRef Filename, LLVMContext &C) {
   auto BufferOrError = setupMemoryBuffer(Filename);
   if (std::error_code EC = BufferOrError.getError())
     return EC;
+  return create(BufferOrError.get(), C);
+}
 
-  auto Buffer = std::move(BufferOrError.get());
+/// \brief Create a sample profile reader based on the format of the input data.
+///
+/// \param B The memory buffer to create the reader from (assumes ownership).
+///
+/// \param Reader The reader to instantiate according to \p Filename's format.
+///
+/// \param C The LLVM context to use to emit diagnostics.
+///
+/// \returns an error code indicating the status of the created reader.
+ErrorOr<std::unique_ptr<SampleProfileReader>>
+SampleProfileReader::create(std::unique_ptr<MemoryBuffer> &B, LLVMContext &C) {
   std::unique_ptr<SampleProfileReader> Reader;
-  if (SampleProfileReaderBinary::hasFormat(*Buffer))
-    Reader.reset(new SampleProfileReaderBinary(std::move(Buffer), C));
-  else if (SampleProfileReaderGCC::hasFormat(*Buffer))
-    Reader.reset(new SampleProfileReaderGCC(std::move(Buffer), C));
-  else if (SampleProfileReaderText::hasFormat(*Buffer))
-    Reader.reset(new SampleProfileReaderText(std::move(Buffer), C));
+  if (SampleProfileReaderBinary::hasFormat(*B))
+    Reader.reset(new SampleProfileReaderBinary(std::move(B), C));
+  else if (SampleProfileReaderGCC::hasFormat(*B))
+    Reader.reset(new SampleProfileReaderGCC(std::move(B), C));
+  else if (SampleProfileReaderText::hasFormat(*B))
+    Reader.reset(new SampleProfileReaderText(std::move(B), C));
   else
     return sampleprof_error::unrecognized_format;
 
   if (std::error_code EC = Reader->readHeader())
     return EC;
 
   return std::move(Reader);
 }

lib/ProfileData/SampleProfWriter.cpp

 /// 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;
std::error_code SampleProfileWriterText::write(StringRef FName,
 
   return sampleprof_error::success;
 }
 
 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);
+  encodeULEB128(ret->second, *OutputStream);
   return sampleprof_error::success;
 }
 
 void SampleProfileWriterBinary::addName(StringRef FName) {
   auto NextIdx = NameTable.size();
   NameTable.insert(std::make_pair(FName, NextIdx));
 }
 
for (const auto &J : S.getCallsiteSamples()) {
     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;
 }
 
 std::error_code SampleProfileWriterBinary::writeBody(StringRef FName,
                                                      const FunctionSamples &S) {
+  auto &OS = *OutputStream;
+
   if (std::error_code EC = writeNameIdx(FName))
     return EC;
 
   encodeULEB128(S.getTotalSamples(), OS);
 
   // Emit all the body samples.
   encodeULEB128(S.getBodySamples().size(), OS);
   for (const auto &I : S.getBodySamples()) {
std::error_code SampleProfileWriterBinary::writeBody(StringRef FName,
   return sampleprof_error::success;
 }
 
 /// \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(), OS);
+  encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(FName, S);
 }
 
-/// \brief Create a sample profile writer based on the specified format.
+/// \brief Create a sample profile file writer based on the specified format.
 ///
 /// \param Filename The file to create.
 ///
 /// \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(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;
 
   if (EC)
     return EC;
 
   return std::move(Writer);
 }

unittests/ProfileData/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Core
   ProfileData
   Support
   )
 
 add_llvm_unittest(ProfileDataTests
   CoverageMappingTest.cpp
   InstrProfTest.cpp
+  SampleProfTest.cpp
   )

unittests/ProfileData/SampleProfTest.cpp

+//===- unittest/ProfileData/SampleProfTest.cpp -------------------*- C++
+//-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ProfileData/SampleProfReader.h"
+#include "llvm/ProfileData/SampleProfWriter.h"
+#include "gtest/gtest.h"
+
+#include <cstdarg>
+
+using namespace llvm;
+using namespace sampleprof;
+
+static ::testing::AssertionResult NoError(std::error_code EC) {
+  if (!EC)
+    return ::testing::AssertionSuccess();
+  return ::testing::AssertionFailure() << "error " << EC.value() << ": "
+                                       << EC.message();
+}
+
+namespace {
+
+struct SampleProfTest : ::testing::Test {
+  std::string Data;
+  std::unique_ptr<raw_ostream> OS;
+  std::unique_ptr<SampleProfileWriter> Writer;
+  std::unique_ptr<SampleProfileReader> Reader;
+
+  SampleProfTest()
+      : Data(), OS(new raw_string_ostream(Data)), Writer(), Reader() {}
+
+  void createWriter(SampleProfileFormat Format) {
+    auto WriterOrErr = SampleProfileWriter::create(OS, Format);
+    ASSERT_TRUE(NoError(WriterOrErr.getError()));
+    Writer = std::move(WriterOrErr.get());
+  }
+
+  void readProfile(std::unique_ptr<MemoryBuffer> &Profile) {
+    auto ReaderOrErr = SampleProfileReader::create(Profile, getGlobalContext());
+    ASSERT_TRUE(NoError(ReaderOrErr.getError()));
+    Reader = std::move(ReaderOrErr.get());
+  }
+
+  void testRoundTrip(SampleProfileFormat Format) {
+    createWriter(Format);
+
+    StringRef FooName("_Z3fooi");
+    FunctionSamples FooSamples;
+    FooSamples.addTotalSamples(7711);
+    FooSamples.addHeadSamples(610);
+    FooSamples.addBodySamples(1, 0, 610);
+
+    StringRef BarName("_Z3bari");
+    FunctionSamples BarSamples;
+    BarSamples.addTotalSamples(20301);
+    BarSamples.addHeadSamples(1437);
+    BarSamples.addBodySamples(1, 0, 1437);
+
+    StringMap<FunctionSamples> Profiles;
+    Profiles[FooName] = std::move(FooSamples);
+    Profiles[BarName] = std::move(BarSamples);
+
+    std::error_code EC;
+    EC = Writer->write(Profiles);
+    ASSERT_TRUE(NoError(EC));
+
+    Writer->getOutputStream().flush();
+
+    auto Profile = MemoryBuffer::getMemBufferCopy(Data);
+    readProfile(Profile);
+
+    EC = Reader->read();
+    ASSERT_TRUE(NoError(EC));
+
+    StringMap<FunctionSamples> &ReadProfiles = Reader->getProfiles();
+    ASSERT_EQ(2u, ReadProfiles.size());
+
+    FunctionSamples &ReadFooSamples = ReadProfiles[FooName];
+    ASSERT_EQ(7711u, ReadFooSamples.getTotalSamples());
+    ASSERT_EQ(610u, ReadFooSamples.getHeadSamples());
+
+    FunctionSamples &ReadBarSamples = ReadProfiles[BarName];
+    ASSERT_EQ(20301u, ReadBarSamples.getTotalSamples());
+    ASSERT_EQ(1437u, ReadBarSamples.getHeadSamples());
+  }
+};
+
+TEST_F(SampleProfTest, roundtrip_text_profile) {
+  testRoundTrip(SampleProfileFormat::SPF_Text);
+}
+
+TEST_F(SampleProfTest, roundtrip_binary_profile) {
+  testRoundTrip(SampleProfileFormat::SPF_Binary);
+}
+
+} // end anonymous namespace