Merge pull request #64854 from DougGregor/top-level-macro-lookup

DougGregor · web-flow · commit f7e479759dbb · 2023-04-03T06:50:39.000-07:00
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
@@ -733,8 +733,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
@@ -906,17 +904,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 {
@@ -1010,6 +1009,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);
 }
@@ -1021,8 +1021,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,
@@ -3980,6 +3981,7 @@ SynthesizedFileUnit::getDiscriminatorForPrivateDecl(const Decl *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
