Merge pull request #62634 from atrick/enclosing-def

Add findEnclosingDefs and findBorrowIntroducers utilities

Author: swift-ci

include/swift/SIL/MemAccessUtils.h

@@ -206,7 +206,12 @@ SILValue findOwnershipReferenceRoot(SILValue ref);
 
 /// Look through all ownership forwarding instructions to find the values which
 /// were originally borrowed.
-void findGuaranteedReferenceRoots(SILValue value, bool lookThroughNestedBorrows,
+///
+/// Note: This treats guaranteed forwarding phis like roots even though they do
+/// not introduce the borrow scope. This ensures that all roots dominate \p
+/// reference Value. But the client will need to handle forwarding phis.
+void findGuaranteedReferenceRoots(SILValue referenceValue,
+                                  bool lookThroughNestedBorrows,
                                   SmallVectorImpl<SILValue> &roots);
 
 /// Find the aggregate containing the first owned root of the

include/swift/SIL/OwnershipUtils.h

@@ -50,6 +50,7 @@
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILValue.h"
+#include "swift/SIL/StackList.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -1299,6 +1300,18 @@ void visitExtendedGuaranteedForwardingPhiBaseValuePairs(
     BorrowedValue borrow, function_ref<void(SILPhiArgument *, SILValue)>
                               visitGuaranteedForwardingPhiBaseValuePair);
 
+/// If \p value is a guaranteed non-phi value forwarded from it's instruction's
+/// operands, visit each forwarded operand.
+///
+/// Returns true if \p visitOperand was called (for convenience).
+///
+/// Precondition: \p value is not a phi. The client must handle phis first by
+/// checking if they are reborrows (using BorrowedValue). Reborrows have no
+/// forwarded operands. Guaranteed forwarding need to be handled by recursing
+/// through the phi operands.
+bool visitForwardedGuaranteedOperands(
+  SILValue value, function_ref<void(Operand *)> visitOperand);
+
 /// Visit the phis in the same block as \p phi which are reborrows of a borrow
 /// of one of the values reaching \p phi.
 ///
@@ -1307,6 +1320,24 @@ void visitExtendedGuaranteedForwardingPhiBaseValuePairs(
 bool visitAdjacentReborrowsOfPhi(SILPhiArgument *phi,
                                  function_ref<bool(SILPhiArgument *)> visitor);
 
+/// Visit each definition of a scope that immediately encloses a guaranteed
+/// value. The guaranteed value effectively keeps these scopes alive.
+///
+/// This means something different depepending on whether \p value is itself a
+/// borrow introducer vs. a forwarded guaranteed value. If \p value is an
+/// introducer, then this discovers the enclosing borrow scope and visits all
+/// introducers of that scope. If \p value is a forwarded value, then this
+/// visits the introducers of the current borrow scope.
+bool visitEnclosingDefs(SILValue value, function_ref<bool(SILValue)> visitor);
+
+/// Visit the values that introduce the borrow scopes that includes \p
+/// value. If value is owned, or introduces a borrow scope, then this only
+/// visits \p value.
+///
+/// Returns false if the visitor returned false and exited early.
+bool visitBorrowIntroducers(SILValue value,
+                            function_ref<bool(SILValue)> visitor);
+
 /// Given a begin of a borrow scope, visit all end_borrow users of the borrow or
 /// its reborrows.
 void visitTransitiveEndBorrows(

include/swift/SIL/SILArgument.h

@@ -205,7 +205,7 @@ class SILArgument : public ValueBase {
   /// Precondition: this->isTerminatorResult()
   ///
   /// TODO: Move this and other APIs into a TerminatorResult abstraction.
-  const Operand *forwardedTerminatorResultOperand() const;
+  Operand *forwardedTerminatorResultOperand() const;
 
   /// Return the SILArgumentKind of this argument.
   SILArgumentKind getKind() const {

include/swift/SIL/SILInstruction.h

@@ -599,6 +599,10 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
     return getTypeDependentOperands().size();
   }
 
+  unsigned getNumRealOperands() const {
+    return getAllOperands().size() - getNumTypeDependentOperands();
+  }
+
   bool isTypeDependentOperand(unsigned i) const {
     return i >= getNumOperands() - getNumTypeDependentOperands();
   }
@@ -8496,9 +8500,14 @@ class TermInst : public NonValueInstruction {
   /// example, a switch forwards ownership of the enum type into ownership of
   /// the payload.
   ///
-  /// Postcondition: each successor has zero or one block arguments which
-  /// represents the forwaded result.
+  /// Postcondition: if the result is non-null, then each successor has zero or
+  /// one block arguments which represents the forwaded result.
   const Operand *forwardedOperand() const;
+
+  Operand *forwardedOperand() {
+    return const_cast<Operand *>(
+      static_cast<const TermInst *>(this)->forwardedOperand());
+  }
 };
 
 // Forwards the first operand to a result in each successor block.

lib/SIL/IR/SILArgument.cpp

@@ -341,7 +341,7 @@ TermInst *SILPhiArgument::getTerminatorForResult() const {
   return nullptr;
 }
 
-const Operand *SILArgument::forwardedTerminatorResultOperand() const {
+Operand *SILArgument::forwardedTerminatorResultOperand() const {
   assert(isTerminatorResult() && "API is invalid for phis");
   return getSingleTerminator()->forwardedOperand();
 }

lib/SIL/Utils/MemAccessUtils.cpp

@@ -16,6 +16,8 @@
 #include "swift/Basic/GraphNodeWorklist.h"
 #include "swift/SIL/Consumption.h"
 #include "swift/SIL/DynamicCasts.h"
+#include "swift/SIL/NodeDatastructures.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBridging.h"
 #include "swift/SIL/SILBridgingUtils.h"
 #include "swift/SIL/SILInstruction.h"
@@ -860,78 +862,42 @@ SILValue swift::findOwnershipReferenceRoot(SILValue ref) {
   return ref;
 }
 
-void swift::findGuaranteedReferenceRoots(SILValue value,
+void swift::findGuaranteedReferenceRoots(SILValue referenceValue,
                                          bool lookThroughNestedBorrows,
                                          SmallVectorImpl<SILValue> &roots) {
-  GraphNodeWorklist<SILValue, 4> worklist;
-  auto addAllOperandsToWorklist = [&worklist](SILInstruction *inst) -> bool {
-    if (inst->getNumOperands() > 0) {
-      for (auto operand : inst->getOperandValues()) {
-        worklist.insert(operand);
-      }
-      return true;
-    }
-    return false;
-  };
-  worklist.initialize(value);
+  ValueWorklist worklist(referenceValue->getFunction());
+  worklist.pushIfNotVisited(referenceValue);
   while (auto value = worklist.pop()) {
-    if (auto *result = SILArgument::isTerminatorResult(value)) {
-      if (auto *forwardedOper = result->forwardedTerminatorResultOperand()) {
-        worklist.insert(forwardedOper->get());
-        continue;
-      }
-    } else if (auto *inst = value->getDefiningInstruction()) {
-      if (auto *bbi = dyn_cast<BeginBorrowInst>(inst)) {
+    // Instructions may forwarded None ownership to guaranteed.
+    if (value->getOwnershipKind() != OwnershipKind::Guaranteed)
+      continue;
+
+    if (SILArgument::asPhi(value)) {
+      roots.push_back(value);
+      continue;
+    }
+
+    if (visitForwardedGuaranteedOperands(value, [&](Operand *operand) {
+        worklist.pushIfNotVisited(operand->get());
+      })) {
+      // This instruction is not a root if any operands were forwarded,
+      // regardless of whether they were already visited.
+      continue;
+    }
+    // Found a potential root.
+    if (lookThroughNestedBorrows) {
+      if (auto *bbi = dyn_cast<BeginBorrowInst>(value)) {
         auto borrowee = bbi->getOperand();
-        if (lookThroughNestedBorrows &&
-            borrowee->getOwnershipKind() == OwnershipKind::Guaranteed) {
+        if (borrowee->getOwnershipKind() == OwnershipKind::Guaranteed) {
           // A nested borrow, the root guaranteed earlier in the use-def chain.
-          worklist.insert(borrowee);
-        }
-        // The borrowee isn't guaranteed or we aren't looking through nested
-        // borrows.  Fall through to add the begin_borrow to roots.
-      } else if (auto *result =
-              dyn_cast<FirstArgOwnershipForwardingSingleValueInst>(inst)) {
-        if (result->getNumOperands() > 0) {
-          worklist.insert(result->getOperand(0));
-          continue;
-        }
-      } else if (auto *result =
-                     dyn_cast<AllArgOwnershipForwardingSingleValueInst>(inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *result = dyn_cast<OwnershipForwardingTermInst>(inst)) {
-        assert(false && "value defined by a terminator?!");
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingConversionInst>(inst)) {
-        worklist.insert(result->getConverted());
-        continue;
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingSelectEnumInstBase>(inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *result =
-                     dyn_cast<OwnershipForwardingMultipleValueInstruction>(
-                         inst)) {
-        if (addAllOperandsToWorklist(result)) {
-          continue;
-        }
-      } else if (auto *mvi =
-                     dyn_cast<MoveOnlyWrapperToCopyableValueInst>(inst)) {
-        if (addAllOperandsToWorklist(mvi)) {
-          continue;
-        }
-      } else if (auto *c = dyn_cast<CopyableToMoveOnlyWrapperValueInst>(inst)) {
-        if (addAllOperandsToWorklist(c)) {
+          worklist.pushIfNotVisited(borrowee);
           continue;
         }
+        // The borrowee isn't guaranteed or we aren't looking through nested
+        // borrows. Fall through to add the begin_borrow to roots.
       }
     }
-
-    if (value->getOwnershipKind() == OwnershipKind::Guaranteed)
-      roots.push_back(value);
+    roots.push_back(value);
   }
 }