[Macros] Ensure that we can find peer macros for members of extensions.

When populating the name lookup tables for a nominal type, we were only
visiting declarations with peer macros within the nominal type
declaration itself. Also visit the members of extensions of the
nominal type.

Without this change, peer-macro-defined members of extensions are not
visible.

Author: DougGregor

include/swift/AST/TypeOrExtensionDecl.h

@@ -23,6 +23,9 @@
 
 namespace swift {
 
+class DeclContext;
+class IterableDeclContext;
+
 /// Describes either a nominal type declaration or an extension
 /// declaration.
 struct TypeOrExtensionDecl {
@@ -38,6 +41,8 @@ struct TypeOrExtensionDecl {
   class Decl *getAsDecl() const;
   /// Return the contained *Decl as the DeclContext superclass.
   DeclContext *getAsDeclContext() const;
+  /// Return the contained *Decl as the DeclContext superclass.
+  IterableDeclContext *getAsIterableDeclContext() const;
   /// Return the contained NominalTypeDecl or that of the extended type
   /// in the ExtensionDecl.
   NominalTypeDecl *getBaseNominal() const;

lib/AST/Decl.cpp

@@ -9891,6 +9891,14 @@ Decl *TypeOrExtensionDecl::getAsDecl() const {
 DeclContext *TypeOrExtensionDecl::getAsDeclContext() const {
   return getAsDecl()->getInnermostDeclContext();
 }
+
+IterableDeclContext *TypeOrExtensionDecl::getAsIterableDeclContext() const {
+  if (auto nominal = Decl.dyn_cast<NominalTypeDecl *>())
+    return nominal;
+
+  return Decl.get<ExtensionDecl *>();
+}
+
 NominalTypeDecl *TypeOrExtensionDecl::getBaseNominal() const {
   return getAsDeclContext()->getSelfNominalTypeDecl();
 }

lib/AST/NameLookup.cpp

@@ -1538,10 +1538,12 @@ static void
 populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
                                             MemberLookupTable &table,
                                             DeclName name,
-                                            NominalTypeDecl *dc) {
+                                            TypeOrExtensionDecl container) {
+  auto dc = container.getAsDeclContext();
   auto *moduleScopeCtx = dc->getModuleScopeContext();
-  auto *module = dc->getModuleContext();
-  for (auto *member : dc->getCurrentMembersWithoutLoading()) {
+  auto *module = dc->getParentModule();
+  auto idc = container.getAsIterableDeclContext();
+  for (auto *member : idc->getCurrentMembersWithoutLoading()) {
     // Collect all macro introduced names, along with its corresponding macro
     // reference. We need the macro reference to prevent adding auxiliary decls
     // that weren't introduced by the macro.
@@ -1773,6 +1775,10 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
   if (!Table.isLazilyCompleteForMacroExpansion(macroExpansionKey)) {
     populateLookupTableEntryFromMacroExpansions(
         ctx, Table, macroExpansionKey, decl);
+    for (auto ext : decl->getExtensions()) {
+      populateLookupTableEntryFromMacroExpansions(
+          ctx, Table, macroExpansionKey, ext);
+    }
     Table.markLazilyCompleteForMacroExpansion(macroExpansionKey);
   }
 

test/Macros/macro_expand_peers.swift

@@ -36,12 +36,26 @@ struct S {
   }
 }
 
-// CHECK-DUMP: @__swiftmacro_18macro_expand_peers1f1a3for_SSSi_SSSdtYaF20addCompletionHandlerfMp_.swift
-// CHECK-DUMP: func f(a: Int, for b: String, _ value: Double, completionHandler: @escaping (String) -> Void) {
-// CHECK-DUMP:   Task {
-// CHECK-DUMP:     completionHandler(await f(a: a, for: b, value))
-// CHECK-DUMP:   }
-// CHECK-DUMP: }
+extension S {
+  @addCompletionHandler
+  func g(a: Int, for b: String, _ value: Double) async -> String {
+    return b
+  }
+
+  // CHECK-DUMP: @__swiftmacro_18macro_expand_peers1SV1g1a3for_SSSi_SSSdtYaF20addCompletionHandlerfMp_.swift
+  // CHECK-DUMP: func f(a: Int, for b: String, _ value: Double, completionHandler: @escaping (String) -> Void) {
+  // CHECK-DUMP:   Task {
+  // CHECK-DUMP:     completionHandler(await f(a: a, for: b, value))
+  // CHECK-DUMP:   }
+  // CHECK-DUMP: }
+
+}
+
+func useCompletionHandlerG(s: S, _ body: @escaping (String) -> Void) {
+  s.g(a: 1, for: "hahaha local", 2.0) {
+    body($0)
+  }
+}
 
 @addCompletionHandler
 func f(a: Int, for b: String, _ value: Double) async -> String {