[Macros] Expand peer macros.

Author: hborla

lib/Sema/LookupVisibleDecls.cpp

@@ -611,6 +611,14 @@ static void synthesizeMemberDeclsForLookup(NominalTypeDecl *NTD,
                     ExpandSynthesizedMemberMacroRequest{NTD},
                     false);
 
+  // Expand peer macros.
+  for (auto *member : NTD->getMembers()) {
+    (void)evaluateOrDefault(
+        ctx.evaluator,
+        ExpandPeerMacroRequest{member},
+        false);
+  }
+
   synthesizePropertyWrapperVariables(NTD);
 }
 

lib/Sema/TypeCheckDecl.cpp

@@ -2802,6 +2802,12 @@ static ArrayRef<Decl *> evaluateMembersRequest(
       nullptr);
 
   for (auto *member : idc->getMembers()) {
+    // Expand peer macros.
+    (void)evaluateOrDefault(
+        ctx.evaluator,
+        ExpandPeerMacroRequest{member},
+        false);
+
     if (auto *var = dyn_cast<VarDecl>(member)) {
       // The projected storage wrapper ($foo) might have
       // dynamically-dispatched accessors, so force them to be synthesized.

lib/Sema/TypeCheckMacros.cpp

@@ -386,6 +386,9 @@ static std::string adjustMacroExpansionBufferName(StringRef name) {
 
 bool ExpandMemberAttributeMacros::evaluate(Evaluator &evaluator,
                                            Decl *decl) const {
+  if (decl->isImplicit())
+    return false;
+
   auto *parentDecl = decl->getDeclContext()->getAsDecl();
   if (!parentDecl)
     return false;

lib/Sema/TypeCheckStorage.cpp

@@ -117,6 +117,12 @@ static void computeLoweredStoredProperties(NominalTypeDecl *decl,
   // Just walk over the members of the type, forcing backing storage
   // for lazy properties and property wrappers to be synthesized.
   for (auto *member : implDecl->getMembers()) {
+    // Expand peer macros.
+    (void)evaluateOrDefault(
+        ctx.evaluator,
+        ExpandPeerMacroRequest{member},
+        false);
+
     auto *var = dyn_cast<VarDecl>(member);
     if (!var || var->isStatic())
       continue;

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -611,6 +611,127 @@ public enum LeftHandOperandFinderMacro: ExpressionMacro {
   }
 }
 
+public struct AddCompletionHandler: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    // Only on functions at the moment. We could handle initializers as well
+    // with a bit of work.
+    guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+      throw CustomError.message("@addCompletionHandler only works on functions")
+    }
+
+    // This only makes sense for async functions.
+    if funcDecl.signature.effectSpecifiers?.asyncSpecifier == nil {
+      throw CustomError.message(
+        "@addCompletionHandler requires an async function"
+      )
+    }
+
+    // Form the completion handler parameter.
+    let resultType: TypeSyntax? = funcDecl.signature.output?.returnType.with(\.leadingTrivia, []).with(\.trailingTrivia, [])
+
+    let completionHandlerParam =
+      FunctionParameterSyntax(
+        firstName: .identifier("completionHandler"),
+        colon: .colonToken(trailingTrivia: .space),
+        type: "(\(resultType ?? "")) -> Void" as TypeSyntax
+      )
+
+    // Add the completion handler parameter to the parameter list.
+    let parameterList = funcDecl.signature.input.parameterList
+    let newParameterList: FunctionParameterListSyntax
+    if let lastParam = parameterList.last {
+      // We need to add a trailing comma to the preceding list.
+      newParameterList = parameterList.removingLast()
+        .appending(
+          lastParam.with(
+            \.trailingComma,
+            .commaToken(trailingTrivia: .space)
+          )
+        )
+        .appending(completionHandlerParam)
+    } else {
+      newParameterList = parameterList.appending(completionHandlerParam)
+    }
+
+    let callArguments: [String] = try parameterList.map { param in
+      guard let argName = param.secondName ?? param.firstName else {
+        throw CustomError.message(
+          "@addCompletionHandler argument must have a name"
+        )
+      }
+
+      if let paramName = param.firstName, paramName.text != "_" {
+        return "\(paramName.text): \(argName.text)"
+      }
+
+      return "\(argName.text)"
+    }
+
+    let call: ExprSyntax =
+      "\(funcDecl.identifier)(\(raw: callArguments.joined(separator: ", ")))"
+
+    // FIXME: We should make CodeBlockSyntax ExpressibleByStringInterpolation,
+    // so that the full body could go here.
+    let newBody: ExprSyntax =
+      """
+
+        Task {
+          completionHandler(await \(call))
+        }
+
+      """
+
+    // Drop the @addCompletionHandler attribute from the new declaration.
+    let newAttributeList = AttributeListSyntax(
+      funcDecl.attributes?.filter {
+        guard case let .attribute(attribute) = $0,
+              let attributeType = attribute.attributeName.as(SimpleTypeIdentifierSyntax.self),
+              let nodeType = node.attributeName.as(SimpleTypeIdentifierSyntax.self)
+        else {
+          return true
+        }
+
+        return attributeType.name.text != nodeType.name.text
+      } ?? []
+    )
+
+    let newFunc =
+      funcDecl
+      .with(
+        \.signature,
+        funcDecl.signature
+          .with(
+            \.effectSpecifiers,
+            funcDecl.signature.effectSpecifiers?.with(\.asyncSpecifier, nil)  // drop async
+          )
+          .with(\.output, nil)  // drop result type
+          .with(
+            \.input,  // add completion handler parameter
+            funcDecl.signature.input.with(\.parameterList, newParameterList)
+              .with(\.trailingTrivia, [])
+          )
+      )
+      .with(
+        \.body,
+        CodeBlockSyntax(
+          leftBrace: .leftBraceToken(leadingTrivia: .space),
+          statements: CodeBlockItemListSyntax(
+            [CodeBlockItemSyntax(item: .expr(newBody))]
+          ),
+          rightBrace: .rightBraceToken(leadingTrivia: .newline)
+        )
+      )
+      .with(\.attributes, newAttributeList)
+      .with(\.leadingTrivia, .newlines(2))
+
+    return [DeclSyntax(newFunc)]
+  }
+}
+
 private extension DeclSyntaxProtocol {
   var isObservableStoredProperty: Bool {
     if let property = self.as(VariableDeclSyntax.self),

test/Macros/macro_expand_peers.swift

@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -swift-version 5 -I %swift-host-lib-dir -L %swift-host-lib-dir -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath
+// RUN: %target-typecheck-verify-swift -swift-version 5 -enable-experimental-feature Macros -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -disable-availability-checking
+// RUN: %target-swift-frontend -swift-version 5 -typecheck -enable-experimental-feature Macros -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir %s -disable-availability-checking -dump-macro-expansions > %t/expansions-dump.txt 2>&1
+// RUN: %FileCheck -check-prefix=CHECK-DUMP %s < %t/expansions-dump.txt
+
+
+// FIXME: Swift parser is not enabled on Linux CI yet.
+// REQUIRES: OS=macosx
+
+@attached(peer)
+macro addCompletionHandler() = #externalMacro(module: "MacroDefinition", type: "AddCompletionHandler")
+
+struct S {
+  @addCompletionHandler
+  func f(a: Int, for b: String, _ value: Double) async -> String {
+    return b
+  }
+
+  // CHECK-DUMP: @__swiftmacro_18macro_expand_peers1SV1f1a3for_SSSi_SSSdtYaF20addCompletionHandlerfMp_.swift
+  // CHECK-DUMP: func f(a: Int, for b: String, _ value: Double, completionHandler: (String) -> Void) {
+  // CHECK-DUMP:   Task {
+  // CHECK-DUMP:     completionHandler(await f(a: a, for: b, value))
+  // CHECK-DUMP:   }
+  // CHECK-DUMP: }
+}