Adjust the referenced function type for @_unsafeSendable and @_unsafeMainActor.

When referencing a function that contains parameters with the hidden
`@_unsafeSendable` or `@_unsafeMainActor` attributes, adjust the
function type to make the types of those parameters `@Sendable` or
`@MainActor`, respectively, based on both the context the expression:

* `@Sendable` will be applied when we are in a context with strict
concurrency checking.
* `@MainActor` will be applied when we are in a context with strict
concurrency checking *or* the function is being directly applied so
that an argument is provided in the immediate expression.

The second part of the rule of `@MainActor` reflects the fact that
making the parameter `@MainActor` doesn't break existing code (because
there is a conversion to add a global actor to a function value), but
it does enable such code to synchronously use a `@MainActor`-qualified
API.

The main effect of this change is that, in a strict concurrency
context, the type of referencing an unapplied function involving
`@_unsafeSendable` or `@_unsafeMainActor` in a strict context will
make those parameters `@Sendable` or `@MainActor`, which ensures that
these constraints properly work with non-closure arguments. The former
solution only applied to closure literals, which left some holes in
Sendable checking.

Fixes rdar://77753021.

Author: DougGregor

lib/Sema/CSApply.cpp

@@ -5825,16 +5825,14 @@ ArgumentList *ExprRewriter::coerceCallArguments(
     // for things like trailing closures and args to property wrapper params.
     arg.setLabel(param.getLabel());
 
-    // Determine whether the parameter is unsafe Sendable or MainActor, and
-    // record it as such.
-    bool isUnsafeSendable = paramInfo.isUnsafeSendable(paramIdx);
-    bool isMainActor = paramInfo.isUnsafeMainActor(paramIdx) ||
-        (isUnsafeSendable && apply && isMainDispatchQueue(apply->getFn()));
+    // Determine whether the closure argument should be treated as being on
+    // the main actor, having implicit self capture, or inheriting actor
+    // context.
     bool isImplicitSelfCapture = paramInfo.isImplicitSelfCapture(paramIdx);
     bool inheritsActorContext = paramInfo.inheritsActorContext(paramIdx);
     applyContextualClosureFlags(
-        argExpr, isUnsafeSendable && contextUsesConcurrencyFeatures(dc),
-        isMainActor, isImplicitSelfCapture, inheritsActorContext);
+        argExpr, false,
+        false, isImplicitSelfCapture, inheritsActorContext);
 
     // If the types exactly match, this is easy.
     auto paramType = param.getOldType();

lib/Sema/ConstraintSystem.cpp

@@ -1255,6 +1255,24 @@ static unsigned getNumRemovedArgumentLabels(ValueDecl *decl,
   llvm_unreachable("Unhandled FunctionRefKind in switch.");
 }
 
+/// Determine the number of applications
+static unsigned getNumApplications(
+    ValueDecl *decl, bool hasAppliedSelf, FunctionRefKind functionRefKind) {
+  switch (functionRefKind) {
+  case FunctionRefKind::Unapplied:
+  case FunctionRefKind::Compound:
+    return 0 + hasAppliedSelf;
+
+  case FunctionRefKind::SingleApply:
+    return 1 + hasAppliedSelf;
+
+  case FunctionRefKind::DoubleApply:
+    return 2;
+  }
+
+  llvm_unreachable("Unhandled FunctionRefKind in switch.");
+}
+
 /// Replaces property wrapper types in the parameter list of the given function type
 /// with the wrapped-value or projected-value types (depending on argument label).
 static FunctionType *
@@ -1336,8 +1354,10 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
 
     AnyFunctionType *funcType = func->getInterfaceType()
         ->castTo<AnyFunctionType>();
-    if (!isRequirementOrWitness(locator))
-      funcType = applyGlobalActorType(funcType, func, useDC);
+    if (!isRequirementOrWitness(locator)) {
+      unsigned numApplies = getNumApplications(value, false, functionRefKind);
+      funcType = applyGlobalActorType(funcType, func, useDC, numApplies, false);
+    }
     auto openedType = openFunctionType(
         funcType, locator, replacements, func->getDeclContext());
 
@@ -1366,8 +1386,12 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
     auto numLabelsToRemove = getNumRemovedArgumentLabels(
         funcDecl, /*isCurriedInstanceReference=*/false, functionRefKind);
 
-    if (!isRequirementOrWitness(locator))
-      funcType = applyGlobalActorType(funcType, funcDecl, useDC);
+    if (!isRequirementOrWitness(locator)) {
+      unsigned numApplies = getNumApplications(
+          funcDecl, false, functionRefKind);
+      funcType = applyGlobalActorType(
+          funcType, funcDecl, useDC, numApplies, false);
+    }
 
     auto openedType = openFunctionType(funcType, locator, replacements,
                                        funcDecl->getDeclContext())
@@ -1634,6 +1658,72 @@ Type constraints::getDynamicSelfReplacementType(
       ->getMetatypeInstanceType();
 }
 
+/// Determine whether the given name is that of a DispatchQueue operation that
+/// takes a closure to be executed on the queue.
+static bool isDispatchQueueOperationName(StringRef name) {
+  return llvm::StringSwitch<bool>(name)
+    .Case("sync", true)
+    .Case("async", true)
+    .Case("asyncAndWait", true)
+    .Case("asyncAfter", true)
+    .Case("concurrentPerform", true)
+    .Default(false);
+}
+
+/// Determine whether this locator refers to a member of "DispatchQueue.main",
+/// which is a special dispatch queue that executes its work on the main actor.
+static bool isMainDispatchQueueMember(ConstraintLocator *locator) {
+  if (!locator)
+    return false;
+
+  if (locator->getPath().size() != 1 ||
+      !locator->isLastElement<LocatorPathElt::Member>())
+    return false;
+
+  auto expr = locator->getAnchor().dyn_cast<Expr *>();
+  if (!expr)
+    return false;
+
+  auto outerUnresolvedDot = dyn_cast<UnresolvedDotExpr>(expr);
+  if (!outerUnresolvedDot)
+    return false;
+
+
+  if (!isDispatchQueueOperationName(
+          outerUnresolvedDot->getName().getBaseName().userFacingName()))
+    return false;
+
+  auto innerUnresolvedDot = dyn_cast<UnresolvedDotExpr>(
+      outerUnresolvedDot->getBase());
+  if (!innerUnresolvedDot)
+    return false;
+
+  if (innerUnresolvedDot->getName().getBaseName().userFacingName() != "main")
+    return false;
+
+  auto typeExpr = dyn_cast<TypeExpr>(innerUnresolvedDot->getBase());
+  if (!typeExpr)
+    return false;
+
+  auto typeRepr = typeExpr->getTypeRepr();
+  if (!typeRepr)
+    return false;
+
+  auto identTypeRepr = dyn_cast<IdentTypeRepr>(typeRepr);
+  if (!identTypeRepr)
+    return false;
+
+  auto components = identTypeRepr->getComponentRange();
+  if (components.empty())
+    return false;
+
+  if (components.back()->getNameRef().getBaseName().userFacingName() !=
+        "DispatchQueue")
+    return false;
+
+  return true;
+}
+
 std::pair<Type, Type>
 ConstraintSystem::getTypeOfMemberReference(
     Type baseTy, ValueDecl *value, DeclContext *useDC,
@@ -1711,8 +1801,13 @@ ConstraintSystem::getTypeOfMemberReference(
     // This is the easy case.
     funcType = value->getInterfaceType()->castTo<AnyFunctionType>();
 
-    if (!isRequirementOrWitness(locator))
-      funcType = applyGlobalActorType(funcType, value, useDC);
+    if (!isRequirementOrWitness(locator)) {
+      unsigned numApplies = getNumApplications(
+          value, hasAppliedSelf, functionRefKind);
+      funcType = applyGlobalActorType(
+          funcType, value, useDC, numApplies,
+          isMainDispatchQueueMember(locator));
+    }
   } else {
     // For a property, build a type (Self) -> PropType.
     // For a subscript, build a type (Self) -> (Indices...) -> ElementType.

lib/Sema/TypeCheckConcurrency.cpp

@@ -4035,8 +4035,119 @@ NormalProtocolConformance *GetImplicitSendableRequest::evaluate(
   return formConformance(nullptr);
 }
 
+/// Apply @Sendable and/or @MainActor to the given parameter type.
+static Type applyUnsafeConcurrencyToParameterType(
+    Type type, bool sendable, bool mainActor) {
+  if (Type objectType = type->getOptionalObjectType()) {
+    return OptionalType::get(
+        applyUnsafeConcurrencyToParameterType(objectType, sendable, mainActor));
+  }
+
+  auto fnType = type->getAs<FunctionType>();
+  if (!fnType)
+    return type;
+
+  Type globalActor;
+  if (mainActor)
+    globalActor = type->getASTContext().getMainActorType();
+
+  return fnType->withExtInfo(fnType->getExtInfo()
+                               .withConcurrent(sendable)
+                               .withGlobalActor(globalActor));
+}
+
+/// Apply @_unsafeSendable and @_unsafeMainActor to the parameters of the
+/// given function.
+static AnyFunctionType *applyUnsafeConcurrencyToFunctionType(
+    AnyFunctionType *fnType, ValueDecl *funcOrEnum,
+    bool inConcurrencyContext, unsigned numApplies, bool isMainDispatchQueue) {
+  // Only functions can have @_unsafeSendable/@_unsafeMainActor parameters.
+  auto func = dyn_cast<AbstractFunctionDecl>(funcOrEnum);
+  if (!func)
+    return fnType;
+
+  AnyFunctionType *outerFnType = nullptr;
+  if (func->hasImplicitSelfDecl()) {
+    outerFnType = fnType;
+    fnType = outerFnType->getResult()->castTo<AnyFunctionType>();
+
+    if (numApplies > 0)
+      --numApplies;
+  }
+
+  SmallVector<AnyFunctionType::Param, 4> newTypeParams;
+  auto typeParams = fnType->getParams();
+  auto paramDecls = func->getParameters();
+  assert(typeParams.size() == paramDecls->size());
+  for (unsigned index : indices(typeParams)) {
+    auto param = typeParams[index];
+    auto paramDecl = (*paramDecls)[index];
+
+    // Determine whether the resulting parameter should be @Sendable or
+    // @MainActor. @Sendable occurs only in concurrency contents, while
+    // @MainActor occurs in concurrency contexts or those where we have an
+    // application.
+    bool isSendable =
+        (paramDecl->getAttrs().hasAttribute<UnsafeSendableAttr>() ||
+         func->hasKnownUnsafeSendableFunctionParams()) &&
+        inConcurrencyContext;
+    bool isMainActor =
+        (paramDecl->getAttrs().hasAttribute<UnsafeMainActorAttr>() ||
+         (isMainDispatchQueue &&
+          func->hasKnownUnsafeSendableFunctionParams())) &&
+        (inConcurrencyContext || numApplies >= 1);
+
+    if (!isSendable && !isMainActor) {
+      // If any prior parameter has changed, record this one.
+      if (!newTypeParams.empty())
+        newTypeParams.push_back(param);
+      continue;
+    }
+
+    // If this is the first parameter to have changed, copy all of the others
+    // over.
+    if (newTypeParams.empty()) {
+      newTypeParams.append(typeParams.begin(), typeParams.begin() + index);
+    }
+
+
+    // Transform the parameter type.
+    Type newParamType = applyUnsafeConcurrencyToParameterType(
+        param.getPlainType(), isSendable, isMainActor);
+    newTypeParams.push_back(param.withType(newParamType));
+  }
+
+  // If we didn't change any parameters, we're done.
+  if (newTypeParams.empty()) {
+    return outerFnType ? outerFnType : fnType;
+  }
+
+  // Rebuild the (inner) function type.
+  fnType = FunctionType::get(
+      newTypeParams, fnType->getResult(), fnType->getExtInfo());
+
+  if (!outerFnType)
+    return fnType;
+
+  // Rebuild the outer function type.
+  if (auto genericFnType = dyn_cast<GenericFunctionType>(outerFnType)) {
+    return GenericFunctionType::get(
+        genericFnType->getGenericSignature(), outerFnType->getParams(),
+        Type(fnType), outerFnType->getExtInfo());
+  }
+
+  return FunctionType::get(
+      outerFnType->getParams(), Type(fnType), outerFnType->getExtInfo());
+}
+
 AnyFunctionType *swift::applyGlobalActorType(
-    AnyFunctionType *fnType, ValueDecl *funcOrEnum, DeclContext *dc) {
+    AnyFunctionType *fnType, ValueDecl *funcOrEnum, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue) {
+  // Apply unsafe concurrency features to the given function type.
+  fnType = applyUnsafeConcurrencyToFunctionType(
+      fnType, funcOrEnum, contextUsesConcurrencyFeatures(dc), numApplies,
+      isMainDispatchQueue);
+
   Type globalActorType;
   switch (auto isolation = getActorIsolation(funcOrEnum)) {
   case ActorIsolation::ActorInstance:

lib/Sema/TypeChecker.h

@@ -1324,7 +1324,8 @@ void bindSwitchCasePatternVars(DeclContext *dc, CaseStmt *stmt);
 /// references to its function type from the given declaration context,
 /// update the given function type to include the global actor.
 AnyFunctionType *applyGlobalActorType(
-    AnyFunctionType *fnType, ValueDecl *funcOrEnum, DeclContext *dc);
+    AnyFunctionType *fnType, ValueDecl *funcOrEnum, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue);
 
 /// If \p attr was added by an access note, wraps the error in
 /// \c diag::wrap_invalid_attr_added_by_access_note and limits it as an access

test/Concurrency/dispatch_inference.swift

@@ -14,10 +14,18 @@ func testMe() {
   }
 }
 
-func testUnsafeSendableInAsync() async {
+func testUnsafeSendableInMainAsync() async {
   var x = 5
   DispatchQueue.main.async {
     x = 17 // expected-error{{mutation of captured var 'x' in concurrently-executing code}}
   }
   print(x)
 }
+
+func testUnsafeSendableInAsync(queue: DispatchQueue) async {
+  var x = 5
+  queue.async {
+    x = 17 // expected-error{{mutation of captured var 'x' in concurrently-executing code}}
+  }
+  print(x)
+}

test/Concurrency/unsafe_param_annotations.swift

@@ -0,0 +1,49 @@
+// RUN: %target-typecheck-verify-swift  -disable-availability-checking
+
+func unsafelySendableClosure(@_unsafeSendable _ closure: () -> Void) { }
+
+func unsafelyMainActorClosure(@_unsafeMainActor _ closure: () -> Void) { }
+
+func unsafelyDoEverythingClosure(@_unsafeSendable @_unsafeMainActor _ closure: () -> Void) { }
+
+struct X {
+  func unsafelyDoEverythingClosure(@_unsafeSendable @_unsafeMainActor _ closure: () -> Void) { }
+}
+
+
+func testInAsync(x: X) async {
+  let _: Int = unsafelySendableClosure // expected-error{{type '(@Sendable () -> Void) -> ()'}}
+  let _: Int = unsafelyMainActorClosure // expected-error{{type '(@MainActor () -> Void) -> ()'}}
+  let _: Int = unsafelyDoEverythingClosure // expected-error{{type '(@MainActor @Sendable () -> Void) -> ()'}}
+  let _: Int = x.unsafelyDoEverythingClosure // expected-error{{type '(@MainActor @Sendable () -> Void) -> ()'}}
+  let _: Int = X.unsafelyDoEverythingClosure // expected-error{{type '(X) -> (@MainActor @Sendable () -> Void) -> ()'}}
+  let _: Int = (X.unsafelyDoEverythingClosure)(x) // expected-error{{type '(@MainActor @Sendable () -> Void) -> ()'}}
+}
+
+func testElsewhere(x: X) {
+  let _: Int = unsafelySendableClosure // expected-error{{type '(() -> Void) -> ()'}}
+  let _: Int = unsafelyMainActorClosure // expected-error{{type '(() -> Void) -> ()'}}
+  let _: Int = unsafelyDoEverythingClosure // expected-error{{type '(() -> Void) -> ()'}}
+  let _: Int = x.unsafelyDoEverythingClosure // expected-error{{type '(() -> Void) -> ()'}}
+  let _: Int = X.unsafelyDoEverythingClosure // expected-error{{type '(X) -> (() -> Void) -> ()'}}
+  let _: Int = (X.unsafelyDoEverythingClosure)(x) // expected-error{{type '(() -> Void) -> ()'}}
+}
+
+@MainActor func onMainActor() { }
+
+func testCalls(x: X) {
+  unsafelyMainActorClosure {
+    onMainActor()
+  }
+
+  unsafelyDoEverythingClosure {
+    onMainActor()
+  }
+
+  x.unsafelyDoEverythingClosure {
+    onMainActor()
+  }
+  (X.unsafelyDoEverythingClosure)(x)( {
+    onMainActor()
+    })
+}

test/SILGen/check_executor.swift

@@ -32,8 +32,8 @@ public actor MyActor {
     }
   }
 
-  // CHECK-RAW-LABEL: sil private [ossa] @$s4test7MyActorC0A10UnsafeMainyyFSiycfU_
-  // CHECK-RAW-NOT: _checkExpectedExecutor
+  // CHECK-RAW-LABEL: sil private [ossa] @$s4test7MyActorC0A10UnsafeMainyyFSiyScMYccfU_
+  // CHECK-RAW: _checkExpectedExecutor
   // CHECK-RAW: onMainActor
   // CHECK-RAW: return
   public func testUnsafeMain() {