Fix test failures when running SourceKit-LSP tests with an Xcode toolchain

We need to jump through a few extra hoops when the SourceKit plugin is located in SourceKit-LSP’s build folder and we are using `sourcekitd_plugin_initialize` instead of `sourcekitd_plugin_initialize_2`.

Author: ahoppen

Sources/SourceKitD/DynamicallyLoadedSourceKitD.swift

@@ -41,6 +41,20 @@ fileprivate extension ThreadSafeBox {
   }
 }
 
+fileprivate func setenv(name: String, value: String, override: Bool) throws {
+  struct FailedToSetEnvError: Error {
+    let errorCode: Int32
+  }
+  try name.withCString { name in
+    try value.withCString { value in
+      let result = setenv(name, value, override ? 0 : 1)
+      if result != 0 {
+        throw FailedToSetEnvError(errorCode: result)
+      }
+    }
+  }
+}
+
 /// Wrapper for sourcekitd, taking care of initialization, shutdown, and notification handler
 /// multiplexing.
 ///
@@ -103,12 +117,21 @@ package actor DynamicallyLoadedSourceKitD: SourceKitD {
     dylibPath: URL,
     pluginPaths: PluginPaths?
   ) async throws -> SourceKitD {
-    try await SourceKitDRegistry.shared
-      .getOrAdd(
-        dylibPath,
-        pluginPaths: pluginPaths,
-        create: { try DynamicallyLoadedSourceKitD(dylib: dylibPath, pluginPaths: pluginPaths) }
-      )
+    try await SourceKitDRegistry.shared.getOrAdd(dylibPath, pluginPaths: pluginPaths) {
+      #if canImport(Darwin)
+      #if compiler(>=6.3)
+      #warning("Remove this when we no longer need to support sourcekitd_plugin_initialize")
+      #endif
+      if let pluginPaths {
+        try setenv(
+          name: "SOURCEKIT_LSP_PLUGIN_SOURCEKITD_PATH_\(pluginPaths.clientPlugin.filePath)",
+          value: dylibPath.filePath,
+          override: false
+        )
+      }
+      #endif
+      return try DynamicallyLoadedSourceKitD(dylib: dylibPath, pluginPaths: pluginPaths)
+    }
   }
 
   package init(dylib path: URL, pluginPaths: PluginPaths?, initialize: Bool = true) throws {

Sources/SourceKitD/dlopen.swift

@@ -34,6 +34,10 @@ package final class DLHandle: Sendable {
     #endif
   }
 
+  #if canImport(Darwin)
+  package static let rtldDefault = DLHandle(rawValue: Handle(handle: UnsafeMutableRawPointer(bitPattern: -2)!))
+  #endif
+
   fileprivate let rawValue: ThreadSafeBox<Handle?>
 
   fileprivate init(rawValue: Handle) {

Sources/SwiftSourceKitClientPlugin/ClientPlugin.swift

@@ -21,6 +21,10 @@ public import Csourcekitd
 import Csourcekitd
 #endif
 
+#if compiler(>=6.3)
+#warning("Remove sourcekitd_plugin_initialize when we no longer support toolchains that call it")
+#endif
+
 /// Legacy plugin initialization logic in which sourcekitd does not inform the plugin about the sourcekitd path it was
 /// loaded from.
 @_cdecl("sourcekitd_plugin_initialize")
@@ -38,6 +42,14 @@ public func sourcekitd_plugin_initialize(_ params: sourcekitd_api_plugin_initial
     .deletingLastPathComponent()
     .appendingPathComponent("sourcekitd.framework")
     .appendingPathComponent("sourcekitd")
+  if !FileManager.default.fileExists(at: url),
+    let sourcekitdPath = ProcessInfo.processInfo.environment["SOURCEKIT_LSP_PLUGIN_SOURCEKITD_PATH_\(path)"]
+  {
+    // When using a SourceKit plugin from the build directory, we can't find sourcekitd relative to the plugin.
+    // Respect the sourcekitd path that was passed to us via an environment variable from
+    // `DynamicallyLoadedSourceKitD.getOrCreate`.
+    url = URL(fileURLWithPath: sourcekitdPath)
+  }
   try! url.filePath.withCString { sourcekitdPath in
     sourcekitd_plugin_initialize_2(params, sourcekitdPath)
   }

Sources/SwiftSourceKitPlugin/Plugin.swift

@@ -143,6 +143,10 @@ final class RequestHandler: Sendable {
   }
 }
 
+#if compiler(>=6.3)
+#warning("Remove sourcekitd_plugin_initialize when we no longer support toolchains that call it")
+#endif
+
 /// Legacy plugin initialization logic in which sourcekitd does not inform the plugin about the sourcekitd path it was
 /// loaded from.
 @_cdecl("sourcekitd_plugin_initialize")
@@ -160,8 +164,16 @@ public func sourcekitd_plugin_initialize(_ params: sourcekitd_api_plugin_initial
     .deletingLastPathComponent()
     .appendingPathComponent("sourcekitd.framework")
     .appendingPathComponent("sourcekitd")
-  try! url.filePath.withCString { sourcekitdPath in
-    sourcekitd_plugin_initialize_2(params, sourcekitdPath)
+  if FileManager.default.fileExists(at: url) {
+    try! url.filePath.withCString { sourcekitdPath in
+      sourcekitd_plugin_initialize_2(params, sourcekitdPath)
+    }
+  } else {
+    // When using a SourceKit plugin from the build directory, we can't find sourcekitd relative to the plugin.
+    // Since sourcekitd_plugin_initialize is only called on Darwin from Xcode toolchains, we know that we are getting
+    // called from an XPC sourcekitd. Thus, all sourcekitd symbols that we need should be loaded in the current process
+    // already and we can use `RTLD_DEFAULT` for the sourcekitd library.
+    sourcekitd_plugin_initialize_2(params, "SOURCEKIT_LSP_PLUGIN_PARENT_LIBRARY_RTLD_DEFAULT")
   }
   #else
   fatalError("sourcekitd_plugin_initialize is not supported on non-Darwin platforms")
@@ -223,15 +235,24 @@ public func sourcekitd_plugin_initialize_2(
   _ params: sourcekitd_api_plugin_initialize_params_t,
   _ parentLibraryPath: UnsafePointer<CChar>
 ) {
+  let parentLibraryPath = String(cString: parentLibraryPath)
   #if canImport(Darwin)
-  // On macOS, we need to find sourcekitdInProc relative to the library the plugin was loaded from.
-  DynamicallyLoadedSourceKitD.forPlugin = try! DynamicallyLoadedSourceKitD.inProcLibrary(
-    relativeTo: URL(fileURLWithPath: String(cString: parentLibraryPath))
-  )
+  if parentLibraryPath == "SOURCEKIT_LSP_PLUGIN_PARENT_LIBRARY_RTLD_DEFAULT" {
+    DynamicallyLoadedSourceKitD.forPlugin = try! DynamicallyLoadedSourceKitD(
+      dlhandle: .rtldDefault,
+      path: URL(string: "rtld-default://")!,
+      pluginPaths: nil,
+      initialize: false
+    )
+  } else {
+    DynamicallyLoadedSourceKitD.forPlugin = try! DynamicallyLoadedSourceKitD.inProcLibrary(
+      relativeTo: URL(fileURLWithPath: String(cString: parentLibraryPath))
+    )
+  }
   #else
   // On other platforms, sourcekitd is always in process, so we can load it straight away.
   DynamicallyLoadedSourceKitD.forPlugin = try! DynamicallyLoadedSourceKitD(
-    dylib: URL(fileURLWithPath: String(cString: parentLibraryPath)),
+    dylib: URL(fileURLWithPath: parentLibraryPath),
     pluginPaths: nil,
     initialize: false
   )