Merge pull request swiftlang#35718 from meg-gupta/dceossafixespr

Fixes for OSSA DCE

Author: meg-gupta

include/swift/SIL/OwnershipUtils.h

@@ -1030,6 +1030,23 @@ bool getAllOwnedValueIntroducers(SILValue value,
 
 OwnedValueIntroducer getSingleOwnedValueIntroducer(SILValue value);
 
+using BaseValueSet = SmallPtrSet<SILValue, 8>;
+
+/// Starting from \p initialScopeOperand, find all reborrows and their
+/// corresponding base values, and run the visitor function \p
+/// visitReborrowBaseValuePair on them.
+///  Note that a reborrow phi, can have different base values based on different
+/// control flow paths.
+void findTransitiveReborrowBaseValuePairs(
+    BorrowingOperand initialScopeOperand, SILValue origBaseValue,
+    function_ref<void(SILPhiArgument *, SILValue)> visitReborrowBaseValuePair);
+
+/// Given a begin_borrow visit all end_borrow users of the borrow or its
+/// reborrows.
+void visitTransitiveEndBorrows(
+    BeginBorrowInst *borrowInst,
+    function_ref<void(EndBorrowInst *)> visitEndBorrow);
+
 } // namespace swift
 
 #endif

lib/SIL/Utils/OwnershipUtils.cpp

@@ -1278,3 +1278,77 @@ bool ForwardingOperand::visitForwardedValues(
     return visitor(args[0]);
   });
 }
+
+void swift::findTransitiveReborrowBaseValuePairs(
+    BorrowingOperand initialScopedOperand, SILValue origBaseValue,
+    function_ref<void(SILPhiArgument *, SILValue)> visitReborrowBaseValuePair) {
+  // We need a SetVector to make sure we don't revisit the same reborrow operand
+  // again.
+  SmallSetVector<std::tuple<Operand *, SILValue>, 4> worklist;
+
+  // Populate the worklist with reborrow and the base value
+  initialScopedOperand.visitScopeEndingUses([&](Operand *op) {
+    if (op->getOperandOwnership() == OperandOwnership::Reborrow) {
+      worklist.insert(std::make_tuple(op, origBaseValue));
+    }
+    return true;
+  });
+
+  // Size of worklist changes in this loop
+  for (unsigned idx = 0; idx < worklist.size(); idx++) {
+    Operand *reborrowOp;
+    SILValue baseValue;
+    std::tie(reborrowOp, baseValue) = worklist[idx];
+
+    BorrowingOperand borrowingOperand(reborrowOp);
+    assert(borrowingOperand.isReborrow());
+
+    auto *branchInst = cast<BranchInst>(reborrowOp->getUser());
+    auto *succBlock = branchInst->getDestBB();
+    auto *phiArg = cast<SILPhiArgument>(
+        succBlock->getArgument(reborrowOp->getOperandNumber()));
+
+    SILValue newBaseVal = baseValue;
+    // If the previous base value was also passed as a phi arg, that will be
+    // the new base value.
+    for (auto *arg : succBlock->getArguments()) {
+      if (arg->getIncomingPhiValue(branchInst->getParent()) == baseValue) {
+        newBaseVal = arg;
+        break;
+      }
+    }
+
+    // Call the visitor function
+    visitReborrowBaseValuePair(phiArg, newBaseVal);
+
+    BorrowedValue scopedValue(phiArg);
+    scopedValue.visitLocalScopeEndingUses([&](Operand *op) {
+      if (op->getOperandOwnership() == OperandOwnership::Reborrow) {
+        worklist.insert(std::make_tuple(op, newBaseVal));
+      }
+      return true;
+    });
+  }
+}
+
+void swift::visitTransitiveEndBorrows(
+    BeginBorrowInst *borrowInst,
+    function_ref<void(EndBorrowInst *)> visitEndBorrow) {
+  SmallSetVector<SILValue, 4> worklist;
+  worklist.insert(borrowInst);
+
+  while (!worklist.empty()) {
+    auto val = worklist.pop_back_val();
+    for (auto *consumingUse : val->getConsumingUses()) {
+      auto *consumingUser = consumingUse->getUser();
+      if (auto *branch = dyn_cast<BranchInst>(consumingUser)) {
+        auto *succBlock = branch->getSingleSuccessorBlock();
+        auto *phiArg = cast<SILPhiArgument>(
+            succBlock->getArgument(consumingUse->getOperandNumber()));
+        worklist.insert(phiArg);
+      } else {
+        visitEndBorrow(cast<EndBorrowInst>(consumingUser));
+      }
+    }
+  }
+}

lib/SIL/Verifier/ReborrowVerifier.cpp

@@ -39,57 +39,18 @@ bool ReborrowVerifier::verifyReborrowLifetime(SILPhiArgument *phiArg,
 
 void ReborrowVerifier::verifyReborrows(BorrowingOperand initialScopedOperand,
                                        SILValue value) {
-  SmallVector<std::tuple<Operand *, SILValue>, 4> worklist;
-  // Initialize the worklist with borrow lifetime ending uses
-  initialScopedOperand.visitScopeEndingUses([&](Operand *op) {
-    worklist.emplace_back(op, value);
-    return true;
-  });
-
-  while (!worklist.empty()) {
-    Operand *borrowLifetimeEndOp;
-    SILValue baseVal;
-    std::tie(borrowLifetimeEndOp, baseVal) = worklist.pop_back_val();
-    auto *borrowLifetimeEndUser = borrowLifetimeEndOp->getUser();
-
-    auto borrowingOperand = BorrowingOperand::get(borrowLifetimeEndOp);
-    if (!borrowingOperand || !borrowingOperand.isReborrow())
-      continue;
-
-    if (isVisitedOp(borrowLifetimeEndOp, baseVal))
-      continue;
-
-    // Process reborrow
-    auto *branchInst = cast<BranchInst>(borrowLifetimeEndUser);
-    for (auto *succBlock : branchInst->getSuccessorBlocks()) {
-      auto *phiArg = cast<SILPhiArgument>(
-          succBlock->getArgument(borrowLifetimeEndOp->getOperandNumber()));
-      assert(phiArg->getOwnershipKind() == OwnershipKind::Guaranteed);
-
-      SILValue newBaseVal = baseVal;
-      // If the previous base value was also passed as a phi arg, that will be
-      // the new base value.
-      for (auto *arg : succBlock->getArguments()) {
-        if (arg->getIncomingPhiValue(branchInst->getParent()) == baseVal) {
-          newBaseVal = arg;
-          break;
-        }
-      }
-
-      if (isVisitedPhiArg(phiArg, newBaseVal))
-        continue;
-      addVisitedPhiArg(phiArg, newBaseVal);
-      verifyReborrowLifetime(phiArg, newBaseVal);
-
-      // Find the scope ending uses of the guaranteed phi arg and add it to the
-      // worklist.
-      auto scopedValue = BorrowedValue(phiArg);
-      assert(scopedValue);
-      scopedValue.visitLocalScopeEndingUses([&](Operand *op) {
-        addVisitedOp(op, newBaseVal);
-        worklist.emplace_back(op, newBaseVal);
-        return true;
-      });
+  auto visitReborrowBaseValuePair = [&](SILPhiArgument *phiArg,
+                                        SILValue baseValue) {
+    // If the (phiArg, baseValue) pair was not visited before, verify the
+    // lifetime.
+    auto it = reborrowToBaseValuesMap.find(phiArg);
+    if (it == reborrowToBaseValuesMap.end() ||
+        it->second.find(baseValue) == it->second.end()) {
+      reborrowToBaseValuesMap[phiArg].insert(baseValue);
+      verifyReborrowLifetime(phiArg, baseValue);
     }
-  }
+  };
+  // For every unique reborrow/base value pair, verify the lifetime
+  findTransitiveReborrowBaseValuePairs(initialScopedOperand, value,
+                                       visitReborrowBaseValuePair);
 }

lib/SIL/Verifier/ReborrowVerifierPrivate.h

@@ -31,45 +31,27 @@ class ReborrowVerifier {
   /// A cache of dead-end basic blocks that we use to determine if we can
   /// ignore "leaks".
   DeadEndBlocks &deadEndBlocks;
-  /// A cache map of borrow lifetime ending operands and their base value
-  llvm::SmallDenseMap<Operand *, SILValue> visitedOps;
-  /// A cache map of guaranteed phi args and their base value
-  llvm::SmallDenseMap<SILPhiArgument *, SILValue> visitedPhiArgs;
   /// The builder that the checker uses to emit error messages, crash if asked
   /// for, or supply back interesting info to the caller.
   LinearLifetimeChecker::ErrorBuilder errorBuilder;
+  /// A map of reborrow phi arg to its base values.
+  /// Note that a reborrow phi arg can have different base values based on
+  /// different control flow paths.
+  llvm::DenseMap<SILPhiArgument *, SmallPtrSet<SILValue, 8>>
+      reborrowToBaseValuesMap;
 
 public:
   ReborrowVerifier(const SILFunction *func, DeadEndBlocks &deadEndBlocks,
                    LinearLifetimeChecker::ErrorBuilder errorBuilder)
       : deadEndBlocks(deadEndBlocks), errorBuilder(errorBuilder) {}
 
+  /// Find all the reborrows of \p initialScopedOperand and verify their
+  /// lifetime.
   void verifyReborrows(BorrowingOperand initialScopedOperand, SILValue value);
 
 private:
   /// Verifies whether the reborrow's lifetime lies within its base value
   bool verifyReborrowLifetime(SILPhiArgument *phiArg, SILValue baseVal);
-
-  /// Check if the operand is visited
-  bool isVisitedOp(Operand *op, SILValue baseVal) {
-    return visitedOps.find(op) != visitedOps.end();
-  }
-
-  /// Check if the phi arg and base value are visited
-  bool isVisitedPhiArg(SILPhiArgument *phiArg, SILValue baseVal) {
-    auto itPhiArg = visitedPhiArgs.find(phiArg);
-    return itPhiArg != visitedPhiArgs.end() && (*itPhiArg).second == baseVal;
-  }
-
-  /// Mark operand as visited
-  void addVisitedOp(Operand *op, SILValue baseVal) {
-    visitedOps.insert({op, baseVal});
-  }
-
-  /// Mark guaranteed phi arg as visited
-  void addVisitedPhiArg(SILPhiArgument *phiArg, SILValue baseVal) {
-    visitedPhiArgs.insert({phiArg, baseVal});
-  }
 };
 
 } // namespace swift