Merge pull request #59963 from DougGregor/constraint-system-preconcurrency

[Constraint solver] Track "isolated by preconcurrency" in the solver.

Author: DougGregor

include/swift/Sema/ConstraintSystem.h

@@ -1320,6 +1320,10 @@ class Solution {
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallVector<ParamDecl *, 2> isolatedParams;
 
+  /// The set of closures that have been inferred to be "isolated by
+  /// preconcurrency".
+  llvm::SmallVector<const ClosureExpr *, 2> preconcurrencyClosures;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;
@@ -2519,6 +2523,13 @@ struct GetClosureType {
   Type operator()(const AbstractClosureExpr *expr) const;
 };
 
+/// Retrieve the closure type from the constraint system.
+struct ClosureIsolatedByPreconcurrency {
+  ConstraintSystem &cs;
+
+  bool operator()(const ClosureExpr *expr) const;
+};
+
 /// Describes the type produced when referencing a declaration.
 struct DeclReferenceType {
   /// The "opened" type, which is the type of the declaration where any
@@ -2573,6 +2584,7 @@ class ConstraintSystem {
   friend class ConjunctionElement;
   friend class RequirementFailure;
   friend class MissingMemberFailure;
+  friend struct ClosureIsolatedByPreconcurrency;
 
   class SolverScope;
 
@@ -2703,6 +2715,10 @@ class ConstraintSystem {
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
 
+  /// The set of closures that have been inferred to be "isolated by
+  /// preconcurrency".
+  llvm::SmallSetVector<const ClosureExpr *, 2> preconcurrencyClosures;
+
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
     OpenedParameterTypes;
@@ -3273,6 +3289,9 @@ class ConstraintSystem {
     /// The length of \c isolatedParams.
     unsigned numIsolatedParams;
 
+    /// The length of \c PreconcurrencyClosures.
+    unsigned numPreconcurrencyClosures;
+
     /// The length of \c ImplicitValueConversions.
     unsigned numImplicitValueConversions;
 
@@ -4553,6 +4572,10 @@ class ConstraintSystem {
       llvm::function_ref<Type(const AbstractClosureExpr *)> getClosureType =
         [](const AbstractClosureExpr *) {
           return Type();
+        },
+      llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency =
+        [](const ClosureExpr *closure) {
+          return closure->isIsolatedByPreconcurrency();
         });
 
   /// Given the opened type and a pile of information about a member reference,

lib/Sema/CSApply.cpp

@@ -5644,25 +5644,22 @@ static bool hasCurriedSelf(ConstraintSystem &cs, ConcreteDeclRef callee,
 
 /// Apply the contextually Sendable flag to the given expression,
 static void applyContextualClosureFlags(
-      Expr *expr, bool implicitSelfCapture, bool inheritActorContext,
-      bool isolatedByPreconcurrency) {
+      Expr *expr, bool implicitSelfCapture, bool inheritActorContext) {
   if (auto closure = dyn_cast<ClosureExpr>(expr)) {
     closure->setAllowsImplicitSelfCapture(implicitSelfCapture);
     closure->setInheritsActorContext(inheritActorContext);
-    closure->setIsolatedByPreconcurrency(isolatedByPreconcurrency);
     return;
   }
 
   if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
     applyContextualClosureFlags(
         captureList->getClosureBody(), implicitSelfCapture,
-        inheritActorContext, isolatedByPreconcurrency);
+        inheritActorContext);
   }
 
   if (auto identity = dyn_cast<IdentityExpr>(expr)) {
     applyContextualClosureFlags(
-        identity->getSubExpr(), implicitSelfCapture, inheritActorContext,
-        isolatedByPreconcurrency);
+        identity->getSubExpr(), implicitSelfCapture, inheritActorContext);
   }
 }
 
@@ -5689,8 +5686,6 @@ ArgumentList *ExprRewriter::coerceCallArguments(
   // Determine the parameter bindings.
   ParameterListInfo paramInfo(params, callee.getDecl(), skipCurriedSelf);
 
-  bool preconcurrency = callee && callee.getDecl()->preconcurrency();
-
   // If this application is an init(wrappedValue:) call that needs an injected
   // wrapped value placeholder, the first non-defaulted argument must be
   // wrapped in an OpaqueValueExpr.
@@ -5713,14 +5708,6 @@ ArgumentList *ExprRewriter::coerceCallArguments(
   auto matches = args->matches(params, [&](Expr *E) { return cs.getType(E); });
   if (matches && !shouldInjectWrappedValuePlaceholder &&
       !paramInfo.anyContextualInfo()) {
-    // Propagate preconcurrency to any closure arguments.
-    if (preconcurrency) {
-      for (const auto &arg : *args) {
-        Expr *argExpr = arg.getExpr();
-        applyContextualClosureFlags(argExpr, false, false, preconcurrency);
-      }
-    }
-
     return args;
   }
 
@@ -5871,7 +5858,7 @@ ArgumentList *ExprRewriter::coerceCallArguments(
     bool isImplicitSelfCapture = paramInfo.isImplicitSelfCapture(paramIdx);
     bool inheritsActorContext = paramInfo.inheritsActorContext(paramIdx);
     applyContextualClosureFlags(
-        argExpr, isImplicitSelfCapture, inheritsActorContext, preconcurrency);
+        argExpr, isImplicitSelfCapture, inheritsActorContext);
 
     // If the types exactly match, this is easy.
     auto paramType = param.getOldType();

lib/Sema/CSClosure.cpp

@@ -1656,6 +1656,9 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
         param->setIsolated(true);
     }
 
+    if (llvm::is_contained(solution.preconcurrencyClosures, closure))
+      closure->setIsolatedByPreconcurrency();
+
     // Coerce the result type, if it was written explicitly.
     if (closure->hasExplicitResultType()) {
       closure->setExplicitResultType(closureFnType->getResult());

lib/Sema/CSSimplify.cpp

@@ -2736,7 +2736,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
         // For a @preconcurrency callee outside of a strict concurrency
         // context, ignore.
         if (hasPreconcurrencyCallee(this, locator) &&
-            !contextRequiresStrictConcurrencyChecking(DC, GetClosureType{*this}))
+            !contextRequiresStrictConcurrencyChecking(DC, GetClosureType{*this}, ClosureIsolatedByPreconcurrency{*this}))
           return FixBehavior::Suppress;
 
         // Otherwise, warn until Swift 6.
@@ -9838,6 +9838,10 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
   auto *closure = castToExpr<ClosureExpr>(closureLocator->getAnchor());
   auto *inferredClosureType = getClosureType(closure);
 
+  // Note if this closure is isolated by preconcurrency.
+  if (hasPreconcurrencyCallee(this, locator))
+    preconcurrencyClosures.insert(closure);
+
   // Let's look through all optionals associated with contextual
   // type to make it possible to infer parameter/result type of
   // the closure faster e.g.:

lib/Sema/CSSolver.cpp

@@ -182,6 +182,8 @@ Solution ConstraintSystem::finalize() {
   solution.solutionApplicationTargets = solutionApplicationTargets;
   solution.caseLabelItems = caseLabelItems;
   solution.isolatedParams.append(isolatedParams.begin(), isolatedParams.end());
+  solution.preconcurrencyClosures.append(preconcurrencyClosures.begin(),
+                                         preconcurrencyClosures.end());
 
   for (const auto &transformed : resultBuilderTransformed) {
     solution.resultBuilderTransformed.insert(transformed);
@@ -295,6 +297,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     isolatedParams.insert(param);
   }
 
+  for (auto closure : solution.preconcurrencyClosures) {
+    preconcurrencyClosures.insert(closure);
+  }
+
   for (const auto &transformed : solution.resultBuilderTransformed) {
     resultBuilderTransformed.insert(transformed);
   }
@@ -535,6 +541,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numSolutionApplicationTargets = cs.solutionApplicationTargets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
   numIsolatedParams = cs.isolatedParams.size();
+  numPreconcurrencyClosures = cs.preconcurrencyClosures.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
   numArgumentLists = cs.ArgumentLists.size();
   numImplicitCallAsFunctionRoots = cs.ImplicitCallAsFunctionRoots.size();
@@ -646,6 +653,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any isolated parameters.
   truncate(cs.isolatedParams, numIsolatedParams);
 
+  // Remove any preconcurrency closures.
+  truncate(cs.preconcurrencyClosures, numPreconcurrencyClosures);
+
   // Remove any implicit value conversions.
   truncate(cs.ImplicitValueConversions, numImplicitValueConversions);
 

lib/Sema/ConstraintSystem.cpp

@@ -1229,6 +1229,12 @@ Type GetClosureType::operator()(const AbstractClosureExpr *expr) const {
   return Type();
 }
 
+bool
+ClosureIsolatedByPreconcurrency::operator()(const ClosureExpr *expr) const {
+  return expr->isIsolatedByPreconcurrency() ||
+      cs.preconcurrencyClosures.count(expr);
+}
+
 Type ConstraintSystem::getUnopenedTypeOfReference(
     VarDecl *value, Type baseType, DeclContext *UseDC,
     ConstraintLocator *memberLocator, bool wantInterfaceType) {
@@ -1244,20 +1250,22 @@ Type ConstraintSystem::getUnopenedTypeOfReference(
 
         return wantInterfaceType ? var->getInterfaceType() : var->getType();
       },
-      memberLocator, wantInterfaceType, GetClosureType{*this});
+      memberLocator, wantInterfaceType, GetClosureType{*this},
+      ClosureIsolatedByPreconcurrency{*this});
 }
 
 Type ConstraintSystem::getUnopenedTypeOfReference(
     VarDecl *value, Type baseType, DeclContext *UseDC,
     llvm::function_ref<Type(VarDecl *)> getType,
     ConstraintLocator *memberLocator, bool wantInterfaceType,
-    llvm::function_ref<Type(const AbstractClosureExpr *)> getClosureType) {
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getClosureType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency) {
   Type requestedType =
       getType(value)->getWithoutSpecifierType()->getReferenceStorageReferent();
 
   // Adjust the type for concurrency.
   requestedType = adjustVarTypeForConcurrency(
-      requestedType, value, UseDC, getClosureType);
+      requestedType, value, UseDC, getClosureType, isolatedByPreconcurrency);
 
   // If we're dealing with contextual types, and we referenced this type from
   // a different context, map the type.
@@ -1436,7 +1444,8 @@ AnyFunctionType *ConstraintSystem::adjustFunctionTypeForConcurrency(
     unsigned numApplies, bool isMainDispatchQueue,
     OpenedTypeMap &replacements) {
   return swift::adjustFunctionTypeForConcurrency(
-      fnType, decl, dc, numApplies, isMainDispatchQueue, GetClosureType{*this},
+      fnType, decl, dc, numApplies, isMainDispatchQueue,
+      GetClosureType{*this}, ClosureIsolatedByPreconcurrency{*this},
       [&](Type type) {
         if (replacements.empty())
           return type;
@@ -2527,7 +2536,8 @@ Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
         type = withDynamicSelfResultReplaced(type, /*uncurryLevel=*/0);
       }
       type = adjustVarTypeForConcurrency(
-          type, var, useDC, GetClosureType{*this});
+          type, var, useDC, GetClosureType{*this},
+          ClosureIsolatedByPreconcurrency{*this});
     } else if (isa<AbstractFunctionDecl>(decl) || isa<EnumElementDecl>(decl)) {
       if (decl->isInstanceMember() &&
           (!overload.getBaseType() ||

lib/Sema/TypeCheckConcurrency.cpp

@@ -609,6 +609,9 @@ static void addSendableFixIt(const GenericTypeParamDecl *genericArgument,
 static bool shouldDiagnoseExistingDataRaces(const DeclContext *dc) {
   return contextRequiresStrictConcurrencyChecking(dc, [](const AbstractClosureExpr *) {
     return Type();
+  },
+  [](const ClosureExpr *closure) {
+    return closure->isIsolatedByPreconcurrency();
   });
 }
 
@@ -4023,7 +4026,8 @@ void swift::checkOverrideActorIsolation(ValueDecl *value) {
 
 bool swift::contextRequiresStrictConcurrencyChecking(
     const DeclContext *dc,
-    llvm::function_ref<Type(const AbstractClosureExpr *)> getType) {
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency) {
   switch (dc->getASTContext().LangOpts.StrictConcurrencyLevel) {
   case StrictConcurrency::Complete:
     return true;
@@ -4044,7 +4048,7 @@ bool swift::contextRequiresStrictConcurrencyChecking(
 
         // Don't take any more cues if this only got its type information by
         // being provided to a `@preconcurrency` operation.
-        if (explicitClosure->isIsolatedByPreconcurrency()) {
+        if (isolatedByPreconcurrency(explicitClosure)) {
           dc = dc->getParent();
           continue;
         }
@@ -4602,11 +4606,13 @@ static bool hasKnownUnsafeSendableFunctionParams(AbstractFunctionDecl *func) {
 
 Type swift::adjustVarTypeForConcurrency(
     Type type, VarDecl *var, DeclContext *dc,
-    llvm::function_ref<Type(const AbstractClosureExpr *)> getType) {
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency) {
   if (!var->preconcurrency())
     return type;
 
-  if (contextRequiresStrictConcurrencyChecking(dc, getType))
+  if (contextRequiresStrictConcurrencyChecking(
+          dc, getType, isolatedByPreconcurrency))
     return type;
 
   bool isLValue = false;
@@ -4727,9 +4733,11 @@ AnyFunctionType *swift::adjustFunctionTypeForConcurrency(
     AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
     unsigned numApplies, bool isMainDispatchQueue,
     llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency,
     llvm::function_ref<Type(Type)> openType) {
   // Apply unsafe concurrency features to the given function type.
-  bool strictChecking = contextRequiresStrictConcurrencyChecking(dc, getType);
+  bool strictChecking = contextRequiresStrictConcurrencyChecking(
+      dc, getType, isolatedByPreconcurrency);
   fnType = applyUnsafeConcurrencyToFunctionType(
       fnType, decl, strictChecking, numApplies, isMainDispatchQueue);
 
@@ -4789,7 +4797,10 @@ bool swift::completionContextUsesConcurrencyFeatures(const DeclContext *dc) {
   return contextRequiresStrictConcurrencyChecking(
       dc, [](const AbstractClosureExpr *) {
         return Type();
-      });
+      },
+    [](const ClosureExpr *closure) {
+      return closure->isIsolatedByPreconcurrency();
+    });
 }
 
 AbstractFunctionDecl const *swift::isActorInitOrDeInitContext(
@@ -4898,6 +4909,9 @@ static ActorIsolation getActorIsolationForReference(
             fromDC,
             [](const AbstractClosureExpr *closure) {
               return closure->getType();
+            },
+            [](const ClosureExpr *closure) {
+              return closure->isIsolatedByPreconcurrency();
             })) {
       declIsolation = ActorIsolation::forGlobalActor(
           declIsolation.getGlobalActor(), /*unsafe=*/false);

lib/Sema/TypeCheckConcurrency.h

@@ -102,7 +102,8 @@ void checkOverrideActorIsolation(ValueDecl *value);
 /// code where strict checking has been enabled.
 bool contextRequiresStrictConcurrencyChecking(
     const DeclContext *dc,
-    llvm::function_ref<Type(const AbstractClosureExpr *)> getType);
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency);
 
 /// Describes a referenced actor variable and whether it is isolated.
 struct ReferencedActor {
@@ -417,7 +418,8 @@ Type getExplicitGlobalActor(ClosureExpr *closure);
 /// Adjust the type of the variable for concurrency.
 Type adjustVarTypeForConcurrency(
     Type type, VarDecl *var, DeclContext *dc,
-    llvm::function_ref<Type(const AbstractClosureExpr *)> getType);
+    llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency);
 
 /// Adjust the given function type to account for concurrency-specific
 /// attributes whose affect on the type might differ based on context.
@@ -428,6 +430,7 @@ AnyFunctionType *adjustFunctionTypeForConcurrency(
     AnyFunctionType *fnType, ValueDecl *decl, DeclContext *dc,
     unsigned numApplies, bool isMainDispatchQueue,
     llvm::function_ref<Type(const AbstractClosureExpr *)> getType,
+    llvm::function_ref<bool(const ClosureExpr *)> isolatedByPreconcurrency,
     llvm::function_ref<Type(Type)> openType);
 
 /// Determine whether the given name is that of a DispatchQueue operation that

test/Concurrency/predates_concurrency.swift

@@ -152,3 +152,30 @@ func testStringPlacement() {
   let fn2 = StringPlacement.position(before:)
   let _: Int = fn2 // expected-error{{cannot convert value of type '(String) -> ([String]) -> Int' to specified type 'Int'}}
 }
+
+// @preconcurrency in an outer closure
+// (https://github.com/apple/swift/issues/59910)
+struct Scheduled<T> { }
+
+@preconcurrency
+func doPreconcurrency(_: @Sendable () -> Void) { }
+
+class EventLoop {
+  @discardableResult
+  @preconcurrency
+  func scheduleTask<T>(deadline: Int, _ task: @escaping @Sendable () throws -> T) -> Scheduled<T> { fatalError("") }
+}
+
+class C {
+  var ev: EventLoop? = nil
+
+  func test(i: Int) {
+    func doNext() {
+      doPreconcurrency {
+        self.ev?.scheduleTask(deadline: i, doNext)
+        return
+      }
+    }
+  }
+}
+