[region-isolation] Make temporary alloc_stack that we form for returning values from a non-final class field take on the class method's isolation.

The reason why we are doing this is that otherwise, we have that the alloc_stack
formed for the result is disconnected and despite the fact that we merge it into
the actor region of the class method, we do not have that the alloc_stack
specifically is marked when we attempt to squelch Please.

This patch fixes that problem by detecting when an alloc_stack is being used as
a temporary for an out parameter and makes the alloc_stack initially isolated as
appropriate. It only does this in the specific cases where we can pattern match
it which in my limited testing has handled everything.

Author: gottesmm

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -1073,27 +1073,23 @@ struct PartitionOpEvaluator {
       }
       std::tie(transferredRegionIsolation, isClosureCapturedElt) = *pairOpt;
 
-      // Before we do anything, see if our dynamic isolation kind is the same as
-      // the isolation info for our partition op. If they match, this is not a
-      // real transfer operation.
-      //
-      // DISCUSSION: We couldn't not emit this earlier since we needed the
-      // dynamic isolation info of our value.
-      if (auto calleeIsolationInfo = getIsolationInfo(op)) {
-        if (transferredRegionIsolation.hasSameIsolation(calleeIsolationInfo)) {
-          return;
-        }
-      }
-
-      // If we merged anything, we need to handle a transfer
-      // non-transferrable. We pass in the dynamic isolation region info of our
-      // region.
-      if (bool(transferredRegionIsolation) &&
+      // If we merged anything, we need to handle a transfer non-transferrable
+      // unless our value has the same isolation info as our callee.
+      auto calleeIsolationInfo = getIsolationInfo(op);
+      if (!(calleeIsolationInfo &&
+            transferredRegionIsolation.hasSameIsolation(calleeIsolationInfo)) &&
           !transferredRegionIsolation.isDisconnected()) {
         return handleTransferNonTransferrableHelper(op, op.getOpArgs()[0],
                                                     transferredRegionIsolation);
       }
 
+      // Next see if we are disconnected and have the same isolation. In such a
+      // case, we do not transfer since the disconnected value is allowed to be
+      // resued after we return.
+      if (transferredRegionIsolation.isDisconnected() && calleeIsolationInfo &&
+          transferredRegionIsolation.hasSameIsolation(calleeIsolationInfo))
+        return;
+
       // Mark op.getOpArgs()[0] as transferred.
       TransferringOperandState &state = operandToStateMap.get(op.getSourceOp());
       state.isClosureCaptured |= isClosureCapturedElt;

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -29,6 +29,7 @@
 #include "swift/SIL/OperandDatastructures.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/PatternMatch.h"
+#include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
@@ -2018,6 +2019,17 @@ class PartitionOpTranslator {
                                    TinyPtrVector<SILValue>());
   }
 
+  void translateSILAssignFresh(SILValue val, SILIsolationInfo info) {
+    auto v = initializeTrackedValue(val, info);
+    if (!v)
+      return translateSILAssignFresh(val);
+
+    if (!v->second)
+      return translateUnknownPatternError(val);
+
+    return translateSILAssignFresh(v->first.getRepresentative().getValue());
+  }
+
   template <typename Collection>
   void translateSILMerge(SILValue dest, Collection collection) {
     auto trackableDest = tryToTrackValue(dest);
@@ -2386,7 +2398,6 @@ CONSTANT_TRANSLATION(AllocBoxInst, AssignFresh)
 CONSTANT_TRANSLATION(AllocPackInst, AssignFresh)
 CONSTANT_TRANSLATION(AllocRefDynamicInst, AssignFresh)
 CONSTANT_TRANSLATION(AllocRefInst, AssignFresh)
-CONSTANT_TRANSLATION(AllocStackInst, AssignFresh)
 CONSTANT_TRANSLATION(AllocVectorInst, AssignFresh)
 CONSTANT_TRANSLATION(KeyPathInst, AssignFresh)
 CONSTANT_TRANSLATION(FunctionRefInst, AssignFresh)
@@ -2741,6 +2752,40 @@ LOOKTHROUGH_IF_NONSENDABLE_RESULT_AND_OPERAND(UncheckedTakeEnumDataAddrInst)
 // Custom Handling
 //
 
+TranslationSemantics
+PartitionOpTranslator::visitAllocStackInst(AllocStackInst *asi) {
+  // Before we do anything, see if asi is Sendable or if it is non-Sendable,
+  // that it is from a var decl. In both cases, we can just return assign fresh
+  // and exit early.
+  if (!SILIsolationInfo::isNonSendableType(asi) || asi->isFromVarDecl())
+    return TranslationSemantics::AssignFresh;
+
+  // Ok at this point we know that our value is a non-Sendable temporary.
+  auto isolationInfo = SILIsolationInfo::get(asi);
+  if (!bool(isolationInfo)) {
+    return TranslationSemantics::AssignFresh;
+  }
+
+  if (isolationInfo.isDisconnected()) {
+    return TranslationSemantics::AssignFresh;
+  }
+
+  // Ok, we can handle this and have a valid isolation. Initialize the value.
+  auto v = initializeTrackedValue(asi, isolationInfo);
+  if (!v)
+    return TranslationSemantics::AssignFresh;
+
+  // If we already had a value for this alloc_stack (which we shouldn't
+  // ever)... emit an unknown pattern error.
+  if (!v->second) {
+    translateUnknownPatternError(asi);
+    return TranslationSemantics::Special;
+  }
+
+  translateSILAssignFresh(v->first.getRepresentative().getValue());
+  return TranslationSemantics::Special;
+}
+
 TranslationSemantics
 PartitionOpTranslator::visitMoveValueInst(MoveValueInst *mvi) {
   if (mvi->isFromVarDecl())

lib/SILOptimizer/Utils/SILIsolationInfo.cpp

@@ -14,10 +14,12 @@
 
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Expr.h"
+#include "swift/SIL/AddressWalker.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/PatternMatch.h"
 #include "swift/SIL/SILGlobalVariable.h"
+#include "swift/SIL/Test.h"
 #include "swift/SILOptimizer/Utils/VariableNameUtils.h"
 
 using namespace swift;
@@ -141,6 +143,204 @@ bool DeclRefExprAnalysis::compute(Expr *expr) {
   return result;
 }
 
+static SILIsolationInfo
+inferIsolationInfoForTempAllocStack(AllocStackInst *asi) {
+  // We want to search for an alloc_stack that is not from a VarDecl and that is
+  // initially isolated along all paths to the same actor isolation. If they
+  // differ, then we emit a we do not understand error.
+  struct AddressWalkerState {
+    AllocStackInst *asi = nullptr;
+    SmallVector<Operand *, 8> indirectResultUses;
+    llvm::SmallSetVector<SILInstruction *, 8> writes;
+    Operand *sameBlockIndirectResultUses = nullptr;
+  };
+
+  struct AddressWalker final : TransitiveAddressWalker<AddressWalker> {
+    AddressWalkerState &state;
+
+    AddressWalker(AddressWalkerState &state) : state(state) {
+      assert(state.asi);
+    }
+
+    bool visitUse(Operand *use) {
+      // If we do not write to memory, then it is harmless.
+      if (!use->getUser()->mayWriteToMemory())
+        return true;
+
+      if (auto fas = FullApplySite::isa(use->getUser())) {
+        if (fas.isIndirectResultOperand(*use)) {
+          // If our indirect result use is in the same block...
+          auto *parentBlock = state.asi->getParent();
+          if (fas.getParent() == parentBlock) {
+            // If we haven't seen any indirect result use yet... just cache it
+            // and return true.
+            if (!state.sameBlockIndirectResultUses) {
+              state.sameBlockIndirectResultUses = use;
+              return true;
+            }
+
+            // If by walking from the alloc stack to the full apply site, we do
+            // not see the current sameBlockIndirectResultUses, we have a new
+            // newest use.
+            if (llvm::none_of(
+                    llvm::make_range(state.asi->getIterator(),
+                                     fas->getIterator()),
+                    [&](const SILInstruction &inst) {
+                      return &inst ==
+                             state.sameBlockIndirectResultUses->getUser();
+                    })) {
+              state.sameBlockIndirectResultUses = use;
+            }
+            return true;
+          }
+
+          // If not, just stash it into the non-same block indirect result use
+          // array.
+          state.indirectResultUses.push_back(use);
+          return true;
+        }
+      }
+
+      state.writes.insert(use->getUser());
+      return true;
+    }
+  };
+
+  AddressWalkerState state;
+  state.asi = asi;
+  AddressWalker walker(state);
+
+  // If we fail to walk, emit an unknown patten error.
+  if (AddressUseKind::Unknown == std::move(walker).walk(asi)) {
+    return SILIsolationInfo();
+  }
+
+  // If we do not have any indirect result uses... we can just assign fresh.
+  if (!state.sameBlockIndirectResultUses && state.indirectResultUses.empty())
+    return SILIsolationInfo::getDisconnected(false /*isUnsafeNonIsolated*/);
+
+  // Otherwise, lets see if we had a same block indirect result.
+  if (state.sameBlockIndirectResultUses) {
+    // If we do not have any writes in between the alloc stack and the
+    // initializer, then we have a good target. Otherwise, we just return
+    // AssignFresh.
+    if (llvm::none_of(
+            llvm::make_range(
+                asi->getIterator(),
+                state.sameBlockIndirectResultUses->getUser()->getIterator()),
+            [&](SILInstruction &inst) { return state.writes.count(&inst); })) {
+      auto isolationInfo =
+          SILIsolationInfo::get(state.sameBlockIndirectResultUses->getUser());
+      if (isolationInfo) {
+        return isolationInfo;
+      }
+    }
+
+    // If we did not find an isolation info, just do a normal assign fresh.
+    return SILIsolationInfo::getDisconnected(false /*is unsafe non isolated*/);
+  }
+
+  // Check if any of our writes are within the first block. This would
+  // automatically stop our search and we should assign fresh. Since we are
+  // going over the writes here, also setup a writeBlocks set.
+  auto *defBlock = asi->getParent();
+  BasicBlockSet writeBlocks(defBlock->getParent());
+  for (auto *write : state.writes) {
+    if (write->getParent() == defBlock)
+      return SILIsolationInfo::getDisconnected(false /*unsafe non isolated*/);
+    writeBlocks.insert(write->getParent());
+  }
+
+  // Ok, at this point we know that we do not have any indirect result uses in
+  // the def block and also we do not have any writes in that initial
+  // block. This sets us up for our global analysis. Our plan is as follows:
+  //
+  // 1. We are going to create a set of writeBlocks and a map from SILBasicBlock
+  // -> first indirect result block if there isn't a write before it.
+  //
+  // 2. We walk from our def block until we reach the first indirect result
+  // block. We stop processing successor if we find a write block successor that
+  // is not also an indirect result block. This makes sense since we earlier
+  // required that any notates indirect result block do not have any writes in
+  // between the indirect result and the beginning of the block.
+  llvm::SmallDenseMap<SILBasicBlock *, Operand *, 2> blockToOperandMap;
+  for (auto *use : state.indirectResultUses) {
+    // If our indirect result use has a write before it in the block, do not
+    // store it. It cannot be our indirect result initializer.
+    if (writeBlocks.contains(use->getParentBlock()) &&
+        llvm::any_of(
+            use->getParentBlock()->getRangeEndingAtInst(use->getUser()),
+            [&](SILInstruction &inst) { return state.writes.contains(&inst); }))
+      continue;
+
+    // Ok, we now know that there aren't any writes before us in the block. Now
+    // try to insert.
+    auto iter = blockToOperandMap.try_emplace(use->getParentBlock(), use);
+
+    // If we actually inserted, then we are done.
+    if (iter.second) {
+      continue;
+    }
+
+    // Otherwise, if we are before the current value, set us to be the value
+    // instead.
+    if (llvm::none_of(
+            use->getParentBlock()->getRangeEndingAtInst(use->getUser()),
+            [&](const SILInstruction &inst) {
+              return &inst == iter.first->second->getUser();
+            })) {
+      iter.first->getSecond() = use;
+    }
+  }
+
+  // Ok, we now have our data all setup.
+  BasicBlockWorklist worklist(asi->getFunction());
+  for (auto *succBlock : asi->getParentBlock()->getSuccessorBlocks()) {
+    worklist.pushIfNotVisited(succBlock);
+  }
+
+  Operand *targetOperand = nullptr;
+  while (auto *next = worklist.pop()) {
+    // First check if this is one of our target blocks.
+    auto iter = blockToOperandMap.find(next);
+
+    // If this is our target blocks...
+    if (iter != blockToOperandMap.end()) {
+      // If we already have an assigned target block, make sure this is the same
+      // one. If it is, just continue. Otherwise, something happened we do not
+      // understand... assign fresh.
+      if (!targetOperand) {
+        targetOperand = iter->second;
+        continue;
+      }
+
+      if (targetOperand->getParentBlock() == iter->first) {
+        continue;
+      }
+
+      return SILIsolationInfo::getDisconnected(
+          false /*is unsafe non isolated*/);
+    }
+
+    // Otherwise, see if this block is a write block. If so, we have a path to a
+    // write block that does not go through one of our blockToOperandMap
+    // blocks... return assign fresh.
+    if (writeBlocks.contains(next))
+      return SILIsolationInfo::getDisconnected(
+          false /*is unsafe non isolated*/);
+
+    // Otherwise, visit this blocks successors if we have not yet visited them.
+    for (auto *succBlock : next->getSuccessorBlocks()) {
+      worklist.pushIfNotVisited(succBlock);
+    }
+  }
+
+  // At this point, we know that we have a single indirect result use that
+  // dominates all writes and other indirect result uses. We can say that our
+  // alloc_stack temporary is that indirect result use's isolation.
+  return SILIsolationInfo::get(targetOperand->getUser());
+}
+
 SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
   if (auto fas = FullApplySite::isa(inst)) {
     if (auto crossing = fas.getIsolationCrossing()) {
@@ -475,6 +675,12 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
               true /*is nonisolated(unsafe)*/);
         }
       }
+    } else {
+      // Ok, we have a temporary. If it is non-Sendable...
+      if (SILIsolationInfo::isNonSendableType(asi)) {
+        if (auto isolation = inferIsolationInfoForTempAllocStack(asi))
+          return isolation;
+      }
     }
   }
 
@@ -842,3 +1048,28 @@ SILDynamicMergedIsolationInfo::merge(SILIsolationInfo other) const {
   // Otherwise, just return other.
   return other;
 }
+
+//===----------------------------------------------------------------------===//
+//                                MARK: Tests
+//===----------------------------------------------------------------------===//
+
+namespace swift::test {
+
+// Arguments:
+// - SILValue: value to emit a name for.
+// Dumps:
+// - The inferred isolation.
+static FunctionTest
+    IsolationInfoInferrence("sil-isolation-info-inference",
+                            [](auto &function, auto &arguments, auto &test) {
+                              auto value = arguments.takeValue();
+
+                              SILIsolationInfo info =
+                                  SILIsolationInfo::get(value);
+                              llvm::outs() << "Input Value: " << *value;
+                              llvm::outs() << "Isolation: ";
+                              info.printForOneLineLogging(llvm::outs());
+                              llvm::outs() << "\n";
+                            });
+
+} // namespace swift::test