Diagnose attempts to use #isolation within an @_unsafeInheritExecutor function

An `@_unsafeInheritExecutor` function is unsafe because it doesn't
really "inherit" the executor, it just avoids immediately hopping off
the executor. That means that using `#isolation` within such a
function is fundamentally broken. Ban the use of `#isolation` within
such a function, providing a Fix-It that removes the
`@_unsafeInheritExecutor` attribute and adds a defaulted parameter

    isolation: (any Actor)? = #isolation

instead. That's the real, safe pattern that want going forward.

We did say it was unsafe, after all. Part of rdar://131151376.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -3997,6 +3997,9 @@ ERROR(reasync_without_async_parameter,none,
 
 ERROR(inherits_executor_without_async,none,
       "non-async functions cannot inherit an executor", ())
+ERROR(isolation_in_inherits_executor,none,
+      "#isolation%select{| (introduced by a default argument)}0 cannot be used "
+      "within an '@_unsafeInheritExecutor' function", (bool))
 
 ERROR(lifetime_invalid_global_scope,none, "%0 is only valid on methods",
       (DeclAttribute))

lib/Sema/TypeCheckConcurrency.cpp

@@ -1999,6 +1999,103 @@ bool swift::isAsyncDecl(ConcreteDeclRef declRef) {
   return false;
 }
 
+/// Find an enclosing function that has @
+static AbstractFunctionDecl *enclosingUnsafeInheritsExecutor(
+    const DeclContext *dc) {
+  for (; dc; dc = dc->getParent()) {
+    if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
+      if (func->getAttrs().hasAttribute<UnsafeInheritExecutorAttr>()) {
+        return const_cast<AbstractFunctionDecl *>(func);
+      }
+
+      return nullptr;
+    }
+
+    if (isa<AbstractClosureExpr>(dc))
+      return nullptr;
+
+    if (dc->isTypeContext())
+      return nullptr;
+  }
+
+  return nullptr;
+}
+
+/// Adjust the location used for diagnostics about #isolation to account for
+/// the fact that they show up in macro expansions.
+///
+/// Returns a pair containing the updated location and whether it's part of
+/// a default argument.
+static std::pair<SourceLoc, bool> adjustPoundIsolationDiagLoc(
+    CurrentContextIsolationExpr *isolationExpr,
+    ModuleDecl *module
+) {
+  // Not part of a macro expansion.
+  SourceLoc diagLoc = isolationExpr->getLoc();
+  auto sourceFile = module->getSourceFileContainingLocation(diagLoc);
+  if (!sourceFile)
+    return { diagLoc, false };
+  auto macroExpansionRange = sourceFile->getMacroInsertionRange();
+  if (macroExpansionRange.Start.isInvalid())
+    return { diagLoc, false };
+
+  diagLoc = macroExpansionRange.Start;
+
+  // If this is from a default argument, note that and go one more
+  // level "out" to the place where the default argument was
+  // introduced.
+  auto expansionSourceFile = module->getSourceFileContainingLocation(diagLoc);
+  if (!expansionSourceFile ||
+      expansionSourceFile->Kind != SourceFileKind::DefaultArgument)
+    return { diagLoc, false };
+
+  return {
+    expansionSourceFile->getNodeInEnclosingSourceFile().getStartLoc(),
+    true
+  };
+}
+
+/// Replace the @_unsafeInheritExecutor with a defaulted isolation
+/// parameter.
+static void replaceUnsafeInheritExecutorWithDefaultedIsolationParam(
+    AbstractFunctionDecl *func, InFlightDiagnostic &diag) {
+  auto attr = func->getAttrs().getAttribute<UnsafeInheritExecutorAttr>();
+  assert(attr && "Caller didn't validate the presence of the attribute");
+
+  // Look for the place where we should insert the new 'isolation' parameter.
+  // We insert toward the back, but skip over any parameters that have function
+  // type.
+  unsigned insertionPos = func->getParameters()->size();
+  while (insertionPos > 0) {
+    Type paramType = func->getParameters()->get(insertionPos - 1)->getInterfaceType();
+    if (paramType->lookThroughSingleOptionalType()->is<AnyFunctionType>()) {
+      --insertionPos;
+      continue;
+    }
+
+    break;
+  }
+
+  // Determine the text to insert. We put the commas before and after, then
+  // slice them away depending on whether we have parameters before or after.
+  StringRef newParameterText = ", isolation: isolated (any Actor)? = #isolation, ";
+  if (insertionPos == 0)
+    newParameterText = newParameterText.drop_front(2);
+  if (insertionPos == func->getParameters()->size())
+    newParameterText = newParameterText.drop_back(2);
+
+  // Determine where to insert the new parameter.
+  SourceLoc insertionLoc;
+  if (insertionPos < func->getParameters()->size()) {
+    insertionLoc = func->getParameters()->get(insertionPos)->getStartLoc();
+  } else {
+    insertionLoc = func->getParameters()->getRParenLoc();
+  }
+
+  diag.fixItRemove(attr->getRangeWithAt());
+  diag.fixItInsert(insertionLoc, newParameterText);
+}
+
 /// Check if it is safe for the \c globalActor qualifier to be removed from
 /// \c ty, when the function value of that type is isolated to that actor.
 ///
@@ -3744,6 +3841,25 @@ namespace {
 
     void recordCurrentContextIsolation(
         CurrentContextIsolationExpr *isolationExpr) {
+      // #isolation does not work within an `@_unsafeInheritExecutor` function.
+      if (auto func = enclosingUnsafeInheritsExecutor(getDeclContext())) {
+        // This expression is always written as a macro #isolation in source,
+        // so find the enclosing macro expansion expression's location.
+        SourceLoc diagLoc;
+        bool inDefaultArgument;
+        std::tie(diagLoc, inDefaultArgument) = adjustPoundIsolationDiagLoc(
+            isolationExpr, getDeclContext()->getParentModule());
+
+        bool inConcurrencyModule = getDeclContext()->getParentModule()->getName()
+            .str().equals("_Concurrency");
+
+        auto diag = ctx.Diags.diagnose(diagLoc,
+                                       diag::isolation_in_inherits_executor,
+                                       inDefaultArgument);
+        diag.limitBehaviorIf(inConcurrencyModule, DiagnosticBehavior::Warning);
+        replaceUnsafeInheritExecutorWithDefaultedIsolationParam(func, diag);
+      }
+
       // If an actor has already been assigned, we're done.
       if (isolationExpr->getActor())
         return;

test/Concurrency/unsafe_inherit_executor.swift

@@ -39,3 +39,31 @@ actor MyActor {
     await useNonSendable(object: self.object)
   }
 }
+
+// Note: the tests below are line-number-sensitive.
+func inheritsIsolationProperly(isolation: isolated (any Actor)? = #isolation) async { }
+
+// @_unsafeInheritExecutor does not work with #isolation
+@_unsafeInheritExecutor
+func unsafeCallerA(x: Int) async {
+  await inheritsIsolationProperly()
+  // expected-error@-1{{#isolation (introduced by a default argument) cannot be used within an '@_unsafeInheritExecutor' function}}{{47:1-24=}}{{48:26-26=, isolation: isolated (any Actor)? = #isolation}}
+}
+
+@_unsafeInheritExecutor
+func unsafeCallerB() async {
+  await inheritsIsolationProperly(isolation: #isolation)
+  // expected-error@-1{{#isolation cannot be used within an '@_unsafeInheritExecutor' function}}{{53:1-24=}}{{54:20-20=isolation: isolated (any Actor)? = #isolation}}
+}
+
+@_unsafeInheritExecutor
+func unsafeCallerC(x: Int, fn: () -> Void, fn2: () -> Void) async {
+  await inheritsIsolationProperly()
+  // expected-error@-1{{#isolation (introduced by a default argument) cannot be used within an '@_unsafeInheritExecutor' function}}{{59:1-24=}}{{60:28-28=, isolation: isolated (any Actor)? = #isolation, }}
+}
+
+@_unsafeInheritExecutor
+func unsafeCallerB(x: some AsyncSequence<Int, Never>) async {
+  for await _ in x { }
+  // expected-error@-1 2{{#isolation cannot be used within an `@_unsafeInheritExecutor` function}}
+}