[memprof] Use IndexedMemProfRecord in MemProfReader (NFC) (#117613)

IndexedMemProfRecord contains a complete package of the MemProf
profile, including frames, call stacks, and records.  This patch
replaces the three member variables of MemProfReader with
IndexedMemProfRecord.

This transition significantly simplies both the constructor and the
final "take" method:

  MemProfReader(IndexedMemProfData MemProfData)
      : MemProfData(std::move(MemProfData)) {}

IndexedMemProfData takeMemProfData() { return std::move(MemProfData); }

Author: kazutakahirata

llvm/include/llvm/ProfileData/MemProfReader.h

@@ -42,43 +42,29 @@ class MemProfReader {
   using Iterator = InstrProfIterator<GuidMemProfRecordPair, MemProfReader>;
   Iterator end() { return Iterator(); }
   Iterator begin() {
-    Iter = FunctionProfileData.begin();
+    Iter = MemProfData.Records.begin();
     return Iterator(this);
   }
 
-  // Take the complete profile data.
-  IndexedMemProfData takeMemProfData() {
-    // TODO: Once we replace the three member variables, namely IdToFrame,
-    // CSIdToCallStack, and FunctionProfileData, with MemProfData, replace the
-    // following code with just "return std::move(MemProfData);".
-    IndexedMemProfData MemProfData;
-    // Copy key-value pairs because IdToFrame uses DenseMap, whereas
-    // IndexedMemProfData::Frames uses MapVector.
-    for (const auto &[FrameId, F] : IdToFrame)
-      MemProfData.Frames.try_emplace(FrameId, F);
-    // Copy key-value pairs because CSIdToCallStack uses DenseMap, whereas
-    // IndexedMemProfData::CallStacks uses MapVector.
-    for (const auto &[CSId, CS] : CSIdToCallStack)
-      MemProfData.CallStacks.try_emplace(CSId, CS);
-    MemProfData.Records = FunctionProfileData;
-    return MemProfData;
-  }
+  // Take the complete profile data.  Once this function is invoked,
+  // MemProfReader no longer owns the MemProf profile.
+  IndexedMemProfData takeMemProfData() { return std::move(MemProfData); }
 
   virtual Error
   readNextRecord(GuidMemProfRecordPair &GuidRecord,
                  std::function<const Frame(const FrameId)> Callback = nullptr) {
-    if (FunctionProfileData.empty())
+    if (MemProfData.Records.empty())
       return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
 
-    if (Iter == FunctionProfileData.end())
+    if (Iter == MemProfData.Records.end())
       return make_error<InstrProfError>(instrprof_error::eof);
 
     if (Callback == nullptr)
       Callback =
           std::bind(&MemProfReader::idToFrame, this, std::placeholders::_1);
 
-    CallStackIdConverter<decltype(CSIdToCallStack)> CSIdConv(CSIdToCallStack,
-                                                             Callback);
+    CallStackIdConverter<decltype(MemProfData.CallStacks)> CSIdConv(
+        MemProfData.CallStacks, Callback);
 
     const IndexedMemProfRecord &IndexedRecord = Iter->second;
     GuidRecord = {
@@ -97,28 +83,18 @@ class MemProfReader {
   virtual ~MemProfReader() = default;
 
   // Initialize the MemProfReader with the given MemProf profile.
-  MemProfReader(IndexedMemProfData MemProfData) {
-    for (const auto &[FrameId, F] : MemProfData.Frames)
-      IdToFrame.try_emplace(FrameId, F);
-    for (const auto &[CSId, CS] : MemProfData.CallStacks)
-      CSIdToCallStack.try_emplace(CSId, CS);
-    FunctionProfileData = std::move(MemProfData.Records);
-  }
+  MemProfReader(IndexedMemProfData &&MemProfData)
+      : MemProfData(std::move(MemProfData)) {}
 
 protected:
   // A helper method to extract the frame from the IdToFrame map.
   const Frame &idToFrame(const FrameId Id) const {
-    auto It = IdToFrame.find(Id);
-    assert(It != IdToFrame.end() && "Id not found in map.");
-    return It->getSecond();
+    auto It = MemProfData.Frames.find(Id);
+    assert(It != MemProfData.Frames.end() && "Id not found in map.");
+    return It->second;
   }
-  // A mapping from FrameId (a hash of the contents) to the frame.
-  llvm::DenseMap<FrameId, Frame> IdToFrame;
-  // A mapping from CallStackId to the call stack.
-  llvm::DenseMap<CallStackId, llvm::SmallVector<FrameId>> CSIdToCallStack;
-  // A mapping from function GUID, hash of the canonical function symbol to the
-  // memprof profile data for that function, i.e allocation and callsite info.
-  llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> FunctionProfileData;
+  // A complete pacakge of the MemProf profile.
+  IndexedMemProfData MemProfData;
   // An iterator to the internal function profile data structure.
   llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord>::iterator Iter;
 };

llvm/lib/ProfileData/MemProfReader.cpp

@@ -304,7 +304,7 @@ bool RawMemProfReader::hasFormat(const MemoryBuffer &Buffer) {
 
 void RawMemProfReader::printYAML(raw_ostream &OS) {
   uint64_t NumAllocFunctions = 0, NumMibInfo = 0;
-  for (const auto &KV : FunctionProfileData) {
+  for (const auto &KV : MemProfData.Records) {
     const size_t NumAllocSites = KV.second.AllocSites.size();
     if (NumAllocSites > 0) {
       NumAllocFunctions++;
@@ -500,13 +500,13 @@ Error RawMemProfReader::mapRawProfileToRecords() {
     }
 
     CallStackId CSId = hashCallStack(Callstack);
-    CSIdToCallStack.insert({CSId, Callstack});
+    MemProfData.CallStacks.insert({CSId, Callstack});
 
     // We attach the memprof record to each function bottom-up including the
     // first non-inline frame.
     for (size_t I = 0; /*Break out using the condition below*/; I++) {
       const Frame &F = idToFrame(Callstack[I]);
-      IndexedMemProfRecord &Record = FunctionProfileData[F.Function];
+      IndexedMemProfRecord &Record = MemProfData.Records[F.Function];
       Record.AllocSites.emplace_back(Callstack, CSId, MIB);
 
       if (!F.IsInlineFrame)
@@ -518,10 +518,10 @@ Error RawMemProfReader::mapRawProfileToRecords() {
   for (const auto &[Id, Locs] : PerFunctionCallSites) {
     // Some functions may have only callsite data and no allocation data. Here
     // we insert a new entry for callsite data if we need to.
-    IndexedMemProfRecord &Record = FunctionProfileData[Id];
+    IndexedMemProfRecord &Record = MemProfData.Records[Id];
     for (LocationPtr Loc : Locs) {
       CallStackId CSId = hashCallStack(*Loc);
-      CSIdToCallStack.insert({CSId, *Loc});
+      MemProfData.CallStacks.insert({CSId, *Loc});
       Record.CallSites.push_back(*Loc);
       Record.CallSiteIds.push_back(CSId);
     }
@@ -585,7 +585,7 @@ Error RawMemProfReader::symbolizeAndFilterStackFrames(
         }
 
         const FrameId Hash = F.hash();
-        IdToFrame.insert({Hash, F});
+        MemProfData.Frames.insert({Hash, F});
         SymbolizedFrame[VAddr].push_back(Hash);
       }
     }

llvm/unittests/ProfileData/MemProfTest.cpp

@@ -440,7 +440,7 @@ TEST(MemProf, BaseMemProfReader) {
   FakeRecord.AllocSites.emplace_back(/*CSId=*/CSId, /*MB=*/Block);
   MemProfData.Records.insert({F1.hash(), FakeRecord});
 
-  MemProfReader Reader(MemProfData);
+  MemProfReader Reader(std::move(MemProfData));
 
   llvm::SmallVector<MemProfRecord, 1> Records;
   for (const auto &KeyRecordPair : Reader) {
@@ -478,7 +478,7 @@ TEST(MemProf, BaseMemProfReaderWithCSIdMap) {
       /*MB=*/Block);
   MemProfData.Records.insert({F1.hash(), FakeRecord});
 
-  MemProfReader Reader(MemProfData);
+  MemProfReader Reader(std::move(MemProfData));
 
   llvm::SmallVector<MemProfRecord, 1> Records;
   for (const auto &KeyRecordPair : Reader) {