[6.1] Fix some type issues around handling closures in SendNonSendable

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -434,17 +434,21 @@ class SILIsolationInfo {
     return {};
   }
 
+  static bool isNonSendableType(SILType type, SILFunction *fn) {
+    return isNonSendableType(type.getASTType(), fn);
+  }
+
+  static bool isNonSendableType(SILValue value) {
+    return isNonSendableType(value->getType(), value->getFunction());
+  }
+
   /// A helper that is used to ensure that we treat certain builtin values as
   /// non-Sendable that the AST level otherwise thinks are non-Sendable.
   ///
   /// E.x.: Builtin.RawPointer and Builtin.NativeObject
   ///
   /// TODO: Fix the type checker.
-  static bool isNonSendableType(SILType type, SILFunction *fn);
-
-  static bool isNonSendableType(SILValue value) {
-    return isNonSendableType(value->getType(), value->getFunction());
-  }
+  static bool isNonSendableType(CanType type, SILFunction *fn);
 
   bool hasSameIsolation(ActorIsolation actorIsolation) const;
 

lib/SILOptimizer/Mandatory/SendNonSendable.cpp

@@ -1661,18 +1661,31 @@ bool SentNeverSendableDiagnosticInferrer::initForSendingPartialApply(
   }
 
   // Ok, we are not tracking an actual isolated value or we do not capture the
-  // isolated value directly... we need to be smarter here. First lets gather up
-  // all non-Sendable values captured by the closure.
+  // isolated value directly, we instead just emit a note on all the potential
+  // non-Sendable captures.
+
+  // First grab the calleeFunction of the closure so we can properly map our
+  // captured types into the closure's context so that we can determine
+  // non-Sendability.
+  auto *calleeFunction = pai->getCalleeFunction();
+  if (!calleeFunction) {
+    // We should always be able to find a callee function for a partial_apply
+    // created from a closure expr. If we don't, we want to flag this as an
+    // error to be reported.
+    diagnosticEmitter.emitUnknownPatternError();
+    return true;
+  }
+
   SmallVector<CapturedValue, 8> nonSendableCaptures;
   for (auto capture : ce->getCaptureInfo().getCaptures()) {
     auto *decl = capture.getDecl();
     auto type = decl->getInterfaceType()->getCanonicalType();
-    auto silType = SILType::getPrimitiveObjectType(type);
-    if (!SILIsolationInfo::isNonSendableType(silType, pai->getFunction()))
+    type = calleeFunction->mapTypeIntoContext(type)->getCanonicalType();
+    if (!SILIsolationInfo::isNonSendableType(type, calleeFunction))
       continue;
 
     auto *fromDC = decl->getInnermostDeclContext();
-    auto *nom = silType.getNominalOrBoundGenericNominal();
+    auto *nom = type.getNominalOrBoundGenericNominal();
     if (nom && fromDC) {
       if (auto diagnosticBehavior =
               getConcurrencyDiagnosticBehaviorLimit(nom, fromDC)) {

lib/SILOptimizer/Utils/SILIsolationInfo.cpp

@@ -1215,18 +1215,18 @@ void SILIsolationInfo::printForOneLineLogging(llvm::raw_ostream &os) const {
 //
 // NOTE: We special case RawPointer and NativeObject to ensure they are
 // treated as non-Sendable and strict checking is applied to it.
-bool SILIsolationInfo::isNonSendableType(SILType type, SILFunction *fn) {
+bool SILIsolationInfo::isNonSendableType(CanType type, SILFunction *fn) {
   // Treat Builtin.NativeObject, Builtin.RawPointer, and Builtin.BridgeObject as
   // non-Sendable.
-  if (type.getASTType()->is<BuiltinNativeObjectType>() ||
-      type.getASTType()->is<BuiltinRawPointerType>() ||
-      type.getASTType()->is<BuiltinBridgeObjectType>()) {
+  if (type->is<BuiltinNativeObjectType>() ||
+      type->is<BuiltinRawPointerType>() ||
+      type->is<BuiltinBridgeObjectType>()) {
     return true;
   }
 
   // Treat Builtin.SILToken as Sendable. It cannot escape from the current
   // function. We should change isSendable to hardwire this.
-  if (type.getASTType()->is<SILTokenType>()) {
+  if (type->is<SILTokenType>()) {
     return false;
   }
 
@@ -1237,12 +1237,17 @@ bool SILIsolationInfo::isNonSendableType(SILType type, SILFunction *fn) {
   // getConcurrencyDiagnosticBehavior could cause us to prevent a
   // "preconcurrency" unneeded diagnostic when just using Sendable values. We
   // only want to trigger that if we analyze a non-Sendable type.
-  if (type.isSendable(fn))
+  if (type->isSendableType())
     return false;
 
   // Grab out behavior. If it is none, then we have a type that we want to treat
   // as non-Sendable.
-  auto behavior = type.getConcurrencyDiagnosticBehavior(fn);
+  auto declRef = fn->getDeclRef();
+  if (!declRef)
+    return true;
+
+  auto *fromDC = declRef.getInnermostDeclContext();
+  auto behavior = type->getConcurrencyDiagnosticBehaviorLimit(fromDC);
   if (!behavior)
     return true;
 

test/Concurrency/transfernonsendable.swift

@@ -1906,3 +1906,34 @@ func offByOneWithImplicitPartialApply() {
       }
   }
 }
+
+protocol Progress { // expected-note {{}}
+  associatedtype AssocType
+  var x: AssocType { get }
+  func checkCancellation() throws
+}
+
+// We used to crash here since the closure diagnostic would not map z's type
+// into the current context.
+func testCaptureDiagnosticMapsTypeIntoContext<T : Progress>(_ x: NonSendableKlass, y: T) async throws {
+  let z = y.x
+  await withTaskGroup(of: Void.self) { group in
+    group.addTask { // expected-tns-warning {{passing closure as a 'sending' parameter risks causing data races between code in the current task and concurrent execution of the closure}}
+      print(z) // expected-tns-note {{closure captures 'z' which is accessible to code in the current task}}
+    }
+  }
+}
+
+// We used to crash here since we were using an API that wanted a SIL Type and
+// we were creating a SILType from a CanType that needed to be lowered.
+func testUseCanTypeNonSendableCheckAPI(y: any Progress) async {
+  @Sendable func test() {
+    print(y) // expected-warning {{capture of 'y' with non-sendable type 'any Progress' in a '@Sendable' local function}}
+  }
+  await withTaskGroup(of: Void.self) { group in
+    group.addTask { // expected-tns-warning {{sending value of non-Sendable type '() async -> ()' risks causing data races}}
+      // expected-tns-note @-1 {{Passing task-isolated value of non-Sendable type '() async -> ()' as a 'sending' parameter risks causing races inbetween task-isolated uses and uses reachable from the callee}}
+      test()
+    }
+  }
+}