[CopyPropagation] Hoist lexical end_borrows over some applies.

Author: nate-chandler

lib/SILOptimizer/Utils/ShrinkBorrowScope.cpp

@@ -54,6 +54,8 @@ class ShrinkBorrowScope {
   SmallPtrSet<SILBasicBlock *, 8> blocksWithReachedTops;
   SmallPtrSet<SILBasicBlock *, 8> blocksToEndAtTop;
 
+  llvm::SmallDenseMap<ApplySite, size_t> transitiveUsesPerApplySite;
+
   // The list of blocks to look for new points at which to insert end_borrows
   // in.  A block must not be processed if all of its successors have not yet
   // been.  For that reason, it is necessary to allow the same block to be
@@ -81,6 +83,12 @@ class ShrinkBorrowScope {
     });
   }
 
+  bool isBarrierApply(SILInstruction *instruction) {
+    // For now, treat every apply (that doesn't use the borrowed value) as a
+    // barrier.
+    return isa<ApplySite>(instruction);
+  }
+
   bool mayAccessPointer(SILInstruction *instruction) {
     if (!instruction->mayReadOrWriteMemory())
       return false;
@@ -100,9 +108,70 @@ class ShrinkBorrowScope {
   }
 
   bool isDeinitBarrier(SILInstruction *instruction) {
-    return users.contains(instruction) || instruction->maySynchronize() ||
-           mayAccessPointer(instruction, value) ||
-           mayLoadWeakOrUnowned(instruction, value);
+    return isBarrierApply(instruction) || instruction->maySynchronize() ||
+           mayAccessPointer(instruction) || mayLoadWeakOrUnowned(instruction);
+  }
+
+  bool canReplaceValueWithBorrowedValue(SILValue value) {
+    while (true) {
+      auto *instruction = value.getDefiningInstruction();
+      if (!instruction)
+        return false;
+      if (auto *cvi = dyn_cast<CopyValueInst>(instruction)) {
+        value = cvi->getOperand();
+        continue;
+      } else if (auto *bbi = dyn_cast<BeginBorrowInst>(instruction)) {
+        if (bbi == introducer) {
+          return true;
+        }
+      }
+      return false;
+    }
+  }
+
+  size_t usesInApply(ApplySite apply) {
+    if (auto count = transitiveUsesPerApplySite.lookup(apply))
+      return count;
+    return 0;
+  }
+
+  bool tryHoistOverInstruction(SILInstruction *instruction) {
+    if (users.contains(instruction)) {
+      if (auto apply = ApplySite::isa(instruction)) {
+        SmallVector<int, 2> rewritableArgumentIndices;
+        auto count = apply.getNumArguments();
+        for (unsigned index = 0; index < count; ++index) {
+          auto argument = apply.getArgument(index);
+          if (canReplaceValueWithBorrowedValue(argument)) {
+            rewritableArgumentIndices.push_back(index);
+          }
+        }
+        if (rewritableArgumentIndices.size() != usesInApply(apply)) {
+          return false;
+        }
+        // We can rewrite all the arguments which are transitive uses of the
+        // borrow.
+        for (auto index : rewritableArgumentIndices) {
+          auto argument = apply.getArgument(index);
+          auto borrowee = introducer->getOperand();
+          if (auto *cvi = dyn_cast<CopyValueInst>(argument)) {
+            cvi->setOperand(borrowee);
+          } else {
+            apply.setArgument(index, borrowee);
+          }
+        }
+        return true;
+      } else if (auto *bbi = dyn_cast<BeginBorrowInst>(instruction)) {
+        if (bbi->isLexical() &&
+            canReplaceValueWithBorrowedValue(bbi->getOperand())) {
+          auto borrowee = introducer->getOperand();
+          bbi->setOperand(borrowee);
+          return true;
+        }
+      }
+      return false;
+    }
+    return !isDeinitBarrier(instruction);
   }
 };
 
@@ -136,6 +205,9 @@ bool ShrinkBorrowScope::populateUsers() {
   for (auto *usePoint : usePoints) {
     auto *user = usePoint->getUser();
     users.insert(user);
+    if (auto apply = ApplySite::isa(user)) {
+      ++transitiveUsesPerApplySite[apply];
+    }
   }
   return true;
 }
@@ -195,7 +267,7 @@ void ShrinkBorrowScope::findBarriers() {
         barrier = instruction;
         break;
       }
-      if (isDeinitBarrier(instruction, borrowedValue.value)) {
+      if (!tryHoistOverInstruction(instruction)) {
         barrier = instruction;
         break;
       }

test/SILOptimizer/shrink_borrow_scope.sil

@@ -42,10 +42,8 @@ sil [ossa] @callee_optional_d_guaranteed: $@convention(thin) (@guaranteed Option
 // Hoist over br.
 // CHECK-LABEL: sil [ossa] @hoist_over_branch_1 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C):
-// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
-// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
-// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[INSTANCE]])
 // CHECK:         br [[EXIT:bb[0-9]+]]
 // CHECK:       [[EXIT]]:
 // CHECK:         return [[INSTANCE]]
@@ -64,10 +62,8 @@ bl1:
 // Hoist over cond_br.
 // CHECK-LABEL: sil [ossa] @hoist_over_branch_2 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C):
-// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
-// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
-// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[INSTANCE]])
 // CHECK:         cond_br undef, [[BL1:bb[0-9]+]], [[BL2:bb[0-9]+]]
 // CHECK:       [[BL1]]:
 // CHECK:         br [[EXIT:bb[0-9]+]]
@@ -95,10 +91,8 @@ exit:
 // Hoist over two brs.
 // CHECK-LABEL: sil [ossa] @hoist_over_branch_3 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C):
-// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
-// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
-// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[INSTANCE]])
 // CHECK:         cond_br undef, [[BL1:bb[0-9]+]], [[BL2:bb[0-9]+]]
 // CHECK:       [[BL1]]:
 // CHECK:         br [[EXIT:bb[0-9]+]]
@@ -125,15 +119,12 @@ exit:
 // Don't hoist over 1 / 2 brs.
 // CHECK-LABEL: sil [ossa] @hoist_over_branch_4 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C):
-// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         cond_br undef, [[BL1:bb[0-9]+]], [[BL2:bb[0-9]+]]
 // CHECK:       [[BL1]]:
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
-// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
-// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         {{%[^,]+}} = apply [[CALLEE_GUARANTEED]]([[INSTANCE]])
 // CHECK:         br [[EXIT:bb[0-9]+]]
 // CHECK:       [[BL2]]:
-// CHECK:         end_borrow [[LIFETIME]]
 // CHECK:         br [[EXIT]]
 // CHECK:       [[EXIT]]:
 // CHECK:         return [[INSTANCE]]
@@ -156,10 +147,8 @@ exit:
 // Hoist over switch_enum destinations.
 // CHECK-LABEL: sil [ossa] @hoist_over_branch_5 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C, [[CASE:%[^,]+]] : $OneOfThree):
-// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
 // CHECK:         [[CALLEE_GUARANTEED:%[^,]+]] = function_ref @callee_guaranteed
-// CHECK:         {{%[0-9]+}} = apply [[CALLEE_GUARANTEED]]([[LIFETIME]])
-// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         {{%[0-9]+}} = apply [[CALLEE_GUARANTEED]]([[INSTANCE]])
 // CHECK:         switch_enum [[CASE]] : $OneOfThree, case #OneOfThree.one!enumelt: [[ONE_DEST:bb[0-9]+]], case #OneOfThree.two!enumelt: [[TWO_DEST:bb[0-9]+]], case #OneOfThree.three!enumelt: [[THREE_DEST:bb[0-9]+]]
 // CHECK:       [[ONE_DEST]]:
 // CHECK:         br [[EXIT:bb[0-9]+]]