[region-isolation] Fix the call site or self error for values used in the same region as a function argument.

This is just good to do and also makes it so that in my test case for
assumeIsolated, I get a better msg.

Author: gottesmm

include/swift/AST/DiagnosticsSIL.def

@@ -907,8 +907,14 @@ WARNING(regionbasedisolation_transfer_yields_race_transferring_parameter, none,
 WARNING(regionbasedisolation_transfer_yields_race_stronglytransferred_binding, none,
         "binding of non-Sendable type %0 accessed after being transferred; later accesses could race",
         (Type))
+WARNING(regionbasedisolation_arg_transferred, none,
+        "task isolated value of type %0 transferred to %1 context; later accesses to value could race",
+        (Type, ActorIsolation))
 NOTE(regionbasedisolation_maybe_race, none,
      "access here could race", ())
+NOTE(regionbasedisolation_isolated_since_in_same_region_basename, none,
+     "value is %0 since it is in the same region as %1",
+     (StringRef, DeclBaseName))
 ERROR(regionbasedisolation_unknown_pattern, none,
       "pattern that the region based isolation checker does not understand how to check. Please file a bug",
       ())

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1604,12 +1604,30 @@ class PartitionOpTranslator {
 
     // First check if our partial_apply is fed into an async let begin. If so,
     // handle it especially.
+    //
+    // NOTE: If it is an async_let, then the closure itself will be Sendable. We
+    // treat passing in a value into the async Sendable closure as transferring
+    // it into the closure.
     if (isAsyncLetBeginPartialApply(pai)) {
       return translateSILPartialApplyAsyncLetBegin(pai);
     }
 
-    // Check if our closure is isolated to a specific actor. If it is, we need
-    // to treat all of the captures as being transferred.
+    // Then check if our partial apply is Sendable. In such a case, we will have
+    // emitted an earlier warning in Sema.
+    //
+    // DISCUSSION: The reason why we can treat values passed into an async let
+    // as transferring safely but it is unsafe to do this for arbitrary Sendable
+    // closures is that we do not know how many times the Sendable closure will
+    // be executed. It is possible to have different invocations of the Sendable
+    // closure to cause races with the captured non-Sendable value. In contrast
+    // since we know an async let runs exactly once, we do not need to worry
+    // about such a possibility. If we had the ability in the language to
+    // specify that a closure is run at most once or that it is always run
+    // serially, we could lift this restriction... so for now we leave in the
+    // Sema warning and just bail here.
+    if (pai->getFunctionType()->isSendableType())
+      return;
+
     if (auto *ace = pai->getLoc().getAsASTNode<AbstractClosureExpr>()) {
       if (ace->getActorIsolation().isActorIsolated()) {
         return translateIsolatedPartialApply(pai);

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -600,19 +600,18 @@ void UseAfterTransferDiagnosticInferrer::init(const Operand *op) {
     }
   }
 
+  auto loc = op->getUser()->getLoc();
+
   // If we have a partial_apply that is actor isolated, see if we found a
   // transfer error due to us transferring a value into it.
-  if (auto pai = dyn_cast<PartialApplyInst>(nonConstOp->getUser())) {
-    if (auto *ace = pai->getLoc().getAsASTNode<AbstractClosureExpr>()) {
-      if (ace->getActorIsolation().isActorIsolated()) {
-        if (initForIsolatedPartialApply(nonConstOp, ace)) {
-          return;
-        }
+  if (auto *ace = loc.getAsASTNode<AbstractClosureExpr>()) {
+    if (ace->getActorIsolation().isActorIsolated()) {
+      if (initForIsolatedPartialApply(nonConstOp, ace)) {
+        return;
       }
     }
   }
 
-  auto loc = op->getUser()->getLoc();
   if (auto *sourceApply = loc.getAsASTNode<ApplyExpr>()) {
     return initForApply(op, sourceApply);
   }
@@ -639,6 +638,132 @@ void UseAfterTransferDiagnosticInferrer::init(const Operand *op) {
   autoClosureExpr->walk(walker);
 }
 
+//===----------------------------------------------------------------------===//
+//            MARK: Transfer NonTransferrable Diagnostic Inference
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct TransferredNonTransferrableInfo {
+  /// The use that actually caused the transfer.
+  Operand *transferredOperand;
+
+  /// The non-transferrable value that is in the same region as \p
+  /// transferredOperand.get().
+  SILValue nonTransferrableValue;
+
+  TransferredNonTransferrableInfo(Operand *transferredOperand,
+                                  SILValue nonTransferrableValue)
+      : transferredOperand(transferredOperand),
+        nonTransferrableValue(nonTransferrableValue) {}
+};
+
+class TransferNonTransferrableDiagnosticInferrer {
+public:
+  enum class UseDiagnosticInfoKind {
+    Invalid = 0,
+
+    /// Used if we have a use that we haven't categorized so we emit the generic
+    /// call site is self error.
+    MiscUse = 1,
+
+    /// Used if we have a function argument that is transferred into an apply.
+    FunctionArgumentApply = 2,
+
+    /// Used if we have a function argument that is transferred into an closure.
+    FunctionArgumentClosure = 3,
+  };
+
+  struct UseDiagnosticInfo {
+    UseDiagnosticInfoKind kind;
+    std::optional<ActorIsolation> transferredIsolation;
+
+    static UseDiagnosticInfo forMiscUse() {
+      return {UseDiagnosticInfoKind::MiscUse, {}};
+    }
+
+    static UseDiagnosticInfo
+    forFunctionArgumentApply(ActorIsolation isolation) {
+      return {UseDiagnosticInfoKind::FunctionArgumentApply, isolation};
+    }
+
+    static UseDiagnosticInfo
+    forFunctionArgumentClosure(ActorIsolation isolation) {
+      return {UseDiagnosticInfoKind::FunctionArgumentClosure, isolation};
+    }
+
+  private:
+    UseDiagnosticInfo(UseDiagnosticInfoKind kind,
+                      std::optional<ActorIsolation> isolation)
+        : kind(kind), transferredIsolation(isolation) {}
+  };
+
+private:
+  TransferredNonTransferrableInfo info;
+  std::optional<UseDiagnosticInfo> diagnosticInfo;
+  SourceLoc loc;
+
+public:
+  TransferNonTransferrableDiagnosticInferrer(
+      TransferredNonTransferrableInfo info)
+      : info(info),
+        loc(info.transferredOperand->getUser()->getLoc().getSourceLoc()) {}
+
+  void run();
+
+  UseDiagnosticInfo getDiagnostic() const { return diagnosticInfo.value(); }
+  SourceLoc getLoc() const { return loc; }
+
+private:
+  bool initForIsolatedPartialApply(Operand *op, AbstractClosureExpr *ace);
+};
+
+} // namespace
+
+bool TransferNonTransferrableDiagnosticInferrer::initForIsolatedPartialApply(
+    Operand *op, AbstractClosureExpr *ace) {
+  SmallVector<std::tuple<CapturedValue, unsigned, ApplyIsolationCrossing>, 8>
+      foundCapturedIsolationCrossing;
+  ace->getIsolationCrossing(foundCapturedIsolationCrossing);
+  if (foundCapturedIsolationCrossing.empty())
+    return false;
+
+  unsigned opIndex = ApplySite(op->getUser()).getAppliedArgIndex(*op);
+  for (auto &p : foundCapturedIsolationCrossing) {
+    if (std::get<1>(p) == opIndex) {
+      loc = std::get<0>(p).getLoc();
+      diagnosticInfo = UseDiagnosticInfo::forFunctionArgumentClosure(
+          std::get<2>(p).getCalleeIsolation());
+      return true;
+    }
+  }
+
+  return false;
+}
+
+void TransferNonTransferrableDiagnosticInferrer::run() {
+  if (isa<SILFunctionArgument>(info.nonTransferrableValue)) {
+    // We need to find the isolation info.
+    auto loc = info.transferredOperand->getUser()->getLoc();
+
+    if (auto *sourceApply = loc.getAsASTNode<ApplyExpr>()) {
+      diagnosticInfo = UseDiagnosticInfo::forFunctionArgumentApply(
+          sourceApply->getIsolationCrossing().value().getCalleeIsolation());
+      return;
+    }
+
+    if (auto *ace = loc.getAsASTNode<AbstractClosureExpr>()) {
+      if (ace->getActorIsolation().isActorIsolated()) {
+        if (initForIsolatedPartialApply(info.transferredOperand, ace)) {
+          return;
+        }
+      }
+    }
+  }
+
+  diagnosticInfo = UseDiagnosticInfo::forMiscUse();
+}
+
 //===----------------------------------------------------------------------===//
 //                         MARK: Diagnostic Emission
 //===----------------------------------------------------------------------===//
@@ -650,13 +775,18 @@ struct DiagnosticEvaluator final
   RegionAnalysisFunctionInfo *info;
   SmallFrozenMultiMap<Operand *, SILInstruction *, 8>
       &transferOpToRequireInstMultiMap;
-  SmallVectorImpl<Operand *> &transferredNonTransferrable;
+
+  /// First value is the operand that was transferred... second value is the
+  /// non-transferrable value in the same region as that value. The second value
+  /// is what is non-transferrable.
+  SmallVectorImpl<TransferredNonTransferrableInfo> &transferredNonTransferrable;
 
   DiagnosticEvaluator(Partition &workingPartition,
                       RegionAnalysisFunctionInfo *info,
                       SmallFrozenMultiMap<Operand *, SILInstruction *, 8>
                           &transferOpToRequireInstMultiMap,
-                      SmallVectorImpl<Operand *> &transferredNonTransferrable)
+                      SmallVectorImpl<TransferredNonTransferrableInfo>
+                          &transferredNonTransferrable)
       : PartitionOpEvaluatorBaseImpl(workingPartition,
                                      info->getOperandSetFactory()),
         info(info),
@@ -706,7 +836,10 @@ struct DiagnosticEvaluator final
                << "        Rep: "
                << *info->getValueMap().getRepresentative(transferredVal));
     auto *self = const_cast<DiagnosticEvaluator *>(this);
-    self->transferredNonTransferrable.push_back(partitionOp.getSourceOp());
+    auto nonTransferrableValue =
+        info->getValueMap().getRepresentative(transferredVal);
+    self->transferredNonTransferrable.emplace_back(partitionOp.getSourceOp(),
+                                                   nonTransferrableValue);
   }
 
   bool isActorDerived(Element element) const {
@@ -729,7 +862,7 @@ class TransferNonSendableImpl {
   RegionAnalysisFunctionInfo *regionInfo;
   SmallFrozenMultiMap<Operand *, SILInstruction *, 8>
       transferOpToRequireInstMultiMap;
-  SmallVector<Operand *, 8> transferredNonTransferrable;
+  SmallVector<TransferredNonTransferrableInfo, 8> transferredNonTransferrable;
 
 public:
   TransferNonSendableImpl(RegionAnalysisFunctionInfo *regionInfo)
@@ -889,8 +1022,66 @@ void TransferNonSendableImpl::emitUseAfterTransferDiagnostics() {
 }
 
 void TransferNonSendableImpl::emitTransferredNonTransferrableDiagnostics() {
-  for (auto *op : transferredNonTransferrable) {
-    diagnose(op->getUser(), diag::regionbasedisolation_selforargtransferred);
+  if (transferredNonTransferrable.empty())
+    return;
+
+  using UseDiagnosticInfoKind =
+      TransferNonTransferrableDiagnosticInferrer::UseDiagnosticInfoKind;
+
+  auto &astContext = regionInfo->getFunction()->getASTContext();
+  for (auto info : transferredNonTransferrable) {
+    auto *op = info.transferredOperand;
+    TransferNonTransferrableDiagnosticInferrer diagnosticInferrer(info);
+    diagnosticInferrer.run();
+
+    auto diagnosticInfo = diagnosticInferrer.getDiagnostic();
+    auto loc = diagnosticInferrer.getLoc();
+
+    switch (diagnosticInfo.kind) {
+    case UseDiagnosticInfoKind::Invalid:
+      llvm_unreachable("Should never see this");
+    case UseDiagnosticInfoKind::MiscUse:
+      diagnose(astContext, loc,
+               diag::regionbasedisolation_selforargtransferred);
+      break;
+    case UseDiagnosticInfoKind::FunctionArgumentApply: {
+      diagnose(astContext, loc, diag::regionbasedisolation_arg_transferred,
+               op->get()->getType().getASTType(),
+               diagnosticInfo.transferredIsolation.value())
+          .highlight(op->getUser()->getLoc().getSourceRange());
+      // Only emit the note if our value is different from the function
+      // argument.
+      auto rep = regionInfo->getValueMap()
+                     .getTrackableValue(op->get())
+                     .getRepresentative();
+      if (rep.maybeGetValue() == info.nonTransferrableValue)
+        continue;
+      auto *fArg = cast<SILFunctionArgument>(info.nonTransferrableValue);
+      diagnose(
+          astContext, fArg->getDecl()->getLoc(),
+          diag::regionbasedisolation_isolated_since_in_same_region_basename,
+          "task isolated", fArg->getDecl()->getBaseName());
+      break;
+    }
+    case UseDiagnosticInfoKind::FunctionArgumentClosure: {
+      diagnose(astContext, loc, diag::regionbasedisolation_arg_transferred,
+               op->get()->getType().getASTType(),
+               diagnosticInfo.transferredIsolation.value())
+          .highlight(op->getUser()->getLoc().getSourceRange());
+      // Only emit the note if our value is different from the function
+      // argument.
+      auto rep = regionInfo->getValueMap()
+                     .getTrackableValue(op->get())
+                     .getRepresentative();
+      if (rep.maybeGetValue() == info.nonTransferrableValue)
+        continue;
+      auto *fArg = cast<SILFunctionArgument>(info.nonTransferrableValue);
+      diagnose(
+          astContext, fArg->getDecl()->getLoc(),
+          diag::regionbasedisolation_isolated_since_in_same_region_basename,
+          "task isolated", fArg->getDecl()->getBaseName());
+    }
+    }
   }
 }
 

test/Concurrency/sendable_checking.swift

@@ -251,11 +251,11 @@ final class NonSendable {
   func call() async {
     await update()
     // expected-targeted-and-complete-warning @-1 {{passing argument of non-sendable type 'NonSendable' into main actor-isolated context may introduce data races}}
-    // expected-tns-warning @-2 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
+    // expected-tns-warning @-2 {{task isolated value of type 'NonSendable' transferred to main actor-isolated context; later accesses to value could race}}
 
     await self.update()
     // expected-targeted-and-complete-warning @-1 {{passing argument of non-sendable type 'NonSendable' into main actor-isolated context may introduce data races}}
-    // expected-tns-warning @-2 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
+    // expected-tns-warning @-2 {{task isolated value of type 'NonSendable' transferred to main actor-isolated context; later accesses to value could race}}
 
     _ = await x
     // expected-warning@-1 {{non-sendable type 'NonSendable' passed in implicitly asynchronous call to main actor-isolated property 'x' cannot cross actor boundary}}
@@ -287,18 +287,18 @@ func globalSendable(_ ns: NonSendable) async {}
 @available(SwiftStdlib 5.1, *)
 @MainActor
 func callNonisolatedAsyncClosure(
-  ns: NonSendable,
+  ns: NonSendable, // expected-tns-note {{value is task isolated since it is in the same region as 'ns'}}
   g: (NonSendable) async -> Void
 ) async {
   await g(ns)
   // expected-targeted-and-complete-warning @-1 {{passing argument of non-sendable type 'NonSendable' outside of main actor-isolated context may introduce data races}}
-  // expected-tns-warning @-2 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
-  // expected-tns-warning @-3 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
+  // expected-tns-warning @-2 {{task isolated value of type 'NonSendable' transferred to nonisolated context; later accesses to value could race}}
+  // expected-tns-warning @-3 {{task isolated value of type '@noescape @async @callee_guaranteed (@guaranteed NonSendable) -> ()' transferred to nonisolated context; later accesses to value could race}}
 
   let f: (NonSendable) async -> () = globalSendable // okay
   await f(ns)
   // expected-targeted-and-complete-warning@-1 {{passing argument of non-sendable type 'NonSendable' outside of main actor-isolated context may introduce data races}}
-  // expected-tns-warning @-2 {{call site passes `self` or a non-sendable argument of this function to another thread, potentially yielding a race with the caller}}
+  // expected-tns-warning @-2 {{task isolated value of type 'NonSendable' transferred to nonisolated context; later accesses to value could race}}
 
 }