Merge pull request #41061 from etcwilde/ewilde/async-top-level-detection

Concurrent top-level detection

Author: etcwilde

include/swift/AST/SourceFile.h

@@ -582,6 +582,9 @@ class SourceFile final : public FileUnit {
 
   ArrayRef<OpaqueTypeDecl *> getOpaqueReturnTypeDecls();
 
+  /// Returns true if the source file contains concurrency in the top-level
+  bool isAsyncTopLevelSourceFile() const;
+
 private:
 
   /// If not \c None, the underlying vector contains the parsed tokens of this

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; }
 };
 

include/swift/AST/TypeCheckRequests.h

@@ -3283,6 +3283,21 @@ class ClosureEffectsRequest
   bool isCached() const { return true; }
 };
 
+class GetSourceFileAsyncNode
+    : public SimpleRequest<GetSourceFileAsyncNode, ASTNode(const SourceFile *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ASTNode evaluate(Evaluator &evaluator, const SourceFile *) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 void simple_display(llvm::raw_ostream &out, Type value);
 void simple_display(llvm::raw_ostream &out, const TypeRepr *TyR);
 void simple_display(llvm::raw_ostream &out, ImplicitMemberAction action);

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -378,3 +378,6 @@ SWIFT_REQUEST(TypeChecker, RenamedDeclRequest,
 SWIFT_REQUEST(TypeChecker, ClosureEffectsRequest,
               FunctionType::ExtInfo(ClosureExpr *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, GetSourceFileAsyncNode,
+              AwaitExpr *(const SourceFile *),
+              Cached, NoLocationInfo)

lib/AST/Decl.cpp

@@ -736,8 +736,8 @@ bool Decl::preconcurrency() const {
   // Variables declared in top-level code are @_predatesConcurrency
   if (const VarDecl *var = dyn_cast<VarDecl>(this)) {
     const LangOptions &langOpts = getASTContext().LangOpts;
-    return !langOpts.isSwiftVersionAtLeast(6) &&
-           langOpts.EnableExperimentalAsyncTopLevel && var->isTopLevelGlobal();
+    return !langOpts.isSwiftVersionAtLeast(6) && var->isTopLevelGlobal() &&
+           var->getDeclContext()->isAsyncContext();
   }
 
   return false;
@@ -8946,9 +8946,8 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
   }
 
   if (auto *tld = dyn_cast<TopLevelCodeDecl>(dc)) {
-    ASTContext &ctx = dc->getASTContext();
-    if (ctx.LangOpts.EnableExperimentalAsyncTopLevel) {
-      if (Type mainActor = ctx.getMainActorType())
+    if (dc->isAsyncContext()) {
+      if (Type mainActor = dc->getASTContext().getMainActorType())
         return ActorIsolation::forGlobalActor(mainActor, /*unsafe=*/false);
     }
   }

lib/AST/DeclContext.cpp

@@ -1228,8 +1228,13 @@ bool DeclContext::isAsyncContext() const {
   case DeclContextKind::GenericTypeDecl:
     return false;
   case DeclContextKind::FileUnit:
+    if (const SourceFile *sf = dyn_cast<SourceFile>(this))
+      return getASTContext().LangOpts.EnableExperimentalAsyncTopLevel &&
+             sf->isAsyncTopLevelSourceFile();
+    return false;
   case DeclContextKind::TopLevelCodeDecl:
-    return getASTContext().LangOpts.EnableExperimentalAsyncTopLevel;
+    return getASTContext().LangOpts.EnableExperimentalAsyncTopLevel &&
+           getParent()->isAsyncContext();
   case DeclContextKind::AbstractClosureExpr:
     return cast<AbstractClosureExpr>(this)->isBodyAsync();
   case DeclContextKind::AbstractFunctionDecl: {

lib/AST/Module.cpp

@@ -3038,6 +3038,24 @@ SourceFile::lookupOpaqueResultType(StringRef MangledName) {
   return nullptr;
 }
 
+bool SourceFile::isAsyncTopLevelSourceFile() const {
+  return isScriptMode() &&
+         (bool)evaluateOrDefault(getASTContext().evaluator,
+                                 GetSourceFileAsyncNode{this}, ASTNode());
+}
+
+ASTNode GetSourceFileAsyncNode::evaluate(Evaluator &eval,
+                                         const SourceFile *sf) const {
+  for (Decl *d : sf->getTopLevelDecls()) {
+    TopLevelCodeDecl *tld = dyn_cast<TopLevelCodeDecl>(d);
+    if (tld && tld->getBody()) {
+      if (ASTNode asyncNode = tld->getBody()->findAsyncNode())
+        return asyncNode;
+    }
+  }
+  return ASTNode();
+}
+
 //===----------------------------------------------------------------------===//
 // SynthesizedFileUnit Implementation
 //===----------------------------------------------------------------------===//

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();
 }

lib/Sema/TypeCheckConcurrency.cpp

@@ -271,8 +271,7 @@ GlobalActorAttributeRequest::evaluate(
   } else if (auto storage = dyn_cast<AbstractStorageDecl>(decl)) {
     // Subscripts and properties are fine...
     if (auto var = dyn_cast<VarDecl>(storage)) {
-      if (var->isTopLevelGlobal() &&
-          var->getASTContext().LangOpts.EnableExperimentalAsyncTopLevel) {
+      if (var->isTopLevelGlobal() && var->getDeclContext()->isAsyncContext()) {
         var->diagnose(diag::global_actor_top_level_var)
             .highlight(globalActorAttr->getRangeWithAt());
         return None;
@@ -3538,10 +3537,8 @@ ActorIsolation ActorIsolationRequest::evaluate(
   }
 
   if (auto var = dyn_cast<VarDecl>(value)) {
-    ASTContext &ctx = var->getASTContext();
-    if (var->isTopLevelGlobal() &&
-        ctx.LangOpts.EnableExperimentalAsyncTopLevel) {
-      if (Type mainActor = ctx.getMainActorType())
+    if (var->isTopLevelGlobal() && var->getDeclContext()->isAsyncContext()) {
+      if (Type mainActor = var->getASTContext().getMainActorType())
         return inferredIsolation(
             ActorIsolation::forGlobalActor(mainActor,
                                            /*unsafe=*/var->preconcurrency()));