Move findAsyncNode to BraceStmt

I need to determine when top-level code contains an `await` to determine
whether to make the source file an async context. This logic is
perfectly encapsulated in the `FindInnerAsync` AST walker.
Unfortunately, that is pushed down in Sema/ConstraintSystem and isn't
available at the AST level. I've pulled it up into the brace statement
so that I can use that as part of determining whether the source file is
async in `DeclContext::isAsyncContext`.

Unfortunately, statements don't have an AST context or evaluator or I
would make this a request.

Author: etcwilde

include/swift/AST/Stmt.h

@@ -161,7 +161,7 @@ class BraceStmt final : public Stmt,
 
   SourceLoc getLBraceLoc() const { return LBLoc; }
   SourceLoc getRBraceLoc() const { return RBLoc; }
-  
+
   SourceRange getSourceRange() const { return SourceRange(LBLoc, RBLoc); }
 
   bool empty() const { return getNumElements() == 0; }
@@ -182,7 +182,9 @@ class BraceStmt final : public Stmt,
   ArrayRef<ASTNode> getElements() const {
     return {getTrailingObjects<ASTNode>(), Bits.BraceStmt.NumElements};
   }
-  
+
+  ASTNode findAsyncNode();
+
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Brace; }
 };
 

lib/AST/Stmt.cpp

@@ -16,6 +16,7 @@
 
 #include "swift/AST/Stmt.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/Pattern.h"
@@ -155,6 +156,65 @@ BraceStmt *BraceStmt::create(ASTContext &ctx, SourceLoc lbloc,
   return ::new(Buffer) BraceStmt(lbloc, elts, rbloc, implicit);
 }
 
+ASTNode BraceStmt::findAsyncNode() {
+  // TODO: Statements don't track their ASTContext/evaluator, so I am not making
+  // this a request. It probably should be a request at some point.
+  //
+  // While we're at it, it would be very nice if this could be a const
+  // operation, but the AST-walking is not a const operation.
+
+  // 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->isAsyncLet())
+          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;
+  walk(asyncFinder);
+
+  return asyncFinder.getAsyncNode();
+}
+
 SourceLoc ReturnStmt::getStartLoc() const {
   if (ReturnLoc.isInvalid() && Result)
     return Result->getStartLoc();

lib/Sema/ConstraintSystem.cpp

@@ -6111,55 +6111,5 @@ 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->isAsyncLet())
-          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();
+  return body->findAsyncNode();
 }