From 68106bd492e294ecf0a7e2829dd9edf6cd72f3ef Mon Sep 17 00:00:00 2001
From: William Junda Huang <williamjhuang@google.com>
Date: Wed, 29 Nov 2023 16:48:55 -0500
Subject: [PATCH] [Sample Profile Loader]  Fix potential invalidated reference
 (#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.
---
 .../llvm/Transforms/IPO/ProfiledCallGraph.h   | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h b/llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h
index 3f986caeb54728..5381ada37fe27e 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,7 +208,9 @@ class ProfiledCallGraph {
   }
 
   ProfiledCallGraphNode Root;
-  HashKeyMap<std::unordered_map, FunctionId, ProfiledCallGraphNode>
+  // backing buffer for ProfiledCallGraphNodes.
+  std::list<ProfiledCallGraphNode> ProfiledCallGraphNodeList;
+  HashKeyMap<llvm::DenseMap, FunctionId, ProfiledCallGraphNode*>
       ProfiledFunctions;
 };