[Concurrency] Infer @Concurrent on closures from contextual type

Author: DougGregor

lib/Sema/CSApply.cpp

@@ -6661,9 +6661,24 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
       }
     }
 
-    // If we have a ClosureExpr, then we can safely propagate the 'no escape'
+    // If we have a ClosureExpr, then we can safely propagate the 'concurrent'
     // bit to the closure without invalidating prior analysis.
     auto fromEI = fromFunc->getExtInfo();
+    if (toEI.isConcurrent() && !fromEI.isConcurrent()) {
+      auto newFromFuncType = fromFunc->withExtInfo(fromEI.withConcurrent());
+      if (applyTypeToClosureExpr(cs, expr, newFromFuncType)) {
+        fromFunc = newFromFuncType->castTo<FunctionType>();
+
+        // Propagating the 'concurrent' bit might have satisfied the entire
+        // conversion. If so, we're done, otherwise keep converting.
+        if (fromFunc->isEqual(toType))
+          return expr;
+      }
+    }
+
+    // If we have a ClosureExpr, then we can safely propagate the 'no escape'
+    // bit to the closure without invalidating prior analysis.
+    fromEI = fromFunc->getExtInfo();
     if (toEI.isNoEscape() && !fromEI.isNoEscape()) {
       auto newFromFuncType = fromFunc->withExtInfo(fromEI.withNoEscape());
       if (!isInDefaultArgumentContext &&

lib/Sema/CSSimplify.cpp

@@ -1894,19 +1894,19 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
     }
   }
 
-    // A @concurrent function can be a subtype of a non-@concurrent function.
-    if (func1->isConcurrent() != func2->isConcurrent()) {
-      // Cannot add '@concurrent'.
-      if (func2->isConcurrent() || kind < ConstraintKind::Subtype) {
-        if (!shouldAttemptFixes())
-          return getTypeMatchFailure(locator);
+  // A @concurrent function can be a subtype of a non-@concurrent function.
+  if (func1->isConcurrent() != func2->isConcurrent()) {
+    // Cannot add '@concurrent'.
+    if (func2->isConcurrent() || kind < ConstraintKind::Subtype) {
+      if (!shouldAttemptFixes())
+        return getTypeMatchFailure(locator);
 
-        auto *fix = AddConcurrentAttribute::create(
-            *this, func1, func2, getConstraintLocator(locator));
-        if (recordFix(fix))
-          return getTypeMatchFailure(locator);
-      }
+      auto *fix = AddConcurrentAttribute::create(
+          *this, func1, func2, getConstraintLocator(locator));
+      if (recordFix(fix))
+        return getTypeMatchFailure(locator);
     }
+  }
 
   // A non-@noescape function type can be a subtype of a @noescape function
   // type.
@@ -7864,9 +7864,16 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
     parameters.push_back(param);
   }
 
+  // Propagate @concurrent from the contextual type to the closure.
+  auto closureExtInfo = inferredClosureType->getExtInfo();
+  if (auto contextualFnType = contextualType->getAs<FunctionType>()) {
+    if (contextualFnType->isConcurrent())
+      closureExtInfo = closureExtInfo.withConcurrent();
+  }
+
   auto closureType =
       FunctionType::get(parameters, inferredClosureType->getResult(),
-                        inferredClosureType->getExtInfo());
+                        closureExtInfo);
   assignFixedType(typeVar, closureType, closureLocator);
 
   // If there is a result builder to apply, do so now.

test/attr/attr_concurrent.swift

@@ -24,3 +24,17 @@ func passingConcurrentOrNot(
   acceptsNonConcurrent(cfn) // okay
   acceptsNonConcurrent(ncfn) // okay
 }
+
+func closures() {
+  // Okay, inferring @concurrent
+  acceptsConcurrent { $0 }
+  acceptsConcurrent({ $0 })
+  acceptsConcurrent({ i in i })
+  acceptsConcurrent({ (i: Int) -> Int in
+      print(i)
+      return i
+    })
+
+  let closure1 = { $0 + 1 } // inferred to be non-concurrent
+  acceptsConcurrent(closure1) // expected-error{{converting non-concurrent function value to '@concurrent (Int) -> Int' may introduce data races}}
+}