[Clang import] Apply main-thread annotations from parameters to types.

Import a main-thread annotation on a function/method parameter in
(Objective-)C as a `@MainActor` function type.

Author: DougGregor

lib/ClangImporter/ImportDecl.cpp

@@ -8010,6 +8010,19 @@ SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
   return *sourceFile;
 }
 
+Optional<bool> swift::importer::isMainActorAttr(
+    ASTContext &ctx, const clang::SwiftAttrAttr *swiftAttr) {
+  if (swiftAttr->getAttribute() == "@MainActor" ||
+      swiftAttr->getAttribute() == "@MainActor(unsafe)" ||
+      swiftAttr->getAttribute() == "@UIActor") {
+    bool isUnsafe = swiftAttr->getAttribute() == "@MainActor(unsafe)" ||
+        !ctx.LangOpts.isSwiftVersionAtLeast(6);
+    return isUnsafe;
+  }
+
+  return None;
+}
+
 /// Import Clang attributes as Swift attributes.
 void ClangImporter::Implementation::importAttributes(
     const clang::NamedDecl *ClangDecl,
@@ -8180,11 +8193,8 @@ void ClangImporter::Implementation::importAttributes(
     if (auto swiftAttr = dyn_cast<clang::SwiftAttrAttr>(*AI)) {
       // FIXME: Hard-core @MainActor and @UIActor, because we don't have a
       // point at which to do name lookup for imported entities.
-      if (swiftAttr->getAttribute() == "@MainActor" ||
-          swiftAttr->getAttribute() == "@MainActor(unsafe)" ||
-          swiftAttr->getAttribute() == "@UIActor") {
-        bool isUnsafe = swiftAttr->getAttribute() == "@MainActor(unsafe)" ||
-            !C.LangOpts.isSwiftVersionAtLeast(6);
+      if (auto isMainActor = isMainActorAttr(SwiftContext, swiftAttr)) {
+        bool isUnsafe = *isMainActor;
         if (Type mainActorType = getMainActorType()) {
           auto typeExpr = TypeExpr::createImplicit(mainActorType, SwiftContext);
           auto attr = CustomAttr::create(SwiftContext, SourceLoc(), typeExpr);

lib/ClangImporter/ImportType.cpp

@@ -1695,26 +1695,50 @@ ImportedType ClangImporter::Implementation::importPropertyType(
 
 /// Apply an attribute to a function type.
 static Type applyToFunctionType(
-    Type type, llvm::function_ref<Type(FunctionType *)> transform) {
+    Type type, llvm::function_ref<ASTExtInfo(ASTExtInfo)> transform) {
   // Recurse into optional types.
   if (Type objectType = type->getOptionalObjectType()) {
     return OptionalType::get(applyToFunctionType(objectType, transform));
   }
 
   // Apply @noescape to function types.
-  if (auto funcType = type->getAs<FunctionType>())
-    return transform(funcType);
+  if (auto funcType = type->getAs<FunctionType>()) {
+    return FunctionType::get(funcType->getParams(), funcType->getResult(),
+                             transform(funcType->getExtInfo()));
+  }
 
   return type;
 }
 
-static Type applyParamAttributes(const clang::ParmVarDecl *param, Type type) {
-  // Map __attribute__((noescape)) to @noescape.
-  if (param->hasAttr<clang::NoEscapeAttr>()) {
-    type = applyToFunctionType(type, [](FunctionType *funcType) {
-      return FunctionType::get(funcType->getParams(), funcType->getResult(),
-                               funcType->getExtInfo().withNoEscape());
-    });
+Type ClangImporter::Implementation::applyParamAttributes(
+    const clang::ParmVarDecl *param, Type type) {
+  if (!param->hasAttrs())
+    return type;
+
+  for (auto attr : param->getAttrs()) {
+    // Map __attribute__((noescape)) to @noescape.
+    if (isa<clang::NoEscapeAttr>(attr)) {
+      type = applyToFunctionType(type, [](ASTExtInfo extInfo) {
+        return extInfo.withNoEscape();
+      });
+
+      continue;
+    }
+
+    auto swiftAttr = dyn_cast<clang::SwiftAttrAttr>(attr);
+    if (!swiftAttr)
+      continue;
+
+    // Map the main-actor attribute.
+    if (isMainActorAttr(SwiftContext, swiftAttr)) {
+      if (Type mainActor = getMainActorType()) {
+        type = applyToFunctionType(type, [&](ASTExtInfo extInfo) {
+          return extInfo.withGlobalActor(mainActor);
+        });
+      }
+
+      continue;
+    }
   }
 
   return type;

lib/ClangImporter/ImporterImpl.h

@@ -830,6 +830,8 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   void importAttributes(const clang::NamedDecl *ClangDecl, Decl *MappedDecl,
                         const clang::ObjCContainerDecl *NewContext = nullptr);
 
+  Type applyParamAttributes(const clang::ParmVarDecl *param, Type type);
+
   /// If we already imported a given decl, return the corresponding Swift decl.
   /// Otherwise, return nullptr.
   Decl *importDeclCached(const clang::NamedDecl *ClangDecl, Version version,
@@ -1562,6 +1564,14 @@ class SwiftNameLookupExtension : public clang::ModuleFileExtension {
                         const llvm::BitstreamCursor &stream) override;
 };
 
+/// Determines whether the given swift_attr attribute describes the main
+/// actor.
+///
+/// \returns None if this is not a main-actor attribute, and a Boolean
+/// indicating whether (unsafe) was provided in the attribute otherwise.
+Optional<bool> isMainActorAttr(
+    ASTContext &ctx, const clang::SwiftAttrAttr *swiftAttr);
+
 }
 }
 

test/ClangImporter/objc_async.swift

@@ -5,6 +5,8 @@
 import Foundation
 import ObjCConcurrency
 
+@MainActor func onlyOnMainActor() { }
+
 func testSlowServer(slowServer: SlowServer) async throws {
   let _: Int = await slowServer.doSomethingSlow("mail")
   let _: Bool = await slowServer.checkAvailability()
@@ -49,6 +51,8 @@ func testSlowServer(slowServer: SlowServer) async throws {
 
 
   _ = await slowServer.operations()
+
+  _ = await slowServer.runOnMainThread()
 }
 
 func testSlowServerSynchronous(slowServer: SlowServer) {
@@ -63,6 +67,11 @@ func testSlowServerSynchronous(slowServer: SlowServer) {
 
   let s = slowServer.operations
   _ = s + []
+
+  slowServer.runOnMainThread { s in
+    print(s)
+    onlyOnMainActor() // okay because runOnMainThread has a @MainActor closure
+  }
 }
 
 func testSlowServerOldSchool(slowServer: SlowServer) {

test/IDE/print_clang_objc_async.swift

@@ -26,6 +26,8 @@ import _Concurrency
 // CHECK-DAG:     func doSomethingFun(_ operation: String) async
 // CHECK-DAG:     func dance(_ step: String) async -> String
 // CHECK-DAG:     func __leap(_ height: Int) async -> String
+// CHECK-DAG:     func runOnMainThread(completionHandler completion: (@MainActor (String) -> Void)? = nil)
+// CHECK-DAG:     func runOnMainThread() async -> String
 // CHECK: {{^[}]$}}
 
 // CHECK-LABEL: protocol RefrigeratorDelegate

test/Inputs/clang-importer-sdk/usr/include/ObjCConcurrency.h

@@ -3,6 +3,8 @@
 
 #pragma clang assume_nonnull begin
 
+#define MAIN_ACTOR __attribute__((__swift_attr__("@MainActor")))
+
 @protocol ServiceProvider
 @property(readonly) NSArray<NSString *> *allOperations;
 -(void)allOperationsWithCompletionHandler:(void (^)(NSArray<NSString *> *))completion;
@@ -66,6 +68,9 @@ typedef void (^CompletionHandler)(NSString * _Nullable, NSString * _Nullable_res
 -(void)doSomethingFlaggyWithCompletionHandler:(void (^)(BOOL, NSString *_Nullable, NSError *_Nullable))completionHandler __attribute__((swift_async_error(nonzero_argument, 1)));
 -(void)doSomethingZeroFlaggyWithCompletionHandler:(void (^)(NSString *_Nullable, BOOL, NSError *_Nullable))completionHandler __attribute__((swift_async_error(zero_argument, 2)));
 -(void)doSomethingMultiResultFlaggyWithCompletionHandler:(void (^)(BOOL, NSString *_Nullable, NSError *_Nullable, NSString *_Nullable))completionHandler __attribute__((swift_async_error(zero_argument, 1)));
+
+
+-(void)runOnMainThreadWithCompletionHandler:(MAIN_ACTOR void (^ _Nullable)(NSString *))completion;
 @end
 
 @protocol RefrigeratorDelegate<NSObject>