[Macros] Resolve macro names using unqualified lookup that ignores expansions

DougGregor · DougGregor · commit 828de17b0098 · 2023-04-02T23:15:38.000-07:00
The macro name resolution in the source lookup cache was only looking at
macros in the current module, meaning that any names introduced by peer
or declaration macros declared in one module but used in another would
not be found by name lookup.

Switch the source lookup cache over to using the same
`forEachPotentialResolvedMacro` API that is used by lookup within
types, so we have consistent name-lookup-level macro resolution in both
places.

... except that would be horribly cyclic, of course, so introduce name
lookup flags to ignore top-level declarations introduced by macro
expansions. This is semantically correct because macro expansions are
not allowed to introduce new macros anyway, because that would have
been a terrible idea.

Fixes rdar://107321469. Peer and declaration macros at module scope
should work a whole lot better now.
diff --git a/include/swift/AST/FileUnit.h b/include/swift/AST/FileUnit.h
@@ -65,6 +65,7 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
   ///
   /// This does a simple local lookup, not recursively looking through imports.
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl*> &result) const = 0;
 
   /// Look up a local type declaration by its mangled name.
@@ -381,6 +382,7 @@ class BuiltinUnit final : public FileUnit {
   explicit BuiltinUnit(ModuleDecl &M);
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl*> &result) const override;
 
   /// Find all Objective-C methods with the given selector.
diff --git a/include/swift/AST/LookupKinds.h b/include/swift/AST/LookupKinds.h
@@ -53,6 +53,9 @@ enum NLOptions : unsigned {
   // Include @usableFromInline and @inlinable
   NL_IncludeUsableFromInline = 1 << 6,
 
+  /// Exclude names introduced by macro expansions in the top-level module.
+  NL_ExcludeMacroExpansions = 1 << 7,
+
   /// The default set of options used for qualified name lookup.
   ///
   /// FIXME: Eventually, add NL_ProtocolMembers to this, once all of the
@@ -81,6 +84,13 @@ static inline NLOptions operator~(NLOptions value) {
 
 void simple_display(llvm::raw_ostream &out, NLOptions options);
 
+
+/// Flags affecting module-level lookup.
+enum class ModuleLookupFlags : unsigned {
+  /// Exclude names introduced by macro expansions in the top-level module.
+  ExcludeMacroExpansions = 1 << 0,
+};
+
 } // end namespace swift
 
 #endif
diff --git a/include/swift/AST/Module.h b/include/swift/AST/Module.h
@@ -728,8 +728,19 @@ class ModuleDecl
   ///
   /// This does a simple local lookup, not recursively looking through imports.
   void lookupValue(DeclName Name, NLKind LookupKind,
+                   OptionSet<ModuleLookupFlags> Flags,
                    SmallVectorImpl<ValueDecl*> &Result) const;
 
+  /// Look up a (possibly overloaded) value set at top-level scope
+  /// (but with the specified access path, which may come from an import decl)
+  /// within the current module.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  void lookupValue(DeclName Name, NLKind LookupKind,
+                   SmallVectorImpl<ValueDecl*> &Result) const {
+    lookupValue(Name, LookupKind, {}, Result);
+  }
+
   /// Look up a local type declaration by its mangled name.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
diff --git a/include/swift/AST/NameLookup.h b/include/swift/AST/NameLookup.h
@@ -240,6 +240,8 @@ enum class UnqualifiedLookupFlags {
   // This lookup should include results that are @inlinable or
   // @usableFromInline.
   IncludeUsableFromInline = 1 << 5,
+  /// This lookup should exclude any names introduced by macro expansions.
+  ExcludeMacroExpansions = 1 << 6,
 };
 
 using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;
@@ -543,6 +545,17 @@ template <typename Result>
 void filterForDiscriminator(SmallVectorImpl<Result> &results,
                             DebuggerClient *debugClient);
 
+/// 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, and should only be used when resolving macro
+/// names for the purposes of (other) name lookup.
+void forEachPotentialResolvedMacro(
+    DeclContext *moduleScopeCtx, DeclNameRef macroName, MacroRole role,
+    llvm::function_ref<void(MacroDecl *, const MacroRoleAttr *)> body
+);
+
 } // end namespace namelookup
 
 /// Describes an inherited nominal entry.
diff --git a/include/swift/AST/SourceFile.h b/include/swift/AST/SourceFile.h
@@ -442,6 +442,7 @@ class SourceFile final : public FileUnit {
   const SmallVectorImpl<ValueDecl *> &getCachedVisibleDecls() const;
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl*> &result) const override;
 
   virtual void lookupVisibleDecls(ImportPath::Access accessPath,
diff --git a/include/swift/AST/SynthesizedFileUnit.h b/include/swift/AST/SynthesizedFileUnit.h
@@ -44,6 +44,7 @@ class SynthesizedFileUnit final : public FileUnit {
   void addTopLevelDecl(Decl *D) { TopLevelDecls.push_back(D); }
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl *> &result) const override;
 
   void lookupObjCMethods(
diff --git a/include/swift/ClangImporter/ClangModule.h b/include/swift/ClangImporter/ClangModule.h
@@ -66,6 +66,7 @@ class ClangModuleUnit final : public LoadedFile {
   virtual bool isSystemModule() const override;
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl*> &results) const override;
 
   virtual TypeDecl *
diff --git a/include/swift/Serialization/SerializedModuleLoader.h b/include/swift/Serialization/SerializedModuleLoader.h
@@ -360,6 +360,7 @@ class SerializedASTFile final : public LoadedFile {
   virtual bool isSystemModule() const override;
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           OptionSet<ModuleLookupFlags> Flags,
                            SmallVectorImpl<ValueDecl*> &results) const override;
 
   virtual StringRef
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -1152,7 +1152,9 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   if (!M)
     return nullptr;
   M->lookupValue(getIdentifier(getProtocolName(kind)),
-                 NLKind::UnqualifiedLookup, results);
+                 NLKind::UnqualifiedLookup,
+                 ModuleLookupFlags::ExcludeMacroExpansions,
+                 results);
 
   for (auto result : results) {
     if (auto protocol = dyn_cast<ProtocolDecl>(result)) {
@@ -6307,7 +6309,7 @@ Type ASTContext::getNamedSwiftType(ModuleDecl *module, StringRef name) {
       if (auto module = clangUnit->getOverlayModule())
         module->lookupValue(identifier, NLKind::UnqualifiedLookup, results);
     } else {
-      file->lookupValue(identifier, NLKind::UnqualifiedLookup, results);
+      file->lookupValue(identifier, NLKind::UnqualifiedLookup, { }, results);
     }
   }
 
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -172,9 +172,6 @@ class swift::SourceLookupCache {
   template<typename Range>
   void addToMemberCache(Range decls);
 
-  using MacroDeclMap = llvm::DenseMap<Identifier, TinyPtrVector<MacroDecl *>>;
-  MacroDeclMap MacroDecls;
-
   using AuxiliaryDeclMap = llvm::DenseMap<DeclName, TinyPtrVector<MissingDecl *>>;
   AuxiliaryDeclMap TopLevelAuxiliaryDecls;
   SmallVector<Decl *, 4> MayHaveAuxiliaryDecls;
@@ -190,6 +187,7 @@ class swift::SourceLookupCache {
   void invalidate();
 
   void lookupValue(DeclName Name, NLKind LookupKind,
+                   OptionSet<ModuleLookupFlags> Flags,
                    SmallVectorImpl<ValueDecl*> &Result);
 
   /// Retrieves all the operator decls. The order of the results is not
@@ -263,6 +261,10 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
       // a malformed context.
       if (ED->isInvalid()) continue;
 
+      if (ED->getAttrs().hasAttribute<CustomAttr>()) {
+        MayHaveAuxiliaryDecls.push_back(ED);
+      }
+
       if (!ED->hasUnparsedMembers() || ED->maybeHasOperatorDeclarations())
         addToUnqualifiedLookupCache(ED->getMembers(), true);
     }
@@ -273,9 +275,6 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
     else if (auto *PG = dyn_cast<PrecedenceGroupDecl>(D))
       PrecedenceGroups[PG->getName()].push_back(PG);
 
-    else if (auto *macro = dyn_cast<MacroDecl>(D))
-      MacroDecls[macro->getBaseIdentifier()].push_back(macro);
-
     else if (auto *MED = dyn_cast<MacroExpansionDecl>(D))
       MayHaveAuxiliaryDecls.push_back(MED);
   }
@@ -353,32 +352,26 @@ void SourceLookupCache::populateAuxiliaryDeclCache() {
     for (auto attrConst : decl->getAttrs().getAttributes<CustomAttr>()) {
       auto *attr = const_cast<CustomAttr *>(attrConst);
       UnresolvedMacroReference macroRef(attr);
-      auto macroName = macroRef.getMacroName().getBaseIdentifier();
-
-      auto found = MacroDecls.find(macroName);
-      if (found == MacroDecls.end())
-        continue;
-
-      for (const auto *macro : found->second) {
-        macro->getIntroducedNames(MacroRole::Peer,
-                                  dyn_cast<ValueDecl>(decl),
-                                  introducedNames[attr]);
-      }
+      namelookup::forEachPotentialResolvedMacro(
+          decl->getDeclContext()->getModuleScopeContext(),
+          macroRef.getMacroName(), MacroRole::Peer,
+          [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+            macro->getIntroducedNames(MacroRole::Peer,
+                                      dyn_cast<ValueDecl>(decl),
+                                      introducedNames[attr]);
+          });
     }
 
     if (auto *med = dyn_cast<MacroExpansionDecl>(decl)) {
       UnresolvedMacroReference macroRef(med);
-      auto macroName = macroRef.getMacroName().getBaseIdentifier();
-
-      auto found = MacroDecls.find(macroName);
-      if (found == MacroDecls.end())
-        continue;
-
-      for (const auto *macro : found->second) {
-        macro->getIntroducedNames(MacroRole::Declaration,
-                                  /*attachedTo*/ nullptr,
-                                  introducedNames[med]);
-      }
+      namelookup::forEachPotentialResolvedMacro(
+          decl->getDeclContext()->getModuleScopeContext(),
+          macroRef.getMacroName(), MacroRole::Declaration,
+          [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+            macro->getIntroducedNames(MacroRole::Declaration,
+                                      /*attachedTo*/ nullptr,
+                                      introducedNames[med]);
+          });
     }
 
     // Add macro-introduced names to the top-level auxiliary decl cache as
@@ -425,6 +418,7 @@ SourceLookupCache::SourceLookupCache(const ModuleDecl &M)
 }
 
 void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
+                                    OptionSet<ModuleLookupFlags> Flags,
                                     SmallVectorImpl<ValueDecl*> &Result) {
   auto I = TopLevelValues.find(Name);
   if (I != TopLevelValues.end()) {
@@ -433,6 +427,10 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
       Result.push_back(Elt);
   }
 
+  // If we aren't supposed to find names introduced by macros, we're done.
+  if (Flags.contains(ModuleLookupFlags::ExcludeMacroExpansions))
+    return;
+
   // Add top-level auxiliary decls to the result.
   //
   // FIXME: We need to not consider auxiliary decls if we're doing lookup
@@ -863,17 +861,18 @@ static bool isParsedModule(const ModuleDecl *mod) {
 }
 
 void ModuleDecl::lookupValue(DeclName Name, NLKind LookupKind,
+                             OptionSet<ModuleLookupFlags> Flags,
                              SmallVectorImpl<ValueDecl*> &Result) const {
   auto *stats = getASTContext().Stats;
   if (stats)
     ++stats->getFrontendCounters().NumModuleLookupValue;
 
   if (isParsedModule(this)) {
-    getSourceLookupCache().lookupValue(Name, LookupKind, Result);
+    getSourceLookupCache().lookupValue(Name, LookupKind, Flags, Result);
     return;
   }
 
-  FORWARD(lookupValue, (Name, LookupKind, Result));
+  FORWARD(lookupValue, (Name, LookupKind, Flags, Result));
 }
 
 TypeDecl * ModuleDecl::lookupLocalType(StringRef MangledName) const {
@@ -967,6 +966,7 @@ void ModuleDecl::lookupImportedSPIGroups(
 }
 
 void BuiltinUnit::lookupValue(DeclName name, NLKind lookupKind,
+                              OptionSet<ModuleLookupFlags> Flags,
                               SmallVectorImpl<ValueDecl*> &result) const {
   getCache().lookupValue(name.getBaseIdentifier(), lookupKind, *this, result);
 }
@@ -978,8 +978,9 @@ void BuiltinUnit::lookupObjCMethods(
 }
 
 void SourceFile::lookupValue(DeclName name, NLKind lookupKind,
+                             OptionSet<ModuleLookupFlags> flags,
                              SmallVectorImpl<ValueDecl*> &result) const {
-  getCache().lookupValue(name, lookupKind, result);
+  getCache().lookupValue(name, lookupKind, flags, result);
 }
 
 void ModuleDecl::lookupVisibleDecls(ImportPath::Access AccessPath,
@@ -3922,6 +3923,7 @@ SynthesizedFileUnit::getDiscriminatorForPrivateValue(const ValueDecl *D) const {
 
 void SynthesizedFileUnit::lookupValue(
     DeclName name, NLKind lookupKind,
+    OptionSet<ModuleLookupFlags> Flags,
     SmallVectorImpl<ValueDecl *> &result) const {
   for (auto *decl : TopLevelDecls) {
     if (auto VD = dyn_cast<ValueDecl>(decl)) {
diff --git a/lib/AST/ModuleNameLookup.cpp b/lib/AST/ModuleNameLookup.cpp
@@ -57,6 +57,10 @@ class ModuleNameLookup {
                       NLOptions options);
 };
 
+// Exclude names introduced by macro expansions.
+enum class LocalLookupFlags {
+  ExcludeMacroExpansions,
+};
 
 /// Encapsulates the work done for a recursive qualified lookup into a module
 /// by full name.
@@ -81,12 +85,13 @@ class LookupByName : public ModuleNameLookup<LookupByName> {
   }
 
   void doLocalLookup(ModuleDecl *module, ImportPath::Access path,
+                     OptionSet<ModuleLookupFlags> flags,
                      SmallVectorImpl<ValueDecl *> &localDecls) {
     // If this import is specific to some named decl ("import Swift.Int")
     // then filter out any lookups that don't match.
     if (!path.matches(name))
       return;
-    module->lookupValue(name, lookupKind, localDecls);
+    module->lookupValue(name, lookupKind, flags, localDecls);
   }
 };
 
@@ -112,6 +117,7 @@ class LookupVisibleDecls : public ModuleNameLookup<LookupVisibleDecls> {
   }
 
   void doLocalLookup(ModuleDecl *module, ImportPath::Access path,
+                     OptionSet<ModuleLookupFlags> flags,
                      SmallVectorImpl<ValueDecl *> &localDecls) {
     VectorDeclConsumer consumer(localDecls);
     module->lookupVisibleDecls(path, consumer, lookupKind);
@@ -167,9 +173,14 @@ void ModuleNameLookup<LookupStrategy>::lookupInModule(
     currentCount = decls.size();
   };
 
+  OptionSet<ModuleLookupFlags> currentModuleLookupFlags = {};
+  if (options & NL_ExcludeMacroExpansions)
+    currentModuleLookupFlags |= ModuleLookupFlags::ExcludeMacroExpansions;
+
   // Do the lookup into the current module.
   auto *module = moduleOrFile->getParentModule();
-  getDerived()->doLocalLookup(module, accessPath, decls);
+  getDerived()->doLocalLookup(
+      module, accessPath, currentModuleLookupFlags, decls);
   updateNewDecls(moduleScopeContext);
 
   bool canReturnEarly = (initialCount != decls.size() &&
@@ -198,7 +209,7 @@ void ModuleNameLookup<LookupStrategy>::lookupInModule(
         return;
 
       getDerived()->doLocalLookup(import.importedModule, import.accessPath,
-                                  decls);
+                                  { }, decls);
       updateNewDecls(moduleScopeContext);
     };
 
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -1539,12 +1539,16 @@ static DeclName adjustLazyMacroExpansionNameKey(
 ///
 /// This routine intentionally avoids calling `forEachAttachedMacro`, which
 /// triggers request cycles.
-static void forEachPotentialResolvedMacro(
+void namelookup::forEachPotentialResolvedMacro(
     DeclContext *moduleScopeCtx, DeclNameRef macroName, MacroRole role,
     llvm::function_ref<void(MacroDecl *, const MacroRoleAttr *)> body
 ) {
   ASTContext &ctx = moduleScopeCtx->getASTContext();
-  UnqualifiedLookupDescriptor lookupDesc{macroName, moduleScopeCtx};
+  UnqualifiedLookupDescriptor lookupDesc{
+    macroName, moduleScopeCtx, SourceLoc(),
+    UnqualifiedLookupFlags::ExcludeMacroExpansions
+  };
+
   auto lookup = evaluateOrDefault(
       ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
   for (auto result : lookup.allResults()) {
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -520,6 +520,8 @@ void UnqualifiedLookupFactory::addImportedResults(const DeclContext *const dc) {
   auto nlOptions = NL_UnqualifiedDefault;
   if (options.contains(Flags::IncludeUsableFromInline))
     nlOptions |= NL_IncludeUsableFromInline;
+  if (options.contains(Flags::ExcludeMacroExpansions))
+    nlOptions |= NL_ExcludeMacroExpansions;
   lookupInModule(dc, Name.getFullName(), CurModuleResults,
                  NLKind::UnqualifiedLookup, resolutionKind, dc, nlOptions);
 
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -3344,6 +3344,7 @@ void ClangModuleUnit::getDisplayDecls(SmallVectorImpl<Decl*> &results, bool recu
 }
 
 void ClangModuleUnit::lookupValue(DeclName name, NLKind lookupKind,
+                                  OptionSet<ModuleLookupFlags> flags,
                                   SmallVectorImpl<ValueDecl*> &results) const {
   // FIXME: Ignore submodules, which are empty for now.
   if (clangModule && clangModule->isSubModule())
diff --git a/lib/ClangImporter/DWARFImporter.cpp b/lib/ClangImporter/DWARFImporter.cpp
@@ -33,6 +33,7 @@ class DWARFModuleUnit final : public LoadedFile {
   /// DWARFimporterDelegate.
   virtual void
   lookupValue(DeclName name, NLKind lookupKind,
+              OptionSet<ModuleLookupFlags> Flags,
               SmallVectorImpl<ValueDecl *> &results) const override {
     Owner.lookupValueDWARF(name, lookupKind,
                            getParentModule()->getName(), results);
diff --git a/lib/Serialization/SerializedModuleLoader.cpp b/lib/Serialization/SerializedModuleLoader.cpp
diff --git a/test/Macros/Inputs/macro_library.swift b/test/Macros/Inputs/macro_library.swift
diff --git a/test/Macros/macro_expand_peers.swift b/test/Macros/macro_expand_peers.swift
