Merge pull request #67874 from gottesmm/pr-179e2ec3ed29bb0d8c35dc99f710a18571a0a24e

[sil] Convert AddressWalker from using virtual methods to use CRTP and add a transitive -> endpoint user API

Author: gottesmm

include/swift/SIL/AddressWalker.h

@@ -21,7 +21,11 @@
 #ifndef SWIFT_SIL_ADDRESSWALKER_H
 #define SWIFT_SIL_ADDRESSWALKER_H
 
+#include "swift/Basic/Defer.h"
 #include "swift/SIL/AddressUseKind.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
+#include "swift/SIL/InstructionUtils.h"
+#include "swift/SIL/Projection.h"
 #include "swift/SIL/SILValue.h"
 
 namespace swift {
@@ -33,15 +37,16 @@ namespace swift {
 /// Validated by the SIL verifier as always being able to visit all addresses
 /// derived from alloc_stack, ref_element_addr, project_box, ref_tail_addr and
 /// all other address roots.
+template <typename Impl>
 class TransitiveAddressWalker {
   /// Whether we could tell if this address use didn't escape, did have a
   /// pointer escape, or unknown if we failed to understand something.
   AddressUseKind result = AddressUseKind::NonEscaping;
 
   unsigned didInvalidate = false;
 
-public:
-  virtual ~TransitiveAddressWalker() {}
+  Impl &asImpl() { return *reinterpret_cast<Impl *>(this); }
+  const Impl &asImpl() const { return *reinterpret_cast<const Impl *>(this); }
 
 protected:
   /// Customization point for visiting uses. Returns true if we should continue
@@ -50,9 +55,17 @@ class TransitiveAddressWalker {
   /// NOTE: Do not call this directly from within
   /// findTransitiveUsesForAddress. Please call callVisitUse. This is intended
   /// just for subclasses to override.
-  virtual bool visitUse(Operand *use) { return true; }
+  bool visitUse(Operand *use) { return true; }
+
+  /// Customization point for visiting operands that we were unable to
+  /// understand. These cause us to return AddressUseKind::Unknown.
+  void onError(Operand *use) {}
 
-  virtual void onError(Operand *use) {}
+  /// Customization point that causes the walker to treat a specific transitive
+  /// use as an end point use.
+  ///
+  /// Example: Visiting a mutable or immutable open_existential_addr.
+  bool visitTransitiveUseAsEndPointUse(Operand *use) { return false; }
 
   void meet(AddressUseKind other) {
     assert(!didInvalidate);
@@ -63,21 +76,219 @@ class TransitiveAddressWalker {
   /// Shim that actually calls visitUse and changes early exit.
   void callVisitUse(Operand *use) {
     assert(!didInvalidate);
-    if (!visitUse(use))
+    if (!asImpl().visitUse(use))
       result = AddressUseKind::Unknown;
   }
 
 public:
   AddressUseKind walk(SILValue address) &&;
 };
 
-/// The algorithm that is used to determine what the verifier will consider to
-/// be transitive uses of the given address. Used to implement \see
-/// findTransitiveUses.
-///
-/// Returns \p AccessUseKind::Unknown on error.
-AddressUseKind findTransitiveUsesForAddress(SILValue address,
-                                            TransitiveAddressWalker &visitor);
+template <typename Impl>
+inline AddressUseKind
+TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
+  assert(!didInvalidate);
+
+  // When we exit, set the result to be invalidated so we can't use this again.
+  SWIFT_DEFER { didInvalidate = true; };
+
+  StackList<Operand *> worklist(projectedAddress->getFunction());
+  SmallPtrSet<Operand *, 32> visitedOperands;
+
+  auto addToWorklist = [&](Operand *use) {
+    if (visitedOperands.insert(use).second)
+      worklist.push_back(use);
+  };
+
+  for (auto *use : projectedAddress->getUses()) {
+    addToWorklist(use);
+  }
+
+  // Record all uses that aren't transitively followed. These are either
+  // instanteneous uses of the addres, or cause a pointer escape.
+  auto transitiveResultUses = [&](Operand *use) {
+    auto *svi = cast<SingleValueInstruction>(use->getUser());
+    if (svi->use_empty()) {
+      return callVisitUse(use);
+    }
+
+    if (asImpl().visitTransitiveUseAsEndPointUse(use))
+      return callVisitUse(use);
+
+    for (auto *use : svi->getUses())
+      addToWorklist(use);
+  };
+
+  while (!worklist.empty()) {
+    if (result == AddressUseKind::Unknown)
+      return AddressUseKind::Unknown;
+
+    auto *op = worklist.pop_back_val();
+
+    // Skip type dependent operands.
+    if (op->isTypeDependent())
+      continue;
+
+    // Then update the worklist with new things to find if we recognize this
+    // inst and then continue. If we fail, we emit an error at the bottom of the
+    // loop that we didn't recognize the user.
+    auto *user = op->getUser();
+
+    if (auto *ti = dyn_cast<TermInst>(user)) {
+      switch (ti->getTermKind()) {
+      case TermKind::BranchInst:
+      case TermKind::CondBranchInst:
+        // We could have an address phi. To be conservative, just treat this as
+        // a point escape.
+        meet(AddressUseKind::PointerEscape);
+        for (auto succBlockArgList : ti->getSuccessorBlockArgumentLists()) {
+          auto *succ = succBlockArgList[op->getOperandNumber()];
+          for (auto *use : succ->getUses())
+            addToWorklist(use);
+        }
+        continue;
+
+      case TermKind::UnreachableInst:
+      case TermKind::UnwindInst:
+        llvm_unreachable("Should never be used");
+      case TermKind::SwitchEnumInst:
+      case TermKind::SwitchValueInst:
+      case TermKind::DynamicMethodBranchInst:
+      case TermKind::AwaitAsyncContinuationInst:
+      case TermKind::CheckedCastBranchInst:
+        llvm_unreachable("Never takes an address");
+      // Point uses.
+      case TermKind::ReturnInst:
+      case TermKind::ThrowInst:
+      case TermKind::YieldInst:
+      case TermKind::TryApplyInst:
+      case TermKind::SwitchEnumAddrInst:
+      case TermKind::CheckedCastAddrBranchInst:
+        callVisitUse(op);
+        continue;
+      }
+    }
+
+    // TODO: Partial apply should be NonEscaping, but then we need to consider
+    // the apply to be a use point.
+    if (isa<PartialApplyInst>(user) || isa<AddressToPointerInst>(user)) {
+      meet(AddressUseKind::PointerEscape);
+      callVisitUse(op);
+      continue;
+    }
+
+    // First, eliminate "end point uses" that we just need to check liveness at
+    // and do not need to check transitive uses of.
+    if (isa<LoadInst>(user) || isa<CopyAddrInst>(user) ||
+        isa<MarkUnresolvedMoveAddrInst>(user) || isIncidentalUse(user) ||
+        isa<StoreInst>(user) || isa<DestroyAddrInst>(user) ||
+        isa<AssignInst>(user) || isa<LoadUnownedInst>(user) ||
+        isa<StoreUnownedInst>(user) || isa<EndApplyInst>(user) ||
+        isa<LoadWeakInst>(user) || isa<StoreWeakInst>(user) ||
+        isa<AssignByWrapperInst>(user) || isa<AssignOrInitInst>(user) ||
+        isa<BeginUnpairedAccessInst>(user) ||
+        isa<EndUnpairedAccessInst>(user) || isa<WitnessMethodInst>(user) ||
+        isa<SelectEnumAddrInst>(user) || isa<InjectEnumAddrInst>(user) ||
+        isa<IsUniqueInst>(user) || isa<ValueMetatypeInst>(user) ||
+        isa<DebugValueInst>(user) || isa<EndBorrowInst>(user) ||
+        isa<ExplicitCopyAddrInst>(user) || isa<DeallocStackInst>(user) ||
+        isa<InitBlockStorageHeaderInst>(user) ||
+        isa<GetAsyncContinuationAddrInst>(user) ||
+        isa<ExistentialMetatypeInst>(user) ||
+        isa<UncheckedRefCastAddrInst>(user) || isa<KeyPathInst>(user) ||
+        isa<RetainValueAddrInst>(user) || isa<ReleaseValueAddrInst>(user) ||
+        isa<PackElementSetInst>(user) || isa<PackElementGetInst>(user) ||
+        isa<DeinitExistentialAddrInst>(user) || isa<LoadBorrowInst>(user)) {
+      callVisitUse(op);
+      continue;
+    }
+
+    if (isa<UnconditionalCheckedCastAddrInst>(user) ||
+        isa<MarkFunctionEscapeInst>(user)) {
+      assert(!user->hasResults());
+      callVisitUse(op);
+      continue;
+    }
+
+    // Then handle users that we need to look at transitive uses of.
+    if (Projection::isAddressProjection(user) ||
+        isa<ProjectBlockStorageInst>(user) ||
+        isa<OpenExistentialAddrInst>(user) ||
+        isa<InitExistentialAddrInst>(user) || isa<InitEnumDataAddrInst>(user) ||
+        isa<BeginAccessInst>(user) || isa<TailAddrInst>(user) ||
+        isa<IndexAddrInst>(user) || isa<StoreBorrowInst>(user) ||
+        isa<UncheckedAddrCastInst>(user) || isa<MarkMustCheckInst>(user) ||
+        isa<MarkUninitializedInst>(user) || isa<DropDeinitInst>(user) ||
+        isa<ProjectBlockStorageInst>(user) || isa<UpcastInst>(user) ||
+        isa<TuplePackElementAddrInst>(user) ||
+        isa<CopyableToMoveOnlyWrapperAddrInst>(user) ||
+        isa<MoveOnlyWrapperToCopyableAddrInst>(user)) {
+      transitiveResultUses(op);
+      continue;
+    }
+
+    if (auto *builtin = dyn_cast<BuiltinInst>(user)) {
+      if (auto kind = builtin->getBuiltinKind()) {
+        switch (*kind) {
+        case BuiltinValueKind::TSanInoutAccess:
+        case BuiltinValueKind::ResumeThrowingContinuationReturning:
+        case BuiltinValueKind::ResumeNonThrowingContinuationReturning:
+        case BuiltinValueKind::Copy:
+        case BuiltinValueKind::GenericAdd:
+        case BuiltinValueKind::GenericFAdd:
+        case BuiltinValueKind::GenericAnd:
+        case BuiltinValueKind::GenericAShr:
+        case BuiltinValueKind::GenericLShr:
+        case BuiltinValueKind::GenericOr:
+        case BuiltinValueKind::GenericFDiv:
+        case BuiltinValueKind::GenericMul:
+        case BuiltinValueKind::GenericFMul:
+        case BuiltinValueKind::GenericSDiv:
+        case BuiltinValueKind::GenericExactSDiv:
+        case BuiltinValueKind::GenericShl:
+        case BuiltinValueKind::GenericSRem:
+        case BuiltinValueKind::GenericSub:
+        case BuiltinValueKind::GenericFSub:
+        case BuiltinValueKind::GenericUDiv:
+        case BuiltinValueKind::GenericExactUDiv:
+        case BuiltinValueKind::GenericURem:
+        case BuiltinValueKind::GenericFRem:
+        case BuiltinValueKind::GenericXor:
+        case BuiltinValueKind::TaskRunInline:
+        case BuiltinValueKind::ZeroInitializer:
+          callVisitUse(op);
+          continue;
+        default:
+          break;
+        }
+      }
+    }
+
+    if (auto fas = FullApplySite::isa(user)) {
+      callVisitUse(op);
+      continue;
+    }
+
+    if (auto *mdi = dyn_cast<MarkDependenceInst>(user)) {
+      // If this is the base, just treat it as a liveness use.
+      if (op->get() == mdi->getBase()) {
+        callVisitUse(op);
+        continue;
+      }
+
+      // If we are the value use, look through it.
+      transitiveResultUses(op);
+      continue;
+    }
+
+    // We were unable to recognize this user, so set AddressUseKind to unknown
+    // and call onError with the specific user that caused the problem.
+    asImpl().onError(op);
+    return AddressUseKind::Unknown;
+  }
+
+  return result;
+}
 
 } // namespace swift
 

include/swift/SIL/OwnershipUtils.h

@@ -736,15 +736,16 @@ inline AddressUseKind findTransitiveUsesForAddress(
   // guaranteed uses to determine if a load_borrow is an escape in OSSA. This
   // is OSSA specific behavior and we should probably create a different API
   // for that. But for now, this lets this APIs users stay the same.
-  struct BasicTransitiveAddressVisitor final : TransitiveAddressWalker {
+  struct BasicTransitiveAddressVisitor
+      : TransitiveAddressWalker<BasicTransitiveAddressVisitor> {
     SmallVectorImpl<Operand *> *foundUses;
     std::function<void(Operand *)> *onErrorFunc;
 
     BasicTransitiveAddressVisitor(SmallVectorImpl<Operand *> *foundUses,
                                   std::function<void(Operand *)> *onErrorFunc)
         : foundUses(foundUses), onErrorFunc(onErrorFunc) {}
 
-    bool visitUse(Operand *use) override {
+    bool visitUse(Operand *use) {
       if (!foundUses)
         return true;
 
@@ -771,7 +772,7 @@ inline AddressUseKind findTransitiveUsesForAddress(
       return true;
     }
 
-    void onError(Operand *use) override {
+    void onError(Operand *use) {
       if (onErrorFunc)
         (*onErrorFunc)(use);
     }