[Macros] Always visit macro-produced decls as auxiliary decls

Always use `Decl::visitAuxiliaryDecls` to visit decls produced by macros, including peer macros and declaration macros. Use name-driven expansion for peer macros. Remove `MacroExpansionDecl::getRewritten()`.

Also make `ExpandMacroExpansionDeclRequest` cache the buffer ID (similar to other macros) instead of an array of decls.

Author: rxwei

include/swift/AST/ASTWalker.h

@@ -95,7 +95,10 @@ enum class MacroWalking {
 
   /// Walk into both the arguments of the macro as written in the source code
   /// and also the macro expansion.
-  ArgumentsAndExpansion
+  ArgumentsAndExpansion,
+
+  /// Don't walk into macros.
+  None
 };
 
 /// An abstract class used to traverse an AST.
@@ -545,6 +548,9 @@ class ASTWalker {
 
     case MacroWalking::ArgumentsAndExpansion:
       return std::make_pair(true, true);
+
+    case MacroWalking::None:
+      return std::make_pair(false, false);
     }
   }
 

include/swift/AST/Decl.h

@@ -8527,7 +8527,6 @@ class MacroExpansionDecl : public Decl {
   DeclNameLoc getMacroNameLoc() const { return MacroNameLoc; }
   DeclNameRef getMacroName() const { return MacroName; }
   ArgumentList *getArgs() const { return ArgList; }
-  ArrayRef<Decl *> getRewritten() const;
   ConcreteDeclRef getMacroRef() const { return macroRef; }
   void setMacroRef(ConcreteDeclRef ref) { macroRef = ref; }
 

include/swift/AST/TypeCheckRequests.h

@@ -3835,16 +3835,16 @@ class MacroDefinitionRequest
 /// Find the definition of a given macro.
 class ExpandMacroExpansionDeclRequest
     : public SimpleRequest<ExpandMacroExpansionDeclRequest,
-                           ArrayRef<Decl *>(MacroExpansionDecl *),
+                           Optional<unsigned>(MacroExpansionDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  ArrayRef<Decl *> evaluate(Evaluator &evaluator,
-                            MacroExpansionDecl *med) const;
+  Optional<unsigned>
+  evaluate(Evaluator &evaluator, MacroExpansionDecl *med) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeMemberVisitor.h

@@ -54,9 +54,17 @@ class TypeMemberVisitor : public DeclVisitor<ImplClass, RetTy> {
     return RetTy();
   }
 
+  RetTy visitMacroExpansionDecl(MacroExpansionDecl *D) {
+    // Expansion already visited as auxiliary decls.
+    return RetTy();
+  }
+
   /// A convenience method to visit all the members.
   void visitMembers(NominalTypeDecl *D) {
     for (Decl *member : D->getMembers()) {
+      member->visitAuxiliaryDecls([&](Decl *decl) {
+        asImpl().visit(decl);
+      });
       asImpl().visit(member);
     }
   }
@@ -78,12 +86,6 @@ class TypeMemberVisitor : public DeclVisitor<ImplClass, RetTy> {
         asImpl().visit(dd);
     }
   }
-
-  /// Visit expanded macros.
-  void visitMacroExpansionDecl(MacroExpansionDecl *D) {
-    for (auto *decl : D->getRewritten())
-      asImpl().visit(decl);
-  }
 };
 
 template<typename ImplClass, typename RetTy = void>

lib/AST/ASTVerifier.cpp

@@ -272,7 +272,7 @@ class Verifier : public ASTWalker {
   }
 
   MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::Expansion;
+    return MacroWalking::None;
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {

lib/AST/ASTWalker.cpp

@@ -433,15 +433,17 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     bool shouldWalkArguments, shouldWalkExpansion;
     std::tie(shouldWalkArguments, shouldWalkExpansion) =
         Walker.shouldWalkMacroArgumentsAndExpansion();
-    if (shouldWalkArguments && MED->getArgs() && doIt(MED->getArgs()))
+    if (shouldWalkArguments && MED->getArgs() && !doIt(MED->getArgs()))
       return true;
-
+    // Visit auxiliary decls, which may be decls from macro expansions.
+    bool alreadyFailed = false;
     if (shouldWalkExpansion) {
-      for (auto *decl : MED->getRewritten())
-        if (doIt(decl))
-          return true;
+      MED->visitAuxiliaryDecls([&](Decl *decl) {
+        if (alreadyFailed) return;
+        alreadyFailed = inherited::visit(decl);
+      });
     }
-    return false;
+    return alreadyFailed;
   }
 
   bool visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {

lib/AST/Decl.cpp

@@ -413,6 +413,16 @@ void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
     }
   }
 
+  else if (auto *med = dyn_cast<MacroExpansionDecl>(mutableThis)) {
+    if (auto bufferID = evaluateOrDefault(
+            ctx.evaluator, ExpandMacroExpansionDeclRequest{med}, {})) {
+      auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
+      auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+      for (auto *decl : sourceFile->getTopLevelDecls())
+        callback(decl);
+    }
+  }
+
   // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
 }
 
@@ -1915,7 +1925,9 @@ StringRef PatternBindingEntry::getInitStringRepresentation(
 
 SourceRange PatternBindingDecl::getSourceRange() const {
   SourceLoc startLoc = getStartLoc();
-  SourceLoc endLoc = getPatternList().back().getSourceRange().End;
+  SourceLoc endLoc = getPatternList().empty()
+      ? SourceLoc()
+      : getPatternList().back().getSourceRange().End;
   if (startLoc.isValid() != endLoc.isValid()) return SourceRange();
   return { startLoc, endLoc };
 }
@@ -10159,13 +10171,6 @@ unsigned MacroExpansionDecl::getDiscriminator() const {
   return getRawDiscriminator();
 }
 
-ArrayRef<Decl *> MacroExpansionDecl::getRewritten() const {
-  auto mutableThis = const_cast<MacroExpansionDecl *>(this);
-  return evaluateOrDefault(
-      getASTContext().evaluator,
-      ExpandMacroExpansionDeclRequest{mutableThis}, {});
-}
-
 NominalTypeDecl *
 ValueDecl::getRuntimeDiscoverableAttrTypeDecl(CustomAttr *attr) const {
   auto &ctx = getASTContext();

lib/AST/NameLookup.cpp

@@ -1510,48 +1510,70 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
                                             MemberLookupTable &table,
                                             DeclName name,
                                             NominalTypeDecl *dc) {
-  auto expandAndPopulate = [&](MacroExpansionDecl *med) {
-    auto expanded = evaluateOrDefault(med->getASTContext().evaluator,
-                                      ExpandMacroExpansionDeclRequest{med},
-                                      nullptr);
-    for (auto *decl : expanded)
-      table.addMember(decl);
-  };
-
+  auto *moduleScopeCtx = dc->getModuleScopeContext();
+  auto *module = dc->getModuleContext();
   for (auto *member : dc->getCurrentMembersWithoutLoading()) {
-    auto *med = dyn_cast<MacroExpansionDecl>(member);
-    if (!med)
-      continue;
-    auto declRef = evaluateOrDefault(
-        ctx.evaluator, ResolveMacroRequest{med, dc},
-        nullptr);
-    auto *macro = dyn_cast_or_null<MacroDecl>(declRef.getDecl());
-    if (!macro)
-      continue;
-    auto *attr = macro->getMacroRoleAttr(MacroRole::Declaration);
-    // If a macro produces arbitrary names, we have to expand it to factor its
-    // expansion results into name lookup.
-    if (attr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
-      expandAndPopulate(med);
-    }
-    // Otherwise, we expand the macro if it has the same decl base name being
-    // looked for.
-    else {
-      auto it = llvm::find_if(attr->getNames(),
-                              [&](const MacroIntroducedDeclName &introName) {
-        // FIXME: The `Named` kind of `MacroIntroducedDeclName` should store a
-        // `DeclName` instead of `Identifier`. This is so that we can compare
-        // base identifiers when the macro specifies a compound name.
-        // Currently only simple names are allowed in a `MacroRoleAttr`.
-        if (!name.isSpecial())
-          return introName.getIdentifier() == name.getBaseIdentifier();
-        else
-          return introName.getIdentifier().str() ==
-              name.getBaseName().userFacingName();
-      });
-      if (it != attr->getNames().end())
-        expandAndPopulate(med);
+    // Collect all macro introduced names, along with its corresponding macro
+    // reference. We need the macro reference to prevent adding auxiliary decls
+    // that weren't introduced by the macro.
+    llvm::SmallSet<DeclName, 4> allIntroducedNames;
+    bool introducesArbitraryNames = false;
+    if (auto *med = dyn_cast<MacroExpansionDecl>(member)) {
+      auto declRef = evaluateOrDefault(
+          ctx.evaluator, ResolveMacroRequest{med, dc},
+          nullptr);
+      if (!declRef)
+        continue;
+      auto *macro = dyn_cast<MacroDecl>(declRef.getDecl());
+      if (macro->getMacroRoleAttr(MacroRole::Declaration)
+          ->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary))
+        introducesArbitraryNames = true;
+      else {
+        SmallVector<DeclName, 4> introducedNames;
+        macro->getIntroducedNames(MacroRole::Declaration, nullptr,
+                                  introducedNames);
+        for (auto name : introducedNames)
+          allIntroducedNames.insert(name);
+      }
+    } else if (auto *vd = dyn_cast<ValueDecl>(member)) {
+      // We intentionally avoid calling `forEachAttachedMacro` in order to avoid
+      // a request cycle.
+      for (auto attrConst : member->getSemanticAttrs().getAttributes<CustomAttr>()) {
+        auto *attr = const_cast<CustomAttr *>(attrConst);
+        UnresolvedMacroReference macroRef(attr);
+        auto macroName = macroRef.getMacroName();
+        UnqualifiedLookupDescriptor lookupDesc{macroName, moduleScopeCtx};
+        auto lookup = evaluateOrDefault(
+            ctx.evaluator, UnqualifiedLookupRequest{lookupDesc}, {});
+        for (auto result : lookup.allResults()) {
+          auto *vd = result.getValueDecl();
+          auto *macro = dyn_cast<MacroDecl>(vd);
+          if (!macro)
+            continue;
+          auto *macroRoleAttr = macro->getMacroRoleAttr(MacroRole::Peer);
+          if (!macroRoleAttr)
+            continue;
+          if (macroRoleAttr->hasNameKind(
+                  MacroIntroducedDeclNameKind::Arbitrary))
+            introducesArbitraryNames = true;
+          else {
+            SmallVector<DeclName, 4> introducedNames;
+            macro->getIntroducedNames(
+                MacroRole::Peer, dyn_cast<ValueDecl>(member), introducedNames);
+            for (auto name : introducedNames)
+              allIntroducedNames.insert(name);
+          }
+        }
+      }
     }
+    // Expand macros based on the name.
+    if (introducesArbitraryNames || allIntroducedNames.contains(name))
+      member->visitAuxiliaryDecls([&](Decl *decl) {
+        auto *sf = module->getSourceFileContainingLocation(decl->getLoc());
+        // Bail out if the auxiliary decl was not produced by a macro.
+        if (!sf || sf->Kind != SourceFileKind::MacroExpansion) return;
+        table.addMember(decl);
+      });
   }
 }
 

lib/IRGen/GenDecl.cpp

@@ -2441,6 +2441,9 @@ void swift::irgen::disableAddressSanitizer(IRGenModule &IGM, llvm::GlobalVariabl
 
 /// Emit a global declaration.
 void IRGenModule::emitGlobalDecl(Decl *D) {
+  D->visitAuxiliaryDecls([&](Decl *decl) {
+    emitGlobalDecl(decl);
+  });
   switch (D->getKind()) {
   case DeclKind::Extension:
     return emitExtension(cast<ExtensionDecl>(D));
@@ -2526,8 +2529,7 @@ void IRGenModule::emitGlobalDecl(Decl *D) {
     return;
 
   case DeclKind::MacroExpansion:
-    for (auto *rewritten : cast<MacroExpansionDecl>(D)->getRewritten())
-      emitGlobalDecl(rewritten);
+    // Expansion already visited as auxiliary decls.
     return;
   }
 
@@ -5452,6 +5454,9 @@ Address IRGenModule::getAddrOfEnumCase(EnumElementDecl *Case,
 
 void IRGenModule::emitNestedTypeDecls(DeclRange members) {
   for (Decl *member : members) {
+    member->visitAuxiliaryDecls([&](Decl *decl) {
+      emitNestedTypeDecls({decl, nullptr});
+    });
     switch (member->getKind()) {
     case DeclKind::Import:
     case DeclKind::TopLevelCode:
@@ -5512,8 +5517,7 @@ void IRGenModule::emitNestedTypeDecls(DeclRange members) {
       emitClassDecl(cast<ClassDecl>(member));
       continue;
     case DeclKind::MacroExpansion:
-      for (auto *decl : cast<MacroExpansionDecl>(member)->getRewritten())
-        emitNestedTypeDecls({decl, nullptr});
+      // Expansion already visited as auxiliary decls.
       continue;
     }
   }

lib/SILGen/SILGen.cpp

@@ -1792,8 +1792,9 @@ void SILGenModule::visitMacroDecl(MacroDecl *d) {
 }
 
 void SILGenModule::visitMacroExpansionDecl(MacroExpansionDecl *d) {
-  // Expanded declaration macros were already added to the parent decl context
-  // for name lookup to work. Nothing to be done here.
+  d->visitAuxiliaryDecls([&](Decl *decl) {
+    visit(decl);
+  });
 }
 
 bool