[Macros] Introduce CompilerPluginMessageListener

Sources/SwiftCompilerPlugin/CompilerPlugin.swift

@@ -177,7 +177,7 @@ extension CompilerPlugin {
     // Handle messages from the host until the input stream is closed,
     // indicating that we're done.
     let provider = MacroProviderAdapter(plugin: Self())
-    let impl = CompilerPluginMessageHandler(connection: connection, provider: provider)
+    let impl = CompilerPluginMessageListener(connection: connection, provider: provider)
     do {
       try impl.main()
     } catch {

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -64,50 +64,52 @@ struct HostCapability {
   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
+/// 'CompilerPluginMessageListener' is a type that listens to the message
+/// connection, delegate them to the message handler, then send back
 /// the response.
 ///
 /// The low level connection and the provider is injected by the client.
 @_spi(PluginMessage)
-public class CompilerPluginMessageHandler<Connection: MessageConnection, Provider: PluginProvider> {
+public class CompilerPluginMessageListener<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
-
-  /// Plugin host capability
-  var hostCapability: HostCapability
+  let handler: CompilerPluginMessageHandler<Provider>
 
   public init(connection: Connection, provider: Provider) {
     self.connection = connection
-    self.provider = provider
-    self.hostCapability = HostCapability()
-  }
-}
-
-extension CompilerPluginMessageHandler {
-  func sendMessage(_ message: PluginToHostMessage) throws {
-    try connection.sendMessage(message)
-  }
-
-  func waitForNextMessage() throws -> HostToPluginMessage? {
-    try connection.waitForNextMessage(HostToPluginMessage.self)
+    self.handler = CompilerPluginMessageHandler(provider: provider)
   }
 
   /// 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)
+    while let message = try connection.waitForNextMessage(HostToPluginMessage.self) {
+      let result = handler.handleMessage(message)
+      try connection.sendMessage(result)
     }
   }
+}
+
+/// 'CompilerPluginMessageHandler' is a type that handle a message and do the
+/// corresponding operation.
+@_spi(PluginMessage)
+public class CompilerPluginMessageHandler<Provider: PluginProvider> {
+  /// Object to provide actual plugin functions.
+  let provider: Provider
+
+  /// Plugin host capability
+  var hostCapability: HostCapability
+
+  public init(provider: Provider) {
+    self.provider = provider
+    self.hostCapability = HostCapability()
+  }
 
   /// Handles a single message received from the plugin host.
-  fileprivate func handleMessage(_ message: HostToPluginMessage) throws {
+  public func handleMessage(_ message: HostToPluginMessage) -> PluginToHostMessage {
     switch message {
     case .getCapability(let hostCapability):
       // Remember the peer capability if provided.
@@ -120,7 +122,7 @@ extension CompilerPluginMessageHandler {
         protocolVersion: PluginMessage.PROTOCOL_VERSION_NUMBER,
         features: provider.features.map({ $0.rawValue })
       )
-      try self.sendMessage(.getCapabilityResult(capability: capability))
+      return .getCapabilityResult(capability: capability)
 
     case .expandFreestandingMacro(
       let macro,
@@ -129,7 +131,7 @@ extension CompilerPluginMessageHandler {
       let expandingSyntax,
       let lexicalContext
     ):
-      try expandFreestandingMacro(
+      return expandFreestandingMacro(
         macro: macro,
         macroRole: macroRole,
         discriminator: discriminator,
@@ -148,7 +150,7 @@ extension CompilerPluginMessageHandler {
       let conformanceListSyntax,
       let lexicalContext
     ):
-      try expandAttachedMacro(
+      return expandAttachedMacro(
         macro: macro,
         macroRole: macroRole,
         discriminator: discriminator,
@@ -176,7 +178,7 @@ extension CompilerPluginMessageHandler {
           )
         )
       }
-      try self.sendMessage(.loadPluginLibraryResult(loaded: diags.isEmpty, diagnostics: diags))
+      return .loadPluginLibraryResult(loaded: diags.isEmpty, diagnostics: diags)
     }
   }
 }

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -55,7 +55,7 @@ extension CompilerPluginMessageHandler {
     discriminator: String,
     expandingSyntax: PluginMessage.Syntax,
     lexicalContext: [PluginMessage.Syntax]?
-  ) throws {
+  ) -> PluginToHostMessage {
     let sourceManager = SourceManager()
     let syntax = sourceManager.add(expandingSyntax, foldingWith: .standardOperators)
 
@@ -105,7 +105,7 @@ extension CompilerPluginMessageHandler {
       // TODO: Remove this  when all compilers have 'hasExpandMacroResult'.
       response = .expandFreestandingMacroResult(expandedSource: expandedSource, diagnostics: diagnostics)
     }
-    try self.sendMessage(response)
+    return response
   }
 
   /// Expand `@attached(XXX)` macros.
@@ -119,7 +119,7 @@ extension CompilerPluginMessageHandler {
     extendedTypeSyntax: PluginMessage.Syntax?,
     conformanceListSyntax: PluginMessage.Syntax?,
     lexicalContext: [PluginMessage.Syntax]?
-  ) throws {
+  ) -> PluginToHostMessage {
     let sourceManager = SourceManager()
     let attributeNode = sourceManager.add(
       attributeSyntax,
@@ -189,7 +189,7 @@ extension CompilerPluginMessageHandler {
     } else {
       response = .expandAttachedMacroResult(expandedSources: expandedSources, diagnostics: diagnostics)
     }
-    try self.sendMessage(response)
+    return response
   }
 }