Merge pull request swiftlang#79736 from atrick/fix-address-walk

Add verification to catch many incorrect liveness checks

Author: atrick

include/swift/SIL/AddressUseKind.h

@@ -15,7 +15,7 @@
 
 namespace swift {
 
-enum class AddressUseKind { NonEscaping, PointerEscape, Unknown };
+enum class AddressUseKind { NonEscaping, Dependent, PointerEscape, Unknown };
 
 inline AddressUseKind meet(AddressUseKind lhs, AddressUseKind rhs) {
   return (lhs > rhs) ? lhs : rhs;

include/swift/SIL/AddressWalker.h

@@ -48,6 +48,7 @@ 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;
+  Operand *escapingUse = nullptr;
 
   unsigned didInvalidate = false;
 
@@ -83,6 +84,11 @@ class TransitiveAddressWalker {
     result = swift::meet(result, other);
   }
 
+  void recordEscape(Operand *op, AddressUseKind kind = AddressUseKind::PointerEscape) {
+    meet(kind);
+    escapingUse = op;
+  }
+
 private:
   /// Shim that actually calls visitUse and changes early exit.
   void callVisitUse(Operand *use) {
@@ -92,12 +98,16 @@ class TransitiveAddressWalker {
   }
 
 public:
-  AddressUseKind walk(SILValue address) &&;
+  AddressUseKind walk(SILValue address);
+
+  // If the result of walk() is not NonEscaping, this returns a non-null
+  // operand that caused the escape or unknown use.
+  Operand *getEscapingUse() const { return escapingUse; }
 };
 
 template <typename Impl>
 inline AddressUseKind
-TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
+TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) {
   assert(!didInvalidate);
 
   // When we exit, set the result to be invalidated so we can't use this again.
@@ -155,7 +165,7 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
       case TermKind::CondBranchInst:
         // We could have an address phi. To be conservative, just treat this as
         // a point escape.
-        meet(AddressUseKind::PointerEscape);
+        recordEscape(op);
         for (auto succBlockArgList : ti->getSuccessorBlockArgumentLists()) {
           auto *succ = succBlockArgList[op->getOperandNumber()];
           for (auto *use : succ->getUses())
@@ -185,10 +195,13 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
       }
     }
 
-    // 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);
+    if (isa<PartialApplyInst>(user)) {
+      recordEscape(op, AddressUseKind::Dependent);
+      callVisitUse(op);
+      continue;
+    }
+    if (isa<AddressToPointerInst>(user)) {
+      recordEscape(op);
       callVisitUse(op);
       continue;
     }
@@ -297,7 +310,7 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
       // address. See AddressUtils.swift. Until that is implemented, this must
       // be considered a pointer escape.
       if (op->get() == mdi->getBase()) {
-        meet(AddressUseKind::PointerEscape);
+        recordEscape(op, AddressUseKind::Dependent);
         callVisitUse(op);
         continue;
       }

include/swift/SIL/OwnershipLiveness.h

@@ -241,6 +241,7 @@ class InteriorLiveness : public OSSALiveness {
 public:
   // Summarize address uses
   AddressUseKind addressUseKind = AddressUseKind::Unknown;
+  Operand *escapingUse = nullptr;
 
 public:
   InteriorLiveness(SILValue def): OSSALiveness(def) {}

include/swift/SIL/OwnershipUseVisitor.h

@@ -494,7 +494,7 @@ bool OwnershipUseVisitor<Impl>::visitInteriorPointerUses(Operand *use) {
   // just assumes that all scopes are incomplete.
   SmallVector<Operand *, 8> interiorUses;
   auto useKind = InteriorPointerOperand(use).findTransitiveUses(&interiorUses);
-  if (useKind == AddressUseKind::PointerEscape) {
+  if (useKind != AddressUseKind::NonEscaping) {
     if (!asImpl().handlePointerEscape(use))
       return false;
   }

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -61,6 +61,13 @@
 
 using namespace swift;
 
+// FIXME: remove this option after fixing:
+// rdar://145994924 (Mem2Reg calls lifetime completion without checking for
+// pointer escapes)
+llvm::cl::opt<bool> VerifyLifetimeCompletion(
+    "verify-lifetime-completion", llvm::cl::init(false),
+    llvm::cl::desc("."));
+
 static SILInstruction *endOSSALifetime(SILValue value,
                                        OSSALifetimeCompletion::LifetimeEnd end,
                                        SILBuilder &builder,
@@ -493,8 +500,16 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(
     }
   };
   Walker walker(*this, scopedAddress, boundary, liveness);
-  std::move(walker).walk(scopedAddress.value);
-
+  AddressUseKind result = walker.walk(scopedAddress.value);
+  if (VerifyLifetimeCompletion && boundary != Boundary::Availability
+      && result != AddressUseKind::NonEscaping) {
+    llvm::errs() << "Incomplete liveness for:\n" << scopedAddress.value;
+    if (auto *escapingUse = walker.getEscapingUse()) {
+      llvm::errs() << "  escapes at:\n";
+      escapingUse->getUser()->printInContext(llvm::errs());
+    }
+    ASSERT(false && "caller must check for pointer escapes");
+  }
   return endLifetimeAtBoundary(scopedAddress.value, liveness, boundary,
                                deadEndBlocks);
 }
@@ -524,6 +539,15 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
   }
   InteriorLiveness liveness(value);
   liveness.compute(domInfo, handleInnerScope);
+  if (VerifyLifetimeCompletion && boundary != Boundary::Availability
+      && liveness.getAddressUseKind() != AddressUseKind::NonEscaping) {
+    llvm::errs() << "Incomplete liveness for: " << value;
+    if (auto *escapingUse = liveness.escapingUse) {
+      llvm::errs() << "  escapes at:\n";
+      escapingUse->getUser()->printInContext(llvm::errs());
+    }
+    ASSERT(false && "caller must check for pointer escapes");
+  }
   return endLifetimeAtBoundary(value, liveness.getLiveness(), boundary,
                                deadEndBlocks);
 }

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -150,6 +150,7 @@ struct InteriorLivenessVisitor :
 
   bool handlePointerEscape(Operand *use) {
     interiorLiveness.addressUseKind = AddressUseKind::PointerEscape;
+    interiorLiveness.escapingUse = use;
     if (!handleUsePoint(use, UseLifetimeConstraint::NonLifetimeEnding))
       return false;
 
@@ -198,6 +199,9 @@ void InteriorLiveness::print(llvm::raw_ostream &OS) const {
   case AddressUseKind::PointerEscape:
     OS << "Incomplete liveness: Escaping address\n";
     break;
+  case AddressUseKind::Dependent:
+    OS << "Incomplete liveness: Dependent value\n";
+    break;
   case AddressUseKind::Unknown:
     OS << "Incomplete liveness: Unknown address use\n";
     break;

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -899,10 +899,16 @@ static SILValue tryRewriteToPartialApplyStack(
 
     OrigUnmodifiedDuringClosureLifetimeWalker origUseWalker(
         closureLiveness, origIsUnmodifiedDuringClosureLifetime);
-    auto walkResult = std::move(origUseWalker).walk(orig);
-
-    if (walkResult == AddressUseKind::Unknown ||
-        !origIsUnmodifiedDuringClosureLifetime) {
+    switch (origUseWalker.walk(orig)) {
+    case AddressUseKind::NonEscaping:
+    case AddressUseKind::Dependent:
+      // Dependent uses are ignored because they cannot modify the original.
+      break;
+    case AddressUseKind::PointerEscape:
+    case AddressUseKind::Unknown:
+      continue;
+    }
+    if (!origIsUnmodifiedDuringClosureLifetime) {
       continue;
     }
 

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -3827,6 +3827,8 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
 
     CopiedLoadBorrowEliminationState state(markedAddress->getFunction());
     CopiedLoadBorrowEliminationVisitor copiedLoadBorrowEliminator(state);
+    // FIXME: should check AddressUseKind::NonEscaping != walk() to handle
+    // PointerEscape.
     if (AddressUseKind::Unknown ==
         std::move(copiedLoadBorrowEliminator).walk(markedAddress)) {
       LLVM_DEBUG(llvm::dbgs() << "Failed copied load borrow eliminator visit: "
@@ -3869,6 +3871,8 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
     RAIILLVMDebug l("main use gathering visitor");
 
     visitor.reset(markedAddress);
+    // FIXME: should check walkResult != AddressUseKind::NonEscaping to handle
+    // PointerEscape.
     if (AddressUseKind::Unknown == std::move(visitor).walk(markedAddress)) {
       LLVM_DEBUG(llvm::dbgs()
                  << "Failed access path visit: " << *markedAddress);

lib/SILOptimizer/Mandatory/MoveOnlyUtils.cpp

@@ -583,6 +583,7 @@ bool siloptimizer::eliminateTemporaryAllocationsFromLet(
   };
   FindCopyAddrWalker walker(copiesToVisit);
   std::move(walker).walk(markedInst);
+  // FIXME: should check walk() == AddressUseKind::NonEscaping.
 
   bool madeChange = false;
 
@@ -612,6 +613,8 @@ bool siloptimizer::eliminateTemporaryAllocationsFromLet(
       nextCAI = nullptr;
       SimpleTemporaryAllocStackElimVisitor visitor(state, cai, nextCAI);
 
+      // FIXME: should check AddressUseKind::NonEscaping != walk() to handle
+      // PointerEscape.
       if (AddressUseKind::Unknown == std::move(visitor).walk(cai->getDest()))
         return false;
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1598,7 +1598,7 @@ void swift::insertDeallocOfCapturedArguments(
     SmallVector<Operand *, 2> uses;
     auto useFinding = findTransitiveUsesForAddress(asi, &uses);
     InstructionSet users(asi->getFunction());
-    if (useFinding != AddressUseKind::Unknown) {
+    if (useFinding == AddressUseKind::NonEscaping) {
       for (auto use : uses) {
         users.insert(use->getUser());
       }
@@ -1609,7 +1609,7 @@ void swift::insertDeallocOfCapturedArguments(
       if (isa<UnreachableInst>(terminator))
         continue;
       SILInstruction *insertionPoint = nullptr;
-      if (useFinding != AddressUseKind::Unknown) {
+      if (useFinding == AddressUseKind::NonEscaping) {
         insertionPoint = &block->front();
         for (auto &instruction : llvm::reverse(*block)) {
           if (users.contains(&instruction)) {
@@ -2463,6 +2463,7 @@ SILValue swift::getInitOfTemporaryAllocStack(AllocStackInst *asi) {
 
   AddressWalkerState state(asi->getFunction());
   AddressWalker walker(state);
+  // Note: ignore pointer escapes for the purpose of finding initializers.
   if (std::move(walker).walk(asi) == AddressUseKind::Unknown ||
       state.foundError)
     return SILValue();