SIL: handle the borrowed-from instruction in the C++ ownership utilities and passes

Author: eeckstein

include/swift/SIL/OwnershipUtils.h

@@ -90,6 +90,10 @@ inline bool isForwardingConsume(SILValue value) {
 //                        Ownership Def-Use Utilities
 //===----------------------------------------------------------------------===//
 
+SILValue lookThroughSingleBorrowedFromUser(SILValue v);
+SILValue lookThroughBorrowedFromDef(SILValue v);
+void getTransitiveConsumingUses(SILValue v, llvm::SmallVectorImpl<Operand *> &results);
+
 /// Whether the specified OSSA-lifetime introducer has a pointer escape.
 ///
 /// precondition: \p value introduces an OSSA-lifetime, either a BorrowedValue

lib/SIL/IR/SILArgument.cpp

@@ -16,6 +16,7 @@
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace swift;
@@ -291,7 +292,8 @@ bool SILPhiArgument::visitTransitiveIncomingPhiOperands(
     argument->getIncomingPhiOperands(operands);
 
     for (auto *operand : operands) {
-      SILPhiArgument *forwarded = dyn_cast<SILPhiArgument>(operand->get());
+      SILValue opVal = lookThroughBorrowedFromDef(operand->get());
+      SILPhiArgument *forwarded = dyn_cast<SILPhiArgument>(opVal);
       if (forwarded && forwarded->isPhi()) {
         worklist.insert(forwarded);
       }

lib/SIL/Utils/OwnershipUtils.cpp

@@ -168,6 +168,7 @@ bool swift::computeIsGuaranteedForwarding(SILValue value) {
   if (value->getOwnershipKind() != OwnershipKind::Guaranteed) {
     return false;
   }
+  value = lookThroughBorrowedFromDef(value);
   // NOTE: canOpcodeForwardInnerGuaranteedValues returns true for transformation
   // terminator results.
   if (canOpcodeForwardInnerGuaranteedValues(value) ||
@@ -185,9 +186,9 @@ bool swift::computeIsGuaranteedForwarding(SILValue value) {
   // OwnershipVerifier.
   bool isGuaranteedForwardingPhi = false;
   phi->visitTransitiveIncomingPhiOperands([&](auto *, auto *op) -> bool {
-    auto opValue = op->get();
-    assert(opValue->getOwnershipKind().isCompatibleWith(
+    assert(op->get()->getOwnershipKind().isCompatibleWith(
         OwnershipKind::Guaranteed));
+    auto opValue = lookThroughBorrowedFromDef(op->get());
     if (canOpcodeForwardInnerGuaranteedValues(opValue) ||
         isa<SILFunctionArgument>(opValue)) {
       isGuaranteedForwardingPhi = true;
@@ -202,6 +203,41 @@ bool swift::computeIsGuaranteedForwarding(SILValue value) {
   return isGuaranteedForwardingPhi;
 }
 
+SILValue swift::lookThroughSingleBorrowedFromUser(SILValue v) {
+  BorrowedFromInst *singleBfi = nullptr;
+  for (auto *use : v->getUses()) {
+    if (auto *bfi = dyn_cast<BorrowedFromInst>(use->getUser())) {
+      if (use->getOperandNumber() == 0) {
+        if (singleBfi)
+          return v;
+        singleBfi = bfi;
+      }
+    }
+  }
+  if (singleBfi)
+    return singleBfi;
+  return v;
+}
+
+SILValue swift::lookThroughBorrowedFromDef(SILValue v) {
+  while (auto *bfi = dyn_cast<BorrowedFromInst>(v)) {
+    v = bfi->getBorrowedValue();
+  }
+  return v;
+}
+
+void swift::getTransitiveConsumingUses(SILValue v, llvm::SmallVectorImpl<Operand *> &results) {
+  for (Operand *use : v->getUses()) {
+    if (use->isLifetimeEnding()) {
+      results.push_back(use);
+    } else if (auto *bfi = dyn_cast<BorrowedFromInst>(use->getUser())) {
+      if (use->getOperandNumber() == 0) {
+        getTransitiveConsumingUses(bfi, results);
+      }
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //                 Guaranteed Use-Point (Lifetime) Discovery
 //===----------------------------------------------------------------------===//
@@ -896,26 +932,28 @@ void BorrowedValue::getLocalScopeEndingInstructions(
 
 // Note: BorrowedLifetimeExtender assumes no intermediate values between a
 // borrow introducer and its reborrow. The borrowed value must be an operand of
-// the reborrow.
+// the reborrow. Exception: it looks through `borrowed-from` of a reborrow phi.
 bool BorrowedValue::visitLocalScopeEndingUses(
     function_ref<bool(Operand *)> visitor) const {
   assert(isLocalScope() && "Should only call this given a local scope");
+  SILValue v = value;
   switch (kind) {
   case BorrowedValueKind::Invalid:
     llvm_unreachable("Using invalid case?!");
   case BorrowedValueKind::SILFunctionArgument:
     llvm_unreachable("Should only call this with a local scope");
-  case BorrowedValueKind::LoadBorrow:
-  case BorrowedValueKind::BeginBorrow:
   case BorrowedValueKind::Phi:
-    for (auto *use : value->getUses()) {
-      if (use->isLifetimeEnding()) {
-        if (!visitor(use))
-          return false;
-      }
+  case BorrowedValueKind::LoadBorrow:
+  case BorrowedValueKind::BeginBorrow: {
+    llvm::SmallVector<Operand *, 16> consumingUses;
+    getTransitiveConsumingUses(v, consumingUses);
+    for (Operand *use : consumingUses) {
+      if (!visitor(use))
+        return false;
     }
     return true;
   }
+  }
   llvm_unreachable("Covered switch isn't covered?!");
 }
 
@@ -1677,7 +1715,7 @@ void swift::visitExtendedGuaranteedForwardingPhiBaseValuePairs(
     // then set newBaseValue to the reborrow.
     for (auto &op : phiOp.getBranch()->getAllOperands()) {
       PhiOperand otherPhiOp(&op);
-      if (otherPhiOp.getSource() != currentBaseValue) {
+      if (lookThroughBorrowedFromDef(otherPhiOp.getSource()) != currentBaseValue) {
         continue;
       }
       newBaseValue = otherPhiOp.getValue();
@@ -1974,6 +2012,8 @@ class FindEnclosingDefs {
 
     assert(incomingValue->getOwnershipKind() == OwnershipKind::Guaranteed);
 
+    incomingValue = lookThroughBorrowedFromDef(incomingValue);
+
     // Avoid repeatedly constructing BorrowedValue during use-def
     // traversal. That would be quadratic if it checks all uses for reborrows.
     if (auto *predPhi = dyn_cast<SILPhiArgument>(incomingValue)) {
@@ -2257,7 +2297,9 @@ void swift::visitTransitiveEndBorrows(
 
   while (!worklist.empty()) {
     auto val = worklist.pop();
-    for (auto *consumingUse : val->getConsumingUses()) {
+    llvm::SmallVector<Operand *, 16> consumingUses;
+    getTransitiveConsumingUses(val, consumingUses);
+    for (auto *consumingUse : consumingUses) {
       auto *consumingUser = consumingUse->getUser();
       if (auto *branch = dyn_cast<BranchInst>(consumingUser)) {
         auto *succBlock = branch->getSingleSuccessorBlock();

lib/SIL/Verifier/GuaranteedPhiVerifier.cpp

@@ -28,12 +28,14 @@ bool GuaranteedPhiVerifier::verifyDependentPhiLifetime(SILPhiArgument *phi,
   dependentPhiToBaseValueMap[phi].insert(baseValue);
 
   // Now, verify whether phi's lifetime is within baseValue's lifetime
-  SmallVector<Operand *, 4> baseValConsumingUses(baseValue->getConsumingUses());
+  SmallVector<Operand *, 4> baseValConsumingUses;
+  getTransitiveConsumingUses(baseValue, baseValConsumingUses);
   // If the baseValue has no consuming uses, there is nothing more to verify
   if (baseValConsumingUses.empty())
     return false;
 
-  SmallVector<Operand *, 4> phiUses(phi->getUses());
+  SmallVector<Operand *, 4> phiUses;
+  getTransitiveConsumingUses(phi, phiUses);
   LinearLifetimeChecker checker(deadEndBlocks);
   // newErrorBuilder is consumed at the end of the checkValue function.
   // Copy initial state from errorBuilder everytime

lib/SIL/Verifier/SILOwnershipVerifier.cpp

@@ -639,7 +639,7 @@ bool SILValueOwnershipChecker::
   bool foundGuaranteedForwardingPhiOperand = false;
   bool foundNonGuaranteedForwardingPhiOperand = false;
   phi->visitTransitiveIncomingPhiOperands([&](auto *, auto *operand) -> bool {
-    auto value = operand->get();
+    auto value = lookThroughBorrowedFromDef(operand->get());
     if (canOpcodeForwardInnerGuaranteedValues(value) ||
         isa<SILFunctionArgument>(value)) {
       foundGuaranteedForwardingPhiOperand = true;

lib/SILOptimizer/SemanticARC/OwnershipLiveRange.cpp

@@ -105,8 +105,10 @@ OwnershipLiveRange::OwnershipLiveRange(SILValue value)
     // the users force the live range to be alive.
     if (!ti) {
       for (SILValue v : user->getResults()) {
-        if (v->getOwnershipKind() != OwnershipKind::Owned)
+        if (v->getOwnershipKind() != OwnershipKind::Owned &&
+            !isa<BorrowedFromInst>(user)) {
           continue;
+        }
         llvm::copy(v->getUses(), std::back_inserter(worklist));
       }
       continue;

lib/SILOptimizer/Transforms/ConditionForwarding.cpp

@@ -18,6 +18,7 @@
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILUndef.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 
@@ -158,13 +159,15 @@ static bool hasNoRelevantSideEffects(SILBasicBlock *BB) {
 bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
   // The switch_enum argument (an Enum) must be a block argument at the merging
   // point of the condition's destinations.
-  auto *Arg = dyn_cast<SILArgument>(SEI->getOperand());
+  auto *Arg = dyn_cast<SILArgument>(lookThroughBorrowedFromDef(SEI->getOperand()));
   if (!Arg)
     return false;
 
+  SILValue argValue = lookThroughSingleBorrowedFromUser(Arg);
+
   // The switch_enum must be the only use of the Enum, except it may be used in
   // SEI's successors.
-  for (Operand *ArgUse : Arg->getUses()) {
+  for (Operand *ArgUse : argValue->getUses()) {
     SILInstruction *ArgUser = ArgUse->getUser();
     if (ArgUser == SEI)
       continue;
@@ -242,8 +245,8 @@ bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
   // First thing to do is to replace all uses of the Enum (= the merging block
   // argument), as this argument gets deleted.
   BasicBlockSet NeedEnumArg(BB->getParent());
-  while (!Arg->use_empty()) {
-    Operand *ArgUse = *Arg->use_begin();
+  while (!argValue->use_empty()) {
+    Operand *ArgUse = *argValue->use_begin();
     SILInstruction *ArgUser = ArgUse->getUser();
     if (ArgUser->isDebugInstruction()) {
       // Don't care about debug instructions. Just remove them.
@@ -318,6 +321,9 @@ bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
       }
     }
   }
+  if (argValue != Arg) {
+    cast<BorrowedFromInst>(argValue)->eraseFromParent();
+  }
 
   // Final step: replace the switch_enum by the condition.
   SILBuilder B(Condition);

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -294,6 +294,10 @@ void DCE::markLive() {
         addReverseDependency(I.getOperand(0), &I);
         break;
       }
+      case SILInstructionKind::BorrowedFromInst: {
+        addReverseDependency(I.getOperand(0), &I);
+        break;
+      }
       case SILInstructionKind::EndLifetimeInst: {
         // The instruction is live only if it's operand value is also live
         addReverseDependency(I.getOperand(0), &I);
@@ -576,7 +580,7 @@ void DCE::endLifetimeOfLiveValue(SILValue value, SILInstruction *insertPt) {
     builder.emitDestroyOperation(RegularLocation::getAutoGeneratedLocation(),
                                  value);
   }
-  BorrowedValue borrow(value);
+  BorrowedValue borrow(lookThroughBorrowedFromDef(value));
   if (borrow && borrow.isLocalScope()) {
     builder.emitEndBorrowOperation(RegularLocation::getAutoGeneratedLocation(),
                                    value);

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -155,6 +155,10 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         defUseWorklist.insert(copy);
         continue;
       }
+      if (auto *bfi = dyn_cast<BorrowedFromInst>(user)) {
+        defUseWorklist.insert(bfi);
+        continue;
+      }
       // Handle debug_value instructions separately.
       if (pruneDebugMode) {
         if (auto *dvi = dyn_cast<DebugValueInst>(user)) {