2 files changed, 40 insertions, 36 deletions

diff --git a/llvm/include/llvm/ProfileData/InstrProfWriter.h b/llvm/include/llvm/ProfileData/InstrProfWriter.h
index 1714f3b..97f6a95 100644
--- a/llvm/include/llvm/ProfileData/InstrProfWriter.h
+++ b/llvm/include/llvm/ProfileData/InstrProfWriter.h
@@ -212,6 +212,12 @@ private:
   void addTemporalProfileTrace(TemporalProfTraceTy Trace);
 
   Error writeImpl(ProfOStream &OS);
+
+  // Writes known header fields and reserves space for fields whose value are
+  // known only after payloads are written. Returns the start byte offset for
+  // back patching.
+  uint64_t writeHeader(const IndexedInstrProf::Header &header,
+                       const bool WritePrevVersion, ProfOStream &OS);
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index b5b1355..101992c 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -639,6 +639,27 @@ static Error writeMemProf(ProfOStream &OS,
               memprof::MaximumSupportedVersion));
 }
 
+uint64_t InstrProfWriter::writeHeader(const IndexedInstrProf::Header &Header,
+                                      const bool WritePrevVersion,
+                                      ProfOStream &OS) {
+  // Only write out the first four fields.
+  for (int I = 0; I < 4; I++)
+    OS.write(reinterpret_cast<const uint64_t *>(&Header)[I]);
+
+  // Remember the offset of the remaining fields to allow back patching later.
+  auto BackPatchStartOffset = OS.tell();
+
+  // Reserve the space for back patching later.
+  OS.write(0); // HashOffset
+  OS.write(0); // MemProfOffset
+  OS.write(0); // BinaryIdOffset
+  OS.write(0); // TemporalProfTracesOffset
+  if (!WritePrevVersion)
+    OS.write(0); // VTableNamesOffset
+
+  return BackPatchStartOffset;
+}
+
 Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   using namespace IndexedInstrProf;
   using namespace support;
@@ -651,7 +672,7 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   InfoObj->CSSummaryBuilder = &CSISB;
 
   // Populate the hash table generator.
-  SmallVector<std::pair<StringRef, const ProfilingData *>, 0> OrderedData;
+  SmallVector<std::pair<StringRef, const ProfilingData *>> OrderedData;
   for (const auto &I : FunctionData)
     if (shouldEncodeData(I.getValue()))
       OrderedData.emplace_back((I.getKey()), &I.getValue());
@@ -693,35 +714,8 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   Header.TemporalProfTracesOffset = 0;
   Header.VTableNamesOffset = 0;
 
-  // Only write out the first four fields. We need to remember the offset of the
-  // remaining fields to allow back patching later.
-  for (int I = 0; I < 4; I++)
-    OS.write(reinterpret_cast<uint64_t *>(&Header)[I]);
-
-  // Save the location of Header.HashOffset field in \c OS.
-  uint64_t HashTableStartFieldOffset = OS.tell();
-  // Reserve the space for HashOffset field.
-  OS.write(0);
-
-  // Save the location of MemProf profile data. This is stored in two parts as
-  // the schema and as a separate on-disk chained hashtable.
-  uint64_t MemProfSectionOffset = OS.tell();
-  // Reserve space for the MemProf table field to be patched later if this
-  // profile contains memory profile information.
-  OS.write(0);
-
-  // Save the location of binary ids section.
-  uint64_t BinaryIdSectionOffset = OS.tell();
-  // Reserve space for the BinaryIdOffset field to be patched later if this
-  // profile contains binary ids.
-  OS.write(0);
-
-  uint64_t TemporalProfTracesOffset = OS.tell();
-  OS.write(0);
-
-  uint64_t VTableNamesOffset = OS.tell();
-  if (!WritePrevVersion)
-    OS.write(0);
+  const uint64_t BackPatchStartOffset =
+      writeHeader(Header, WritePrevVersion, OS);
 
   // Reserve space to write profile summary data.
   uint32_t NumEntries = ProfileSummaryBuilder::DefaultCutoffs.size();
@@ -850,16 +844,20 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   }
   InfoObj->CSSummaryBuilder = nullptr;
 
+  const size_t MemProfOffset = BackPatchStartOffset + sizeof(uint64_t);
+  const size_t BinaryIdOffset = MemProfOffset + sizeof(uint64_t);
+  const size_t TemporalProfTracesOffset = BinaryIdOffset + sizeof(uint64_t);
+  const size_t VTableNamesOffset = TemporalProfTracesOffset + sizeof(uint64_t);
   if (!WritePrevVersion) {
     // Now do the final patch:
     PatchItem PatchItems[] = {
         // Patch the Header.HashOffset field.
-        {HashTableStartFieldOffset, &HashTableStart, 1},
+        {BackPatchStartOffset, &HashTableStart, 1},
         // Patch the Header.MemProfOffset (=0 for profiles without MemProf
         // data).
-        {MemProfSectionOffset, &MemProfSectionStart, 1},
+        {MemProfOffset, &MemProfSectionStart, 1},
         // Patch the Header.BinaryIdSectionOffset.
-        {BinaryIdSectionOffset, &BinaryIdSectionStart, 1},
+        {BinaryIdOffset, &BinaryIdSectionStart, 1},
         // Patch the Header.TemporalProfTracesOffset (=0 for profiles without
         // traces).
         {TemporalProfTracesOffset, &TemporalProfTracesSectionStart, 1},
@@ -875,12 +873,12 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
     // Now do the final patch:
     PatchItem PatchItems[] = {
         // Patch the Header.HashOffset field.
-        {HashTableStartFieldOffset, &HashTableStart, 1},
+        {BackPatchStartOffset, &HashTableStart, 1},
         // Patch the Header.MemProfOffset (=0 for profiles without MemProf
         // data).
-        {MemProfSectionOffset, &MemProfSectionStart, 1},
+        {MemProfOffset, &MemProfSectionStart, 1},
         // Patch the Header.BinaryIdSectionOffset.
-        {BinaryIdSectionOffset, &BinaryIdSectionStart, 1},
+        {BinaryIdOffset, &BinaryIdSectionStart, 1},
         // Patch the Header.TemporalProfTracesOffset (=0 for profiles without
         // traces).
         {TemporalProfTracesOffset, &TemporalProfTracesSectionStart, 1},