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

This problem comes up with the following example:

```swift
class A {
    var description = ""
}

class B {
    let a = A()

    func b() {
        let asdf = ""
        Task { @mainactor in
            a.description = asdf // Sending 'asdf' risks causing data races
        }
    }
}
```

The specific issue is that the closure we generate actually includes an
implicit(any) parameter at the SIL level which occurs after the callee operand
but before the captures. This caused the captured variable index from the AST
and the one we compute from the partial_apply to differ by 1. So we need to
subtract 1 in such a case. That is why we used to print 'asdf' instead of 'a'
above.

DISCUSSION: This shows an interesting difference between SIL applied arg indices
and AST indices. SIL applied arg indices would include the implicit(any)
parameter since it is a parameter in the SIL function type. In contrast, this
doesn't show up in the formal AST parameters or captures. To make it easier to
reason about this, I added a new API to ApplySite called
ApplySite::getASTAppliedArgIndex and added large comments to
getASTAppliedArgIndex and getAppliedArgIndex that explains the issue.

rdar://136593706
swiftlang#76648

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

@@ -4416,6 +4416,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}}
+          }
+      }
+  }
+}