[memprof] Add a MemProfReader base class.

[snehasish]

Summary

Add a MemProfReader base class which can be used directly where
symbolization and processing a raw profile is unnecessary.

[tejohnson]

Should there be a test to cover the new base class? In particular the constructor isn't called and the readNextRecord with a null Callback case.

llvm/include/llvm/ProfileData/RawMemProfReader.h

69

nit: I think the braces can be removed in this case

86

I don't see this constructor being called anywhere in this patch?

147

Should this call the base class constructor?

[snehasish]

Add unit test, address comments, cleanup unused includes in test.

[snehasish]

In D159141#4626637, @tejohnson wrote:

Should there be a test to cover the new base class? In particular the constructor isn't called and the readNextRecord with a null Callback case.

Added a test which uses the constructor and relies on the default null Callback behaviour.

llvm/include/llvm/ProfileData/RawMemProfReader.h

86

Added a test which uses this constructor.

147

No we can't call the base class constructor from here since we need to do additional processing (symbolizeAndFilterStackFrames and mapRawProfileToRecords) before populating the function profile data structure.

[tejohnson]

lgtm

llvm/include/llvm/ProfileData/RawMemProfReader.h

 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/MemProf.h"
 #include "llvm/ProfileData/MemProfData.inc"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 
-#include <cstddef>
+#include <functional>
 
 namespace llvm {
 namespace memprof {
+// A class for memprof profile data populated directly from external
+// sources.
+// TODO: Rename this file to MemProfReader.h to better reflect the contents.
+class MemProfReader {
+public:
+  // The MemProfReader only holds memory profile information.
+  InstrProfKind getProfileKind() const { return InstrProfKind::MemProf; }
+
+  using GuidMemProfRecordPair = std::pair<GlobalValue::GUID, MemProfRecord>;
+  using Iterator = InstrProfIterator<GuidMemProfRecordPair, MemProfReader>;
+  Iterator end() { return Iterator(); }
+  Iterator begin() {
+    Iter = FunctionProfileData.begin();
+    return Iterator(this);
+  }
+
+  // Return a const reference to the internal Id to Frame mappings.
+  const llvm::DenseMap<FrameId, Frame> &getFrameMapping() const {
+    return IdToFrame;
+  }
+
+  // Return a const reference to the internal function profile data.
+  const llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> &
+  getProfileData() const {
+    return FunctionProfileData;
+  }
+
+  virtual Error
+  readNextRecord(GuidMemProfRecordPair &GuidRecord,
+                 std::function<const Frame(const FrameId)> Callback = nullptr) {
+    if (FunctionProfileData.empty())
+      return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
+
+    if (Iter == FunctionProfileData.end())
+      return make_error<InstrProfError>(instrprof_error::eof);
+
+    if (Callback == nullptr) {
+      Callback =
+          std::bind(&MemProfReader::idToFrame, this, std::placeholders::_1);
+    }
+
+    const IndexedMemProfRecord &IndexedRecord = Iter->second;
+    GuidRecord = {Iter->first, MemProfRecord(IndexedRecord, Callback)};
+    Iter++;
+    return Error::success();
+  }
+
+  // Allow default construction for derived classes which can populate the
+  // contents after construction.
+  MemProfReader() = default;
+  virtual ~MemProfReader() = default;
+
+  // Initialize the MemProfReader with the frame mappings and profile contents.
+  MemProfReader(
+      llvm::DenseMap<FrameId, Frame> FrameIdMap,
+      llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> ProfData)
+      : IdToFrame(std::move(FrameIdMap)),
+        FunctionProfileData(std::move(ProfData)) {}
+
+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();
+  }
+  // A mapping from FrameId (a hash of the contents) to the frame.
+  llvm::DenseMap<FrameId, Frame> IdToFrame;
+  // 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;
+  // An iterator to the internal function profile data structure.
+  llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord>::iterator Iter;
+};
 
 // Map from id (recorded from sanitizer stack depot) to virtual addresses for
 // each program counter address in the callstack.
 using CallStackMap = llvm::DenseMap<uint64_t, llvm::SmallVector<uint64_t>>;
 
-class RawMemProfReader {
+// Specializes the MemProfReader class to populate the contents from raw binary
+// memprof profiles from instrumentation based profiling.
+class RawMemProfReader final : public MemProfReader {
 public:
   RawMemProfReader(const RawMemProfReader &) = delete;
   RawMemProfReader &operator=(const RawMemProfReader &) = delete;
+  virtual ~RawMemProfReader() override = default;
 
   // Prints the contents of the profile in YAML format.
   void printYAML(raw_ostream &OS);
 
   // Return true if the \p DataBuffer starts with magic bytes indicating it is
   // a raw binary memprof profile.
   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.
   static Expected<std::unique_ptr<RawMemProfReader>>
   create(const Twine &Path, StringRef ProfiledBinary, bool KeepName = false);
   static Expected<std::unique_ptr<RawMemProfReader>>
   create(std::unique_ptr<MemoryBuffer> Buffer, StringRef ProfiledBinary,
          bool KeepName = false);
 
   // Returns a list of build ids recorded in the segment information.
   static std::vector<std::string> peekBuildIds(MemoryBuffer *DataBuffer);
 
-  using GuidMemProfRecordPair = std::pair<GlobalValue::GUID, MemProfRecord>;
-  using Iterator = InstrProfIterator<GuidMemProfRecordPair, RawMemProfReader>;
-  Iterator end() { return Iterator(); }
-  Iterator begin() {
-    Iter = FunctionProfileData.begin();
-    return Iterator(this);
-  }
-
-  Error readNextRecord(GuidMemProfRecordPair &GuidRecord);
-
-  // The RawMemProfReader only holds memory profile information.
-  InstrProfKind getProfileKind() const { return InstrProfKind::MemProf; }
+  virtual Error
+  readNextRecord(GuidMemProfRecordPair &GuidRecord,
+                 std::function<const Frame(const FrameId)> Callback) override;
 
   // Constructor for unittests only.
   RawMemProfReader(std::unique_ptr<llvm::symbolize::SymbolizableModule> Sym,
                    llvm::SmallVectorImpl<SegmentEntry> &Seg,
                    llvm::MapVector<uint64_t, MemInfoBlock> &Prof,
                    CallStackMap &SM, bool KeepName = false)
       : Symbolizer(std::move(Sym)), SegmentInfo(Seg.begin(), Seg.end()),
         CallstackProfileData(Prof), StackMap(SM), KeepSymbolName(KeepName) {
     // We don't call initialize here since there is no raw profile to read. The
     // test should pass in the raw profile as structured data.
 
     // If there is an error here then the mock symbolizer has not been
     // initialized properly.
     if (Error E = symbolizeAndFilterStackFrames())
       report_fatal_error(std::move(E));
     if (Error E = mapRawProfileToRecords())
       report_fatal_error(std::move(E));
   }
 
-  // Return a const reference to the internal Id to Frame mappings.
-  const llvm::DenseMap<FrameId, Frame> &getFrameMapping() const {
-    return IdToFrame;
-  }
-
-  // Return a const reference to the internal function profile data.
-  const llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> &
-  getProfileData() const {
-    return FunctionProfileData;
-  }
-
 private:
   RawMemProfReader(object::OwningBinary<object::Binary> &&Bin, bool KeepName)
       : Binary(std::move(Bin)), KeepSymbolName(KeepName) {}
   // Initializes the RawMemProfReader with the contents in `DataBuffer`.
   Error initialize(std::unique_ptr<MemoryBuffer> DataBuffer);
   // Read and parse the contents of the `DataBuffer` as a binary format profile.
   Error readRawProfile(std::unique_ptr<MemoryBuffer> DataBuffer);
   // Initialize the segment mapping information for symbolization.
   Error setupForSymbolization();
   // Symbolize and cache all the virtual addresses we encounter in the
   // callstacks from the raw profile. Also prune callstack frames which we can't
   // symbolize or those that belong to the runtime. For profile entries where
   // the entire callstack is pruned, we drop the entry from the profile.
   Error symbolizeAndFilterStackFrames();
   // Construct memprof records for each function and store it in the
   // `FunctionProfileData` map. A function may have allocation profile data or
   // callsite data or both.
   Error mapRawProfileToRecords();
 
-  // 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();
-  }
-
   object::SectionedAddress getModuleOffset(uint64_t VirtualAddress);
 
   // The profiled binary.
   object::OwningBinary<object::Binary> Binary;
   // A symbolizer to translate virtual addresses to code locations.
   std::unique_ptr<llvm::symbolize::SymbolizableModule> Symbolizer;
   // The preferred load address of the executable segment.
   uint64_t PreferredTextSegmentAddress = 0;
   // The base address of the text segment in the process during profiling.
   uint64_t ProfiledTextSegmentStart = 0;
   // The limit address of the text segment in the process during profiling.
   uint64_t ProfiledTextSegmentEnd = 0;
 
   // The memory mapped segment information for all executable segments in the
   // profiled binary (filtered from the raw profile using the build id).
   llvm::SmallVector<SegmentEntry, 2> SegmentInfo;
 
   // A map from callstack id (same as key in CallStackMap below) to the heap
   // information recorded for that allocation context.
   llvm::MapVector<uint64_t, MemInfoBlock> CallstackProfileData;
   CallStackMap StackMap;
 
   // Cached symbolization from PC to Frame.
   llvm::DenseMap<uint64_t, llvm::SmallVector<FrameId>> SymbolizedFrame;
-  llvm::DenseMap<FrameId, Frame> IdToFrame;
-
-  llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> FunctionProfileData;
-  llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord>::iterator Iter;
 
   // Whether to keep the symbol name for each frame after hashing.
   bool KeepSymbolName = false;
   // A mapping of the hash to symbol name, only used if KeepSymbolName is true.
   llvm::DenseMap<uint64_t, std::string> GuidToSymbolName;
 };
 } // namespace memprof
 } // namespace llvm
 
 #endif // LLVM_PROFILEDATA_RAWMEMPROFREADER_H_

llvm/lib/ProfileData/RawMemProfReader.cpp

if (VirtualAddress > ProfiledTextSegmentStart &&
     return object::SectionedAddress{AdjustedAddress};
   }
   // Addresses which do not originate from the profiled text segment in the
   // binary are not adjusted. These will fail symbolization and be filtered out
   // during processing.
   return object::SectionedAddress{VirtualAddress};
 }
 
-Error RawMemProfReader::readNextRecord(GuidMemProfRecordPair &GuidRecord) {
-  if (FunctionProfileData.empty())
-    return make_error<InstrProfError>(instrprof_error::empty_raw_profile);
-
-  if (Iter == FunctionProfileData.end())
-    return make_error<InstrProfError>(instrprof_error::eof);
-
+Error RawMemProfReader::readNextRecord(
+    GuidMemProfRecordPair &GuidRecord,
+    std::function<const Frame(const FrameId)> Callback) {
+  // Create a new callback for the RawMemProfRecord iterator so that we can
+  // provide the symbol name if the reader was initialized with KeepSymbolName =
+  // true. This is useful for debugging and testing.
   auto IdToFrameCallback = [this](const FrameId Id) {
     Frame F = this->idToFrame(Id);
     if (!this->KeepSymbolName)
       return F;
     auto Iter = this->GuidToSymbolName.find(F.Function);
     assert(Iter != this->GuidToSymbolName.end());
     F.SymbolName = Iter->getSecond();
     return F;
   };
-
-  const IndexedMemProfRecord &IndexedRecord = Iter->second;
-  GuidRecord = {Iter->first, MemProfRecord(IndexedRecord, IdToFrameCallback)};
-  Iter++;
-  return Error::success();
+  return MemProfReader::readNextRecord(GuidRecord, IdToFrameCallback);
 }
 } // namespace memprof
 } // namespace llvm