Merge pull request swiftlang#22901 from gottesmm/pr-3ccb97d3295cc1c1a850d2e03daab424fcbdfcef

[closure-lifetime-fixup] Use the SSAUpdater rather than memory to ext…

Author: gottesmm

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -21,6 +21,7 @@
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CFG.h"
 #include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 #include "swift/SILOptimizer/Utils/StackNesting.h"
 
 #include "llvm/Support/CommandLine.h"
@@ -88,8 +89,33 @@ static SILInstruction *getDeinitSafeClosureDestructionPoint(TermInst *Term) {
   return Term;
 }
 
+static void findReachableExitBlocks(SILInstruction *i,
+                                    SmallVectorImpl<SILBasicBlock *> &result) {
+  SmallVector<SILBasicBlock *, 32> worklist;
+  SmallPtrSet<SILBasicBlock *, 8> visitedBlocks;
+
+  visitedBlocks.insert(i->getParent());
+  worklist.push_back(i->getParent());
+
+  while (!worklist.empty()) {
+    auto *bb = worklist.pop_back_val();
+    if (bb->getTerminator()->isFunctionExiting()) {
+      result.push_back(bb);
+      continue;
+    }
+    copy_if(bb->getSuccessorBlocks(), std::back_inserter(worklist),
+            [&](SILBasicBlock *bb) { return visitedBlocks.insert(bb).second; });
+  }
+}
+
 /// Extend the lifetime of the convert_escape_to_noescape's operand to the end
 /// of the function.
+///
+/// NOTE: Since we are lifetime extending a copy that we have introduced, we do
+/// not need to consider destroy_value emitted by SILGen unlike
+/// copy_block_without_escaping which consumes its sentinel parameter. Unlike
+/// that case where we have to consider that destroy_value, we have a simpler
+/// time here.
 static void extendLifetimeToEndOfFunction(SILFunction &Fn,
                                           ConvertEscapeToNoEscapeInst *Cvt) {
   auto EscapingClosure = Cvt->getOperand();
@@ -104,40 +130,78 @@ static void extendLifetimeToEndOfFunction(SILFunction &Fn,
   Cvt->eraseFromParent();
   Cvt = NewCvt;
 
-  // Create an alloc_stack Optional<() -> ()> at the beginning of the function.
-  AllocStackInst *Slot;
-  auto &Context = Cvt->getModule().getASTContext();
-  {
-    SILBuilderWithScope B(Fn.getEntryBlock()->begin());
-    Slot = B.createAllocStack(loc, OptionalEscapingClosureTy);
-    auto *NoneDecl = Context.getOptionalNoneDecl();
-    // Store None to it.
-    B.createStore(
-        loc, B.createEnum(loc, SILValue(), NoneDecl, OptionalEscapingClosureTy),
-        Slot, StoreOwnershipQualifier::Init);
-  }
-  // Insert a copy before the convert_escape_to_noescape and store it to the
-  // alloc_stack location.
-  {
+  // If our Cvt is in the initial block, we do not need to use the SSA updater
+  // since we know Cvt can not be in a loop and must dominate all exits
+  // (*). Just insert a copy of the escaping closure at the Cvt and destroys at
+  // the exit blocks of the function.
+  //
+  // (*) In fact we can't use the SILSSAUpdater::GetValueInMiddleOfBlock.
+  if (Cvt->getParent() == Cvt->getFunction()->getEntryBlock()) {
     SILBuilderWithScope B(Cvt);
-    auto *SomeDecl = Context.getOptionalSomeDecl();
-    B.createDestroyAddr(loc, Slot);
-    auto ClosureCopy = B.createCopyValue(loc, EscapingClosure);
-    B.createStore(
-        loc,
-        B.createEnum(loc, ClosureCopy, SomeDecl, OptionalEscapingClosureTy),
-        Slot, StoreOwnershipQualifier::Init);
+    CopyValueInst *InnerCVI = B.createCopyValue(loc, EscapingClosure);
+    SmallVector<SILBasicBlock *, 4> ExitingBlocks;
+    Fn.findExitingBlocks(ExitingBlocks);
+
+    for (auto *Exit : ExitingBlocks) {
+      auto *Term = Exit->getTerminator();
+      auto *SafeClosureDestructionPt =
+          getDeinitSafeClosureDestructionPoint(Term);
+      SILBuilderWithScope B(SafeClosureDestructionPt);
+      B.createDestroyValue(loc, InnerCVI);
+    }
+    return;
   }
-  // Insert destroys at the function exits.
+
+  // Ok. At this point we know that Cvt is not in the entry block... so we can
+  // use SILSSAUpdater::GetValueInMiddleOfBlock() to extend the object's
+  // lifetime respecting loops.
+  SILSSAUpdater Updater(Cvt->getModule());
+  Updater.Initialize(OptionalEscapingClosureTy);
+
+  // Create an Optional<() -> ()>.none in the entry block of the function and
+  // add it as an available value to the SSAUpdater.
+  //
+  // Since we know that Cvt is not in the entry block and this must be, we know
+  // that it is safe to use the SSAUpdater's getValueInMiddleOfBlock with this
+  // value.
+  Updater.AddAvailableValue(Fn.getEntryBlock(), [&]() -> SILValue {
+    SILBuilderWithScope B(Fn.getEntryBlock()->begin());
+    return B.createOptionalNone(loc, OptionalEscapingClosureTy);
+  }());
+
+  // Create a copy of the convert_escape_to_no_escape and add it as an available
+  // value to the SSA updater.
+  //
+  // NOTE: The SSAUpdater does not support providing multiple values in the same
+  // block without extra work. So the fact that Cvt is not in the entry block
+  // means that we don't have to worry about overwriting the .none value.
+  auto *CVI = [&]() -> CopyValueInst * {
+    SILBuilderWithScope B(Cvt);
+    CopyValueInst *InnerCVI = B.createCopyValue(loc, EscapingClosure);
+    Updater.AddAvailableValue(
+        Cvt->getParent(),
+        B.createOptionalSome(loc, InnerCVI, OptionalEscapingClosureTy));
+    return InnerCVI;
+  }();
+
+  // Then we use the SSA updater to insert a destroy_value before the cvt and at
+  // the reachable exit blocks.
   SmallVector<SILBasicBlock *, 4> ExitingBlocks;
-  Fn.findExitingBlocks(ExitingBlocks);
+  findReachableExitBlocks(Cvt, ExitingBlocks);
+
+  {
+    // Before the copy value, insert an extra destroy_value to handle
+    // loops. Since we used our enum value this is safe.
+    SILBuilderWithScope B(CVI);
+    B.createDestroyValue(loc,
+                         Updater.GetValueInMiddleOfBlock(CVI->getParent()));
+  }
+
   for (auto *Exit : ExitingBlocks) {
     auto *Term = Exit->getTerminator();
     auto *SafeClosureDestructionPt = getDeinitSafeClosureDestructionPoint(Term);
     SILBuilderWithScope B(SafeClosureDestructionPt);
-    B.createDestroyAddr(loc, Slot);
-    SILBuilderWithScope B2(Term);
-    B2.createDeallocStack(loc, Slot);
+    B.createDestroyValue(loc, Updater.GetValueAtEndOfBlock(Exit));
   }
 }
 
@@ -538,88 +602,171 @@ static SILInstruction *getOnlyDestroy(CopyBlockWithoutEscapingInst *CB) {
 ///      %e = is_escaping %closure
 ///      cond_fail %e
 ///      destroy_value %closure
-static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *CB) {
-  SmallVector<SILInstruction *, 4> LifetimeEndPoints;
-
+static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *CB,
+                                          bool &modifiedCFG) {
   // Find the end of the lifetime of the copy_block_without_escaping
   // instruction.
   auto &Fn  = *CB->getFunction();
+
+  // If we find a single destroy, this destroy is going to be a destroy that may
+  // be in the same block as CB. It is important that we make sure that the
+  // destroy is in a different block than CB or any terminating blocks to ensure
+  // that we can use the SSAUpdater if needed.
   auto *SingleDestroy = getOnlyDestroy(CB);
-  SmallVector<SILBasicBlock *, 4> ExitingBlocks;
-  Fn.findExitingBlocks(ExitingBlocks);
-  if (SingleDestroy) {
-    LifetimeEndPoints.push_back(
-        &*std::next(SILBasicBlock::iterator(SingleDestroy)));
-  } else {
-    // Otherwise, conservatively insert verification at the end of the function.
-    for (auto *Exit : ExitingBlocks)
-      LifetimeEndPoints.push_back(Exit->getTerminator());
+  if (SingleDestroy && SingleDestroy->getParent() == CB->getParent()) {
+    modifiedCFG = true;
+    {
+      SILBuilderWithScope B(SingleDestroy);
+      splitBasicBlockAndBranch(B, SingleDestroy, nullptr, nullptr);
+    }
+
+    {
+      SILBuilderWithScope B(SingleDestroy);
+      auto *Term = SingleDestroy->getParent()->getTerminator();
+      if (Term->isFunctionExiting()) {
+        splitBasicBlockAndBranch(B, &*std::next(SingleDestroy->getIterator()),
+                                 nullptr, nullptr);
+      }
+    }
   }
 
   auto SentinelClosure = CB->getClosure();
   auto Loc = CB->getLoc();
 
+  // At this point, we transform our copy_block_without_escaping into a
+  // copy_block. This has a few important implications:
+  //
+  // 1. copy_block_without_escaping takes the sentinel value at +1. We will need
+  //    to balance that +1.
+  // 2. The destroy_value associated with the copy_block_without_escaping will
+  //    be on the copy_block value.
   SILBuilderWithScope B(CB);
   auto *NewCB = B.createCopyBlock(Loc, CB->getBlock());
   CB->replaceAllUsesWith(NewCB);
   CB->eraseFromParent();
 
-  // Create an stack slot for the closure sentinel and store the sentinel (or
-  // none on other paths.
   auto generatedLoc = RegularLocation::getAutoGeneratedLocation();
-  AllocStackInst *Slot;
-  auto &Context = NewCB->getModule().getASTContext();
+
+  // If CB is in the entry block, we know that our definition of SentinelClosure
+  // must be as well. Thus we know that we do not need to worry about loops or
+  // dominance issues and can just insert destroy_values for the sentinel at the
+  // lifetime end points.
+  if (NewCB->getParent() == NewCB->getFunction()->getEntryBlock()) {
+    // Our single destroy must not be in the entry block since if so, we would
+    // have inserted an edge to appease the SSA updater.
+    if (SingleDestroy) {
+      SILBuilderWithScope B(std::next(SingleDestroy->getIterator()));
+      SILValue V = SentinelClosure;
+      SILValue isEscaping = B.createIsEscapingClosure(
+          Loc, V, IsEscapingClosureInst::ObjCEscaping);
+      B.createCondFail(Loc, isEscaping);
+      B.createDestroyValue(Loc, V);
+      return true;
+    }
+
+    // If we couldn't find a specific destroy_value, lifetime extend to the end
+    // of the function.
+    SmallVector<SILBasicBlock *, 4> ExitingBlocks;
+    Fn.findExitingBlocks(ExitingBlocks);
+    for (auto *Block : ExitingBlocks) {
+      SILBuilderWithScope B(Block->getTerminator());
+      SILValue V = SentinelClosure;
+      SILValue isEscaping = B.createIsEscapingClosure(
+          Loc, V, IsEscapingClosureInst::ObjCEscaping);
+      B.createCondFail(Loc, isEscaping);
+      B.createDestroyValue(Loc, V);
+    }
+
+    return true;
+  }
+
+  // Otherwise, we need to be more careful since we can have loops and may not
+  // transitively dominate all uses of the closure. So we:
+  //
+  // 1. Create an Optional<T>.none at the entry.
+  // 2. Create a destroy_value(val), val = Optional<T>.some(sentinel) in the cvt
+  // block.
+  // 3. Create a destroy_value at all exits of the value.
+  //
+  // and then use the SSAUpdater to ensure that we handle loops correctly.
   auto OptionalEscapingClosureTy =
       SILType::getOptionalType(SentinelClosure->getType());
-  auto *NoneDecl = Context.getOptionalNoneDecl();
+
+  SILSSAUpdater Updater(Fn.getModule());
+  Updater.Initialize(OptionalEscapingClosureTy);
+
+  // Create the Optional.none as the beginning available value.
   {
     SILBuilderWithScope B(Fn.getEntryBlock()->begin());
-    Slot = B.createAllocStack(generatedLoc, OptionalEscapingClosureTy);
-    // Store None to it.
-    B.createStore(generatedLoc,
-                  B.createEnum(generatedLoc, SILValue(), NoneDecl,
-                               OptionalEscapingClosureTy),
-                  Slot, StoreOwnershipQualifier::Init);
+    Updater.AddAvailableValue(
+        Fn.getEntryBlock(),
+        B.createOptionalNone(generatedLoc, OptionalEscapingClosureTy));
   }
-  {
+
+  // Then create the Optional.some(closure sentinel).
+  //
+  // NOTE: We return the appropriate insertion point to insert the destroy_value
+  // before the value (to ensure we handle loops). We need to get all available
+  // values first though.
+  auto *InitialValue = [&]() -> EnumInst * {
     SILBuilderWithScope B(NewCB);
-    // Store the closure sentinel (the copy_block_without_escaping closure
+    // Create the closure sentinel (the copy_block_without_escaping closure
     // operand consumed at +1, so we don't need a copy) to it.
-    B.createDestroyAddr(generatedLoc, Slot); // We could be in a loop.
-    auto *SomeDecl = Context.getOptionalSomeDecl();
-    B.createStore(generatedLoc,
-                  B.createEnum(generatedLoc, SentinelClosure, SomeDecl,
-                               OptionalEscapingClosureTy),
-                  Slot, StoreOwnershipQualifier::Init);
+    auto *Result = B.createOptionalSome(generatedLoc, SentinelClosure,
+                                        OptionalEscapingClosureTy);
+    Updater.AddAvailableValue(Result->getParent(), Result);
+    return Result;
+  }();
+
+  // If we had a single destroy, creating a .none after it and add that as a
+  // value to the SSA updater.
+  if (SingleDestroy) {
+    SILBuilderWithScope B(std::next(SingleDestroy->getIterator()));
+    auto *Result =
+        B.createOptionalNone(generatedLoc, OptionalEscapingClosureTy);
+    Updater.AddAvailableValue(Result->getParent(), Result);
   }
 
-  for (auto LifetimeEndPoint : LifetimeEndPoints) {
-    SILBuilderWithScope B(LifetimeEndPoint);
-    SILValue isEscaping;
-    B.emitScopedBorrowOperation(generatedLoc, Slot, [&](SILValue value) {
-      isEscaping = B.createIsEscapingClosure(
-          Loc, value, IsEscapingClosureInst::ObjCEscaping);
-    });
+  // Now that we have all of our available values, insert a destroy_value before
+  // the initial Optional.some value using the SSA updater to ensure that we
+  // handle loops correctly.
+  {
+    SILBuilderWithScope B(InitialValue);
+    B.createDestroyValue(generatedLoc, Updater.GetValueInMiddleOfBlock(
+                                           InitialValue->getParent()));
+  }
+
+  // And insert an is_escaping_closure, cond_fail, destroy_value at each of the
+  // lifetime end points. This ensures we do not expand our lifetime too much.
+  if (SingleDestroy) {
+    SILBuilderWithScope B(std::next(SingleDestroy->getIterator()));
+    SILValue V = Updater.GetValueInMiddleOfBlock(SingleDestroy->getParent());
+    SILValue isEscaping =
+        B.createIsEscapingClosure(Loc, V, IsEscapingClosureInst::ObjCEscaping);
     B.createCondFail(Loc, isEscaping);
-    B.createDestroyAddr(generatedLoc, Slot);
-    // Store None to it.
-    B.createStore(generatedLoc,
-                  B.createEnum(generatedLoc, SILValue(), NoneDecl,
-                               OptionalEscapingClosureTy),
-                  Slot, StoreOwnershipQualifier::Init);
+    B.createDestroyValue(Loc, V);
   }
 
-  // Insert the dealloc_stack in all exiting blocks.
-  for (auto *ExitBlock: ExitingBlocks) {
-    auto *Terminator = ExitBlock->getTerminator();
-    SILBuilderWithScope B(Terminator);
-    B.createDeallocStack(generatedLoc, Slot);
+  // Then to be careful with regards to loops, insert at each of the destroy
+  // blocks destroy_value to ensure that we obey ownership invariants.
+  {
+    SmallVector<SILBasicBlock *, 4> ExitingBlocks;
+    findReachableExitBlocks(NewCB, ExitingBlocks);
+
+    for (auto *Exit : ExitingBlocks) {
+      auto *Term = Exit->getTerminator();
+      auto *SafeClosureDestructionPt =
+          getDeinitSafeClosureDestructionPoint(Term);
+      SILBuilderWithScope B(SafeClosureDestructionPt);
+      B.createDestroyValue(generatedLoc, Updater.GetValueInMiddleOfBlock(Exit));
+    }
   }
 
   return true;
 }
 
-static bool fixupClosureLifetimes(SILFunction &Fn, bool &checkStackNesting) {
+static bool fixupClosureLifetimes(SILFunction &Fn, bool &checkStackNesting,
+                                  bool &modifiedCFG) {
   bool Changed = false;
 
   // tryExtendLifetimeToLastUse uses a cache of recursive instruction use
@@ -634,7 +781,7 @@ static bool fixupClosureLifetimes(SILFunction &Fn, bool &checkStackNesting) {
 
       // Handle, copy_block_without_escaping instructions.
       if (auto *CB = dyn_cast<CopyBlockWithoutEscapingInst>(Inst)) {
-        Changed |= fixupCopyBlockWithoutEscaping(CB);
+        Changed |= fixupCopyBlockWithoutEscaping(CB, modifiedCFG);
         continue;
       }
 
@@ -688,7 +835,7 @@ class ClosureLifetimeFixup : public SILFunctionTransform {
     bool checkStackNesting = false;
     bool modifiedCFG = false;
 
-    if (fixupClosureLifetimes(*getFunction(), checkStackNesting)) {
+    if (fixupClosureLifetimes(*getFunction(), checkStackNesting, modifiedCFG)) {
       if (checkStackNesting){
         StackNesting SN;
         modifiedCFG =