[CompilerPlugin] Factor out message handling logic to a module

Package.swift

@@ -44,6 +44,7 @@ let package = Package(
     .library(name: "SwiftSyntaxBuilder", type: .static, targets: ["SwiftSyntaxBuilder"]),
     .library(name: "SwiftSyntaxMacros", type: .static, targets: ["SwiftSyntaxMacros"]),
     .library(name: "SwiftCompilerPlugin", type: .static, targets: ["SwiftCompilerPlugin"]),
+    .library(name: "SwiftCompilerPluginMessageHandling", type: .static, targets: ["SwiftCompilerPluginMessageHandling"]),
     .library(name: "SwiftRefactor", type: .static, targets: ["SwiftRefactor"]),
   ],
   targets: [
@@ -124,6 +125,12 @@ let package = Package(
     ),
     .target(
       name: "SwiftCompilerPlugin",
+      dependencies: [
+        "SwiftCompilerPluginMessageHandling", "SwiftSyntaxMacros",
+      ]
+    ),
+    .target(
+      name: "SwiftCompilerPluginMessageHandling",
       dependencies: [
         "SwiftSyntax", "SwiftParser", "SwiftDiagnostics", "SwiftSyntaxMacros", "SwiftOperators",
       ]

Sources/SwiftCompilerPlugin/CompilerPlugin.swift

@@ -13,6 +13,7 @@
 // https://github.com/apple/swift-package-manager/blob/main/Sources/PackagePlugin/Plugin.swift
 
 import SwiftSyntaxMacros
+@_implementationOnly import SwiftCompilerPluginMessageHandling
 
 @_implementationOnly import Foundation
 #if os(Windows)
@@ -61,6 +62,37 @@ public protocol CompilerPlugin {
   var providingMacros: [Macro.Type] { get }
 }
 
+extension CompilerPlugin {
+  func resolveMacro(moduleName: String, typeName: String) -> Macro.Type? {
+    let qualifedName = "\(moduleName).\(typeName)"
+
+    for type in providingMacros {
+      // FIXME: Is `String(reflecting:)` stable?
+      // Getting the module name and type name should be more robust.
+      let name = String(reflecting: type)
+      if name == qualifedName {
+        return type
+      }
+    }
+    return nil
+  }
+
+  // @testable
+  public func _resolveMacro(moduleName: String, typeName: String) -> Macro.Type? {
+    resolveMacro(moduleName: moduleName, typeName: typeName)
+  }
+}
+
+struct MacroProviderAdapter<Plugin: CompilerPlugin>: PluginProvider {
+  let plugin: Plugin
+  init(plugin: Plugin) {
+    self.plugin = plugin
+  }
+  func resolveMacro(moduleName: String, typeName: String) -> Macro.Type? {
+    plugin.resolveMacro(moduleName: moduleName, typeName: typeName)
+  }
+}
+
 extension CompilerPlugin {
 
   /// Main entry point of the plugin — sets up a communication channel with
@@ -105,18 +137,17 @@ extension CompilerPlugin {
     #endif
 
     // Open a message channel for communicating with the plugin host.
-    pluginHostConnection = PluginHostConnection(
+    let connection = PluginHostConnection(
       inputStream: FileHandle(fileDescriptor: inputFD),
       outputStream: FileHandle(fileDescriptor: outputFD)
     )
 
     // Handle messages from the host until the input stream is closed,
     // indicating that we're done.
-    let instance = Self()
+    let provider = MacroProviderAdapter(plugin: Self())
+    let impl = CompilerPluginMessageHandler(connection: connection, provider: provider)
     do {
-      while let message = try pluginHostConnection.waitForNextMessage() {
-        try instance.handleMessage(message)
-      }
+      try impl.main()
     } catch {
       // Emit a diagnostic and indicate failure to the plugin host,
       // and exit with an error code.
@@ -135,46 +166,13 @@ extension CompilerPlugin {
     if let cStr = strerror(errno) { return String(cString: cStr) }
     return String(describing: errno)
   }
-
-  /// Handles a single message received from the plugin host.
-  fileprivate func handleMessage(_ message: HostToPluginMessage) throws {
-    switch message {
-    case .getCapability:
-      try pluginHostConnection.sendMessage(
-        .getCapabilityResult(capability: PluginMessage.capability)
-      )
-      break
-
-    case .expandFreestandingMacro(let macro, let discriminator, let expandingSyntax):
-      try expandFreestandingMacro(
-        macro: macro,
-        discriminator: discriminator,
-        expandingSyntax: expandingSyntax
-      )
-
-    case .expandAttachedMacro(let macro, let macroRole, let discriminator, let attributeSyntax, let declSyntax, let parentDeclSyntax):
-      try expandAttachedMacro(
-        macro: macro,
-        macroRole: macroRole,
-        discriminator: discriminator,
-        attributeSyntax: attributeSyntax,
-        declSyntax: declSyntax,
-        parentDeclSyntax: parentDeclSyntax
-      )
-    }
-  }
 }
 
-/// Message channel for bidirectional communication with the plugin host.
-internal fileprivate(set) var pluginHostConnection: PluginHostConnection!
-
-typealias PluginHostConnection = MessageConnection<PluginToHostMessage, HostToPluginMessage>
-
-internal struct MessageConnection<TX, RX> where TX: Encodable, RX: Decodable {
+internal struct PluginHostConnection: MessageConnection {
   let inputStream: FileHandle
   let outputStream: FileHandle
 
-  func sendMessage(_ message: TX) throws {
+  func sendMessage<TX: Encodable>(_ message: TX) throws {
     // Encode the message as JSON.
     let payload = try JSONEncoder().encode(message)
 
@@ -188,7 +186,7 @@ internal struct MessageConnection<TX, RX> where TX: Encodable, RX: Decodable {
     try outputStream._write(contentsOf: payload)
   }
 
-  func waitForNextMessage() throws -> RX? {
+  func waitForNextMessage<RX: Decodable>(_ ty: RX.Type) throws -> RX? {
     // Read the header (a 64-bit length field in little endian byte order).
     guard
       let header = try inputStream._read(upToCount: 8),

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -0,0 +1,92 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntaxMacros
+
+/// A type that provides the actual plugin functions.
+public protocol PluginProvider {
+  func resolveMacro(moduleName: String, typeName: String) -> Macro.Type?
+}
+
+/// Low level message connection to the plugin host.
+/// This encapsulates the connection and the message serialization.
+public protocol MessageConnection {
+  /// Send a message to the peer.
+  func sendMessage<TX: Encodable>(_ message: TX) throws
+  /// Wait until receiving a message from the peer, and return it.
+  func waitForNextMessage<RX: Decodable>(_ type: RX.Type) throws -> RX?
+}
+
+/// 'CompilerPluginMessageHandler' is a type that listens to the message
+/// connection and dispatches them to the actual plugin provider, then send back
+/// the response.
+///
+/// The low level connection and the provider is injected by the client.
+public class CompilerPluginMessageHandler<Connection: MessageConnection, Provider: PluginProvider> {
+  /// Message channel for bidirectional communication with the plugin host.
+  let connection: Connection
+
+  /// Object to provide actual plugin functions.
+  let provider: Provider
+
+  public init(connection: Connection, provider: Provider) {
+    self.connection = connection
+    self.provider = provider
+  }
+}
+
+extension CompilerPluginMessageHandler {
+  func sendMessage(_ message: PluginToHostMessage) throws {
+    try connection.sendMessage(message)
+  }
+
+  func waitForNextMessage() throws -> HostToPluginMessage? {
+    try connection.waitForNextMessage(HostToPluginMessage.self)
+  }
+
+  /// Run the main message listener loop.
+  /// Returns when the message connection was closed.
+  /// Throws an error when it failed to send/receive the message, or failed
+  /// to serialize/deserialize the message.
+  public func main() throws {
+    while let message = try self.waitForNextMessage() {
+      try handleMessage(message)
+    }
+  }
+
+  /// Handles a single message received from the plugin host.
+  fileprivate func handleMessage(_ message: HostToPluginMessage) throws {
+    switch message {
+    case .getCapability:
+      try self.sendMessage(
+        .getCapabilityResult(capability: PluginMessage.capability)
+      )
+
+    case .expandFreestandingMacro(let macro, let discriminator, let expandingSyntax):
+      try expandFreestandingMacro(
+        macro: macro,
+        discriminator: discriminator,
+        expandingSyntax: expandingSyntax
+      )
+
+    case .expandAttachedMacro(let macro, let macroRole, let discriminator, let attributeSyntax, let declSyntax, let parentDeclSyntax):
+      try expandAttachedMacro(
+        macro: macro,
+        macroRole: macroRole,
+        discriminator: discriminator,
+        attributeSyntax: attributeSyntax,
+        declSyntax: declSyntax,
+        parentDeclSyntax: parentDeclSyntax
+      )
+    }
+  }
+}

Sources/SwiftCompilerPlugin/Macros.swift → Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -14,25 +14,10 @@ import SwiftDiagnostics
 import SwiftSyntax
 import SwiftSyntaxMacros
 
-/// Implementation for `CompilerPlugin` macro related request processing.
-extension CompilerPlugin {
-  private func resolveMacro(moduleName: String, typeName: String) -> Macro.Type? {
-    let qualifedName = "\(moduleName).\(typeName)"
-
-    for type in self.providingMacros {
-      // FIXME: Is `String(reflecting:)` stable?
-      // Getting the module name and type name should be more robust.
-      let name = String(reflecting: type)
-      if name == qualifedName {
-        return type
-      }
-    }
-    return nil
-  }
-
+extension CompilerPluginMessageHandler {
   /// Get concrete macro type from a pair of module name and type name.
   private func resolveMacro(_ ref: PluginMessage.MacroReference) -> Macro.Type? {
-    resolveMacro(moduleName: ref.moduleName, typeName: ref.typeName)
+    provider.resolveMacro(moduleName: ref.moduleName, typeName: ref.typeName)
   }
 
   /// Expand `@freestainding(XXX)` macros.
@@ -90,7 +75,7 @@ extension CompilerPlugin {
     let diagnostics = context.diagnostics.map {
       PluginMessage.Diagnostic(from: $0, in: sourceManager)
     }
-    try pluginHostConnection.sendMessage(
+    try self.sendMessage(
       .expandFreestandingMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
     )
   }
@@ -249,19 +234,12 @@ extension CompilerPlugin {
     let diagnostics = context.diagnostics.map {
       PluginMessage.Diagnostic(from: $0, in: sourceManager)
     }
-    try pluginHostConnection.sendMessage(
+    try self.sendMessage(
       .expandAttachedMacroResult(expandedSources: expandedSources, diagnostics: diagnostics)
     )
   }
 }
 
-extension CompilerPlugin {
-  // @testable
-  public func _resolveMacro(moduleName: String, typeName: String) -> Macro.Type? {
-    resolveMacro(moduleName: moduleName, typeName: typeName)
-  }
-}
-
 /// Diagnostic message used for thrown errors.
 fileprivate struct ThrownErrorDiagnostic: DiagnosticMessage {
   let message: String