[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.

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

@@ -8497,7 +8497,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

@@ -115,7 +115,15 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   LLVM_NODISCARD
   bool visit(Decl *D) {
     SetParentRAII SetParent(Walker, D);
-    return inherited::visit(D);
+    // Visit auxiliary decls, which may be decls from macro expansions.
+    bool alreadyFailed = false;
+    if (Walker.shouldWalkMacroArgumentsAndExpansion().second) {
+      D->visitAuxiliaryDecls([&](Decl *decl) {
+        if (alreadyFailed) return;
+        alreadyFailed = inherited::visit(decl);
+      });
+    }
+    return alreadyFailed || inherited::visit(D);
   }
 
   LLVM_NODISCARD
@@ -430,18 +438,9 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool visitMacroExpansionDecl(MacroExpansionDecl *MED) {
-    bool shouldWalkArguments, shouldWalkExpansion;
-    std::tie(shouldWalkArguments, shouldWalkExpansion) =
-        Walker.shouldWalkMacroArgumentsAndExpansion();
-    if (shouldWalkArguments && MED->getArgs() && doIt(MED->getArgs()))
-      return true;
-
-    if (shouldWalkExpansion) {
-      for (auto *decl : MED->getRewritten())
-        if (doIt(decl))
-          return true;
-    }
-    return false;
+    bool shouldWalkArguments =
+        Walker.shouldWalkMacroArgumentsAndExpansion().first;
+    return shouldWalkArguments && MED->getArgs() && !doIt(MED->getArgs());
   }
 
   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.
 }
 
@@ -10073,13 +10083,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,14 +1510,6 @@ 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);
-  };
-
   for (auto *member : dc->getCurrentMembersWithoutLoading()) {
     auto *med = dyn_cast<MacroExpansionDecl>(member);
     if (!med)
@@ -1532,7 +1524,9 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
     // 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);
+      med->visitAuxiliaryDecls([&](Decl *decl) {
+        table.addMember(decl);
+      });
     }
     // Otherwise, we expand the macro if it has the same decl base name being
     // looked for.
@@ -1550,7 +1544,9 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
               name.getBaseName().userFacingName();
       });
       if (it != attr->getNames().end())
-        expandAndPopulate(med);
+        med->visitAuxiliaryDecls([&](Decl *decl) {
+          table.addMember(decl);
+        });
     }
   }
 }

lib/IRGen/GenDecl.cpp

@@ -2457,6 +2457,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));
@@ -2542,8 +2545,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;
   }
 
@@ -5466,6 +5468,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:
@@ -5526,8 +5531,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/Sema/MiscDiagnostics.cpp

@@ -3805,7 +3805,7 @@ void VarDeclUsageChecker::handleIfConfig(IfConfigDecl *ICD) {
       : VDUC(VDUC), SF(VDUC.DC->getParentSourceFile()) {}
 
     MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::ArgumentsAndExpansion;
+      return MacroWalking::Arguments;
     }
 
     PostWalkResult<Expr *> walkToExprPost(Expr *E) override {

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -2063,12 +2063,7 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
     // Assign a discriminator.
     (void)MED->getDiscriminator();
-
-    auto rewritten = evaluateOrDefault(
-        Ctx.evaluator, ExpandMacroExpansionDeclRequest{MED}, {});
-
-    for (auto *decl : rewritten)
-      visit(decl);
+    // Expansion already visited as auxiliary decls.
   }
 
   void visitBoundVariable(VarDecl *VD) {
@@ -3763,7 +3758,7 @@ void TypeChecker::checkParameterList(ParameterList *params,
   }
 }
 
-ArrayRef<Decl *>
+Optional<unsigned>
 ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
                                           MacroExpansionDecl *MED) const {
   auto &ctx = MED->getASTContext();
@@ -3785,8 +3780,5 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
   MED->setMacroRef(macro);
 
   // Expand the macro.
-  SmallVector<Decl *, 2> expandedTemporary;
-  if (!expandFreestandingDeclarationMacro(MED, expandedTemporary))
-    return {};
-  return ctx.AllocateCopy(expandedTemporary);
+  return expandFreestandingDeclarationMacro(MED);
 }