[Macros] Top-level freestanding macros (#63553)

Allow freestanding macros to be used at top-level.
- Parse top-level `#…` as `MacroExpansionDecl` when we are not in scripting mode.
- Add macro expansion decls to the source lookup cache with name-driven lazy expansion. Not supporting arbitrary name yet.
- Experimental support for script mode and brace-level declaration macro expansions: When type-checking a `MacroExpansionExpr`, assign it a substitute `MacroExpansionDecl` if the macro reference resolves to a declaration macro. This doesn’t work quite fully yet and will be enabled in a future fix.

Author: rxwei

include/swift/AST/Decl.h

@@ -8306,19 +8306,21 @@ class PostfixOperatorDecl : public OperatorDecl {
   }
 };
 
+class MacroExpansionDecl;
+
 /// Represents a missing declaration in the source code.
 ///
 /// This is used for parser recovery, e.g. when parsing a floating
 /// attribute list, and to represent placeholders for unexpanded
 /// declarations generated by macros.
 class MissingDecl : public Decl {
-  /// If this missing decl represents an unexpanded peer generated
-  /// by a macro, \c unexpandedPeer contains the macro custom attribute
-  /// and the declaration the macro is attached to.
+  /// If this missing decl represents an unexpanded peer generated by a macro,
+  /// \c unexpandedMacro contains the macro reference and the base declaration
+  /// where the macro expansion applies.
   struct {
-    CustomAttr *macroAttr = nullptr;
-    ValueDecl *attachedTo = nullptr;
-  } unexpandedPeer;
+    llvm::PointerUnion<MacroExpansionDecl *, CustomAttr *> macroRef;
+    Decl *baseDecl;
+  } unexpandedMacro;
 
   /// The location that the decl would be if it wasn't missing.
   SourceLoc Loc;
@@ -8339,19 +8341,21 @@ class MissingDecl : public Decl {
   SourceRange getSourceRange() const { return SourceRange(Loc); }
 
   static MissingDecl *
-  forUnexpandedPeer(CustomAttr *macro, ValueDecl *attachedTo) {
-    auto &ctx = attachedTo->getASTContext();
-    auto *dc = attachedTo->getDeclContext();
+  forUnexpandedMacro(
+      llvm::PointerUnion<MacroExpansionDecl *, CustomAttr *> macroRef,
+      Decl *baseDecl) {
+    auto &ctx = baseDecl->getASTContext();
+    auto *dc = baseDecl->getDeclContext();
     auto *missing = new (ctx) MissingDecl(dc, SourceLoc());
 
-    missing->unexpandedPeer.macroAttr = macro;
-    missing->unexpandedPeer.attachedTo = attachedTo;
+    missing->unexpandedMacro.macroRef = macroRef;
+    missing->unexpandedMacro.baseDecl = baseDecl;
 
     return missing;
   }
 
-  using ExpandedPeerCallback = llvm::function_ref<void(ValueDecl *)>;
-  void forEachExpandedPeer(ExpandedPeerCallback callback);
+  using MacroExpandedDeclCallback = llvm::function_ref<void(ValueDecl *)>;
+  void forEachMacroExpandedDecl(MacroExpandedDeclCallback callback);
 
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Missing;
@@ -8512,13 +8516,7 @@ class MacroExpansionDecl : public Decl {
                      SourceLoc leftAngleLoc,
                      ArrayRef<TypeRepr *> genericArgs,
                      SourceLoc rightAngleLoc,
-                     ArgumentList *args)
-      : Decl(DeclKind::MacroExpansion, dc), PoundLoc(poundLoc),
-        MacroName(macro), MacroNameLoc(macroLoc),
-        LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
-        GenericArgs(genericArgs), ArgList(args) {
-    Bits.MacroExpansionDecl.Discriminator = InvalidDiscriminator;
-  }
+                     ArgumentList *args);
 
   ArrayRef<TypeRepr *> getGenericArgs() const { return GenericArgs; }
 
@@ -8532,6 +8530,9 @@ class MacroExpansionDecl : public Decl {
   DeclNameLoc getMacroNameLoc() const { return MacroNameLoc; }
   DeclNameRef getMacroName() const { return MacroName; }
   ArgumentList *getArgs() const { return ArgList; }
+  void setArgs(ArgumentList *args) { ArgList = args; }
+  using ExprOrStmtExpansionCallback = llvm::function_ref<void(ASTNode)>;
+  void forEachExpandedExprOrStmt(ExprOrStmtExpansionCallback) const;
   ConcreteDeclRef getMacroRef() const { return macroRef; }
   void setMacroRef(ConcreteDeclRef ref) { macroRef = ref; }
 

include/swift/AST/DiagnosticsSema.def

@@ -6928,6 +6928,9 @@ ERROR(literal_type_in_macro_expansion,none,
 ERROR(invalid_macro_introduced_name,none,
       "declaration name %0 is not covered by macro %1",
       (DeclName, DeclName))
+ERROR(global_freestanding_macro_script,none,
+      "global freestanding macros not yet supported in script mode",
+      ())
 
 //------------------------------------------------------------------------------
 // MARK: Move Only Errors

include/swift/AST/Expr.h

@@ -6197,6 +6197,7 @@ class MacroExpansionExpr final : public Expr {
   ArgumentList *ArgList;
   Expr *Rewritten;
   MacroRoles Roles;
+  MacroExpansionDecl *SubstituteDecl;
 
   /// The referenced macro.
   ConcreteDeclRef macroRef;
@@ -6219,7 +6220,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
@@ -6276,6 +6277,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;
   }

include/swift/AST/FileUnit.h

@@ -210,6 +210,14 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
   /// The order of the results is not guaranteed to be meaningful.
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {}
 
+  /// Finds all top-level decls in this file with their auxiliary decls such as
+  /// macro expansions.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  /// The order of the results is not guaranteed to be meaningful.
+  void getTopLevelDeclsWithAuxiliaryDecls(
+      SmallVectorImpl<Decl*> &results) const;
+
   virtual void
   getExportedPrespecializations(SmallVectorImpl<Decl *> &results) const {}
 

lib/AST/ASTDumper.cpp

@@ -1311,12 +1311,9 @@ namespace {
     void visitMacroExpansionDecl(MacroExpansionDecl *MED) {
       printCommon(MED, "macro_expansion_decl ");
       OS << MED->getMacroName();
-      if (MED->getArgs()) {
-        OS << '\n';
-        OS.indent(Indent + 2);
-        printArgumentList(OS, MED->getArgs(), Indent,
-                          [&](Expr *E) { printRec(E); });
-      }
+      OS << '\n';
+      printArgumentList(OS, MED->getArgs(), Indent,
+                        [&](Expr *E) { printRec(E); });
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
   };
@@ -3050,7 +3047,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
       printArgumentList(E->getArgs());
     }
     if (auto rewritten = E->getRewritten()) {
-      OS << '\n';
+      OS << " rewritten=\n";
       printRec(rewritten);
     }
     PrintWithColorRAII(OS, ParenthesisColor) << ')';

lib/AST/ASTPrinter.cpp

@@ -4683,20 +4683,18 @@ void PrintAST::visitMacroDecl(MacroDecl *decl) {
 
 void PrintAST::visitMacroExpansionDecl(MacroExpansionDecl *decl) {
   Printer << '#' << decl->getMacroName();
-  if (decl->getArgs()) {
-    Printer << '(';
-    auto args = decl->getArgs()->getOriginalArgs();
-    bool isFirst = true;
-    // FIXME: handle trailing closures.
-    for (auto arg : *args) {
-      if (!isFirst) {
-        Printer << ", ";
-      }
-      printArgument(arg);
-      isFirst = false;
+  Printer << '(';
+  auto args = decl->getArgs()->getOriginalArgs();
+  bool isFirst = true;
+  // FIXME: handle trailing closures.
+  for (auto arg : *args) {
+    if (!isFirst) {
+      Printer << ", ";
     }
-    Printer << ')';
+    printArgument(arg);
+    isFirst = false;
   }
+  Printer << ')';
 }
 
 void PrintAST::visitIntegerLiteralExpr(IntegerLiteralExpr *expr) {

lib/AST/ASTWalker.cpp

@@ -433,15 +433,29 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     bool shouldWalkArguments, shouldWalkExpansion;
     std::tie(shouldWalkArguments, shouldWalkExpansion) =
         Walker.shouldWalkMacroArgumentsAndExpansion();
-    if (shouldWalkArguments && MED->getArgs() && !doIt(MED->getArgs()))
-      return true;
+    if (shouldWalkArguments && MED->getArgs()) {
+      if (auto *argList = doIt(MED->getArgs()))
+        MED->setArgs(argList);
+      else
+        return true;
+    }
     // Visit auxiliary decls, which may be decls from macro expansions.
     bool alreadyFailed = false;
     if (shouldWalkExpansion) {
       MED->visitAuxiliaryDecls([&](Decl *decl) {
         if (alreadyFailed) return;
         alreadyFailed = inherited::visit(decl);
       });
+      MED->forEachExpandedExprOrStmt([&](ASTNode expandedNode) {
+        if (alreadyFailed) return;
+        if (auto *expr = expandedNode.dyn_cast<Expr *>()) {
+          if (!doIt(expr))
+            alreadyFailed = true;
+        } else if (auto *stmt = expandedNode.dyn_cast<Stmt *>()) {
+          if (!doIt(stmt))
+            alreadyFailed = true;
+        }
+      });
     }
     return alreadyFailed;
   }

lib/AST/Decl.cpp

@@ -9986,26 +9986,28 @@ MacroRoles swift::getAttachedMacroRoles() {
   return attachedMacroRoles;
 }
 
-void
-MissingDecl::forEachExpandedPeer(ExpandedPeerCallback callback) {
-  auto *macro = unexpandedPeer.macroAttr;
-  auto *attachedTo = unexpandedPeer.attachedTo;
-  if (!macro || !attachedTo)
+void MissingDecl::forEachMacroExpandedDecl(MacroExpandedDeclCallback callback) {
+  auto macroRef = unexpandedMacro.macroRef;
+  auto *baseDecl = unexpandedMacro.baseDecl;
+  if (!macroRef || !baseDecl)
     return;
 
-  attachedTo->visitAuxiliaryDecls(
-      [&](Decl *auxiliaryDecl) {
-        auto *sf = auxiliaryDecl->getInnermostDeclContext()->getParentSourceFile();
-        auto *macroAttr = sf->getAttachedMacroAttribute();
-        if (macroAttr != unexpandedPeer.macroAttr)
-          return;
-
-        auto *value = dyn_cast<ValueDecl>(auxiliaryDecl);
-        if (!value)
-          return;
-
-        callback(value);
-      });
+  baseDecl->visitAuxiliaryDecls([&](Decl *auxiliaryDecl) {
+    auto *sf = auxiliaryDecl->getInnermostDeclContext()->getParentSourceFile();
+    // We only visit auxiliary decls that are macro expansions associated with
+    // this macro reference.
+    if (auto *med = macroRef.dyn_cast<MacroExpansionDecl *>()) {
+     if (med != sf->getMacroExpansion().dyn_cast<Decl *>())
+       return;
+    } else if (auto *attr = macroRef.dyn_cast<CustomAttr *>()) {
+     if (attr != sf->getAttachedMacroAttribute())
+       return;
+    } else {
+      return;
+    }
+    if (auto *vd = dyn_cast<ValueDecl>(auxiliaryDecl))
+      callback(vd);
+  });
 }
 
 MacroDecl::MacroDecl(
@@ -10142,10 +10144,24 @@ Optional<BuiltinMacroKind> MacroDecl::getBuiltinKind() const {
   return def.getBuiltinKind();
 }
 
+MacroExpansionDecl::MacroExpansionDecl(
+    DeclContext *dc, SourceLoc poundLoc, DeclNameRef macro,
+    DeclNameLoc macroLoc, SourceLoc leftAngleLoc,
+    ArrayRef<TypeRepr *> genericArgs, SourceLoc rightAngleLoc,
+    ArgumentList *args)
+    : Decl(DeclKind::MacroExpansion, dc), PoundLoc(poundLoc),
+      MacroName(macro), MacroNameLoc(macroLoc),
+      LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
+      GenericArgs(genericArgs),
+      ArgList(args ? args
+                   : ArgumentList::createImplicit(dc->getASTContext(), {})) {
+  Bits.MacroExpansionDecl.Discriminator = InvalidDiscriminator;
+}
+
 SourceRange MacroExpansionDecl::getSourceRange() const {
   SourceLoc endLoc;
-  if (ArgList)
-    endLoc = ArgList->getEndLoc();
+  if (auto argsEndList = ArgList->getEndLoc())
+    endLoc = argsEndList;
   else if (RightAngleLoc.isValid())
     endLoc = RightAngleLoc;
   else
@@ -10171,6 +10187,23 @@ unsigned MacroExpansionDecl::getDiscriminator() const {
   return getRawDiscriminator();
 }
 
+void MacroExpansionDecl::forEachExpandedExprOrStmt(
+    ExprOrStmtExpansionCallback callback) const {
+  auto mutableThis = const_cast<MacroExpansionDecl *>(this);
+  auto bufferID = evaluateOrDefault(
+      getASTContext().evaluator,
+      ExpandMacroExpansionDeclRequest{mutableThis}, {});
+  auto &sourceMgr = getASTContext().SourceMgr;
+  auto *moduleDecl = getModuleContext();
+  if (!bufferID)
+    return;
+  auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
+  auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+  for (auto node : sourceFile->getTopLevelItems())
+    if (node.is<Expr *>() || node.is<Stmt *>())
+      callback(node);
+}
+
 NominalTypeDecl *
 ValueDecl::getRuntimeDiscoverableAttrTypeDecl(CustomAttr *attr) const {
   auto &ctx = getASTContext();

lib/AST/Expr.cpp

@@ -2745,6 +2745,23 @@ unsigned MacroExpansionExpr::getDiscriminator() const {
   return getRawDiscriminator();
 }
 
+MacroExpansionDecl *MacroExpansionExpr::createSubstituteDecl() const {
+  auto dc = DC;
+  if (auto *tlcd = dyn_cast_or_null<TopLevelCodeDecl>(dc->getAsDecl()))
+    dc = tlcd->getDeclContext();
+  return new (DC->getASTContext()) MacroExpansionDecl(
+      dc, 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)";