Merge pull request #76136 from xedin/rdar-131321053

 [ConstraintSystem] InferSendableFromCaptures: Mark unapplied operator references as `@Sendable`

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -5605,8 +5605,6 @@ class ConstraintSystem {
     return range.isValid() ? range : std::optional<SourceRange>();
   }
 
-  bool isPartialApplication(ConstraintLocator *locator);
-
   bool isTooComplex(size_t solutionMemory) {
     if (isAlreadyTooComplex.first)
       return true;
@@ -6554,6 +6552,11 @@ Type getPatternTypeOfSingleUnlabeledPackExpansionTuple(Type type);
 /// Check whether this is a reference to one of the special result builder
 /// methods prefixed with `build*` i.e. `buildBlock`, `buildExpression` etc.
 bool isResultBuilderMethodReference(ASTContext &, UnresolvedDotExpr *);
+
+/// Determine the number of applications applied to the given overload.
+unsigned getNumApplications(ValueDecl *decl, bool hasAppliedSelf,
+                            FunctionRefKind functionRefKind);
+
 } // end namespace constraints
 
 template<typename ...Args>

lib/Sema/CSGen.cpp

@@ -4607,13 +4607,14 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
         new (ctx) DeclRefExpr(makeIteratorVar, DeclNameLoc(stmt->getForLoc()),
                               /*Implicit=*/true),
         nextId, labels);
-    nextRef->setFunctionRefKind(FunctionRefKind::Compound);
+    nextRef->setFunctionRefKind(FunctionRefKind::SingleApply);
 
     ArgumentList *nextArgs;
     if (nextFn && nextFn->getParameters()->size() == 1) {
       auto isolationArg =
         new (ctx) CurrentContextIsolationExpr(stmt->getForLoc(), Type());
-      nextArgs = ArgumentList::forImplicitUnlabeled(ctx, { isolationArg });
+      nextArgs = ArgumentList::createImplicit(
+          ctx, {Argument(SourceLoc(), ctx.Id_isolation, isolationArg)});
     } else {
       nextArgs = ArgumentList::createImplicit(ctx, {});
     }

lib/Sema/CSSimplify.cpp

@@ -10293,15 +10293,16 @@ performMemberLookup(ConstraintKind constraintKind, DeclNameRef memberName,
       if (!Context.getProtocol(KnownProtocolKind::Sendable))
         return false;
 
-      // Static members are always sendable because they only capture
-      // metatypes which are Sendable.
-      if (baseObjTy->is<AnyMetatypeType>())
-        return false;
+      return llvm::any_of(lookup, [&](const auto &result) {
+        auto decl = result.getValueDecl();
+        if (!isa_and_nonnull<FuncDecl>(decl))
+          return false;
 
-      return isPartialApplication(memberLocator) &&
-             llvm::any_of(lookup, [&](const auto &result) {
-               return isa_and_nonnull<FuncDecl>(result.getValueDecl());
-             });
+        auto hasAppliedSelf = decl->hasCurriedSelf() &&
+                              doesMemberRefApplyCurriedSelf(baseObjTy, decl);
+        return getNumApplications(decl, hasAppliedSelf, functionRefKind) <
+               decl->getNumCurryLevels();
+      });
     };
 
     if (shouldCheckSendabilityOfBase()) {
@@ -10945,7 +10946,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
 
         return simplifyValueWitnessConstraint(
             ConstraintKind::ValueWitness, baseTy, makeIterator, memberTy, useDC,
-            FunctionRefKind::Compound, flags, locator);
+            FunctionRefKind::SingleApply, flags, locator);
       }
 
       // Handle `next` reference.
@@ -10963,7 +10964,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
 
         return simplifyValueWitnessConstraint(
             ConstraintKind::ValueWitness, baseTy, next, memberTy, useDC,
-            FunctionRefKind::Compound, flags, locator);
+            FunctionRefKind::SingleApply, flags, locator);
       }
     }
   }

lib/Sema/ConstraintSystem.cpp

@@ -1655,8 +1655,8 @@ static unsigned getNumRemovedArgumentLabels(ValueDecl *decl,
 }
 
 /// Determine the number of applications
-static unsigned getNumApplications(
-    ValueDecl *decl, bool hasAppliedSelf, FunctionRefKind functionRefKind) {
+unsigned constraints::getNumApplications(ValueDecl *decl, bool hasAppliedSelf,
+                                         FunctionRefKind functionRefKind) {
   switch (functionRefKind) {
   case FunctionRefKind::Unapplied:
   case FunctionRefKind::Compound:
@@ -1774,18 +1774,36 @@ FunctionType *ConstraintSystem::adjustFunctionTypeForConcurrency(
           adjustedTy =
               adjustedTy->withExtInfo(adjustedTy->getExtInfo().withSendable());
         }
-      } else if (isPartialApplication(getConstraintLocator(locator))) {
-        if (baseType &&
-            (baseType->is<AnyMetatypeType>() || baseType->isSendableType())) {
-          auto referenceTy = adjustedTy->getResult()->castTo<FunctionType>();
+      } else if (numApplies < decl->getNumCurryLevels() &&
+                 decl->hasCurriedSelf() ) {
+        auto shouldMarkMemberTypeSendable = [&]() {
+          // Static member types are @Sendable on both levels because
+          // they only capture a metatype "base" that is always Sendable.
+          // For example, `(S.Type) -> () -> Void`.
+          if (!decl->isInstanceMember())
+            return true;
+
+          // For instance members we need to check whether instance type
+          // is Sendable because @Sendable function values cannot capture
+          // non-Sendable values (base instance type in this case).
+          // For example, `(C) -> () -> Void` where `C` should be Sendable
+          // for the inner function type to be Sendable as well.
+          return baseType &&
+                 baseType->getMetatypeInstanceType()->isSendableType();
+        };
+
+        auto referenceTy = adjustedTy->getResult()->castTo<FunctionType>();
+        if (shouldMarkMemberTypeSendable()) {
           referenceTy =
-              referenceTy->withExtInfo(referenceTy->getExtInfo().withSendable())
+              referenceTy
+                  ->withExtInfo(referenceTy->getExtInfo().withSendable())
                   ->getAs<FunctionType>();
-
-          adjustedTy =
-              FunctionType::get(adjustedTy->getParams(), referenceTy,
-                                adjustedTy->getExtInfo().withSendable());
         }
+
+        // @Sendable since fully uncurried type doesn't capture anything.
+        adjustedTy =
+            FunctionType::get(adjustedTy->getParams(), referenceTy,
+                              adjustedTy->getExtInfo().withSendable());
       }
     }
   }
@@ -2636,26 +2654,6 @@ static unsigned getApplicationLevel(ConstraintSystem &CS, Type baseTy,
   return level;
 }
 
-bool ConstraintSystem::isPartialApplication(ConstraintLocator *locator) {
-  // If this is a compiler synthesized implicit conversion, let's skip
-  // the check because the base of `UDE` is not the base of the injected
-  // initializer.
-  if (locator->isLastElement<LocatorPathElt::ConstructorMember>() &&
-      locator->findFirst<LocatorPathElt::ImplicitConversion>())
-    return false;
-
-  auto *UDE = getAsExpr<UnresolvedDotExpr>(locator->getAnchor());
-  if (UDE == nullptr)
-    return false;
-
-  auto baseTy =
-      simplifyType(getType(UDE->getBase()))->getWithoutSpecifierType();
-  auto level = getApplicationLevel(*this, baseTy, UDE);
-  // Static members have base applied implicitly which means that their
-  // application level is lower.
-  return level < (baseTy->is<MetatypeType>() ? 1 : 2);
-}
-
 bool IsInLeftHandSideOfAssignment::operator()(Expr *expr) const {
   // Walk up the parent tree.
   auto parent = cs.getParentExpr(expr);

test/Concurrency/sendable_methods.swift

@@ -286,3 +286,36 @@ func testPatternMatch(ge: [GenericE<Int>]) {
     _ = a
   }
 }
+
+// rdar://131321053 - cannot pass an operator to parameter that expectes a @Sendable type
+do {
+  func test(_: @Sendable (Int, Int) -> Bool) {
+  }
+
+  test(<) // Ok
+}
+
+// Partially applied instance method
+do {
+  struct S {
+    func foo() {}
+  }
+
+  func bar(_ x: @Sendable () -> Void) {}
+
+  let fn = S.foo(S())
+  bar(fn) // Ok
+
+  let _: @Sendable (S) -> @Sendable () -> Void = S.foo // Ok
+
+  let classFn = NonSendable.f(NonSendable())
+  bar(classFn) // expected-warning {{converting non-sendable function value to '@Sendable () -> Void' may introduce data races}}
+
+  let _: @Sendable (NonSendable) -> () -> Void = NonSendable.f // Ok
+
+  class Test {
+    static func staticFn() {}
+  }
+
+  bar(Test.staticFn) // Ok
+}