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

[oota-llvm.git] / lib / ProfileData / SampleProfWriter.cpp

diff --git a/lib/ProfileData/SampleProfWriter.cpp b/lib/ProfileData/SampleProfWriter.cpp
index c9f8923346863b8190ad246d82507eabe26d2ad1..51feee5ad7d1d22dd84b678e7a6bd47057a1cbdf 100644 (file)
--- a/lib/ProfileData/SampleProfWriter.cpp
+++ b/lib/ProfileData/SampleProfWriter.cpp
@@ -39,6 +39,8 @@ using namespace llvm;
 /// 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();
@@ -84,7 +86,7 @@ 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;
 }
 
@@ -112,6 +114,8 @@ void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
 
 std::error_code SampleProfileWriterBinary::writeHeader(
     const StringMap<FunctionSamples> &ProfileMap) {
+  auto &OS = *OutputStream;
+
   // Write file magic identifier.
   encodeULEB128(SPMagic(), OS);
   encodeULEB128(SPVersion(), OS);
@@ -134,6 +138,8 @@ std::error_code SampleProfileWriterBinary::writeHeader(
 
 std::error_code SampleProfileWriterBinary::writeBody(StringRef FName,
                                                      const FunctionSamples &S) {
+  auto &OS = *OutputStream;
+
   if (std::error_code EC = writeNameIdx(FName))
     return EC;
 
@@ -176,11 +182,11 @@ std::error_code SampleProfileWriterBinary::writeBody(StringRef FName,
 /// \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.
 ///
@@ -192,12 +198,36 @@ std::error_code 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