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

[oota-llvm.git] / include / llvm / ProfileData / SampleProfWriter.h

diff --git a/include/llvm/ProfileData/SampleProfWriter.h b/include/llvm/ProfileData/SampleProfWriter.h
index d1cd506d157cbe857604838147e6f9aa19cff148..029dd2ebacb07b9f1c858a4084ad767292f4e725 100644 (file)
--- a/include/llvm/ProfileData/SampleProfWriter.h
+++ b/include/llvm/ProfileData/SampleProfWriter.h
@@ -29,9 +29,6 @@ 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) {}
   virtual ~SampleProfileWriter() {}
 
   /// Write sample profiles in \p S for function \p FName.
@@ -55,30 +52,40 @@ public:
     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:
+  SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
+      : OutputStream(std::move(OS)) {}
+
   /// \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) {}
-
   std::error_code write(StringRef FName, const FunctionSamples &S) override;
 
 protected:
+  SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriter(OS), Indent(0) {}
+
   std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override {
     return sampleprof_error::success;
@@ -89,17 +96,21 @@ private:
   ///
   /// This is used when printing inlined callees.
   unsigned Indent;
+
+  friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
+  SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
+                              SampleProfileFormat Format);
 };
 
 /// \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() {}
-
   std::error_code write(StringRef F, const FunctionSamples &S) override;
 
 protected:
+  SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriter(OS), NameTable() {}
+
   std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeNameIdx(StringRef FName);
@@ -110,6 +121,10 @@ private:
   void addNames(const FunctionSamples &S);
 
   MapVector<StringRef, uint32_t> NameTable;
+
+  friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
+  SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
+                              SampleProfileFormat Format);
 };
 
 } // End namespace sampleprof