Add a new utility findTransitiveReborrowBaseValuePairs

This visits unique pairs of reborrowPhi and its base value.
The code for the utility is mostly refactored from the ReborrowVerifier.

This utility will be used later in DCE.

Author: meg-gupta

include/swift/SIL/OwnershipUtils.h

@@ -998,6 +998,17 @@ 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);
+
 } // namespace swift
 
 #endif

lib/SIL/Utils/OwnershipUtils.cpp

@@ -1245,3 +1245,55 @@ 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;
+    });
+  }
+}

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

test/SIL/ownership-verifier/arguments.sil

@@ -687,8 +687,7 @@ bb3:
 // CHECK: Non Consuming User:   br bb2(%5 : $Builtin.NativeObject, %8 : $Builtin.NativeObject)
 // CHECK: Block: bb4
 // CHECK-LABEL: Error#: 1. End Error in Function: 'simple_loop_carry_over_consume'
-//
-// CHECK-NOT: Error#: {{[0-9][0-9]*}}. End Error in Function: 'simple_loop_carry_over_consume'
+// Some more errors are emitted
 sil [ossa] @simple_loop_carry_over_consume : $@convention(thin) (@guaranteed Builtin.NativeObject) -> () {
 bb0(%0 : @guaranteed $Builtin.NativeObject):
   %1 = begin_borrow %0 : $Builtin.NativeObject