Infer 'isolated' closure parameters from context.

include/swift/Sema/ConstraintSystem.h

@@ -1200,6 +1200,9 @@ class Solution {
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
       caseLabelItems;
 
+  /// The set of parameters that have been inferred to be 'isolated'.
+  llvm::SmallVector<ParamDecl *, 2> isolatedParams;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;
@@ -2360,6 +2363,9 @@ class ConstraintSystem {
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
       caseLabelItems;
 
+  /// The set of parameters that have been inferred to be 'isolated'.
+  llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
+
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
     OpenedParameterTypes;
@@ -2924,6 +2930,9 @@ class ConstraintSystem {
     /// The length of \c caseLabelItems.
     unsigned numCaseLabelItems;
 
+    /// The length of \c isolatedParams.
+    unsigned numIsolatedParams;
+
     /// The length of \c ImplicitValueConversions.
     unsigned numImplicitValueConversions;
 

lib/Sema/CSClosure.cpp

@@ -1370,6 +1370,12 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     auto *params = closure->getParameters();
     TypeChecker::coerceParameterListToType(params, closureFnType);
 
+    // Find any isolated parameters in this closure and mark them as isolated.
+    for (auto param : solution.isolatedParams) {
+      if (param->getDeclContext() == closure)
+        param->setIsolated(true);
+    }
+
     // Coerce the result type, if it was written explicitly.
     if (closure->hasExplicitResultType()) {
       closure->setExplicitResultType(closureFnType->getResult());

lib/Sema/CSSimplify.cpp

@@ -1854,7 +1854,8 @@ static bool isSingleTupleParam(ASTContext &ctx,
     return false;
 
   const auto &param = params.front();
-  if (param.isVariadic() || param.isInOut() || param.hasLabel())
+  if (param.isVariadic() || param.isInOut() || param.hasLabel() ||
+      param.isIsolated())
     return false;
 
   auto paramType = param.getPlainType();
@@ -8988,15 +8989,22 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
     auto param = inferredClosureType->getParams()[i];
     auto *paramDecl = paramList->get(i);
 
-    // In case of anonymous or name-only parameters, let's infer inout/variadic
-    // flags from context, that helps to propagate type information into the
-    // internal type of the parameter and reduces inference solver has to make.
+    // In case of anonymous or name-only parameters, let's infer
+    // inout/variadic/isolated flags from context, that helps to propagate
+    // type information into the internal type of the parameter and reduces
+    // inference solver has to make.
     if (!paramDecl->getTypeRepr()) {
       if (auto contextualParam = getContextualParamAt(i)) {
         auto flags = param.getParameterFlags();
+
+        // Note when a parameter is inferred to be isolated.
+        if (contextualParam->isIsolated() && !flags.isIsolated() && paramDecl)
+          isolatedParams.insert(paramDecl);
+
         param =
             param.withFlags(flags.withInOut(contextualParam->isInOut())
-                                 .withVariadic(contextualParam->isVariadic()));
+                                 .withVariadic(contextualParam->isVariadic())
+                                 .withIsolated(contextualParam->isIsolated()));
       }
     }
 

lib/Sema/CSSolver.cpp

@@ -186,6 +186,7 @@ Solution ConstraintSystem::finalize() {
 
   solution.solutionApplicationTargets = solutionApplicationTargets;
   solution.caseLabelItems = caseLabelItems;
+  solution.isolatedParams.append(isolatedParams.begin(), isolatedParams.end());
 
   for (const auto &transformed : resultBuilderTransformed) {
     solution.resultBuilderTransformed.insert(transformed);
@@ -292,6 +293,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
       setCaseLabelItemInfo(info.first, info.second);
   }
 
+  for (auto param : solution.isolatedParams) {
+    isolatedParams.insert(param);
+  }
+
   for (const auto &transformed : solution.resultBuilderTransformed) {
     resultBuilderTransformed.push_back(transformed);
   }
@@ -521,6 +526,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numContextualTypes = cs.contextualTypes.size();
   numSolutionApplicationTargets = cs.solutionApplicationTargets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
+  numIsolatedParams = cs.isolatedParams.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
   numArgumentLists = cs.ArgumentLists.size();
 
@@ -629,6 +635,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any case label item infos.
   truncate(cs.caseLabelItems, numCaseLabelItems);
 
+  // Remove any isolated parameters.
+  truncate(cs.isolatedParams, numIsolatedParams);
+
   // Remove any implicit value conversions.
   truncate(cs.ImplicitValueConversions, numImplicitValueConversions);
 

test/Concurrency/isolated_parameters.swift

@@ -111,3 +111,23 @@ func testTuplingIsolated(_ a: isolated A, _ b: isolated A) {
   // expected-error@-1 {{generic parameter 'Ts' could not be inferred}}
   // expected-error@-2 {{cannot convert value of type '(isolated A, isolated A) -> ()' to expected argument type '(Ts) -> Void'}}
 }
+
+// Inference of "isolated" on closure parameters.
+@available(SwiftStdlib 5.5, *)
+func testIsolatedClosureInference(a: A) {
+  let _: (isolated A) -> Void = {
+    $0.f()
+  }
+
+  let _: (isolated A) -> Void = {
+    $0.f()
+  }
+
+  let _: (isolated A) -> Void = { a2 in
+    a2.f()
+  }
+
+  let _: (isolated A) -> Void = { (a2: isolated A) in
+    a2.f()
+  }
+}