[Macros] Ensure that name lookup can find uniquely-generated macro names

include/swift/AST/Decl.h

@@ -8476,6 +8476,22 @@ class MacroDecl : public GenericContext, public ValueDecl {
     return DeclName();
   }
 
+  /// Returns a DeclName that acts as a stand-in for all unique names that
+  /// are manufactured by the macro expansion context's `makeUniqueName`.
+  static DeclName getUniqueNamePlaceholder(ASTContext &ctx);
+
+  /// Determine whether the given name is the unique-name placeholder
+  /// produced by `getUniqueNamePlaceholder`.
+  static bool isUniqueNamePlaceholder(DeclName name);
+
+  /// Determine whether the given name is one of the unique names manufactured
+  /// by the macro expansion context's `makeUniqueName`.
+  static bool isUniqueMacroName(StringRef name);
+
+  /// Determine whether the given name is one of the unique names manufactured
+  /// by the macro expansion context's `makeUniqueName`.
+  static bool isUniqueMacroName(DeclBaseName name);
+
   /// Retrieve the definition of this macro.
   MacroDefinition getDefinition() const;
 

lib/AST/Decl.cpp

@@ -10073,6 +10073,39 @@ const MacroRoleAttr *MacroDecl::getMacroRoleAttr(MacroRole role) const {
   return nullptr;
 }
 
+DeclName MacroDecl::getUniqueNamePlaceholder(ASTContext &ctx) {
+  return ctx.getIdentifier("$");
+}
+
+bool MacroDecl::isUniqueNamePlaceholder(DeclName name) {
+  return name.getBaseName().userFacingName() == "$";
+}
+
+bool MacroDecl::isUniqueMacroName(StringRef name) {
+  // Unique macro names are mangled names, which always start with "$s".
+  if (!name.startswith("$s"))
+    return false;
+
+  // Unique macro names end with fMu<digits>_. Match that.
+
+  // Strip off the trailing _.
+  if (name.back() != '_')
+    return false;
+  name = name.drop_back();
+
+  // Strip off trailing digits. This is the discriminator.
+  while (isdigit(name.back()))
+    name = name.drop_back();
+
+  // Check for fMu.
+  return name.endswith("fMu");
+}
+
+bool MacroDecl::isUniqueMacroName(DeclBaseName name) {
+  return isUniqueMacroName(name.userFacingName());
+}
+
+
 void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
                                    SmallVectorImpl<DeclName> &names) const {
   ASTContext &ctx = getASTContext();
@@ -10134,6 +10167,21 @@ void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
       break;
     }
   }
+
+  // Add the unique name, if the macro can introduce declarations anywhere.
+  switch (role) {
+  case MacroRole::Expression:
+  case MacroRole::Declaration:
+  case MacroRole::Member:
+  case MacroRole::Peer:
+    names.push_back(MacroDecl::getUniqueNamePlaceholder(getASTContext()));
+    break;
+
+  case MacroRole::Accessor:
+  case MacroRole::Conformance:
+  case MacroRole::MemberAttribute:
+    break;
+  }
 }
 
 MacroDefinition MacroDecl::getDefinition() const {

lib/AST/Module.cpp

@@ -160,6 +160,7 @@ class swift::SourceLookupCache {
   ValueDeclMap TopLevelValues;
   ValueDeclMap ClassMembers;
   bool MemberCachePopulated = false;
+  DeclName UniqueMacroNamePlaceholder;
 
   template<typename T>
   using OperatorMap = llvm::DenseMap<Identifier, TinyPtrVector<T *>>;
@@ -179,6 +180,8 @@ class swift::SourceLookupCache {
   SmallVector<Decl *, 4> MayHaveAuxiliaryDecls;
   void populateAuxiliaryDeclCache();
 
+  SourceLookupCache(ASTContext &ctx);
+
 public:
   SourceLookupCache(const SourceFile &SF);
   SourceLookupCache(const ModuleDecl &Mod);
@@ -392,15 +395,22 @@ void SourceLookupCache::populateAuxiliaryDeclCache() {
   MayHaveAuxiliaryDecls.clear();
 }
 
+SourceLookupCache::SourceLookupCache(ASTContext &ctx)
+  : UniqueMacroNamePlaceholder(MacroDecl::getUniqueNamePlaceholder(ctx)) { }
+
 /// Populate our cache on the first name lookup.
-SourceLookupCache::SourceLookupCache(const SourceFile &SF) {
+SourceLookupCache::SourceLookupCache(const SourceFile &SF)
+    : SourceLookupCache(SF.getASTContext())
+{
   FrontendStatsTracer tracer(SF.getASTContext().Stats,
                              "source-file-populate-cache");
   addToUnqualifiedLookupCache(SF.getTopLevelDecls(), false);
   addToUnqualifiedLookupCache(SF.getHoistedDecls(), false);
 }
 
-SourceLookupCache::SourceLookupCache(const ModuleDecl &M) {
+SourceLookupCache::SourceLookupCache(const ModuleDecl &M)
+  : SourceLookupCache(M.getASTContext())
+{
   FrontendStatsTracer tracer(M.getASTContext().Stats,
                              "module-populate-cache");
   for (const FileUnit *file : M.getFiles()) {
@@ -428,7 +438,10 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
   // FIXME: We need to not consider auxiliary decls if we're doing lookup
   // from inside a macro argument at module scope.
   populateAuxiliaryDeclCache();
-  auto auxDecls = TopLevelAuxiliaryDecls.find(Name);
+  DeclName keyName = MacroDecl::isUniqueMacroName(Name.getBaseName())
+    ? UniqueMacroNamePlaceholder
+    : Name;
+  auto auxDecls = TopLevelAuxiliaryDecls.find(keyName);
   if (auxDecls == TopLevelAuxiliaryDecls.end())
     return;
 

lib/AST/NameLookup.cpp

@@ -1255,6 +1255,7 @@ class swift::MemberLookupTable : public ASTAllocated<swift::MemberLookupTable> {
   }
 
   bool isLazilyCompleteForMacroExpansion(DeclName name) const {
+    assert(!MacroDecl::isUniqueMacroName(name.getBaseName()));
     // If we've already expanded macros for a simple name, we must have expanded
     // all macros that produce names with the same base identifier.
     bool isBaseNameComplete = name.isCompoundName() &&
@@ -1264,6 +1265,7 @@ class swift::MemberLookupTable : public ASTAllocated<swift::MemberLookupTable> {
   }
 
   void markLazilyCompleteForMacroExpansion(DeclName name) {
+    assert(!MacroDecl::isUniqueMacroName(name.getBaseName()));
     LazilyCompleteNamesForMacroExpansion.insert(name);
   }
 
@@ -1521,6 +1523,17 @@ populateLookupTableEntryFromExtensions(ASTContext &ctx,
   }
 }
 
+/// Adjust the given name to make it a proper key for the lazy macro expansion
+/// cache, which maps all uniquely-generated names down to a single placeholder
+/// key.
+static DeclName adjustLazyMacroExpansionNameKey(
+    ASTContext &ctx, DeclName name) {
+  if (MacroDecl::isUniqueMacroName(name.getBaseName()))
+    return MacroDecl::getUniqueNamePlaceholder(ctx);
+
+  return name;
+}
+
 static void
 populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
                                             MemberLookupTable &table,
@@ -1756,9 +1769,11 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
     Table.markLazilyComplete(baseName);
   }
 
-  if (!Table.isLazilyCompleteForMacroExpansion(name)) {
-    populateLookupTableEntryFromMacroExpansions(ctx, Table, name, decl);
-    Table.markLazilyCompleteForMacroExpansion(name);
+  DeclName macroExpansionKey = adjustLazyMacroExpansionNameKey(ctx, name);
+  if (!Table.isLazilyCompleteForMacroExpansion(macroExpansionKey)) {
+    populateLookupTableEntryFromMacroExpansions(
+        ctx, Table, macroExpansionKey, decl);
+    Table.markLazilyCompleteForMacroExpansion(macroExpansionKey);
   }
 
   // Look for a declaration with this name.

lib/Serialization/Serialization.cpp

@@ -6137,7 +6137,8 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
 
     for (auto D : fileDecls) {
       if (isa<ImportDecl>(D) || isa<IfConfigDecl>(D) ||
-          isa<PoundDiagnosticDecl>(D) || isa<TopLevelCodeDecl>(D)) {
+          isa<PoundDiagnosticDecl>(D) || isa<TopLevelCodeDecl>(D) ||
+          isa<MacroExpansionDecl>(D)) {
         continue;
       }
 

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -456,6 +456,8 @@ public struct AddMembers: MemberMacro {
     providingMembersOf decl: some DeclGroupSyntax,
     in context: some MacroExpansionContext
   ) throws -> [DeclSyntax] {
+    let uniqueClassName = context.createUniqueName("uniqueClass")
+
     let storageStruct: DeclSyntax =
       """
       struct Storage {}
@@ -481,7 +483,12 @@ public struct AddMembers: MemberMacro {
 
     let initDecl: DeclSyntax =
       """
-      init() {}
+      init() { _ = \(uniqueClassName)() }
+      """
+
+    let classDecl: DeclSyntax =
+      """
+      class \(uniqueClassName) { }
       """
 
     return [
@@ -490,6 +497,7 @@ public struct AddMembers: MemberMacro {
       instanceMethod,
       staticMethod,
       initDecl,
+      classDecl,
     ]
   }
 }

test/Macros/macro_expand_synthesized_members.swift

@@ -2,7 +2,7 @@
 // RUN: %target-build-swift -swift-version 5 -I %swift-host-lib-dir -L %swift-host-lib-dir -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath
 // RUNx: %target-swift-frontend -dump-ast -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir %s -module-name MacroUser 2>&1 | %FileCheck --check-prefix CHECK-AST %s
 // RUN: %target-typecheck-verify-swift -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -module-name MacroUser -DTEST_DIAGNOSTICS -swift-version 5
-// RUN: %target-build-swift -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -L %swift-host-lib-dir %s -o %t/main -module-name MacroUser -swift-version 5
+// RUN: %target-build-swift -g -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -L %swift-host-lib-dir %s -o %t/main -module-name MacroUser -swift-version 5
 // RUN: %target-run %t/main | %FileCheck %s
 // REQUIRES: executable_test
 

test/Macros/top_level_freestanding.swift

@@ -2,7 +2,7 @@
 // RUN: %target-build-swift -swift-version 5 -parse-as-library -enable-experimental-feature FreestandingMacros -I %swift-host-lib-dir -L %swift-host-lib-dir -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath
 // RUNx: %target-swift-frontend -enable-experimental-feature FreestandingMacros -parse-as-library -emit-sil -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir %s -module-name MacroUser 2>&1 | %FileCheck --check-prefix CHECK-SIL %s
 // RUN: %target-typecheck-verify-swift -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -module-name MacroUser -DTEST_DIAGNOSTICS -swift-version 5
-// RUN: %target-build-swift -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -L %swift-host-lib-dir %s -o %t/main -module-name MacroUser -swift-version 5
+// RUN: %target-build-swift -g -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -L %swift-host-lib-dir %s -o %t/main -module-name MacroUser -swift-version 5
 // RUN: %target-run %t/main | %FileCheck %s
 // REQUIRES: executable_test