Merge pull request #65911 from gottesmm/fix-closures

[move-only] Fix noncopyable semantics for let closures to match how we process inout captured by closures

Author: gottesmm

include/swift/SILOptimizer/PassManager/Passes.def

@@ -112,9 +112,6 @@ PASS(AccessMarkerElimination, "access-marker-elim",
      "Access Marker Elimination.")
 PASS(AddressLowering, "address-lowering",
      "SIL Address Lowering")
-PASS(EarlyAllocBoxToStack, "early-allocbox-to-stack",
-     "Stack Promotion of Box Objects. Doesn't attempt to promote noncopyable "
-     "types captured by escaping closures")
 PASS(AllocBoxToStack, "allocbox-to-stack",
      "Stack Promotion of Box Objects")
 IRGEN_PASS(AllocStackHoisting, "alloc-stack-hoisting",

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerTester.cpp

@@ -105,20 +105,18 @@ class MoveOnlyAddressCheckerTesterPass : public SILFunctionTransform {
                << '\n');
 
     bool madeChange = false;
-    unsigned diagCount = 0;
     if (moveIntroducersToProcess.empty()) {
       LLVM_DEBUG(llvm::dbgs() << "No move introducers found?!\n");
     } else {
       borrowtodestructure::IntervalMapAllocator allocator;
       MoveOnlyAddressChecker checker{getFunction(), diagnosticEmitter,
                                      allocator, domTree, poa};
       madeChange = checker.check(moveIntroducersToProcess);
-      diagCount = checker.diagnosticEmitter.getDiagnosticCount();
     }
 
     // If we did not emit any diagnostics, emit a diagnostic if we missed any
     // copies.
-    if (!diagCount) {
+    if (!diagnosticEmitter.emittedDiagnostic()) {
       emitCheckerMissedCopyOfNonCopyableTypeErrors(getFunction(),
                                                    diagnosticEmitter);
     }

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -225,6 +225,7 @@
 #include "swift/AST/AccessScope.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSIL.h"
+#include "swift/AST/SemanticAttrs.h"
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/FrozenMultiMap.h"
@@ -492,6 +493,92 @@ static bool isInOutDefThatNeedsEndOfFunctionLiveness(MarkMustCheckInst *markedAd
   return false;
 }
 
+//===----------------------------------------------------------------------===//
+//                       MARK: Partial Apply Utilities
+//===----------------------------------------------------------------------===//
+
+static bool findNonEscapingPartialApplyUses(
+    PartialApplyInst *pai, TypeTreeLeafTypeRange leafRange,
+    llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4>
+        &livenessUses) {
+  StackList<Operand *> worklist(pai->getFunction());
+  for (auto *use : pai->getUses())
+    worklist.push_back(use);
+
+  LLVM_DEBUG(llvm::dbgs() << "Searching for partial apply uses!\n");
+  while (!worklist.empty()) {
+    auto *use = worklist.pop_back_val();
+
+    if (use->isTypeDependent())
+      continue;
+
+    auto *user = use->getUser();
+
+    // These instructions do not cause us to escape.
+    if (isIncidentalUse(user) || isa<DestroyValueInst>(user))
+      continue;
+
+    // Look through these instructions.
+    if (isa<BeginBorrowInst>(user) || isa<CopyValueInst>(user) ||
+        isa<MoveValueInst>(user) ||
+        // If we capture this partial_apply in another partial_apply, then we
+        // know that said partial_apply must not have escaped the value since
+        // otherwise we could not have an inout_aliasable argument or be
+        // on_stack. Process it recursively so that we treat uses of that
+        // partial_apply and applies of that partial_apply as uses of our
+        // partial_apply.
+        //
+        // We have this separately from the other look through sections so that
+        // we can make it clearer what we are doing here.
+        isa<PartialApplyInst>(user)) {
+      for (auto *use : cast<SingleValueInstruction>(user)->getUses())
+        worklist.push_back(use);
+      continue;
+    }
+
+    // If we have a mark_dependence and are the value, look through the
+    // mark_dependence.
+    if (auto *mdi = dyn_cast<MarkDependenceInst>(user)) {
+      if (mdi->getValue() == use->get()) {
+        for (auto *use : mdi->getUses())
+          worklist.push_back(use);
+        continue;
+      }
+    }
+
+    if (auto apply = FullApplySite::isa(user)) {
+      // If we apply the function or pass the function off to an apply, then we
+      // need to treat the function application as a liveness use of the
+      // variable since if the partial_apply is invoked within the function
+      // application, we may access the captured variable.
+      livenessUses.insert({user, leafRange});
+      if (apply.beginsCoroutineEvaluation()) {
+        // If we have a coroutine, we need to treat the abort_apply and
+        // end_apply as liveness uses since once we execute one of those
+        // instructions, we have returned control to the coroutine which means
+        // that we could then access the captured variable again.
+        auto *bai = cast<BeginApplyInst>(user);
+        SmallVector<EndApplyInst *, 4> endApplies;
+        SmallVector<AbortApplyInst *, 4> abortApplies;
+        bai->getCoroutineEndPoints(endApplies, abortApplies);
+        for (auto *eai : endApplies)
+          livenessUses.insert({eai, leafRange});
+        for (auto *aai : abortApplies)
+          livenessUses.insert({aai, leafRange});
+      }
+      continue;
+    }
+
+    LLVM_DEBUG(
+        llvm::dbgs()
+        << "Found instruction we did not understand... returning false!\n");
+    LLVM_DEBUG(llvm::dbgs() << "Instruction: " << *user);
+    return false;
+  }
+
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 //                   MARK: Find Candidate Mark Must Checks
 //===----------------------------------------------------------------------===//
@@ -1634,10 +1721,15 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
                    << "Found mark must check [nocopy] error: " << *user);
         auto *fArg = dyn_cast<SILFunctionArgument>(
             stripAccessMarkers(markedValue->getOperand()));
-        if (fArg && fArg->isClosureCapture() && fArg->getType().isAddress()) {
-          moveChecker.diagnosticEmitter.emitPromotedBoxArgumentError(
-              markedValue, fArg);
+        // If we have a closure captured that we specialized, we should have a
+        // no consume or assign and should emit a normal guaranteed diagnostic.
+        if (fArg && fArg->isClosureCapture() &&
+            fArg->getArgumentConvention().isInoutConvention()) {
+          assert(checkKind == MarkMustCheckInst::CheckKind::NoConsumeOrAssign);
+          moveChecker.diagnosticEmitter.emitObjectGuaranteedDiagnostic(
+              markedValue);
         } else {
+          // Otherwise, we need to emit an escaping closure error.
           moveChecker.diagnosticEmitter
               .emitAddressEscapingClosureCaptureLoadedAndConsumed(markedValue);
         }
@@ -1792,6 +1884,19 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
   }
 
   if (auto *pas = dyn_cast<PartialApplyInst>(user)) {
+    if (auto *fArg = dyn_cast<SILFunctionArgument>(
+            stripAccessMarkers(markedValue->getOperand()))) {
+      // If we are processing an inout convention and we emitted an error on the
+      // partial_apply, we shouldn't process this mark_must_check, but squelch
+      // the compiler doesn't understand error.
+      if (fArg->getArgumentConvention().isInoutConvention() &&
+          pas->getCalleeFunction()->hasSemanticsAttr(
+              semantics::NO_MOVEONLY_DIAGNOSTICS)) {
+        diagnosticEmitter.emitEarlierPassEmittedDiagnostic(markedValue);
+        return false;
+      }
+    }
+
     if (pas->isOnStack() ||
         ApplySite(pas).getArgumentConvention(*op).isInoutConvention()) {
       LLVM_DEBUG(llvm::dbgs() << "Found on stack partial apply or inout usage!\n");
@@ -1803,13 +1908,18 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
         return false;
       }
 
-      useState.livenessUses.insert({user, *leafRange});
-      for (auto *use : pas->getConsumingUses()) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "Adding consuming use of partial apply as liveness use: "
-                   << *use->getUser());
-        useState.livenessUses.insert({use->getUser(), *leafRange});
+      // Attempt to find calls of the non-escaping partial apply and places
+      // where the partial apply is passed to a function. We treat those as
+      // liveness uses. If we find a use we don't understand, we return false
+      // here.
+      if (!findNonEscapingPartialApplyUses(pas, *leafRange,
+                                           useState.livenessUses)) {
+        LLVM_DEBUG(
+            llvm::dbgs()
+            << "Failed to understand use of a non-escaping partial apply?!\n");
+        return false;
       }
+
       return true;
     }
   }
@@ -2534,9 +2644,19 @@ bool MoveOnlyAddressChecker::check(
     LLVM_DEBUG(llvm::dbgs() << "Visiting: " << *markedValue);
 
     // Perform our address check.
+    unsigned diagnosticEmittedByEarlierPassCount =
+      diagnosticEmitter.getDiagnosticEmittedByEarlierPassCount();
     if (!pimpl.performSingleCheck(markedValue)) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "Failed to perform single check! Emitting error!\n");
+      if (diagnosticEmittedByEarlierPassCount !=
+          diagnosticEmitter.getDiagnosticEmittedByEarlierPassCount()) {
+        LLVM_DEBUG(
+            llvm::dbgs()
+            << "Failed to perform single check but found earlier emitted "
+               "error. Not emitting checker doesn't understand diagnostic!\n");
+        continue;
+      }
+      LLVM_DEBUG(llvm::dbgs() << "Failed to perform single check! Emitting "
+                                 "compiler doesn't understand diagnostic!\n");
       // If we fail the address check in some way, set the diagnose!
       diagnosticEmitter.emitCheckerDoesntUnderstandDiagnostic(markedValue);
     }

lib/SILOptimizer/Mandatory/MoveOnlyChecker.cpp

@@ -179,7 +179,7 @@ class MoveOnlyCheckerPass : public SILFunctionTransform {
     // remain. If we emitted a diagnostic, we just want to rewrite all of the
     // non-copyable copies into explicit variants below and let the user
     // recompile.
-    if (!checker.diagnosticEmitter.getDiagnosticCount()) {
+    if (!checker.diagnosticEmitter.emittedDiagnostic()) {
       emitCheckerMissedCopyOfNonCopyableTypeErrors(getFunction(),
                                                    checker.diagnosticEmitter);
     }

lib/SILOptimizer/Mandatory/MoveOnlyDiagnostics.h

@@ -52,6 +52,23 @@ class DiagnosticEmitter {
 
   bool emittedCheckerDoesntUnderstandDiagnostic = false;
 
+  /// This is incremented every time that the checker determines that an earlier
+  /// pass emitted a diagnostic while processing a mark_must_check. In such a
+  /// case, we want to suppress:
+  ///
+  /// 1. Emitting the compiler doesn't understand how to check error for the
+  ///    specific mark_must_check.
+  ///
+  /// 2. The "copy of noncopyable type" error over the entire function since us
+  ///    stopping processing at some point may have left copies.
+  ///
+  /// We use a counter rather than a boolean here so that a caller that is
+  /// processing an individual mark_must_check can determine if the checker
+  /// identified such an earlier pass diagnostic for the specific allocation so
+  /// that we can still emit "compiler doesn't understand" errors for other
+  /// allocations.
+  unsigned diagnosticEmittedByEarlierPassCount = 0;
+
 public:
   DiagnosticEmitter(SILFunction *inputFn) : fn(inputFn) {}
 
@@ -67,11 +84,37 @@ class DiagnosticEmitter {
     return *canonicalizer.get();
   }
 
+  /// Returns true if when processing any allocation in the current function:
+  ///
+  /// 1. This diagnostic emitter emitted a diagnostic.
+  /// 2. The user of the diagnostic emitter signaled to the diagnostic emitter
+  ///    that it detected an earlier diagnostic was emitted that prevented it
+  ///    from performing checking.
+  ///
+  /// DISCUSSION: This is used by the checker to decide whether or not it should
+  /// emit "found copy of a noncopyable type" error. If the checker emitted one
+  /// of these diagnostics, then the checker may have stopped processing early
+  /// and left copies since it was no longer able to check. In such a case, we
+  /// want the user to fix the pre-existing errors and re-run.
+  bool emittedDiagnostic() const {
+    return getDiagnosticCount() || getDiagnosticEmittedByEarlierPassCount();
+  }
+
   unsigned getDiagnosticCount() const { return diagnosticCount; }
+
   bool didEmitCheckerDoesntUnderstandDiagnostic() const {
     return emittedCheckerDoesntUnderstandDiagnostic;
   }
 
+  bool getDiagnosticEmittedByEarlierPassCount() const {
+    return diagnosticEmittedByEarlierPassCount;
+  }
+
+  void emitEarlierPassEmittedDiagnostic(MarkMustCheckInst *mmci) {
+    ++diagnosticEmittedByEarlierPassCount;
+    registerDiagnosticEmitted(mmci);
+  }
+
   /// Used at the end of the MoveOnlyAddressChecker to tell the user in a nice
   /// way to file a bug.
   void emitCheckedMissedCopyError(SILInstruction *copyInst);

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -125,7 +125,7 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // This guarantees that stack-promotable boxes have [static] enforcement.
   P.addAccessEnforcementSelection();
 
-  P.addEarlyAllocBoxToStack();
+  P.addAllocBoxToStack();
   P.addNoReturnFolding();
   addDefiniteInitialization(P);
 
@@ -169,11 +169,6 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // End Ownership Optimizations
   //===---
 
-  // Now that we have finished checking noncopyable types, run later alloc box
-  // to stack, so that we promote to the stack any heap boxes that are captured
-  // by escaping closures where the closure does not actually escape.
-  P.addAllocBoxToStack();
-
 #ifndef NDEBUG
   // Add a verification pass to check our work when skipping
   // function bodies.