[SIL] Verify extend_lifetime instruction.

Author: nate-chandler

lib/SIL/Verifier/LinearLifetimeChecker.cpp

@@ -90,6 +90,9 @@ struct State {
   /// The list of passed in consuming uses.
   ArrayRef<Operand *> consumingUses;
 
+  /// extend_lifetime uses.
+  ArrayRef<Operand *> extendLifetimeUses;
+
   /// The list of passed in non consuming uses.
   ArrayRef<Operand *> nonConsumingUses;
 
@@ -117,32 +120,43 @@ struct State {
         std::optional<function_ref<void(SILBasicBlock *)>> leakingBlockCallback,
         std::optional<function_ref<void(Operand *)>>
             nonConsumingUseOutsideLifetimeCallback,
-        ArrayRef<Operand *> consumingUses, ArrayRef<Operand *> nonConsumingUses)
+        ArrayRef<Operand *> consumingUses,
+        ArrayRef<Operand *> extendLifetimeUses,
+        ArrayRef<Operand *> nonConsumingUses)
       : value(value), beginInst(value->getDefiningInsertionPoint()),
         errorBuilder(errorBuilder), visitedBlocks(value->getFunction()),
         leakingBlockCallback(leakingBlockCallback),
         nonConsumingUseOutsideLifetimeCallback(
             nonConsumingUseOutsideLifetimeCallback),
-        consumingUses(consumingUses), nonConsumingUses(nonConsumingUses),
+        consumingUses(consumingUses), extendLifetimeUses(extendLifetimeUses),
+        nonConsumingUses(nonConsumingUses),
         blocksWithConsumingUses(value->getFunction()) {}
 
   State(SILBasicBlock *beginBlock,
         LinearLifetimeChecker::ErrorBuilder &errorBuilder,
         std::optional<function_ref<void(SILBasicBlock *)>> leakingBlockCallback,
         std::optional<function_ref<void(Operand *)>>
             nonConsumingUseOutsideLifetimeCallback,
-        ArrayRef<Operand *> consumingUses, ArrayRef<Operand *> nonConsumingUses)
+        ArrayRef<Operand *> consumingUses,
+        ArrayRef<Operand *> linearLifetimeEndingUses,
+        ArrayRef<Operand *> nonConsumingUses)
       : value(), beginInst(&*beginBlock->begin()), errorBuilder(errorBuilder),
         visitedBlocks(beginBlock->getParent()),
         leakingBlockCallback(leakingBlockCallback),
         nonConsumingUseOutsideLifetimeCallback(
             nonConsumingUseOutsideLifetimeCallback),
-        consumingUses(consumingUses), nonConsumingUses(nonConsumingUses),
+        consumingUses(consumingUses),
+        extendLifetimeUses(linearLifetimeEndingUses),
+        nonConsumingUses(nonConsumingUses),
         blocksWithConsumingUses(beginBlock->getParent()) {}
 
   SILBasicBlock *getBeginBlock() const { return beginInst->getParent(); }
 
   void initializeAllNonConsumingUses(ArrayRef<Operand *> nonConsumingUsers);
+  void initializeAllExtendLifetimeUses(
+      ArrayRef<Operand *> linearLifetimeEndingUses,
+      SmallVectorImpl<std::pair<Operand *, SILBasicBlock *>>
+          &predsToAddToWorklist);
   void initializeAllConsumingUses(
       ArrayRef<Operand *> consumingUsers,
       SmallVectorImpl<std::pair<Operand *, SILBasicBlock *>>
@@ -245,6 +259,53 @@ void State::initializeAllNonConsumingUses(
 //                        Consuming Use Initialization
 //===----------------------------------------------------------------------===//
 
+void State::initializeAllExtendLifetimeUses(
+    ArrayRef<Operand *> linearLifetimeEndingUses,
+    SmallVectorImpl<std::pair<Operand *, SILBasicBlock *>>
+        &predsToAddToWorklist) {
+  // TODO: Efficiency.
+  for (auto *use : linearLifetimeEndingUses) {
+    auto *user = use->getUser();
+    auto *block = user->getParent();
+
+    auto iter = blocksWithNonConsumingUses.find(block);
+    if (!iter.has_value()) {
+      continue;
+    }
+
+    // Remove the block from `blocksWithNonConsumingUses` if all non-consuming
+    // uses within the block occur before `user`.
+    bool allBefore = true;
+    for (auto *nonConsumingUse : *iter) {
+      auto *nonConsumingUser = nonConsumingUse->getUser();
+      assert(nonConsumingUser != user);
+
+      bool sawNonConsumingUser = false;
+      for (auto *inst = user->getNextInstruction(); inst;
+           inst = inst->getNextInstruction()) {
+        if (inst == nonConsumingUser) {
+          sawNonConsumingUser = true;
+          break;
+        }
+      }
+      if (sawNonConsumingUser) {
+        allBefore = false;
+        break;
+      }
+    }
+    if (allBefore) {
+      blocksWithNonConsumingUses.erase(block);
+    }
+
+    // Add relevant predecessors to the worklist.
+    if (block == getBeginBlock())
+      continue;
+    for (auto *predecessor : block->getPredecessorBlocks()) {
+      predsToAddToWorklist.push_back({use, predecessor});
+    }
+  }
+}
+
 void State::initializeAllConsumingUses(
     ArrayRef<Operand *> consumingUses,
     SmallVectorImpl<std::pair<Operand *, SILBasicBlock *>>
@@ -555,7 +616,7 @@ void State::checkDataflowEndState(DeadEndBlocks *deBlocks) {
 
 LinearLifetimeChecker::Error LinearLifetimeChecker::checkValueImpl(
     SILValue value, ArrayRef<Operand *> consumingUses,
-    ArrayRef<Operand *> nonConsumingUses, ErrorBuilder &errorBuilder,
+    ArrayRef<Operand *> regularUses, ErrorBuilder &errorBuilder,
     std::optional<function_ref<void(SILBasicBlock *)>> leakingBlockCallback,
     std::optional<function_ref<void(Operand *)>>
         nonConsumingUseOutsideLifetimeCallback) {
@@ -571,9 +632,19 @@ LinearLifetimeChecker::Error LinearLifetimeChecker::checkValueImpl(
   //        || deadEndBlocks.isDeadEnd(value->getParentBlock())) &&
   //       "Must have at least one consuming user?!");
 
+  SmallVector<Operand *, 32> nonConsumingUses;
+  SmallVector<Operand *, 32> extendLifetimeUses;
+  for (auto *use : regularUses) {
+    if (isa<ExtendLifetimeInst>(use->getUser())) {
+      extendLifetimeUses.push_back(use);
+    } else {
+      nonConsumingUses.push_back(use);
+    }
+  }
+
   State state(value, errorBuilder, leakingBlockCallback,
               nonConsumingUseOutsideLifetimeCallback, consumingUses,
-              nonConsumingUses);
+              extendLifetimeUses, nonConsumingUses);
 
   // First add our non-consuming uses and their blocks to the
   // blocksWithNonConsumingUses map. While we do this, if we have multiple uses
@@ -595,6 +666,9 @@ LinearLifetimeChecker::Error LinearLifetimeChecker::checkValueImpl(
   SmallVector<std::pair<Operand *, SILBasicBlock *>, 32> predsToAddToWorklist;
   state.initializeAllConsumingUses(consumingUses, predsToAddToWorklist);
 
+  state.initializeAllExtendLifetimeUses(extendLifetimeUses,
+                                        predsToAddToWorklist);
+
   // If we have a singular consuming use and it is in the same block as value's
   // def, we bail early. Any use-after-frees due to non-consuming uses would
   // have been detected by initializing our consuming uses. So we are done.
@@ -633,6 +707,10 @@ LinearLifetimeChecker::Error LinearLifetimeChecker::checkValueImpl(
     state.visitedBlocks.insert(i->getUser()->getParent());
   }
 
+  for (const auto &i : extendLifetimeUses) {
+    state.visitedBlocks.insert(i->getUser()->getParent());
+  }
+
   // Now that we have marked all of our producing blocks, we go through our
   // predsToAddToWorklist list and add our preds, making sure that none of these
   // preds are in blocksWithConsumingUses. This is important so that we do not

lib/SIL/Verifier/SILOwnershipVerifier.cpp

@@ -99,6 +99,12 @@ class SILValueOwnershipChecker {
   /// successful.
   SmallVector<Operand *, 16> lifetimeEndingUsers;
 
+  /// extend_lifetime users of `value`.
+  ///
+  /// Collected separately because such users do "end the lifetime" but are not
+  /// consuming users.
+  SmallVector<Operand *, 16> extendLifetimeUses;
+
   /// The list of non lifetime ending users that we found. Only valid if check
   /// is successful.
   SmallVector<Operand *, 16> regularUsers;
@@ -129,16 +135,26 @@ class SILValueOwnershipChecker {
   bool isCompatibleDefUse(Operand *op, ValueOwnershipKind ownershipKind);
 
   bool gatherUsers(SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+                   SmallVectorImpl<Operand *> &extendLifetimeUses,
                    SmallVectorImpl<Operand *> &regularUsers);
 
   bool gatherNonGuaranteedUsers(SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+                                SmallVectorImpl<Operand *> &extendLifetimeUses,
                                 SmallVectorImpl<Operand *> &regularUsers);
 
-  bool checkValueWithoutLifetimeEndingUses(ArrayRef<Operand *> regularUsers);
+  bool
+  checkValueWithoutLifetimeEndingUses(ArrayRef<Operand *> regularUsers,
+                                      ArrayRef<Operand *> extendLifetimeUses);
 
-  bool checkFunctionArgWithoutLifetimeEndingUses(SILFunctionArgument *arg);
-  bool checkYieldWithoutLifetimeEndingUses(MultipleValueInstructionResult *yield,
-                                           ArrayRef<Operand *> regularUsers);
+  bool checkFunctionArgWithoutLifetimeEndingUses(
+      SILFunctionArgument *arg, ArrayRef<Operand *> regularUsers,
+      ArrayRef<Operand *> extendLifetimeUses);
+  bool
+  checkYieldWithoutLifetimeEndingUses(MultipleValueInstructionResult *yield,
+                                      ArrayRef<Operand *> regularUsers,
+                                      ArrayRef<Operand *> extendLifetimeUses);
+  bool checkDeadEnds(SILBasicBlock *block, ArrayRef<Operand *> regularUsers,
+                     ArrayRef<Operand *> extendLifetimeUses);
 
   bool isGuaranteedFunctionArgWithLifetimeEndingUses(
       SILFunctionArgument *arg,
@@ -166,6 +182,7 @@ bool SILValueOwnershipChecker::check() {
   llvm::copy(lifetimeEndingUsers, std::back_inserter(allLifetimeEndingUsers));
   SmallVector<Operand *, 32> allRegularUsers;
   llvm::copy(regularUsers, std::back_inserter(allRegularUsers));
+  llvm::copy(extendLifetimeUses, std::back_inserter(allRegularUsers));
 
   LinearLifetimeChecker checker(deadEndBlocks);
   auto linearLifetimeResult = checker.checkValue(value, allLifetimeEndingUsers,
@@ -199,6 +216,7 @@ bool SILValueOwnershipChecker::isCompatibleDefUse(
 
 bool SILValueOwnershipChecker::gatherNonGuaranteedUsers(
     SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+    SmallVectorImpl<Operand *> &extendLifetimeUses,
     SmallVectorImpl<Operand *> &nonLifetimeEndingUsers) {
   bool foundError = false;
 
@@ -209,6 +227,11 @@ bool SILValueOwnershipChecker::gatherNonGuaranteedUsers(
   for (auto *op : value->getUses()) {
     auto *user = op->getUser();
 
+    if (isa<ExtendLifetimeInst>(user)) {
+      extendLifetimeUses.push_back(op);
+      continue;
+    }
+
     // For example, type dependent operands are non-use. It is not interesting
     // from an ownership perspective.
     if (op->getOperandOwnership() == OperandOwnership::NonUse)
@@ -268,14 +291,15 @@ bool SILValueOwnershipChecker::gatherNonGuaranteedUsers(
 
 bool SILValueOwnershipChecker::gatherUsers(
     SmallVectorImpl<Operand *> &lifetimeEndingUsers,
+    SmallVectorImpl<Operand *> &extendLifetimeUses,
     SmallVectorImpl<Operand *> &nonLifetimeEndingUsers) {
 
   // See if Value is guaranteed. If we are guaranteed and not forwarding, then
   // we need to look through subobject uses for more uses. Otherwise, if we are
   // forwarding, we do not create any lifetime ending users/non lifetime ending
   // users since we verify against our base.
   if (value->getOwnershipKind() != OwnershipKind::Guaranteed) {
-    return !gatherNonGuaranteedUsers(lifetimeEndingUsers,
+    return !gatherNonGuaranteedUsers(lifetimeEndingUsers, extendLifetimeUses,
                                      nonLifetimeEndingUsers);
   }
 
@@ -474,8 +498,39 @@ bool SILValueOwnershipChecker::gatherUsers(
   return !foundError;
 }
 
+/// Precondition: value has no lifetimeEndingUsers
+bool SILValueOwnershipChecker::checkDeadEnds(
+    SILBasicBlock *block, ArrayRef<Operand *> regularUses,
+    ArrayRef<Operand *> extendedLifetimeUses) {
+  if (deadEndBlocks && deadEndBlocks->isDeadEnd(block))
+    return true;
+
+  if (extendLifetimeUses.size() == 0)
+    return false;
+
+  SSAPrunedLiveness liveness(value->getFunction());
+  liveness.initializeDef(value);
+  for (auto *use : extendedLifetimeUses) {
+    liveness.updateForUse(use->getUser(), /*lifetimeEnding=*/true);
+  }
+  auto allWithinBoundary = true;
+  for (auto *use : regularUses) {
+    if (!liveness.isWithinBoundary(use->getUser())) {
+      allWithinBoundary |= errorBuilder.handleMalformedSIL([&] {
+        llvm::errs()
+            << "Owned value without lifetime ending uses whose regular use "
+               "isn't enclosed within extend_lifetime instructions:\n"
+            << "Value: " << value << '\n'
+            << "User: " << use->getUser() << '\n';
+      });
+    }
+  }
+  return allWithinBoundary;
+}
+
 bool SILValueOwnershipChecker::checkFunctionArgWithoutLifetimeEndingUses(
-    SILFunctionArgument *arg) {
+    SILFunctionArgument *arg, ArrayRef<Operand *> regularUses,
+    ArrayRef<Operand *> extendLifetimeUses) {
   switch (arg->getOwnershipKind()) {
   case OwnershipKind::Any:
     llvm_unreachable("Value can not have any ownership kind?!");
@@ -487,7 +542,7 @@ bool SILValueOwnershipChecker::checkFunctionArgWithoutLifetimeEndingUses(
     break;
   }
 
-  if (deadEndBlocks && deadEndBlocks->isDeadEnd(arg->getParent()))
+  if (checkDeadEnds(arg->getParent(), regularUses, extendLifetimeUses))
     return true;
 
   return !errorBuilder.handleMalformedSIL([&] {
@@ -497,16 +552,17 @@ bool SILValueOwnershipChecker::checkFunctionArgWithoutLifetimeEndingUses(
 }
 
 bool SILValueOwnershipChecker::checkYieldWithoutLifetimeEndingUses(
-    MultipleValueInstructionResult *yield, ArrayRef<Operand *> regularUses) {
+    MultipleValueInstructionResult *yield, ArrayRef<Operand *> regularUses,
+    ArrayRef<Operand *> extendLifetimeUses) {
   switch (yield->getOwnershipKind()) {
   case OwnershipKind::Any:
     llvm_unreachable("value with any ownership kind?!");
   case OwnershipKind::Unowned:
   case OwnershipKind::None:
     return true;
   case OwnershipKind::Owned:
-    if (deadEndBlocks
-        && deadEndBlocks->isDeadEnd(yield->getParent()->getParent())) {
+    if (checkDeadEnds(yield->getParent()->getParent(), regularUses,
+                      extendLifetimeUses)) {
       return true;
     }
     return !errorBuilder.handleMalformedSIL([&] {
@@ -549,16 +605,21 @@ bool SILValueOwnershipChecker::checkYieldWithoutLifetimeEndingUses(
 }
 
 bool SILValueOwnershipChecker::checkValueWithoutLifetimeEndingUses(
-    ArrayRef<Operand *> regularUses) {
+    ArrayRef<Operand *> regularUses, ArrayRef<Operand *> extendLifetimeUses) {
+  if (extendLifetimeUses.size()) {
+  }
+
   LLVM_DEBUG(llvm::dbgs() << "No lifetime ending users?! Bailing early.\n");
   if (auto *arg = dyn_cast<SILFunctionArgument>(value)) {
-    if (checkFunctionArgWithoutLifetimeEndingUses(arg)) {
+    if (checkFunctionArgWithoutLifetimeEndingUses(arg, regularUses,
+                                                  extendLifetimeUses)) {
       return true;
     }
   }
 
   if (auto *yield = isaResultOf<BeginApplyInst>(value)) {
-    return checkYieldWithoutLifetimeEndingUses(yield, regularUses);
+    return checkYieldWithoutLifetimeEndingUses(yield, regularUses,
+                                               extendLifetimeUses);
   }
 
   // Check if we are a guaranteed subobject. In such a case, we should never
@@ -575,7 +636,7 @@ bool SILValueOwnershipChecker::checkValueWithoutLifetimeEndingUses(
     return true;
 
   if (auto *parentBlock = value->getParentBlock()) {
-    if (deadEndBlocks && deadEndBlocks->isDeadEnd(parentBlock)) {
+    if (checkDeadEnds(parentBlock, regularUses, extendLifetimeUses)) {
       LLVM_DEBUG(llvm::dbgs() << "Ignoring transitively unreachable value "
                               << "without users!\n"
                               << "    Value: " << *value << '\n');
@@ -673,7 +734,7 @@ bool SILValueOwnershipChecker::checkUses() {
   // 1. Verify that none of the uses are in the same block. This would be an
   // overconsume so in this case we assert.
   // 2. Verify that the uses are compatible with our ownership convention.
-  if (!gatherUsers(lifetimeEndingUsers, regularUsers)) {
+  if (!gatherUsers(lifetimeEndingUsers, extendLifetimeUses, regularUsers)) {
     // Silently return false if this fails.
     //
     // If the user pass in a ErrorBehaviorKind that will assert, we
@@ -701,7 +762,7 @@ bool SILValueOwnershipChecker::checkUses() {
   // In the case of a yielded guaranteed value, we need to validate that all
   // regular uses of the value are within the coroutine.
   if (lifetimeEndingUsers.empty()) {
-    if (checkValueWithoutLifetimeEndingUses(regularUsers))
+    if (checkValueWithoutLifetimeEndingUses(regularUsers, extendLifetimeUses))
       return false;
     return true;
   }