[Module-scope lookup] Find arbitrary names in macro expansions.

include/swift/AST/Decl.h

@@ -3835,6 +3835,8 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
     /// Whether to include @_implements members.
     /// Used by conformance-checking to find special @_implements members.
     IncludeAttrImplements = 1 << 0,
+    /// Whether to exclude members of macro expansions.
+    ExcludeMacroExpansions = 1 << 1,
   };
 
   /// Find all of the declarations with the given name within this nominal type

include/swift/AST/Evaluator.h

@@ -208,6 +208,18 @@ class Evaluator {
   /// is treated as a stack and is used to detect cycles.
   llvm::SetVector<ActiveRequest> activeRequests;
 
+  /// How many `ResolveMacroRequest` requests are active.
+  ///
+  /// This allows us to quickly determine whether there is any
+  /// `ResolveMacroRequest` active in the active request stack.
+  /// It saves us from a linear scan through `activeRequests` when
+  /// we need to determine this information.
+  ///
+  /// Why on earth would we need to determine this information?
+  /// Please see the extended comment that goes with the constructor
+  /// of `UnqualifiedLookupRequest`.
+  unsigned numActiveResolveMacroRequests = 0;
+
   /// A cache that stores the results of requests.
   evaluator::RequestCache cache;
 
@@ -324,6 +336,16 @@ class Evaluator {
     return activeRequests.count(ActiveRequest(request));
   }
 
+  /// Determine whether there is any active "resolve macro" request
+  /// on the request stack.
+  ///
+  /// Why on earth would we need to determine this information?
+  /// Please see the extended comment that goes with the constructor
+  /// of `UnqualifiedLookupRequest`.
+  bool hasActiveResolveMacroRequest() const {
+    return numActiveResolveMacroRequests > 0;
+  }
+
 private:
   /// Diagnose a cycle detected in the evaluation of the given
   /// request.
@@ -337,6 +359,10 @@ class Evaluator {
   /// request to the \c activeRequests stack.
   bool checkDependency(const ActiveRequest &request);
 
+  /// Note that we have finished this request, popping it from the
+  /// \c activeRequests stack.
+  void finishedRequest(const ActiveRequest &request);
+
   /// Produce the result of the request without caching.
   template<typename Request>
   llvm::Expected<typename Request::OutputType>
@@ -366,8 +392,7 @@ class Evaluator {
 
     // Make sure we remove this from the set of active requests once we're
     // done.
-    assert(activeRequests.back() == activeReq);
-    activeRequests.pop_back();
+    finishedRequest(activeReq);
 
     return std::move(result);
   }

include/swift/AST/NameLookupRequests.h

@@ -430,7 +430,7 @@ class UnqualifiedLookupRequest
                            RequestFlags::Uncached |
                                RequestFlags::DependencySink> {
 public:
-  using SimpleRequest::SimpleRequest;
+  UnqualifiedLookupRequest(UnqualifiedLookupDescriptor);
 
 private:
   friend SimpleRequest;
@@ -456,7 +456,10 @@ class LookupInModuleRequest
                                 NLOptions),
           RequestFlags::Uncached | RequestFlags::DependencySink> {
 public:
-  using SimpleRequest::SimpleRequest;
+  LookupInModuleRequest(
+      const DeclContext *, DeclName, NLKind,
+      namelookup::ResolutionKind, const DeclContext *,
+      NLOptions);
 
 private:
   friend SimpleRequest;
@@ -504,7 +507,9 @@ class ModuleQualifiedLookupRequest
                            RequestFlags::Uncached |
                               RequestFlags::DependencySink> {
 public:
-  using SimpleRequest::SimpleRequest;
+  ModuleQualifiedLookupRequest(const DeclContext *,
+                               ModuleDecl *, DeclNameRef,
+                               NLOptions);
 
 private:
   friend SimpleRequest;
@@ -528,7 +533,9 @@ class QualifiedLookupRequest
                                                  DeclNameRef, NLOptions),
                            RequestFlags::Uncached> {
 public:
-  using SimpleRequest::SimpleRequest;
+  QualifiedLookupRequest(const DeclContext *,
+                         SmallVector<NominalTypeDecl *, 4>,
+                         DeclNameRef, NLOptions);
 
 private:
   friend SimpleRequest;
@@ -579,7 +586,7 @@ class DirectLookupRequest
                            TinyPtrVector<ValueDecl *>(DirectLookupDescriptor),
                            RequestFlags::Uncached|RequestFlags::DependencySink> {
 public:
-  using SimpleRequest::SimpleRequest;
+  DirectLookupRequest(DirectLookupDescriptor);
 
 private:
   friend SimpleRequest;

lib/AST/Evaluator.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/Evaluator.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/DiagnosticEngine.h"
+#include "swift/AST/TypeCheckRequests.h" // for ResolveMacroRequest
 #include "swift/Basic/LangOptions.h"
 #include "swift/Basic/Range.h"
 #include "swift/Basic/SourceManager.h"
@@ -61,14 +62,25 @@ Evaluator::Evaluator(DiagnosticEngine &diags, const LangOptions &opts)
 
 bool Evaluator::checkDependency(const ActiveRequest &request) {
   // Record this as an active request.
-  if (activeRequests.insert(request))
+  if (activeRequests.insert(request)) {
+    if (request.getAs<ResolveMacroRequest>())
+      ++numActiveResolveMacroRequests;
     return false;
+  }
 
   // Diagnose cycle.
   diagnoseCycle(request);
   return true;
 }
 
+void Evaluator::finishedRequest(const ActiveRequest &request) {
+  if (request.getAs<ResolveMacroRequest>())
+    --numActiveResolveMacroRequests;
+
+  assert(activeRequests.back() == request);
+  activeRequests.pop_back();
+}
+
 void Evaluator::diagnoseCycle(const ActiveRequest &request) {
   if (debugDumpCycles) {
     const auto printIndent = [](llvm::raw_ostream &OS, unsigned indent) {

lib/AST/Module.cpp

@@ -174,6 +174,10 @@ class swift::SourceLookupCache {
 
   using AuxiliaryDeclMap = llvm::DenseMap<DeclName, TinyPtrVector<MissingDecl *>>;
   AuxiliaryDeclMap TopLevelAuxiliaryDecls;
+
+  /// Top-level macros that produce arbitrary names.
+  SmallVector<MissingDecl *, 4> TopLevelArbitraryMacros;
+
   SmallVector<Decl *, 4> MayHaveAuxiliaryDecls;
   void populateAuxiliaryDeclCache();
 
@@ -352,26 +356,46 @@ void SourceLookupCache::populateAuxiliaryDeclCache() {
     for (auto attrConst : decl->getAttrs().getAttributes<CustomAttr>()) {
       auto *attr = const_cast<CustomAttr *>(attrConst);
       UnresolvedMacroReference macroRef(attr);
+      bool introducesArbitraryNames = false;
       namelookup::forEachPotentialResolvedMacro(
           decl->getDeclContext()->getModuleScopeContext(),
           macroRef.getMacroName(), MacroRole::Peer,
           [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+            // First check for arbitrary names.
+            if (roleAttr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
+              introducesArbitraryNames = true;
+            }
+
             macro->getIntroducedNames(MacroRole::Peer,
                                       dyn_cast<ValueDecl>(decl),
                                       introducedNames[attr]);
           });
+
+      // Record this macro where appropriate.
+      if (introducesArbitraryNames)
+        TopLevelArbitraryMacros.push_back(MissingDecl::forUnexpandedMacro(attr, decl));
     }
 
     if (auto *med = dyn_cast<MacroExpansionDecl>(decl)) {
       UnresolvedMacroReference macroRef(med);
+      bool introducesArbitraryNames = false;
       namelookup::forEachPotentialResolvedMacro(
           decl->getDeclContext()->getModuleScopeContext(),
           macroRef.getMacroName(), MacroRole::Declaration,
           [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+            // First check for arbitrary names.
+            if (roleAttr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
+              introducesArbitraryNames = true;
+            }
+
             macro->getIntroducedNames(MacroRole::Declaration,
                                       /*attachedTo*/ nullptr,
                                       introducedNames[med]);
           });
+
+      // Record this macro where appropriate.
+      if (introducesArbitraryNames)
+        TopLevelArbitraryMacros.push_back(MissingDecl::forUnexpandedMacro(med, decl));
     }
 
     // Add macro-introduced names to the top-level auxiliary decl cache as
@@ -440,15 +464,30 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
     ? UniqueMacroNamePlaceholder
     : Name;
   auto auxDecls = TopLevelAuxiliaryDecls.find(keyName);
-  if (auxDecls == TopLevelAuxiliaryDecls.end())
+
+  // Check macro expansions that could produce this name.
+  SmallVector<MissingDecl *, 4> unexpandedDecls;
+  if (auxDecls != TopLevelAuxiliaryDecls.end()) {
+    unexpandedDecls.insert(
+      unexpandedDecls.end(), auxDecls->second.begin(), auxDecls->second.end());
+  }
+
+  // Check macro expansions that can produce arbitrary names.
+  unexpandedDecls.insert(
+      unexpandedDecls.end(),
+      TopLevelArbitraryMacros.begin(), TopLevelArbitraryMacros.end());
+
+  if (unexpandedDecls.empty())
     return;
 
-  for (auto *unexpandedDecl : auxDecls->second) {
-    // Add expanded peers and freestanding declarations to the result.
+  // Add matching expanded peers and freestanding declarations to the results.
+  SmallPtrSet<ValueDecl *, 4> macroExpandedDecls;
+  for (auto *unexpandedDecl : unexpandedDecls) {
     unexpandedDecl->forEachMacroExpandedDecl(
         [&](ValueDecl *decl) {
           if (decl->getName().matchesRef(Name)) {
-            Result.push_back(decl);
+            if (macroExpandedDecls.insert(decl).second)
+              Result.push_back(decl);
           }
         });
   }

lib/AST/NameLookup.cpp

@@ -1713,13 +1713,24 @@ void NominalTypeDecl::prepareLookupTable() {
 }
 
 static TinyPtrVector<ValueDecl *>
-maybeFilterOutAttrImplements(TinyPtrVector<ValueDecl *> decls,
-                             DeclName name,
-                             bool includeAttrImplements) {
-  if (includeAttrImplements)
+maybeFilterOutUnwantedDecls(TinyPtrVector<ValueDecl *> decls,
+                            DeclName name,
+                            bool includeAttrImplements,
+                            bool excludeMacroExpansions) {
+  if (includeAttrImplements && !excludeMacroExpansions)
     return decls;
   TinyPtrVector<ValueDecl*> result;
   for (auto V : decls) {
+    // If we're supposed to exclude anything that comes from a macro expansion,
+    // check whether the source location of the declaration is in a macro
+    // expansion, and skip this declaration if it does.
+    if (excludeMacroExpansions) {
+      auto sourceFile =
+          V->getModuleContext()->getSourceFileContainingLocation(V->getLoc());
+      if (sourceFile && sourceFile->Kind == SourceFileKind::MacroExpansion)
+        continue;
+    }
+
     // Filter-out any decl that doesn't have the name we're looking for
     // (asserting as a consistency-check that such entries all have
     // @_implements attrs for the name!)
@@ -1755,12 +1766,16 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
                                           decl->hasLazyMembers());
   const bool includeAttrImplements =
       flags.contains(NominalTypeDecl::LookupDirectFlags::IncludeAttrImplements);
+  const bool excludeMacroExpansions =
+      flags.contains(NominalTypeDecl::LookupDirectFlags::ExcludeMacroExpansions);
 
   LLVM_DEBUG(llvm::dbgs() << decl->getNameStr() << ".lookupDirect("
                           << name << ")"
                           << ", hasLazyMembers()=" << decl->hasLazyMembers()
                           << ", useNamedLazyMemberLoading="
                           << useNamedLazyMemberLoading
+                          << ", excludeMacroExpansions="
+                          << excludeMacroExpansions
                           << "\n");
 
   decl->prepareLookupTable();
@@ -1797,8 +1812,9 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
       if (!allFound.empty()) {
         auto known = Table.find(name);
         if (known != Table.end()) {
-          auto swiftLookupResult = maybeFilterOutAttrImplements(
-              known->second, name, includeAttrImplements);
+          auto swiftLookupResult = maybeFilterOutUnwantedDecls(
+              known->second, name, includeAttrImplements,
+              excludeMacroExpansions);
           for (auto foundSwiftDecl : swiftLookupResult) {
             allFound.push_back(foundSwiftDecl);
           }
@@ -1828,7 +1844,8 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
   }
 
   DeclName macroExpansionKey = adjustLazyMacroExpansionNameKey(ctx, name);
-  if (!Table.isLazilyCompleteForMacroExpansion(macroExpansionKey)) {
+  if (!excludeMacroExpansions &&
+      !Table.isLazilyCompleteForMacroExpansion(macroExpansionKey)) {
     populateLookupTableEntryFromMacroExpansions(
         ctx, Table, macroExpansionKey, decl);
     for (auto ext : decl->getExtensions()) {
@@ -1845,8 +1862,9 @@ DirectLookupRequest::evaluate(Evaluator &evaluator,
   }
 
   // We found something; return it.
-  return maybeFilterOutAttrImplements(known->second, name,
-                                      includeAttrImplements);
+  return maybeFilterOutUnwantedDecls(known->second, name,
+                                     includeAttrImplements,
+                                     excludeMacroExpansions);
 }
 
 bool NominalTypeDecl::createObjCMethodLookup() {
@@ -2201,6 +2219,8 @@ QualifiedLookupRequest::evaluate(Evaluator &eval, const DeclContext *DC,
     auto flags = OptionSet<NominalTypeDecl::LookupDirectFlags>();
     if (options & NL_IncludeAttributeImplements)
       flags |= NominalTypeDecl::LookupDirectFlags::IncludeAttrImplements;
+    if (options & NL_ExcludeMacroExpansions)
+      flags |= NominalTypeDecl::LookupDirectFlags::ExcludeMacroExpansions;
     for (auto decl : current->lookupDirect(member.getFullName(), flags)) {
       // If we're performing a type lookup, don't even attempt to validate
       // the decl if its not a type.