1 files changed, 108 insertions, 10 deletions
diff --git a/include/llvm/ProfileData/SampleProfWriter.h b/include/llvm/ProfileData/SampleProfWriter.h
index 81e6e3ab0b4a..cc951594c9e2 100644
--- a/include/llvm/ProfileData/SampleProfWriter.h
+++ b/include/llvm/ProfileData/SampleProfWriter.h
@@ -36,7 +36,7 @@ public:
   /// Write sample profiles in \p S.
   ///
   /// \returns status code of the file update operation.
-  virtual std::error_code write(const FunctionSamples &S) = 0;
+  virtual std::error_code writeSample(const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
@@ -56,6 +56,8 @@ public:
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(std::unique_ptr<raw_ostream> &OS, SampleProfileFormat Format);
 
+  virtual void setProfileSymbolList(ProfileSymbolList *PSL) {}
+
 protected:
   SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
       : OutputStream(std::move(OS)) {}
@@ -64,6 +66,10 @@ protected:
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) = 0;
 
+  // Write function profiles to the profile file.
+  virtual std::error_code
+  writeFuncProfiles(const StringMap<FunctionSamples> &ProfileMap);
+
   /// Output stream where to emit the profile to.
   std::unique_ptr<raw_ostream> OutputStream;
 
@@ -72,12 +78,15 @@ protected:
 
   /// Compute summary for this profile.
   void computeSummary(const StringMap<FunctionSamples> &ProfileMap);
+
+  /// Profile format.
+  SampleProfileFormat Format;
 };
 
 /// Sample-based profile writer (text format).
 class SampleProfileWriterText : public SampleProfileWriter {
 public:
-  std::error_code write(const FunctionSamples &S) override;
+  std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
   SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
@@ -102,13 +111,14 @@ private:
 /// Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
-  virtual std::error_code write(const FunctionSamples &S) override;
   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS) {}
 
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
+
 protected:
-  virtual std::error_code writeNameTable() = 0;
-  virtual std::error_code writeMagicIdent() = 0;
+  virtual std::error_code writeMagicIdent(SampleProfileFormat Format);
+  virtual std::error_code writeNameTable();
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeSummary();
@@ -118,10 +128,10 @@ protected:
 
   MapVector<StringRef, uint32_t> NameTable;
 
-private:
   void addName(StringRef FName);
   void addNames(const FunctionSamples &S);
 
+private:
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
@@ -129,10 +139,99 @@ private:
 
 class SampleProfileWriterRawBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
+};
+
+class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
+  using SampleProfileWriterBinary::SampleProfileWriterBinary;
+public:
+  virtual std::error_code
+  write(const StringMap<FunctionSamples> &ProfileMap) override;
+
+  void setToCompressAllSections();
+  void setToCompressSection(SecType Type);
 
 protected:
-  virtual std::error_code writeNameTable() override;
-  virtual std::error_code writeMagicIdent() override;
+  uint64_t markSectionStart(SecType Type);
+  std::error_code addNewSection(SecType Sec, uint64_t SectionStart);
+  virtual void initSectionHdrLayout() = 0;
+  virtual std::error_code
+  writeSections(const StringMap<FunctionSamples> &ProfileMap) = 0;
+
+  // Specifiy the order of sections in section header table. Note
+  // the order of sections in the profile may be different that the
+  // order in SectionHdrLayout. sample Reader will follow the order
+  // in SectionHdrLayout to read each section.
+  SmallVector<SecHdrTableEntry, 8> SectionHdrLayout;
+
+private:
+  void allocSecHdrTable();
+  std::error_code writeSecHdrTable();
+  virtual std::error_code
+  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  void addSectionFlags(SecType Type, SecFlags Flags);
+  SecHdrTableEntry &getEntryInLayout(SecType Type);
+  std::error_code compressAndOutput();
+
+  // We will swap the raw_ostream held by LocalBufStream and that
+  // held by OutputStream if we try to add a section which needs
+  // compression. After the swap, all the data written to output
+  // will be temporarily buffered into the underlying raw_string_ostream
+  // originally held by LocalBufStream. After the data writing for the
+  // section is completed, compress the data in the local buffer,
+  // swap the raw_ostream back and write the compressed data to the
+  // real output.
+  std::unique_ptr<raw_ostream> LocalBufStream;
+  // The location where the output stream starts.
+  uint64_t FileStart;
+  // The location in the output stream where the SecHdrTable should be
+  // written to.
+  uint64_t SecHdrTableOffset;
+  // Initial Section Flags setting.
+  std::vector<SecHdrTableEntry> SecHdrTable;
+};
+
+class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
+public:
+  SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriterExtBinaryBase(OS) {
+    initSectionHdrLayout();
+  }
+
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
+  virtual void setProfileSymbolList(ProfileSymbolList *PSL) override {
+    ProfSymList = PSL;
+  };
+
+private:
+  virtual void initSectionHdrLayout() override {
+    // Note that SecFuncOffsetTable section is written after SecLBRProfile
+    // in the profile, but is put before SecLBRProfile in SectionHdrLayout.
+    //
+    // This is because sample reader follows the order of SectionHdrLayout to
+    // read each section, to read function profiles on demand sample reader
+    // need to get the offset of each function profile first.
+    //
+    // SecFuncOffsetTable section is written after SecLBRProfile in the
+    // profile because FuncOffsetTable needs to be populated while section
+    // SecLBRProfile is written.
+    SectionHdrLayout = {{SecProfSummary, 0, 0, 0},
+                        {SecNameTable, 0, 0, 0},
+                        {SecFuncOffsetTable, 0, 0, 0},
+                        {SecLBRProfile, 0, 0, 0},
+                        {SecProfileSymbolList, 0, 0, 0}};
+  };
+  virtual std::error_code
+  writeSections(const StringMap<FunctionSamples> &ProfileMap) override;
+  ProfileSymbolList *ProfSymList = nullptr;
+
+  // Save the start of SecLBRProfile so we can compute the offset to the
+  // start of SecLBRProfile for each Function's Profile and will keep it
+  // in FuncOffsetTable.
+  uint64_t SecLBRProfileStart;
+  // FuncOffsetTable maps function name to its profile offset in SecLBRProfile
+  // section. It is used to load function profile on demand.
+  MapVector<StringRef, uint64_t> FuncOffsetTable;
+  std::error_code writeFuncOffsetTable();
 };
 
 // CompactBinary is a compact format of binary profile which both reduces
@@ -169,7 +268,7 @@ class SampleProfileWriterCompactBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 
 public:
-  virtual std::error_code write(const FunctionSamples &S) override;
+  virtual std::error_code writeSample(const FunctionSamples &S) override;
   virtual std::error_code
   write(const StringMap<FunctionSamples> &ProfileMap) override;
 
@@ -181,7 +280,6 @@ protected:
   /// towards profile start.
   uint64_t TableOffset;
   virtual std::error_code writeNameTable() override;
-  virtual std::error_code writeMagicIdent() override;
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeFuncOffsetTable();