emit error on implicit destruction of self in discard context

kavon · kavon · commit 6dfc9c54225b · 2023-06-05T19:39:07.000-07:00
As part of SE-390, you're required to write either: - `consume self` - pass self as a `consuming` parameter to a function - `discard self` before the function ends in a context that contains a `discard self` somewhere. This prevents people from accidentally invoking the deinit due to implicit destruction of `self` before exiting the function. rdar://106099027 (cherry picked from commit 88d35a0)
diff --git a/include/swift/AST/DiagnosticsSIL.def b/include/swift/AST/DiagnosticsSIL.def
@@ -776,7 +776,11 @@ ERROR(sil_movechecking_notconsumable_but_assignable_was_consumed, none,
 ERROR(sil_movechecking_cannot_destructure_has_deinit, none,
       "cannot partially consume '%0' when it has a deinitializer",
       (StringRef))
+ERROR(sil_movechecking_discard_missing_consume_self, none,
+      "must consume 'self' before exiting method that discards self", ())
 
+NOTE(sil_movechecking_discard_self_here, none,
+     "discarded self here", ())
 NOTE(sil_movechecking_partial_consume_here, none,
      "partially consumed here", ())
 NOTE(sil_movechecking_consuming_use_here, none,
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp
@@ -691,6 +691,9 @@ struct UseState {
   /// [assign] that are reinits that we will convert to inits and true reinits.
   llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4> reinitInsts;
 
+  /// The set of drop_deinits of this mark_must_check
+  SmallSetVector<SILInstruction *, 2> dropDeinitInsts;
+
   /// A "inout terminator use" is an implicit liveness use of the entire value
   /// placed on a terminator. We use this both so we add liveness for the
   /// terminator user and so that we can use the set to quickly identify later
@@ -712,6 +715,31 @@ struct UseState {
     return inoutTermUsers.count(inst);
   }
 
+  /// Returns true if the given instruction is within the same block as a reinit
+  /// and precedes a reinit instruction in that block.
+  bool precedesReinitInSameBlock(SILInstruction *inst) const {
+    SILBasicBlock *block = inst->getParent();
+    SmallSetVector<SILInstruction *, 8> sameBlockReinits;
+
+    // First, search for all reinits that are within the same block.
+    for (auto &reinit : reinitInsts) {
+      if (reinit.first->getParent() != block)
+        continue;
+      sameBlockReinits.insert(reinit.first);
+    }
+
+    if (sameBlockReinits.empty())
+      return false;
+
+    // Walk down from the given instruction to see if we encounter a reinit.
+    for (auto ii = std::next(inst->getIterator()); ii != block->end(); ++ii) {
+      if (sameBlockReinits.contains(&*ii))
+        return true;
+    }
+
+    return false;
+  }
+
   void clear() {
     address = nullptr;
     destroys.clear();
@@ -721,6 +749,7 @@ struct UseState {
     takeInsts.clear();
     initInsts.clear();
     reinitInsts.clear();
+    dropDeinitInsts.clear();
     inoutTermUsers.clear();
     debugValue = nullptr;
   }
@@ -755,6 +784,10 @@ struct UseState {
     for (auto pair : reinitInsts) {
       llvm::dbgs() << *pair.first;
     }
+    llvm::dbgs() << "DropDeinits:\n";
+    for (auto *inst : dropDeinitInsts) {
+      llvm::dbgs() << *inst;
+    }
     llvm::dbgs() << "InOut Term Users:\n";
     for (auto *inst : inoutTermUsers) {
       llvm::dbgs() << *inst;
@@ -1737,6 +1770,12 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
     LLVM_DEBUG(llvm::dbgs() << "Running copy propagation!\n");
     moveChecker.changed |= moveChecker.canonicalizer.canonicalize();
 
+    // Export the drop_deinit's discovered by the ObjectChecker into the
+    // AddressChecker to preserve it for later use. We need to do this since
+    // the ObjectChecker's state gets cleared after running on this LoadInst.
+    for (auto *dropDeinit : moveChecker.canonicalizer.getDropDeinitUses())
+      moveChecker.addressUseState.dropDeinitInsts.insert(dropDeinit);
+
     // If we are asked to perform no_consume_or_assign checking or
     // assignable_but_not_consumable checking, if we found any consumes of our
     // load, then we need to emit an error.
@@ -2458,10 +2497,41 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
     FieldSensitiveMultiDefPrunedLiveRange &liveness,
     const FieldSensitivePrunedLivenessBoundary &boundary) {
   LLVM_DEBUG(llvm::dbgs() << "MoveOnlyAddressChecker Rewrite Uses!\n");
-  // First remove all destroy_addr that have not been claimed.
+
+  /// Whether the marked value appeared in a discard statement.
+  const bool isDiscardingContext = !addressUseState.dropDeinitInsts.empty();
+
+  // Process destroys
   for (auto destroyPair : addressUseState.destroys) {
-    if (!consumes.claimConsume(destroyPair.first, destroyPair.second)) {
+    /// Is this destroy instruction a final consuming use?
+    bool isFinalConsume =
+        consumes.claimConsume(destroyPair.first, destroyPair.second);
+
+    // Remove destroys that are not the final consuming use.
+    if (!isFinalConsume) {
       destroyPair.first->eraseFromParent();
+      continue;
+    }
+
+    // Otherwise, if we're in a discarding context, flag this final destroy_addr
+    // as a point where we're missing an explicit `consume self`. The reasoning
+    // here is that if a destroy of self is the final consuming use,
+    // then these are the points where we implicitly destroy self to clean-up
+    // that self var before exiting the scope. An explicit 'consume self'
+    // that is thrown away is a consume of this mark_must_check'd var and not a
+    // destroy of it, according to the use classifier.
+    if (isDiscardingContext) {
+
+      // Since the boundary computations treat a newly-added destroy prior to
+      // a reinit within that same block as a "final consuming use", exclude
+      // such destroys-before-reinit. We are only interested in the final
+      // destroy of a var, not intermediate destroys of the var.
+      if (addressUseState.precedesReinitInSameBlock(destroyPair.first))
+        continue;
+
+      auto *dropDeinit = addressUseState.dropDeinitInsts.front();
+      diagnosticEmitter.emitMissingConsumeInDiscardingContext(destroyPair.first,
+                                                              dropDeinit);
     }
   }
 
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp
@@ -15,6 +15,7 @@
 #include "MoveOnlyDiagnostics.h"
 
 #include "swift/AST/DiagnosticsSIL.h"
+#include "swift/AST/Stmt.h"
 #include "swift/Basic/Defer.h"
 #include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
@@ -190,6 +191,125 @@ void DiagnosticEmitter::emitCheckedMissedCopyError(SILInstruction *copyInst) {
            diag::sil_movechecking_bug_missed_copy);
 }
 
+void DiagnosticEmitter::emitMissingConsumeInDiscardingContext(
+    SILInstruction *leftoverDestroy,
+    SILInstruction *discard) {
+  assert(isa<DropDeinitInst>(discard));
+
+  // A good location is one that has some connection with the original source
+  // and corresponds to an exit of the function.
+  auto hasGoodLocation = [](SILInstruction *si) -> bool {
+    if (!si)
+      return false;
+
+    SILLocation loc = si->getLoc();
+    if (loc.isNull())
+      return false;
+
+    switch (loc.getKind()) {
+    case SILLocation::ReturnKind:
+    case SILLocation::ImplicitReturnKind:
+      return true;
+
+    case SILLocation::RegularKind: {
+      Stmt *stmt = loc.getAsASTNode<Stmt>();
+      if (!stmt)
+        return true; // For non-statements, assume it is exiting the func.
+
+      // Prefer statements that can possibly lead to an exit of the function.
+      // This is determined by whether the statement causes an exit of a
+      // lexical scope; so a 'break' counts but not a 'continue'.
+      switch (stmt->getKind()) {
+      case StmtKind::Throw:
+      case StmtKind::Return:
+      case StmtKind::Yield:
+      case StmtKind::Break:
+      case StmtKind::Fail:
+      case StmtKind::PoundAssert:
+        return true;
+
+      case StmtKind::Continue:
+      case StmtKind::Brace:
+      case StmtKind::Defer:
+      case StmtKind::If:
+      case StmtKind::Guard:
+      case StmtKind::While:
+      case StmtKind::Do:
+      case StmtKind::DoCatch:
+      case StmtKind::RepeatWhile:
+      case StmtKind::ForEach:
+      case StmtKind::Switch:
+      case StmtKind::Case:
+      case StmtKind::Fallthrough:
+      case StmtKind::Discard:
+        return false;
+      };
+    }
+
+    case SILLocation::InlinedKind:
+    case SILLocation::MandatoryInlinedKind:
+    case SILLocation::CleanupKind:
+    case SILLocation::ArtificialUnreachableKind:
+      return false;
+    };
+  };
+
+  // An instruction corresponding to the logical place where the value is
+  // destroyed. Ideally an exit point of the function reachable from here or
+  // some relevant statement.
+  SILInstruction *destroyPoint = leftoverDestroy;
+  if (!hasGoodLocation(destroyPoint)) {
+    // Search for a nearby function exit reachable from this destroy. We do this
+    // because the move checker may have injected or hoisted an existing
+    // destroy from leaf blocks to some earlier point. For example, if 'd'
+    // represents a destroy of self, then we may have this CFG:
+    //
+    //        before:      after:
+    //           .          d
+    //          / \        / \
+    //         d   d      .   .
+    //
+    BasicBlockSet visited(destroyPoint->getFunction());
+    std::deque<SILBasicBlock *> bfsWorklist = {destroyPoint->getParent()};
+    while (auto *bb = bfsWorklist.front()) {
+      visited.insert(bb);
+      bfsWorklist.pop_front();
+
+      TermInst *term = bb->getTerminator();
+
+      // Looking for a block that exits the function or terminates the program.
+      if (term->isFunctionExiting() || term->isProgramTerminating()) {
+        SILInstruction *candidate = term;
+
+        // Walk backwards until we find an instruction with any source location.
+        // Sometimes a terminator like 'unreachable' may not have one, but one
+        // of the preceding instructions will.
+        while (candidate && candidate->getLoc().isNull())
+          candidate = candidate->getPreviousInstruction();
+
+        if (candidate && candidate->getLoc()) {
+          destroyPoint = candidate;
+          break;
+        }
+      }
+
+      for (auto *nextBB : term->getSuccessorBlocks())
+        if (!visited.contains(nextBB))
+          bfsWorklist.push_back(nextBB);
+    }
+  }
+
+  assert(destroyPoint->getLoc() && "missing loc!");
+  assert(discard->getLoc() && "missing loc!");
+
+  diagnose(leftoverDestroy->getFunction()->getASTContext(),
+           destroyPoint,
+           diag::sil_movechecking_discard_missing_consume_self);
+
+  diagnose(discard->getFunction()->getASTContext(), discard,
+            diag::sil_movechecking_discard_self_here);
+}
+
 //===----------------------------------------------------------------------===//
 //                          MARK: Object Diagnostics
 //===----------------------------------------------------------------------===//
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.h b/lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.h
@@ -119,6 +119,12 @@ class DiagnosticEmitter {
   /// way to file a bug.
   void emitCheckedMissedCopyError(SILInstruction *copyInst);
 
+  /// Assuming the given instruction represents the implicit destruction of
+  /// 'self', emits an error saying that you needed to explicitly 'consume self'
+  /// here because you're in a discarding context.
+  void emitMissingConsumeInDiscardingContext(SILInstruction *leftoverDestroy,
+                                             SILInstruction *dropDeinit);
+
   void emitCheckerDoesntUnderstandDiagnostic(MarkMustCheckInst *markedValue);
   void emitObjectGuaranteedDiagnostic(MarkMustCheckInst *markedValue);
   void emitObjectOwnedDiagnostic(MarkMustCheckInst *markedValue);
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.h b/lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.h
@@ -116,6 +116,20 @@ struct OSSACanonicalizer {
              return !isa<PartialApplyInst>(user);
            });
   }
+
+  struct DropDeinitFilter {
+    bool operator()(SILInstruction *inst) const {
+      return isa<DropDeinitInst>(inst);
+    }
+  };
+  using DropDeinitIter =
+      llvm::filter_iterator<SILInstruction *const *, DropDeinitFilter>;
+  using DropDeinitRange = iterator_range<DropDeinitIter>;
+
+  /// Returns a range of final uses of the mark_must_check that are drop_deinit
+  DropDeinitRange getDropDeinitUses() const {
+    return llvm::make_filter_range(consumingBoundaryUsers, DropDeinitFilter());
+  }
 };
 
 /// Search for candidate object mark_must_checks. If we find one that does not
diff --git a/test/SILGen/discard.swift b/test/SILGen/discard.swift
@@ -68,6 +68,8 @@ func invokedDeinit() {}
   consuming func tryDestroy(doDiscard: Bool) throws {
     if doDiscard {
       discard self
+    } else {
+     _ = consume self
     }
     throw E.err
   }
@@ -136,17 +138,19 @@ final class Wallet {
 
   consuming func changeTicket(inWallet wallet: Wallet? = nil) {
     if let existingWallet = wallet {
-      discard self
       self = .within(existingWallet)
+      _ = consume self
+    } else {
+      discard self
     }
   }
-  // As of now, we allow reinitialization after discard. Not sure if this is intended.
+
   // CHECK-LABEL: sil hidden [ossa] @$s4test6TicketO06changeB08inWalletyAA0E0CSg_tF : $@convention(method) (@guaranteed Optional<Wallet>, @owned Ticket) -> () {
   // CHECK:    [[SELF_REF:%.*]] = project_box [[SELF_BOX:%.*]] : ${ var Ticket }, 0
-  // CHECK:    switch_enum {{.*}} : $Optional<Wallet>, case #Optional.some!enumelt: [[HAVE_WALLET_BB:bb.*]], case #Optional.none!enumelt: {{.*}}
+  // CHECK:    switch_enum {{.*}} : $Optional<Wallet>, case #Optional.some!enumelt: {{.*}}, case #Optional.none!enumelt: [[NO_WALLET_BB:bb[0-9]+]]
   //
   // >> now we begin the destruction sequence, which involves pattern matching on self to destroy its innards
-  // CHECK:  [[HAVE_WALLET_BB]]({{%.*}} : @owned $Wallet):
+  // CHECK:  [[NO_WALLET_BB]]
   // CHECK:    [[SELF_ACCESS:%.*]] = begin_access [read] [unknown] {{%.*}} : $*Ticket
   // CHECK:    [[SELF_MMC:%.*]] = mark_must_check [no_consume_or_assign] [[SELF_ACCESS]]
   // CHECK:    [[SELF_COPY:%.*]] = load [copy] [[SELF_MMC]] : $*Ticket
@@ -157,13 +161,7 @@ final class Wallet {
   // CHECK:  [[TICKET_WITHIN]]([[PREV_SELF_WALLET:%.*]] : @owned $Wallet):
   // CHECK:    destroy_value [[PREV_SELF_WALLET]] : $Wallet
   // CHECK:    br [[JOIN_POINT]]
-  // >> from here on we are reinitializing self.
   // CHECK:  [[JOIN_POINT]]:
-  // CHECK:    [[NEW_SELF_VAL:%.*]] = enum $Ticket, #Ticket.within!enumelt, {{.*}} : $Wallet
-  // CHECK:    [[SELF_ACCESS2:%.*]] = begin_access [modify] [unknown] [[SELF_REF]] : $*Ticket
-  // CHECK:    [[SELF_MMC2:%.*]] = mark_must_check [assignable_but_not_consumable] [[SELF_ACCESS2]] : $*Ticket
-  // CHECK:    assign [[NEW_SELF_VAL]] to [[SELF_MMC2]] : $*Ticket
-  // CHECK:    end_access [[SELF_ACCESS2]] : $*Ticket
 
   deinit {
     print("destroying ticket")
diff --git a/test/SILOptimizer/discard_checking.swift b/test/SILOptimizer/discard_checking.swift
