Merge pull request #35774 from gottesmm/pr-aed4f36ace8254f8ea9b635f05a8c903fd964553

[capture-promotion] Emit a warning when the pass fails to promote a capture of a concurrent function to a by value capture instead of by ref capture.

Author: gottesmm

include/swift/AST/DiagnosticsSIL.def

@@ -645,5 +645,15 @@ NOTE(box_to_stack_cannot_promote_box_to_stack_due_to_escape_location, none,
 
 WARNING(semantic_function_improper_nesting, none, "'@_semantics' function calls non-'@_semantics' function with nested '@_semantics' calls", ())
 
+// Capture promotion diagnostics
+WARNING(capturepromotion_concurrentcapture_mutation, none,
+        "'%0' mutated after capture by concurrent closure", (StringRef))
+NOTE(capturepromotion_concurrentcapture_closure_here, none,
+     "variable captured by concurrent closure", ())
+NOTE(capturepromotion_concurrentcapture_capturinguse_here, none,
+     "capturing use", ())
+NOTE(capturepromotion_variable_defined_here,none,
+     "variable defined here", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/SIL/SILType.h

@@ -565,6 +565,13 @@ class SILType {
   /// Returns true if this SILType is a differentiable type.
   bool isDifferentiable(SILModule &M) const;
 
+  /// If this is a SILBoxType, return getSILBoxFieldType(). Otherwise, return
+  /// SILType().
+  ///
+  /// \p field Return the type of the ith field of the box. Default set to 0
+  /// since we only support one field today. This is just future proofing.
+  SILType getSILBoxFieldType(const SILFunction *f, unsigned field = 0);
+
   /// Returns the hash code for the SILType.
   llvm::hash_code getHashCode() const {
     return llvm::hash_combine(*this);

lib/SIL/IR/SILType.cpp

@@ -699,3 +699,11 @@ bool SILType::isEffectivelyExhaustiveEnumType(SILFunction *f) {
   return decl->isEffectivelyExhaustive(f->getModule().getSwiftModule(),
                                        f->getResilienceExpansion());
 }
+
+SILType SILType::getSILBoxFieldType(const SILFunction *f, unsigned field) {
+  auto *boxTy = getASTType()->getAs<SILBoxType>();
+  if (!boxTy)
+    return SILType();
+  return ::getSILBoxFieldType(f->getTypeExpansionContext(), boxTy,
+                              f->getModule().Types, field);
+}

lib/SIL/Verifier/SILVerifier.cpp

@@ -1671,13 +1671,24 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     SILFunctionConventions substConv(substTy, F.getModule());
     unsigned appliedArgStartIdx =
         substConv.getNumSILArguments() - PAI->getNumArguments();
-    for (unsigned i = 0, size = PAI->getArguments().size(); i < size; ++i) {
+    bool isConcurrentAndStageIsCanonical =
+        PAI->getFunctionType()->isConcurrent() &&
+        F.getModule().getStage() >= SILStage::Canonical;
+    for (auto p : llvm::enumerate(PAI->getArguments())) {
       requireSameType(
-          PAI->getArguments()[i]->getType(),
-          substConv.getSILArgumentType(appliedArgStartIdx + i,
+          p.value()->getType(),
+          substConv.getSILArgumentType(appliedArgStartIdx + p.index(),
                                        F.getTypeExpansionContext()),
           "applied argument types do not match suffix of function type's "
           "inputs");
+
+      // TODO: Expand this to also be true for address only types.
+      if (isConcurrentAndStageIsCanonical)
+        require(
+            !p.value()->getType().getASTType()->is<SILBoxType>() ||
+                p.value()->getType().getSILBoxFieldType(&F).isAddressOnly(F),
+            "Concurrent partial apply in canonical SIL with a loadable box "
+            "type argument?!");
     }
 
     // The arguments to the result function type must match the prefix of the

lib/SILOptimizer/Mandatory/CapturePromotion.cpp

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -44,7 +44,9 @@
 
 #define DEBUG_TYPE "sil-capture-promotion"
 
+#include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/Basic/FrozenMultiMap.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/TypeSubstCloner.h"
@@ -56,6 +58,7 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <tuple>
 
 using namespace swift;
@@ -765,6 +768,15 @@ struct EscapeMutationScanningState {
   /// found.
   SmallVector<Operand *, 8> accumulatedEscapes;
 
+  /// A multimap that maps partial applies to the set of operands in the partial
+  /// applies referenced function that the pass has identified as being the use
+  /// that caused the partial apply to capture our box.
+  ///
+  /// We use a frozen multi-map since our algorithm first accumulates this info
+  /// and then wants to use it, perfect for the 2-stage frozen multi map.
+  SmallFrozenMultiMap<PartialApplyInst *, Operand *, 16>
+      accumulatedCaptureCausingUses;
+
   /// A flag that we use to ensure that we only ever see 1 project_box on an
   /// alloc_box.
   bool sawProjectBoxInst;
@@ -797,15 +809,17 @@ static bool isNonMutatingLoad(SILInstruction *inst) {
 /// address of the box's contents), return true if this box has mutating
 /// captures. Return false otherwise. All of the mutating captures that we find
 /// are placed into \p accumulatedMutatingUses.
-static bool getPartialApplyArgMutationsAndEscapes(
-    SILArgument *boxArg, SmallVectorImpl<Operand *> &accumulatedMutatingUses,
-    SmallVectorImpl<Operand *> &accumulatedEscapes) {
+static bool
+getPartialApplyArgMutationsAndEscapes(PartialApplyInst *pai,
+                                      SILArgument *boxArg,
+                                      EscapeMutationScanningState &state) {
   SmallVector<ProjectBoxInst *, 2> projectBoxInsts;
 
   // Conservatively do not allow any use of the box argument other than a
   // strong_release or projection, since this is the pattern expected from
   // SILGen.
   SmallVector<Operand *, 32> incrementalEscapes;
+  SmallVector<Operand *, 32> incrementalCaptureCausingUses;
   for (auto *use : boxArg->getUses()) {
     if (isa<StrongReleaseInst>(use->getUser()) ||
         isa<DestroyValueInst>(use->getUser()))
@@ -827,18 +841,25 @@ static bool getPartialApplyArgMutationsAndEscapes(
   // function that mirrors isNonEscapingUse.
   auto checkIfAddrUseMutating = [&](Operand *addrUse) -> bool {
     unsigned initSize = incrementalEscapes.size();
-    auto *addrInst = addrUse->getUser();
-    if (auto *seai = dyn_cast<StructElementAddrInst>(addrInst)) {
+    auto *addrUser = addrUse->getUser();
+    if (auto *seai = dyn_cast<StructElementAddrInst>(addrUser)) {
       for (auto *seaiUse : seai->getUses()) {
-        if (!isNonMutatingLoad(seaiUse->getUser())) {
+        if (isNonMutatingLoad(seaiUse->getUser())) {
+          incrementalCaptureCausingUses.push_back(seaiUse);
+        } else {
           incrementalEscapes.push_back(seaiUse);
         }
       }
       return incrementalEscapes.size() != initSize;
     }
 
-    if (isNonMutatingLoad(addrInst) || isa<DebugValueAddrInst>(addrInst) ||
-        isa<MarkFunctionEscapeInst>(addrInst) || isa<EndAccessInst>(addrInst)) {
+    if (isNonMutatingLoad(addrUser)) {
+      incrementalCaptureCausingUses.push_back(addrUse);
+      return false;
+    }
+
+    if (isa<DebugValueAddrInst>(addrUser) ||
+        isa<MarkFunctionEscapeInst>(addrUser) || isa<EndAccessInst>(addrUser)) {
       return false;
     }
 
@@ -859,10 +880,15 @@ static bool getPartialApplyArgMutationsAndEscapes(
     }
   }
 
+  auto &accCaptureCausingUses = state.accumulatedCaptureCausingUses;
+  while (!incrementalCaptureCausingUses.empty())
+    accCaptureCausingUses.insert(pai,
+                                 incrementalCaptureCausingUses.pop_back_val());
+
   if (incrementalEscapes.empty())
     return false;
   while (!incrementalEscapes.empty())
-    accumulatedEscapes.push_back(incrementalEscapes.pop_back_val());
+    state.accumulatedEscapes.push_back(incrementalEscapes.pop_back_val());
   return true;
 }
 
@@ -925,8 +951,7 @@ bool isPartialApplyNonEscapingUser(Operand *currentOp, PartialApplyInst *pai,
   // Verify that this closure is known not to mutate the captured value; if
   // it does, then conservatively refuse to promote any captures of this
   // value.
-  if (getPartialApplyArgMutationsAndEscapes(boxArg, state.accumulatedMutations,
-                                            state.accumulatedEscapes)) {
+  if (getPartialApplyArgMutationsAndEscapes(pai, boxArg, state)) {
     LLVM_DEBUG(llvm::dbgs() << "        FAIL: Have a mutation or escape of a "
                                "partial apply arg?!\n");
     return false;
@@ -1195,6 +1220,57 @@ static bool findEscapeOrMutationUses(Operand *op,
   return isNonEscapingUse(op, state);
 }
 
+/// We found a capture of \p abi in concurrent closure \p pai that we can not
+/// promote to a by value capture. Emit a nice warning (FIXME: error) to warn
+/// the user and provide the following information in the compiler feedback:
+///
+/// 1. The source loc where the variable's box is written to.
+///
+/// 2. The source loc of the captured variable's declaration.
+///
+/// 3. The source loc of the start of the concurrent closure that caused the
+///    variable to be captured.
+///
+/// 4. All places in the concurrent closure that triggered the box's
+///    capture. NOTE: For objects these are load points. For address only things
+///    it is still open for debate at this point.
+static void diagnoseInvalidCaptureByConcurrentClosure(
+    AllocBoxInst *abi, PartialApplyInst *pai,
+    const EscapeMutationScanningState &state, SILInstruction *mutatingUser) {
+  auto captureCausingUses = state.accumulatedCaptureCausingUses.find(pai);
+  if (!captureCausingUses) {
+    llvm::errs() << "Didn't find capture causing use of partial apply: "
+                 << *pai;
+    llvm::errs() << "Original Func: " << pai->getFunction()->getName() << '\n';
+    llvm::errs() << "Partial Applied Func: "
+                 << pai->getReferencedFunctionOrNull()->getName() << '\n';
+    llvm::report_fatal_error("standard compiler error");
+  }
+
+  auto &astCtx = pai->getFunction()->getASTContext();
+  auto &de = astCtx.Diags;
+  auto varInfo = abi->getVarInfo();
+  StringRef name = "<unknown>";
+  if (varInfo) {
+    name = varInfo->Name;
+  }
+
+  de.diagnoseWithNotes(
+      de.diagnose(mutatingUser->getLoc().getSourceLoc(),
+                  diag::capturepromotion_concurrentcapture_mutation, name),
+      [&]() {
+        de.diagnose(abi->getLoc().getSourceLoc(),
+                    diag::capturepromotion_variable_defined_here);
+        de.diagnose(pai->getLoc().getSourceLoc(),
+                    diag::capturepromotion_concurrentcapture_closure_here);
+        for (auto *use : *captureCausingUses) {
+          de.diagnose(
+              use->getUser()->getLoc().getSourceLoc(),
+              diag::capturepromotion_concurrentcapture_capturinguse_here);
+        }
+      });
+}
+
 /// Examine an alloc_box instruction, returning true if at least one
 /// capture of the boxed variable is promotable.  If so, then the pair of the
 /// partial_apply instruction and the index of the box argument in the closure's
@@ -1203,7 +1279,7 @@ static bool
 examineAllocBoxInst(AllocBoxInst *abi, ReachabilityInfo &ri,
                     llvm::DenseMap<PartialApplyInst *, unsigned> &im) {
   LLVM_DEBUG(llvm::dbgs() << "Visiting alloc box: " << *abi);
-  EscapeMutationScanningState state{{}, {}, false, im};
+  EscapeMutationScanningState state{{}, {}, {}, false, im};
 
   // Scan the box for escaping or mutating uses.
   for (auto *use : abi->getUses()) {
@@ -1220,6 +1296,7 @@ examineAllocBoxInst(AllocBoxInst *abi, ReachabilityInfo &ri,
     return false;
   }
 
+  state.accumulatedCaptureCausingUses.setFrozen();
   LLVM_DEBUG(llvm::dbgs() << "We can optimize this alloc box!\n");
 
   // Helper lambda function to determine if instruction b is strictly after
@@ -1249,6 +1326,13 @@ examineAllocBoxInst(AllocBoxInst *abi, ReachabilityInfo &ri,
       // block is after the partial_apply.
       if (ri.isReachable(pai->getParent(), user->getParent()) ||
           (pai->getParent() == user->getParent() && isAfter(pai, user))) {
+        // If our partial apply is concurrent and we can not promote this, emit
+        // a warning that shows the variable, where the variable is captured,
+        // and the mutation that we found.
+        if (pai->getFunctionType()->isConcurrent()) {
+          diagnoseInvalidCaptureByConcurrentClosure(abi, pai, state, user);
+        }
+
         LLVM_DEBUG(llvm::dbgs() << "    Invalidating: " << *pai);
         LLVM_DEBUG(llvm::dbgs() << "    Because of user: " << *user);
         auto prev = iter++;
@@ -1257,6 +1341,7 @@ examineAllocBoxInst(AllocBoxInst *abi, ReachabilityInfo &ri,
       }
       ++iter;
     }
+
     // If there are no valid captures left, then stop.
     if (im.empty()) {
       LLVM_DEBUG(llvm::dbgs() << "    Ran out of valid captures... bailing!\n");