Merge pull request #79980 from hborla/closure-body-macro

[Macros] Implement support for function body macros on closures.

include/swift/AST/ASTMangler.h

@@ -412,10 +412,18 @@ public:
   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);
 

include/swift/AST/AnyFunctionRef.h

@@ -243,6 +243,54 @@ public:
     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;
    }

include/swift/AST/DiagnosticsSema.def

@@ -7843,6 +7843,9 @@ ERROR(conformance_macro,none,
 ERROR(experimental_macro,none,
       "macro %0 is experimental",
       (DeclName))
+ERROR(experimental_closure_body_macro,none,
+      "function body macros on closures is experimental",
+      ())
 
 ERROR(macro_resolve_circular_reference, none,
       "circular reference resolving %select{freestanding|attached}0 macro %1",

include/swift/AST/Expr.h

@@ -4305,6 +4305,8 @@ public:
   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 @@ public:
     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.
   ///

include/swift/AST/TypeCheckRequests.h

@@ -4708,7 +4708,7 @@ public:
 
 class ExpandBodyMacroRequest
     : public SimpleRequest<ExpandBodyMacroRequest,
-                           std::optional<unsigned>(AbstractFunctionDecl *),
+                           std::optional<unsigned>(AnyFunctionRef),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -4717,7 +4717,7 @@ private:
   friend SimpleRequest;
 
   std::optional<unsigned> evaluate(Evaluator &evaluator,
-                                   AbstractFunctionDecl *fn) const;
+                                   AnyFunctionRef fn) const;
 
 public:
   bool isCached() const { return true; }

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 *),

include/swift/Basic/Features.def

@@ -512,6 +512,9 @@ EXPERIMENTAL_FEATURE(ConcurrencySyntaxSugar, true)
 /// Allow declaration of compile-time values
 EXPERIMENTAL_FEATURE(CompileTimeValues, true)
 
+/// Allow function body macros applied to closures.
+EXPERIMENTAL_FEATURE(ClosureBodyMacro, true)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE