Fix OSSA Outliner for scoped guaranteed values

include/swift/SIL/BasicBlockUtils.h

@@ -91,8 +91,11 @@ class DeadEndBlocks {
   /// the set of reachable blocks.
   void updateForReachableBlock(SILBasicBlock *reachableBB);
 
+  /// Add new blocks to the set of reachable blocks.
+  void updateForNewBlock(SILBasicBlock *newBB);
+
   const SILFunction *getFunction() const { return f; }
-  
+
   /// Performs a simple check if \p block (or its single successor) ends in an
   /// "unreachable".
   ///

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -73,6 +73,11 @@ void extendLocalBorrow(BeginBorrowInst *beginBorrow,
 /// newly created phis do not yet have a borrow scope.
 bool createBorrowScopeForPhiOperands(SILPhiArgument *newPhi);
 
+SILValue
+makeGuaranteedValueAvailable(SILValue value, SILInstruction *user,
+                             DeadEndBlocks &deBlocks,
+                             InstModCallbacks callbacks = InstModCallbacks());
+
 //===----------------------------------------------------------------------===//
 //                        GuaranteedOwnershipExtension
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/BasicBlockUtils.cpp

@@ -404,6 +404,16 @@ void DeadEndBlocks::updateForReachableBlock(SILBasicBlock *reachableBB) {
   propagateNewlyReachableBlocks(numReachable);
 }
 
+void DeadEndBlocks::updateForNewBlock(SILBasicBlock *newBB) {
+  if (!didComputeValue)
+    return;
+
+  assert(reachableBlocks.count(newBB) == 0);
+  unsigned numReachable = reachableBlocks.size();
+  reachableBlocks.insert(newBB);
+  propagateNewlyReachableBlocks(numReachable);
+}
+
 bool DeadEndBlocks::triviallyEndsInUnreachable(SILBasicBlock *block) {
   // Handle the case where a single "unreachable" block (e.g. containing a call
   // to fatalError()), is jumped to from multiple source blocks.

lib/SILOptimizer/Transforms/Outliner.cpp

@@ -19,16 +19,19 @@
 #include "swift/Demangling/Demangler.h"
 #include "swift/Demangling/ManglingMacros.h"
 #include "swift/SIL/ApplySite.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
+#include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/DynamicCasts.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/Support/CommandLine.h"
@@ -118,9 +121,14 @@ namespace {
 class OutlinePattern {
 protected:
   SILOptFunctionBuilder &FuncBuilder;
+  InstModCallbacks callbacks;
+  DeadEndBlocks *deBlocks;
 
 public:
-  OutlinePattern(SILOptFunctionBuilder &FuncBuilder) : FuncBuilder(FuncBuilder) {}
+  OutlinePattern(SILOptFunctionBuilder &FuncBuilder,
+                 InstModCallbacks callbacks,
+                 DeadEndBlocks *deBlocks)
+      : FuncBuilder(FuncBuilder), callbacks(callbacks), deBlocks(deBlocks) {}
 
   /// Match the instruction sequence.
   virtual bool matchInstSequence(SILBasicBlock::iterator I) = 0;
@@ -233,7 +241,10 @@ class BridgedProperty : public OutlinePattern {
   std::pair<SILFunction *, SILBasicBlock::iterator>
   outline(SILModule &M) override;
 
-  BridgedProperty(SILOptFunctionBuilder &FuncBuilder) : OutlinePattern(FuncBuilder) {
+  BridgedProperty(SILOptFunctionBuilder &FuncBuilder,
+                  InstModCallbacks callbacks,
+                  DeadEndBlocks *deBlocks)
+      : OutlinePattern(FuncBuilder, callbacks, deBlocks) {
     clearState();
   }
 
@@ -352,7 +363,9 @@ BridgedProperty::outline(SILModule &M) {
   auto *OutlinedEntryBB = StartBB->split(SILBasicBlock::iterator(FirstInst));
   auto *OldMergeBB = switchInfo.Br->getDestBB();
   auto *NewTailBB = OldMergeBB->split(OldMergeBB->begin());
-
+  if (deBlocks) {
+    deBlocks->updateForNewBlock(NewTailBB);
+  }
   // Call the outlined function.
   {
     SILBuilder Builder(StartBB);
@@ -900,8 +913,11 @@ namespace {
 //
 // bb3(%32 : $Optional<String>):
 class BridgedReturn {
+  DeadEndBlocks *deBlocks;
   SwitchInfo switchInfo;
 public:
+  BridgedReturn(DeadEndBlocks *deBlocks) : deBlocks(deBlocks) {}
+
   bool match(ApplyInst *BridgedCall) {
     switchInfo = SwitchInfo();
     auto *SwitchBB = BridgedCall->getParent();
@@ -941,7 +957,10 @@ void BridgedReturn::outline(SILFunction *Fun, ApplyInst *NewOutlinedCall) {
   auto *OutlinedEntryBB = StartBB->split(SILBasicBlock::iterator(switchInfo.SwitchEnum));
   auto *OldMergeBB = switchInfo.Br->getDestBB();
   auto *NewTailBB = OldMergeBB->split(OldMergeBB->begin());
-	auto Loc = switchInfo.SwitchEnum->getLoc();
+  if (deBlocks) {
+    deBlocks->updateForNewBlock(NewTailBB);
+  }
+  auto Loc = switchInfo.SwitchEnum->getLoc();
 
   {
     SILBuilder Builder(StartBB);
@@ -992,8 +1011,10 @@ class ObjCMethodCall : public OutlinePattern {
   std::pair<SILFunction *, SILBasicBlock::iterator>
   outline(SILModule &M) override;
 
-  ObjCMethodCall(SILOptFunctionBuilder &FuncBuilder)
-      : OutlinePattern(FuncBuilder) {}
+  ObjCMethodCall(SILOptFunctionBuilder &FuncBuilder,
+                 InstModCallbacks callbacks,
+                 DeadEndBlocks *deBlocks)
+      : OutlinePattern(FuncBuilder, callbacks, deBlocks), BridgedReturn(deBlocks) {}
   ~ObjCMethodCall();
 
 private:
@@ -1043,7 +1064,12 @@ ObjCMethodCall::outline(SILModule &M) {
     for (auto Arg : BridgedCall->getArguments()) {
       if (BridgedArgIdx < BridgedArguments.size() &&
           BridgedArguments[BridgedArgIdx].Idx == OrigSigIdx) {
-        Args.push_back(BridgedArguments[BridgedArgIdx].bridgedValue());
+        auto bridgedArgValue = BridgedArguments[BridgedArgIdx].bridgedValue();
+        if (bridgedArgValue.getOwnershipKind() == OwnershipKind::Guaranteed) {
+          bridgedArgValue = makeGuaranteedValueAvailable(
+              bridgedArgValue, BridgedCall, *deBlocks);
+        }
+        Args.push_back(bridgedArgValue);
         ++BridgedArgIdx;
       } else {
         // Otherwise, use the original type convention.
@@ -1263,9 +1289,11 @@ class OutlinePatterns {
     return nullptr;
   }
 
-  OutlinePatterns(SILOptFunctionBuilder &FuncBuilder)
-      : BridgedPropertyPattern(FuncBuilder),
-        ObjCMethodCallPattern(FuncBuilder) {}
+  OutlinePatterns(SILOptFunctionBuilder &FuncBuilder,
+                  InstModCallbacks callbacks,
+                  DeadEndBlocks *deBlocks)
+      : BridgedPropertyPattern(FuncBuilder, callbacks, deBlocks),
+        ObjCMethodCallPattern(FuncBuilder, callbacks, deBlocks) {}
   ~OutlinePatterns() {}
 
   OutlinePatterns(const OutlinePatterns&) = delete;
@@ -1277,9 +1305,11 @@ class OutlinePatterns {
 /// Perform outlining on the function and return any newly created outlined
 /// functions.
 bool tryOutline(SILOptFunctionBuilder &FuncBuilder, SILFunction *Fun,
-                SmallVectorImpl<SILFunction *> &FunctionsAdded) {
+                SmallVectorImpl<SILFunction *> &FunctionsAdded,
+                InstModCallbacks callbacks = InstModCallbacks(),
+                DeadEndBlocks *deBlocks = nullptr) {
   BasicBlockWorklist Worklist(Fun->getEntryBlock());
-  OutlinePatterns patterns(FuncBuilder);
+  OutlinePatterns patterns(FuncBuilder, callbacks, deBlocks);
   bool changed = false;
 
   // Traverse the function.
@@ -1329,9 +1359,21 @@ class Outliner : public SILFunctionTransform {
       Fun->dump();
     }
 
+    DeadEndBlocksAnalysis *deBlocksAnalysis =
+        PM->getAnalysis<DeadEndBlocksAnalysis>();
+    DeadEndBlocks *deBlocks = deBlocksAnalysis->get(Fun);
+    InstModCallbacks callbacks;
+
     SILOptFunctionBuilder FuncBuilder(*this);
     SmallVector<SILFunction *, 16> FunctionsAdded;
-    bool Changed = tryOutline(FuncBuilder, Fun, FunctionsAdded);
+    bool Changed = false;
+
+    if (Fun->hasOwnership()) {
+      Changed =
+          tryOutline(FuncBuilder, Fun, FunctionsAdded, callbacks, deBlocks);
+    } else {
+      Changed = tryOutline(FuncBuilder, Fun, FunctionsAdded);
+    }
 
     if (!FunctionsAdded.empty()) {
       // Notify the pass manager of any new functions we outlined.

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -643,6 +643,12 @@ struct OwnershipLifetimeExtender {
                                SILBasicBlock::iterator borrowPoint,
                                RangeTy guaranteedUsePoints);
 
+  template <typename RangeTy>
+  BeginBorrowInst *
+  borrowCopyOverGuaranteedUsers(SILValue newValue,
+                                SILBasicBlock::iterator borrowPoint,
+                                RangeTy guaranteedUsers);
+
   /// Borrow \p newValue over the lifetime of \p guaranteedValue. Return the
   /// new guaranteed value.
   SILValue borrowOverValue(SILValue newValue, SILValue guaranteedValue);
@@ -654,6 +660,10 @@ struct OwnershipLifetimeExtender {
   /// the BorrowedValue that begins the scope.
   SILValue borrowOverSingleUse(SILValue newValue,
                                Operand *singleGuaranteedUse);
+
+  SILValue
+  borrowOverSingleNonLifetimeEndingUser(SILValue newValue,
+                                        SILInstruction *nonLifetimeEndingUser);
 };
 
 } // end anonymous namespace
@@ -780,9 +790,9 @@ OwnershipLifetimeExtender::borrowCopyOverScope(SILValue newValue,
 ///
 /// Precondition: None of \p guaranteedUses are lifetime ending.
 template <typename RangeTy>
-BeginBorrowInst *OwnershipLifetimeExtender::borrowCopyOverGuaranteedUses(
+BeginBorrowInst *OwnershipLifetimeExtender::borrowCopyOverGuaranteedUsers(
     SILValue newValue, SILBasicBlock::iterator borrowPoint,
-    RangeTy guaranteedUsePoints) {
+    RangeTy guaranteedUsers) {
 
   auto loc = RegularLocation::getAutoGeneratedLocation(borrowPoint->getLoc());
 
@@ -794,12 +804,11 @@ BeginBorrowInst *OwnershipLifetimeExtender::borrowCopyOverGuaranteedUses(
 
   // Create end_borrows at the end of the borrow's lifetime.
   {
-    // We don't expect an empty guaranteedUsePoints. If it happens, then the
+    // We don't expect an empty guaranteedUsers. If it happens, then the
     // newly created copy will never be destroyed.
-    assert(!guaranteedUsePoints.empty());
+    assert(!guaranteedUsers.empty());
 
-    auto opRange = makeUserRange(guaranteedUsePoints);
-    ValueLifetimeAnalysis lifetimeAnalysis(borrow, opRange);
+    ValueLifetimeAnalysis lifetimeAnalysis(borrow, guaranteedUsers);
     ValueLifetimeBoundary borrowBoundary;
     lifetimeAnalysis.computeLifetimeBoundary(borrowBoundary);
 
@@ -834,6 +843,14 @@ BeginBorrowInst *OwnershipLifetimeExtender::borrowCopyOverGuaranteedUses(
   return borrow;
 }
 
+template <typename RangeTy>
+BeginBorrowInst *OwnershipLifetimeExtender::borrowCopyOverGuaranteedUses(
+    SILValue newValue, SILBasicBlock::iterator borrowPoint,
+    RangeTy guaranteedUsePoints) {
+  return borrowCopyOverGuaranteedUsers(newValue, borrowPoint,
+                                       makeUserRange(guaranteedUsePoints));
+}
+
 // Return the borrow position when replacing guaranteedValue with newValue.
 //
 // Precondition: newValue's block dominates and reaches guaranteedValue's block.
@@ -944,6 +961,27 @@ OwnershipLifetimeExtender::borrowOverSingleUse(SILValue newValue,
   return newBeginBorrow;
 }
 
+SILValue OwnershipLifetimeExtender::borrowOverSingleNonLifetimeEndingUser(
+    SILValue newValue, SILInstruction *nonLifetimeEndingUser) {
+  // Avoid borrowing guaranteed function arguments.
+  if (isa<SILFunctionArgument>(newValue) &&
+      newValue.getOwnershipKind() == OwnershipKind::Guaranteed) {
+    return newValue;
+  }
+  auto borrowPt = newValue->getNextInstruction()->getIterator();
+  return borrowCopyOverGuaranteedUsers(
+      newValue, borrowPt, ArrayRef<SILInstruction *>(nonLifetimeEndingUser));
+}
+
+SILValue swift::makeGuaranteedValueAvailable(SILValue value,
+                                             SILInstruction *user,
+                                             DeadEndBlocks &deBlocks,
+                                             InstModCallbacks callbacks) {
+  OwnershipFixupContext ctx{callbacks, deBlocks};
+  OwnershipLifetimeExtender extender{ctx};
+  return extender.borrowOverSingleNonLifetimeEndingUser(value, user);
+}
+
 //===----------------------------------------------------------------------===//
 //                OwnershipRAUWUtility - RAUW + fix ownership
 //===----------------------------------------------------------------------===//