[PGO] Commonize (more) index profile file and buffer writer.

The file and buffer writer code are mostly shared except for the
stream back-patching. This is because raw_string_ostream does not
support seek like interface. The result is that the data patching
code needs to be pushed to the caller which is not quite readable 
(passing around offset, value etc). This also makes future enhancement
(which needs more patching) more difficult (and can make impl messy).

In this patch, two types of streams needed by the writer are now
unified with same set of interfaces under ProfOStream class. The patch
method is added so that common implementation becomes cleaner. It
also enables future enhancement. Should be NFC.

llvm-svn: 257921

1 files changed, 70 insertions, 33 deletions

diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index 44ad016f1a8..f684dde336a 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -20,6 +20,58 @@
 
 using namespace llvm;
 
+// A struct to define how the data stream should be patched. For Indexed
+// profiling, only uint64_t data type is needed.
+struct PatchItem {
+  uint64_t Pos; // Where to patch.
+  uint64_t *D;  // Pointer to an array of source data.
+  int N;        // Number of elements in \c D array.
+};
+
+namespace llvm {
+// A wrapper class to abstract writer stream with support of bytes
+// back patching.
+struct ProfOStream {
+
+  ProfOStream(llvm::raw_fd_ostream &FD) : IsFDOStream(true), OS(FD), LE(FD) {}
+  ProfOStream(llvm::raw_string_ostream &STR)
+      : IsFDOStream(false), OS(STR), LE(STR) {}
+
+  uint64_t tell() { return OS.tell(); }
+  void write(uint64_t V) { LE.write<uint64_t>(V); }
+  // \c patch can only be called when all data is written and flushed.
+  // For raw_string_ostream, the patch is done on the target string
+  // directly and it won't be reflected in the stream's internal buffer.
+  void patch(PatchItem *P, int NItems) {
+    using namespace support;
+    if (IsFDOStream) {
+      llvm::raw_fd_ostream &FDOStream = static_cast<llvm::raw_fd_ostream &>(OS);
+      for (int K = 0; K < NItems; K++) {
+        FDOStream.seek(P[K].Pos);
+        for (int I = 0; I < P[K].N; I++)
+          write(P[K].D[I]);
+      }
+    } else {
+      llvm::raw_string_ostream &SOStream =
+          static_cast<llvm::raw_string_ostream &>(OS);
+      std::string &Data = SOStream.str(); // with flush
+      for (int K = 0; K < NItems; K++) {
+        for (int I = 0; I < P[K].N; I++) {
+          uint64_t Bytes = endian::byte_swap<uint64_t, little>(P[K].D[I]);
+          Data.replace(P[K].Pos + I * sizeof(uint64_t), sizeof(uint64_t),
+                       (const char *)&Bytes, sizeof(uint64_t));
+        }
+      }
+    }
+  }
+  // If \c OS is an instance of \c raw_fd_ostream, this field will be
+  // true. Otherwise, \c OS will be an raw_string_ostream.
+  bool IsFDOStream;
+  raw_ostream &OS;
+  support::endian::Writer<support::little> LE;
+};
+}
+
 namespace {
 static support::endianness ValueProfDataEndianness = support::little;
 
@@ -127,16 +179,11 @@ std::error_code InstrProfWriter::addRecord(InstrProfRecord &&I,
   return Result;
 }
 
-std::pair<uint64_t, uint64_t> InstrProfWriter::writeImpl(raw_ostream &OS) {
+void InstrProfWriter::writeImpl(ProfOStream &OS) {
   OnDiskChainedHashTableGenerator<InstrProfRecordTrait> Generator;
-
   // Populate the hash table generator.
   for (const auto &I : FunctionData)
     Generator.insert(I.getKey(), &I.getValue());
-
-  using namespace llvm::support;
-  endian::Writer<little> LE(OS);
-
   // Write the header.
   IndexedInstrProf::Header Header;
   Header.Magic = IndexedInstrProf::Magic;
@@ -150,27 +197,34 @@ std::pair<uint64_t, uint64_t> InstrProfWriter::writeImpl(raw_ostream &OS) {
   // to remember the offset of that field to allow back patching
   // later.
   for (int I = 0; I < N - 1; I++)
-    LE.write<uint64_t>(reinterpret_cast<uint64_t *>(&Header)[I]);
+    OS.write(reinterpret_cast<uint64_t *>(&Header)[I]);
 
   // Save a space to write the hash table start location.
   uint64_t HashTableStartLoc = OS.tell();
   // Reserve the space for HashOffset field.
-  LE.write<uint64_t>(0);
+  OS.write(0);
   // Write the hash table.
-  uint64_t HashTableStart = Generator.Emit(OS);
+  uint64_t HashTableStart = Generator.Emit(OS.OS);
 
-  return std::make_pair(HashTableStartLoc, HashTableStart);
+  // Now do the final patch:
+  PatchItem PatchItems[1] = {{HashTableStartLoc, &HashTableStart, 1}};
+  OS.patch(PatchItems, sizeof(PatchItems) / sizeof(*PatchItems));
 }
 
 void InstrProfWriter::write(raw_fd_ostream &OS) {
   // Write the hash table.
-  auto TableStart = writeImpl(OS);
+  ProfOStream POS(OS);
+  writeImpl(POS);
+}
 
-  // Go back and fill in the hash table start.
-  using namespace support;
-  OS.seek(TableStart.first);
-  // Now patch the HashOffset field previously reserved.
-  endian::Writer<little>(OS).write<uint64_t>(TableStart.second);
+std::unique_ptr<MemoryBuffer> InstrProfWriter::writeBuffer() {
+  std::string Data;
+  llvm::raw_string_ostream OS(Data);
+  ProfOStream POS(OS);
+  // Write the hash table.
+  writeImpl(POS);
+  // Return this in an aligned memory buffer.
+  return MemoryBuffer::getMemBufferCopy(Data);
 }
 
 static const char *ValueProfKindStr[] = {
@@ -227,20 +281,3 @@ void InstrProfWriter::writeText(raw_fd_ostream &OS) {
     for (const auto &Func : I.getValue())
       writeRecordInText(Func.second, Symtab, OS);
 }
-
-std::unique_ptr<MemoryBuffer> InstrProfWriter::writeBuffer() {
-  std::string Data;
-  llvm::raw_string_ostream OS(Data);
-  // Write the hash table.
-  auto TableStart = writeImpl(OS);
-  OS.flush();
-
-  // Go back and fill in the hash table start.
-  using namespace support;
-  uint64_t Bytes = endian::byte_swap<uint64_t, little>(TableStart.second);
-  Data.replace(TableStart.first, sizeof(uint64_t), (const char *)&Bytes,
-               sizeof(uint64_t));
-
-  // Return this in an aligned memory buffer.
-  return MemoryBuffer::getMemBufferCopy(Data);
-}