From 491448cb48e7977fde1891fd52d228ec57b860c6 Mon Sep 17 00:00:00 2001
From: Holly Borla <hborla@apple.com>
Date: Mon, 5 Feb 2024 19:38:43 -0800
Subject: [PATCH] [Macros] Expand nested macros in qualified name lookup.

---
 lib/AST/NameLookup.cpp                        | 47 ++++++++++++-------
 .../Inputs/syntax_macro_definitions.swift     | 29 ++++++++++++
 .../macro_expand_synthesized_members.swift    | 16 +++++++
 3 files changed, 75 insertions(+), 17 deletions(-)

diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
index 6177852144e67..cc168d84808e2 100644
--- a/lib/AST/NameLookup.cpp
+++ b/lib/AST/NameLookup.cpp
@@ -1903,24 +1903,37 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
     // 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.
-    MacroIntroducedNameTracker nameTracker;
-    if (auto *med = dyn_cast<MacroExpansionDecl>(member)) {
-      forEachPotentialResolvedMacro(
-          dc->getModuleScopeContext(), med->getMacroName(),
-          MacroRole::Declaration, nameTracker);
-    } else if (auto *vd = dyn_cast<ValueDecl>(member)) {
-      nameTracker.attachedTo = dyn_cast<ValueDecl>(member);
-      forEachPotentialAttachedMacro(member, MacroRole::Peer, nameTracker);
-    }
 
-    // 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);
-      });
+    std::deque<Decl *> mightIntroduceNames;
+    mightIntroduceNames.push_back(member);
+
+    while (!mightIntroduceNames.empty()) {
+      auto *member = mightIntroduceNames.front();
+
+      MacroIntroducedNameTracker nameTracker;
+      if (auto *med = dyn_cast<MacroExpansionDecl>(member)) {
+        forEachPotentialResolvedMacro(
+            dc->getModuleScopeContext(), med->getMacroName(),
+            MacroRole::Declaration, nameTracker);
+      } else if (auto *vd = dyn_cast<ValueDecl>(member)) {
+        nameTracker.attachedTo = dyn_cast<ValueDecl>(member);
+        forEachPotentialAttachedMacro(member, MacroRole::Peer, nameTracker);
+      }
+
+      // 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;
+
+          mightIntroduceNames.push_back(decl);
+          table.addMember(decl);
+        });
+      }
+
+      mightIntroduceNames.pop_front();
     }
   }
 }
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
index fb7bd5f71c3a9..af5d4e7286d72 100644
--- a/test/Macros/Inputs/syntax_macro_definitions.swift
+++ b/test/Macros/Inputs/syntax_macro_definitions.swift
@@ -2091,6 +2091,35 @@ public struct SendableMacro: ExtensionMacro {
   }
 }
 
+public struct GenerateStubMemberMacro: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf declaration: some DeclGroupSyntax,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["#generateMemberStubs"]
+  }
+}
+
+public struct GenerateStubsFreestandingMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["#generateMember"]
+  }
+}
+
+public struct SingleMemberStubMacro: DeclarationMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["static func member() {}"]
+  }
+}
+
 public struct FakeCodeItemMacro: DeclarationMacro, PeerMacro {
   public static func expansion(
     of node: some FreestandingMacroExpansionSyntax,
diff --git a/test/Macros/macro_expand_synthesized_members.swift b/test/Macros/macro_expand_synthesized_members.swift
index 1121baff1499d..eba4113e96180 100644
--- a/test/Macros/macro_expand_synthesized_members.swift
+++ b/test/Macros/macro_expand_synthesized_members.swift
@@ -141,3 +141,19 @@ func testC2() {
   // CHECK: deinit was called
 }
 testC2()
+
+@attached(member, names: arbitrary)
+macro GenerateStubs() = #externalMacro(module: "MacroDefinition", type: "GenerateStubMemberMacro")
+
+@freestanding(declaration, names: arbitrary)
+macro generateMemberStubs() = #externalMacro(module: "MacroDefinition", type: "GenerateStubsFreestandingMacro")
+
+@freestanding(declaration, names: named(member()))
+macro generateMember() = #externalMacro(module: "MacroDefinition", type: "SingleMemberStubMacro")
+
+@GenerateStubs
+struct NestedMacroExpansion {}
+
+func callNestedExpansionMember() {
+  NestedMacroExpansion.member()
+}