[Macro] Handle error during initializing executable plugins

Previously, the compiler fails to read the result from 'getCapability'
IPC message, it used to just emit a fatalError and crashed.
Instead, propagate the error status and diagnose it.

rdar://108022847
(cherry picked from commit e93e4ba)

Author: rintaro

lib/AST/PluginRegistry.cpp

@@ -178,7 +178,8 @@ LoadedExecutablePlugin::PluginProcess::~PluginProcess() {
 
 LoadedExecutablePlugin::~LoadedExecutablePlugin() {
   // Let ASTGen to cleanup things.
-  this->cleanup();
+  if (this->cleanup)
+    this->cleanup();
 }
 
 ssize_t LoadedExecutablePlugin::PluginProcess::read(void *buf,

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -24,10 +24,18 @@ enum PluginError: Error {
 
 @_cdecl("swift_ASTGen_initializePlugin")
 public func _initializePlugin(
-  opaqueHandle: UnsafeMutableRawPointer
-) {
+  opaqueHandle: UnsafeMutableRawPointer,
+  cxxDiagnosticEngine: UnsafeMutablePointer<UInt8>?
+) -> Bool {
   let plugin = CompilerPlugin(opaqueHandle: opaqueHandle)
-  plugin.initialize()
+  do {
+    try plugin.initialize()
+    return true
+  } catch {
+    // Diagnostics are emitted in the caller.
+    // FIXME: Return what happened or emit diagnostics here.
+    return false
+  }
 }
 
 @_cdecl("swift_ASTGen_deinitializePlugin")
@@ -52,12 +60,7 @@ func swift_ASTGen_pluginServerLoadLibraryPlugin(
   let libraryPath = String(cString: libraryPath)
   let moduleName = String(cString: moduleName)
 
-  let diagEngine: PluginDiagnosticsEngine?
-  if let cxxDiagnosticEngine = cxxDiagnosticEngine {
-    diagEngine = PluginDiagnosticsEngine(cxxDiagnosticEngine: cxxDiagnosticEngine)
-  } else {
-    diagEngine = nil
-  }
+  let diagEngine = PluginDiagnosticsEngine(cxxDiagnosticEngine: cxxDiagnosticEngine)
 
   do {
     let result = try plugin.sendMessageAndWaitWithoutLock(
@@ -136,20 +139,15 @@ struct CompilerPlugin {
   }
 
   /// Initialize the plugin. This should be called inside lock.
-  func initialize() {
-    do {
-      // Get capability.
-      let response = try self.sendMessageAndWaitWithoutLock(.getCapability)
-      guard case .getCapabilityResult(let capability) = response else {
-        throw PluginError.invalidReponseKind
-      }
-      let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
-      ptr.initialize(to: .init(capability))
-      Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
-    } catch {
-      assertionFailure(String(describing: error))
-      return
+  func initialize() throws {
+    // Get capability.
+    let response = try self.sendMessageAndWaitWithoutLock(.getCapability)
+    guard case .getCapabilityResult(let capability) = response else {
+      throw PluginError.invalidReponseKind
     }
+    let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
+    ptr.initialize(to: .init(capability))
+    Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
   }
 
   func deinitialize() {
@@ -179,6 +177,14 @@ class PluginDiagnosticsEngine {
     self.cxxDiagnosticEngine = cxxDiagnosticEngine
   }
 
+  /// Failable convenience initializer for optional cxx engine pointer.
+  convenience init?(cxxDiagnosticEngine: UnsafeMutablePointer<UInt8>?) {
+    guard let cxxDiagnosticEngine = cxxDiagnosticEngine else {
+      return nil
+    }
+    self.init(cxxDiagnosticEngine: cxxDiagnosticEngine)
+  }
+
   /// Register an 'ExportedSourceFile' to the engine. So the engine can get
   /// C++ SourceLoc from a pair of filename and offset.
   func add(exportedSourceFile: UnsafePointer<ExportedSourceFile>) {

lib/Sema/TypeCheckMacros.cpp

@@ -21,6 +21,7 @@
 #include "swift/AST/ASTMangler.h"
 #include "swift/AST/ASTNode.h"
 #include "swift/AST/CASTBridging.h"
+#include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/Expr.h"
 #include "../AST/InlinableText.h"
 #include "swift/AST/MacroDefinition.h"
@@ -74,7 +75,7 @@ extern "C" ptrdiff_t swift_ASTGen_expandAttachedMacro(
     void *parentDeclSourceFile, const void *parentDeclSourceLocation,
     const char **evaluatedSource, ptrdiff_t *evaluatedSourceLength);
 
-extern "C" void swift_ASTGen_initializePlugin(void *handle);
+extern "C" bool swift_ASTGen_initializePlugin(void *handle, void *diagEngine);
 extern "C" void swift_ASTGen_deinitializePlugin(void *handle);
 extern "C" bool swift_ASTGen_pluginServerLoadLibraryPlugin(
     void *handle, const char *libraryPath, const char *moduleName,
@@ -308,7 +309,11 @@ loadExecutablePluginByName(ASTContext &ctx, Identifier moduleName) {
   // But plugin loading is in libAST and it can't link ASTGen symbols.
   if (!executablePlugin->isInitialized()) {
 #if SWIFT_SWIFT_PARSER
-    swift_ASTGen_initializePlugin(executablePlugin);
+    if (!swift_ASTGen_initializePlugin(executablePlugin, &ctx.Diags)) {
+      ctx.Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded,
+                         executablePluginPath, "failed to initialize plugin");
+      return nullptr;
+    }
     executablePlugin->setCleanup([executablePlugin] {
       swift_ASTGen_deinitializePlugin(executablePlugin);
     });
@@ -320,7 +325,9 @@ loadExecutablePluginByName(ASTContext &ctx, Identifier moduleName) {
 #if SWIFT_SWIFT_PARSER
     llvm::SmallString<128> resolvedLibraryPath;
     auto fs = ctx.SourceMgr.getFileSystem();
-    if (fs->getRealPath(libraryPath, resolvedLibraryPath)) {
+    if (auto err = fs->getRealPath(libraryPath, resolvedLibraryPath)) {
+      ctx.Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded,
+                         executablePluginPath, err.message());
       return nullptr;
     }
     std::string resolvedLibraryPathStr(resolvedLibraryPath);

test/Macros/macro_plugin_broken.swift

@@ -0,0 +1,40 @@
+// REQUIRES: swift_swift_parser
+
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %host-build-swift \
+// RUN:   -swift-version 5 -o %t/broken-plugin \
+// RUN:   -module-name=TestPlugin \
+// RUN:   %t/broken_plugin.swift \
+// RUN:   -g -no-toolchain-stdlib-rpath -swift-version 5
+
+// RUN: not %swift-target-frontend \
+// RUN:   -typecheck \
+// RUN:   -swift-version 5 -enable-experimental-feature Macros \
+// RUN:   -load-plugin-executable %t/broken-plugin#TestPlugin \
+// RUN:   -module-name MyApp \
+// RUN:   -serialize-diagnostics-path %t/macro_expand.dia \
+// RUN:   %t/test.swift
+
+// RUN: c-index-test -read-diagnostics %t/macro_expand.dia 2>&1 | %FileCheck -check-prefix CHECK %s
+
+// CHECK: (null):0:0: warning: compiler plugin not loaded: {{.+}}broken-plugin; loader error: failed to initialize plugin
+// CHECK: test.swift:1:33: warning: external macro implementation type 'TestPlugin.FooMacro' could not be found for macro 'fooMacro';
+// CHECK: test.swift:4:7: error: external macro implementation type 'TestPlugin.FooMacro' could not be found for macro 'fooMacro';
+// CHECK: +-{{.+}}test.swift:1:33: note: 'fooMacro' declared here
+
+//--- test.swift
+@freestanding(expression) macro fooMacro(_: Any) -> String = #externalMacro(module: "TestPlugin", type: "FooMacro")
+
+func test() {
+  _ = #fooMacro(1)
+}
+
+//--- broken_plugin.swift
+func minusTen(value: UInt) -> UInt {
+  // Causes crash.
+  return value - 10
+}
+
+_ = minusTen(value: 5)