[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

@@ -433,15 +433,16 @@ 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;
-
+    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.
 }
 
@@ -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/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);
 }

lib/Sema/TypeCheckMacros.cpp

@@ -676,17 +676,17 @@ Expr *swift::expandMacroExpr(
   return expandedExpr;
 }
 
-/// Expands the given macro expansion declaration and type-check the result.
-bool swift::expandFreestandingDeclarationMacro(
-    MacroExpansionDecl *med, SmallVectorImpl<Decl *> &results) {
+/// Expands the given macro expansion declaration.
+Optional<unsigned>
+swift::expandFreestandingDeclarationMacro(MacroExpansionDecl *med) {
   auto *dc = med->getDeclContext();
   ASTContext &ctx = dc->getASTContext();
   SourceManager &sourceMgr = ctx.SourceMgr;
 
   auto moduleDecl = dc->getParentModule();
   auto sourceFile = moduleDecl->getSourceFileContainingLocation(med->getLoc());
   if (!sourceFile)
-    return false;
+    return None;
 
   // Evaluate the macro.
   NullTerminatedStringRef evaluatedSource;
@@ -697,21 +697,21 @@ bool swift::expandFreestandingDeclarationMacro(
 
   if (isFromExpansionOfMacro(sourceFile, macro, MacroRole::Declaration)) {
     med->diagnose(diag::macro_recursive, macro->getName());
-    return false;
+    return None;
   }
 
   auto macroDef = macro->getDefinition();
   switch (macroDef.kind) {
   case MacroDefinition::Kind::Undefined:
   case MacroDefinition::Kind::Invalid:
     // Already diagnosed as an error elsewhere.
-    return false;
+    return None;
 
   case MacroDefinition::Kind::Builtin: {
     switch (macroDef.getBuiltinKind()) {
     case BuiltinMacroKind::ExternalMacro:
       // FIXME: Error here.
-      return false;
+      return None;
     }
   }
 
@@ -729,15 +729,15 @@ bool swift::expandFreestandingDeclarationMacro(
                     macro->getName()
       );
       macro->diagnose(diag::decl_declared_here, macro->getName());
-      return false;
+      return None;
     }
 
     // Make sure freestanding macros are enabled before we expand.
     if (!ctx.LangOpts.hasFeature(Feature::FreestandingMacros) &&
         !macro->getMacroRoles().contains(MacroRole::Expression)) {
       med->diagnose(
           diag::macro_experimental, "freestanding", "FreestandingMacros");
-      return false;
+      return None;
     }
 
 #if SWIFT_SWIFT_PARSER
@@ -746,7 +746,7 @@ bool swift::expandFreestandingDeclarationMacro(
     // Builtin macros are handled via ASTGen.
     auto astGenSourceFile = sourceFile->exportedSourceFile;
     if (!astGenSourceFile)
-      return false;
+      return None;
 
     Mangle::ASTMangler mangler;
     auto discriminator = mangler.mangleMacroExpansion(med);
@@ -760,13 +760,13 @@ bool swift::expandFreestandingDeclarationMacro(
         med->getStartLoc().getOpaquePointerValue(), &evaluatedSourceAddress,
         &evaluatedSourceLength);
     if (!evaluatedSourceAddress)
-      return false;
+      return None;
     evaluatedSource = NullTerminatedStringRef(evaluatedSourceAddress,
                                               (size_t)evaluatedSourceLength);
     break;
 #else
     med->diagnose(diag::macro_unsupported);
-    return false;
+    return None;
 #endif
   }
   }
@@ -820,12 +820,11 @@ bool swift::expandFreestandingDeclarationMacro(
     if (!decl) {
       ctx.Diags.diagnose(
           macroBufferRange.getStart(), diag::expected_macro_expansion_decls);
-      return false;
+      return None;
     }
     decl->setDeclContext(dc);
-    results.push_back(decl);
   }
-  return true;
+  return macroBufferID;
 }
 
 // If this storage declaration is a variable with an explicit initializer,
@@ -1219,17 +1218,6 @@ swift::expandPeers(CustomAttr *attr, MacroDecl *macro, Decl *decl) {
     return None;
 
   PrettyStackTraceDecl debugStack("applying expanded peer macro", decl);
-
-  auto *parent = decl->getDeclContext();
-  auto topLevelDecls = macroSourceFile->getTopLevelDecls();
-  for (auto peer : topLevelDecls) {
-    if (auto *nominal = dyn_cast<NominalTypeDecl>(parent)) {
-      nominal->addMember(peer);
-    } else if (auto *extension = dyn_cast<ExtensionDecl>(parent)) {
-      extension->addMember(peer);
-    }
-  }
-
   return macroSourceFile->getBufferID();
 }
 

lib/Sema/TypeCheckMacros.h

@@ -43,9 +43,8 @@ Expr *expandMacroExpr(
 /// Expands the given macro expansion declaration, type-checks the replacement
 /// declarations, and adds them to \p results.
 ///
-/// \returns true if expansion succeeded, false if failed.
-bool expandFreestandingDeclarationMacro(
-    MacroExpansionDecl *med, SmallVectorImpl<Decl *> &results);
+/// \returns Expansion buffer ID if expansion succeeded, \p None if failed.
+Optional<unsigned> expandFreestandingDeclarationMacro(MacroExpansionDecl *med);
 
 /// Expand the accessors for the given storage declaration based on the
 /// custom attribute that references the given macro.