[NameLookup] Move macro-related name lookup operations into the namelookup

namespace.

This moves the `isInMacroArgument` predicate and `lookupMacros` into `namelookup`.
ASTScope still encapsulates the scope tree and contains the operation to lookup
the enclosing macro scope, which then invokes a callback to determine whether a
potential macro scope is indeed a macro, because answering this question requires
name lookup.

Author: hborla

include/swift/AST/ASTScope.h

@@ -218,6 +218,10 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
     return nullptr;
   }
 
+  virtual NullablePtr<MacroExpansionDecl> getFreestandingMacro() const {
+    return nullptr;
+  }
+
 #pragma mark - debugging and printing
 
 public:
@@ -278,16 +282,9 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   static std::pair<CaseStmt *, CaseStmt *>
   lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
 
-  /// Returns \c true if the given source location is inside an attached
-  /// or freestanding macro argument.
-  static bool isInMacroArgument(SourceFile *sourceFile, SourceLoc loc);
-
-  /// Returns \c true if this scope contains macro arguments.
-  ///
-  /// This is always true within macro expansion decl scopes, and it's
-  /// true within custom attribute scopes if the attribute name is a
-  /// potential macro reference.
-  virtual bool isMacroArgumentScope() const { return false; }
+  static void lookupEnclosingMacroScope(
+      SourceFile *sourceFile, SourceLoc loc,
+      llvm::function_ref<bool(ASTScope::PotentialMacro)> consume);
 
   /// Scopes that cannot bind variables may set this to true to create more
   /// compact scope tree in the debug info.
@@ -879,13 +876,6 @@ class CustomAttributeScope final : public ASTScopeImpl {
     return attr;
   }
   bool ignoreInDebugInfo() const override { return true; }
-
-  bool isMacroArgumentScope() const override {
-    // FIXME: This should check whether the attribute name is
-    // a macro. Otherwise, macro expansion will be suppressed
-    // for property wrapper arguments.
-    return true;
-  }
 
 private:
   void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
@@ -1285,8 +1275,8 @@ class MacroExpansionDeclScope final : public ASTScopeImpl {
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
 
-  bool isMacroArgumentScope() const override {
-    return true;
+  NullablePtr<MacroExpansionDecl> getFreestandingMacro() const override {
+    return decl;
   }
 
 protected:

include/swift/AST/NameLookup.h

@@ -545,6 +545,15 @@ template <typename Result>
 void filterForDiscriminator(SmallVectorImpl<Result> &results,
                             DebuggerClient *debugClient);
 
+/// \returns The set of macro declarations with the given name that
+/// fulfill any of the given macro roles.
+SmallVector<MacroDecl *, 1>
+lookupMacros(DeclContext *dc, DeclNameRef macroName, MacroRoles roles);
+
+/// \returns Whether the given source location is inside an attached
+/// or freestanding macro argument.
+bool isInMacroArgument(SourceFile *sourceFile, SourceLoc loc);
+
 /// Call the given function body with each macro declaration and its associated
 /// role attribute for the given role.
 ///
@@ -816,9 +825,21 @@ class ASTScope : public ASTAllocated<ASTScope> {
   static std::pair<CaseStmt *, CaseStmt *>
   lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
 
-  /// Returns \c true if the given source location is inside an attached
-  /// or freestanding macro argument.
-  static bool isInMacroArgument(SourceFile *sourceFile, SourceLoc loc);
+  using PotentialMacro =
+      llvm::PointerUnion<FreestandingMacroExpansion *, CustomAttr *>;
+
+  /// Look up the scope tree for the nearest enclosing macro scope at
+  /// the given source location.
+  ///
+  /// \param sourceFile The source file containing the given location.
+  /// \param loc        The source location to start lookup from.
+  /// \param consume    A function that is called when a potential macro
+  ///                   scope is found. If \c consume returns \c true, lookup
+  ///                   will stop. If \c consume returns \c false, lookup will
+  ///                   continue up the scope tree.
+  static void lookupEnclosingMacroScope(
+      SourceFile *sourceFile, SourceLoc loc,
+      llvm::function_ref<bool(PotentialMacro macro)> consume);
 
   SWIFT_DEBUG_DUMP;
   void print(llvm::raw_ostream &) const;

include/swift/AST/TypeCheckRequests.h

@@ -3249,6 +3249,9 @@ class UnresolvedMacroReference {
 public:
   UnresolvedMacroReference(FreestandingMacroExpansion *exp) : pointer(exp) {}
   UnresolvedMacroReference(CustomAttr *attr) : pointer(attr) {}
+  UnresolvedMacroReference(
+      llvm::PointerUnion<FreestandingMacroExpansion *, CustomAttr *> pointer)
+      : pointer(pointer) {}
 
   FreestandingMacroExpansion *getFreestanding() const {
     return pointer.dyn_cast<FreestandingMacroExpansion *>();

lib/AST/ASTScope.cpp

@@ -60,9 +60,10 @@ std::pair<CaseStmt *, CaseStmt *> ASTScope::lookupFallthroughSourceAndDest(
   return ASTScopeImpl::lookupFallthroughSourceAndDest(sourceFile, loc);
 }
 
-bool ASTScope::isInMacroArgument(SourceFile *sourceFile,
-                                 SourceLoc loc) {
-  return ASTScopeImpl::isInMacroArgument(sourceFile, loc);
+void ASTScope::lookupEnclosingMacroScope(
+    SourceFile *sourceFile, SourceLoc loc,
+    llvm::function_ref<bool(PotentialMacro)> body) {
+  return ASTScopeImpl::lookupEnclosingMacroScope(sourceFile, loc, body);
 }
 
 #if SWIFT_COMPILER_IS_MSVC

lib/AST/ASTScopeLookup.cpp

@@ -676,19 +676,26 @@ std::pair<CaseStmt *, CaseStmt *> ASTScopeImpl::lookupFallthroughSourceAndDest(
   return { nullptr, nullptr };
 }
 
-bool ASTScopeImpl::isInMacroArgument(SourceFile *sourceFile,
-                                     SourceLoc loc) {
+void ASTScopeImpl::lookupEnclosingMacroScope(
+    SourceFile *sourceFile, SourceLoc loc,
+    llvm::function_ref<bool(ASTScope::PotentialMacro)> consume) {
   if (!sourceFile || sourceFile->Kind == SourceFileKind::Interface)
-    return false;
+    return;
 
   if (loc.isInvalid())
-    return false;
+    return;
 
   auto *fileScope = sourceFile->getScope().impl;
   auto *scope = fileScope->findInnermostEnclosingScope(loc, nullptr);
   do {
-    if (scope->isMacroArgumentScope())
-      return true;
+    auto *freestanding = scope->getFreestandingMacro().getPtrOrNull();
+    if (freestanding && consume(freestanding))
+      return;
+
+    auto *attr = scope->getDeclAttributeIfAny().getPtrOrNull();
+    auto *potentialAttached = dyn_cast_or_null<CustomAttr>(attr);
+    if (potentialAttached && consume(potentialAttached))
+      return;
 
     // If we've reached a source file scope, we can't be inside of
     // a macro argument. Either this is a top-level source file, or
@@ -697,9 +704,7 @@ bool ASTScopeImpl::isInMacroArgument(SourceFile *sourceFile,
     // MacroExpansionDeclScope, and child scopes of freestanding macros
     // are otherwise inside the macro argument.
     if (scope->getClassName() == "ASTSourceFileScope")
-      return false;
+      return;
 
   } while ((scope = scope->getParent().getPtrOrNull()));
-
-  return false;
 }

lib/AST/NameLookup.cpp

@@ -1594,6 +1594,60 @@ static DeclName adjustLazyMacroExpansionNameKey(
   return name;
 }
 
+SmallVector<MacroDecl *, 1>
+namelookup::lookupMacros(DeclContext *dc, DeclNameRef macroName,
+                         MacroRoles roles) {
+  SmallVector<MacroDecl *, 1> choices;
+  auto moduleScopeDC = dc->getModuleScopeContext();
+  ASTContext &ctx = moduleScopeDC->getASTContext();
+
+  // Macro lookup should always exclude macro expansions; macro
+  // expansions cannot introduce new macro declarations. Note that
+  // the source location here doesn't matter.
+  UnqualifiedLookupDescriptor descriptor{
+    macroName, moduleScopeDC, SourceLoc(),
+    UnqualifiedLookupFlags::ExcludeMacroExpansions
+  };
+
+  auto lookup = evaluateOrDefault(
+      ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
+  for (const auto &found : lookup.allResults()) {
+    if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
+      auto candidateRoles = macro->getMacroRoles();
+      if ((candidateRoles && roles.contains(candidateRoles)) ||
+          // FIXME: `externalMacro` should have all roles.
+          macro->getBaseIdentifier().str() == "externalMacro") {
+        choices.push_back(macro);
+      }
+    }
+  }
+  return choices;
+}
+
+bool
+namelookup::isInMacroArgument(SourceFile *sourceFile, SourceLoc loc) {
+  bool inMacroArgument = false;
+
+  ASTScope::lookupEnclosingMacroScope(
+      sourceFile, loc,
+      [&](auto potentialMacro) -> bool {
+        UnresolvedMacroReference macro(potentialMacro);
+
+        if (macro.getFreestanding()) {
+          inMacroArgument = true;
+        } else if (auto *attr = macro.getAttr()) {
+          auto *moduleScope = sourceFile->getModuleScopeContext();
+          auto results = lookupMacros(moduleScope, macro.getMacroName(),
+                                      getAttachedMacroRoles());
+          inMacroArgument = !results.empty();
+        }
+
+        return inMacroArgument;
+      });
+
+  return inMacroArgument;
+}
+
 /// Call the given function body with each macro declaration and its associated
 /// role attribute for the given role.
 ///
@@ -2631,7 +2685,7 @@ directReferencesForUnqualifiedTypeLookup(DeclNameRef name,
 
   // Manually exclude macro expansions here since the source location
   // is overridden below.
-  if (ASTScope::isInMacroArgument(dc->getParentSourceFile(), loc))
+  if (namelookup::isInMacroArgument(dc->getParentSourceFile(), loc))
     options |= UnqualifiedLookupFlags::ExcludeMacroExpansions;
 
   // In a protocol or protocol extension, the 'where' clause can refer to

lib/AST/NameLookupRequests.cpp

@@ -545,7 +545,7 @@ static UnqualifiedLookupDescriptor excludeMacrosIfNeeded(
           UnqualifiedLookupFlags::ExcludeMacroExpansions))
     return descriptor;
 
-  auto isInMacroArgument = ASTScope::isInMacroArgument(
+  auto isInMacroArgument = namelookup::isInMacroArgument(
       descriptor.DC->getParentSourceFile(), descriptor.Loc);
 
   if (!isInMacroArgument)
@@ -562,7 +562,7 @@ static DirectLookupDescriptor excludeMacrosIfNeeded(
           NominalTypeDecl::LookupDirectFlags::ExcludeMacroExpansions))
     return descriptor;
 
-  auto isInMacroArgument = ASTScope::isInMacroArgument(
+  auto isInMacroArgument = namelookup::isInMacroArgument(
       descriptor.DC->getParentSourceFile(), loc);
 
   if (!isInMacroArgument)
@@ -581,7 +581,7 @@ excludeMacrosIfNeeded(const DeclContext *dc, SourceLoc loc,
   if (options & NL_ExcludeMacroExpansions)
     return options;
 
-  auto isInMacroArgument = ASTScope::isInMacroArgument(
+  auto isInMacroArgument = namelookup::isInMacroArgument(
       dc->getParentSourceFile(), loc);
 
   if (!isInMacroArgument)

lib/Sema/CSGen.cpp

@@ -3900,7 +3900,7 @@ namespace {
                  FunctionRefKind functionRefKind,
                  MacroRoles roles) {
       SmallVector<OverloadChoice, 1> choices;
-      auto results = TypeChecker::lookupMacros(
+      auto results = namelookup::lookupMacros(
           CurDC, DeclNameRef(macroName), roles);
       for (const auto &result : results) {
         OverloadChoice choice = OverloadChoice(Type(), result, functionRefKind);

lib/Sema/PreCheckExpr.cpp

@@ -606,7 +606,7 @@ Expr *TypeChecker::resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE,
       // FIXME: Don't perform typo correction inside macro arguments, because it
       // will invoke synthesizing declarations in this scope, which will attempt to
       // expand this macro which leads to circular reference errors.
-      if (!ASTScope::isInMacroArgument(DC->getParentSourceFile(), UDRE->getLoc())) {
+      if (!namelookup::isInMacroArgument(DC->getParentSourceFile(), UDRE->getLoc())) {
         TypeChecker::performTypoCorrection(DC, UDRE->getRefKind(), Type(),
                                            lookupOptions, corrections);
       }

lib/Sema/TypeCheckDecl.cpp

@@ -1609,36 +1609,6 @@ TypeChecker::lookupPrecedenceGroup(DeclContext *dc, Identifier name,
   return PrecedenceGroupLookupResult(dc, name, std::move(groups));
 }
 
-SmallVector<MacroDecl *, 1>
-TypeChecker::lookupMacros(DeclContext *dc, DeclNameRef macroName,
-                          MacroRoles roles) {
-  SmallVector<MacroDecl *, 1> choices;
-  auto moduleScopeDC = dc->getModuleScopeContext();
-  ASTContext &ctx = moduleScopeDC->getASTContext();
-
-  // Macro lookup should always exclude macro expansions; macro
-  // expansions cannot introduce new macro declarations. Note that
-  // the source location here doesn't matter.
-  UnqualifiedLookupDescriptor descriptor{
-    macroName, moduleScopeDC, SourceLoc(),
-    UnqualifiedLookupFlags::ExcludeMacroExpansions
-  };
-
-  auto lookup = evaluateOrDefault(
-      ctx.evaluator, UnqualifiedLookupRequest{descriptor}, {});
-  for (const auto &found : lookup.allResults()) {
-    if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
-      auto candidateRoles = macro->getMacroRoles();
-      if ((candidateRoles && roles.contains(candidateRoles)) ||
-          // FIXME: `externalMacro` should have all roles.
-          macro->getBaseIdentifier().str() == "externalMacro") {
-        choices.push_back(macro);
-      }
-    }
-  }
-  return choices;
-}
-
 /// Validate the given operator declaration.
 ///
 /// This establishes key invariants, such as an InfixOperatorDecl's

lib/Sema/TypeCheckMacros.cpp

@@ -1520,7 +1520,7 @@ ConcreteDeclRef ResolveMacroRequest::evaluate(Evaluator &evaluator,
   // When a macro is not found for a custom attribute, it may be a non-macro.
   // So bail out to prevent diagnostics from the contraint system.
   if (macroRef.getAttr()) {
-    auto foundMacros = TypeChecker::lookupMacros(
+    auto foundMacros = namelookup::lookupMacros(
         dc, macroRef.getMacroName(), roles);
     if (foundMacros.empty())
       return ConcreteDeclRef();

lib/Sema/TypeChecker.h

@@ -941,9 +941,6 @@ lookupPrecedenceGroupForInfixOperator(DeclContext *dc, Expr *op, bool diagnose);
 PrecedenceGroupLookupResult
 lookupPrecedenceGroup(DeclContext *dc, Identifier name, SourceLoc nameLoc);
 
-SmallVector<MacroDecl *, 1>
-lookupMacros(DeclContext *dc, DeclNameRef macroName, MacroRoles contexts);
-
 enum class UnsupportedMemberTypeAccessKind : uint8_t {
   None,
   TypeAliasOfUnboundGeneric,