[Macros] Attached macro expasnions return single string

* Move collapse(expansions:for:attachedTo:) to SwiftSyntaxMacroExpansion
* SwiftSyntaxMacroExpansion.expandAttachedMacro() now perform collapsing
* SwiftSyntaxMacroExpansion.expandAttachedMacroWithoutCollapsing() for
  to keep old behavior
* IPC message 'getCapability' now send the host protocol version
* Unified IPC response 'macroExpansionResult' that returns single string
  for both 'expandFreestandingMacro' and 'expandAttachedMacro'
* Compiler accepts old 'expandFreestandingMacroResult' and
  'expandAttachedMacroResult' to keep compatibility
* Plugins check the compiler's protcol version to see if it suppports
  'macroExpansionResult', and fall back to old behavior if necessary

Author: rintaro

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -40,6 +40,23 @@ public protocol MessageConnection {
   func waitForNextMessage<RX: Decodable>(_ type: RX.Type) throws -> RX?
 }
 
+/// Represent the capability of the plugin host (i.e. compiler).
+struct HostCapability {
+  var protocolVersion: Int
+
+  // Create an "oldest" capability.
+  init() {
+    self.protocolVersion = 0
+  }
+
+  init(_ message: PluginMessage.HostCapability) {
+    self.protocolVersion = message.protocolVersion
+  }
+
+  /// Compiler accept 'expandMacroResult' response message.
+  var hasExpandMacroResult: Bool { protocolVersion >= 5 }
+}
+
 /// 'CompilerPluginMessageHandler' is a type that listens to the message
 /// connection and dispatches them to the actual plugin provider, then send back
 /// the response.
@@ -52,9 +69,13 @@ public class CompilerPluginMessageHandler<Connection: MessageConnection, Provide
   /// Object to provide actual plugin functions.
   let provider: Provider
 
+  /// Plugin host capability
+  var hostCapability: HostCapability
+
   public init(connection: Connection, provider: Provider) {
     self.connection = connection
     self.provider = provider
+    self.hostCapability = HostCapability()
   }
 }
 
@@ -80,7 +101,13 @@ extension CompilerPluginMessageHandler {
   /// Handles a single message received from the plugin host.
   fileprivate func handleMessage(_ message: HostToPluginMessage) throws {
     switch message {
-    case .getCapability:
+    case .getCapability(let hostCapability):
+      // Remember the peer capability if provided.
+      if let hostCapability = hostCapability {
+        self.hostCapability = .init(hostCapability)
+      }
+
+      // Return the plugin capability.
       let capability = PluginMessage.PluginCapability(
         protocolVersion: PluginMessage.PROTOCOL_VERSION_NUMBER,
         features: provider.features.map({ $0.rawValue })

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -68,9 +68,15 @@ extension CompilerPluginMessageHandler {
     let diagnostics = context.diagnostics.map {
       PluginMessage.Diagnostic(from: $0, in: sourceManager)
     }
-    try self.sendMessage(
-      .expandFreestandingMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
-    )
+
+    let response: PluginToHostMessage
+    if (hostCapability.hasExpandMacroResult) {
+      response = .expandMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
+    } else {
+      // TODO: Remove this  when all compilers have 'hasExpandMacroResult'.
+      response = .expandFreestandingMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
+    }
+    try self.sendMessage(response)
   }
 
   /// Expand `@attached(XXX)` macros.
@@ -95,20 +101,38 @@ extension CompilerPluginMessageHandler {
     let declarationNode = sourceManager.add(declSyntax).cast(DeclSyntax.self)
     let parentDeclNode = parentDeclSyntax.map { sourceManager.add($0).cast(DeclSyntax.self) }
 
+    // TODO: Make this a 'String?' and remove non-'hasExpandMacroResult' branches.
     let expandedSources: [String]?
-    do {
+    MAIN: do {
       guard let macroDefinition = resolveMacro(macro) else {
         throw MacroExpansionError.macroTypeNotFound(macro)
       }
+      let role = MacroRole(messageMacroRole: macroRole)
 
-      expandedSources = SwiftSyntaxMacroExpansion.expandAttachedMacro(
+      let expansions = SwiftSyntaxMacroExpansion.expandAttachedMacroWithoutCollapsing(
         definition: macroDefinition,
-        macroRole: MacroRole(messageMacroRole: macroRole),
+        macroRole: role,
         attributeNode: attributeNode,
         declarationNode: declarationNode,
         parentDeclNode: parentDeclNode,
         in: context
       )
+      guard let expansions = expansions else {
+        expandedSources = nil
+        break MAIN
+      }
+      if hostCapability.hasExpandMacroResult {
+        // Make a single element array by collapsing the results into a string.
+        expandedSources = [
+          SwiftSyntaxMacroExpansion.collapse(
+            expansions: expansions,
+            for: role,
+            attachedTo: declarationNode
+          )
+        ]
+      } else {
+        expandedSources = expansions
+      }
     } catch {
       context.addDiagnostics(from: error, node: attributeNode)
       expandedSources = nil
@@ -117,9 +141,14 @@ extension CompilerPluginMessageHandler {
     let diagnostics = context.diagnostics.map {
       PluginMessage.Diagnostic(from: $0, in: sourceManager)
     }
-    try self.sendMessage(
-      .expandAttachedMacroResult(expandedSources: expandedSources, diagnostics: diagnostics)
-    )
+
+    let response: PluginToHostMessage
+    if (hostCapability.hasExpandMacroResult) {
+      response = .expandMacroResult(expandedSource: expandedSources?.first, diagnostics: diagnostics)
+    } else {
+      response = .expandAttachedMacroResult(expandedSources: expandedSources, diagnostics: diagnostics)
+    }
+    try self.sendMessage(response)
   }
 }
 

Sources/SwiftCompilerPluginMessageHandling/PluginMessages.swift

@@ -14,8 +14,10 @@
 // NOTE: Types in this file should be self-contained and should not depend on any non-stdlib types.
 
 internal enum HostToPluginMessage: Codable {
-  /// Get capability of this plugin.
-  case getCapability
+  /// Send capability of the host, and get capability of the plugin.
+  case getCapability(
+    capability: PluginMessage.HostCapability?
+  )
 
   /// Expand a '@freestanding' macro.
   case expandFreestandingMacro(
@@ -49,11 +51,19 @@ internal enum PluginToHostMessage: Codable {
     capability: PluginMessage.PluginCapability
   )
 
+  /// Unified response for freestanding/attached macro expansion.
+  case expandMacroResult(
+    expandedSource: String?,
+    diagnostics: [PluginMessage.Diagnostic]
+  )
+
+  // @available(*, deprecated: "use expandMacroResult() instead")
   case expandFreestandingMacroResult(
     expandedSource: String?,
     diagnostics: [PluginMessage.Diagnostic]
   )
 
+  // @available(*, deprecated: "use expandMacroResult() instead")
   case expandAttachedMacroResult(
     expandedSources: [String]?,
     diagnostics: [PluginMessage.Diagnostic]
@@ -66,7 +76,11 @@ internal enum PluginToHostMessage: Codable {
 }
 
 /*namespace*/ internal enum PluginMessage {
-  static var PROTOCOL_VERSION_NUMBER: Int { 4 }  // Added 'loadPluginLibrary'.
+  static var PROTOCOL_VERSION_NUMBER: Int { 5 }  // Added 'expandMacroResult'.
+
+  struct HostCapability: Codable {
+    var protocolVersion: Int
+  }
 
   struct PluginCapability: Codable {
     var protocolVersion: Int

Sources/SwiftSyntaxMacroExpansion/MacroExpansion.swift

@@ -1,3 +1,15 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
 import SwiftSyntax
 @_spi(MacroExpansion) import SwiftSyntaxMacros
 
@@ -160,7 +172,7 @@ public func expandFreestandingMacro(
 /// - Returns: A list of expanded source text. Upon failure (i.e.
 ///   `defintion.expansion()` throws) returns `nil`, and the diagnostics
 ///   representing the `Error` are guaranteed to be added to context.
-public func expandAttachedMacro<Context: MacroExpansionContext>(
+public func expandAttachedMacroWithoutCollapsing<Context: MacroExpansionContext>(
   definition: Macro.Type,
   macroRole: MacroRole,
   attributeNode: AttributeSyntax,
@@ -292,6 +304,40 @@ public func expandAttachedMacro<Context: MacroExpansionContext>(
   }
 }
 
+/// Expand `@attached(XXX)` macros.
+///
+/// - Parameters:
+///   - definition: a type that conforms to one or more attached `Macro` protocols.
+///   - macroRole: indicates which `Macro` protocol expansion should be performed
+///   - attributeNode: attribute syntax node (e.g. `@macroName(argument)`).
+///   - declarationNode: target declaration syntax node to apply the expansion.
+///   - parentDeclNode: Only used for `MacroRole.memberAttribute`. The parent
+///     context node of `declarationNode`.
+///   - in: context of the expansion.
+/// - Returns: expanded source text. Upon failure (i.e. `defintion.expansion()`
+///   throws) returns `nil`, and the diagnostics representing the `Error` are
+///   guaranteed to be added to context.
+public func expandAttachedMacro<Context: MacroExpansionContext>(
+  definition: Macro.Type,
+  macroRole: MacroRole,
+  attributeNode: AttributeSyntax,
+  declarationNode: DeclSyntax,
+  parentDeclNode: DeclSyntax?,
+  in context: Context
+) -> String? {
+  let expandedSources = expandAttachedMacroWithoutCollapsing(
+    definition: definition,
+    macroRole: macroRole,
+    attributeNode: attributeNode,
+    declarationNode: declarationNode,
+    parentDeclNode: parentDeclNode,
+    in: context
+  )
+  return expandedSources.map {
+    collapse(expansions: $0, for: macroRole, attachedTo: declarationNode)
+  }
+}
+
 fileprivate extension SyntaxProtocol {
   /// Perform a format if required and then trim any leading/trailing
   /// whitespace.
@@ -306,3 +352,56 @@ fileprivate extension SyntaxProtocol {
     return formatted.trimmedDescription(matching: { $0.isWhitespace })
   }
 }
+
+/// Join `expansions`
+public func collapse<Node: SyntaxProtocol>(
+  expansions: [String],
+  for role: MacroRole,
+  attachedTo declarationNode: Node
+) -> String {
+  if expansions.isEmpty {
+    return ""
+  }
+
+  var expansions = expansions
+  var separator: String = "\n\n"
+
+  if role == .accessor,
+    let varDecl = declarationNode.as(VariableDeclSyntax.self),
+    let binding = varDecl.bindings.first,
+    binding.accessor == nil
+  {
+    let indentation = String(repeating: " ", count: 4)
+
+    expansions = expansions.map({ indent($0, with: indentation) })
+    expansions[0] = "{\n" + expansions[0]
+    expansions[expansions.count - 1] += "\n}"
+  } else if role == .memberAttribute {
+    separator = " "
+  }
+
+  return expansions.joined(separator: separator)
+}
+
+fileprivate func indent(_ source: String, with indentation: String) -> String {
+  if source.isEmpty || indentation.isEmpty {
+    return source
+  }
+
+  var indented = ""
+  var remaining = source[...]
+  while let nextNewline = remaining.firstIndex(where: { $0.isNewline }) {
+    if nextNewline != remaining.startIndex {
+      indented += indentation
+    }
+    indented += remaining[...nextNewline]
+    remaining = remaining[remaining.index(after: nextNewline)...]
+  }
+
+  if !remaining.isEmpty {
+    indented += indentation
+    indented += remaining
+  }
+
+  return indented
+}