1 files changed, 110 insertions, 18 deletions

diff --git a/llvm/lib/ProfileData/PGOCtxProfReader.cpp b/llvm/lib/ProfileData/PGOCtxProfReader.cpp
index 1b42d8c765f2..e1363cfafdfd 100644
--- a/llvm/lib/ProfileData/PGOCtxProfReader.cpp
+++ b/llvm/lib/ProfileData/PGOCtxProfReader.cpp
@@ -16,8 +16,11 @@
 #include "llvm/Bitstream/BitstreamReader.h"
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
-#include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/YAMLTraits.h"
+#include <iterator>
+#include <utility>
 
 using namespace llvm;
 
@@ -32,25 +35,17 @@ using namespace llvm;
   if (auto Err = (EXPR))                                                       \
     return Err;
 
-Expected<PGOContextualProfile &>
-PGOContextualProfile::getOrEmplace(uint32_t Index, GlobalValue::GUID G,
-                                   SmallVectorImpl<uint64_t> &&Counters) {
-  auto [Iter, Inserted] = Callsites[Index].insert(
-      {G, PGOContextualProfile(G, std::move(Counters))});
+Expected<PGOCtxProfContext &>
+PGOCtxProfContext::getOrEmplace(uint32_t Index, GlobalValue::GUID G,
+                                SmallVectorImpl<uint64_t> &&Counters) {
+  auto [Iter, Inserted] =
+      Callsites[Index].insert({G, PGOCtxProfContext(G, std::move(Counters))});
   if (!Inserted)
     return make_error<InstrProfError>(instrprof_error::invalid_prof,
                                       "Duplicate GUID for same callsite.");
   return Iter->second;
 }
 
-void PGOContextualProfile::getContainedGuids(
-    DenseSet<GlobalValue::GUID> &Guids) const {
-  Guids.insert(GUID);
-  for (const auto &[_, Callsite] : Callsites)
-    for (const auto &[_, Callee] : Callsite)
-      Callee.getContainedGuids(Guids);
-}
-
 Expected<BitstreamEntry> PGOCtxProfileReader::advance() {
   return Cursor.advance(BitstreamCursor::AF_DontAutoprocessAbbrevs);
 }
@@ -73,7 +68,7 @@ bool PGOCtxProfileReader::canReadContext() {
          Blk->ID == PGOCtxProfileBlockIDs::ContextNodeBlockID;
 }
 
-Expected<std::pair<std::optional<uint32_t>, PGOContextualProfile>>
+Expected<std::pair<std::optional<uint32_t>, PGOCtxProfContext>>
 PGOCtxProfileReader::readContext(bool ExpectIndex) {
   RET_ON_ERR(Cursor.EnterSubBlock(PGOCtxProfileBlockIDs::ContextNodeBlockID));
 
@@ -124,7 +119,7 @@ PGOCtxProfileReader::readContext(bool ExpectIndex) {
     }
   }
 
-  PGOContextualProfile Ret(*Guid, std::move(*Counters));
+  PGOCtxProfContext Ret(*Guid, std::move(*Counters));
 
   while (canReadContext()) {
     EXPECT_OR_RET(SC, readContext(true));
@@ -139,6 +134,20 @@ PGOCtxProfileReader::readContext(bool ExpectIndex) {
 }
 
 Error PGOCtxProfileReader::readMetadata() {
+  if (Magic.size() < PGOCtxProfileWriter::ContainerMagic.size() ||
+      Magic != PGOCtxProfileWriter::ContainerMagic)
+    return make_error<InstrProfError>(instrprof_error::invalid_prof,
+                                      "Invalid magic");
+
+  BitstreamEntry Entry;
+  RET_ON_ERR(Cursor.advance().moveInto(Entry));
+  if (Entry.Kind != BitstreamEntry::SubBlock ||
+      Entry.ID != bitc::BLOCKINFO_BLOCK_ID)
+    return unsupported("Expected Block ID");
+  // We don't need the blockinfo to read the rest, it's metadata usable for e.g.
+  // llvm-bcanalyzer.
+  RET_ON_ERR(Cursor.SkipBlock());
+
   EXPECT_OR_RET(Blk, advance());
   if (Blk->Kind != BitstreamEntry::SubBlock)
     return unsupported("Expected Version record");
@@ -159,9 +168,9 @@ Error PGOCtxProfileReader::readMetadata() {
   return Error::success();
 }
 
-Expected<std::map<GlobalValue::GUID, PGOContextualProfile>>
+Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>>
 PGOCtxProfileReader::loadContexts() {
-  std::map<GlobalValue::GUID, PGOContextualProfile> Ret;
+  std::map<GlobalValue::GUID, PGOCtxProfContext> Ret;
   RET_ON_ERR(readMetadata());
   while (canReadContext()) {
     EXPECT_OR_RET(E, readContext(false));
@@ -171,3 +180,86 @@ PGOCtxProfileReader::loadContexts() {
   }
   return std::move(Ret);
 }
+
+namespace {
+// We want to pass `const` values PGOCtxProfContext references to the yaml
+// converter, and the regular yaml mapping APIs are designed to handle both
+// serialization and deserialization, which prevents using const for
+// serialization. Using an intermediate datastructure is overkill, both
+// space-wise and design complexity-wise. Instead, we use the lower-level APIs.
+void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx);
+
+void toYaml(yaml::Output &Out,
+            const PGOCtxProfContext::CallTargetMapTy &CallTargets) {
+  Out.beginSequence();
+  size_t Index = 0;
+  void *SaveData = nullptr;
+  for (const auto &[_, Ctx] : CallTargets) {
+    Out.preflightElement(Index++, SaveData);
+    toYaml(Out, Ctx);
+    Out.postflightElement(nullptr);
+  }
+  Out.endSequence();
+}
+
+void toYaml(yaml::Output &Out,
+            const PGOCtxProfContext::CallsiteMapTy &Callsites) {
+  auto AllCS = ::llvm::make_first_range(Callsites);
+  auto MaxIt = ::llvm::max_element(AllCS);
+  assert(MaxIt != AllCS.end() && "We should have a max value because the "
+                                 "callsites collection is not empty.");
+  void *SaveData = nullptr;
+  Out.beginSequence();
+  for (auto I = 0U; I <= *MaxIt; ++I) {
+    Out.preflightElement(I, SaveData);
+    auto It = Callsites.find(I);
+    if (It == Callsites.end()) {
+      // This will produce a `[ ]` sequence, which is what we want here.
+      Out.beginFlowSequence();
+      Out.endFlowSequence();
+    } else {
+      toYaml(Out, It->second);
+    }
+    Out.postflightElement(nullptr);
+  }
+  Out.endSequence();
+}
+
+void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx) {
+  yaml::EmptyContext Empty;
+  Out.beginMapping();
+  void *SaveInfo = nullptr;
+  bool UseDefault = false;
+  {
+    Out.preflightKey("Guid", /*Required=*/true, /*SameAsDefault=*/false,
+                     UseDefault, SaveInfo);
+    auto Guid = Ctx.guid();
+    yaml::yamlize(Out, Guid, true, Empty);
+    Out.postflightKey(nullptr);
+  }
+  {
+    Out.preflightKey("Counters", true, false, UseDefault, SaveInfo);
+    Out.beginFlowSequence();
+    for (size_t I = 0U, E = Ctx.counters().size(); I < E; ++I) {
+      Out.preflightFlowElement(I, SaveInfo);
+      uint64_t V = Ctx.counters()[I];
+      yaml::yamlize(Out, V, true, Empty);
+      Out.postflightFlowElement(SaveInfo);
+    }
+    Out.endFlowSequence();
+    Out.postflightKey(nullptr);
+  }
+  if (!Ctx.callsites().empty()) {
+    Out.preflightKey("Callsites", true, false, UseDefault, SaveInfo);
+    toYaml(Out, Ctx.callsites());
+    Out.postflightKey(nullptr);
+  }
+  Out.endMapping();
+}
+} // namespace
+
+void llvm::convertCtxProfToYaml(
+    raw_ostream &OS, const PGOCtxProfContext::CallTargetMapTy &Profiles) {
+  yaml::Output Out(OS);
+  toYaml(Out, Profiles);
+}
\ No newline at end of file