Switch back to parsing top-level MacroExpansionExpr.

Use a "substitute decl" in `MacroExpansionExpr` to represent top-level or local macro expansions that can be interpreted as a declaration macro or code item macro expansion.

Author: rxwei

include/swift/AST/Decl.h

@@ -8475,7 +8475,6 @@ class MacroExpansionDecl : public Decl {
   ArrayRef<ASTNode> getRewritten() const;
   ConcreteDeclRef getMacroRef() const { return macroRef; }
   void setMacroRef(ConcreteDeclRef ref) { macroRef = ref; }
-  MacroExpansionExpr *createExpr() const;
 
   /// Returns a discriminator which determines this macro expansion's index
   /// in the sequence of macro expansions within the current function.

include/swift/AST/Expr.h

@@ -6156,6 +6156,7 @@ class MacroExpansionExpr final : public Expr {
   ArgumentList *ArgList;
   Expr *Rewritten;
   MacroRoles Roles;
+  MacroExpansionDecl *SubstituteDecl;
 
   /// The referenced macro.
   ConcreteDeclRef macroRef;
@@ -6178,7 +6179,7 @@ class MacroExpansionExpr final : public Expr {
         MacroName(macroName), MacroNameLoc(macroNameLoc),
         LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
         GenericArgs(genericArgs),
-        Rewritten(nullptr), Roles(roles) {
+        Rewritten(nullptr), Roles(roles), SubstituteDecl(nullptr) {
     Bits.MacroExpansionExpr.Discriminator = InvalidDiscriminator;
 
     // Macro expansions always have an argument list. If one is not provided, create
@@ -6235,6 +6236,10 @@ class MacroExpansionExpr final : public Expr {
 
   SourceRange getSourceRange() const;
 
+  MacroExpansionDecl *createSubstituteDecl() const;
+  MacroExpansionDecl *getSubstituteDecl() const;
+  void setSubstituteDecl(MacroExpansionDecl *decl);
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::MacroExpansion;
   }

lib/AST/ASTVerifier.cpp

@@ -357,11 +357,6 @@ class Verifier : public ASTWalker {
       llvm_unreachable("not all cases handled!");
     }
 
-    // Don't eagerly expand macros.
-    bool shouldWalkIntoMacroExpansions() override {
-      return false;
-    }
-
     PreWalkAction walkToDeclPre(Decl *D) override {
       switch (D->getKind()) {
 #define DISPATCH(ID) return dispatchVisitPre(static_cast<ID##Decl*>(D))

lib/AST/Decl.cpp

@@ -398,6 +398,28 @@ void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
     }
   }
 
+  if (auto *med = dyn_cast<MacroExpansionDecl>(this)) {
+    for (auto node : med->getRewritten())
+      if (auto *decl = node.dyn_cast<Decl *>())
+        callback(decl);
+  }
+
+  // Top-level macro expansion expressions are treated as top-level decls.
+  else if (auto *tlcd = dyn_cast<TopLevelCodeDecl>(this)) {
+    for (auto node : tlcd->getBody()->getElements()) {
+      if (auto *expr = node.dyn_cast<Expr *>()) {
+        if (auto *mee = dyn_cast<MacroExpansionExpr>(expr)) {
+          auto *med = mee->getSubstituteDecl();
+          if (!med) {
+            med = mee->createSubstituteDecl();
+            mee->setSubstituteDecl(med);
+          }
+          callback(med);
+        }
+      }
+    }
+  }
+
   // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
 }
 
@@ -9966,13 +9988,6 @@ SourceRange MacroExpansionDecl::getSourceRange() const {
   return SourceRange(PoundLoc, endLoc);
 }
 
-MacroExpansionExpr *MacroExpansionDecl::createExpr() const {
-  return new (getASTContext()) MacroExpansionExpr(
-      getDeclContext(), PoundLoc, MacroName, MacroNameLoc, LeftAngleLoc,
-      GenericArgs, RightAngleLoc, ArgList,
-      MacroRoles(MacroRole::Expression) | MacroRole::Declaration);
-}
-
 unsigned MacroExpansionDecl::getDiscriminator() const {
   if (getRawDiscriminator() != InvalidDiscriminator)
     return getRawDiscriminator();

lib/AST/Expr.cpp

@@ -2709,6 +2709,23 @@ unsigned MacroExpansionExpr::getDiscriminator() const {
   return getRawDiscriminator();
 }
 
+MacroExpansionDecl *MacroExpansionExpr::createSubstituteDecl() const {
+  auto dc = getDeclContext();
+  if (auto *tlcd = dyn_cast_or_null<TopLevelCodeDecl>(dc->getAsDecl()))
+    dc = tlcd->getDeclContext();
+  return new (DC->getASTContext()) MacroExpansionDecl(
+      getDeclContext(), SigilLoc, MacroName, MacroNameLoc, LeftAngleLoc,
+      GenericArgs, RightAngleLoc, ArgList);
+}
+
+MacroExpansionDecl *MacroExpansionExpr::getSubstituteDecl() const {
+  return SubstituteDecl;
+}
+
+void MacroExpansionExpr::setSubstituteDecl(MacroExpansionDecl *decl) {
+  SubstituteDecl = decl;
+}
+
 void swift::simple_display(llvm::raw_ostream &out, const ClosureExpr *CE) {
   if (!CE) {
     out << "(null)";

lib/AST/Module.cpp

@@ -263,6 +263,12 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
 
     else if (auto *MED = dyn_cast<MacroExpansionDecl>(D))
       DelayedMacroExpansions.push_back(MED);
+
+    else if (isa<TopLevelCodeDecl>(D)) {
+      D->visitAuxiliaryDecls([&](Decl *auxDecl) {
+        addToUnqualifiedLookupCache(ArrayRef<Decl *>{auxDecl}, false);
+      });
+    }
   }
 }
 
@@ -339,13 +345,12 @@ SourceLookupCache::SourceLookupCache(const ModuleDecl &M) {
 }
 
 void SourceLookupCache::addDelayedMacroExpansions() {
-  for (unsigned i = 0; i != DelayedMacroExpansions.size(); ++i) {
-    auto *MED = DelayedMacroExpansions[i];
+  while (!DelayedMacroExpansions.empty()) {
+    auto *MED = DelayedMacroExpansions.pop_back_val();
     for (auto node : MED->getRewritten())
       if (auto *decl = node.dyn_cast<Decl *>())
         addToUnqualifiedLookupCache(TinyPtrVector<Decl *>{decl}, false);
   }
-  DelayedMacroExpansions.clear();
 }
 
 void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,

lib/AST/NameLookup.cpp

@@ -1336,7 +1336,6 @@ MemberLookupTable::MemberLookupTable(ASTContext &ctx) {
 }
 
 void MemberLookupTable::addMember(Decl *member) {
-
   // Only value declarations matter.
   auto vd = dyn_cast<ValueDecl>(member);
   if (!vd)
@@ -1514,6 +1513,7 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
                                             MemberLookupTable &table,
                                             DeclName name,
                                             NominalTypeDecl *dc) {
+  // FIXME: Recurse
   auto expandAndPopulate = [&](MacroExpansionDecl *med) {
     for (auto node : med->getRewritten())
       if (auto *decl = node.dyn_cast<Decl *>())

lib/AST/TypeCheckRequests.cpp

@@ -1743,10 +1743,7 @@ ArgumentList *UnresolvedMacroReference::getArgs() const {
 }
 
 MacroRoles UnresolvedMacroReference::getMacroRoles() const {
-  if (pointer.is<MacroExpansionExpr *>())
-    return MacroRole::Expression;
-
-  if (pointer.is<MacroExpansionDecl *>())
+  if (pointer.is<MacroExpansionExpr *>() || pointer.is<MacroExpansionDecl *>())
     return getFreestandingMacroRoles();
 
   if (pointer.is<CustomAttr *>())

lib/Parse/ParseDecl.cpp

@@ -4682,7 +4682,6 @@ bool swift::isKeywordPossibleDeclStart(const Token &Tok) {
   case tok::kw_subscript:
   case tok::kw_typealias:
   case tok::kw_var:
-  case tok::pound:
   case tok::pound_if:
   case tok::pound_warning:
   case tok::pound_error:
@@ -4794,10 +4793,6 @@ bool Parser::isStartOfSwiftDecl(bool allowPoundIfAttributes) {
     return isStartOfSwiftDecl(allowPoundIfAttributes);
   }
 
-  // Macro expansion starts with a pound and a custom identifier.
-  if (Tok.is(tok::pound) && peekToken().is(tok::identifier))
-    return true;
-
   // Skip a #if that contains only attributes in all branches. These will be
   // parsed as attributes of a declaration, not as separate declarations.
   if (Tok.is(tok::pound_if) && allowPoundIfAttributes) {
@@ -9593,9 +9588,6 @@ Parser::parseDeclMacroExpansion(ParseDeclOptions flags,
     }
   }
 
-  if (!argList)
-    argList = ArgumentList::createImplicit(Context, {});
-
   return makeParserResult(
       status,
       new (Context) MacroExpansionDecl(

lib/Parse/ParseExpr.cpp

@@ -3391,7 +3391,9 @@ ParserResult<Expr> Parser::parseExprMacroExpansion(bool isExprBasic) {
       new (Context) MacroExpansionExpr(
           CurDeclContext, poundLoc, macroNameRef, macroNameLoc, leftAngleLoc,
           Context.AllocateCopy(genericArgs), rightAngleLoc, argList,
-          MacroRole::Expression));
+          CurDeclContext->isTypeContext()
+              ? MacroRole::Declaration
+              : getFreestandingMacroRoles()));
 }
 
 /// parseExprCollection - Parse a collection literal expression.

lib/Parse/ParseRequests.cpp

@@ -173,6 +173,13 @@ SourceFileParsingResult ParseSourceFileRequest::evaluate(Evaluator &evaluator,
 
     switch (generatedInfo->kind) {
     case GeneratedSourceInfo::FreestandingDeclMacroExpansion:
+      if (parser.CurDeclContext->isTypeContext()) {
+        parser.parseExpandedMemberList(items);
+      } else {
+        parser.parseTopLevelItems(items);
+      }
+      break;
+
     case GeneratedSourceInfo::ExpressionMacroExpansion:
     case GeneratedSourceInfo::ReplacedFunctionBody:
     case GeneratedSourceInfo::PrettyPrinted: {

lib/SILGen/SILGen.cpp

@@ -1792,8 +1792,7 @@ 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.
+  // Expanded declaration macros were already visited
 }
 
 bool

lib/Sema/CSApply.cpp

@@ -5365,16 +5365,23 @@ namespace {
       ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
       E->setMacroRef(macroRef);
 
-      // For now, only expand macro expansion expressions that fulfill
-      // `MacroRole::Expression` exactly. Freestanding code item macros
-      // have a `getMacroRoles()` value equal to `getFreestandingMacroRoles()`,
-      // which includes expression macros, and they are expanded in a separate
-      // request.
-      if (E->getMacroRoles() == MacroRole::Expression &&
-          !cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions)) {
-        if (auto newExpr = expandMacroExpr(dc, E, macroRef, expandedType)) {
-          E->setRewritten(newExpr);
+      if (!cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions)) {
+        if (macro->getMacroRoles().contains(MacroRole::Expression)) {
+          if (auto newExpr = expandMacroExpr(dc, E, macroRef, expandedType)) {
+            E->setRewritten(newExpr);
+          }
+        }
+        // For a non-expression macro, expand it as a declaration.
+        else if (macro->getMacroRoles().contains(MacroRole::Declaration)) {
+          if (!E->getSubstituteDecl()) {
+            auto *med = E->createSubstituteDecl();
+            E->setSubstituteDecl(med);
+            TypeChecker::typeCheckDecl(med);
+          }
+          cs.setType(E, cs.getASTContext().getVoidType());
         }
+        // Other macro roles may also be encountered here, as they use
+        // `MacroExpansionExpr` for resolution. In those cases, do not expand.
       }
 
       cs.cacheExprTypes(E);

lib/Sema/MiscDiagnostics.cpp

@@ -18,7 +18,6 @@
 #include "TypeCheckAvailability.h"
 #include "TypeCheckConcurrency.h"
 #include "TypeChecker.h"
-#include "swift/AST/ASTScope.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/ExistentialLayout.h"
@@ -3708,8 +3707,8 @@ ASTWalker::PreWalkResult<Expr *> VarDeclUsageChecker::walkToExprPre(Expr *E) {
     return Action::SkipChildren(E);
   }
 
-  // FIXME: `MacroExpansionDecl` may expand to a `MacroExpansionExpr` with the
-  // same argument expressions.
+  // TODO: `MacroExpansionExpr` may have a substitute `MacroExpansionDecl` with
+  // the same argument expressions. How should we assert here?
   // assert(AllExprsSeen.insert(E).second && "duplicate traversal");
 
   // If this is a DeclRefExpr found in a random place, it is a load of the

lib/Sema/TypeCheckDecl.cpp

@@ -1617,7 +1617,9 @@ TypeChecker::lookupMacros(DeclContext *dc, DeclNameRef macroName,
   for (const auto &found : lookup.allResults()) {
     if (auto macro = dyn_cast<MacroDecl>(found.getValueDecl())) {
       auto candidateRoles = macro->getMacroRoles();
-      if (candidateRoles && roles.contains(candidateRoles)) {
+      if ((candidateRoles && roles.contains(candidateRoles)) ||
+          // FIXME: `externalMacro` should have all roles.
+          macro->getBaseIdentifier().str() == "externalMacro") {
         foundRoles |= candidateRoles;
         results.push_back(macro);
       }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -2054,23 +2054,9 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
   }
 
   void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
+    // Expansion already visited as auxiliary decls.
     // Assign a discriminator.
     (void)MED->getDiscriminator();
-
-    auto *args = MED->getArgs();
-    for (auto argIndex : range(args->size())) {
-      auto expr = args->getExpr(argIndex);
-      TypeChecker::typeCheckExpression(expr, MED->getDeclContext());
-      args->setExpr(argIndex, expr);
-    }
-
-    for (auto node : MED->getRewritten()) {
-      if (auto *decl = node.dyn_cast<Decl *>()) {
-        if (auto *tlcd = dyn_cast<TopLevelCodeDecl>(decl))
-          continue; // Already handled.
-        visit(decl);
-      }
-    }
   }
 
   void visitBoundVariable(VarDecl *VD) {
@@ -3779,14 +3765,6 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
                                           MacroExpansionDecl *MED) const {
   auto &ctx = MED->getASTContext();
   auto *dc = MED->getDeclContext();
-  MacroRoles viableRoles(MacroRole::Declaration);
-
-  // In a code block context, a macro expansion decl can also expand to an
-  // arbitrary code item.
-  if (dc->isLocalContext() ||
-      dc->getContextKind() == DeclContextKind::FileUnit) {
-    viableRoles |= MacroRole::Expression;
-  }
 
   // Resolve macro candidates.
   auto macro = evaluateOrDefault(
@@ -3796,31 +3774,11 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
   MED->setMacroRef(macro);
 
   auto roles = macro->getMacroRoles();
-
-  // If the resolved macro is an expression-only macro, expand to a macro
-  // expansion expr.
-  if (roles.containsOnly(MacroRole::Expression)) {
-    // Make an equivalent `MacroExpansionExpr` as the expansion.
-    Expr *mee = MED->createExpr();
-    ASTNode expansionNode;
-
-    // If top-level, wrap it in a `TopLevelCodeDecl(BraceStmt(...))`.
-    if (dc->getContextKind() == DeclContextKind::FileUnit) {
-      auto *tlcd = new (ctx) TopLevelCodeDecl(
-          dc, BraceStmt::createImplicit(ctx, {mee}));
-      tlcd->setDeclContext(dc);
-      TypeChecker::typeCheckTopLevelCodeDecl(tlcd);
-      TypeChecker::contextualizeTopLevelCode(tlcd);
-      expansionNode = tlcd;
-    }
-    else {
-      TypeChecker::typeCheckExpression(
-          mee, dc, {}, TypeCheckExprFlags::IsDiscarded);
-      TypeChecker::postTypeCheckMacroExpansion(mee, dc);
-      expansionNode = mee;
-    }
-    return ctx.AllocateCopy(ArrayRef<ASTNode>{expansionNode});
-  }
+  // If it's not a declaration macro or a code item macro, it must have been
+  // parsed as an expression macro, and this decl is just its substitute decl.
+  // So there's no thing to be done here.
+  if (!roles.contains(MacroRole::Declaration))
+    return {};
 
   // Otherwise, we treat it as a declaration macro.
   assert(roles.contains(MacroRole::Declaration));