Merge pull request #27833 from eeckstein/fix-basic-callee-analysis

BasicCalleeAnalysis: fix computation of class method callees if an overridden function has a different ABI

Author: eeckstein

include/swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h

@@ -54,6 +54,11 @@ class CalleeList {
       : CalleeFunctions(llvm::makeArrayRef(List.begin(), List.end())),
         IsIncomplete(IsIncomplete) {}
 
+  LLVM_ATTRIBUTE_DEPRECATED(void dump() const LLVM_ATTRIBUTE_USED,
+                            "Only for use in the debugger");
+
+  void print(llvm::raw_ostream &os) const;
+
   /// Return an iterator for the beginning of the list.
   ArrayRef<SILFunction *>::iterator begin() const {
     return CalleeFunctions.begin();
@@ -67,7 +72,7 @@ class CalleeList {
   bool isIncomplete() const { return IsIncomplete; }
 
   /// Returns true if all callees are known and not external.
-  bool allCalleesVisible();
+  bool allCalleesVisible() const;
 };
 
 /// CalleeCache is a helper class that builds lists of potential
@@ -106,16 +111,16 @@ class CalleeCache {
   /// given instruction. E.g. it could be destructors.
   CalleeList getCalleeList(SILInstruction *I) const;
 
+  CalleeList getCalleeList(SILDeclRef Decl) const;
+
 private:
   void enumerateFunctionsInModule();
   void sortAndUniqueCallees();
   CalleesAndCanCallUnknown &getOrCreateCalleesForMethod(SILDeclRef Decl);
-  void computeClassMethodCalleesForClass(ClassDecl *CD);
-  void computeClassMethodCallees(ClassDecl *CD, SILDeclRef Method);
+  void computeClassMethodCallees();
   void computeWitnessMethodCalleesForWitnessTable(SILWitnessTable &WT);
   void computeMethodCallees();
   SILFunction *getSingleCalleeForWitnessMethod(WitnessMethodInst *WMI) const;
-  CalleeList getCalleeList(SILDeclRef Decl) const;
   CalleeList getCalleeList(WitnessMethodInst *WMI) const;
   CalleeList getCalleeList(ClassMethodInst *CMI) const;
   CalleeList getCalleeListForCalleeKind(SILValue Callee) const;
@@ -162,17 +167,23 @@ class BasicCalleeAnalysis : public SILAnalysis {
     Cache.reset();
   }
 
-  CalleeList getCalleeList(FullApplySite FAS) {
+  LLVM_ATTRIBUTE_DEPRECATED(void dump() const LLVM_ATTRIBUTE_USED,
+                            "Only for use in the debugger");
+
+  void print(llvm::raw_ostream &os) const;
+
+  void updateCache() {
     if (!Cache)
       Cache = llvm::make_unique<CalleeCache>(M);
+  }
 
+  CalleeList getCalleeList(FullApplySite FAS) {
+    updateCache();
     return Cache->getCalleeList(FAS);
   }
 
   CalleeList getCalleeList(SILInstruction *I) {
-    if (!Cache)
-      Cache = llvm::make_unique<CalleeCache>(M);
-
+    updateCache();
     return Cache->getCalleeList(I);
   }
 };

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -316,6 +316,10 @@ void replaceLoadSequence(SILInstruction *inst, SILValue value);
 /// be reached by calling the function represented by Decl?
 bool calleesAreStaticallyKnowable(SILModule &module, SILDeclRef decl);
 
+/// Do we have enough information to determine all callees that could
+/// be reached by calling the function represented by Decl?
+bool calleesAreStaticallyKnowable(SILModule &module, AbstractFunctionDecl *afd);
+
 // Attempt to get the instance for , whose static type is the same as
 // its exact dynamic type, returning a null SILValue() if we cannot find it.
 // The information that a static type is the same as the exact dynamic,

lib/SILOptimizer/Analysis/BasicCalleeAnalysis.cpp

@@ -25,7 +25,24 @@
 
 using namespace swift;
 
-bool CalleeList::allCalleesVisible() {
+void CalleeList::dump() const {
+  print(llvm::errs());
+}
+
+void CalleeList::print(llvm::raw_ostream &os) const {
+  os << "Incomplete callee list? : "
+               << (isIncomplete() ? "Yes" : "No");
+  if (!allCalleesVisible())
+    os <<", not all callees visible";
+  os << '\n';
+  os << "Known callees:\n";
+  for (auto *CalleeFn : *this) {
+    os << "  " << CalleeFn->getName() << "\n";
+  }
+  os << "\n";
+}
+
+bool CalleeList::allCalleesVisible() const {
   if (isIncomplete())
     return false;
 
@@ -81,48 +98,45 @@ CalleeCache::getOrCreateCalleesForMethod(SILDeclRef Decl) {
   return It->second;
 }
 
-/// Update the callees for each method of a given class, along with
-/// all the overridden methods from superclasses.
-void CalleeCache::computeClassMethodCalleesForClass(ClassDecl *CD) {
-  assert(!CD->hasClangNode());
-
-  for (auto *Member : CD->getMembers()) {
-    auto *AFD = dyn_cast<AbstractFunctionDecl>(Member);
-    if (!AFD)
-      continue;
-
-    if (auto *ConstrDecl = dyn_cast<ConstructorDecl>(AFD)) {
-      computeClassMethodCallees(CD, SILDeclRef(AFD,
-                                               SILDeclRef::Kind::Initializer));
-      if (ConstrDecl->isRequired()) {
-        computeClassMethodCallees(CD, SILDeclRef(AFD,
-                                                 SILDeclRef::Kind::Allocator));
+/// Update the callees for each method of a given vtable.
+void CalleeCache::computeClassMethodCallees() {
+  SmallPtrSet<AbstractFunctionDecl *, 16> unknownCallees;
+
+  // First mark all method declarations which might be overridden in another
+  // translation unit, i.e. outside the visibility of the optimizer.
+  // This is a little bit more complicated than to just check the VTable
+  // entry.Method itself, because an overridden method might be more accessible
+  // than the base method (e.g. a public method overrides a private method).
+  for (auto &VTable : M.getVTableList()) {
+    assert(!VTable.getClass()->hasClangNode());
+
+    for (Decl *member : VTable.getClass()->getMembers()) {
+      if (auto *afd = dyn_cast<AbstractFunctionDecl>(member)) {
+        // If a method implementation might be overridden in another translation
+        // unit, also mark all the base methods as 'unknown'.
+        bool unknown = false;
+        do {
+          if (!calleesAreStaticallyKnowable(M, afd))
+            unknown = true;
+          if (unknown)
+            unknownCallees.insert(afd);
+          afd = afd->getOverriddenDecl();
+        } while (afd);
       }
-    } else {
-      computeClassMethodCallees(CD, SILDeclRef(AFD));
     }
   }
-}
-
-void CalleeCache::computeClassMethodCallees(ClassDecl *CD, SILDeclRef Method) {
-  auto *CalledFn = M.lookUpFunctionInVTable(CD, Method);
-  if (!CalledFn)
-    return;
-
-  bool canCallUnknown = !calleesAreStaticallyKnowable(M, Method);
-
-  // Update the callees for this method and all the methods it
-  // overrides by adding this function to their lists.
-  do {
-    auto &TheCallees = getOrCreateCalleesForMethod(Method);
-    assert(TheCallees.getPointer() && "Unexpected null callees!");
-
-    TheCallees.getPointer()->push_back(CalledFn);
-    if (canCallUnknown)
-      TheCallees.setInt(true);
 
-    Method = Method.getNextOverriddenVTableEntry();
-  } while (Method);
+  // Second step: collect all implementations of a method.
+  for (auto &VTable : M.getVTableList()) {
+    for (const SILVTable::Entry &entry : VTable.getEntries()) {
+      if (auto *afd = entry.Method.getAbstractFunctionDecl()) {
+        CalleesAndCanCallUnknown &callees = getOrCreateCalleesForMethod(entry.Method);
+        if (unknownCallees.count(afd) != 0)
+          callees.setInt(1);
+        callees.getPointer()->push_back(entry.Implementation);
+      }
+    }
+  }
 }
 
 void CalleeCache::computeWitnessMethodCalleesForWitnessTable(
@@ -183,8 +197,8 @@ void CalleeCache::computeWitnessMethodCalleesForWitnessTable(
 /// Witness Table.
 void CalleeCache::computeMethodCallees() {
   SWIFT_FUNC_STAT;
-  for (auto &VTable : M.getVTableList())
-    computeClassMethodCalleesForClass(VTable.getClass());
+
+  computeClassMethodCallees();
 
   for (auto &WTable : M.getWitnessTableList())
     computeWitnessMethodCalleesForWitnessTable(WTable);
@@ -284,3 +298,22 @@ CalleeList CalleeCache::getCalleeList(SILInstruction *I) const {
   SILDeclRef Destructor = SILDeclRef(Class->getDestructor());
   return getCalleeList(Destructor);
 }
+
+void BasicCalleeAnalysis::dump() const {
+  print(llvm::errs());
+}
+
+void BasicCalleeAnalysis::print(llvm::raw_ostream &os) const {
+  if (!Cache) {
+    os << "<no cache>\n";
+  }
+  llvm::DenseSet<SILDeclRef> printed;
+  for (auto &VTable : M.getVTableList()) {
+    for (const SILVTable::Entry &entry : VTable.getEntries()) {
+      if (printed.insert(entry.Method).second) {
+        os << "callees for " << entry.Method << ":\n";
+        Cache->getCalleeList(entry.Method).print(os);
+      }
+    }
+  }
+}

lib/SILOptimizer/UtilityPasses/BasicCalleePrinter.cpp

@@ -40,12 +40,7 @@ class BasicCalleePrinterPass : public SILModuleTransform {
     llvm::outs() << *FAS.getInstruction();
 
     auto Callees = BCA->getCalleeList(FAS);
-    llvm::outs() << "Incomplete callee list? : "
-                 << (Callees.isIncomplete() ? "Yes" : "No") << "\n";
-    llvm::outs() << "Known callees:\n";
-    for (auto *CalleeFn : Callees)
-      llvm::outs() << CalleeFn->getName() << "\n";
-    llvm::outs() << "\n";
+    Callees.print(llvm::outs());
   }
 
   /// The entry point to the transformation.

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1433,13 +1433,19 @@ bool swift::calleesAreStaticallyKnowable(SILModule &module, SILDeclRef decl) {
   if (decl.isForeign)
     return false;
 
+  auto *afd = decl.getAbstractFunctionDecl();
+  assert(afd && "Expected abstract function decl!");
+  return calleesAreStaticallyKnowable(module, afd);
+}
+
+/// Are the callees that could be called through Decl statically
+/// knowable based on the Decl and the compilation mode?
+bool swift::calleesAreStaticallyKnowable(SILModule &module,
+                                         AbstractFunctionDecl *afd) {
   const DeclContext *assocDC = module.getAssociatedContext();
   if (!assocDC)
     return false;
 
-  auto *afd = decl.getAbstractFunctionDecl();
-  assert(afd && "Expected abstract function decl!");
-
   // Only handle members defined within the SILModule's associated context.
   if (!afd->isChildContextOf(assocDC))
     return false;

test/SILOptimizer/callee_analysis_crash.swift

@@ -0,0 +1,41 @@
+// RUN: %target-swift-frontend -O %s -emit-ir -o /dev/null
+
+// Check that we don't crash here.
+
+enum E<T> {
+  case A
+  case B(T)
+}
+
+public protocol P {
+  associatedtype A
+}
+
+internal class Base<T: P> {
+  func foo() -> E<T.A> {
+    return .A
+  }
+}
+
+struct Outer<T: P> where T.A == Int {
+  private class Inner : Base<T> {
+
+    // This overridden function has a different ABI than the base implementation.
+    // The BasicCalleeAnalysis put both methods in the callee list, which let
+    // some optimizations crash.
+    override func foo() -> E<T.A> {
+      return .A
+    }
+  }
+}
+
+@inline(never)
+func getit<T: P>(_ t: T) -> Base<T> {
+  return Base<T>()
+}
+
+public func testit<T: P>(_ t: T) {
+  let b = getit(t)
+  b.foo()
+}
+