[memprof] Pass FrameIdConverter and CallStackIdConverter by reference (#92327)

CallStackIdConverter sets LastUnmappedId when a mapping failure
occurs.  Now, since toMemProfRecord takes an instance of
CallStackIdConverter by value, namely std::function, the caller of
toMemProfRecord never receives the mapping failure that occurs inside
toMemProfRecord.  The same problem applies to FrameIdConverter.

The patch fixes the problem by passing FrameIdConverter and
CallStackIdConverter by reference, namely llvm::function_ref.

While I am it, this patch deletes the copy constructor and copy
assignment operator to avoid accidental copies.

Author: kazutakahirata

llvm/include/llvm/ProfileData/MemProf.h

@@ -426,8 +426,8 @@ struct IndexedMemProfRecord {
   // Convert IndexedMemProfRecord to MemProfRecord.  Callback is used to
   // translate CallStackId to call stacks with frames inline.
   MemProfRecord toMemProfRecord(
-      std::function<const llvm::SmallVector<Frame>(const CallStackId)> Callback)
-      const;
+      llvm::function_ref<const llvm::SmallVector<Frame>(const CallStackId)>
+          Callback) const;
 
   // Returns the GUID for the function name after canonicalization. For
   // memprof, we remove any .llvm suffix added by LTO. MemProfRecords are
@@ -784,6 +784,12 @@ template <typename MapTy> struct FrameIdConverter {
   FrameIdConverter() = delete;
   FrameIdConverter(MapTy &Map) : Map(Map) {}
 
+  // Delete the copy constructor and copy assignment operator to avoid a
+  // situation where a copy of FrameIdConverter gets an error in LastUnmappedId
+  // while the original instance doesn't.
+  FrameIdConverter(const FrameIdConverter &) = delete;
+  FrameIdConverter &operator=(const FrameIdConverter &) = delete;
+
   Frame operator()(FrameId Id) {
     auto Iter = Map.find(Id);
     if (Iter == Map.end()) {
@@ -798,12 +804,19 @@ template <typename MapTy> struct FrameIdConverter {
 template <typename MapTy> struct CallStackIdConverter {
   std::optional<CallStackId> LastUnmappedId;
   MapTy &Map;
-  std::function<Frame(FrameId)> FrameIdToFrame;
+  llvm::function_ref<Frame(FrameId)> FrameIdToFrame;
 
   CallStackIdConverter() = delete;
-  CallStackIdConverter(MapTy &Map, std::function<Frame(FrameId)> FrameIdToFrame)
+  CallStackIdConverter(MapTy &Map,
+                       llvm::function_ref<Frame(FrameId)> FrameIdToFrame)
       : Map(Map), FrameIdToFrame(FrameIdToFrame) {}
 
+  // Delete the copy constructor and copy assignment operator to avoid a
+  // situation where a copy of CallStackIdConverter gets an error in
+  // LastUnmappedId while the original instance doesn't.
+  CallStackIdConverter(const CallStackIdConverter &) = delete;
+  CallStackIdConverter &operator=(const CallStackIdConverter &) = delete;
+
   llvm::SmallVector<Frame> operator()(CallStackId CSId) {
     llvm::SmallVector<Frame> Frames;
     auto CSIter = Map.find(CSId);

llvm/lib/ProfileData/MemProf.cpp

@@ -243,8 +243,8 @@ IndexedMemProfRecord::deserialize(const MemProfSchema &Schema,
 }
 
 MemProfRecord IndexedMemProfRecord::toMemProfRecord(
-    std::function<const llvm::SmallVector<Frame>(const CallStackId)> Callback)
-    const {
+    llvm::function_ref<const llvm::SmallVector<Frame>(const CallStackId)>
+        Callback) const {
   MemProfRecord Record;
 
   for (const memprof::IndexedAllocationInfo &IndexedAI : AllocSites) {

llvm/unittests/ProfileData/MemProfTest.cpp

@@ -596,4 +596,70 @@ TEST(MemProf, IndexedMemProfRecordToMemProfRecord) {
   EXPECT_EQ(Record.CallSites[1][0].hash(), F2.hash());
   EXPECT_EQ(Record.CallSites[1][1].hash(), F4.hash());
 }
+
+using FrameIdMapTy =
+    llvm::DenseMap<::llvm::memprof::FrameId, ::llvm::memprof::Frame>;
+using CallStackIdMapTy =
+    llvm::DenseMap<::llvm::memprof::CallStackId,
+                   ::llvm::SmallVector<::llvm::memprof::FrameId>>;
+
+// Populate those fields returned by getHotColdSchema.
+MemInfoBlock makePartialMIB() {
+  MemInfoBlock MIB;
+  MIB.AllocCount = 1;
+  MIB.TotalSize = 5;
+  MIB.TotalLifetime = 10;
+  MIB.TotalLifetimeAccessDensity = 23;
+  return MIB;
+}
+
+TEST(MemProf, MissingCallStackId) {
+  // Use a non-existent CallStackId to trigger a mapping error in
+  // toMemProfRecord.
+  llvm::memprof::IndexedAllocationInfo AI({}, 0xdeadbeefU, makePartialMIB(),
+                                          llvm::memprof::getHotColdSchema());
+
+  IndexedMemProfRecord IndexedMR;
+  IndexedMR.AllocSites.push_back(AI);
+
+  // Create empty maps.
+  const FrameIdMapTy IdToFrameMap;
+  const CallStackIdMapTy CSIdToCallStackMap;
+  llvm::memprof::FrameIdConverter<decltype(IdToFrameMap)> FrameIdConv(
+      IdToFrameMap);
+  llvm::memprof::CallStackIdConverter<decltype(CSIdToCallStackMap)> CSIdConv(
+      CSIdToCallStackMap, FrameIdConv);
+
+  // We are only interested in errors, not the return value.
+  (void)IndexedMR.toMemProfRecord(CSIdConv);
+
+  ASSERT_TRUE(CSIdConv.LastUnmappedId.has_value());
+  EXPECT_EQ(*CSIdConv.LastUnmappedId, 0xdeadbeefU);
+  EXPECT_EQ(FrameIdConv.LastUnmappedId, std::nullopt);
+}
+
+TEST(MemProf, MissingFrameId) {
+  llvm::memprof::IndexedAllocationInfo AI({}, 0x222, makePartialMIB(),
+                                          llvm::memprof::getHotColdSchema());
+
+  IndexedMemProfRecord IndexedMR;
+  IndexedMR.AllocSites.push_back(AI);
+
+  // An empty map to trigger a mapping error.
+  const FrameIdMapTy IdToFrameMap;
+  CallStackIdMapTy CSIdToCallStackMap;
+  CSIdToCallStackMap.insert({0x222, {2, 3}});
+
+  llvm::memprof::FrameIdConverter<decltype(IdToFrameMap)> FrameIdConv(
+      IdToFrameMap);
+  llvm::memprof::CallStackIdConverter<decltype(CSIdToCallStackMap)> CSIdConv(
+      CSIdToCallStackMap, FrameIdConv);
+
+  // We are only interested in errors, not the return value.
+  (void)IndexedMR.toMemProfRecord(CSIdConv);
+
+  EXPECT_EQ(CSIdConv.LastUnmappedId, std::nullopt);
+  ASSERT_TRUE(FrameIdConv.LastUnmappedId.has_value());
+  EXPECT_EQ(*FrameIdConv.LastUnmappedId, 3U);
+}
 } // namespace