[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.

(cherry picked from commit 74ac12c)

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)
@@ -2740,6 +2751,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