[wip] emit error on implicit destruction of self in discard context

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

Author: kavon

include/swift/AST/DiagnosticsSIL.def

@@ -777,7 +777,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,

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
@@ -721,6 +724,7 @@ struct UseState {
     takeInsts.clear();
     initInsts.clear();
     reinitInsts.clear();
+    dropDeinitInsts.clear();
     inoutTermUsers.clear();
     debugValue = nullptr;
   }
@@ -755,6 +759,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 +1745,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 +2472,26 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
     FieldSensitiveMultiDefPrunedLiveRange &liveness,
     const FieldSensitivePrunedLivenessBoundary &boundary) {
   LLVM_DEBUG(llvm::dbgs() << "MoveOnlyAddressChecker Rewrite Uses!\n");
+
+  /// whether the marked value appeared in a discard statement.
+  const bool isDiscardingContext = !addressUseState.dropDeinitInsts.empty();
+
   // First remove all destroy_addr that have not been claimed.
   for (auto destroyPair : addressUseState.destroys) {
     if (!consumes.claimConsume(destroyPair.first, destroyPair.second)) {
       destroyPair.first->eraseFromParent();
+      continue;
+    }
+
+    // Otherwise, if we're in a discarding context and checking 'self', flag
+    // the claimed destroy_addr as a point where we're missing an explicit
+    // `consume self`. The reasoning here is that if we leave behind one of
+    // these destroys, then the 'self' didn't get explicitly consumed or
+    // discarded along this path.
+    if (isDiscardingContext) {
+      auto *dropDeinit = addressUseState.dropDeinitInsts.front();
+      diagnosticEmitter.emitMissingConsumeInDiscardingContext(destroyPair.first,
+                                                              dropDeinit);
     }
   }
 
@@ -2657,7 +2687,6 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
   liveness.computeBoundary(boundary);
   insertDestroysOnBoundary(markedAddress, liveness, boundary);
   rewriteUses(markedAddress, liveness, boundary);
-
   return true;
 }
 

lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp

@@ -190,6 +190,17 @@ void DiagnosticEmitter::emitCheckedMissedCopyError(SILInstruction *copyInst) {
            diag::sil_movechecking_bug_missed_copy);
 }
 
+void DiagnosticEmitter::emitMissingConsumeInDiscardingContext(
+    SILInstruction *leftoverDestroy,
+    SILInstruction *discard) {
+  assert(isa<DropDeinitInst>(discard));
+
+  diagnose(leftoverDestroy->getFunction()->getASTContext(), leftoverDestroy,
+           diag::sil_movechecking_discard_missing_consume_self);
+  diagnose(discard->getFunction()->getASTContext(), discard,
+            diag::sil_movechecking_discard_self_here);
+}
+
 //===----------------------------------------------------------------------===//
 //                          MARK: Object Diagnostics
 //===----------------------------------------------------------------------===//

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);

lib/SILOptimizer/Mandatory/MoveOnlyObjectCheckerUtils.h

@@ -117,6 +117,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