Use the closure type during type checking for establishing use of new features

When evaluating whether code is within a closure that uses concurrency
features, use the type of the closure as it's known during type checking,
so that contextual information (e.g., it's passed to a `@Sendable` or
`async` parameter of function type) can affect the result. This
corrects the definition for doing strict checking within a minimal
context for the end result of the type-check, rather than it's initial
state, catching more issues.

Fixes SR-15131 / rdar://problem/82535088.

(cherry picked from commit cc7904c)

Author: DougGregor

include/swift/Sema/ConstraintSystem.h

@@ -2217,6 +2217,13 @@ enum class SolutionApplicationToFunctionResult {
   Delay,
 };
 
+/// Retrieve the closure type from the constraint system.
+struct GetClosureType {
+  ConstraintSystem &cs;
+
+  Type operator()(const AbstractClosureExpr *expr) const;
+};
+
 /// Describes a system of constraints on type variables, the
 /// solution of which assigns concrete types to each of the type variables.
 /// Constraint systems are typically generated given an (untyped) expression.
@@ -3096,9 +3103,16 @@ class ConstraintSystem {
   }
 
   FunctionType *getClosureType(const ClosureExpr *closure) const {
+    auto result = getClosureTypeIfAvailable(closure);
+    assert(result);
+    return result;
+  }
+
+  FunctionType *getClosureTypeIfAvailable(const ClosureExpr *closure) const {
     auto result = ClosureTypes.find(closure);
-    assert(result != ClosureTypes.end());
-    return result->second;
+    if (result != ClosureTypes.end())
+      return result->second;
+    return nullptr;
   }
 
   TypeBase* getFavoredType(Expr *E) {
@@ -4116,6 +4130,12 @@ class ConstraintSystem {
          ConstraintLocatorBuilder locator,
          const OpenedTypeMap &replacements);
 
+  /// Wrapper over swift::adjustFunctionTypeForConcurrency that passes along
+  /// the appropriate closure-type extraction function.
+  AnyFunctionType *adjustFunctionTypeForConcurrency(
+    AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue);
+
   /// Retrieve the type of a reference to the given value declaration.
   ///
   /// For references to polymorphic function types, this routine "opens up"
@@ -4164,10 +4184,15 @@ class ConstraintSystem {
   ///
   /// \param getType Optional callback to extract a type for given declaration.
   static Type
-  getUnopenedTypeOfReference(VarDecl *value, Type baseType, DeclContext *UseDC,
-                             llvm::function_ref<Type(VarDecl *)> getType,
-                             ConstraintLocator *memberLocator = nullptr,
-                             bool wantInterfaceType = false);
+  getUnopenedTypeOfReference(
+      VarDecl *value, Type baseType, DeclContext *UseDC,
+      llvm::function_ref<Type(VarDecl *)> getType,
+      ConstraintLocator *memberLocator = nullptr,
+      bool wantInterfaceType = false,
+      llvm::function_ref<Type(const AbstractClosureExpr *)> getClosureType =
+        [](const AbstractClosureExpr *) {
+          return Type();
+        });
 
   /// Retrieve the type of a reference to the given value declaration,
   /// as a member with a base of the given type.

lib/Sema/CSSimplify.cpp

@@ -2171,7 +2171,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
 
   /// Whether to downgrade to a concurrency warning.
   auto isConcurrencyWarning = [&] {
-    if (contextRequiresStrictConcurrencyChecking(DC))
+    if (contextRequiresStrictConcurrencyChecking(DC, GetClosureType{*this}))
       return false;
 
     switch (kind) {

lib/Sema/ConstraintSystem.cpp

@@ -1098,6 +1098,21 @@ doesStorageProduceLValue(AbstractStorageDecl *storage, Type baseType,
           !storage->isSetterMutating());
 }
 
+Type GetClosureType::operator()(const AbstractClosureExpr *expr) const {
+  if (auto closure = dyn_cast<ClosureExpr>(expr)) {
+    // Look through type bindings, if we have them.
+    auto mutableClosure = const_cast<ClosureExpr *>(closure);
+    if (cs.hasType(mutableClosure)) {
+      return cs.getFixedTypeRecursive(
+          cs.getType(mutableClosure), /*wantRValue=*/true);
+    }
+
+    return cs.getClosureTypeIfAvailable(closure);
+  }
+
+  return Type();
+}
+
 Type ConstraintSystem::getUnopenedTypeOfReference(
     VarDecl *value, Type baseType, DeclContext *UseDC,
     ConstraintLocator *memberLocator, bool wantInterfaceType) {
@@ -1113,18 +1128,20 @@ Type ConstraintSystem::getUnopenedTypeOfReference(
 
         return wantInterfaceType ? var->getInterfaceType() : var->getType();
       },
-      memberLocator, wantInterfaceType);
+      memberLocator, wantInterfaceType, GetClosureType{*this});
 }
 
 Type ConstraintSystem::getUnopenedTypeOfReference(
     VarDecl *value, Type baseType, DeclContext *UseDC,
     llvm::function_ref<Type(VarDecl *)> getType,
-    ConstraintLocator *memberLocator, bool wantInterfaceType) {
+    ConstraintLocator *memberLocator, bool wantInterfaceType,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getClosureType) {
   Type requestedType =
       getType(value)->getWithoutSpecifierType()->getReferenceStorageReferent();
 
   // Adjust the type for concurrency.
-  requestedType = adjustVarTypeForConcurrency(requestedType, value, UseDC);
+  requestedType = adjustVarTypeForConcurrency(
+      requestedType, value, UseDC, getClosureType);
 
   // If we're dealing with contextual types, and we referenced this type from
   // a different context, map the type.
@@ -1298,6 +1315,13 @@ static bool isRequirementOrWitness(const ConstraintLocatorBuilder &locator) {
   return false;
 }
 
+AnyFunctionType *ConstraintSystem::adjustFunctionTypeForConcurrency(
+    AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue) {
+  return swift::adjustFunctionTypeForConcurrency(
+      fnType, decl, dc, numApplies, isMainDispatchQueue, GetClosureType{*this});
+}
+
 std::pair<Type, Type>
 ConstraintSystem::getTypeOfReference(ValueDecl *value,
                                      FunctionRefKind functionRefKind,
@@ -1314,7 +1338,8 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
         ->castTo<AnyFunctionType>();
     if (!isRequirementOrWitness(locator)) {
       unsigned numApplies = getNumApplications(value, false, functionRefKind);
-      funcType = adjustFunctionTypeForConcurrency(funcType, func, useDC, numApplies, false);
+      funcType = adjustFunctionTypeForConcurrency(
+          funcType, func, useDC, numApplies, false);
     }
     auto openedType = openFunctionType(
         funcType, locator, replacements, func->getDeclContext());
@@ -2085,7 +2110,8 @@ Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
       } else if (type->hasDynamicSelfType()) {
         type = withDynamicSelfResultReplaced(type, /*uncurryLevel=*/0);
       }
-      type = adjustVarTypeForConcurrency(type, var, useDC);
+      type = adjustVarTypeForConcurrency(
+          type, var, useDC, GetClosureType{*this});
     } else if (isa<AbstractFunctionDecl>(decl) || isa<EnumElementDecl>(decl)) {
       if (decl->isInstanceMember() &&
           (!overload.getBaseType() ||
@@ -2536,8 +2562,14 @@ bool ConstraintSystem::isAsynchronousContext(DeclContext *dc) {
   if (auto func = dyn_cast<AbstractFunctionDecl>(dc))
     return func->isAsyncContext();
 
-  if (auto closure = dyn_cast<AbstractClosureExpr>(dc))
-    return closure->isBodyAsync();
+  if (auto abstractClosure = dyn_cast<AbstractClosureExpr>(dc)) {
+    if (Type type = GetClosureType{*this}(abstractClosure)) {
+      if (auto fnType = type->getAs<AnyFunctionType>())
+        return fnType->isAsync();
+    }
+
+    return abstractClosure->isBodyAsync();
+  }
 
   return false;
 }

lib/Sema/TypeCheckConcurrency.cpp

@@ -615,7 +615,9 @@ static bool shouldDiagnoseExistingDataRaces(const DeclContext *dc) {
   if (dc->getParentModule()->isConcurrencyChecked())
     return true;
 
-  return contextRequiresStrictConcurrencyChecking(dc);
+  return contextRequiresStrictConcurrencyChecking(dc, [](const AbstractClosureExpr *) {
+    return Type();
+  });
 }
 
 /// Determine the default diagnostic behavior for this language mode.
@@ -1843,20 +1845,6 @@ namespace {
       }
     }
 
-    bool isInAsynchronousContext() const {
-      auto *dc = getDeclContext();
-
-      if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
-        return func->isAsyncContext();
-      }
-
-      if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
-        return closure->isBodyAsync();
-      }
-
-      return false;
-    }
-
     enum class AsyncMarkingResult {
       FoundAsync, // successfully marked an implicitly-async operation
       NotFound,  // fail: no valid implicitly-async operation was found
@@ -1917,7 +1905,7 @@ namespace {
 
         if (auto declRef = dyn_cast_or_null<DeclRefExpr>(context)) {
           if (usageEnv(declRef) == VarRefUseEnv::Read) {
-            if (!isInAsynchronousContext())
+            if (!getDeclContext()->isAsyncContext())
               return AsyncMarkingResult::SyncContext;
 
             declRef->setImplicitlyAsync(target);
@@ -1926,7 +1914,7 @@ namespace {
         } else if (auto lookupExpr = dyn_cast_or_null<LookupExpr>(context)) {
           if (usageEnv(lookupExpr) == VarRefUseEnv::Read) {
 
-            if (!isInAsynchronousContext())
+            if (!getDeclContext()->isAsyncContext())
               return AsyncMarkingResult::SyncContext;
 
             lookupExpr->setImplicitlyAsync(target);
@@ -1939,7 +1927,7 @@ namespace {
         // actor-isolated non-isolated-self calls are implicitly async
         // and thus OK.
 
-        if (!isInAsynchronousContext())
+        if (!getDeclContext()->isAsyncContext())
           return AsyncMarkingResult::SyncContext;
 
         isAsyncCall = true;
@@ -1955,7 +1943,7 @@ namespace {
           auto concDecl = memberRef->first;
           if (decl == concDecl.getDecl() && !apply->isImplicitlyAsync()) {
 
-            if (!isInAsynchronousContext())
+            if (!getDeclContext()->isAsyncContext())
               return AsyncMarkingResult::SyncContext;
 
             // then this ValueDecl appears as the called value of the ApplyExpr.
@@ -2068,7 +2056,7 @@ namespace {
         return false;
 
       // If we are not in an asynchronous context, complain.
-      if (!isInAsynchronousContext()) {
+      if (!getDeclContext()->isAsyncContext()) {
         if (auto calleeDecl = apply->getCalledValue()) {
           ctx.Diags.diagnose(
               apply->getLoc(), diag::actor_isolated_call_decl,
@@ -3562,7 +3550,9 @@ void swift::checkOverrideActorIsolation(ValueDecl *value) {
   overridden->diagnose(diag::overridden_here);
 }
 
-bool swift::contextRequiresStrictConcurrencyChecking(const DeclContext *dc) {
+bool swift::contextRequiresStrictConcurrencyChecking(
+    const DeclContext *dc,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType) {
   // If Swift >= 6, everything uses strict concurrency checking.
   if (dc->getASTContext().LangOpts.isSwiftVersionAtLeast(6))
     return true;
@@ -3574,6 +3564,12 @@ bool swift::contextRequiresStrictConcurrencyChecking(const DeclContext *dc) {
       if (auto explicitClosure = dyn_cast<ClosureExpr>(closure)) {
         if (getExplicitGlobalActor(const_cast<ClosureExpr *>(explicitClosure)))
           return true;
+
+        if (auto type = getType(closure)) {
+          if (auto fnType = type->getAs<AnyFunctionType>())
+            if (fnType->isAsync() || fnType->isSendable())
+              return true;
+        }
       }
 
       // Async and @Sendable closures use concurrency features.
@@ -4049,11 +4045,12 @@ static bool hasKnownUnsafeSendableFunctionParams(AbstractFunctionDecl *func) {
 }
 
 Type swift::adjustVarTypeForConcurrency(
-    Type type, VarDecl *var, DeclContext *dc) {
+    Type type, VarDecl *var, DeclContext *dc,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType) {
   if (!var->predatesConcurrency())
     return type;
 
-  if (contextRequiresStrictConcurrencyChecking(dc))
+  if (contextRequiresStrictConcurrencyChecking(dc, getType))
     return type;
 
   bool isLValue = false;
@@ -4172,11 +4169,12 @@ static AnyFunctionType *applyUnsafeConcurrencyToFunctionType(
 
 AnyFunctionType *swift::adjustFunctionTypeForConcurrency(
     AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
-    unsigned numApplies, bool isMainDispatchQueue) {
+    unsigned numApplies, bool isMainDispatchQueue,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType) {
   // Apply unsafe concurrency features to the given function type.
+  bool strictChecking = contextRequiresStrictConcurrencyChecking(dc, getType);
   fnType = applyUnsafeConcurrencyToFunctionType(
-      fnType, decl, contextRequiresStrictConcurrencyChecking(dc), numApplies,
-      isMainDispatchQueue);
+      fnType, decl, strictChecking, numApplies, isMainDispatchQueue);
 
   Type globalActorType;
   if (decl) {
@@ -4190,7 +4188,7 @@ AnyFunctionType *swift::adjustFunctionTypeForConcurrency(
     case ActorIsolation::GlobalActorUnsafe:
       // Only treat as global-actor-qualified within code that has adopted
       // Swift Concurrency features.
-      if (!contextRequiresStrictConcurrencyChecking(dc))
+      if (!strictChecking)
         return fnType;
 
       LLVM_FALLTHROUGH;
@@ -4232,7 +4230,10 @@ AnyFunctionType *swift::adjustFunctionTypeForConcurrency(
 }
 
 bool swift::completionContextUsesConcurrencyFeatures(const DeclContext *dc) {
-  return contextRequiresStrictConcurrencyChecking(dc);
+  return contextRequiresStrictConcurrencyChecking(
+      dc, [](const AbstractClosureExpr *) {
+        return Type();
+      });
 }
 
 AbstractFunctionDecl const *swift::isActorInitOrDeInitContext(

lib/Sema/TypeCheckConcurrency.h

@@ -214,7 +214,9 @@ void checkOverrideActorIsolation(ValueDecl *value);
 /// Determine whether the given context requires strict concurrency checking,
 /// e.g., because it uses concurrency features directly or because it's in
 /// code where strict checking has been enabled.
-bool contextRequiresStrictConcurrencyChecking(const DeclContext *dc);
+bool contextRequiresStrictConcurrencyChecking(
+    const DeclContext *dc,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType);
 
 /// Diagnose the presence of any non-sendable types when referencing a
 /// given declaration from a particular declaration context.
@@ -337,22 +339,19 @@ checkGlobalActorAttributes(
 Type getExplicitGlobalActor(ClosureExpr *closure);
 
 /// Adjust the type of the variable for concurrency.
-Type adjustVarTypeForConcurrency(Type type, VarDecl *var, DeclContext *dc);
-
-/// Adjust the function type of a function / subscript / enum case for
-/// concurrency.
-AnyFunctionType *adjustFunctionTypeForConcurrency(
-    AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
-    unsigned numApplies, bool isMainDispatchQueue);
+Type adjustVarTypeForConcurrency(
+    Type type, VarDecl *var, DeclContext *dc,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType);
 
 /// Adjust the given function type to account for concurrency-specific
 /// attributes whose affect on the type might differ based on context.
 /// This includes adjustments for unsafe parameter attributes like
 /// `@_unsafeSendable` and `@_unsafeMainActor` as well as a global actor
 /// on the declaration itself.
 AnyFunctionType *adjustFunctionTypeForConcurrency(
-    AnyFunctionType *fnType, ValueDecl *funcOrEnum, DeclContext *dc,
-    unsigned numApplies, bool isMainDispatchQueue);
+    AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue,
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType);
 
 /// Determine whether the given name is that of a DispatchQueue operation that
 /// takes a closure to be executed on the queue.

test/Concurrency/global_actor_from_ordinary_context.swift

@@ -150,7 +150,8 @@ class Sub: Super {
 
   func g2() {
     Task.detached {
-      self.f() // EXPECTED ERROR because f is on @SomeGlobalActor
+      self.f() // expected-error{{expression is 'async' but is not marked with 'await'}}
+      // expected-note@-1{{calls to instance method 'f()' from outside of its actor context are implicitly asynchronous}}
     }
   }
 }