Add implementation of _unavailableFromAsync check

The core part of the check runs as follows
 1. Are we in an async context
 2. Does the decl we're calling have the unavailableFromAsync attribute
 3. Emit a diagnostic

There are a couple challenges that muddy the implementation. First,
single-expression closures are typechecked differently than
multi-expression closures. The single-expression closure is never added
to the closure stack. We have to record it separately upon entry to
verify separately. This is necessary for `_ = { foo() }()` where `foo`
is unavailable from async, and the surrounding context is async.

The second challenge is regarding when things get typechecked. A type
must be assigned to AbstractClosureExprs to determine whether they are
an async context or not. Unfortunately, availability checking runs
before types are fully assigned in some cases. This specifically happens
when working with result builders in specific conditions. I haven't been
able to figure out how to reduce the issue further.

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 |
+  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/DiagnosticsSema.def

@@ -4722,6 +4722,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
 //------------------------------------------------------------------------------

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,8 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
   }
 
   auto where = ExportContext::forDeclSignature(D);
-  diagnoseDeclAvailability(mainFunction, attr->getRange(), nullptr,
-                           where, None);
+  diagnoseDeclAvailability(mainFunction, attr->getRange(), nullptr, where, None,
+                           /*containingClosure*/ nullptr);
 
   auto *const func = FuncDecl::createImplicit(
       context, StaticSpellingKind::KeywordStatic,
@@ -5672,6 +5674,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

@@ -2767,6 +2767,7 @@ class ExprAvailabilityWalker : public ASTWalker {
   ASTContext &Context;
   MemberAccessContext AccessContext = MemberAccessContext::Getter;
   SmallVector<const Expr *, 16> ExprStack;
+  SmallVector<ClosureExpr *, 4> singleExprClosureStack;
   const ExportContext &Where;
 
 public:
@@ -2847,24 +2848,47 @@ class ExprAvailabilityWalker : public ASTWalker {
                                E->getLoc(), Where);
     }
 
+    // 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.
+    // As a result, we never see single-statement closures as the decl context
+    // in out 'where' here. By keeping a stack of closures, we can see if
+    // we're inside of one of these closures and go from there.
+    if (ClosureExpr *closure = dyn_cast<ClosureExpr>(E)) {
+      if (closure->hasSingleExpressionBody())
+        singleExprClosureStack.push_back(closure);
+    }
+
     return visitChildren();
   }
 
   Expr *walkToExprPost(Expr *E) override {
     assert(ExprStack.back() == E);
     ExprStack.pop_back();
 
+    if (ClosureExpr *closure = dyn_cast<ClosureExpr>(E)) {
+      if (closure->hasSingleExpressionBody()) {
+        assert(closure == singleExprClosureStack.back() &&
+               "Popping wrong closure");
+        singleExprClosureStack.pop_back();
+      }
+    }
+
     return E;
   }
 
   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);
+    DeclContext *dc = singleExprClosureStack.empty()
+                          ? Where.getDeclContext()
+                          : singleExprClosureStack.back();
+    diagnoseStmtAvailability(S, dc, /*walkRecursively=*/true);
     return std::make_pair(false, S);
   }
 
@@ -3043,8 +3067,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, nullptr))
       return;
   }
 };
@@ -3067,7 +3091,9 @@ bool ExprAvailabilityWalker::diagnoseDeclRefAvailability(
       return true;
   }
 
-  diagnoseDeclAvailability(D, R, call, Where, Flags);
+  const ClosureExpr *closure =
+      singleExprClosureStack.empty() ? nullptr : singleExprClosureStack.back();
+  diagnoseDeclAvailability(D, R, call, Where, Flags, closure);
 
   if (R.isValid()) {
     if (diagnoseSubstitutionMapAvailability(R.Start, declRef.getSubstitutions(),
@@ -3079,12 +3105,71 @@ 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,
+                                 const ClosureExpr *singleExprClosure) {
+  // 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're directly in a sync closure, then we are not in an async context.
+  // If we're directly in a sync closure, or we're not in a closure and the
+  // outer context is not async, we are in a sync context.
+  const bool inSingleClosure = singleExprClosure;
+  const bool inSyncClosure =
+      inSingleClosure && !singleExprClosure->isAsyncContext();
+  const bool inSyncContext =
+      !inSingleClosure && !Where.getDeclContext()->isAsyncContext();
+  if (inSyncClosure || inSyncContext)
+    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,
                                      const Expr *call,
                                      const ExportContext &Where,
-                                     DeclAvailabilityFlags Flags) {
+                                     DeclAvailabilityFlags Flags,
+                                     const ClosureExpr *singleExprClosure) {
   assert(!Where.isImplicit());
 
   // Generic parameters are always available.
@@ -3112,6 +3197,9 @@ bool swift::diagnoseDeclAvailability(const ValueDecl *D, SourceRange R,
   if (diagnoseExplicitUnavailability(D, R, Where, call, Flags))
     return true;
 
+  if (diagnoseDeclUnavailableFromAsync(D, R, call, Where, singleExprClosure))
+    return true;
+
   // Make sure not to diagnose an accessor's deprecation if we already
   // complained about the property/subscript.
   bool isAccessorWithDeprecatedStorage =
@@ -3360,7 +3448,8 @@ class TypeReprAvailabilityWalker : public ASTWalker {
   bool checkComponentIdentTypeRepr(ComponentIdentTypeRepr *ITR) {
     if (auto *typeDecl = ITR->getBoundDecl()) {
       auto range = ITR->getNameLoc().getSourceRange();
-      if (diagnoseDeclAvailability(typeDecl, range, nullptr, where, flags))
+      if (diagnoseDeclAvailability(typeDecl, range, nullptr, where, flags,
+                                   nullptr))
         return true;
     }
 

lib/Sema/TypeCheckAvailability.h

@@ -27,6 +27,7 @@ namespace swift {
   class ApplyExpr;
   class AvailableAttr;
   class Expr;
+  class ClosureExpr;
   class InFlightDiagnostic;
   class Decl;
   class ProtocolConformanceRef;
@@ -232,7 +233,8 @@ diagnoseSubstitutionMapAvailability(SourceLoc loc,
 /// was emitted.
 bool diagnoseDeclAvailability(const ValueDecl *D, SourceRange R,
                               const Expr *call, const ExportContext &where,
-                              DeclAvailabilityFlags flags = None);
+                              DeclAvailabilityFlags flags = None,
+                              const ClosureExpr *singleExprClosure = nullptr);
 
 void diagnoseUnavailableOverride(ValueDecl *override,
                                  const ValueDecl *base,

lib/Sema/TypeCheckDeclOverride.cpp

@@ -1559,6 +1559,7 @@ namespace  {
     UNINTERESTING_ATTR(InheritActorContext)
     UNINTERESTING_ATTR(Isolated)
     UNINTERESTING_ATTR(NoImplicitCopy)
+    UNINTERESTING_ATTR(UnavailableFromAsync)
 
     UNINTERESTING_ATTR(TypeSequence)
     UNINTERESTING_ATTR(CompileTimeConst)