From 5617fb1411f765667c016b5b75daa9d1110c36af Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin@google.com>
Date: Fri, 4 Nov 2022 16:05:10 -0700
Subject: [PATCH] [MLGO][NFC] Use std::map instead of DenseMap to avoid use
 after free

In `MLInlineAdvisor::getAdviceImpl`, we call `getCachedFPI` twice, once
for the caller, once for the callee, so the second may invalidate the
reference obtained by the first because the underlying implementation of
the cache is a `DenseMap`. `std::map` doesn't have that problem.
---
 llvm/include/llvm/Analysis/MLInlineAdvisor.h | 2 +-
 llvm/lib/Analysis/MLInlineAdvisor.cpp        | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MLInlineAdvisor.h b/llvm/include/llvm/Analysis/MLInlineAdvisor.h
index 00e8d7d7dd4d..3db948d365c7 100644
--- a/llvm/include/llvm/Analysis/MLInlineAdvisor.h
+++ b/llvm/include/llvm/Analysis/MLInlineAdvisor.h
@@ -69,7 +69,7 @@ private:
   getSkipAdviceIfUnreachableCallsite(CallBase &CB);
   void print(raw_ostream &OS) const override;
 
-  mutable DenseMap<const Function *, FunctionPropertiesInfo> FPICache;
+  mutable std::map<const Function *, FunctionPropertiesInfo> FPICache;
 
   LazyCallGraph &CG;
 
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index f55de71ea98a..a20c05243b77 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -415,8 +415,8 @@ void MLInlineAdvisor::print(raw_ostream &OS) const {
      << " EdgesOfLastSeenNodes: " << EdgesOfLastSeenNodes << "\n";
   OS << "[MLInlineAdvisor] FPI:\n";
   for (auto I : FPICache) {
-    OS << I.getFirst()->getName() << ":\n";
-    I.getSecond().print(OS);
+    OS << I.first->getName() << ":\n";
+    I.second.print(OS);
     OS << "\n";
   }
   OS << "\n";