[sil-isolation-info] When determining isolation of a function arg, use its VarDecl.

Otherwise, we can be inconsistent with isolations returned by other parts of the
code. Previously we were just treating it always as self + nom decl, which is
clearly wrong if a type is not self (e.x.: if it is an isolated parameter).

rdar://135459885

include/swift/AST/ActorIsolation.h

@@ -217,6 +217,10 @@ public:
     }
   }
 
+  /// In the debugger return the index for the stored actorInstance pointer
+  /// union index. Asserts if not an actor instance.
+  SWIFT_DEBUG_HELPER(unsigned getActorInstanceUnionIndex() const);
+
   NominalTypeDecl *getActor() const;
 
   VarDecl *getActorInstance() const;

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -71,7 +71,14 @@ public:
   ActorInstance() : ActorInstance(SILValue(), Kind::Value) {}
 
   static ActorInstance getForValue(SILValue value) {
+    if (!value)
+      return ActorInstance();
     value = lookThroughInsts(value);
+    if (!value->getType()
+             .getASTType()
+             ->lookThroughAllOptionalTypes()
+             ->isAnyActorType())
+      return ActorInstance();
     return ActorInstance(value, Kind::Value);
   }
 
@@ -96,7 +103,7 @@ public:
     return value.getPointer();
   }
 
-  SILValue maybeGetValue() const {
+  LLVM_ATTRIBUTE_USED SILValue maybeGetValue() const {
     if (getKind() != Kind::Value)
       return SILValue();
     return getValue();
@@ -132,6 +139,10 @@ public:
   bool operator!=(const ActorInstance &other) const {
     return !(*this == other);
   }
+
+  void print(llvm::raw_ostream &os) const;
+
+  SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 };
 
 /// The isolation info inferred for a specific SILValue. Use
@@ -372,6 +383,21 @@ public:
             ActorIsolation::forActorInstanceSelf(typeDecl)};
   }
 
+  static SILIsolationInfo
+  getActorInstanceIsolated(SILValue isolatedValue,
+                           const SILFunctionArgument *actorInstance) {
+    assert(actorInstance);
+    auto *varDecl =
+        llvm::dyn_cast_if_present<VarDecl>(actorInstance->getDecl());
+    if (!varDecl)
+      return {};
+    return {isolatedValue, actorInstance,
+            actorInstance->isSelf()
+                ? ActorIsolation::forActorInstanceSelf(varDecl)
+                : ActorIsolation::forActorInstanceParameter(
+                      varDecl, actorInstance->getIndex())};
+  }
+
   static SILIsolationInfo getActorInstanceIsolated(SILValue isolatedValue,
                                                    ActorInstance actorInstance,
                                                    NominalTypeDecl *typeDecl) {

lib/AST/TypeCheckRequests.cpp

@@ -2022,6 +2022,17 @@ void ActorIsolation::dumpForDiagnostics() const {
   llvm::dbgs() << '\n';
 }
 
+unsigned ActorIsolation::getActorInstanceUnionIndex() const {
+  assert(isActorInstanceIsolated());
+  if (actorInstance.is<NominalTypeDecl *>())
+    return 0;
+  if (actorInstance.is<VarDecl *>())
+    return 1;
+  if (actorInstance.is<Expr *>())
+    return 2;
+  llvm_unreachable("Unhandled");
+}
+
 void swift::simple_display(
     llvm::raw_ostream &out, const ActorIsolation &state) {
   if (state.preconcurrency())

lib/SILOptimizer/Utils/SILIsolationInfo.cpp

@@ -415,6 +415,12 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
           return SILIsolationInfo::getGlobalActorIsolated(SILValue(), selfASTType);
         }
 
+        if (auto *fArg = dyn_cast<SILFunctionArgument>(actualIsolatedValue)) {
+          if (auto info =
+                  SILIsolationInfo::getActorInstanceIsolated(fArg, fArg))
+            return info;
+        }
+
         // TODO: We really should be doing this based off of an Operand. Then
         // we would get the SILValue() for the first element. Today this can
         // only mess up isolation history.
@@ -449,7 +455,21 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
           }
         }
 
+        // Ok, we found an actor instance. Look for our actual isolated value.
         if (actorInstance) {
+          if (auto actualIsolatedValue =
+                  ActorInstance::getForValue(actorInstance)) {
+            // See if we have a function parameter. In that case, we want to see
+            // if we have a function argument. In such a case, we need to use
+            // the right parameter and the var decl.
+            if (auto *fArg = dyn_cast<SILFunctionArgument>(
+                    actualIsolatedValue.getValue())) {
+              if (auto info =
+                      SILIsolationInfo::getActorInstanceIsolated(pai, fArg))
+                return info;
+            }
+          }
+
           return SILIsolationInfo::getActorInstanceIsolated(
               pai, actorInstance, actorIsolation.getActor());
         }
@@ -480,6 +500,14 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
   if (auto *rei = dyn_cast<RefElementAddrInst>(inst)) {
     auto varIsolation = swift::getActorIsolation(rei->getField());
 
+    if (auto instance = ActorInstance::getForValue(rei->getOperand())) {
+      if (auto *fArg = llvm::dyn_cast_or_null<SILFunctionArgument>(
+              instance.maybeGetValue())) {
+        if (auto info = SILIsolationInfo::getActorInstanceIsolated(rei, fArg))
+          return info.withUnsafeNonIsolated(varIsolation.isNonisolatedUnsafe());
+      }
+    }
+
     auto *nomDecl =
         rei->getOperand()->getType().getNominalOrBoundGenericNominal();
 
@@ -868,11 +896,11 @@ SILIsolationInfo SILIsolationInfo::get(SILArgument *arg) {
 
   // Before we do anything further, see if we have an isolated parameter. This
   // handles isolated self and specifically marked isolated.
-  if (auto *isolatedArg = fArg->getFunction()->maybeGetIsolatedArgument()) {
+  if (auto *isolatedArg = llvm::cast_or_null<SILFunctionArgument>(
+          fArg->getFunction()->maybeGetIsolatedArgument())) {
     auto astType = isolatedArg->getType().getASTType();
     if (auto *nomDecl = astType->lookThroughAllOptionalTypes()->getAnyActor()) {
-      return SILIsolationInfo::getActorInstanceIsolated(fArg, isolatedArg,
-                                                        nomDecl);
+      return SILIsolationInfo::getActorInstanceIsolated(fArg, isolatedArg);
     }
   }
 
@@ -1017,10 +1045,10 @@ void SILIsolationInfo::print(llvm::raw_ostream &os) const {
   }
 }
 
-bool SILIsolationInfo::hasSameIsolation(ActorIsolation actorIsolation) const {
+bool SILIsolationInfo::hasSameIsolation(ActorIsolation other) const {
   if (getKind() != Kind::Actor)
     return false;
-  return getActorIsolation() == actorIsolation;
+  return getActorIsolation() == other;
 }
 
 bool SILIsolationInfo::hasSameIsolation(const SILIsolationInfo &other) const {
@@ -1362,6 +1390,25 @@ SILDynamicMergedIsolationInfo::merge(SILIsolationInfo other) const {
   return {other};
 }
 
+void ActorInstance::print(llvm::raw_ostream &os) const {
+  os << "Actor Instance. Kind: ";
+  switch (getKind()) {
+  case Kind::Value:
+    os << "Value.";
+    break;
+  case Kind::ActorAccessorInit:
+    os << "ActorAccessorInit.";
+    break;
+  case Kind::CapturedActorSelf:
+    os << "CapturedActorSelf.";
+    break;
+  }
+
+  if (auto value = maybeGetValue()) {
+    os << " Value: " << value;
+  };
+}
+
 //===----------------------------------------------------------------------===//
 //                                MARK: Tests
 //===----------------------------------------------------------------------===//

test/Concurrency/async_task_locals_basic_warnings_bug_isolation.swift

@@ -13,9 +13,6 @@ actor Test {
   func withTaskLocal(isolation: isolated (any Actor)? = #isolation,
                      _ body: (consuming NonSendableValue, isolated (any Actor)?) -> Void) async {
     Self.$local.withValue(12) {
-      // Unexpected errors here:
-      //  error: unexpected warning produced: sending 'body' risks causing data races; this is an error in the Swift 6 language mode
-      //  error: unexpected note produced: actor-isolated 'body' is captured by a actor-isolated closure. actor-isolated uses in closure may race against later nonisolated uses
       body(NonSendableValue(), isolation)
     }
   }

test/Concurrency/silisolationinfo_inference.sil

@@ -876,3 +876,22 @@ bb2(%10 : @owned $any Error):
   dealloc_stack %2 : $*T
   throw %10 : $any Error
 }
+
+// CHECK-LABEL: begin running test 1 of 1 on optional_test: sil_isolation_info_inference with: @trace[0]
+// CHECK-NEXT: Input Value:   %5 = ref_element_addr %4 : $MyActor, #MyActor.ns
+// CHECK-NEXT: Isolation: 'argument'-isolated
+// CHECK-NEXT: end running test 1 of 1 on optional_test: sil_isolation_info_inference with: @trace[0]
+sil [ossa] @optional_test : $@convention(thin) (@sil_isolated @guaranteed Optional<MyActor>) -> () {
+bb0(%0 : @guaranteed $Optional<MyActor>):
+  specify_test "sil_isolation_info_inference @trace[0]"
+  debug_value %0 : $Optional<MyActor>, let, name "argument"
+  %1a = copy_value %0 : $Optional<MyActor>
+  %1b = begin_borrow %1a
+  %1c = unchecked_enum_data %1b : $Optional<MyActor>, #Optional.some!enumelt
+  %1d = ref_element_addr %1c : $MyActor, #MyActor.ns
+  debug_value [trace] %1d
+  end_borrow %1b
+  destroy_value %1a
+  %9999 = tuple ()
+  return %9999 : $()
+}

test/Concurrency/transfernonsendable.swift

@@ -2007,3 +2007,23 @@ nonisolated func localCaptureDataRace4() {
 
   x = 2 // expected-tns-note {{access can happen concurrently}}
 }
+
+// We shouldn't error here since every time around, we are using the same
+// isolation.
+func testIsolatedParamInference() {
+  class S : @unchecked Sendable {}
+
+  final class B {
+    var s = S()
+    var b: Bool = false
+
+    func foo(isolation: isolated Actor = #isolation) async {
+      while !b {
+        await withTaskGroup(of: Int.self) { group in
+          _ = isolation
+          self.s = S()
+        }
+      }
+    }
+  }
+}