1 files changed, 26 insertions, 4 deletions
diff --git a/include/llvm/ProfileData/InstrProfWriter.h b/include/llvm/ProfileData/InstrProfWriter.h
index e7f53de051c3c..7d292731cccb8 100644
--- a/include/llvm/ProfileData/InstrProfWriter.h
+++ b/include/llvm/ProfileData/InstrProfWriter.h
@@ -24,21 +24,29 @@
 namespace llvm {
 
 /// Writer for instrumentation based profile data.
+class ProfOStream;
+class InstrProfRecordWriterTrait;
+
 class InstrProfWriter {
 public:
   typedef SmallDenseMap<uint64_t, InstrProfRecord, 1> ProfilingData;
+  enum ProfKind { PF_Unknown = 0, PF_FE, PF_IRLevel };
 
 private:
+  bool Sparse;
   StringMap<ProfilingData> FunctionData;
-  uint64_t MaxFunctionCount;
+  ProfKind ProfileKind;
+  // Use raw pointer here for the incomplete type object.
+  InstrProfRecordWriterTrait *InfoObj;
 
 public:
-  InstrProfWriter() : MaxFunctionCount(0) {}
+  InstrProfWriter(bool Sparse = false);
+  ~InstrProfWriter();
 
   /// Add function counts for the given function. If there are already counts
   /// for this function and the hash and number of counts match, each counter is
   /// summed. Optionally scale counts by \p Weight.
-  std::error_code addRecord(InstrProfRecord &&I, uint64_t Weight = 1);
+  Error addRecord(InstrProfRecord &&I, uint64_t Weight = 1);
   /// Write the profile to \c OS
   void write(raw_fd_ostream &OS);
   /// Write the profile in text format to \c OS
@@ -49,11 +57,25 @@ public:
   /// Write the profile, returning the raw data. For testing.
   std::unique_ptr<MemoryBuffer> writeBuffer();
 
+  /// Set the ProfileKind. Report error if mixing FE and IR level profiles.
+  Error setIsIRLevelProfile(bool IsIRLevel) {
+    if (ProfileKind == PF_Unknown) {
+      ProfileKind = IsIRLevel ? PF_IRLevel: PF_FE;
+      return Error::success();
+    }
+    return (IsIRLevel == (ProfileKind == PF_IRLevel))
+               ? Error::success()
+               : make_error<InstrProfError>(
+                     instrprof_error::unsupported_version);
+  }
+
   // Internal interface for testing purpose only.
   void setValueProfDataEndianness(support::endianness Endianness);
+  void setOutputSparse(bool Sparse);
 
 private:
-  std::pair<uint64_t, uint64_t> writeImpl(raw_ostream &OS);
+  bool shouldEncodeData(const ProfilingData &PD);
+  void writeImpl(ProfOStream &OS);
 };
 
 } // end namespace llvm