Use @_unsafeInheritExecutor forms of with*Continuation from @_unsafeInheritExecutor functions

The move from `@_unsafeInheritExecutor` to `#isolation` for the
with*Continuation breaks code that is using `@_unsafeInheritExecutor` and
calling these APIs. This originally caused silent breakage (which manifest
as runtime crashes), and is now detected by the compiler as an error.

However, despite `@_unsafeInheritExecutor` being an unsafe,
not-intended-to-be-user-facing feature, it is indeed being used, along
with these APIs. Introduce _unsafeInheritExecutor_-prefixed versions of
the `with*Continuation` and `withTaskCancellationHandler` APIs into
the _Concurrency library that use `@_unsafeInheritExecutor`. Then,
teach the type checker to swap in these
_unsafeInheritExecutor_-prefixed versions in lieu of the originals
when they are called from an `@_unsafeInheritExecutor` function. This
allows existing code using `@_unsafeInheritExecutor` with these APIs
to continue working as it has before, albeit with a warning that
`@_unsafeInheritExecutor` has been removed.

Fixes rdar://131151376.

Author: DougGregor

lib/Sema/PreCheckExpr.cpp

@@ -16,6 +16,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "TypeChecker.h"
+#include "TypeCheckConcurrency.h"
 #include "TypeCheckType.h"
 #include "TypoCorrection.h"
 #include "swift/AST/ASTVisitor.h"
@@ -737,6 +738,14 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
       }
     }
 
+    // If we are in an @_unsafeInheritExecutor context, swap out
+    // declarations for their _unsafeInheritExecutor_ counterparts if they
+    // exist.
+    if (enclosingUnsafeInheritsExecutor(DC)) {
+      introduceUnsafeInheritExecutorReplacements(
+          DC, UDRE->getNameLoc().getBaseNameLoc(), ResultValues);
+    }
+
     return buildRefExpr(ResultValues, DC, UDRE->getNameLoc(),
                         UDRE->isImplicit(), UDRE->getFunctionRefKind());
   }

lib/Sema/TypeCheckConcurrency.cpp

@@ -2097,6 +2097,62 @@ void swift::replaceUnsafeInheritExecutorWithDefaultedIsolationParam(
   diag.fixItInsert(insertionLoc, newParameterText);
 }
 
+/// Whether this declaration context is in the _Concurrency module.
+static bool inConcurrencyModule(const DeclContext *dc) {
+  return dc->getParentModule()->getName().str().equals("_Concurrency");
+}
+
+void swift::introduceUnsafeInheritExecutorReplacements(
+    const DeclContext *dc, SourceLoc loc, SmallVectorImpl<ValueDecl *> &decls) {
+  if (decls.empty())
+    return;
+
+  auto isReplaceable = [&](ValueDecl *decl) {
+    return isa<FuncDecl>(decl) && inConcurrencyModule(decl->getDeclContext()) &&
+        decl->getDeclContext()->isModuleScopeContext();
+  };
+
+  // Make sure at least some of the entries are functions in the _Concurrency
+  // module.
+  ModuleDecl *concurrencyModule = nullptr;
+  for (auto decl: decls) {
+    if (isReplaceable(decl)) {
+      concurrencyModule = decl->getDeclContext()->getParentModule();
+      break;
+    }
+  }
+  if (!concurrencyModule)
+    return;
+
+  // Dig out the name.
+  auto baseName = decls.front()->getName().getBaseName();
+  if (baseName.isSpecial())
+    return;
+
+  // Look for entities with the _unsafeInheritExecutor_ prefix on the name.
+  ASTContext &ctx = decls.front()->getASTContext();
+  Identifier newIdentifier = ctx.getIdentifier(
+      ("_unsafeInheritExecutor_" + baseName.getIdentifier().str()).str());
+
+  NameLookupOptions lookupOptions = defaultUnqualifiedLookupOptions;
+  LookupResult lookup = TypeChecker::lookupUnqualified(
+      const_cast<DeclContext *>(dc), DeclNameRef(newIdentifier), loc,
+      lookupOptions);
+  if (!lookup)
+    return;
+
+  // Drop all of the _Concurrency entries in favor of the ones found by this
+  // lookup.
+  decls.erase(std::remove_if(decls.begin(), decls.end(), [&](ValueDecl *decl) {
+                return isReplaceable(decl);
+              }),
+              decls.end());
+  for (const auto &lookupResult: lookup) {
+    if (auto decl = lookupResult.getValueDecl())
+      decls.push_back(decl);
+  }
+}
+
 /// 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.
 ///
@@ -3845,9 +3901,7 @@ namespace {
         std::tie(diagLoc, inDefaultArgument) = adjustPoundIsolationDiagLoc(
             isolationExpr, getDeclContext()->getParentModule());
 
-        bool inConcurrencyModule = getDeclContext()->getParentModule()->getName()
-            .str().equals("_Concurrency");
-
+        bool inConcurrencyModule = ::inConcurrencyModule(getDeclContext());
         auto diag = ctx.Diags.diagnose(diagLoc,
                                        diag::isolation_in_inherits_executor,
                                        inDefaultArgument);

lib/Sema/TypeCheckConcurrency.h

@@ -656,6 +656,18 @@ AbstractFunctionDecl *enclosingUnsafeInheritsExecutor(const DeclContext *dc);
 void replaceUnsafeInheritExecutorWithDefaultedIsolationParam(
     AbstractFunctionDecl *func, InFlightDiagnostic &diag);
 
+/// Replace any functions in this list that were found in the _Concurrency
+/// module and have _unsafeInheritExecutor_-prefixed versions with those
+/// _unsafeInheritExecutor_-prefixed versions.
+///
+/// This function is an egregious hack that allows us to introduce the
+/// #isolation-based versions of functions into the concurrency library
+/// without breaking clients that use @_unsafeInheritExecutor. Since those
+/// clients can't use #isolation (it doesn't work with @_unsafeInheritExecutor),
+/// we route them to the @_unsafeInheritExecutor versions implicitly.
+void introduceUnsafeInheritExecutorReplacements(
+    const DeclContext *dc, SourceLoc loc, SmallVectorImpl<ValueDecl *> &decls);
+
 } // end namespace swift
 
 namespace llvm {

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -300,11 +300,10 @@ public func withCheckedContinuation<T>(
 }
 
 @available(SwiftStdlib 5.1, *)
-@usableFromInline
 @_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @_unavailableInEmbedded
 @_silgen_name("$ss23withCheckedContinuation8function_xSS_yScCyxs5NeverOGXEtYalF")
-internal func __abi_withCheckedContinuation<T>(
+public func _unsafeInheritExecutor_withCheckedContinuation<T>(
   function: String = #function,
   _ body: (CheckedContinuation<T, Never>) -> Void
 ) async -> T {
@@ -361,11 +360,10 @@ public func withCheckedThrowingContinuation<T>(
 }
 
 @available(SwiftStdlib 5.1, *)
-@usableFromInline
 @_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @_unavailableInEmbedded
 @_silgen_name("$ss31withCheckedThrowingContinuation8function_xSS_yScCyxs5Error_pGXEtYaKlF")
-internal func __abi_withCheckedThrowingContinuation<T>(
+public func _unsafeInheritExecutor_withCheckedThrowingContinuation<T>(
   function: String = #function,
   _ body: (CheckedContinuation<T, Error>) -> Void
 ) async throws -> T {

stdlib/public/Concurrency/PartialAsyncTask.swift

@@ -730,6 +730,32 @@ public func withUnsafeThrowingContinuation<T>(
   }
 }
 
+// HACK: a version of withUnsafeContinuation that uses the unsafe
+// @_unsafeInheritExecutor.
+@available(SwiftStdlib 5.1, *)
+@_alwaysEmitIntoClient
+@_unsafeInheritExecutor
+public func _unsafeInheritExecutor_withUnsafeContinuation<T>(
+  _ fn: (UnsafeContinuation<T, Never>) -> Void
+) async -> T {
+  return await Builtin.withUnsafeContinuation {
+    fn(UnsafeContinuation<T, Never>($0))
+  }
+}
+
+// HACK: a version of withUnsafeThrowingContinuation that uses the unsafe
+// @_unsafeInheritExecutor.
+@available(SwiftStdlib 5.1, *)
+@_alwaysEmitIntoClient
+@_unsafeInheritExecutor
+public func _unsafeInheritExecutor_withUnsafeThrowingContinuation<T>(
+  _ fn: (UnsafeContinuation<T, Error>) -> Void
+) async throws -> sending T {
+  return try await Builtin.withUnsafeThrowingContinuation {
+    fn(UnsafeContinuation<T, Error>($0))
+  }
+}
+
 /// A hack to mark an SDK that supports swift_continuation_await.
 @available(SwiftStdlib 5.1, *)
 @_alwaysEmitIntoClient

stdlib/public/Concurrency/SourceCompatibilityShims.swift

@@ -62,6 +62,17 @@ public func withTaskCancellationHandler<T>(
   try await withTaskCancellationHandler(operation: operation, onCancel: handler)
 }
 
+@available(SwiftStdlib 5.1, *)
+@_alwaysEmitIntoClient
+@_unsafeInheritExecutor
+@available(*, deprecated, renamed: "withTaskCancellationHandler(operation:onCancel:)")
+public func _unsafeInheritExecutor_withTaskCancellationHandler<T>(
+  handler: @Sendable () -> Void,
+  operation: () async throws -> T
+) async rethrows -> T {
+  try await withTaskCancellationHandler(operation: operation, onCancel: handler)
+}
+
 @available(SwiftStdlib 5.1, *)
 extension Task where Success == Never, Failure == Never {
   @available(*, deprecated, message: "`Task.withCancellationHandler` has been replaced by `withTaskCancellationHandler` and will be removed shortly.")

stdlib/public/Concurrency/TaskCancellation.swift

@@ -84,9 +84,8 @@ public func withTaskCancellationHandler<T>(
 
 @_unsafeInheritExecutor // ABI compatibility with Swift 5.1
 @available(SwiftStdlib 5.1, *)
-@usableFromInline
 @_silgen_name("$ss27withTaskCancellationHandler9operation8onCancelxxyYaKXE_yyYbXEtYaKlF")
-internal func __abi_withTaskCancellationHandler<T>(
+public func _unsafeInheritExecutor_withTaskCancellationHandler<T>(
   operation: () async throws -> T,
   onCancel handler: @Sendable () -> Void
 ) async rethrows -> T {

test/Concurrency/unsafe_inherit_executor.swift

@@ -78,3 +78,34 @@ func unsafeCallerAvoidsNewLoop(x: some AsyncSequence<Int, Never>) async throws {
 
   for try await _ in x { }
 }
+
+// -------------------------------------------------------------------------
+// Type checker hack to use _unsafeInheritExecutor_-prefixed versions of
+// some concurrency library functions.
+// -------------------------------------------------------------------------
+
+@_unsafeInheritExecutor
+func unsafeCallerAvoidsNewLoop() async throws {
+  // expected-warning@-1{{@_unsafeInheritExecutor attribute is deprecated; consider an 'isolated' parameter defaulted to '#isolation' instead}}
+
+  _ = await withUnsafeContinuation { (continuation: UnsafeContinuation<Int, Never>) in
+    continuation.resume(returning: 5)
+  }
+
+  _ = try await withUnsafeThrowingContinuation { (continuation: UnsafeContinuation<Int, any Error>) in
+    continuation.resume(returning: 5)
+  }
+
+  _ = await withCheckedContinuation { (continuation: CheckedContinuation<Int, Never>) in
+    continuation.resume(returning: 5)
+  }
+
+  _ = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Int, any Error>) in
+    continuation.resume(returning: 5)
+  }
+
+  _ = await withTaskCancellationHandler {
+    5
+  } onCancel: {
+  }
+}