[Concurrency] Allow #isolation to have a more specific contextual type.

lib/Sema/CSGen.cpp

@@ -3923,6 +3923,13 @@ namespace {
     }
 
     Type visitCurrentContextIsolationExpr(CurrentContextIsolationExpr *E) {
+      // If this was expanded from the builtin `#isolation` macro, it
+      // already has a type.
+      if (auto type = E->getType())
+        return type;
+
+      // Otherwise, this was created for a `for await` loop, where its
+      // type is always `(any Actor)?`.
       auto actorProto = CS.getASTContext().getProtocol(
           KnownProtocolKind::Actor);
       return OptionalType::get(actorProto->getDeclaredExistentialType());

lib/Sema/ConstraintSystem.cpp

@@ -25,6 +25,7 @@
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
+#include "swift/AST/MacroDefinition.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
@@ -3964,6 +3965,16 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
         // Record a fix here
         (void)recordFix(MacroMissingPound::create(*this, macro, locator));
       }
+
+      // The default type of the #isolation builtin macro is `(any Actor)?`
+      if (macro->getBuiltinKind() == BuiltinMacroKind::IsolationMacro) {
+        auto *fnType = openedType->getAs<FunctionType>();
+        auto actor = getASTContext().getProtocol(KnownProtocolKind::Actor);
+        addConstraint(
+            ConstraintKind::Defaultable, fnType->getResult(),
+            OptionalType::get(actor->getDeclaredExistentialType()),
+            locator);
+      }
     }
   }
 

stdlib/public/Concurrency/Actor.swift

@@ -76,8 +76,12 @@ internal func _enqueueOnMain(_ job: UnownedJob)
 #if $Macros
 /// Produce a reference to the actor to which the enclosing code is
 /// isolated, or `nil` if the code is nonisolated.
+///
+/// If the type annotation provided for `#isolation` is not `(any Actor)?`,
+/// the type must match the enclosing actor type. If no type annotation is
+/// provided, the type defaults to `(any Actor)?`.
 @available(SwiftStdlib 5.1, *)
 @freestanding(expression)
-public macro isolation() -> (any Actor)? = Builtin.IsolationMacro
+public macro isolation<T>() -> T = Builtin.IsolationMacro
 #endif
 

test/Concurrency/isolation_macro.swift

@@ -1,5 +1,10 @@
+// RUN: %empty-directory(%t)
 // RUN: %target-swift-frontend -dump-ast %s -enable-experimental-feature OptionalIsolatedParameters | %FileCheck %s
 
+// Diagnostics testing
+// RUN: not %target-swift-frontend -swift-version 5 -typecheck -enable-experimental-feature OptionalIsolatedParameters -DTEST_DIAGNOSTICS %s > %t/diagnostics.txt 2>&1
+// RUN: %FileCheck %s --check-prefix CHECK-DIAGS < %t/diagnostics.txt
+
 // REQUIRES: concurrency
 // REQUIRES: asserts
 // REQUIRES: swift_swift_parser
@@ -94,3 +99,18 @@ extension A {
 
   func g() {}
 }
+
+#if TEST_DIAGNOSTICS
+@available(SwiftStdlib 5.1, *)
+@MainActor
+func testContextualType() {
+  let _: any Actor = #isolation
+  let _: MainActor = #isolation
+  let _: MainActor? = #isolation
+
+  // CHECK-DIAGS: error: cannot convert value of type 'MainActor' to expected argument type 'Int'
+  // CHECK-DIAGS: note: in expansion of macro 'isolation' here
+  // CHECK-DIAGS: let _: Int = #isolation
+  let _: Int = #isolation
+}
+#endif