[llvm] [Sample Profile Loader] Fix potential invalidated reference (PR #73181)

[William Junda Huang via llvm-commits]

https://github.com/huangjd created https://github.com/llvm/llvm-project/pull/73181

There is a potential issue in ProfiledCallGraph where pointers to ProfiledCallGraphNode are used to construct edges, while ProfiledCallGraphNode instances are being added to a hash map  ProfiledFunctions simultaneously. If rehash happens, those pointers are invalidated, resulting in undefined behavior/crash. Previously (before md5phase2) ProfiledFunctions is a llvm::StringMap, which also have the same issue theoretically when rehashing but was not observed. This patch fixes this potential issue by using a backing buffer for ProrfiledCallGraphNode that does not relocate. 

>From 9265d2ae3ef1b4d12f57b879b5f1e512ec507618 Mon Sep 17 00:00:00 2001
From: William Huang <williamjhuang at google.com>
Date: Wed, 22 Nov 2023 22:04:12 +0000
Subject: [PATCH] [Sample Profile Loader]  Fix potential invalidated reference

---
 .../llvm/Transforms/IPO/ProfiledCallGraph.h   | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
index 3f986caeb547287..3d844cc92aa0a8e 100644
--- a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
+++ b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
@@ -140,8 +140,11 @@ class ProfiledCallGraph {
     if (!ProfiledFunctions.count(Name)) {
       // Link to synthetic root to make sure every node is reachable
       // from root. This does not affect SCC order.
-      ProfiledFunctions[Name] = ProfiledCallGraphNode(Name);
-      Root.Edges.emplace(&Root, &ProfiledFunctions[Name], 0);
+      // Store the pointer of the node because the map can be rehashed.
+      auto &Node =
+          ProfiledCallGraphNodeList.emplace_back(ProfiledCallGraphNode(Name));
+      ProfiledFunctions[Name] = &Node;
+      Root.Edges.emplace(&Root, ProfiledFunctions[Name], 0);
     }
   }
 
@@ -152,9 +155,9 @@ class ProfiledCallGraph {
     auto CalleeIt = ProfiledFunctions.find(CalleeName);
     if (CalleeIt == ProfiledFunctions.end())
       return;
-    ProfiledCallGraphEdge Edge(&ProfiledFunctions[CallerName],
-                               &CalleeIt->second, Weight);
-    auto &Edges = ProfiledFunctions[CallerName].Edges;
+    ProfiledCallGraphEdge Edge(ProfiledFunctions[CallerName],
+                               CalleeIt->second, Weight);
+    auto &Edges = ProfiledFunctions[CallerName]->Edges;
     auto EdgeIt = Edges.find(Edge);
     if (EdgeIt == Edges.end()) {
       Edges.insert(Edge);
@@ -193,7 +196,7 @@ class ProfiledCallGraph {
       return;
 
     for (auto &Node : ProfiledFunctions) {
-      auto &Edges = Node.second.Edges;
+      auto &Edges = Node.second->Edges;
       auto I = Edges.begin();
       while (I != Edges.end()) {
         if (I->Weight <= Threshold)
@@ -205,8 +208,9 @@ class ProfiledCallGraph {
   }
 
   ProfiledCallGraphNode Root;
-  HashKeyMap<std::unordered_map, FunctionId, ProfiledCallGraphNode>
-      ProfiledFunctions;
+  // backing buffer for ProfiledCallGraphNodes.
+  std::list<ProfiledCallGraphNode> ProfiledCallGraphNodeList;
+  llvm::DenseMap<FunctionId, ProfiledCallGraphNode*> ProfiledFunctions;
 };
 
 } // end namespace sampleprof