Merge pull request #41826 from atrick/fix-recursive-closure

Rewrite ClosureScopeAnalysis for generality.

Author: atrick

include/swift/SIL/ApplySite.h

@@ -169,6 +169,10 @@ class ApplySite {
     FOREACH_IMPL_RETURN(getCalleeFunction());
   }
 
+  bool isCalleeDynamicallyReplaceable() const {
+    FOREACH_IMPL_RETURN(isCalleeDynamicallyReplaceable());
+  }
+
   /// Return the referenced function if the callee is a function_ref
   /// instruction.
   SILFunction *getReferencedFunctionOrNull() const {

include/swift/SIL/SILFunction.h

@@ -224,6 +224,10 @@ class SILFunction
   /// For details see BasicBlockBitfield::bitfieldID;
   unsigned currentBitfieldID = 1;
 
+  /// Unique identifier for vector indexing and deterministic sorting.
+  /// May be reused when zombie functions are recovered.
+  unsigned index;
+
   /// The function's set of semantics attributes.
   ///
   /// TODO: Why is this using a std::string? Why don't we use uniqued
@@ -446,6 +450,8 @@ class SILFunction
     return SILFunctionConventions(fnType, getModule());
   }
 
+  unsigned getIndex() const { return index; }
+
   SILProfiler *getProfiler() const { return Profiler; }
 
   SILFunction *getDynamicallyReplacedFunction() const {

include/swift/SIL/SILModule.h

@@ -345,6 +345,10 @@ class SILModule {
   /// The options passed into this SILModule.
   const SILOptions &Options;
 
+  /// The number of functions created in this module, which will be the index of
+  /// the next function.
+  unsigned nextFunctionIndex = 0;
+
   /// Set if the SILModule was serialized already. It is used
   /// to ensure that the module is serialized only once.
   bool serialized;
@@ -449,6 +453,12 @@ class SILModule {
   /// Called after an instruction is moved from one function to another.
   void notifyMovedInstruction(SILInstruction *inst, SILFunction *fromFunction);
 
+  unsigned getNewFunctionIndex() { return nextFunctionIndex++; }
+
+  // This may be larger that the number of live functions in the 'functions'
+  // linked list because it includes the indices of zombie functions.
+  unsigned getNumFunctionIndices() const { return nextFunctionIndex; }
+
   /// Set a serialization action.
   void setSerializeSILAction(ActionCallback SerializeSILAction);
   ActionCallback getSerializeSILAction() const;

include/swift/SILOptimizer/Analysis/ClosureScope.h

@@ -52,6 +52,7 @@
 
 #include "swift/Basic/BlotSetVector.h"
 #include "swift/SIL/SILFunction.h"
+#include "swift/SIL/SILModule.h"
 #include "swift/SILOptimizer/Analysis/Analysis.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/iterator.h"
@@ -80,40 +81,15 @@ inline bool isNonEscapingClosure(CanSILFunctionType funcTy) {
   return llvm::any_of(funcTy->getParameters(), isInoutAliasable);
 }
 
-class ClosureScopeData;
+class ClosureGraph;
 
 class ClosureScopeAnalysis : public SILAnalysis {
-  friend class ClosureScopeData;
-
-  // Get a closure's scope function from its index. This functor is compatible
-  // with OptionalTransformRange. Unfortunately it exposes the internals of the
-  // analysis data.
-  struct IndexLookupFunc {
-    // A reference to all closure parent scopes ordered by their index.
-    const std::vector<SILFunction *> &indexedScopes;
-
-    IndexLookupFunc(const std::vector<SILFunction *> &indexedScopes)
-        : indexedScopes(indexedScopes) {}
-
-    Optional<SILFunction *> operator()(int idx) const {
-      if (auto funcPtr = indexedScopes[idx]) {
-        return funcPtr;
-      }
-      return None;
-    }
-  };
-  using IndexRange = iterator_range<int *>;
-
-public:
-  // A range of SILFunction scopes converted from their scope indices and
-  // filtered to remove any erased functions.
-  using ScopeRange = OptionalTransformRange<IndexRange, IndexLookupFunc, int *>;
+  friend class ClosureGraph;
 
-private:
   SILModule *M;
 
   // The analysis data. nullptr if it has never been computed.
-  std::unique_ptr<ClosureScopeData> scopeData;
+  std::unique_ptr<ClosureGraph> scopeGraph;
 
 public:
   ClosureScopeAnalysis(SILModule *M);
@@ -125,13 +101,20 @@ class ClosureScopeAnalysis : public SILAnalysis {
 
   SILModule *getModule() const { return M; }
 
-  // Return true if the given function is the parent scope for any closures.
-  bool isClosureScope(SILFunction *scopeFunc);
+  /// Visit the parent scopes of \p closure if it has any. If \p visitor returns
+  /// false, exit early and return false. Otherwise return true.
+  bool visitClosureScopes(SILFunction *closure,
+                          std::function<bool(SILFunction *scopeFunc)> visitor);
 
-  // Return a range of scopes for the given closure. The elements of the
-  // returned range have type `SILFunction *` and are non-null. Returns an
-  // empty range for a SILFunction that is not a closure or is a dead closure.
-  ScopeRange getClosureScopes(SILFunction *closureFunc);
+  /// Visit the closures directly referenced by \p scopeFunc.
+  bool visitClosures(SILFunction *scopeFunc,
+                     std::function<bool(SILFunction *closure)> visitor);
+
+  /// Return true if this function is a reachable closure.
+  bool isReachableClosure(SILFunction *function) {
+    // This visitor returns false immediately on the first scope.
+    return !visitClosureScopes(function, [](SILFunction *) { return false; });
+  }
 
   /// Invalidate all information in this analysis.
   virtual void invalidate() override;
@@ -159,25 +142,50 @@ class ClosureScopeAnalysis : public SILAnalysis {
   }
 
 protected:
-  ClosureScopeData *getOrComputeScopeData();
+  ClosureScopeAnalysis(const ClosureScopeAnalysis &) = delete;
+  ClosureScopeAnalysis &operator=(const ClosureScopeAnalysis &) = delete;
+
+  ClosureGraph *getOrComputeGraph();
 };
 
-// ClosureScopeAnalysis utility for visiting functions top down in closure scope
-// order.
-class TopDownClosureFunctionOrder {
-  ClosureScopeAnalysis *CSA;
+// ClosureScopeAnalysis utility for visiting functions top-down or bottom-up in
+// closure scope order.
+class ClosureFunctionOrder {
+  class ClosureDFS;
+  friend class ClosureDFS;
+
+  ClosureScopeAnalysis *csa;
 
-  llvm::SmallSet<SILFunction *, 16> visited;
+  // All functions in this module in top-down order (RPO) following the closure
+  // scope graph. Functions that define a closure occur before the closure.
+  std::vector<SILFunction *> topDownFunctions;
 
-  BlotSetVector<SILFunction *> closureWorklist;
+  // If the closure scope graph has any cycles, record each function at the
+  // head of a cycle. This does not include all the functions in the
+  // strongly-connected component. This is extremely rare. It is always a local
+  // function that refers to itself either directly or indirectly.
+  SmallPtrSet<SILFunction *, 4> closureCycleHeads;
 
 public:
-  TopDownClosureFunctionOrder(ClosureScopeAnalysis *CSA) : CSA(CSA) {}
+  ClosureFunctionOrder(ClosureScopeAnalysis *csa) : csa(csa) {}
 
-  // Visit all functions in a module, visiting each closure scope function
-  // before
-  // the closure function itself.
-  void visitFunctions(llvm::function_ref<void(SILFunction *)> visitor);
+  void compute();
+
+  ArrayRef<SILFunction *> getTopDownFunctions() const {
+    assert(!topDownFunctions.empty()
+           || llvm::empty(csa->getModule()->getFunctions()));
+    return topDownFunctions;
+  }
+
+  bool isHeadOfClosureCycle(SILFunction *function) const {
+    return closureCycleHeads.contains(function);
+  }
+
+  SWIFT_ASSERT_ONLY_DECL(void dump());
+
+protected:
+  ClosureFunctionOrder(const ClosureFunctionOrder &) = delete;
+  ClosureFunctionOrder &operator=(const ClosureFunctionOrder &) = delete;
 };
 
 } // end namespace swift

lib/SIL/IR/SILFunction.cpp

@@ -150,8 +150,9 @@ SILFunction::SILFunction(SILModule &Module, SILLinkage Linkage, StringRef Name,
                          IsDynamicallyReplaceable_t isDynamic,
                          IsExactSelfClass_t isExactSelfClass,
                          IsDistributed_t isDistributed)
-    : SwiftObjectHeader(functionMetatype),
-      Module(Module), Availability(AvailabilityContext::alwaysAvailable())  {
+    : SwiftObjectHeader(functionMetatype), Module(Module),
+      index(Module.getNewFunctionIndex()),
+      Availability(AvailabilityContext::alwaysAvailable()) {
   init(Linkage, Name, LoweredType, genericEnv, Loc, isBareSILFunction, isTrans,
        isSerialized, entryCount, isThunk, classSubclassScope, inlineStrategy,
        E, DebugScope, isDynamic, isExactSelfClass, isDistributed);

lib/SILOptimizer/Analysis/AccessSummaryAnalysis.cpp

@@ -162,6 +162,9 @@ static bool hasExpectedUsesOfNoEscapePartialApply(Operand *partialApplyUse) {
   if (user->getModule().getASTContext().hadError())
     return true;
 
+  if (isIncidentalUse(user))
+    return true;
+
   // It is fine to call the partial apply
   switch (user->getKind()) {
   case SILInstructionKind::ApplyInst:
@@ -206,10 +209,6 @@ static bool hasExpectedUsesOfNoEscapePartialApply(Operand *partialApplyUse) {
     return llvm::all_of(cast<CopyValueInst>(user)->getUses(),
                         hasExpectedUsesOfNoEscapePartialApply);
 
-  // End borrow is always ok.
-  case SILInstructionKind::EndBorrowInst:
-    return true;
-
   case SILInstructionKind::IsEscapingClosureInst:
   case SILInstructionKind::StoreInst:
   case SILInstructionKind::DestroyValueInst: