CallerAnalysis: be more tolerant about missing FunctionInfos

include/swift/SILOptimizer/Analysis/CallerAnalysis.h

@@ -91,12 +91,7 @@ class CallerAnalysis final : public SILAnalysis {
 
   /// Notify the analysis about a function which will be deleted from the
   /// module.
-  void notifyWillDeleteFunction(SILFunction *f) override {
-    invalidateAllInfo(f);
-    recomputeFunctionList.remove(f);
-    // Now that we have invalidated all references to the function, delete it.
-    funcInfos.erase(f);
-  }
+  void notifyWillDeleteFunction(SILFunction *f) override;
 
   /// Notify the analysis about changed witness or vtables.
   ///
@@ -107,6 +102,9 @@ class CallerAnalysis final : public SILAnalysis {
 
   /// Look up the function info that we have stored for f, recomputing all
   /// invalidating parts of the call graph.
+  ///
+  /// Warning: The returned FunctionInfo is only alive until the next call to
+  /// `getFunctionInfo`.
   const FunctionInfo &getFunctionInfo(SILFunction *f) const;
 
   SWIFT_DEBUG_DUMP;
@@ -146,7 +144,7 @@ class CallerAnalysis final : public SILAnalysis {
   void invalidateKnownCallees(SILFunction *caller, FunctionInfo &callerInfo);
 
   /// Invalidate both the known callees of f and the known callers of f.
-  void invalidateAllInfo(SILFunction *f);
+  void invalidateAllInfo(SILFunction *f, FunctionInfo &fInfo);
 
   /// Helper method that reprocesses all elements of recomputeFunctionList and
   /// then clears the function list.

lib/SILOptimizer/Analysis/CallerAnalysis.cpp

@@ -48,16 +48,14 @@ struct CallerAnalysis::ApplySiteFinderVisitor
     : SILInstructionVisitor<ApplySiteFinderVisitor, bool> {
   CallerAnalysis *analysis;
   SILFunction *callerFn;
-  FunctionInfo &callerInfo;
 
 #ifndef NDEBUG
   SmallPtrSet<SILInstruction *, 8> visitedCallSites;
   SmallSetVector<SILInstruction *, 8> callSitesThatMustBeVisited;
 #endif
 
   ApplySiteFinderVisitor(CallerAnalysis *analysis, SILFunction *callerFn)
-      : analysis(analysis), callerFn(callerFn),
-        callerInfo(analysis->unsafeGetFunctionInfo(callerFn)) {}
+      : analysis(analysis), callerFn(callerFn) {}
   ~ApplySiteFinderVisitor();
 
   bool visitSILInstruction(SILInstruction *) { return false; }
@@ -120,14 +118,14 @@ bool CallerAnalysis::ApplySiteFinderVisitor::visitFunctionRefBaseInst(
     FunctionRefBaseInst *fri) {
   auto optResult = findLocalApplySites(fri);
   auto *calleeFn = fri->getInitiallyReferencedFunction();
-  FunctionInfo &calleeInfo = analysis->unsafeGetFunctionInfo(calleeFn);
 
   // First make an edge from our callerInfo to our calleeState for invalidation
   // purposes.
-  callerInfo.calleeStates.insert(calleeFn);
+  analysis->getOrInsertFunctionInfo(callerFn).calleeStates.insert(calleeFn);
 
   // Then grab our callee state and update it with state for this caller.
-  auto iter = calleeInfo.callerStates.insert({callerFn, {}});
+  auto iter = analysis->getOrInsertFunctionInfo(calleeFn).callerStates.
+                  insert({callerFn, {}});
   // If we succeeded in inserting a new value, put in an optimistic
   // value for escaping.
   if (iter.second) {
@@ -222,6 +220,10 @@ const FunctionInfo &CallerAnalysis::getFunctionInfo(SILFunction *f) const {
   // Recompute every function in the invalidated function list and empty the
   // list.
   auto &self = const_cast<CallerAnalysis &>(*this);
+  if (funcInfos.find(f) == funcInfos.end()) {
+    (void)self.getOrInsertFunctionInfo(f);
+    self.recomputeFunctionList.insert(f);
+  }
   self.processRecomputeFunctionList();
   return self.unsafeGetFunctionInfo(f);
 }
@@ -262,11 +264,7 @@ void CallerAnalysis::processFunctionCallSites(SILFunction *callerFn) {
   visitor.process();
 }
 
-void CallerAnalysis::invalidateAllInfo(SILFunction *f) {
-  // Look up the callees that our caller refers to and invalidate any
-  // values that point back at the caller.
-  FunctionInfo &fInfo = unsafeGetFunctionInfo(f);
-
+void CallerAnalysis::invalidateAllInfo(SILFunction *f, FunctionInfo &fInfo) {
   // Then we first eliminate any callees that we point at.
   invalidateKnownCallees(f, fInfo);
 
@@ -303,9 +301,12 @@ void CallerAnalysis::invalidateKnownCallees(SILFunction *caller,
 }
 
 void CallerAnalysis::invalidateKnownCallees(SILFunction *caller) {
-  // Look up the callees that our caller refers to and invalidate any
-  // values that point back at the caller.
-  invalidateKnownCallees(caller, unsafeGetFunctionInfo(caller));
+  auto iter = funcInfos.find(caller);
+  if (iter != funcInfos.end()) {
+    // Look up the callees that our caller refers to and invalidate any
+    // values that point back at the caller.
+    invalidateKnownCallees(caller, iter->second);
+  }
 }
 
 void CallerAnalysis::verify(SILFunction *caller) const {
@@ -377,6 +378,17 @@ void CallerAnalysis::invalidate() {
   }
 }
 
+void CallerAnalysis::notifyWillDeleteFunction(SILFunction *f) {
+  auto iter = funcInfos.find(f);
+  if (iter == funcInfos.end())
+    return;
+
+  invalidateAllInfo(f, iter->second);
+  recomputeFunctionList.remove(f);
+  // Now that we have invalidated all references to the function, delete it.
+  funcInfos.erase(iter);
+}
+
 //===----------------------------------------------------------------------===//
 //                         CallerAnalysis YAML Dumper
 //===----------------------------------------------------------------------===//