[Macros] Support function body macros on closures.

hborla · hborla · commit 31d4f661fc06 · 2025-03-18T20:50:47.000-07:00
diff --git a/include/swift/AST/ASTMangler.h b/include/swift/AST/ASTMangler.h
@@ -408,10 +408,18 @@ class ASTMangler : public Mangler {
   std::string mangleMacroExpansion(const FreestandingMacroExpansion *expansion);
   std::string mangleAttachedMacroExpansion(
       const Decl *decl, CustomAttr *attr, MacroRole role);
+  std::string mangleAttachedMacroExpansion(
+      ClosureExpr *attachedTo, CustomAttr *attr, MacroDecl *macro);
 
   void appendMacroExpansion(const FreestandingMacroExpansion *expansion);
   void appendMacroExpansionContext(SourceLoc loc, DeclContext *origDC,
-                                   const FreestandingMacroExpansion *expansion);
+                                   Identifier macroName,
+                                   unsigned discriminator);
+
+  void appendMacroExpansion(ClosureExpr *attachedTo,
+                            CustomAttr *attr,
+                            MacroDecl *macro);
+
   void appendMacroExpansionOperator(
       StringRef macroName, MacroRole role, unsigned discriminator);
 
diff --git a/include/swift/AST/AnyFunctionRef.h b/include/swift/AST/AnyFunctionRef.h
@@ -243,6 +243,54 @@ class AnyFunctionRef {
     llvm_unreachable("unexpected AnyFunctionRef representation");
   }
 
+  DeclAttributes getDeclAttributes() const {
+    if (auto afd = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
+      return afd->getExpandedAttrs();
+    }
+
+    if (auto ace = TheFunction.dyn_cast<AbstractClosureExpr *>()) {
+      if (auto *ce = dyn_cast<ClosureExpr>(ace)) {
+        return ce->getAttrs();
+      }
+    }
+
+    return DeclAttributes();
+  }
+
+  MacroDecl *getResolvedMacro(CustomAttr *attr) const {
+    if (auto afd = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
+      return afd->getResolvedMacro(attr);
+    }
+
+    if (auto ace = TheFunction.dyn_cast<AbstractClosureExpr *>()) {
+      if (auto *ce = dyn_cast<ClosureExpr>(ace)) {
+        return ce->getResolvedMacro(attr);
+      }
+    }
+
+    return nullptr;
+  }
+
+  using MacroCallback = llvm::function_ref<void(CustomAttr *, MacroDecl *)>;
+
+  void
+  forEachAttachedMacro(MacroRole role,
+                       MacroCallback macroCallback) const {
+    auto attrs = getDeclAttributes();
+    for (auto customAttrConst : attrs.getAttributes<CustomAttr>()) {
+      auto customAttr = const_cast<CustomAttr *>(customAttrConst);
+      auto *macroDecl = getResolvedMacro(customAttr);
+
+      if (!macroDecl)
+        continue;
+
+      if (!macroDecl->getMacroRoles().contains(role))
+        continue;
+
+      macroCallback(customAttr, macroDecl);
+    }
+  }
+
   friend bool operator==(AnyFunctionRef lhs, AnyFunctionRef rhs) {
      return lhs.TheFunction == rhs.TheFunction;
    }
diff --git a/include/swift/AST/Expr.h b/include/swift/AST/Expr.h
@@ -4305,6 +4305,8 @@ class ClosureExpr : public AbstractClosureExpr {
   BraceStmt *getBody() const { return Body; }
   void setBody(BraceStmt *S) { Body = S; }
 
+  BraceStmt *getExpandedBody();
+
   DeclAttributes &getAttrs() { return Attributes; }
   const DeclAttributes &getAttrs() const { return Attributes; }
 
@@ -4422,6 +4424,10 @@ class ClosureExpr : public AbstractClosureExpr {
     return ExplicitResultTypeAndBodyState.getPointer()->getTypeRepr();
   }
 
+  /// Returns the resolved macro for the given custom attribute
+  /// attached to this closure expression.
+  MacroDecl *getResolvedMacro(CustomAttr *customAttr);
+
   /// Determine whether the closure has a single expression for its
   /// body.
   ///
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -4708,7 +4708,7 @@ class ExpandPreambleMacroRequest
 
 class ExpandBodyMacroRequest
     : public SimpleRequest<ExpandBodyMacroRequest,
-                           std::optional<unsigned>(AbstractFunctionDecl *),
+                           std::optional<unsigned>(AnyFunctionRef),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -4717,7 +4717,7 @@ class ExpandBodyMacroRequest
   friend SimpleRequest;
 
   std::optional<unsigned> evaluate(Evaluator &evaluator,
-                                   AbstractFunctionDecl *fn) const;
+                                   AnyFunctionRef fn) const;
 
 public:
   bool isCached() const { return true; }
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -531,7 +531,7 @@ SWIFT_REQUEST(TypeChecker, ExpandPreambleMacroRequest,
               ArrayRef<unsigned>(AbstractFunctionDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandBodyMacroRequest,
-              std::optional<unsigned>(AbstractFunctionDecl *),
+              std::optional<unsigned>(AnyFunctionRef),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, LocalDiscriminatorsRequest,
               unsigned(DeclContext *),
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -4759,7 +4759,8 @@ static Identifier encodeLocalPrecheckedDiscriminator(
 
 void ASTMangler::appendMacroExpansionContext(
     SourceLoc loc, DeclContext *origDC,
-    const FreestandingMacroExpansion *expansion
+    Identifier macroName,
+    unsigned macroDiscriminator
 ) {
   origDC = MacroDiscriminatorContext::getInnermostMacroContext(origDC);
   BaseEntitySignature nullBase(nullptr);
@@ -4768,9 +4769,9 @@ void ASTMangler::appendMacroExpansionContext(
     if (auto outermostLocalDC = getOutermostLocalContext(origDC)) {
       auto innermostNonlocalDC = outermostLocalDC->getParent();
       appendContext(innermostNonlocalDC, nullBase, StringRef());
-      Identifier name = expansion->getMacroName().getBaseIdentifier();
+      Identifier name = macroName;
       ASTContext &ctx = origDC->getASTContext();
-      unsigned discriminator = expansion->getDiscriminator();
+      unsigned discriminator = macroDiscriminator;
       name = encodeLocalPrecheckedDiscriminator(ctx, name, discriminator);
       appendIdentifier(name.str());
     } else {
@@ -4875,7 +4876,10 @@ void ASTMangler::appendMacroExpansionContext(
     return appendMacroExpansionLoc();
 
   // Append our own context and discriminator.
-  appendMacroExpansionContext(outerExpansionLoc, origDC, expansion);
+  appendMacroExpansionContext(
+      outerExpansionLoc, origDC,
+      macroName,
+      macroDiscriminator);
   appendMacroExpansionOperator(
       baseName.userFacingName(), role, discriminator);
 }
@@ -4902,16 +4906,14 @@ void ASTMangler::appendMacroExpansionOperator(
 }
 
 static StringRef getPrivateDiscriminatorIfNecessary(
-      const MacroExpansionExpr *expansion) {
-  auto dc = MacroDiscriminatorContext::getInnermostMacroContext(
-      expansion->getDeclContext());
-  auto decl = dc->getAsDecl();
+      const DeclContext *macroDC) {
+  auto decl = macroDC->getAsDecl();
   if (decl && !decl->isOutermostPrivateOrFilePrivateScope())
     return StringRef();
 
   // Mangle non-local private declarations with a textual discriminator
   // based on their enclosing file.
-  auto topLevelSubcontext = dc->getModuleScopeContext();
+  auto topLevelSubcontext = macroDC->getModuleScopeContext();
   SourceFile *sf = dyn_cast<SourceFile>(topLevelSubcontext);
   if (!sf)
     return StringRef();
@@ -4927,6 +4929,13 @@ static StringRef getPrivateDiscriminatorIfNecessary(
   return discriminator.str();
 }
 
+static StringRef getPrivateDiscriminatorIfNecessary(
+    const MacroExpansionExpr *expansion) {
+  auto dc = MacroDiscriminatorContext::getInnermostMacroContext(
+      expansion->getDeclContext());
+  return getPrivateDiscriminatorIfNecessary(dc);
+}
+
 static StringRef getPrivateDiscriminatorIfNecessary(
     const FreestandingMacroExpansion *expansion) {
   switch (expansion->getFreestandingMacroKind()) {
@@ -4943,7 +4952,8 @@ void
 ASTMangler::appendMacroExpansion(const FreestandingMacroExpansion *expansion) {
   appendMacroExpansionContext(expansion->getPoundLoc(),
                               expansion->getDeclContext(),
-                              expansion);
+                              expansion->getMacroName().getBaseIdentifier(),
+                              expansion->getDiscriminator());
   auto privateDiscriminator = getPrivateDiscriminatorIfNecessary(expansion);
   if (!privateDiscriminator.empty()) {
     appendIdentifier(privateDiscriminator);
@@ -4955,6 +4965,42 @@ ASTMangler::appendMacroExpansion(const FreestandingMacroExpansion *expansion) {
       expansion->getDiscriminator());
 }
 
+void ASTMangler::appendMacroExpansion(ClosureExpr *attachedTo,
+                                      CustomAttr *attr,
+                                      MacroDecl *macro) {
+  auto &ctx = attachedTo->getASTContext();
+  auto discriminator =
+      ctx.getNextMacroDiscriminator(attachedTo,
+                                    macro->getBaseName());
+
+  appendMacroExpansionContext(
+      attr->getLocation(),
+      attachedTo,
+      macro->getBaseName().getIdentifier(),
+      discriminator);
+
+  auto privateDiscriminator =
+      getPrivateDiscriminatorIfNecessary(attachedTo);
+  if (!privateDiscriminator.empty()) {
+    appendIdentifier(privateDiscriminator);
+    appendOperator("Ll");
+  }
+
+  appendMacroExpansionOperator(
+      macro->getBaseName().userFacingName(),
+      MacroRole::Body,
+      discriminator);
+}
+
+std::string
+ASTMangler::mangleAttachedMacroExpansion(ClosureExpr *attachedTo,
+                                         CustomAttr *attr,
+                                         MacroDecl *macro) {
+  beginMangling();
+  appendMacroExpansion(attachedTo, attr, macro);
+  return finalize();
+}
+
 std::string
 ASTMangler::mangleMacroExpansion(const FreestandingMacroExpansion *expansion) {
   beginMangling();
diff --git a/lib/AST/ASTScopeCreation.cpp b/lib/AST/ASTScopeCreation.cpp
@@ -309,13 +309,22 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
       break;
     }
     case MacroRole::Body: {
-      // Use the end location of the function decl itself as the parentLoc
-      // for the new function body scope. This is different from the end
-      // location of the original source range, which is after the end of the
-      // function decl.
       auto expansion = SF->getMacroExpansion();
-      parentLoc = expansion.getEndLoc();
-      bodyForDecl = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
+      if (expansion.is<Decl *>()) {
+        // Use the end location of the function decl itself as the parentLoc
+        // for the new function body scope. This is different from the end
+        // location of the original source range, which is after the end of the
+        // function decl.
+        bodyForDecl = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
+        parentLoc = expansion.getEndLoc();
+        break;
+      }
+
+      // Otherwise, we have a closure body macro.
+      auto insertionRange = SF->getMacroInsertionRange();
+      parentLoc = insertionRange.End;
+      if (insertionRange.Start != insertionRange.End)
+        parentLoc = parentLoc.getAdvancedLoc(-1);
       break;
     }
     }
diff --git a/lib/AST/Expr.cpp b/lib/AST/Expr.cpp
@@ -18,6 +18,7 @@
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Statistic.h"
 #include "swift/Basic/Unicode.h"
+#include "swift/Basic/SourceManager.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/Decl.h" // FIXME: Bad dependency
@@ -2016,6 +2017,29 @@ BraceStmt * AbstractClosureExpr::getBody() const {
   llvm_unreachable("Unknown closure expression");
 }
 
+BraceStmt *ClosureExpr::getExpandedBody() {
+  auto &ctx = getASTContext();
+
+  // Expand a body macro, if there is one.
+  BraceStmt *macroExpandedBody = nullptr;
+  if (auto bufferID = evaluateOrDefault(
+          ctx.evaluator,
+          ExpandBodyMacroRequest{this},
+          std::nullopt)) {
+    CharSourceRange bufferRange = ctx.SourceMgr.getRangeForBuffer(*bufferID);
+    auto bufferStart = bufferRange.getStart();
+    auto module = getParentModule();
+    auto macroSourceFile = module->getSourceFileContainingLocation(bufferStart);
+
+    if (macroSourceFile->getTopLevelItems().size() == 1) {
+      auto stmt = macroSourceFile->getTopLevelItems()[0].dyn_cast<Stmt *>();
+      macroExpandedBody = dyn_cast<BraceStmt>(stmt);
+    }
+  }
+
+  return macroExpandedBody;
+}
+
 bool AbstractClosureExpr::bodyHasExplicitReturnStmt() const {
   return AnyFunctionRef(const_cast<AbstractClosureExpr *>(this))
       .bodyHasExplicitReturnStmt();
@@ -2175,6 +2199,17 @@ void ClosureExpr::setExplicitResultType(Type ty) {
       ->setType(MetatypeType::get(ty));
 }
 
+MacroDecl *
+ClosureExpr::getResolvedMacro(CustomAttr *customAttr) {
+  auto &ctx = getASTContext();
+  auto declRef = evaluateOrDefault(
+      ctx.evaluator,
+      ResolveMacroRequest{customAttr, this},
+      ConcreteDeclRef());
+
+  return dyn_cast_or_null<MacroDecl>(declRef.getDecl());
+}
+
 FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)
 
 void AutoClosureExpr::setBody(Expr *E) {
diff --git a/lib/AST/TypeCheckRequests.cpp b/lib/AST/TypeCheckRequests.cpp
@@ -2681,19 +2681,17 @@ void ExpandPreambleMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
 //----------------------------------------------------------------------------//
 
 void ExpandBodyMacroRequest::diagnoseCycle(DiagnosticEngine &diags) const {
-  auto decl = std::get<0>(getStorage());
-  diags.diagnose(decl->getLoc(),
-                 diag::macro_expand_circular_reference_entity,
-                 "body",
-                 decl->getName());
+  auto fn = std::get<0>(getStorage());
+  diags.diagnose(fn.getLoc(),
+                 diag::macro_expand_circular_reference_unnamed,
+                 "body");
 }
 
 void ExpandBodyMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
-  auto decl = std::get<0>(getStorage());
-  diags.diagnose(decl->getLoc(),
-                 diag::macro_expand_circular_reference_entity_through,
-                 "body",
-                 decl->getName());
+  auto fn = std::get<0>(getStorage());
+  diags.diagnose(fn.getLoc(),
+                 diag::macro_expand_circular_reference_unnamed_through,
+                 "body");
 }
 
 //----------------------------------------------------------------------------//
diff --git a/lib/ASTGen/Sources/MacroEvaluation/Macros.swift b/lib/ASTGen/Sources/MacroEvaluation/Macros.swift
@@ -589,7 +589,7 @@ func expandAttachedMacro(
     let declarationNode = findSyntaxNodeInSourceFile(
       sourceFilePtr: declarationSourceFilePtr,
       sourceLocationPtr: declarationSourceLocPointer,
-      type: DeclSyntax.self
+      type: Syntax.self
     )
   else {
     return 1
@@ -657,7 +657,7 @@ func expandAttachedMacroImpl(
   customAttrSourceFilePtr: UnsafePointer<ExportedSourceFile>,
   customAttrNode: AttributeSyntax,
   declarationSourceFilePtr: UnsafePointer<ExportedSourceFile>,
-  attachedTo declarationNode: DeclSyntax,
+  attachedTo declarationNode: Syntax,
   parentDeclSourceFilePtr: UnsafePointer<ExportedSourceFile>?,
   parentDeclNode: DeclSyntax?
 ) -> String? {
@@ -689,7 +689,7 @@ func expandAttachedMacroImpl(
   )!
 
   let declSyntax = PluginMessage.Syntax(
-    syntax: Syntax(declarationNode),
+    syntax: declarationNode,
     in: declarationSourceFilePtr
   )!
 
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -4478,7 +4478,7 @@ NodePointer Demangler::demangleMacroExpansion() {
     context = popContext();
   NodePointer discriminator = demangleIndexAsNode();
   NodePointer result;
-  if (isAttached) {
+  if (isAttached && attachedName) {
     result = createWithChildren(
         kind, context, attachedName, macroName, discriminator);
   } else {
diff --git a/lib/Macros/Sources/SwiftMacros/StartTaskMacro.swift b/lib/Macros/Sources/SwiftMacros/StartTaskMacro.swift
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
diff --git a/lib/Sema/TypeCheckStmt.cpp b/lib/Sema/TypeCheckStmt.cpp
diff --git a/test/Macros/start_task.swift b/test/Macros/start_task.swift
