From 446ca950b1ba16ebee4f6d207ad10a416b5f4285 Mon Sep 17 00:00:00 2001
From: Rintaro Ishizaki <rishizaki@apple.com>
Date: Sun, 28 Apr 2024 22:21:30 -0700
Subject: [PATCH] [CompilerPlugin] Sink error exit logic to "message listener"

---
 .../SwiftCompilerPlugin/CompilerPlugin.swift  |  8 +------
 .../CompilerPluginMessageHandler.swift        | 22 ++++++++++++++-----
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/Sources/SwiftCompilerPlugin/CompilerPlugin.swift b/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
index 431b8ad2f3b..3ced3a43de1 100644
--- a/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
+++ b/Sources/SwiftCompilerPlugin/CompilerPlugin.swift
@@ -112,12 +112,6 @@ extension CompilerPlugin {
     let connection = try StandardIOMessageConnection()
     let provider = MacroProviderAdapter(plugin: Self())
     let impl = CompilerPluginMessageListener(connection: connection, provider: provider)
-    do {
-      try impl.main()
-    } catch {
-      // Emit a diagnostic and indicate failure to the plugin host,
-      // and exit with an error code.
-      fatalError("Internal Error: \(error)")
-    }
+    impl.main()
   }
 }
diff --git a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
index 8051af7e549..ca18cb418c8 100644
--- a/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
+++ b/Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift
@@ -11,8 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 #if swift(>=6)
+private import _SwiftSyntaxCShims
 public import SwiftSyntaxMacros
 #else
+@_implementationOnly import _SwiftSyntaxCShims
 import SwiftSyntaxMacros
 #endif
 
@@ -83,12 +85,20 @@ public class CompilerPluginMessageListener<Connection: MessageConnection, Provid
 
   /// 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 connection.waitForNextMessage(HostToPluginMessage.self) {
-      let result = handler.handleMessage(message)
-      try connection.sendMessage(result)
+  ///
+  /// On internal errors, such as I/O errors or JSON serialization errors, print
+  /// an error message and `exit(1)`
+  public func main() {
+    do {
+      while let message = try connection.waitForNextMessage(HostToPluginMessage.self) {
+        let result = handler.handleMessage(message)
+        try connection.sendMessage(result)
+      }
+    } catch {
+      // Emit a diagnostic and indicate failure to the plugin host,
+      // and exit with an error code.
+      fputs("Internal Error: \(error)\n", _stderr)
+      exit(1)
     }
   }
 }