[rbi] Make it so that we correctly do not error on uses of Sendable values that are projected from non-Sendable bases.

Specifically, we only do this if the base is a let or if it is a var but not
captured by reference.

rdar://149019222
(cherry picked from commit 23b6937)

Author: gottesmm

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -368,6 +368,12 @@ class RegionAnalysisValueMap {
     }
 
     operator bool() const { return value; }
+
+    void print(llvm::raw_ostream &os) const;
+    SWIFT_DEBUG_DUMP {
+      print(llvm::dbgs());
+      llvm::dbgs() << '\n';
+    }
   };
 
   /// A map from a SILValue to its equivalence class representative.

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -1461,13 +1461,6 @@ struct PartitionOpEvaluator {
           return;
       }
 
-      // If this is a require of a mutable base of a Sendable value and we were
-      // not closure captured, we can bail.
-      if (op.getOptions().containsOnly(PartitionOp::Flag::RequireOfMutableBaseOfSendableValue) &&
-          regionHasClosureCapturedElt) {
-        return;
-      }
-
       // Mark op.getOpArg1() as sent.
       SendingOperandState &state = operandToStateMap.get(op.getSourceOp());
       state.isClosureCaptured |= regionHasClosureCapturedElt;

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -310,6 +310,8 @@ static bool isStaticallyLookThroughInst(SILInstruction *inst) {
   case SILInstructionKind::UnmanagedToRefInst:
   case SILInstructionKind::InitExistentialValueInst:
   case SILInstructionKind::UncheckedEnumDataInst:
+  case SILInstructionKind::StructElementAddrInst:
+  case SILInstructionKind::TupleElementAddrInst:
     return true;
   case SILInstructionKind::MoveValueInst:
     // Look through if it isn't from a var decl.
@@ -334,9 +336,6 @@ static bool isLookThroughIfOperandAndResultNonSendable(SILInstruction *inst) {
   case SILInstructionKind::UncheckedTrivialBitCastInst:
   case SILInstructionKind::UncheckedBitwiseCastInst:
   case SILInstructionKind::UncheckedValueCastInst:
-  case SILInstructionKind::StructElementAddrInst:
-  case SILInstructionKind::TupleElementAddrInst:
-  case SILInstructionKind::UncheckedTakeEnumDataAddrInst:
   case SILInstructionKind::ConvertEscapeToNoEscapeInst:
   case SILInstructionKind::ConvertFunctionInst:
   case SILInstructionKind::RefToRawPointerInst:
@@ -735,6 +734,7 @@ TrackableValueLookupResult RegionAnalysisValueMap::getTrackableValue(
 
   auto trackedValue =
       getTrackableValueHelper(value, isAddressCapturedByPartialApply);
+
   std::optional<TrackableValue> trackedBase;
   if (base)
     trackedBase =
@@ -1130,6 +1130,17 @@ void TrackableValueLookupResult::print(llvm::raw_ostream &os) const {
   }
 }
 
+//===----------------------------------------------------------------------===//
+//                      MARK: UnderlyingTrackedValueInfo
+//===----------------------------------------------------------------------===//
+
+void RegionAnalysisValueMap::UnderlyingTrackedValueInfo::print(
+    llvm::raw_ostream &os) const {
+  os << "RegionAnalysisValueMap.\nValue: " << value;
+  if (base)
+    os << "Base: " << base;
+}
+
 //===----------------------------------------------------------------------===//
 //                      MARK: Partial Apply Reachability
 //===----------------------------------------------------------------------===//
@@ -3204,10 +3215,16 @@ void PartitionOpBuilder::addRequire(TrackableValueLookupResult value) {
         if (!boxType->getLayout()->isMutable())
           return;
 
+        options |= PartitionOp::Flag::RequireOfMutableBaseOfSendableValue;
+      } else if (auto *asi = dyn_cast<AllocStackInst>(rep)) {
+        if (asi->isLet())
+          return;
+
         options |= PartitionOp::Flag::RequireOfMutableBaseOfSendableValue;
       }
+
+      addRequire(value.base, options);
     }
-    addRequire(value.base, options);
   }
 }
 
@@ -3355,6 +3372,9 @@ CONSTANT_TRANSLATION(RefToUnmanagedInst, LookThrough)
 CONSTANT_TRANSLATION(UnmanagedToRefInst, LookThrough)
 CONSTANT_TRANSLATION(InitExistentialValueInst, LookThrough)
 CONSTANT_TRANSLATION(UncheckedEnumDataInst, LookThrough)
+CONSTANT_TRANSLATION(TupleElementAddrInst, LookThrough)
+CONSTANT_TRANSLATION(StructElementAddrInst, LookThrough)
+CONSTANT_TRANSLATION(UncheckedTakeEnumDataAddrInst, LookThrough)
 
 //===---
 // Store
@@ -3604,9 +3624,6 @@ IGNORE_IF_SENDABLE_RESULT_ASSIGN_OTHERWISE(StructExtractInst)
 LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(UncheckedTrivialBitCastInst)
 LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(UncheckedBitwiseCastInst)
 LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(UncheckedValueCastInst)
-LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(TupleElementAddrInst)
-LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(StructElementAddrInst)
-LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(UncheckedTakeEnumDataAddrInst)
 LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(ConvertEscapeToNoEscapeInst)
 LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(ConvertFunctionInst)
 
@@ -3701,20 +3718,42 @@ PartitionOpTranslator::visitBeginBorrowInst(BeginBorrowInst *bbi) {
   return TranslationSemantics::LookThrough;
 }
 
-/// LoadInst is technically a statically look through instruction, but we want
-/// to handle it especially in the infrastructure, so we cannot mark it as
-/// such. This makes marking it as a normal lookthrough instruction impossible
-/// since the routine checks that invariant.
-TranslationSemantics PartitionOpTranslator::visitLoadInst(LoadInst *limvi) {
+/// LoadInst has two different semantics:
+///
+/// 1. If the load produces a non-Sendable value, we want to treat it as a
+/// statically look through instruction, but we want to handle it especially in
+/// the infrastructure, so we cannot mark it as such. This makes marking it as a
+/// normal lookthrough instruction impossible since the routine checks that
+/// invariant.
+///
+/// 2. If the load produces a Sendable value, we want to perform a require on
+/// the address so that if we have a load from a non-Sendable base, we properly
+/// require the base.
+TranslationSemantics PartitionOpTranslator::visitLoadInst(LoadInst *li) {
+  if (SILIsolationInfo::isSendableType(li->getOperand())) {
+    translateSILRequire(li->getOperand());
+  }
+
   return TranslationSemantics::Special;
 }
 
-/// LoadBorrowInst is technically a statically look through instruction, but we
-/// want to handle it especially in the infrastructure, so we cannot mark it as
-/// such. This makes marking it as a normal lookthrough instruction impossible
-/// since the routine checks that invariant.
+/// LoadBorrowInst has two different semantics:
+///
+/// 1. If the load produces a non-Sendable value, we want to treat it as a
+/// statically look through instruction, but we want to handle it especially in
+/// the infrastructure, so we cannot mark it as such. This makes marking it as a
+/// normal lookthrough instruction impossible since the routine checks that
+/// invariant.
+///
+/// 2. If the load produces a Sendable value, we want to perform a require on
+/// the address so that if we have a load from a non-Sendable base, we properly
+/// require the base.
 TranslationSemantics
 PartitionOpTranslator::visitLoadBorrowInst(LoadBorrowInst *lbi) {
+  if (SILIsolationInfo::isSendableType(lbi->getOperand())) {
+    translateSILRequire(lbi->getOperand());
+  }
+
   return TranslationSemantics::Special;
 }
 

lib/SILOptimizer/Mandatory/SendNonSendable.cpp

@@ -2702,6 +2702,14 @@ struct DiagnosticEvaluator final
   void handleLocalUseAfterSend(LocalUseAfterSendError error) const {
     const auto &partitionOp = *error.op;
     REGIONBASEDISOLATION_LOG(error.print(llvm::dbgs(), info->getValueMap()));
+
+    // Ignore this if we are erroring on a mutable base of a Sendable value and
+    // if when we sent the value's region was not closure captured.
+    if (error.op->getOptions().containsOnly(
+            PartitionOp::Flag::RequireOfMutableBaseOfSendableValue) &&
+        !operandToStateMap.get(error.sendingOp).isClosureCaptured)
+      return;
+
     sendingOpToRequireInstMultiMap.insert(
         error.sendingOp, RequireInst::forUseAfterSend(partitionOp.getSourceInst()));
   }

lib/SILOptimizer/Utils/PartitionUtils.cpp

@@ -142,6 +142,9 @@ void PartitionOp::print(llvm::raw_ostream &os, bool extraSpace) const {
   }
   case PartitionOpKind::Require: {
     os << "require ";
+    if (getOptions().containsOnly(
+            PartitionOp::Flag::RequireOfMutableBaseOfSendableValue))
+      os << "[mutable_base_of_sendable_val] ";
     if (extraSpace)
       os << extraSpaceLiteral;
     os << "%%" << getOpArg1();

test/Concurrency/transfernonsendable.swift

@@ -1968,3 +1968,32 @@ extension NonIsolatedFinalKlass {
     // expected-tns-note @-1 {{sending task-isolated 'self.ns' to main actor-isolated global function 'transferToMain' risks causing data races between main actor-isolated and task-isolated uses}}
   }
 }
+
+func mutableLocalCaptureDataRace() async {
+  var x = 0
+  x = 0
+  _ = x
+
+  Task.detached { x = 1 } // expected-tns-warning {{sending value of non-Sendable type '() async -> ()' risks causing data races}}
+  // expected-tns-note @-1 {{Passing value of non-Sendable type '() async -> ()' as a 'sending' argument to static method 'detached(priority:operation:)' risks causing races in between local and caller code}}
+
+  x = 2 // expected-tns-note {{access can happen concurrently}}
+}
+
+func mutableLocalCaptureDataRace2() async {
+  var x = 0
+  x = 0
+
+  Task.detached { x = 1 } // expected-tns-warning {{sending value of non-Sendable type '() async -> ()' risks causing data races}}
+  // expected-tns-note @-1 {{Passing value of non-Sendable type '() async -> ()' as a 'sending' argument to static method 'detached(priority:operation:)' risks causing races in between local and caller code}}
+
+  print(x) // expected-tns-note {{access can happen concurrently}}
+}
+
+func localCaptureDataRace3() async {
+  let x = 0
+
+  Task.detached { print(x) }
+
+  print(x)
+}