Extend type checker hack for @_unsafeInheritExecutor functions to methods

With the re-introduction of `@_unsafeInheritExecutor` for `TaskLocal.withValue`,
we need to extend the type checker trick with `_unsafeInheritExecutor_`-prefixed
functions to work with methods. Do so to make `TaskLocal.withValue` actually
work this way.

Author: DougGregor

lib/Sema/ConstraintSystem.cpp

@@ -292,6 +292,13 @@ LookupResult &ConstraintSystem::lookupMember(Type base, DeclNameRef name,
   result = TypeChecker::lookupMember(DC, base, name, loc,
                                      defaultMemberLookupOptions);
 
+  // If we are in an @_unsafeInheritExecutor context, swap out
+  // declarations for their _unsafeInheritExecutor_ counterparts if they
+  // exist.
+  if (enclosingUnsafeInheritsExecutor(DC)) {
+    introduceUnsafeInheritExecutorReplacements(DC, base, loc, *result);
+  }
+
   // If we aren't performing dynamic lookup, we're done.
   if (!*result || !base->isAnyObject())
     return *result;

lib/Sema/TypeCheckConcurrency.cpp

@@ -2115,17 +2115,18 @@ void swift::introduceUnsafeInheritExecutorReplacements(
   // Make sure at least some of the entries are functions in the _Concurrency
   // module.
   ModuleDecl *concurrencyModule = nullptr;
+  DeclBaseName baseName;
   for (auto decl: decls) {
     if (isReplaceable(decl)) {
       concurrencyModule = decl->getDeclContext()->getParentModule();
+      baseName = decl->getName().getBaseName();
       break;
     }
   }
   if (!concurrencyModule)
     return;
 
-  // Dig out the name.
-  auto baseName = decls.front()->getName().getBaseName();
+  // Ignore anything with a special name.
   if (baseName.isSpecial())
     return;
 
@@ -2153,6 +2154,60 @@ void swift::introduceUnsafeInheritExecutorReplacements(
   }
 }
 
+void swift::introduceUnsafeInheritExecutorReplacements(
+    const DeclContext *dc, Type base, SourceLoc loc, LookupResult &lookup) {
+  if (lookup.empty())
+    return;
+
+  auto baseNominal = base->getAnyNominal();
+  if (!baseNominal || !inConcurrencyModule(baseNominal))
+    return;
+
+  auto isReplaceable = [&](ValueDecl *decl) {
+    return isa<FuncDecl>(decl) && inConcurrencyModule(decl->getDeclContext());
+  };
+
+  // Make sure at least some of the entries are functions in the _Concurrency
+  // module.
+  ModuleDecl *concurrencyModule = nullptr;
+  DeclBaseName baseName;
+  for (auto &result: lookup) {
+    auto decl = result.getValueDecl();
+    if (isReplaceable(decl)) {
+      concurrencyModule = decl->getDeclContext()->getParentModule();
+      baseName = decl->getBaseName();
+      break;
+    }
+  }
+  if (!concurrencyModule)
+    return;
+
+  // Ignore anything with a special name.
+  if (baseName.isSpecial())
+    return;
+
+  // Look for entities with the _unsafeInheritExecutor_ prefix on the name.
+  ASTContext &ctx = base->getASTContext();
+  Identifier newIdentifier = ctx.getIdentifier(
+      ("_unsafeInheritExecutor_" + baseName.getIdentifier().str()).str());
+
+  LookupResult replacementLookup = TypeChecker::lookupMember(
+      const_cast<DeclContext *>(dc), base, DeclNameRef(newIdentifier), loc,
+      defaultMemberLookupOptions);
+  if (replacementLookup.innerResults().empty())
+    return;
+
+  // Drop all of the _Concurrency entries in favor of the ones found by this
+  // lookup.
+  lookup.filter([&](const LookupResultEntry &entry, bool) {
+    return !isReplaceable(entry.getValueDecl());
+  });
+
+  for (const auto &entry: replacementLookup.innerResults()) {
+    lookup.add(entry, /*isOuter=*/false);
+  }
+}
+
 /// 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.
 ///

lib/Sema/TypeCheckConcurrency.h

@@ -44,6 +44,7 @@ class EnumElementDecl;
 class Expr;
 class FuncDecl;
 class Initializer;
+class LookupResult;
 class PatternBindingDecl;
 class ProtocolConformance;
 class TopLevelCodeDecl;
@@ -668,6 +669,18 @@ void replaceUnsafeInheritExecutorWithDefaultedIsolationParam(
 void introduceUnsafeInheritExecutorReplacements(
     const DeclContext *dc, SourceLoc loc, SmallVectorImpl<ValueDecl *> &decls);
 
+/// Replace any functions in this list that were found in the _Concurrency
+/// module as a member on "base" 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, Type base, SourceLoc loc, LookupResult &result);
+
 } // end namespace swift
 
 namespace llvm {

test/Concurrency/unsafe_inherit_executor.swift

@@ -113,7 +113,10 @@ func unsafeCallerAvoidsNewLoop() async throws {
   } onCancel: {
   }
 
-  TL.$string.withValue("hello") {
+  await TL.$string.withValue("hello") {
     print(TL.string)
   }
+
+  func operation() async throws -> Int { 7 }
+  try await TL.$string.withValue("hello", operation: operation)
 }