Merge pull request #40149 from etcwilde/ewilde/concurrency/underscored-unavailablefromasync

Add `_unavailableFromAsync` attribute

Author: etcwilde

docs/ReferenceGuides/UnderscoredAttributes.md

@@ -634,3 +634,8 @@ calls a runtime function which allocates memory or locks, respectively.
 The `@_noLocks` attribute implies `@_noAllocation` because a memory allocation
 also locks.
 
+## `@_unavailableFromAsync`
+
+Marks a synchronous API as being unavailable from asynchronous contexts. Direct
+usage of annotated API from asynchronous contexts will result in a warning from
+the compiler.

include/swift/AST/Attr.def

@@ -699,6 +699,12 @@ CONTEXTUAL_SIMPLE_DECL_ATTR(_const, CompileTimeConst,
   ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIStableToRemove,
   126)
 
+SIMPLE_DECL_ATTR(_unavailableFromAsync, UnavailableFromAsync,
+  OnFunc | OnConstructor | UserInaccessible |
+  ABIStableToAdd | ABIStableToRemove |
+  APIBreakingToAdd | APIStableToRemove,
+  127)
+
 // If you're adding a new underscored attribute here, please document it in
 // docs/ReferenceGuides/UnderscoredAttributes.md.
 

include/swift/AST/DeclContext.h

@@ -296,6 +296,9 @@ class alignas(1 << DeclContextAlignInBits) DeclContext
   /// Returns the kind of context this is.
   DeclContextKind getContextKind() const;
 
+  /// Returns whether this context asynchronous
+  bool isAsyncContext() const;
+
   /// Returns whether this context has value semantics.
   bool hasValueSemantics() const;
 

include/swift/AST/DiagnosticsSema.def

@@ -4705,6 +4705,15 @@ ERROR(actor_isolation_superclass_mismatch,none,
       "%0 class %1 has different actor isolation from %2 superclass %3",
       (ActorIsolation, DeclName, ActorIsolation, DeclName))
 
+ERROR(async_decl_must_be_available_from_async,none,
+        "asynchronous %0 must be available from asynchronous contexts",
+        (DescriptiveDeclKind))
+ERROR(async_named_decl_must_be_available_from_async,none,
+        "asynchronous %0 %1 must be available from asynchronous contexts",
+        (DescriptiveDeclKind, DeclName))
+ERROR(async_unavailable_decl,none,
+      "%0 %1 is unavailable from asynchronous contexts", (DescriptiveDeclKind, DeclBaseName))
+
 //------------------------------------------------------------------------------
 // MARK: Type Check Types
 //------------------------------------------------------------------------------

include/swift/AST/Expr.h

@@ -3601,6 +3601,13 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// Only valid when \c hasSingleExpressionBody() is true.
   Expr *getSingleExpressionBody() const;
 
+  /// Whether this closure has a body
+  bool hasBody() const;
+
+  /// Returns the body of closures that have a body
+  /// returns nullptr if the closure doesn't have a body
+  BraceStmt *getBody() const;
+
   ClosureActorIsolation getActorIsolation() const { return actorIsolation; }
 
   void setActorIsolation(ClosureActorIsolation actorIsolation) {

lib/AST/DeclContext.cpp

@@ -1208,6 +1208,32 @@ bool DeclContext::isClassConstrainedProtocolExtension() const {
   return false;
 }
 
+bool DeclContext::isAsyncContext() const {
+  switch (getContextKind()) {
+  case DeclContextKind::Initializer:
+  case DeclContextKind::TopLevelCodeDecl:
+  case DeclContextKind::EnumElementDecl:
+  case DeclContextKind::ExtensionDecl:
+  case DeclContextKind::SerializedLocal:
+  case DeclContextKind::Module:
+  case DeclContextKind::FileUnit:
+  case DeclContextKind::GenericTypeDecl:
+    return false;
+  case DeclContextKind::AbstractClosureExpr:
+    return cast<AbstractClosureExpr>(this)->isBodyAsync();
+  case DeclContextKind::AbstractFunctionDecl: {
+    const AbstractFunctionDecl *function = cast<AbstractFunctionDecl>(this);
+    return function->hasAsync();
+  }
+  case DeclContextKind::SubscriptDecl: {
+    AccessorDecl *getter =
+        cast<SubscriptDecl>(this)->getAccessor(AccessorKind::Get);
+    return getter != nullptr && getter->hasAsync();
+  }
+  }
+  llvm_unreachable("Unhandled DeclContextKind switch");
+}
+
 SourceLoc swift::extractNearestSourceLoc(const DeclContext *dc) {
   switch (dc->getContextKind()) {
   case DeclContextKind::Module:

lib/AST/Expr.cpp

@@ -1674,6 +1674,26 @@ void AbstractClosureExpr::setParameterList(ParameterList *P) {
     P->setDeclContextOfParamDecls(this);
 }
 
+bool AbstractClosureExpr::hasBody() const {
+  switch (getKind()) {
+    case ExprKind::Closure:
+    case ExprKind::AutoClosure:
+      return true;
+    default:
+      return false;
+  }
+}
+
+BraceStmt * AbstractClosureExpr::getBody() const {
+  if (!hasBody())
+    return nullptr;
+  if (const AutoClosureExpr *autocls = dyn_cast<AutoClosureExpr>(this))
+    return autocls->getBody();
+  if (const ClosureExpr *cls = dyn_cast<ClosureExpr>(this))
+    return cls->getBody();
+  llvm_unreachable("Unknown closure expression");
+}
+
 Type AbstractClosureExpr::getResultType(
     llvm::function_ref<Type(Expr *)> getType) const {
   auto *E = const_cast<AbstractClosureExpr *>(this);

lib/Sema/TypeCheckAttr.cpp

@@ -270,6 +270,8 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   void visitNonisolatedAttr(NonisolatedAttr *attr);
 
   void visitNoImplicitCopyAttr(NoImplicitCopyAttr *attr);
+
+  void visitUnavailableFromAsyncAttr(UnavailableFromAsyncAttr *attr);
 };
 
 } // end anonymous namespace
@@ -2081,8 +2083,7 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
   }
 
   auto where = ExportContext::forDeclSignature(D);
-  diagnoseDeclAvailability(mainFunction, attr->getRange(), nullptr,
-                           where, None);
+  diagnoseDeclAvailability(mainFunction, attr->getRange(), nullptr, where, None);
 
   auto *const func = FuncDecl::createImplicit(
       context, StaticSpellingKind::KeywordStatic,
@@ -5678,6 +5679,23 @@ void AttributeChecker::visitReasyncAttr(ReasyncAttr *attr) {
   attr->setInvalid();
 }
 
+void AttributeChecker::visitUnavailableFromAsyncAttr(
+    UnavailableFromAsyncAttr *attr) {
+  if (DeclContext *dc = dyn_cast<DeclContext>(D)) {
+    if (dc->isAsyncContext()) {
+      if (ValueDecl *vd = dyn_cast<ValueDecl>(D)) {
+        D->getASTContext().Diags.diagnose(
+            D->getLoc(), diag::async_named_decl_must_be_available_from_async,
+            D->getDescriptiveKind(), vd->getName());
+      } else {
+        D->getASTContext().Diags.diagnose(
+            D->getLoc(), diag::async_decl_must_be_available_from_async,
+            D->getDescriptiveKind());
+      }
+    }
+  }
+}
+
 namespace {
 
 class ClosureAttributeChecker

lib/Sema/TypeCheckAvailability.cpp

@@ -2842,6 +2842,22 @@ class ExprAvailabilityWalker : public ASTWalker {
                                E->getLoc(), Where);
     }
 
+    if (AbstractClosureExpr *closure = dyn_cast<AbstractClosureExpr>(E)) {
+      // Multi-statement closures are collected by ExprWalker::rewriteFunction
+      // and checked by ExprWalker::processDelayed in CSApply.cpp.
+      // Single-statement closures only have the attributes checked
+      // by TypeChecker::checkClosureAttributes in that rewriteFunction.
+      // Multi-statement closures will be checked explicitly later (as the decl
+      // context in the Where). Single-expression closures will not be
+      // revisited, and are not automatically set as the context of the 'where'.
+      // Don't double-check multi-statement closures, but do check
+      // single-statement closures, setting the closure as the decl context.
+      if (closure->hasSingleExpressionBody()) {
+        walkAbstractClosure(closure);
+        return skipChildren();
+      }
+    }
+
     return visitChildren();
   }
 
@@ -2853,13 +2869,13 @@ class ExprAvailabilityWalker : public ASTWalker {
   }
 
   std::pair<bool, Stmt *> walkToStmtPre(Stmt *S) override {
+
     // We end up here when checking the output of the result builder transform,
     // which includes closures that are not "separately typechecked" and yet
     // contain statements and declarations. We need to walk them recursively,
     // since these availability for these statements is not diagnosed from
     // typeCheckStmt() as usual.
-    diagnoseStmtAvailability(S, Where.getDeclContext(),
-                             /*walkRecursively=*/true);
+    diagnoseStmtAvailability(S, Where.getDeclContext(), /*walkRecursively=*/true);
     return std::make_pair(false, S);
   }
 
@@ -2978,6 +2994,21 @@ class ExprAvailabilityWalker : public ASTWalker {
     walkInContext(E, E->getSubExpr(), MemberAccessContext::InOut);
   }
 
+  /// Walk an abstract closure expression, checking for availability
+  void walkAbstractClosure(AbstractClosureExpr *closure) {
+    // Do the walk with the closure set as the decl context of the 'where'
+    auto where = ExportContext::forFunctionBody(closure, closure->getStartLoc());
+    if (where.isImplicit())
+      return;
+    ExprAvailabilityWalker walker(where);
+
+    // Manually dive into the body
+    closure->getBody()->walk(walker);
+
+    return;
+  }
+
+
   /// Walk the given expression in the member access context.
   void walkInContext(Expr *baseExpr, Expr *E,
                      MemberAccessContext AccessContext) {
@@ -3038,8 +3069,8 @@ class ExprAvailabilityWalker : public ASTWalker {
 
     Flags &= DeclAvailabilityFlag::ForInout;
     Flags |= DeclAvailabilityFlag::ContinueOnPotentialUnavailability;
-    if (diagnoseDeclAvailability(D, ReferenceRange, /*call*/nullptr,
-                                 Where, Flags))
+    if (diagnoseDeclAvailability(D, ReferenceRange, /*call*/ nullptr, Where,
+                                 Flags))
       return;
   }
 };
@@ -3074,6 +3105,56 @@ bool ExprAvailabilityWalker::diagnoseDeclRefAvailability(
   return false;
 }
 
+/// Diagnose uses of API annotated '@unavailableFromAsync' when used from
+/// asynchronous contexts.
+/// Returns true if a diagnostic was emitted, false otherwise.
+static bool
+diagnoseDeclUnavailableFromAsync(const ValueDecl *D, SourceRange R,
+                                 const Expr *call, const ExportContext &Where) {
+  // FIXME: I don't think this is right, but I don't understand the issue well
+  //        enough to fix it properly. If the decl context is an abstract
+  //        closure, we need it to have a type assigned to it before we can
+  //        determine whether it is an asynchronous context. It will crash
+  //        when we go to check without one. In TypeChecker::typeCheckExpression
+  //        (TypeCheckConstraints.cpp:403), we apply a solution before calling
+  //        `performSyntacticDiagnosticsForTarget`, which eventually calls
+  //        down to this function. Under most circumstances, the context that
+  //        we're in is typechecked at that point and has a type assigned.
+  //        When working with specific result builders, the solution applied
+  //        results in an expression with an unset type. In these cases, the
+  //        application makes its way into `ConstraintSystem::applySolution` for
+  //        closures (CSClosure.cpp:1356). The type is computed, but is
+  //        squirreled away in the constrain system to be applied once the
+  //        checks (including this one) approve of the decls within the decl
+  //        context before applying the type to the expression. It might be
+  //        possible to drive the constraint solver through the availability
+  //        checker and into us so that we can ask for it, but that feels wrong
+  //        too.
+  //        This behavior is demonstrated by the first use of the `tuplify`
+  //        function in `testExistingPatternsInCaseStatements` in
+  //        `test/Constraints/result_builder.swift`.
+  const AbstractClosureExpr *declCtxAsExpr =
+      dyn_cast<AbstractClosureExpr>(Where.getDeclContext());
+  if (declCtxAsExpr && !declCtxAsExpr->getType()) {
+    return false;
+  }
+
+  // If we are in a synchronous context, don't check it
+  if (!Where.getDeclContext()->isAsyncContext())
+    return false;
+  if (!D->getAttrs().hasAttribute<UnavailableFromAsyncAttr>())
+    return false;
+
+  ASTContext &ctx = Where.getDeclContext()->getASTContext();
+  SourceLoc diagLoc = call ? call->getLoc() : R.Start;
+  ctx.Diags
+      .diagnose(diagLoc, diag::async_unavailable_decl, D->getDescriptiveKind(),
+                D->getBaseName())
+      .warnUntilSwiftVersion(6);
+  D->diagnose(diag::decl_declared_here, D->getName());
+  return true;
+}
+
 /// Diagnose uses of unavailable declarations. Returns true if a diagnostic
 /// was emitted.
 bool swift::diagnoseDeclAvailability(const ValueDecl *D, SourceRange R,
@@ -3107,6 +3188,9 @@ bool swift::diagnoseDeclAvailability(const ValueDecl *D, SourceRange R,
   if (diagnoseExplicitUnavailability(D, R, Where, call, Flags))
     return true;
 
+  if (diagnoseDeclUnavailableFromAsync(D, R, call, Where))
+    return true;
+
   // Make sure not to diagnose an accessor's deprecation if we already
   // complained about the property/subscript.
   bool isAccessorWithDeprecatedStorage =

lib/Sema/TypeCheckAvailability.h

@@ -27,6 +27,7 @@ namespace swift {
   class ApplyExpr;
   class AvailableAttr;
   class Expr;
+  class ClosureExpr;
   class InFlightDiagnostic;
   class Decl;
   class ProtocolConformanceRef;