[Clang importer + macros] Handle name lookup and type checking for expanded macros

Introduce a number of fixes to allow us to fully use declarations that
are produced by applying a peer macro to an imported declarations.
These changes include:
* Ensuring that we have the right set of imports in the source file
containing the macro expansion, because it depends only on the module
it comes from
* Ensuring that name lookup looks in that file even when the
DeclContext hierarchy doesn't contain the source file (because it's
based on the Clang module structure)

Expand testing to be sure that we're getting the right calls,
diagnostics, and generated IR symbols.

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -808,6 +808,8 @@ class AbstractFunctionDeclScope final : public ASTScopeImpl {
 protected:
   void printSpecifics(llvm::raw_ostream &out) const override;
 
+  bool lookupLocalsOrMembers(DeclConsumer) const override;
+
 public:
   Decl *getDecl() const { return decl; }
 

include/swift/AST/NameLookup.h

@@ -528,6 +528,15 @@ void lookupVisibleMemberDecls(VisibleDeclConsumer &Consumer,
                               bool includeProtocolExtensionMembers,
                               GenericSignature genericSig = GenericSignature());
 
+/// Determine the module-scope context from which lookup should proceed.
+///
+/// In the common case, module-scope context is the source file in which
+/// the declaration context is nested. However, when declaration context is
+/// part of an imported Clang declaration context, it won't be nested within a
+/// source file. Rather, the source file will be on the side, and will be
+/// provided here because it contains information about the available imports.
+DeclContext *getModuleScopeLookupContext(DeclContext *dc);
+
 namespace namelookup {
 
 /// Add semantic members to \p type before attempting a semantic lookup.

include/swift/Subsystems.h

@@ -120,6 +120,12 @@ namespace swift {
   /// This walks the AST to resolve imports.
   void performImportResolution(SourceFile &SF);
 
+  /// Resolve imports for a source file generated to adapt a given
+  /// Clang module.
+  void performImportResolutionForClangMacroBuffer(
+    SourceFile &SF, ModuleDecl *clangModule
+  );
+
   /// Once type-checking is complete, this instruments code with calls to an
   /// intrinsic that record the expected values of local variables so they can
   /// be compared against the results from the debugger.

lib/AST/ASTScopeLookup.cpp

@@ -30,6 +30,7 @@
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/ClangImporter/ClangModule.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/Support/Compiler.h"
 
@@ -264,6 +265,20 @@ bool ASTScopeImpl::lookupLocalsOrMembers(DeclConsumer) const {
   return false; // many kinds of scopes have none
 }
 
+bool AbstractFunctionDeclScope::lookupLocalsOrMembers(DeclConsumer consumer) const {
+  // Special case: if we're within a function inside a type context, but the
+  // parent context is within a Clang module unit, we need to make sure to
+  // look for members in it.
+  auto dc = decl->getDeclContext();
+  if (!dc->isTypeContext())
+    return false;
+
+  if (!isa<ClangModuleUnit>(dc->getModuleScopeContext()))
+    return false;
+
+  return consumer.lookInMembers(cast<GenericContext>(dc->getAsDecl()));
+}
+
 bool GenericTypeOrExtensionScope::lookupLocalsOrMembers(
     ASTScopeImpl::DeclConsumer consumer) const {
   return portion->lookupMembersOf(this, consumer);

lib/AST/AvailabilityScope.cpp

@@ -55,8 +55,10 @@ AvailabilityScope::createForSourceFile(SourceFile *SF,
   case SourceFileKind::DefaultArgument: {
     // Look up the parent context in the enclosing file that this file's
     // root context should be nested under.
-    if (auto parentScope =
-            SF->getEnclosingSourceFile()->getAvailabilityScope()) {
+    auto enclosingSF = SF->getEnclosingSourceFile();
+    if (!enclosingSF)
+      break;
+    if (auto parentScope = enclosingSF->getAvailabilityScope()) {
       auto charRange = Ctx.SourceMgr.getRangeForBuffer(SF->getBufferID());
       range = SourceRange(charRange.getStart(), charRange.getEnd());
       auto originalNode = SF->getNodeInEnclosingSourceFile();

lib/AST/NameLookup.cpp

@@ -1698,7 +1698,7 @@ SmallVector<MacroDecl *, 1> namelookup::lookupMacros(DeclContext *dc,
                                                      DeclNameRef macroName,
                                                      MacroRoles roles) {
   SmallVector<MacroDecl *, 1> choices;
-  auto moduleScopeDC = dc->getModuleScopeContext();
+  auto moduleScopeDC = getModuleScopeLookupContext(dc);
   ASTContext &ctx = moduleScopeDC->getASTContext();
 
   auto addChoiceIfApplicable = [&](ValueDecl *decl) {

lib/AST/UnqualifiedLookup.cpp

@@ -30,6 +30,7 @@
 #include "swift/Basic/STLExtras.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Basic/Statistic.h"
+#include "swift/ClangImporter/ClangModule.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TinyPtrVector.h"
@@ -265,7 +266,8 @@ void UnqualifiedLookupFactory::performUnqualifiedLookup() {
   if (!isFirstResultEnough()) {
     // If no result has been found yet, the dependency must be on a top-level
     // name, since up to now, the search has been for non-top-level names.
-    auto *moduleScopeContext = DC->getModuleScopeContext();
+    auto *moduleScopeContext = getModuleScopeLookupContext(DC);
+
     lookUpTopLevelNamesInModuleScopeContext(moduleScopeContext);
   }
 }
@@ -939,3 +941,19 @@ ValueDecl *ASTScope::lookupSingleLocalDecl(SourceFile *sf, DeclName name,
     return nullptr;
   return result[0];
 }
+
+DeclContext *swift::getModuleScopeLookupContext(DeclContext *dc) {
+  auto moduleScopeContext = dc->getModuleScopeContext();
+
+  // When the module scope context is in a Clang module but we actually
+  // have a parent source file, it's because we are within a macro
+  // expansion triggered for the imported declaration. In such cases,
+  // we want to use the parent source file as the lookup context, becauae
+  // it has the appropriate resolved imports.
+  if (isa<ClangModuleUnit>(moduleScopeContext)) {
+    if (auto parentSF = dc->getParentSourceFile())
+      return parentSF;
+  }
+
+  return moduleScopeContext;
+}

lib/ClangImporter/ClangImporter.cpp

@@ -2798,7 +2798,11 @@ ClangModuleUnit *ClangImporter::Implementation::getWrapperForModule(
   Identifier name = underlying->Name == "std"
                         ? SwiftContext.Id_CxxStdlib
                         : SwiftContext.getIdentifier(underlying->Name);
-  auto wrapper = ModuleDecl::create(name, SwiftContext);
+  ImplicitImportInfo implicitImportInfo;
+  if (auto mainModule = SwiftContext.MainModule) {
+    implicitImportInfo = mainModule->getImplicitImportInfo();
+  }
+  auto wrapper = ModuleDecl::create(name, SwiftContext, implicitImportInfo);
   wrapper->setIsSystemModule(underlying->IsSystem);
   wrapper->setIsNonSwiftModule();
   wrapper->setHasResolvedImports();

lib/Parse/ParseRequests.cpp

@@ -114,11 +114,23 @@ ParseAbstractFunctionBodyRequest::evaluate(Evaluator &evaluator,
 
   case BodyKind::Unparsed: {
     // FIXME: How do we configure code completion?
-    SourceFile &sf = *afd->getDeclContext()->getParentSourceFile();
-    SourceManager &sourceMgr = sf.getASTContext().SourceMgr;
+    SourceManager &sourceMgr = afd->getASTContext().SourceMgr;
     unsigned bufferID =
         sourceMgr.findBufferContainingLoc(afd->getBodySourceRange().Start);
-    Parser parser(bufferID, sf, /*SIL*/ nullptr);
+    SourceFile *sf = afd->getDeclContext()->getParentSourceFile();
+    if (!sf) {
+      auto sourceFiles = sourceMgr.getSourceFilesForBufferID(bufferID);
+      auto expectedModule = afd->getParentModule();
+      for (auto checkSF: sourceFiles) {
+        if (checkSF->getParentModule() == expectedModule) {
+          sf = checkSF;
+          break;
+        }
+      }
+      assert(sf && "Could not find source file containing parsed body");
+    }
+
+    Parser parser(bufferID, *sf, /*SIL*/ nullptr);
     auto result = parser.parseAbstractFunctionBodyDelayed(afd);
     afd->setBodyKind(BodyKind::Parsed);
     return result;

lib/Sema/ImportResolution.cpp

@@ -192,6 +192,11 @@ class ImportResolver final : public DeclVisitor<ImportResolver> {
 
   void addImplicitImports();
 
+  void addImplicitImport(ModuleDecl *module) {
+    boundImports.push_back(ImportedModule(module));
+    bindPendingImports();
+  }
+
   /// Retrieve the finalized imports.
   ArrayRef<AttributedImport<ImportedModule>> getFinishedImports() const {
     return boundImports;
@@ -301,6 +306,22 @@ void swift::performImportResolution(SourceFile &SF) {
   verify(SF);
 }
 
+void swift::performImportResolutionForClangMacroBuffer(
+    SourceFile &SF, ModuleDecl *clangModule
+) {
+  // If we've already performed import resolution, bail.
+  if (SF.ASTStage == SourceFile::ImportsResolved)
+    return;
+
+  ImportResolver resolver(SF);
+  resolver.addImplicitImport(clangModule);
+
+  SF.setImports(resolver.getFinishedImports());
+  SF.setImportedUnderlyingModule(resolver.getUnderlyingClangModule());
+
+  SF.ASTStage = SourceFile::ImportsResolved;
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: Import handling generally
 //===----------------------------------------------------------------------===//

lib/Sema/TypeCheckDecl.cpp

@@ -52,9 +52,11 @@
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Bridging/ASTGen.h"
+#include "swift/ClangImporter/ClangModule.h"
 #include "swift/Sema/IDETypeChecking.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
 #include "swift/Strings.h"
+#include "swift/Subsystems.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/APSInt.h"

lib/Sema/TypeCheckMacros.cpp

@@ -1033,8 +1033,11 @@ createMacroSourceFile(std::unique_ptr<llvm::MemoryBuffer> buffer,
       /*parsingOpts=*/{}, /*isPrimary=*/false);
   if (auto parentSourceFile = dc->getParentSourceFile())
     macroSourceFile->setImports(parentSourceFile->getImports());
-  else if (isa<ClangModuleUnit>(dc->getModuleScopeContext()))
-    macroSourceFile->setImports({});
+  else if (auto clangModuleUnit =
+               dyn_cast<ClangModuleUnit>(dc->getModuleScopeContext())) {
+    auto clangModule = clangModuleUnit->getParentModule();
+    performImportResolutionForClangMacroBuffer(*macroSourceFile, clangModule);
+  }
   return macroSourceFile;
 }
 

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

@@ -1,15 +1,17 @@
 #if __SWIFT_ATTR_SUPPORTS_MACROS
 #define ADD_ASYNC __attribute__((swift_attr("@macro_library.AddAsync")))
+#define ADD_ASYNC_FINAL __attribute__((swift_attr("@macro_library.AddAsyncFinal")))
 #else
 #define ADD_ASYNC
+#define ADD_ASYNC_FINAL
 #endif
 
-void async_divide(double x, double y, void (* _Nonnull completionHandler)(double x)) ADD_ASYNC;
+void async_divide(double x, double y, void (^ _Nonnull completionHandler)(double x)) ADD_ASYNC;
 
 typedef struct SlowComputer {
 } SlowComputer;
 
-void computer_divide(const SlowComputer *computer, double x, double y, void (* _Nonnull completionHandler)(double x))
+void computer_divide(const SlowComputer *computer, double x, double y, void (^ _Nonnull completionHandler)(double x))
   ADD_ASYNC
   __attribute__((swift_name("SlowComputer.divide(self:_:_:completionHandler:)")));
 
@@ -19,7 +21,7 @@ void computer_divide(const SlowComputer *computer, double x, double y, void (* _
 
 @interface Computer: NSObject
 -(void)multiply:(double)x by:(double)y afterDone:(void (^ _Nonnull)(double x))afterDone
-  ADD_ASYNC
+  ADD_ASYNC_FINAL
   __attribute__((swift_async(none)));
 @end
 #endif

test/Macros/Inputs/macro_library.swift

@@ -56,3 +56,6 @@ public macro declareVarValuePeerShadowed() = #externalMacro(module: "MacroDefini
 
 @attached(peer, names: overloaded)
 public macro AddAsync() = #externalMacro(module: "MacroDefinition", type: "AddAsyncMacro")
+
+@attached(peer, names: overloaded)
+public macro AddAsyncFinal() = #externalMacro(module: "MacroDefinition", type: "AddAsyncMacro")

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -935,7 +935,6 @@ public struct AddAsyncMacro: PeerMacro {
     providingPeersOf declaration: Declaration,
     in context: Context
   ) throws -> [DeclSyntax] {
-
     // Only on functions at the moment.
     guard var funcDecl = declaration.as(FunctionDeclSyntax.self) else {
       throw CustomError.message("@addAsync only works on functions")
@@ -1084,6 +1083,37 @@ public struct AddAsyncMacro: PeerMacro {
 
     funcDecl.attributes = newAttributeList
 
+    // If this declaration needs to be final, make it so.
+    let isFinal = node.attributeName.trimmedDescription.contains("AddAsyncFinal")
+    if isFinal {
+      var allModifiers = Array(funcDecl.modifiers)
+      if let openIndex = allModifiers.firstIndex(where: { $0.name.text == "open" }) {
+        allModifiers[openIndex].name = .keyword(.public)
+      } else {
+        allModifiers.append(
+          DeclModifierSyntax(
+            name: .keyword(.public),
+            trailingTrivia: .space
+          )
+        )
+      }
+
+      allModifiers.append(
+        DeclModifierSyntax(
+          name: .keyword(.final),
+          trailingTrivia: .space
+        )
+      )
+
+      funcDecl.modifiers = DeclModifierListSyntax(allModifiers)
+
+      var allAttributes = Array(funcDecl.attributes)
+        allAttributes.append(
+          .attribute("@_alwaysEmitIntoClient ")
+        )
+      funcDecl.attributes = AttributeListSyntax(allAttributes)
+    }
+
     funcDecl.leadingTrivia = .newlines(2)
 
     return [DeclSyntax(funcDecl)]

test/Macros/expand_on_imported.swift

@@ -10,20 +10,17 @@
 // Diagnostics testing
 // RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -typecheck -verify -swift-version 5 -enable-experimental-feature MacrosOnImports -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name ModuleUser %s -I %t
 
+// Emit IR to ensure that we are handling the created body end-to-end.
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -emit-ir -swift-version 5 -g -enable-experimental-feature MacrosOnImports -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name ModuleUser %s -I %t -validate-tbd-against-ir=none | %FileCheck %s
+
 import CompletionHandlerGlobals
 import macro_library
 
-// Make sure that @AddAsync works at all.
-@AddAsync
-@available(SwiftStdlib 5.1, *)
-func asyncTest(_ value: Int, completionHandler: @escaping (String) -> Void) {
-  completionHandler(String(value))
-}
-
 @available(SwiftStdlib 5.1, *)
 func testAll(x: Double, y: Double, computer: SlowComputer) async {
-  _ = await asyncTest(17)
-
   let _: Double = await async_divide(1.0, 2.0)
   let _: Double = await computer.divide(x, y)
 }
+
+// CHECK: define swifttailcc void @"$sSC12async_divideyS2d_SdtYaF"
+// CHECK: define swifttailcc void @"$sSo12SlowComputerV6divideyS2d_SdtYaF"

test/Macros/expand_on_imported_objc.swift

@@ -8,7 +8,10 @@
 // RUN: %target-swift-frontend -swift-version 5 -emit-module -o %t/macro_library.swiftmodule %S/Inputs/macro_library.swift -module-name macro_library -load-plugin-library %t/%target-library-name(MacroDefinition)
 
 // Diagnostics testing
-// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -typecheck -verify -swift-version 5 -enable-experimental-feature MacrosOnImports -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name ModuleUser %s -I %t 
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -typecheck -verify -swift-version 5 -enable-experimental-feature MacrosOnImports -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name ModuleUser %s -I %t -DTEST_DIAGNOSTICS
+
+// Emit IR just to make sure nothing else fails.
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -emit-ir -swift-version 5 -g -enable-experimental-feature MacrosOnImports -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name ModuleUser %s -I %t | %FileCheck %s
 
 import CompletionHandlerGlobals
 import macro_library
@@ -17,6 +20,10 @@ import macro_library
 func testAll(x: Double, y: Double, computer: Computer, untypedComputer: AnyObject) async {
   let _: Double = await computer.multiply(x, by: y)
 
+  #if TEST_DIAGNOSTICS
   // expected-error@+1{{missing argument for parameter 'afterDone' in call}}
   untypedComputer.multiply(x, by: y)
+  #endif
 }
+
+// CHECK: define linkonce_odr hidden swifttailcc void @"$sSo8ComputerC8multiply_2byS2d_SdtYaF"

test/Macros/print_clang_expand_on_imported.swift

@@ -14,7 +14,7 @@
 
 import CompletionHandlerGlobals
 
-// CHECK: func async_divide(_ x: Double, _ y: Double, _ completionHandler: @convention(c) (Double) -> Void)
+// CHECK: func async_divide(_ x: Double, _ y: Double, _ completionHandler: @escaping (Double) -> Void)
 
 // CHECK: extension SlowComputer
 // CHECK: public func divide(_ x: Double, _ y: Double) async -> Double

test/Macros/print_clang_expand_on_imported_objc.swift

@@ -13,5 +13,4 @@
 import CompletionHandlerGlobals
 
 // CHECK: class Computer
-// FIXME: The "open" is odd here. We want this to be "final", but can't yet.
-// CHECK: open func multiply(_ x: Double, by y: Double) async -> Double
+// CHECK: @_alwaysEmitIntoClient public final func multiply(_ x: Double, by y: Double) async -> Double