[rbi] Teach RegionIsolation how to properly error when 'inout sending' params are returned.

include/swift/AST/DiagnosticsSIL.def

@@ -1119,6 +1119,32 @@ ERROR(regionbasedisolation_inout_sending_cannot_be_actor_isolated, none,
 NOTE(regionbasedisolation_inout_sending_cannot_be_actor_isolated_note, none,
       "%1 %0 risks causing races in between %1 uses and caller uses since caller assumes value is not actor isolated",
       (Identifier, StringRef))
+ERROR(regionbasedisolation_inout_sending_cannot_be_returned_param, none,
+      "'inout sending' parameter %0 cannot be returned",
+      (Identifier))
+ERROR(regionbasedisolation_inout_sending_cannot_be_returned_value, none,
+      "%0 cannot be returned",
+      (Identifier))
+ERROR(regionbasedisolation_inout_sending_cannot_be_returned_value_result, none,
+      "result of %kind0 cannot be returned",
+      (const ValueDecl *))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_param, none,
+     "returning 'inout sending' parameter %0 risks concurrent access as caller assumes %0 and result can be sent to different isolation domains",
+     (Identifier))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_value, none,
+     "returning %0 risks concurrent access to 'inout sending' parameter %1 as caller assumes %1 and result can be sent to different isolation domains",
+     (Identifier, Identifier))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_return_value, none,
+     "returning result of %kind0 risks concurrent access to 'inout sending' parameter %1 as caller assumes %1 and result can be sent to different isolation domains",
+     (const ValueDecl *, Identifier))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_actor_param, none,
+     "returning %0 risks concurrent access as caller assumes %0 is not actor-isolated and result is %1", (Identifier, StringRef))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_actor_value, none,
+     "returning %0 risks concurrent access to 'inout sending' parameter %1 as caller assumes %1 is not actor-isolated and result is %2",
+     (Identifier, Identifier, StringRef))
+NOTE(regionbasedisolation_inout_sending_cannot_be_returned_note_actor_return_value, none,
+     "returning result of %kind0 risks concurrent access to 'inout sending' parameter %1 as caller assumes %1 is not actor-isolated and result is %2",
+     (const ValueDecl *, Identifier, StringRef))
 
 //===
 // Out Sending

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -544,6 +544,10 @@ class RegionAnalysisFunctionInfo {
 
   bool isSupportedFunction() const { return supportedFunction; }
 
+  NullablePtr<BlockPartitionState> getBlockState(SILBasicBlock *block) const {
+    return blockStates->get(block);
+  }
+
   using iterator = BasicBlockData::iterator;
   using const_iterator = BasicBlockData::const_iterator;
   using reverse_iterator = BasicBlockData::reverse_iterator;

include/swift/SILOptimizer/Utils/PartitionOpError.def

@@ -52,6 +52,9 @@ PARTITION_OP_ERROR(InOutSendingNotInitializedAtExit)
 /// at end of function without being reinitialized with something disconnected.
 PARTITION_OP_ERROR(InOutSendingNotDisconnectedAtExit)
 
+/// This is emitted when a concrete inout sending parameter is returned.
+PARTITION_OP_ERROR(InOutSendingReturned)
+
 /// Used to signify an "unknown code pattern" has occured while performing
 /// dataflow.
 ///

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -1126,6 +1126,40 @@ class PartitionOpError {
     }
   };
 
+  struct InOutSendingReturnedError {
+    const PartitionOp *op;
+
+    /// The 'inout sending' parameter that caused the error to be emitted.
+    Element inoutSendingElement;
+
+    /// The actual returned value. If we return the 'inout sending' parameter
+    /// then this will equal inoutSendingElement.
+    Element returnedValue;
+
+    /// If we are actor isolated due to being in the same region as the out
+    /// parameter of a actor isolated method... this is that actor isolation.
+    SILDynamicMergedIsolationInfo isolationInfo;
+
+    InOutSendingReturnedError(const PartitionOp &op,
+                              Element inoutSendingElement,
+                              Element returnedValue,
+                              SILDynamicMergedIsolationInfo isolationInfo = {})
+        : op(&op), inoutSendingElement(inoutSendingElement),
+          returnedValue(returnedValue), isolationInfo(isolationInfo) {}
+
+    InOutSendingReturnedError(const PartitionOp &op,
+                              Element inoutSendingElement,
+                              SILDynamicMergedIsolationInfo isolationInfo = {})
+        : InOutSendingReturnedError(op, inoutSendingElement,
+                                    inoutSendingElement, isolationInfo) {}
+
+    void print(llvm::raw_ostream &os, RegionAnalysisValueMap &valueMap) const;
+
+    SWIFT_DEBUG_DUMPER(dump(RegionAnalysisValueMap &valueMap)) {
+      print(llvm::dbgs(), valueMap);
+    }
+  };
+
   struct NonSendableIsolationCrossingResultError {
     const PartitionOp *op;
 
@@ -1297,6 +1331,42 @@ struct PartitionOpEvaluator {
     return SILDynamicMergedIsolationInfo();
   }
 
+  /// If \p region is not disconnected only because of an out parameter, return
+  /// that out parameter. We do not care about other non-disconnected parameters
+  /// since we always want to prefer the return error.
+  SILValue findNonDisconnectedOutParameterInRegion(Region region) const {
+    for (const auto &pair : p.range()) {
+      if (pair.second != region)
+        continue;
+      auto info = getIsolationRegionInfo(pair.first);
+
+      // If we do not have any isolation info, then something bad
+      // happened. Return SILValue().
+      if (!info)
+        return {};
+
+      // If we have something that is disconnected... continue. We do not care
+      // about this.
+      if (info.isDisconnected())
+        continue;
+
+      // See if we ahve an indirect out parameter...
+      if (auto value = asImpl().getIndirectOutParameter(pair.first)) {
+        return value;
+      }
+
+      // If this was not an out parameter, return {}.
+      return {};
+    }
+
+    // After going over our loop, if result is not SILValue(), then we found
+    // that our region is not disconnected due only to a single 'indirect out'
+    // parameter. If SILValue() then we had all disconnected values. If we found
+    // multiple non-disconnected things, we would have bailed earlier in the
+    // loop itself.
+    return SILValue();
+  }
+
   bool isTaskIsolatedDerived(Element elt) const {
     return asImpl().isTaskIsolatedDerived(elt);
   }
@@ -1534,7 +1604,9 @@ struct PartitionOpEvaluator {
              "Require PartitionOp's argument should already be tracked");
 
       // First check if the region of our 'inout sending' element has been
-      // sent. In that case, we emit a special use after free error.
+      // sent. In that case, we emit a special use after free error so that the
+      // user knows that they need to reinitialize the 'inout sending'
+      // parameter.
       if (auto *sentOperandSet = p.getSentOperandSet(op.getOpArg1())) {
         for (auto sentOperand : sentOperandSet->data()) {
           handleError(InOutSendingNotInitializedAtExitError(op, op.getOpArg1(),
@@ -1543,7 +1615,7 @@ struct PartitionOpEvaluator {
         return;
       }
 
-      // If we were not sent, check if our region is actor isolated. If so,
+      // If we were not sent, check if our region is not disconnected. If so,
       // error since we need a disconnected value in the inout parameter.
       Region inoutSendingRegion = p.getRegion(op.getOpArg1());
       auto dynamicRegionIsolation = getIsolationRegionInfo(inoutSendingRegion);
@@ -1554,12 +1626,82 @@ struct PartitionOpEvaluator {
         return;
       }
 
-      // Otherwise, emit the error if the dynamic region isolation is not
-      // disconnected.
+      auto handleDirectReturn = [&]() -> bool {
+        bool emittedDiagnostic = false;
+
+        // For each value we are returning...
+        for (auto value : op.getSourceInst()->getOperandValues()) {
+          // Find its element and check if that element is in the same region as
+          // our inout sending param...
+          auto elt = asImpl().getElement(value);
+          if (!elt || inoutSendingRegion != p.getRegion(*elt))
+            continue;
+
+          emittedDiagnostic = true;
+          // Then check if our element is the same as our inoutSendingParam. In
+          // that case, we emit a special error that only refers to the
+          // inoutSendingParam as one entity.
+          //
+          // NOTE: This is taking advantage of us looking through loads when
+          // determining element identity. When that changes, this code will
+          // need to be updated to look through loads.
+          if (*elt == op.getOpArg1()) {
+            handleError(InOutSendingReturnedError(op, op.getOpArg1(),
+                                                  dynamicRegionIsolation));
+            continue;
+          }
+
+          // Otherwise, we need to refer to a different value in the same region
+          // as the 'inout sending' parameter. Emit a special error. For that.
+          handleError(InOutSendingReturnedError(op, op.getOpArg1(), *elt,
+                                                dynamicRegionIsolation));
+        }
+
+        return emittedDiagnostic;
+      };
+
+      // Then check if our region isolated value is not disconnected...
       if (!dynamicRegionIsolation.isDisconnected()) {
+        // Before we emit the more general "inout sending" not disconnected at
+        // exit error... check if our dynamic region isolation is not
+        // disconnected since it is in the same region an out parameter. In that
+        // case we want to emit a special 'cannot return value' error.
+        if (auto outParam =
+                findNonDisconnectedOutParameterInRegion(inoutSendingRegion)) {
+          handleError(InOutSendingReturnedError(op, op.getOpArg1(),
+                                                getElement(outParam).value(),
+                                                dynamicRegionIsolation));
+          return;
+        }
+
+        // If we did not have an out parameter, see if we are returning a
+        // value that is in the region as our 'inout sending' parameter.
+        if (handleDirectReturn())
+          return;
+
+        // Otherwise, we emit the normal not disconnected at exit error.
         handleError(InOutSendingNotDisconnectedAtExitError(
             op, op.getOpArg1(), dynamicRegionIsolation));
+        return;
       }
+
+      // Ok, we have a disconnected value. We could still be returning a
+      // disconnected value in the same region as an 'inout sending'
+      // parameter. We cannot allow that since the caller considers 'inout
+      // sending' values to be in their own region on function return. So it
+      // would assume that it could potentially send the value in the 'inout
+      // sending' parameter and the return value to different isolation
+      // domains... allowing for races.
+
+      // Even though we are disconnected, see if our SILFunction has an actor
+      // isolation. In that case, we want to use that since the direct return
+      // value will be inferred in the caller to have that isolation.
+      if (auto isolation = SILIsolationInfo::getFunctionIsolation(
+              op.getSourceInst()->getFunction())) {
+        dynamicRegionIsolation = {isolation};
+      }
+
+      handleDirectReturn();
       return;
     }
     case PartitionOpKind::UnknownPatternError:
@@ -1569,10 +1711,6 @@ struct PartitionOpEvaluator {
       // Then emit an unknown code pattern error.
       return handleError(UnknownCodePatternError(op));
     case PartitionOpKind::NonSendableIsolationCrossingResult: {
-      // Grab the dynamic dataflow isolation information for our element's
-      // region.
-      Region region = p.getRegion(op.getOpArg1());
-
       // Then emit the error.
       return handleError(
           NonSendableIsolationCrossingResultError(op, op.getOpArg1()));
@@ -1757,6 +1895,14 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
   /// isolated value.
   bool isTaskIsolatedDerived(Element elt) const { return false; }
 
+  /// If \p elt maps to a representative that is an indirect out parameter,
+  /// return that value.
+  SILValue getIndirectOutParameter(Element elt) const { return {}; }
+
+  /// If \p elt maps to a representative that is an 'inout sending' parameter
+  /// that returns that value.
+  SILValue getInOutSendingParameter(Element elt) const { return {}; }
+
   /// By default, do nothing upon error.
   void handleError(PartitionOpError error) {}
 

include/swift/SILOptimizer/Utils/SILIsolationInfo.h

@@ -546,6 +546,11 @@ class SILIsolationInfo {
   /// Infer isolation of conformances for the given instruction.
   static SILIsolationInfo getConformanceIsolation(SILInstruction *inst);
 
+  /// Return SILIsolationInfo based off of the attached ActorIsolation of \p
+  /// fn. If \p fn does not have an actor isolation set, returns an invalid
+  /// SILIsolationInfo.
+  static SILIsolationInfo getFunctionIsolation(SILFunction *fn);
+
   /// A helper that is used to ensure that we treat certain builtin values as
   /// non-Sendable that the AST level otherwise thinks are non-Sendable.
   ///

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -4115,7 +4115,10 @@ bool BlockPartitionState::recomputeExitFromEntry(
     }
 
     std::optional<Element> getElement(SILValue value) const {
-      return translator.getValueMap().getTrackableValue(value).value.getID();
+      auto trackableValue = translator.getValueMap().getTrackableValue(value);
+      if (trackableValue.value.isSendable())
+        return {};
+      return trackableValue.value.getID();
     }
 
     SILValue getRepresentative(SILValue value) const {