prevent reinitialization of self after discard

The value `self` is mutable (i.e., var-bound) in
a `consuming` method. Since you're allowed to
reinitialize a var after consuming, that means
you were also naturally allowed to reinitialize
self after `discard self`. But that capability was
not intended; after you discard self you shouldn't
be reinitializing it, as that's probably a mistake.

This change makes reinitialization of `self`
reachable from a `discard self` statement an error.

rdar://106098163

Author: kavon

include/swift/AST/DiagnosticsSIL.def

@@ -778,6 +778,8 @@ ERROR(sil_movechecking_cannot_destructure_has_deinit, none,
       (StringRef))
 ERROR(sil_movechecking_discard_missing_consume_self, none,
       "must consume 'self' before exiting method that discards self", ())
+ERROR(sil_movechecking_reinit_after_discard, none,
+      "cannot reinitialize 'self' after 'discard self'", ())
 
 NOTE(sil_movechecking_discard_self_here, none,
      "discarded self here", ())

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -1300,6 +1300,10 @@ struct MoveOnlyAddressCheckerPImpl {
                    FieldSensitiveMultiDefPrunedLiveRange &liveness,
                    const FieldSensitivePrunedLivenessBoundary &boundary);
 
+  /// Identifies and diagnoses reinitializations that are reachable from a
+  /// discard statement.
+  void checkForReinitAfterDiscard();
+
   void handleSingleBlockDestroy(SILInstruction *destroy, bool isReinit);
 };
 
@@ -2629,6 +2633,80 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
 #endif
 }
 
+void MoveOnlyAddressCheckerPImpl::checkForReinitAfterDiscard() {
+  auto const &dropDeinits = addressUseState.dropDeinitInsts;
+  auto const &reinits = addressUseState.reinitInsts;
+
+  if (dropDeinits.empty() || reinits.empty())
+    return;
+
+  using BasicBlockMap = llvm::DenseMap<SILBasicBlock *,
+                                       llvm::SmallPtrSet<SILInstruction *, 2>>;
+  BasicBlockMap blocksWithReinit;
+  for (auto const &info : reinits) {
+    auto *reinit = info.first;
+    blocksWithReinit[reinit->getParent()].insert(reinit);
+  }
+
+  // Starting from each drop_deinit instruction, can we reach a reinit of self?
+  for (auto *dropInst : dropDeinits) {
+    auto *dropBB = dropInst->getParent();
+
+    // First, if the block containing this drop_deinit also contains a reinit,
+    // check if that reinit happens after this drop_deinit.
+    auto result = blocksWithReinit.find(dropBB);
+    if (result != blocksWithReinit.end()) {
+      auto &blockReinits = result->second;
+      for (auto ii = std::next(dropInst->getIterator()); ii != dropBB->end();
+           ++ii) {
+        SILInstruction *current = &*ii;
+        if (blockReinits.contains(current)) {
+          // Then the drop_deinit can reach a reinit immediately after it in the
+          // same block.
+          diagnosticEmitter.emitReinitAfterDiscardError(current, dropInst);
+          return;
+        }
+      }
+    }
+
+    BasicBlockWorklist worklist(fn);
+
+    // Seed the search with the successors of the drop_init block, so that if we
+    // visit the drop_deinit block again, we'll know the reinits _before_ the
+    // drop_deinit are reachable via some back-edge / cycle.
+    for (auto *succ : dropBB->getSuccessorBlocks())
+      worklist.pushIfNotVisited(succ);
+
+    // Determine reachability across blocks.
+    while (auto *bb = worklist.pop()) {
+      // Set-up next iteration.
+      for (auto *succ : bb->getSuccessorBlocks())
+        worklist.pushIfNotVisited(succ);
+
+      auto result = blocksWithReinit.find(bb);
+      if (result == blocksWithReinit.end())
+        continue;
+
+      // We found a reachable reinit! Identify the earliest reinit in this block
+      // for diagnosis.
+      auto &blockReinits = result->second;
+      SILInstruction *firstBadReinit = nullptr;
+      for (auto &inst : *bb) {
+        if (blockReinits.contains(&inst)) {
+          firstBadReinit = &inst;
+          break;
+        }
+      }
+
+      if (!firstBadReinit)
+        llvm_unreachable("bug");
+
+      diagnosticEmitter.emitReinitAfterDiscardError(firstBadReinit, dropInst);
+      return;
+    }
+  }
+}
+
 bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
     MarkMustCheckInst *markedAddress) {
   SWIFT_DEFER { diagnosticEmitter.clearUsesWithDiagnostic(); };
@@ -2724,6 +2802,7 @@ bool MoveOnlyAddressCheckerPImpl::performSingleCheck(
   FieldSensitivePrunedLivenessBoundary boundary(liveness.getNumSubElements());
   liveness.computeBoundary(boundary);
   insertDestroysOnBoundary(markedAddress, liveness, boundary);
+  checkForReinitAfterDiscard();
   rewriteUses(markedAddress, liveness, boundary);
 
   return true;

lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.cpp

@@ -191,6 +191,20 @@ void DiagnosticEmitter::emitCheckedMissedCopyError(SILInstruction *copyInst) {
            diag::sil_movechecking_bug_missed_copy);
 }
 
+void DiagnosticEmitter::emitReinitAfterDiscardError(SILInstruction *badReinit,
+                                                    SILInstruction *discard) {
+  assert(isa<DropDeinitInst>(discard));
+  assert(badReinit->getLoc() && "missing loc!");
+  assert(discard->getLoc() && "missing loc!");
+
+  diagnose(badReinit->getFunction()->getASTContext(),
+           badReinit,
+           diag::sil_movechecking_reinit_after_discard);
+
+  diagnose(discard->getFunction()->getASTContext(), discard,
+           diag::sil_movechecking_discard_self_here);
+}
+
 void DiagnosticEmitter::emitMissingConsumeInDiscardingContext(
     SILInstruction *leftoverDestroy,
     SILInstruction *discard) {

lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.h

@@ -119,6 +119,11 @@ class DiagnosticEmitter {
   /// way to file a bug.
   void emitCheckedMissedCopyError(SILInstruction *copyInst);
 
+  /// Given a drop_deinit of self and an instruction reinitializing self,
+  /// emits an error saying that you cannot reinitialize self after a discard.
+  void emitReinitAfterDiscardError(SILInstruction *badReinit,
+                                   SILInstruction *dropDeinit);
+
   /// 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.

test/SILOptimizer/discard_checking.swift

@@ -10,6 +10,7 @@ enum E: Error {
 }
 
 func globalConsumingFn(_ b: consuming Basics) {}
+func globalThrowingFn() throws {}
 
 struct Basics: ~Copyable {
   consuming func test1(_ b: Bool) {
@@ -441,6 +442,110 @@ struct Basics: ~Copyable {
     return
   }
 
+
+  ////////////////////////////////////////
+  /// Checking for reinit-after-discard
+  ////////////////////////////////////////
+  consuming func reinitAfterDiscard1(_ c: Color) throws {
+    discard self // expected-note {{discarded self here}}
+    self = Basics() // expected-error {{cannot reinitialize 'self' after 'discard self'}}
+    _ = consume self
+  }
+
+  consuming func reinitAfterDiscard2_bad(_ c: Color) throws {
+    discard self // expected-note {{discarded self here}}
+
+    if case .red = c {
+      throw E.someError
+    }
+
+    self = Basics() // expected-error {{cannot reinitialize 'self' after 'discard self'}}
+    _ = consume self
+  }
+
+  consuming func reinitAfterDiscard2_ok(_ c: Color) throws {
+    if case .red = c {
+      discard self
+      throw E.someError
+    }
+
+    self = Basics()
+    _ = consume self
+  }
+
+  consuming func reinitAfterDiscard3_bad(_ c: Color) throws {
+    // expected-error@-1 {{must consume 'self' before exiting method that discards self}}
+    // FIXME: ^ this error is related to rdar://110239087
+
+    repeat {
+      self = Basics() // expected-error {{cannot reinitialize 'self' after 'discard self'}}
+      discard self // expected-note 2{{discarded self here}}
+    } while false
+  }
+
+  consuming func reinitAfterDiscard3_ok(_ c: Color) throws {
+    self = Basics()
+    discard self
+  }
+
+  consuming func reinitAfterDiscard4_bad(_ c: Color) throws {
+    do {
+      discard self // expected-note {{discarded self here}}
+      if case .red = c {
+        try globalThrowingFn()
+      }
+    } catch {
+      self = Basics() // expected-error {{cannot reinitialize 'self' after 'discard self'}}
+      _ = consume self
+    }
+  }
+
+  consuming func reinitAfterDiscard4_ok(_ c: Color) throws {
+    do {
+      if case .red = c {
+        try globalThrowingFn()
+      }
+      discard self
+    } catch {
+      self = Basics()
+      _ = consume self
+    }
+  }
+
+  consuming func reinitAfterDiscard4_ok_v2(_ c: Color) throws {
+    do {
+      if case .red = c {
+        try? globalThrowingFn()
+      }
+      discard self
+    } catch { // expected-warning {{'catch' block is unreachable because no errors are thrown in 'do' block}}
+      self = Basics()
+      _ = consume self
+    }
+  }
+
+  // FIXME move checker is treating the defer like a closure capture (rdar://100468597)
+  // There is a different expected error here, where this reinit in a defer
+  // should be illegal because it's a reinit after discard.
+  consuming func reinitAfterDiscard5(_ c: Color) throws {
+    // expected-error@-1 {{'self' used after consume}}
+
+    defer { self = Basics() } // expected-note {{used here}}
+    discard self // expected-note {{consumed here}}
+  }
+
+  consuming func reinitAfterDiscard6(_ c: Color) throws {
+    self = Basics()
+    discard self // expected-note 3{{discarded self here}}
+    switch c {
+    case .red, .blue:
+      self = Basics()  // expected-error {{must consume 'self' before exiting method that discards self}}
+    default:
+      self = Basics()  // expected-error {{must consume 'self' before exiting method that discards self}}
+                       // expected-error@-1 {{cannot reinitialize 'self' after 'discard self'}}
+    }
+  }
+
   mutating func mutator() { self = Basics() }
   borrowing func borrower() {}
   borrowing func thrower() throws {}