Add TopDownClosureFunctionOrder based on ClosureScopeAnalysis.

Simple utility for transfersing functions such that parent scopes are always
visited before noescape closures.

Note that recursion is disallowed. Noescape closures are not reentrant.

Author: atrick

include/swift/SILOptimizer/Analysis/ClosureScope.h

@@ -48,6 +48,7 @@
 #ifndef SWIFT_SILOPTIMIZER_ANALYSIS_CLOSURESCOPE_H
 #define SWIFT_SILOPTIMIZER_ANALYSIS_CLOSURESCOPE_H
 
+#include "swift/Basic/BlotSetVector.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SILOptimizer/Analysis/Analysis.h"
 #include "llvm/ADT/SmallSet.h"
@@ -120,6 +121,8 @@ class ClosureScopeAnalysis : public SILAnalysis {
     return S->getKind() == AnalysisKind::ClosureScope;
   }
 
+  SILModule *getModule() const { return M; }
+
   // Return true if the given function is the parent scope for any closures.
   bool isClosureScope(SILFunction *scopeFunc);
 
@@ -157,6 +160,24 @@ class ClosureScopeAnalysis : public SILAnalysis {
   ClosureScopeData *getOrComputeScopeData();
 };
 
+// ClosureScopeAnalysis utility for visiting functions top down in closure scope
+// order.
+class TopDownClosureFunctionOrder {
+  ClosureScopeAnalysis *CSA;
+
+  llvm::SmallSet<SILFunction *, 16> visited;
+
+  BlotSetVector<SILFunction *> closureWorklist;
+
+public:
+  TopDownClosureFunctionOrder(ClosureScopeAnalysis *CSA) : CSA(CSA) {}
+
+  // Visit all functions in a module, visiting each closure scope function
+  // before
+  // the closure function itself.
+  void visitFunctions(std::function<void(SILFunction *)> visitor);
+};
+
 } // end namespace swift
 
 #endif

lib/SILOptimizer/Analysis/ClosureScope.cpp

@@ -14,6 +14,7 @@
 
 #define DEBUG_TYPE "closure-scope"
 
+#include "swift/SIL/SILModule.h"
 #include "swift/SILOptimizer/Analysis/ClosureScope.h"
 
 namespace swift {
@@ -154,4 +155,42 @@ SILAnalysis *createClosureScopeAnalysis(SILModule *M) {
   return new ClosureScopeAnalysis(M);
 }
 
+void TopDownClosureFunctionOrder::visitFunctions(
+    std::function<void(SILFunction *)> visitor) {
+  auto markVisited = [&](SILFunction *F) {
+    bool visitOnce = visited.insert(F).second;
+    assert(visitOnce);
+    (void)visitOnce;
+  };
+  auto allScopesVisited = [&](SILFunction *closureF) {
+    return llvm::all_of(CSA->getClosureScopes(closureF),
+                        [this](SILFunction *F) { return visited.count(F); });
+  };
+  for (auto &F : *CSA->getModule()) {
+    if (!allScopesVisited(&F)) {
+      closureWorklist.insert(&F);
+      continue;
+    }
+    markVisited(&F);
+    visitor(&F);
+  }
+  unsigned numClosures = closureWorklist.size();
+  while (numClosures) {
+    for (auto &closureNode : closureWorklist) {
+      // skip erased closures.
+      if (!closureNode)
+        continue;
+
+      auto closureF = closureNode.getValue();
+      if (!allScopesVisited(closureF))
+        continue;
+
+      markVisited(closureF);
+      visitor(closureF);
+      closureWorklist.erase(closureF);
+      --numClosures;
+    }
+  }
+}
+
 } // namespace swift