Fix lookup into macros that can introduce members into types and extensions

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

@@ -1534,75 +1534,129 @@ static DeclName adjustLazyMacroExpansionNameKey(
   return name;
 }
 
+/// Call the given function body with each macro declaration and its associated
+/// role attribute for the given role.
+///
+/// This routine intentionally avoids calling `forEachAttachedMacro`, which
+/// triggers request cycles.
+static void forEachPotentialResolvedMacro(
+    DeclContext *moduleScopeCtx, DeclNameRef macroName, MacroRole role,
+    llvm::function_ref<void(MacroDecl *, const MacroRoleAttr *)> body
+) {
+  ASTContext &ctx = moduleScopeCtx->getASTContext();
+  UnqualifiedLookupDescriptor lookupDesc{macroName, moduleScopeCtx};
+  auto lookup = evaluateOrDefault(
+      ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
+  for (auto result : lookup.allResults()) {
+    auto *vd = result.getValueDecl();
+    auto *macro = dyn_cast<MacroDecl>(vd);
+    if (!macro)
+      continue;
+
+    auto *macroRoleAttr = macro->getMacroRoleAttr(role);
+    if (!macroRoleAttr)
+      continue;
+
+    body(macro, macroRoleAttr);
+  }
+}
+
+/// For each macro with the given role that might be attached to the given
+/// declaration, call the body.
+static void forEachPotentialAttachedMacro(
+    Decl *decl, MacroRole role,
+    llvm::function_ref<void(MacroDecl *macro, const MacroRoleAttr *)> body
+) {
+  // We intentionally avoid calling `forEachAttachedMacro` in order to avoid
+  // a request cycle.
+  auto moduleScopeCtx = decl->getDeclContext()->getModuleScopeContext();
+  for (auto attrConst : decl->getSemanticAttrs().getAttributes<CustomAttr>()) {
+    auto *attr = const_cast<CustomAttr *>(attrConst);
+    UnresolvedMacroReference macroRef(attr);
+    auto macroName = macroRef.getMacroName();
+    forEachPotentialResolvedMacro(moduleScopeCtx, macroName, role, body);
+  }
+}
+
+namespace {
+  /// Function object that tracks macro-introduced names.
+  struct MacroIntroducedNameTracker {
+    ValueDecl *attachedTo = nullptr;
+
+    llvm::SmallSet<DeclName, 4> allIntroducedNames;
+    bool introducesArbitraryNames = false;
+
+    /// Augment the set of names with those introduced by the given macro.
+    void operator()(MacroDecl *macro, const MacroRoleAttr *attr) {
+      // First check for arbitrary names.
+      if (attr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
+        introducesArbitraryNames = true;
+      }
+
+      // If this introduces arbitrary names, there's nothing more to do.
+      if (introducesArbitraryNames)
+        return;
+
+      SmallVector<DeclName, 4> introducedNames;
+      macro->getIntroducedNames(
+          attr->getMacroRole(), attachedTo, introducedNames);
+      for (auto name : introducedNames)
+        allIntroducedNames.insert(name.getBaseName());
+    }
+
+    bool shouldExpandForName(DeclName name) const {
+      return introducesArbitraryNames ||
+          allIntroducedNames.contains(name.getBaseName());
+    }
+  };
+}
+
 static void
 populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
                                             MemberLookupTable &table,
                                             DeclName name,
-                                            NominalTypeDecl *dc) {
-  auto *moduleScopeCtx = dc->getModuleScopeContext();
-  auto *module = dc->getModuleContext();
-  for (auto *member : dc->getCurrentMembersWithoutLoading()) {
+                                            TypeOrExtensionDecl container) {
+
+  // Trigger the expansion of member macros on the container, if any of the
+  // names match.
+  {
+    MacroIntroducedNameTracker nameTracker;
+    auto decl = container.getAsDecl();
+    forEachPotentialAttachedMacro(decl, MacroRole::Member, nameTracker);
+    if (nameTracker.shouldExpandForName(name)) {
+      (void)evaluateOrDefault(
+          ctx.evaluator,
+          ExpandSynthesizedMemberMacroRequest{decl},
+          false);
+    }
+  }
+
+  auto dc = container.getAsDeclContext();
+  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.
-    llvm::SmallSet<DeclName, 4> allIntroducedNames;
-    bool introducesArbitraryNames = false;
+    MacroIntroducedNameTracker nameTracker;
     if (auto *med = dyn_cast<MacroExpansionDecl>(member)) {
-      auto declRef = evaluateOrDefault(
-          ctx.evaluator, ResolveMacroRequest{med, dc},
-          nullptr);
-      if (!declRef)
-        continue;
-      auto *macro = dyn_cast<MacroDecl>(declRef.getDecl());
-      if (macro->getMacroRoleAttr(MacroRole::Declaration)
-          ->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary))
-        introducesArbitraryNames = true;
-      else {
-        SmallVector<DeclName, 4> introducedNames;
-        macro->getIntroducedNames(MacroRole::Declaration, nullptr,
-                                  introducedNames);
-        for (auto name : introducedNames)
-          allIntroducedNames.insert(name);
-      }
+      forEachPotentialResolvedMacro(
+          member->getModuleContext(), med->getMacroName(),
+          MacroRole::Declaration, nameTracker);
     } else if (auto *vd = dyn_cast<ValueDecl>(member)) {
-      // We intentionally avoid calling `forEachAttachedMacro` in order to avoid
-      // a request cycle.
-      for (auto attrConst : member->getSemanticAttrs().getAttributes<CustomAttr>()) {
-        auto *attr = const_cast<CustomAttr *>(attrConst);
-        UnresolvedMacroReference macroRef(attr);
-        auto macroName = macroRef.getMacroName();
-        UnqualifiedLookupDescriptor lookupDesc{macroName, moduleScopeCtx};
-        auto lookup = evaluateOrDefault(
-            ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
-        for (auto result : lookup.allResults()) {
-          auto *vd = result.getValueDecl();
-          auto *macro = dyn_cast<MacroDecl>(vd);
-          if (!macro)
-            continue;
-          auto *macroRoleAttr = macro->getMacroRoleAttr(MacroRole::Peer);
-          if (!macroRoleAttr)
-            continue;
-          if (macroRoleAttr->hasNameKind(
-                  MacroIntroducedDeclNameKind::Arbitrary))
-            introducesArbitraryNames = true;
-          else {
-            SmallVector<DeclName, 4> introducedNames;
-            macro->getIntroducedNames(
-                MacroRole::Peer, dyn_cast<ValueDecl>(member), introducedNames);
-            for (auto name : introducedNames)
-              allIntroducedNames.insert(name);
-          }
-        }
-      }
+      nameTracker.attachedTo = dyn_cast<ValueDecl>(member);
+      forEachPotentialAttachedMacro(member, MacroRole::Peer, nameTracker);
     }
-    // Expand macros based on the name.
-    if (introducesArbitraryNames || allIntroducedNames.contains(name))
+
+    // Expand macros on this member.
+    if (nameTracker.shouldExpandForName(name)) {
       member->visitAuxiliaryDecls([&](Decl *decl) {
         auto *sf = module->getSourceFileContainingLocation(decl->getLoc());
         // Bail out if the auxiliary decl was not produced by a macro.
         if (!sf || sf->Kind != SourceFileKind::MacroExpansion) return;
         table.addMember(decl);
       });
+    }
   }
 }
 
@@ -1773,6 +1827,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);
   }
 
@@ -2134,12 +2192,6 @@ QualifiedLookupRequest::evaluate(Evaluator &eval, const DeclContext *DC,
     // Make sure we've resolved property wrappers, if we need them.
     installPropertyWrapperMembersIfNeeded(current, member);
 
-    // Expand synthesized member macros.
-    auto &ctx = current->getASTContext();
-    (void)evaluateOrDefault(ctx.evaluator,
-                            ExpandSynthesizedMemberMacroRequest{current},
-                            false);
-
     // Look for results within the current nominal type and its extensions.
     bool currentIsProtocol = isa<ProtocolDecl>(current);
     auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();

lib/Sema/LookupVisibleDecls.cpp

@@ -253,6 +253,21 @@ static void doGlobalExtensionLookup(Type BaseType,
 
     synthesizePropertyWrapperVariables(extension);
 
+    // Expand member macros.
+    ASTContext &ctx = nominal->getASTContext();
+    (void)evaluateOrDefault(
+        ctx.evaluator,
+        ExpandSynthesizedMemberMacroRequest{extension},
+        false);
+
+    // Expand peer macros.
+    for (auto *member : extension->getMembers()) {
+      (void)evaluateOrDefault(
+          ctx.evaluator,
+          ExpandPeerMacroRequest{member},
+          {});
+    }
+
     collectVisibleMemberDecls(CurrDC, LS, BaseType, extension, FoundDecls);
   }
 

lib/Sema/TypeCheckMacros.cpp

@@ -1258,7 +1258,7 @@ evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo, CustomAttr *attr,
     if (auto nominal = dyn_cast<NominalTypeDecl>(attachedTo)) {
       rightBraceLoc = nominal->getBraces().End;
     } else {
-      auto ext = cast<ExtensionDecl>(parentDecl);
+      auto ext = cast<ExtensionDecl>(attachedTo);
       rightBraceLoc = ext->getBraces().End;
     }
 

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -554,6 +554,37 @@ public struct AddMembers: MemberMacro {
   }
 }
 
+public struct AddExtMembers: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf decl: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    let uniqueClassName = context.createUniqueName("uniqueClass")
+
+    let instanceMethod: DeclSyntax =
+      """
+      func extInstanceMethod() {}
+      """
+
+    let staticMethod: DeclSyntax =
+      """
+      static func extStaticMethod() {}
+      """
+
+    let classDecl: DeclSyntax =
+      """
+      class \(uniqueClassName) { }
+      """
+
+    return [
+      instanceMethod,
+      staticMethod,
+      classDecl,
+    ]
+  }
+}
+
 public struct AddArbitraryMembers: MemberMacro {
   public static func expansion(
     of node: AttributeSyntax,

test/Macros/macro_expand.swift

@@ -335,6 +335,12 @@ func testFreestandingMacroExpansion() {
 }
 testFreestandingMacroExpansion()
 
+// Explicit structs to force macros to be parsed as decl.
+struct ContainerOfNumberedStructs {
+  #bitwidthNumberedStructs("MyIntOne")
+  #bitwidthNumberedStructs("MyIntTwo")
+}
+
 // Avoid re-type-checking declaration macro arguments.
 @freestanding(declaration)
 macro freestandingWithClosure<T>(_ value: T, body: (T) -> T) = #externalMacro(module: "MacroDefinition", type: "EmptyDeclarationMacro")

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 {

test/Macros/macro_expand_synthesized_members.swift

@@ -23,12 +23,25 @@ struct S {
   }
 }
 
+@attached(
+  member,
+  names: named(extInstanceMethod), named(extStaticMethod)
+)
+macro addExtMembers() = #externalMacro(module: "MacroDefinition", type: "AddExtMembers")
+
+@addExtMembers
+extension S { }
+
 let s = S()
 
 // CHECK: synthesized method
 // CHECK: Storage
 s.useSynthesized()
 
+// Members added via extension.
+s.extInstanceMethod()
+S.extStaticMethod()
+
 @attached(member, names: arbitrary)
 macro addArbitraryMembers() = #externalMacro(module: "MacroDefinition", type: "AddArbitraryMembers")
 

test/Macros/top_level_freestanding.swift

@@ -46,3 +46,8 @@ macro freestandingWithClosure<T>(_ value: T, body: (T) -> T) = #externalMacro(mo
   let x = $0
   return x
 }
+
+struct HasInnerClosure {
+  #freestandingWithClosure(0) { x in x }
+  #freestandingWithClosure(1) { x in x }
+}