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

We want 'inout sending' parameters to have the semantics that not only are they
disconnected on return from the function but additionally they are guaranteed to
be in their own disconnected region on return. This implies that we must emit
errors when an 'inout sending' parameter or any element that is in the same
region as the current value within an 'inout sending' parameter is
returned. This commit contains a new diagnostic for RegionIsolation that adds
specific logic for detecting and emitting errors in these situations.

To implement this, we introduce 3 new diagnostics with each individual
diagnostic being slightly different to reflect the various ways that this error
can come up in source:

* Returning 'inout sending' directly:

```swift
func returnInOutSendingDirectly(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
  return x // expected-warning {{cannot return 'inout sending' parameter 'x' from global function 'returnInOutSendingDirectly'}}
  // expected-note @-1 {{returning 'x' risks concurrent access since caller assumes that 'x' and the result of global function 'returnInOutSendingDirectly' can be safely sent to different isolation domains}}
}
```

* Returning a value in the same region as an 'inout sending' parameter. E.x.:

```swift
func returnInOutSendingRegionVar(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
  var y = x
  y = x
  return y // expected-warning {{cannot return 'y' from global function 'returnInOutSendingRegionVar'}}
  // expected-note @-1 {{returning 'y' risks concurrent access to 'inout sending' parameter 'x' since the caller assumes that 'x' and the result of global function 'returnInOutSendingRegionVar' can be safely sent to different isolation domains}}
}
```

* Returning the result of a function or computed property that is in the same
region as the 'inout parameter'.

```swift
func returnInOutSendingViaHelper(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
  let y = x
  return useNonSendableKlassAndReturn(y) // expected-warning {{cannot return result of global function 'useNonSendableKlassAndReturn' from global function 'returnInOutSendingViaHelper'}}
  // expected-note @-1 {{returning result of global function 'useNonSendableKlassAndReturn' risks concurrent access to 'inout sending' parameter 'x' since the caller assumes that 'x' and the result of global function 'returnInOutSendingViaHelper' can be safely sent to different isolation domains}}
}
```

Additionally, I had to introduce a specific variant for each of these
diagnostics for cases where due to us being in a method, we are actually in our
caller causing the 'inout sending' parameter to be in the same region as an
actor isolated value:

* Returning 'inout sending' directly:

```swift
extension MyActor {
  func returnInOutSendingDirectly(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
    return x // expected-warning {{cannot return 'inout sending' parameter 'x' from instance method 'returnInOutSendingDirectly'}}
    // expected-note @-1 {{returning 'x' risks concurrent access since caller assumes that 'x' is not actor-isolated and the result of instance method 'returnInOutSendingDirectly' is 'self'-isolated}}
  }
}
```

* Returning a value in the same region as an 'inout sending' parameter. E.x.:

```swift
extension MyActor {
  func returnInOutSendingRegionLet(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
    let y = x
    return y // expected-warning {{cannot return 'y' from instance method 'returnInOutSendingRegionLet'}}
    // expected-note @-1 {{returning 'y' risks concurrent access to 'inout sending' parameter 'x' since the caller assumes that 'x' is not actor-isolated and the result of instance method 'returnInOutSendingRegionLet' is 'self'-isolated}}
  }
}
```

* Returning the result of a function or computed property that is in the same region as the 'inout parameter'.

```swift
extension MyActor {
  func returnInOutSendingViaHelper(_ x: inout sending NonSendableKlass) -> NonSendableKlass {
    let y = x
    return useNonSendableKlassAndReturn(y) // expected-warning {{cannot return result of global function 'useNonSendableKlassAndReturn' from instance method 'returnInOutSendingViaHelper'; this is an error in the Swift 6 language mode}}
    // expected-note @-1 {{returning result of global function 'useNonSendableKlassAndReturn' risks concurrent access to 'inout sending' parameter 'x' since the caller assumes that 'x' is not actor-isolated and the result of instance method 'returnInOutSendingViaHelper' is 'self'-isolated}}
  }
}
```

To implement this, I used two different approaches depending on whether or not
the returned value was generic or not.

* Concrete

In the case where we had a concrete value, I was able to in simple cases emit
diagnostics based off of the values returned by the return inst. In cases where
we phied together results due to multiple results in the same function, we
determine which of the incoming phied values caused the error by grabbing the
exit partition information of each of the incoming value predecessors and seeing
if an InOutSendingAtFunctionExit would emit an error.

* Generic

In the case of generic code, it is a little more interesting since the result is
a value stored in an our parameter instead of being a value directly returned by
a return inst. To work around this, I use PrunedLiveness to determine the last
values stored into the out parameter in the function to avoid having to do a
full dataflow. Then I take the exit blocks where we assign each of those values
and run the same check as we do in the direct phi case to emit the appropriate
error.

rdar://152454571

Author: gottesmm

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 {