[Macros] Enable freestanding macros at module scope in script mode

include/swift/AST/Decl.h

@@ -8449,7 +8449,7 @@ class MissingDecl : public Decl {
   /// \c unexpandedMacro contains the macro reference and the base declaration
   /// where the macro expansion applies.
   struct {
-    llvm::PointerUnion<MacroExpansionDecl *, CustomAttr *> macroRef;
+    llvm::PointerUnion<FreestandingMacroExpansion *, CustomAttr *> macroRef;
     Decl *baseDecl;
   } unexpandedMacro;
 
@@ -8473,7 +8473,7 @@ class MissingDecl : public Decl {
 
   static MissingDecl *
   forUnexpandedMacro(
-      llvm::PointerUnion<MacroExpansionDecl *, CustomAttr *> macroRef,
+      llvm::PointerUnion<FreestandingMacroExpansion *, CustomAttr *> macroRef,
       Decl *baseDecl) {
     auto &ctx = baseDecl->getASTContext();
     auto *dc = baseDecl->getDeclContext();

include/swift/AST/DiagnosticsSema.def

@@ -7244,9 +7244,6 @@ 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",
-      ())
 ERROR(invalid_macro_role_for_macro_syntax,none,
       "invalid macro role for %{a freestanding|an attached}0 macro",
       (unsigned))

lib/AST/Decl.cpp

@@ -408,7 +408,25 @@ void Decl::visitAuxiliaryDecls(
   }
 
   if (visitFreestandingExpanded) {
-    if (auto *med = dyn_cast<MacroExpansionDecl>(mutableThis)) {
+    Decl *thisDecl = mutableThis;
+
+    // If this is a top-level code decl consisting of a macro expansion
+    // expression that substituted with a macro expansion declaration, use
+    // that instead.
+    if (auto *tlcd = dyn_cast<TopLevelCodeDecl>(thisDecl)) {
+      if (auto body = tlcd->getBody()) {
+        if (body->getNumElements() == 1) {
+          if (auto expr = body->getFirstElement().dyn_cast<Expr *>()) {
+            if (auto expansion = dyn_cast<MacroExpansionExpr>(expr)) {
+              if (auto substitute = expansion->getSubstituteDecl())
+                thisDecl = substitute;
+            }
+          }
+        }
+      }
+    }
+
+    if (auto *med = dyn_cast<MacroExpansionDecl>(thisDecl)) {
       if (auto bufferID = evaluateOrDefault(
               ctx.evaluator, ExpandMacroExpansionDeclRequest{med}, {})) {
         auto startLoc = sourceMgr.getLocForBufferStart(*bufferID);
@@ -10492,6 +10510,45 @@ bool swift::isMacroSupported(MacroRole role, ASTContext &ctx) {
 void MissingDecl::forEachMacroExpandedDecl(MacroExpandedDeclCallback callback) {
   auto macroRef = unexpandedMacro.macroRef;
   auto *baseDecl = unexpandedMacro.baseDecl;
+
+  // If the macro itself is a macro expansion expression, it should come with
+  // a top-level code declaration that we can use for resolution. For such
+  // cases, resolve the macro to determine whether it is a declaration or
+  // code-item macro, meaning that it can produce declarations. In such cases,
+  // expand the macro and use its substituted declaration (a MacroExpansionDecl)
+  // instead.
+  if (auto freestanding = macroRef.dyn_cast<FreestandingMacroExpansion *>()) {
+    if (auto expr = dyn_cast<MacroExpansionExpr>(freestanding)) {
+      bool replacedWithDecl = false;
+      if (auto tlcd = dyn_cast_or_null<TopLevelCodeDecl>(baseDecl)) {
+        ASTContext &ctx = tlcd->getASTContext();
+        if (auto macro = evaluateOrDefault(
+                ctx.evaluator,
+                ResolveMacroRequest{macroRef, tlcd->getDeclContext()},
+                nullptr)) {
+          auto macroDecl = cast<MacroDecl>(macro.getDecl());
+          auto roles = macroDecl->getMacroRoles();
+          if (roles.contains(MacroRole::Declaration) ||
+              roles.contains(MacroRole::CodeItem)) {
+            (void)evaluateOrDefault(ctx.evaluator,
+                                    ExpandMacroExpansionExprRequest{expr},
+                                    llvm::None);
+            if (auto substituted = expr->getSubstituteDecl()) {
+              macroRef = substituted;
+              baseDecl = substituted;
+              replacedWithDecl = true;
+            }
+          }
+        }
+      }
+
+      // If we didn't end up replacing the macro expansion expression with
+      // a declaration, we're done.
+      if (!replacedWithDecl)
+        return;
+    }
+  }
+
   if (!macroRef || !baseDecl)
     return;
   auto *module = getModuleContext();
@@ -10502,9 +10559,12 @@ void MissingDecl::forEachMacroExpandedDecl(MacroExpandedDeclCallback callback) {
         : 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;
+    if (auto *med = macroRef.dyn_cast<FreestandingMacroExpansion *>()) {
+      auto medAsDecl = dyn_cast<MacroExpansionDecl>(med);
+      auto medAsExpr = dyn_cast<MacroExpansionExpr>(med);
+      if ((!medAsDecl || medAsDecl != sf->getMacroExpansion().dyn_cast<Decl *>()) &&
+          (!medAsExpr || medAsExpr != sf->getMacroExpansion().dyn_cast<Expr *>()))
+        return;
     } else if (auto *attr = macroRef.dyn_cast<CustomAttr *>()) {
      if (attr != sf->getAttachedMacroAttribute())
        return;

lib/AST/Module.cpp

@@ -179,9 +179,9 @@ class swift::SourceLookupCache {
   /// Top-level macros that produce arbitrary names.
   SmallVector<MissingDecl *, 4> TopLevelArbitraryMacros;
 
-  SmallVector<Decl *, 4> MayHaveAuxiliaryDecls;
+  SmallVector<llvm::PointerUnion<Decl *, MacroExpansionExpr *>, 4>
+      MayHaveAuxiliaryDecls;
   void populateAuxiliaryDeclCache();
-
   SourceLookupCache(ASTContext &ctx);
 
 public:
@@ -238,10 +238,30 @@ SourceLookupCache &SourceFile::getCache() const {
   return *Cache;
 }
 
+static Expr *getAsExpr(Decl *decl) { return nullptr; }
+static Decl *getAsDecl(Decl *decl) { return decl; }
+
+static Expr *getAsExpr(ASTNode node) { return node.dyn_cast<Expr *>(); }
+static Decl *getAsDecl(ASTNode node) { return node.dyn_cast<Decl *>(); }
+
 template<typename Range>
-void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
+void SourceLookupCache::addToUnqualifiedLookupCache(Range items,
                                                     bool onlyOperators) {
-  for (Decl *D : decls) {
+  for (auto item : items) {
+    // In script mode, we'll see macro expansion expressions for freestanding
+    // macros.
+    if (Expr *E = getAsExpr(item)) {
+      if (auto MEE = dyn_cast<MacroExpansionExpr>(E)) {
+        if (!onlyOperators)
+          MayHaveAuxiliaryDecls.push_back(MEE);
+      }
+      continue;
+    }
+
+    Decl *D = getAsDecl(item);
+    if (!D)
+      continue;
+
     if (auto *VD = dyn_cast<ValueDecl>(D)) {
       if (onlyOperators ? VD->isOperator() : VD->hasName()) {
         // Cache the value under both its compound name and its full name.
@@ -280,6 +300,10 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
     else if (auto *MED = dyn_cast<MacroExpansionDecl>(D)) {
       if (!onlyOperators)
         MayHaveAuxiliaryDecls.push_back(MED);
+    } else if (auto TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
+      if (auto body = TLCD->getBody()){
+        addToUnqualifiedLookupCache(body->getElements(), onlyOperators);
+      }
     }
   }
 }
@@ -335,75 +359,113 @@ void SourceLookupCache::addToMemberCache(Range decls) {
 }
 
 void SourceLookupCache::populateAuxiliaryDeclCache() {
-  using MacroRef = llvm::PointerUnion<MacroExpansionDecl *, CustomAttr *>;
-  for (auto *decl : MayHaveAuxiliaryDecls) {
+  using MacroRef = llvm::PointerUnion<FreestandingMacroExpansion *, CustomAttr *>;
+  for (auto item : MayHaveAuxiliaryDecls) {
+    TopLevelCodeDecl *topLevelCodeDecl = nullptr;
+
     // Gather macro-introduced peer names.
     llvm::SmallDenseMap<MacroRef, llvm::SmallVector<DeclName, 2>>
         introducedNames;
 
-    // This code deliberately avoids `forEachAttachedMacro`, because it
-    // will perform overload resolution and possibly invoke unqualified
-    // lookup for macro arguments, which will recursively populate the
-    // auxiliary decl cache and cause request cycles.
-    //
-    // We do not need a fully resolved macro until expansion. Instead, we
-    // conservatively consider peer names for all macro declarations with a
-    // custom attribute name. Unqualified lookup for that name will later
-    // invoke expansion of the macro, and will yield no results if the resolved
-    // macro does not produce the requested name, so the only impact is possibly
-    // expanding earlier than needed / unnecessarily looking in the top-level
-    // auxiliary decl cache.
-    for (auto attrConst : decl->getAttrs().getAttributes<CustomAttr>()) {
-      auto *attr = const_cast<CustomAttr *>(attrConst);
-      UnresolvedMacroReference macroRef(attr);
-      bool introducesArbitraryNames = false;
-      namelookup::forEachPotentialResolvedMacro(
-          decl->getDeclContext()->getModuleScopeContext(),
-          macroRef.getMacroName(), MacroRole::Peer,
-          [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
-            // First check for arbitrary names.
-            if (roleAttr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
-              introducesArbitraryNames = true;
-            }
+    /// Introduce names for a freestanding macro.
+    auto introduceNamesForFreestandingMacro =
+      [&](FreestandingMacroExpansion *macroRef, Decl *decl, MacroRole role) {
+        bool introducesArbitraryNames = false;
+        namelookup::forEachPotentialResolvedMacro(
+            decl->getDeclContext()->getModuleScopeContext(),
+            macroRef->getMacroName(), role,
+            [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+              // First check for arbitrary names.
+              if (roleAttr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
+                introducesArbitraryNames = true;
+              }
 
-            macro->getIntroducedNames(MacroRole::Peer,
-                                      dyn_cast<ValueDecl>(decl),
-                                      introducedNames[attr]);
-          });
+              macro->getIntroducedNames(role,
+                                        /*attachedTo*/ nullptr,
+                                        introducedNames[macroRef]);
+            });
 
-      // Record this macro where appropriate.
-      if (introducesArbitraryNames)
-        TopLevelArbitraryMacros.push_back(MissingDecl::forUnexpandedMacro(attr, decl));
+        return introducesArbitraryNames;
+      };
+
+    // Handle macro expansion expressions, which show up in when we have
+    // freestanding macros in "script" mode.
+    if (auto expr = item.dyn_cast<MacroExpansionExpr *>()) {
+      topLevelCodeDecl = dyn_cast<TopLevelCodeDecl>(expr->getDeclContext());
+      if (topLevelCodeDecl) {
+        bool introducesArbitraryNames = false;
+        if (introduceNamesForFreestandingMacro(
+                expr, topLevelCodeDecl, MacroRole::Declaration))
+          introducesArbitraryNames = true;
+
+        if (introduceNamesForFreestandingMacro(
+                expr, topLevelCodeDecl, MacroRole::CodeItem))
+          introducesArbitraryNames = true;
+
+        // Record this macro if it introduces arbitrary names.
+        if (introducesArbitraryNames) {
+          TopLevelArbitraryMacros.push_back(
+              MissingDecl::forUnexpandedMacro(expr, topLevelCodeDecl));
+        }
+      }
     }
 
-    if (auto *med = dyn_cast<MacroExpansionDecl>(decl)) {
-      UnresolvedMacroReference macroRef(med);
-      bool introducesArbitraryNames = false;
-      namelookup::forEachPotentialResolvedMacro(
-          decl->getDeclContext()->getModuleScopeContext(),
-          macroRef.getMacroName(), MacroRole::Declaration,
-          [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
-            // First check for arbitrary names.
-            if (roleAttr->hasNameKind(MacroIntroducedDeclNameKind::Arbitrary)) {
-              introducesArbitraryNames = true;
-            }
+    auto *decl = item.dyn_cast<Decl *>();
+    if (decl) {
+      // This code deliberately avoids `forEachAttachedMacro`, because it
+      // will perform overload resolution and possibly invoke unqualified
+      // lookup for macro arguments, which will recursively populate the
+      // auxiliary decl cache and cause request cycles.
+      //
+      // We do not need a fully resolved macro until expansion. Instead, we
+      // conservatively consider peer names for all macro declarations with a
+      // custom attribute name. Unqualified lookup for that name will later
+      // invoke expansion of the macro, and will yield no results if the resolved
+      // macro does not produce the requested name, so the only impact is possibly
+      // expanding earlier than needed / unnecessarily looking in the top-level
+      // auxiliary decl cache.
+      for (auto attrConst : decl->getAttrs().getAttributes<CustomAttr>()) {
+        auto *attr = const_cast<CustomAttr *>(attrConst);
+        UnresolvedMacroReference macroRef(attr);
+        bool introducesArbitraryNames = false;
+        namelookup::forEachPotentialResolvedMacro(
+            decl->getDeclContext()->getModuleScopeContext(),
+            macroRef.getMacroName(), MacroRole::Peer,
+            [&](MacroDecl *macro, const MacroRoleAttr *roleAttr) {
+              // First check for arbitrary names.
+              if (roleAttr->hasNameKind(
+                      MacroIntroducedDeclNameKind::Arbitrary)) {
+                introducesArbitraryNames = true;
+              }
+
+              macro->getIntroducedNames(MacroRole::Peer,
+                                        dyn_cast<ValueDecl>(decl),
+                                        introducedNames[attr]);
+            });
+
+        // Record this macro where appropriate.
+        if (introducesArbitraryNames)
+          TopLevelArbitraryMacros.push_back(
+              MissingDecl::forUnexpandedMacro(attr, decl));
+      }
+    }
 
-            macro->getIntroducedNames(MacroRole::Declaration,
-                                      /*attachedTo*/ nullptr,
-                                      introducedNames[med]);
-          });
+    if (auto *med = dyn_cast_or_null<MacroExpansionDecl>(decl)) {
+      bool introducesArbitraryNames =
+          introduceNamesForFreestandingMacro(med, decl, MacroRole::Declaration);
 
-      // Record this macro where appropriate.
+      // Note whether this macro produces arbitrary names.
       if (introducesArbitraryNames)
         TopLevelArbitraryMacros.push_back(MissingDecl::forUnexpandedMacro(med, decl));
     }
 
     // Add macro-introduced names to the top-level auxiliary decl cache as
     // unexpanded decls represented by a MissingDecl.
+    auto anchorDecl = decl ? decl : topLevelCodeDecl;
     for (auto macroNames : introducedNames) {
       auto macroRef = macroNames.getFirst();
       for (auto name : macroNames.getSecond()) {
-        auto *placeholder = MissingDecl::forUnexpandedMacro(macroRef, decl);
+        auto *placeholder = MissingDecl::forUnexpandedMacro(macroRef, anchorDecl);
         name.addToLookupTable(TopLevelAuxiliaryDecls, placeholder);
       }
     }
@@ -421,7 +483,7 @@ SourceLookupCache::SourceLookupCache(const SourceFile &SF)
 {
   FrontendStatsTracer tracer(SF.getASTContext().Stats,
                              "source-file-populate-cache");
-  addToUnqualifiedLookupCache(SF.getTopLevelDecls(), false);
+  addToUnqualifiedLookupCache(SF.getTopLevelItems(), false);
   addToUnqualifiedLookupCache(SF.getHoistedDecls(), false);
 }
 
@@ -432,7 +494,7 @@ SourceLookupCache::SourceLookupCache(const ModuleDecl &M)
                              "module-populate-cache");
   for (const FileUnit *file : M.getFiles()) {
     auto *SF = cast<SourceFile>(file);
-    addToUnqualifiedLookupCache(SF->getTopLevelDecls(), false);
+    addToUnqualifiedLookupCache(SF->getTopLevelItems(), false);
     addToUnqualifiedLookupCache(SF->getHoistedDecls(), false);
 
     if (auto *SFU = file->getSynthesizedFile()) {

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -3926,11 +3926,5 @@ ExpandMacroExpansionDeclRequest::evaluate(Evaluator &evaluator,
       !roles.contains(MacroRole::CodeItem))
     return llvm::None;
 
-  // For now, restrict global freestanding macros in script mode.
-  if (dc->isModuleScopeContext() &&
-      dc->getParentSourceFile()->isScriptMode()) {
-    MED->diagnose(diag::global_freestanding_macro_script);
-  }
-
   return expandFreestandingMacro(MED);
 }

lib/Sema/TypeCheckMacros.cpp

@@ -469,8 +469,7 @@ ExpandMacroExpansionExprRequest::evaluate(Evaluator &evaluator,
   else if (macro->getMacroRoles().contains(MacroRole::Declaration) ||
            macro->getMacroRoles().contains(MacroRole::CodeItem)) {
     if (!mee->getSubstituteDecl()) {
-      auto *med = mee->createSubstituteDecl();
-      TypeChecker::typeCheckDecl(med);
+      (void)mee->createSubstituteDecl();
     }
     // Return the expanded buffer ID.
     return evaluateOrDefault(