1 files changed, 7 insertions, 2 deletions
diff --git a/include/llvm/ProfileData/InstrProfWriter.h b/include/llvm/ProfileData/InstrProfWriter.h
index e76f668..ce0bb52 100644
--- a/include/llvm/ProfileData/InstrProfWriter.h
+++ b/include/llvm/ProfileData/InstrProfWriter.h
@@ -20,8 +20,8 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
-
 #include <vector>
 
 namespace llvm {
@@ -42,8 +42,13 @@ public:
   std::error_code addFunctionCounts(StringRef FunctionName,
                                     uint64_t FunctionHash,
                                     ArrayRef<uint64_t> Counters);
-  /// Ensure that all data is written to disk.
+  /// Write the profile to \c OS
   void write(raw_fd_ostream &OS);
+  /// Write the profile, returning the raw data. For testing.
+  std::unique_ptr<MemoryBuffer> writeBuffer();
+
+private:
+  std::pair<uint64_t, uint64_t> writeImpl(raw_ostream &OS);
 };
 
 } // end namespace llvm