[semantic-arc] Optimize more lifetime joining when a copy_value, destroy_value are in the same block.

include/swift/SIL/OwnershipUtils.h

@@ -75,14 +75,34 @@ class ForwardingOperand {
 public:
   static Optional<ForwardingOperand> get(Operand *use);
 
+  Operand *getUse() const { return use; }
   ValueOwnershipKind getOwnershipKind() const;
   void setOwnershipKind(ValueOwnershipKind newKind) const;
   void replaceOwnershipKind(ValueOwnershipKind oldKind,
                             ValueOwnershipKind newKind) const;
 
-  OwnershipForwardingInst *getUser() const {
+  const OwnershipForwardingInst *operator->() const {
     return cast<OwnershipForwardingInst>(use->getUser());
   }
+  OwnershipForwardingInst *operator->() {
+    return cast<OwnershipForwardingInst>(use->getUser());
+  }
+  const OwnershipForwardingInst &operator*() const {
+    return *cast<OwnershipForwardingInst>(use->getUser());
+  }
+  OwnershipForwardingInst &operator*() {
+    return *cast<OwnershipForwardingInst>(use->getUser());
+  }
+
+  /// Call \p visitor with each value that contains the final forwarded
+  /// ownership of. E.x.: result of a unchecked_ref_cast, phi arguments of a
+  /// switch_enum.
+  bool visitForwardedValues(function_ref<bool(SILValue)> visitor);
+
+  /// If statically this forwarded operand has a single forwarded value that the
+  /// operand forwards ownership into, return that value. Return false
+  /// otherwise.
+  SILValue getSingleForwardedValue() const;
 };
 
 /// Returns true if the instruction is a 'reborrow'.
@@ -179,7 +199,10 @@ struct BorrowingOperand {
   /// Example: An apply performs an instantaneous recursive borrow of a
   /// guaranteed value but a begin_apply borrows the value over the entire
   /// region of code corresponding to the coroutine.
-  void visitLocalEndScopeInstructions(function_ref<void(Operand *)> func) const;
+  ///
+  /// NOTE: Return false from func to stop iterating. Returns false if the
+  /// closure requested to stop early.
+  bool visitLocalEndScopeUses(function_ref<bool(Operand *)> func) const;
 
   /// Returns true if this borrow scope operand consumes guaranteed
   /// values and produces a new scope afterwards.

lib/SIL/Utils/OwnershipUtils.cpp

@@ -161,31 +161,33 @@ llvm::raw_ostream &swift::operator<<(llvm::raw_ostream &os,
   return os;
 }
 
-void BorrowingOperand::visitLocalEndScopeInstructions(
-    function_ref<void(Operand *)> func) const {
+bool BorrowingOperand::visitLocalEndScopeUses(
+    function_ref<bool(Operand *)> func) const {
   switch (kind) {
   case BorrowingOperandKind::BeginBorrow:
     for (auto *use : cast<BeginBorrowInst>(op->getUser())->getUses()) {
       if (use->isLifetimeEnding()) {
-        func(use);
+        if (!func(use))
+          return false;
       }
     }
-    return;
+    return true;
   case BorrowingOperandKind::BeginApply: {
     auto *user = cast<BeginApplyInst>(op->getUser());
     for (auto *use : user->getTokenResult()->getUses()) {
-      func(use);
+      if (!func(use))
+        return false;
     }
-    return;
+    return true;
   }
   // These are instantaneous borrow scopes so there aren't any special end
   // scope instructions.
   case BorrowingOperandKind::Apply:
   case BorrowingOperandKind::TryApply:
   case BorrowingOperandKind::Yield:
-    return;
+    return true;
   case BorrowingOperandKind::Branch:
-    return;
+    return true;
   }
 }
 
@@ -267,7 +269,10 @@ void BorrowingOperand::visitUserResultConsumingUses(
 void BorrowingOperand::getImplicitUses(
     SmallVectorImpl<Operand *> &foundUses,
     std::function<void(Operand *)> *errorFunction) const {
-  visitLocalEndScopeInstructions([&](Operand *op) { foundUses.push_back(op); });
+  visitLocalEndScopeUses([&](Operand *op) {
+    foundUses.push_back(op);
+    return true;
+  });
 }
 
 //===----------------------------------------------------------------------===//
@@ -842,7 +847,7 @@ Optional<ForwardingOperand> ForwardingOperand::get(Operand *use) {
 }
 
 ValueOwnershipKind ForwardingOperand::getOwnershipKind() const {
-  return getUser()->getOwnershipKind();
+  return (*this)->getOwnershipKind();
 }
 
 void ForwardingOperand::setOwnershipKind(ValueOwnershipKind newKind) const {
@@ -968,3 +973,45 @@ void ForwardingOperand::replaceOwnershipKind(ValueOwnershipKind oldKind,
   }
   llvm_unreachable("Out of sync with ForwardingOperand::get?!");
 }
+
+SILValue ForwardingOperand::getSingleForwardedValue() const {
+  assert(isGuaranteedForwardingUse(use));
+  if (auto *svi = dyn_cast<SingleValueInstruction>(use->getUser()))
+    return svi;
+  return SILValue();
+}
+
+bool ForwardingOperand::visitForwardedValues(
+    function_ref<bool(SILValue)> visitor) {
+  auto *user = use->getUser();
+
+  assert(isGuaranteedForwardingUse(use));
+
+  // See if we have a single value instruction... if we do that is always the
+  // transitive result.
+  if (auto *svi = dyn_cast<SingleValueInstruction>(user)) {
+    return visitor(svi);
+  }
+
+  if (auto *mvri = dyn_cast<MultipleValueInstruction>(user)) {
+    return llvm::all_of(mvri->getResults(), [&](SILValue value) {
+      if (value.getOwnershipKind() == OwnershipKind::None)
+        return true;
+      return visitor(value);
+    });
+  }
+
+  // This is an instruction like switch_enum and checked_cast_br that are
+  // "transforming terminators"... We know that this means that we should at
+  // most have a single phi argument.
+  auto *ti = cast<TermInst>(user);
+  return llvm::all_of(ti->getSuccessorBlocks(), [&](SILBasicBlock *succBlock) {
+    // If we do not have any arguments, then continue.
+    if (succBlock->args_empty())
+      return true;
+
+    auto args = succBlock->getSILPhiArguments();
+    assert(args.size() == 1 && "Transforming terminator with multiple args?!");
+    return visitor(args[0]);
+  });
+}

lib/SIL/Verifier/ReborrowVerifier.cpp

@@ -42,8 +42,9 @@ void ReborrowVerifier::verifyReborrows(BorrowingOperand initialScopedOperand,
                                        SILValue value) {
   SmallVector<std::tuple<Operand *, SILValue>, 4> worklist;
   // Initialize the worklist with borrow lifetime ending uses
-  initialScopedOperand.visitLocalEndScopeInstructions([&](Operand *op) {
+  initialScopedOperand.visitLocalEndScopeUses([&](Operand *op) {
     worklist.emplace_back(op, value);
+    return true;
   });
 
   while (!worklist.empty()) {

lib/SILOptimizer/SemanticARC/BorrowScopeOpts.cpp

@@ -18,12 +18,17 @@
 ///
 //===----------------------------------------------------------------------===//
 
+#include "Context.h"
 #include "SemanticARCOptVisitor.h"
 
 using namespace swift;
 using namespace swift::semanticarc;
 
 bool SemanticARCOptVisitor::visitBeginBorrowInst(BeginBorrowInst *bbi) {
+  // Quickly check if we are supposed to perform this transformation.
+  if (!ctx.shouldPerform(ARCTransformKind::RedundantBorrowScopeElimPeephole))
+    return false;
+
   auto kind = bbi->getOperand().getOwnershipKind();
   SmallVector<EndBorrowInst *, 16> endBorrows;
   for (auto *op : bbi->getUses()) {

lib/SILOptimizer/SemanticARC/Context.h

@@ -14,6 +14,7 @@
 #define SWIFT_SILOPTIMIZER_SEMANTICARC_CONTEXT_H
 
 #include "OwnershipLiveRange.h"
+#include "SemanticARCOpts.h"
 
 #include "swift/Basic/BlotSetVector.h"
 #include "swift/Basic/FrozenMultiMap.h"
@@ -30,6 +31,7 @@ namespace semanticarc {
 
 struct LLVM_LIBRARY_VISIBILITY Context {
   SILFunction &fn;
+  ARCTransformKind transformKind = ARCTransformKind::All;
   Optional<DeadEndBlocks> deadEndBlocks;
   ValueLifetimeAnalysis::Frontier lifetimeFrontier;
   SmallMultiMapCache<SILValue, Operand *> addressToExhaustiveWriteListCache;
@@ -91,6 +93,25 @@ struct LLVM_LIBRARY_VISIBILITY Context {
 
   void verify() const;
 
+  bool shouldPerform(ARCTransformKind testKind) const {
+    // When asserts are enabled, we allow for specific arc transforms to be
+    // turned on/off via LLVM args. So check that if we have asserts, perform
+    // all optimizations otherwise.
+#ifndef NDEBUG
+    if (transformKind == ARCTransformKind::Invalid)
+      return false;
+    return bool(testKind & transformKind);
+#else
+    return true;
+#endif
+  }
+
+  void reset() {
+    lifetimeFrontier.clear();
+    addressToExhaustiveWriteListCache.clear();
+    joinedOwnedIntroducerToConsumedOperands.reset();
+  }
+
 private:
   static bool
   constructCacheValue(SILValue initialValue,