Fix DCE of load_borrow when it is derived from another borrow scope

Fixes rdar://138663452

Author: meg-gupta

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -12,12 +12,13 @@
 
 #define DEBUG_TYPE "sil-dce"
 #include "swift/Basic/Assertions.h"
+#include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/MemAccessUtils.h"
+#include "swift/SIL/NodeBits.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBasicBlock.h"
-#include "swift/SIL/BasicBlockBits.h"
-#include "swift/SIL/NodeBits.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILUndef.h"
@@ -158,6 +159,8 @@ class DCE {
   bool CallsChanged = false;
 
   bool precomputeControlInfo();
+  /// Populates borrow dependencies and disables DCE if needed.
+  void processBorrow(BorrowedValue borrow);
   void markLive();
   /// Record a reverse dependency from \p from to \p to meaning \p to is live
   /// if \p from is also live.
@@ -252,6 +255,49 @@ static BuiltinInst *getProducer(CondFailInst *CFI) {
   return nullptr;
 }
 
+void DCE::processBorrow(BorrowedValue borrow) {
+  // Populate guaranteedPhiDependencies for this borrow
+  findGuaranteedPhiDependencies(borrow);
+  if (!borrow.hasReborrow()) {
+    return;
+  }
+
+  // If the borrow was not computed from another
+  // borrow, return.
+  SILValue baseValue;
+  if (auto *beginBorrow = dyn_cast<BeginBorrowInst>(*borrow)) {
+    auto borrowOp = beginBorrow->getOperand();
+    if (borrowOp->getOwnershipKind() != OwnershipKind::Guaranteed) {
+      return;
+    }
+    baseValue = borrowOp;
+  } else {
+    auto *loadBorrow = cast<LoadBorrowInst>(*borrow);
+    auto accessBase = AccessBase::compute(loadBorrow->getOperand());
+    if (!accessBase.isReference()) {
+      return;
+    }
+    baseValue = accessBase.getReference();
+  }
+  // If the borrow was computed from another
+  // borrow, disable DCE of the outer borrow.
+  // This is because, when a reborrow is dead, DCE has to insert
+  // end_borrows in predecessor blocks and it cannot yet handle borrow
+  // nesting.
+  // TODO: Instead of disabling DCE of outer borrow, consider inserting
+  // end_borrows inside-out.
+  SmallVector<SILValue, 4> roots;
+  findGuaranteedReferenceRoots(baseValue,
+                               /*lookThroughNestedBorrows=*/false, roots);
+  // Visit the end_borrows of all the borrow scopes that this
+  // begin_borrow could be borrowing, and mark them live.
+  for (auto root : roots) {
+    visitTransitiveEndBorrows(root, [&](EndBorrowInst *endBorrow) {
+      markInstructionLive(endBorrow);
+    });
+  }
+}
+
 // Determine which instructions from this function we need to keep.
 void DCE::markLive() {
   // Find the initial set of instructions in this function that appear
@@ -327,32 +373,12 @@ void DCE::markLive() {
       }
       case SILInstructionKind::BeginBorrowInst: {
         auto *borrowInst = cast<BeginBorrowInst>(&I);
-        // Populate guaranteedPhiDependencies for this borrowInst
-        findGuaranteedPhiDependencies(BorrowedValue(borrowInst));
-        auto disableBorrowDCE = [&](SILValue borrow) {
-          visitTransitiveEndBorrows(borrow, [&](EndBorrowInst *endBorrow) {
-            markInstructionLive(endBorrow);
-          });
-        };
-        // If we have a begin_borrow of a @guaranteed operand, disable DCE'ing
-        // of parent borrow scopes. Dead reborrows needs complex handling, which
-        // is why it is disabled for now.
-        if (borrowInst->getOperand()->getOwnershipKind() ==
-            OwnershipKind::Guaranteed) {
-          SmallVector<SILValue, 4> roots;
-          findGuaranteedReferenceRoots(borrowInst->getOperand(),
-                                       /*lookThroughNestedBorrows=*/false,
-                                       roots);
-          // Visit the end_borrows of all the borrow scopes that this
-          // begin_borrow could be borrowing, and mark them live.
-          for (auto root : roots) {
-            disableBorrowDCE(root);
-          }
-        }
+        processBorrow(BorrowedValue(borrowInst));
         break;
       }
       case SILInstructionKind::LoadBorrowInst: {
-        findGuaranteedPhiDependencies(BorrowedValue(cast<LoadBorrowInst>(&I)));
+        auto *loadBorrowInst = cast<LoadBorrowInst>(&I);
+        processBorrow(BorrowedValue(loadBorrowInst));
         break;
       }
       default:

test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil

@@ -854,20 +854,7 @@ exit(%outer_lifetime_3 : @guaranteed $Klass, %inner_lifetime_2 : @guaranteed $Kl
 }
 
 // CHECK-LABEL: sil [ossa] @reborrow_guaranteed_phi2 : {{.*}} {
-// CHECK:       {{bb[0-9]+}}([[EITHER:%[^,]+]] : @owned $EitherNoneOrAnyObject):
-// CHECK:         [[BORROW_EITHER:%[^,]+]] = begin_borrow [[EITHER]]
-// CHECK:         switch_enum [[BORROW_EITHER]] : $EitherNoneOrAnyObject, case #EitherNoneOrAnyObject.none!enumelt: [[NONE_BLOCK:bb[0-9]+]], case #EitherNoneOrAnyObject.any!enumelt: [[SOME_BLOCK:bb[0-9]+]]
-// CHECK:       [[SOME_BLOCK]]([[PAYLOAD:%[^,]+]] : @guaranteed $AnyObject):
-// CHECK:         end_borrow [[BORROW_EITHER]]
-// CHECK:         destroy_value [[EITHER]]
-// CHECK:         br [[EXIT:bb[0-9]+]]
-// CHECK:       [[NONE_BLOCK]]:
-// CHECK:         end_borrow [[BORROW_EITHER]]
-// CHECK:         destroy_value [[EITHER]]
-// CHECK:         br [[EXIT]]
-// CHECK:       [[EXIT]]:
-// CHECK:         [[RETVAL:%[^,]+]] = tuple ()
-// CHECK:         return [[RETVAL]]
+// CHECK-NOT: switch_enum
 // CHECK-LABEL: } // end sil function 'reborrow_guaranteed_phi2'
 sil [ossa] @reborrow_guaranteed_phi2 : $@convention(thin) (@owned EitherNoneOrAnyObject) -> () {
 entry(%0 : @owned $EitherNoneOrAnyObject):
@@ -952,6 +939,62 @@ bb1(%4 : @guaranteed $Klass, %5 : @guaranteed $Klass):
   return %8 : $()
 }
 
+class Wrapper {
+  let val: Klass
+}
+
+class DoubleWrapper {
+  let s: StructWrapper
+}
+
+struct StructWrapper {
+  let val: Klass
+}
+
+// CHECK-LABEL: sil [ossa] @reborrow_load_borrow3 : $@convention(method) (@owned Wrapper) -> () {
+// CHECK: load_borrow
+// CHECK: end_borrow
+// CHECK: br bb1
+// CHECK-LABEL: } // end sil function 'reborrow_load_borrow3'
+sil [ossa] @reborrow_load_borrow3 : $@convention(method) (@owned Wrapper) -> () {
+bb0(%0 : @owned $Wrapper):
+  %1 = begin_borrow %0 : $Wrapper
+  %2 = ref_element_addr %1 : $Wrapper, #Wrapper.val
+  %3 = load_borrow %2 : $*Klass
+  %f = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %f(%3) : $@convention(thin) (@guaranteed Klass) -> ()
+  br bb1(%1 : $Wrapper, %3 : $Klass)
+
+bb1(%4 :  @guaranteed $Wrapper, %5 :  @guaranteed $Klass):
+  end_borrow %5 : $Klass
+  end_borrow %4 : $Wrapper
+  destroy_value %0 : $Wrapper
+  %8 = tuple ()
+  return %8 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @reborrow_load_borrow4 : $@convention(method) (@owned DoubleWrapper) -> () {
+// CHECK: load_borrow
+// CHECK: end_borrow
+// CHECK: br bb1
+// CHECK-LABEL: } // end sil function 'reborrow_load_borrow4'
+sil [ossa] @reborrow_load_borrow4 : $@convention(method) (@owned DoubleWrapper) -> () {
+bb0(%0 : @owned $DoubleWrapper):
+  %1 = begin_borrow %0 : $DoubleWrapper
+  %2 = ref_element_addr %1 : $DoubleWrapper, #DoubleWrapper.s
+  %4 = struct_element_addr %2 : $*StructWrapper, #StructWrapper.val
+  %5 = load_borrow %4 : $*Klass
+  %f = function_ref @use_klass : $@convention(thin) (@guaranteed Klass) -> ()
+  apply %f(%5) : $@convention(thin) (@guaranteed Klass) -> ()
+  br bb1(%1 : $DoubleWrapper, %5 : $Klass)
+
+bb1(%6 : @guaranteed $DoubleWrapper, %7 : @guaranteed $Klass):
+  end_borrow %7 : $Klass
+  end_borrow %6 : $DoubleWrapper
+  destroy_value %0 : $DoubleWrapper
+  %8 = tuple ()
+  return %8 : $()
+}
 
 // CHECK-LABEL: sil [ossa] @borrow_guaranteed_tuple : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE_1:%[^,]+]] : @owned $Klass, [[INSTANCE_2:%[^,]+]] : @owned $Klass):