Fix visitNonEscapingLifetimeEnds to handle mark_dependence uses

atrick · atrick · commit d7b9149ee516 · 2024-03-22T14:29:57.000-07:00
Now it visits unknown uses separately rather than asserting.
diff --git a/SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowUtils.swift b/SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowUtils.swift
@@ -206,6 +206,11 @@ enum BorrowingInstruction : CustomStringConvertible, Hashable {
   /// incoming value dominates or is consumed by an outer adjacent
   /// phi. See InteriorLiveness.
   ///
+  /// FIXME: To generate conservatively correct liveness, this should return
+  /// .abortWalk if this is a mark_dependence and the scope-ending use is not
+  /// the last in the function (e.g. a store rather than a destroy or return).
+  /// The client needs to use LifetimeDependenceDefUseWalker to do better.
+  ///
   /// TODO: to hande reborrow-extended uses, migrate ExtendedLiveness
   /// to SwiftCompilerSources.
   ///
diff --git a/include/swift/SIL/OwnershipUseVisitor.h b/include/swift/SIL/OwnershipUseVisitor.h
@@ -258,9 +258,13 @@ bool OwnershipUseVisitor<Impl>::visitInnerBorrow(Operand *borrowingOperand) {
     return false;
 
   return BorrowingOperand(borrowingOperand)
-    .visitScopeEndingUses([&](Operand *borrowEnd) {
-      return visitInnerBorrowScopeEnd(borrowEnd);
-    });
+    .visitScopeEndingUses(
+      [&](Operand *borrowEnd) {
+        return visitInnerBorrowScopeEnd(borrowEnd);
+      },
+      [&](Operand *unknownUse) {
+        return asImpl().handlePointerEscape(unknownUse);
+      });
 }
 
 template <typename Impl>
diff --git a/include/swift/SIL/OwnershipUtils.h b/include/swift/SIL/OwnershipUtils.h
@@ -360,12 +360,16 @@ struct BorrowingOperand {
   /// over a region of code instead of just for a single instruction, visit
   /// those uses.
   ///
-  /// Returns false and early exits if the visitor \p func returns false.
+  /// Returns false and early exits if the \p visitScopeEnd or \p
+  /// visitUnknownUse returns false.
   ///
   /// For an instantaneous borrow, such as apply, this visits no uses. For
   /// begin_apply it visits the end_apply uses. For borrow introducers, it
   /// visits the end of the introduced borrow scope.
-  bool visitScopeEndingUses(function_ref<bool(Operand *)> func) const;
+  bool visitScopeEndingUses(function_ref<bool(Operand *)> visitScopeEnd,
+                            function_ref<bool(Operand *)> visitUnknownUse
+                            = [](Operand *){ return false; })
+    const;
 
   /// Returns true for borrows that create a local borrow scope but have no
   /// scope-ending uses (presumably all paths are dead-end blocks). This does
@@ -381,7 +385,10 @@ struct BorrowingOperand {
   /// BorrowingOperand.
   ///
   /// Returns false and early exits if the visitor \p func returns false.
-  bool visitExtendedScopeEndingUses(function_ref<bool(Operand *)> func) const;
+  bool visitExtendedScopeEndingUses(
+    function_ref<bool(Operand *)> func,
+    function_ref<bool(Operand *)> visitUnknownUse
+    = [](Operand *){ return false; }) const;
 
   /// Returns true if this borrow scope operand consumes guaranteed
   /// values and produces a new scope afterwards.
diff --git a/include/swift/SIL/PrunedLiveness.h b/include/swift/SIL/PrunedLiveness.h
@@ -289,8 +289,9 @@ class PrunedLiveBlocks {
 /// scope. Reborrows within nested borrows scopes are already summarized by the
 /// outer borrow scope.
 enum class InnerBorrowKind {
-  Contained, // any borrows are fully contained within this live range
-  Reborrowed // at least one immediately nested borrow is reborrowed
+  Contained,  // any borrows are fully contained within this live range
+  Reborrowed, // at least one immediately nested borrow is reborrowed
+  Escaped     // the end of the borrow scope is indeterminate
 };
 
 inline InnerBorrowKind meet(InnerBorrowKind lhs, InnerBorrowKind rhs) {
diff --git a/include/swift/SIL/SILInstruction.h b/include/swift/SIL/SILInstruction.h
@@ -8546,8 +8546,9 @@ class MarkDependenceInst
   }
 
   /// Visit the instructions that end the lifetime of an OSSA on-stack closure.
-  bool visitNonEscapingLifetimeEnds(llvm::function_ref<bool (Operand*)> func)
-    const;
+  bool visitNonEscapingLifetimeEnds(
+    llvm::function_ref<bool (Operand*)> visitScopeEnd,
+    llvm::function_ref<bool (Operand*)> visitUnknownUse) const;
 };
 
 /// Promote an Objective-C block that is on the stack to the heap, or simply
diff --git a/lib/SIL/IR/SILInstruction.cpp b/lib/SIL/IR/SILInstruction.cpp
@@ -1796,33 +1796,27 @@ bool SILInstruction::maySuspend() const {
 }
 
 static bool
-visitRecursivelyLifetimeEndingUses(SILValue i, bool &noUsers,
-                                   llvm::function_ref<bool(Operand *)> func,
-                                   bool allowPhis) {
+visitRecursivelyLifetimeEndingUses(
+  SILValue i, bool &noUsers,
+  llvm::function_ref<bool(Operand *)> visitScopeEnd,
+  llvm::function_ref<bool(Operand *)> visitUnknownUse) {
+
   for (Operand *use : i->getConsumingUses()) {
     noUsers = false;
     if (isa<DestroyValueInst>(use->getUser())) {
-      if (!func(use)) {
+      if (!visitScopeEnd(use)) {
         return false;
       }
       continue;
     }
     if (auto *ret = dyn_cast<ReturnInst>(use->getUser())) {
       auto fnTy = ret->getFunction()->getLoweredFunctionType();
       assert(!fnTy->getLifetimeDependenceInfo().empty());
-      if (!func(use)) {
+      if (!visitScopeEnd(use)) {
         return false;
       }
       continue;
     }
-    if (allowPhis) {
-      if (PhiOperand(use)) {
-        if (!func(use)) {
-          return false;
-        }
-        continue;
-      }
-    }
     // FIXME: Handle store to indirect result
     
     // There shouldn't be any dead-end consumptions of a nonescaping
@@ -1836,17 +1830,11 @@ visitRecursivelyLifetimeEndingUses(SILValue i, bool &noUsers,
     // the structural requirements of on-stack partial_apply uses.
     auto *user = use->getUser();
     if (user->getNumResults() == 0) {
-      llvm::errs() << "partial_apply [on_stack] use:\n";
-      user->printInContext(llvm::errs());
-      if (isa<BranchInst>(user)) {
-        llvm::report_fatal_error("partial_apply [on_stack] cannot be cloned");
-      }
-      llvm::report_fatal_error("partial_apply [on_stack] must be directly "
-                               "forwarded to a destroy_value");
+      return visitUnknownUse(use);
     }
     for (auto res : use->getUser()->getResults()) {
-      if (!visitRecursivelyLifetimeEndingUses(res, noUsers, func,
-                                              allowPhis)) {
+      if (!visitRecursivelyLifetimeEndingUses(res, noUsers, visitScopeEnd,
+                                              visitUnknownUse)) {
         return false;
       }
     }
@@ -1862,21 +1850,36 @@ PartialApplyInst::visitOnStackLifetimeEnds(
          && "only meaningful for OSSA stack closures");
   bool noUsers = true;
 
+  auto visitUnknownUse = [](Operand *unknownUse){
+    llvm::errs() << "partial_apply [on_stack] use:\n";
+    auto *user = unknownUse->getUser();
+    user->printInContext(llvm::errs());
+    if (isa<BranchInst>(user)) {
+      llvm::report_fatal_error("partial_apply [on_stack] cannot be cloned");
+    }
+    llvm::report_fatal_error("partial_apply [on_stack] must be directly "
+                             "forwarded to a destroy_value");
+    return false;
+  };
   if (!visitRecursivelyLifetimeEndingUses(this, noUsers, func,
-                                          /*allowPhis*/ false)) {
+                                          visitUnknownUse)) {
     return false;
   }
   return !noUsers;
 }
 
-bool MarkDependenceInst::
-visitNonEscapingLifetimeEnds(llvm::function_ref<bool (Operand *)> func) const {
+// FIXME: Rather than recursing through all results, this should only recurse
+// through ForwardingInstruction and OwnershipTransitionInstruction and the
+// client should prove that any other uses cannot be upstream from a consume of
+// the dependent value.
+bool MarkDependenceInst::visitNonEscapingLifetimeEnds(
+  llvm::function_ref<bool (Operand *)> visitScopeEnd,
+  llvm::function_ref<bool (Operand *)> visitUnknownUse) const {
   assert(getFunction()->hasOwnership() && isNonEscaping()
          && "only meaningful for nonescaping dependencies");
   bool noUsers = true;
-
-  if (!visitRecursivelyLifetimeEndingUses(this, noUsers, func,
-                                          /*allowPhis*/ true)) {
+  if (!visitRecursivelyLifetimeEndingUses(this, noUsers, visitScopeEnd,
+                                          visitUnknownUse)) {
     return false;
   }
   return !noUsers;
diff --git a/lib/SIL/Utils/OwnershipUtils.cpp b/lib/SIL/Utils/OwnershipUtils.cpp
@@ -327,18 +327,19 @@ bool swift::findInnerTransitiveGuaranteedUses(
       //
       // FIXME: visit[Extended]ScopeEndingUses can't return false here once dead
       // borrows are disallowed.
-      if (!BorrowingOperand(use).visitScopeEndingUses([&](Operand *endUse) {
-            if (endUse->getOperandOwnership() == OperandOwnership::Reborrow) {
-              foundPointerEscape = true;
-            }
-            if (PhiOperand(endUse)) {
-              assert(endUse->getOperandOwnership() ==
-                     OperandOwnership::ForwardingConsume);
+      if (!BorrowingOperand(use).visitScopeEndingUses(
+            [&](Operand *endUse) {
+              if (endUse->getOperandOwnership() == OperandOwnership::Reborrow) {
+                foundPointerEscape = true;
+              }
+              leafUse(endUse);
+              return true;
+            },
+            [&](Operand *unknownUse) {
               foundPointerEscape = true;
-            }
-            leafUse(endUse);
-            return true;
-          })) {
+              leafUse(unknownUse);
+              return true;
+            })) {
         // Special case for dead borrows. This is dangerous because clients
         // don't expect a begin_borrow to be in the use list.
         leafUse(use);
@@ -451,16 +452,12 @@ bool swift::findExtendedUsesOfSimpleBorrowedValue(
       break;
     }
     case OperandOwnership::Borrow:
-      // FIXME: visitExtendedScopeEndingUses can't return false here once dead
-      // borrows are disallowed.
       if (!BorrowingOperand(use).visitExtendedScopeEndingUses(
             [&](Operand *endUse) {
               recordUse(endUse);
               return true;
             })) {
-        // Special case for dead borrows. This is dangerous because clients
-        // don't expect a begin_borrow to be in the use list.
-        recordUse(use);
+        return false;
       }
       break;
     }
@@ -480,11 +477,6 @@ bool swift::findUsesOfSimpleValue(SILValue value,
         if (end->getOperandOwnership() == OperandOwnership::Reborrow) {
           return false;
         }
-        if (PhiOperand(use)) {
-          assert(use->getOperandOwnership() ==
-                 OperandOwnership::ForwardingConsume);
-          return false;
-        }
         usePoints->push_back(end);
         return true;
       })) {
@@ -694,7 +686,8 @@ bool BorrowingOperand::hasEmptyRequiredEndingUses() const {
 }
 
 bool BorrowingOperand::visitScopeEndingUses(
-    function_ref<bool(Operand *)> func) const {
+  function_ref<bool(Operand *)> visitScopeEnd,
+  function_ref<bool(Operand *)> visitUnknownUse) const {
   switch (kind) {
   case BorrowingOperandKind::Invalid:
     llvm_unreachable("Using invalid case");
@@ -703,7 +696,7 @@ bool BorrowingOperand::visitScopeEndingUses(
     for (auto *use : cast<BeginBorrowInst>(op->getUser())->getUses()) {
       if (use->isLifetimeEnding()) {
         deadBorrow = false;
-        if (!func(use))
+        if (!visitScopeEnd(use))
           return false;
       }
     }
@@ -717,7 +710,7 @@ bool BorrowingOperand::visitScopeEndingUses(
     for (auto *use : cast<StoreBorrowInst>(op->getUser())->getUses()) {
       if (isa<EndBorrowInst>(use->getUser())) {
         deadBorrow = false;
-        if (!func(use))
+        if (!visitScopeEnd(use))
           return false;
       }
     }
@@ -731,7 +724,7 @@ bool BorrowingOperand::visitScopeEndingUses(
     auto *user = cast<BeginApplyInst>(op->getUser());
     for (auto *use : user->getTokenResult()->getUses()) {
       deadApply = false;
-      if (!func(use))
+      if (!visitScopeEnd(use))
         return false;
     }
     return !deadApply;
@@ -742,12 +735,12 @@ bool BorrowingOperand::visitScopeEndingUses(
     // The closure's borrow lifetimes end when the closure itself ends its
     // lifetime. That may happen transitively through conversions that forward
     // ownership of the closure.
-    return user->visitOnStackLifetimeEnds(func);
+    return user->visitOnStackLifetimeEnds(visitScopeEnd);
   }
   case BorrowingOperandKind::MarkDependenceNonEscaping: {
     auto *user = cast<MarkDependenceInst>(op->getUser());
     assert(user->isNonEscaping() && "escaping dependencies don't borrow");
-    return user->visitNonEscapingLifetimeEnds(func);
+    return user->visitNonEscapingLifetimeEnds(visitScopeEnd, visitUnknownUse);
   }
   case BorrowingOperandKind::BeginAsyncLet: {
     auto user = cast<BuiltinInst>(op->getUser());
@@ -760,7 +753,7 @@ bool BorrowingOperand::visitScopeEndingUses(
           || builtinUser->getBuiltinKind() != BuiltinValueKind::EndAsyncLetLifetime)
         continue;
 
-      if (!func(use)) {
+      if (!visitScopeEnd(use)) {
         return false;
       }
     }
@@ -778,7 +771,7 @@ bool BorrowingOperand::visitScopeEndingUses(
     for (auto *use : br->getArgForOperand(op)->getUses()) {
       if (use->isLifetimeEnding()) {
         deadBranch = false;
-        if (!func(use))
+        if (!visitScopeEnd(use))
           return false;
       }
     }
@@ -789,13 +782,14 @@ bool BorrowingOperand::visitScopeEndingUses(
 }
 
 bool BorrowingOperand::visitExtendedScopeEndingUses(
-    function_ref<bool(Operand *)> visitor) const {
+  function_ref<bool(Operand *)> visitor,
+  function_ref<bool(Operand *)> visitUnknownUse) const {
 
   if (hasBorrowIntroducingUser()) {
     auto borrowedValue = getBorrowIntroducingUserResult();
     return borrowedValue.visitExtendedScopeEndingUses(visitor);
   }
-  return visitScopeEndingUses(visitor);
+  return visitScopeEndingUses(visitor, visitUnknownUse);
 }
 
 BorrowedValue BorrowingOperand::getBorrowIntroducingUserResult() const {
@@ -857,10 +851,11 @@ void BorrowingOperand::getImplicitUses(
     SmallVectorImpl<Operand *> &foundUses) const {
   // FIXME: this visitScopeEndingUses should never return false once dead
   // borrows are disallowed.
-  if (!visitScopeEndingUses([&](Operand *endOp) {
+  auto handleUse = [&](Operand *endOp) {
     foundUses.push_back(endOp);
     return true;
-  })) {
+  };
+  if (!visitScopeEndingUses(handleUse, handleUse)) {
     // Special-case for dead borrows.
     foundUses.push_back(op);
   }
@@ -988,9 +983,6 @@ bool BorrowedValue::visitExtendedScopeEndingUses(
       reborrows.insert(borrowedValue.value);
       return true;
     }
-    if (auto phiOp = PhiOperand(scopeEndingUse)) {
-      return false;
-    }
     return visitor(scopeEndingUse);
   };
 
diff --git a/lib/SIL/Utils/PrunedLiveness.cpp b/lib/SIL/Utils/PrunedLiveness.cpp
@@ -234,21 +234,23 @@ PrunedLiveRange<LivenessWithDefs>::updateForBorrowingOperand(Operand *operand) {
   //
   // Note: Ownership liveness should follow reborrows that are dominated by the
   // ownership definition.
-  if (!BorrowingOperand(operand).visitScopeEndingUses([this](Operand *end) {
-        if (end->getOperandOwnership() == OperandOwnership::Reborrow) {
-          return false;
-        }
-        if (PhiOperand(end)) {
-          assert(end->getOperandOwnership() ==
-                 OperandOwnership::ForwardingConsume);
-          return false;
-        }
-        updateForUse(end->getUser(), /*lifetimeEnding*/ false);
-        return true;
-      })) {
-    return InnerBorrowKind::Reborrowed;
+  auto innerBorrowKind = InnerBorrowKind::Contained;
+  if (!BorrowingOperand(operand).visitScopeEndingUses(
+        [&](Operand *end) {
+          if (end->getOperandOwnership() == OperandOwnership::Reborrow) {
+            innerBorrowKind = InnerBorrowKind::Reborrowed;
+          }
+          updateForUse(end->getUser(), /*lifetimeEnding*/ false);
+          return true;
+        }, [&](Operand *unknownUse) {
+          updateForUse(unknownUse->getUser(), /*lifetimeEnding*/ false);
+          innerBorrowKind = InnerBorrowKind::Escaped;
+          return true;
+        })) {
+    // Handle dead borrows.
+    updateForUse(operand->getUser(), /*lifetimeEnding*/ false);
   }
-  return InnerBorrowKind::Contained;
+  return innerBorrowKind;
 }
 
 template <typename LivenessWithDefs>
diff --git a/lib/SILOptimizer/Utils/CanonicalizeBorrowScope.cpp b/lib/SILOptimizer/Utils/CanonicalizeBorrowScope.cpp
diff --git a/test/SILOptimizer/copy_propagation.sil b/test/SILOptimizer/copy_propagation.sil
diff --git a/test/SILOptimizer/liveness_unit.sil b/test/SILOptimizer/liveness_unit.sil
