diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 3ef6ca8586fb2..60bff76d0e464 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/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);
diff --git a/llvm/lib/ProfileData/MemProf.cpp b/llvm/lib/ProfileData/MemProf.cpp
index 4667778ca11dd..f5789186094c6 100644
--- a/llvm/lib/ProfileData/MemProf.cpp
+++ b/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) {
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 8b97866e403fc..a913718d0fe06 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/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