[ConstraintSystem] NFC: Extract async node identification into a function

It's useful not only for effect determination but for diagnostics as well.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -5592,6 +5592,11 @@ Type getConcreteReplacementForProtocolSelfType(ValueDecl *member);
 /// of operator overload choices.
 bool isOperatorDisjunction(Constraint *disjunction);
 
+/// Find out whether closure body has any `async` or `await` expressions,
+/// declarations, or statements directly in its body (no in other closures
+/// or nested declarations).
+ASTNode findAsyncNode(ClosureExpr *closure);
+
 } // end namespace constraints
 
 template<typename ...Args>

lib/Sema/ConstraintSystem.cpp

@@ -2486,52 +2486,6 @@ FunctionType::ExtInfo ConstraintSystem::closureEffects(ClosureExpr *expr) {
     bool foundThrow() { return FoundThrow; }
   };
 
-  // A walker that looks for 'async' and 'await' expressions
-  // that aren't nested within closures or nested declarations.
-  class FindInnerAsync : public ASTWalker {
-    bool FoundAsync = false;
-
-    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      // If we've found an 'await', record it and terminate the traversal.
-      if (isa<AwaitExpr>(expr)) {
-        FoundAsync = true;
-        return { false, nullptr };
-      }
-
-      // Do not recurse into other closures.
-      if (isa<ClosureExpr>(expr))
-        return { false, expr };
-
-      return { true, expr };
-    }
-
-    bool walkToDeclPre(Decl *decl) override {
-      // Do not walk into function or type declarations.
-      if (auto *patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
-        if (patternBinding->isSpawnLet())
-          FoundAsync = true;
-
-        return true;
-      }
-
-      return false;
-    }
-
-    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override { 
-      if (auto forEach = dyn_cast<ForEachStmt>(stmt)) {
-        if (forEach->getAwaitLoc().isValid()) {
-          FoundAsync = true;
-          return { false, nullptr };
-        }
-      }
-
-      return { true, stmt };
-    }
-
-  public:
-    bool foundAsync() { return FoundAsync; }
-  };
-
   // If either 'throws' or 'async' was explicitly specified, use that
   // set of effects.
   bool throws = expr->getThrowsLoc().isValid();
@@ -2552,13 +2506,11 @@ FunctionType::ExtInfo ConstraintSystem::closureEffects(ClosureExpr *expr) {
 
   auto throwFinder = FindInnerThrows(*this, expr);
   body->walk(throwFinder);
-  auto asyncFinder = FindInnerAsync();
-  body->walk(asyncFinder);
   auto result = ASTExtInfoBuilder()
-    .withThrows(throwFinder.foundThrow())
-    .withAsync(asyncFinder.foundAsync())
-    .withConcurrent(concurrent)
-    .build();
+                    .withThrows(throwFinder.foundThrow())
+                    .withAsync(bool(findAsyncNode(expr)))
+                    .withConcurrent(concurrent)
+                    .build();
   closureEffectsCache[expr] = result;
   return result;
 }
@@ -5708,3 +5660,60 @@ bool constraints::isOperatorDisjunction(Constraint *disjunction) {
   auto *decl = getOverloadChoiceDecl(choices.front());
   return decl ? decl->isOperator() : false;
 }
+
+ASTNode constraints::findAsyncNode(ClosureExpr *closure) {
+  auto *body = closure->getBody();
+  if (!body)
+    return ASTNode();
+
+  // A walker that looks for 'async' and 'await' expressions
+  // that aren't nested within closures or nested declarations.
+  class FindInnerAsync : public ASTWalker {
+    ASTNode AsyncNode;
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+      // If we've found an 'await', record it and terminate the traversal.
+      if (isa<AwaitExpr>(expr)) {
+        AsyncNode = expr;
+        return {false, nullptr};
+      }
+
+      // Do not recurse into other closures.
+      if (isa<ClosureExpr>(expr))
+        return {false, expr};
+
+      return {true, expr};
+    }
+
+    bool walkToDeclPre(Decl *decl) override {
+      // Do not walk into function or type declarations.
+      if (auto *patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
+        if (patternBinding->isSpawnLet())
+          AsyncNode = patternBinding;
+
+        return true;
+      }
+
+      return false;
+    }
+
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
+      if (auto forEach = dyn_cast<ForEachStmt>(stmt)) {
+        if (forEach->getAwaitLoc().isValid()) {
+          AsyncNode = forEach;
+          return {false, nullptr};
+        }
+      }
+
+      return {true, stmt};
+    }
+
+  public:
+    ASTNode getAsyncNode() { return AsyncNode; }
+  };
+
+  FindInnerAsync asyncFinder;
+  body->walk(asyncFinder);
+
+  return asyncFinder.getAsyncNode();
+}