[ConstraintSystem] Move @sendable inference for partially applied members to adjustFunctionTypeForConcurrency

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -4371,7 +4371,7 @@ class ConstraintSystem {
   /// Wrapper over swift::adjustFunctionTypeForConcurrency that passes along
   /// the appropriate closure-type and opening extraction functions.
   FunctionType *adjustFunctionTypeForConcurrency(
-      FunctionType *fnType, ValueDecl *decl, DeclContext *dc,
+      FunctionType *fnType, Type baseType, ValueDecl *decl, DeclContext *dc,
       unsigned numApplies, bool isMainDispatchQueue,
       OpenedTypeMap &replacements, ConstraintLocatorBuilder locator);
 

lib/Sema/ConstraintSystem.cpp

@@ -1696,8 +1696,8 @@ static bool isRequirementOrWitness(const ConstraintLocatorBuilder &locator) {
 }
 
 FunctionType *ConstraintSystem::adjustFunctionTypeForConcurrency(
-    FunctionType *fnType, ValueDecl *decl, DeclContext *dc, unsigned numApplies,
-    bool isMainDispatchQueue, OpenedTypeMap &replacements,
+    FunctionType *fnType, Type baseType, ValueDecl *decl, DeclContext *dc,
+    unsigned numApplies, bool isMainDispatchQueue, OpenedTypeMap &replacements,
     ConstraintLocatorBuilder locator) {
 
   auto *adjustedTy = swift::adjustFunctionTypeForConcurrency(
@@ -1713,10 +1713,23 @@ FunctionType *ConstraintSystem::adjustFunctionTypeForConcurrency(
     if (auto *FD = dyn_cast<AbstractFunctionDecl>(decl)) {
       auto *DC = FD->getDeclContext();
       // All global functions should be @Sendable
-      if (DC->isModuleScopeContext() &&
-          !adjustedTy->getExtInfo().isSendable()) {
-        adjustedTy =
-            adjustedTy->withExtInfo(adjustedTy->getExtInfo().withSendable());
+      if (DC->isModuleScopeContext()) {
+        if (!adjustedTy->getExtInfo().isSendable()) {
+          adjustedTy =
+              adjustedTy->withExtInfo(adjustedTy->getExtInfo().withSendable());
+        }
+      } else if (isPartialApplication(getConstraintLocator(locator))) {
+        if (baseType &&
+            (baseType->is<AnyMetatypeType>() || baseType->isSendableType())) {
+          auto referenceTy = adjustedTy->getResult()->castTo<FunctionType>();
+          referenceTy =
+              referenceTy->withExtInfo(referenceTy->getExtInfo().withSendable())
+                  ->getAs<FunctionType>();
+
+          adjustedTy =
+              FunctionType::get(adjustedTy->getParams(), referenceTy,
+                                adjustedTy->getExtInfo().withSendable());
+        }
       }
     }
   }
@@ -1794,9 +1807,9 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
     auto origOpenedType = openedType;
     if (!isRequirementOrWitness(locator)) {
       unsigned numApplies = getNumApplications(value, false, functionRefKind);
-      openedType = adjustFunctionTypeForConcurrency(origOpenedType, func, useDC,
-                                                    numApplies, false,
-                                                    replacements, locator);
+      openedType = adjustFunctionTypeForConcurrency(
+          origOpenedType, /*baseType=*/Type(), func, useDC, numApplies, false,
+          replacements, locator);
     }
 
     // The reference implicitly binds 'self'.
@@ -1824,8 +1837,8 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
       unsigned numApplies = getNumApplications(
           funcDecl, false, functionRefKind);
       openedType = adjustFunctionTypeForConcurrency(
-          origOpenedType->castTo<FunctionType>(), funcDecl, useDC, numApplies,
-          false, replacements, locator);
+          origOpenedType->castTo<FunctionType>(), /*baseType=*/Type(), funcDecl,
+          useDC, numApplies, false, replacements, locator);
     }
 
     if (isForCodeCompletion() && openedType->hasError()) {
@@ -2793,20 +2806,6 @@ ConstraintSystem::getTypeOfMemberReference(
     // FIXME: Verify ExtInfo state is correct, not working by accident.
     FunctionType::ExtInfo info;
 
-    if (Context.LangOpts.hasFeature(Feature::InferSendableFromCaptures)) {
-      if (isPartialApplication(locator) &&
-          (resolvedBaseTy->is<AnyMetatypeType>() ||
-           baseOpenedTy->isSendableType())) {
-        // Add @Sendable to functions without conditional conformances
-        functionType =
-            functionType
-                ->withExtInfo(functionType->getExtInfo().withSendable())
-                ->getAs<FunctionType>();
-      }
-      // Unapplied values should always be Sendable
-      info = info.withSendable();
-    }
-
     // We'll do other adjustment later, but we need to handle parameter
     // isolation to avoid assertions.
     if (fullFunctionType->getIsolation().isParameter())
@@ -2824,11 +2823,12 @@ ConstraintSystem::getTypeOfMemberReference(
     unsigned numApplies = getNumApplications(
         value, hasAppliedSelf, functionRefKind);
     openedType = adjustFunctionTypeForConcurrency(
-        origOpenedType->castTo<FunctionType>(), value, useDC, numApplies,
-        isMainDispatchQueueMember(locator), replacements, locator);
+        origOpenedType->castTo<FunctionType>(), resolvedBaseTy, value, useDC,
+        numApplies, isMainDispatchQueueMember(locator), replacements, locator);
   } else if (auto subscript = dyn_cast<SubscriptDecl>(value)) {
     openedType = adjustFunctionTypeForConcurrency(
-        origOpenedType->castTo<FunctionType>(), subscript, useDC,
+        origOpenedType->castTo<FunctionType>(), resolvedBaseTy, subscript,
+        useDC,
         /*numApplies=*/2, /*isMainDispatchQueue=*/false, replacements, locator);
   } else if (auto var = dyn_cast<VarDecl>(value)) {
     // Adjust the function's result type, since that's the Var's actual type.
@@ -2957,7 +2957,8 @@ Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
       // FIXME: Verify ExtInfo state is correct, not working by accident.
       FunctionType::ExtInfo info;
       type = adjustFunctionTypeForConcurrency(
-          FunctionType::get(indices, elementTy, info), subscript, useDC,
+          FunctionType::get(indices, elementTy, info), overload.getBaseType(),
+          subscript, useDC,
           /*numApplies=*/1, /*isMainDispatchQueue=*/false, emptyReplacements,
           locator);
     } else if (auto var = dyn_cast<VarDecl>(decl)) {
@@ -3005,7 +3006,8 @@ Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
           decl, hasAppliedSelf, overload.getFunctionRefKind());
 
       type = adjustFunctionTypeForConcurrency(
-                 type->castTo<FunctionType>(), decl, useDC, numApplies,
+                 type->castTo<FunctionType>(), overload.getBaseType(), decl,
+                 useDC, numApplies,
                  /*isMainDispatchQueue=*/false, emptyReplacements, locator)
                  ->getResult();
     }

test/Concurrency/sendable_functions.swift

@@ -46,19 +46,3 @@ extension S: Sendable where T: Sendable {
 
 @available(SwiftStdlib 5.1, *)
 @MainActor @Sendable func globalActorFuncAsync() async { }
-
-func test_initializer_ref() {
-  func test<T>(_: @Sendable (T, T) -> Array<T>) {
-  }
-
-  // Type of `initRef` should be @Sendable but due to implicitly injected autoclosure it isn't
-  let initRef = Array.init as (Int, Int) -> Array<Int>
-
-  // FIXME: incorrect non-Sendable diagnostic is produced due to `autoclosure` wrapping `Array.init`
-  test(initRef)
-  // expected-warning@-1 {{converting non-sendable function value to '@Sendable (Int, Int) -> Array<Int>' may introduce data races}}
-
-  // FIXME: Same here
-  test(Array.init as (Int, Int) -> Array<Int>)
-  // expected-warning@-1 {{converting non-sendable function value to '@Sendable (Int, Int) -> Array<Int>' may introduce data races}}
-}

test/Concurrency/sendable_methods.swift

@@ -244,3 +244,14 @@ do {
 
   let _: () -> Void = forward(Test.fn) // Ok
 }
+
+
+func test_initializer_ref() {
+  func test<T>(_: @Sendable (T, T) -> Array<T>) {
+  }
+
+  let initRef: @Sendable (Int, Int) -> Array<Int> = Array<Int>.init // Ok
+
+  test(initRef) // Ok
+  test(Array<Int>.init) // Ok
+}