Address comments.

Author: kazutakahirata

llvm/include/llvm/ProfileData/InstrProfReader.h

@@ -717,8 +717,8 @@ class IndexedMemProfReader {
   DenseMap<uint64_t, SmallVector<memprof::CallEdgeTy, 0>>
   getMemProfCallerCalleePairs() const;
 
-  // Return a vector of all GUIDs that we have corresponding MemProfRecords for.
-  SmallVector<uint64_t, 0> getMemProfRecordKeys() const;
+  // Return the entire MemProf profile.
+  memprof::AllMemProfData getAllMemProfData() const;
 };
 
 /// Reader for the indexed binary instrprof format.
@@ -826,8 +826,8 @@ class IndexedInstrProfReader : public InstrProfReader {
     return MemProfReader.getMemProfCallerCalleePairs();
   }
 
-  SmallVector<uint64_t, 0> getMemProfRecordKeys() {
-    return MemProfReader.getMemProfRecordKeys();
+  memprof::AllMemProfData getAllMemProfData() const {
+    return MemProfReader.getAllMemProfData();
   }
 
   /// Fill Counts with the profile data for the given function name.

llvm/include/llvm/ProfileData/MemProfReader.h

@@ -214,16 +214,14 @@ class YAMLMemProfReader final : public MemProfReader {
 public:
   YAMLMemProfReader() = default;
 
-  // Return true if the \p DataBuffer starts with magic bytes indicating it is
-  // a raw binary memprof profile.
+  // Return true if the \p DataBuffer starts "---" indicating it is a YAML file.
   static bool hasFormat(const MemoryBuffer &DataBuffer);
   // Return true if the file at \p Path starts with magic bytes indicating it is
   // a raw binary memprof profile.
   static bool hasFormat(const StringRef Path);
 
-  // Create a RawMemProfReader after sanity checking the contents of the file at
-  // \p Path or the \p Buffer. The binary from which the profile has been
-  // collected is specified via a path in \p ProfiledBinary.
+  // Create a YAMLMemProfReader after sanity checking the contents of the file at
+  // \p Path or the \p Buffer.
   static Expected<std::unique_ptr<YAMLMemProfReader>> create(const Twine &Path);
   static Expected<std::unique_ptr<YAMLMemProfReader>>
   create(std::unique_ptr<MemoryBuffer> Buffer);

llvm/lib/ProfileData/InstrProfReader.cpp

@@ -1664,12 +1664,19 @@ IndexedMemProfReader::getMemProfCallerCalleePairs() const {
   return Pairs;
 }
 
-SmallVector<uint64_t, 0> IndexedMemProfReader::getMemProfRecordKeys() const {
-  SmallVector<uint64_t, 0> Keys;
-  Keys.reserve(MemProfRecordTable->getNumEntries());
-  for (uint64_t Key : MemProfRecordTable->keys())
-    Keys.push_back(Key);
-  return Keys;
+memprof::AllMemProfData IndexedMemProfReader::getAllMemProfData() const {
+  memprof::AllMemProfData AllMemProfData;
+  AllMemProfData.HeapProfileRecords.reserve(MemProfRecordTable->getNumEntries());
+  for (uint64_t Key : MemProfRecordTable->keys()) {
+    auto Record = getMemProfRecord(Key);
+    if (Record.takeError())
+      continue;
+    memprof::GUIDMemProfRecordPair Pair;
+    Pair.GUID = Key;
+    Pair.Record = std::move(*Record);
+    AllMemProfData.HeapProfileRecords.push_back(std::move(Pair));
+  }
+  return AllMemProfData;
 }
 
 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,

llvm/tools/llvm-profdata/llvm-profdata.cpp

@@ -726,17 +726,16 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
   using ::llvm::memprof::YAMLMemProfReader;
   if (YAMLMemProfReader::hasFormat(Input.Filename)) {
     auto ReaderOrErr = YAMLMemProfReader::create(Input.Filename);
-    if (!ReaderOrErr) {
+    if (!ReaderOrErr)
       exitWithError(ReaderOrErr.takeError(), Input.Filename);
-    }
     std::unique_ptr<YAMLMemProfReader> Reader = std::move(ReaderOrErr.get());
     // Check if the profile types can be merged, e.g. clang frontend profiles
     // should not be merged with memprof profiles.
     if (Error E = WC->Writer.mergeProfileKind(Reader->getProfileKind())) {
       consumeError(std::move(E));
       WC->Errors.emplace_back(
           make_error<StringError>(
-              "Cannot merge MemProf profile with Clang generated profile.",
+              "Cannot merge MemProf profile with incompatible profile.",
               std::error_code()),
           Filename);
       return;
@@ -3303,19 +3302,7 @@ static int showMemProfProfile(ShowFormat SFormat, raw_fd_ostream &OS) {
     exitWithError(std::move(E), Filename);
 
   auto Reader = std::move(ReaderOrErr.get());
-
-  // Build pairs of GUID and MemProfRecord.
-  memprof::AllMemProfData Data;
-  for (const uint64_t Key : Reader->getMemProfRecordKeys()) {
-    auto Record = Reader->getMemProfRecord(Key);
-    if (Record.takeError())
-      continue;
-    memprof::GUIDMemProfRecordPair Pair;
-    Pair.GUID = Key;
-    Pair.Record = std::move(*Record);
-    Data.HeapProfileRecords.push_back(std::move(Pair));
-  }
-
+  memprof::AllMemProfData Data = Reader->getAllMemProfData();
   yaml::Output Yout(OS);
   Yout << Data;