Merge pull request #77354 from gottesmm/rdar136593706

[region-isolation] Fix an off by one error when mapping AST capture indices to SIL level parameter indices.

Author: gottesmm

include/swift/AST/Expr.h

@@ -4284,7 +4284,8 @@ class ClosureExpr : public AbstractClosureExpr {
   }
 
   /// Whether this closure should implicitly inherit the actor context from
-  /// where it was formed. This only affects @Sendable async closures.
+  /// where it was formed. This only affects @Sendable and sendable async
+  /// closures.
   bool inheritsActorContext() const {
     return Bits.ClosureExpr.InheritActorContext;
   }

include/swift/AST/Types.h

@@ -4393,6 +4393,10 @@ class SILParameterInfo {
     Sending = 0x2,
 
     /// Set if the given parameter is isolated.
+    ///
+    /// This means that the value provides the functions isolation. This implies
+    /// that the parameter must be an Optional actor or something that conforms
+    /// to AnyActor.
     Isolated = 0x4,
   };
 

include/swift/SIL/ApplySite.h

@@ -338,13 +338,45 @@ class ApplySite {
   }
 
   /// Return the applied argument index for the given operand.
+  ///
+  /// This corresponds to the index of the operand in the actual
+  /// SILFunctionType.
+  ///
+  /// As a result this does not include indirect parameters, but it does include
+  /// implicit parameters (e.x.: the implicit isolated(any) parameter) that are
+  /// in the SIL type but are not in the AST type. To ignore such implicit
+  /// parameters, use getASTAppliedArgIndex.
   unsigned getAppliedArgIndex(const Operand &oper) const {
     assert(oper.getUser() == Inst);
     assert(isArgumentOperand(oper));
 
     return oper.getOperandNumber() - getOperandIndexOfFirstArgument();
   }
 
+  /// Returns the applied AST argument index for the given operand.
+  ///
+  /// This is guaranteed to return an index that can be used with AST function
+  /// types. This means it looks past the callee, indirect function types, as
+  /// well as any other implicit parameters like isolated(any).
+  ///
+  /// NOTE: If we add more implicit parameters to SILFunction types that do not
+  /// appear at the AST level, this code needs to be updated.
+  unsigned getASTAppliedArgIndex(const Operand &oper) const {
+    // We rely on the assertions in getAppliedArgIndex to allow for us to assume
+    // that we have an argument operand. We check later that we do not have an
+    // isolated(any) parameter.
+    unsigned appliedArgIndex = getAppliedArgIndex(oper);
+    if (auto *pai = dyn_cast<PartialApplyInst>(Inst)) {
+      if (pai->getFunctionType()->getIsolation() ==
+          SILFunctionTypeIsolation::Erased) {
+        assert(appliedArgIndex != 0 &&
+               "isolation(any) does not correspond to an AST argument");
+        appliedArgIndex -= 1;
+      }
+    }
+    return appliedArgIndex;
+  }
+
   /// Return the callee's function argument index corresponding to the first
   /// applied argument: 0 for full applies; >= 0 for partial applies.
   unsigned getCalleeArgIndexOfFirstAppliedArg() const {

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -891,7 +891,7 @@ bool UseAfterTransferDiagnosticInferrer::initForIsolatedPartialApply(
   if (foundCapturedIsolationCrossing.empty())
     return false;
 
-  unsigned opIndex = ApplySite(op->getUser()).getAppliedArgIndex(*op);
+  unsigned opIndex = ApplySite(op->getUser()).getASTAppliedArgIndex(*op);
   bool emittedDiagnostic = false;
   for (auto &p : foundCapturedIsolationCrossing) {
     if (std::get<1>(p) != opIndex)
@@ -1131,7 +1131,7 @@ void UseAfterTransferDiagnosticInferrer::infer() {
 
   auto *i = transferOp->getUser();
   auto pai = ApplySite::isa(i);
-  unsigned captureIndex = pai.getAppliedArgIndex(*transferOp);
+  unsigned captureIndex = pai.getASTAppliedArgIndex(*transferOp);
 
   auto &state = transferringOpToStateMap.get(transferOp);
   auto captureInfo =
@@ -1640,11 +1640,7 @@ class SentNeverSendableDiagnosticInferrer {
     for (auto &paiOp : ApplySite(pai).getArgumentOperands()) {
       if (valueMap.getTrackableValue(paiOp.get()).getRepresentative() ==
           isolatedValue) {
-        // isolated_any causes all partial apply parameters to be shifted by 1
-        // due to the implicit isolated any parameter.
-        unsigned isIsolatedAny = pai->getFunctionType()->getIsolation() ==
-                                 SILFunctionTypeIsolation::Erased;
-        return ApplySite(pai).getAppliedArgIndex(paiOp) - isIsolatedAny;
+        return ApplySite(pai).getASTAppliedArgIndex(paiOp);
       }
     }
 
@@ -1728,7 +1724,10 @@ bool SentNeverSendableDiagnosticInferrer::initForIsolatedPartialApply(
   if (foundCapturedIsolationCrossing.empty())
     return false;
 
-  unsigned opIndex = ApplySite(op->getUser()).getAppliedArgIndex(*op);
+  // We use getASTAppliedArgIndex instead of getAppliedArgIndex to ensure that
+  // we ignore for our indexing purposes any implicit initial parameters like
+  // isolated(any).
+  unsigned opIndex = ApplySite(op->getUser()).getASTAppliedArgIndex(*op);
   for (auto &p : foundCapturedIsolationCrossing) {
     if (std::get<1>(p) == opIndex) {
       auto loc = RegularLocation(std::get<0>(p).getLoc());
@@ -1944,7 +1943,7 @@ bool SentNeverSendableDiagnosticInferrer::run() {
     if (autoClosureExpr->getThunkKind() == AutoClosureExpr::Kind::AsyncLet) {
       auto *i = op->getUser();
       auto pai = ApplySite::isa(i);
-      unsigned captureIndex = pai.getAppliedArgIndex(*op);
+      unsigned captureIndex = pai.getASTAppliedArgIndex(*op);
       auto captureInfo =
           autoClosureExpr->getCaptureInfo().getCaptures()[captureIndex];
       auto loc = RegularLocation(captureInfo.getLoc(), false /*implicit*/);

test/Concurrency/transfernonsendable.swift

@@ -1888,3 +1888,21 @@ func nonSendableAllocBoxConsumingVar() async throws {
     try await group.waitForAll()
   }
 }
+
+func offByOneWithImplicitPartialApply() {
+  class A {
+      var description = ""
+  }
+
+  class B {
+      let a = A()
+
+      func b() {
+          let asdf = ""
+          Task { @MainActor in
+            a.description = asdf // expected-tns-warning {{sending 'self' risks causing data races}}
+            // expected-tns-note @-1 {{task-isolated 'self' is captured by a main actor-isolated closure. main actor-isolated uses in closure may race against later nonisolated uses}}
+          }
+      }
+  }
+}