Support semantic functionality in generated interfaces if the client supports getReferenceDocument

Sources/SKTestSupport/IndexedSingleSwiftFileTestProject.swift

@@ -51,6 +51,7 @@ package struct IndexedSingleSwiftFileTestProject {
   ///   - cleanUp: Whether to remove the temporary directory when the SourceKit-LSP server shuts down.
   package init(
     _ markedText: String,
+    capabilities: ClientCapabilities = ClientCapabilities(),
     indexSystemModules: Bool = false,
     allowBuildFailure: Bool = false,
     workspaceDirectory: URL? = nil,
@@ -153,6 +154,7 @@ package struct IndexedSingleSwiftFileTestProject {
     )
     self.testClient = try await TestSourceKitLSPClient(
       options: options,
+      capabilities: capabilities,
       workspaceFolders: [
         WorkspaceFolder(uri: DocumentURI(testWorkspaceDirectory))
       ],

Sources/SourceKitLSP/CMakeLists.txt

@@ -48,11 +48,12 @@ target_sources(SourceKitLSP PRIVATE
   Swift/DocumentSymbols.swift
   Swift/ExpandMacroCommand.swift
   Swift/FoldingRange.swift
+  Swift/GeneratedInterfaceManager.swift
   Swift/MacroExpansion.swift
   Swift/MacroExpansionReferenceDocumentURLData.swift
   Swift/OpenInterface.swift
-  Swift/RefactoringResponse.swift
   Swift/RefactoringEdit.swift
+  Swift/RefactoringResponse.swift
   Swift/ReferenceDocumentURL.swift
   Swift/RelatedIdentifiers.swift
   Swift/RewriteSourceKitPlaceholders.swift

Sources/SourceKitLSP/MessageHandlingDependencyTracker.swift

@@ -74,9 +74,9 @@ package enum MessageHandlingDependencyTracker: QueueBasedMessageHandlerDependenc
     case (.documentUpdate(let selfUri), .documentUpdate(let otherUri)):
       return selfUri == otherUri
     case (.documentUpdate(let selfUri), .documentRequest(let otherUri)):
-      return selfUri == otherUri
+      return selfUri.buildSettingsFile == otherUri.buildSettingsFile
     case (.documentRequest(let selfUri), .documentUpdate(let otherUri)):
-      return selfUri == otherUri
+      return selfUri.buildSettingsFile == otherUri.buildSettingsFile
 
     // documentRequest
     case (.documentRequest, .documentRequest):

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -242,7 +242,7 @@ package actor SourceKitLSPServer {
   }
 
   package func workspaceForDocument(uri: DocumentURI) async -> Workspace? {
-    let uri = uri.primaryFile ?? uri
+    let uri = uri.buildSettingsFile
     if let cachedWorkspace = self.workspaceForUri[uri]?.value {
       return cachedWorkspace
     }
@@ -1590,14 +1590,14 @@ extension SourceKitLSPServer {
   }
 
   func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
-    let primaryFileURI = try ReferenceDocumentURL(from: req.uri).primaryFile
+    let buildSettingsUri = try ReferenceDocumentURL(from: req.uri).buildSettingsFile
 
-    guard let workspace = await workspaceForDocument(uri: primaryFileURI) else {
-      throw ResponseError.workspaceNotOpen(primaryFileURI)
+    guard let workspace = await workspaceForDocument(uri: buildSettingsUri) else {
+      throw ResponseError.workspaceNotOpen(buildSettingsUri)
     }
 
-    guard let languageService = workspace.documentService(for: primaryFileURI) else {
-      throw ResponseError.unknown("No Language Service for URI: \(primaryFileURI)")
+    guard let languageService = workspace.documentService(for: buildSettingsUri) else {
+      throw ResponseError.unknown("No Language Service for URI: \(buildSettingsUri)")
     }
 
     return try await languageService.getReferenceDocument(req)

Sources/SourceKitLSP/Swift/GeneratedInterfaceDocumentURLData.swift

@@ -0,0 +1,87 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import LanguageServerProtocol
+
+/// Represents url of generated interface reference document.
+
+package struct GeneratedInterfaceDocumentURLData: Hashable, ReferenceURLData {
+  package static let documentType = "generated-swift-interface"
+
+  private struct Parameters {
+    static let moduleName = "moduleName"
+    static let groupName = "groupName"
+    static let sourcekitdDocumentName = "sourcekitdDocument"
+    static let buildSettingsFrom = "buildSettingsFrom"
+  }
+
+  /// The module that should be shown in this generated interface.
+  let moduleName: String
+
+  /// The group that should be shown in this generated interface, if applicable.
+  let groupName: String?
+
+  /// The name by which this document is referred to in sourcekitd.
+  let sourcekitdDocumentName: String
+
+  /// The document from which the build settings for the generated interface should be inferred.
+  let buildSettingsFrom: DocumentURI
+
+  var displayName: String {
+    if let groupName {
+      return "\(moduleName).\(groupName.replacing("/", with: ".")).swiftinterface"
+    }
+    return "\(moduleName).swiftinterface"
+  }
+
+  var queryItems: [URLQueryItem] {
+    var result = [
+      URLQueryItem(name: Parameters.moduleName, value: moduleName)
+    ]
+    if let groupName {
+      result.append(URLQueryItem(name: Parameters.groupName, value: groupName))
+    }
+    result += [
+      URLQueryItem(name: Parameters.sourcekitdDocumentName, value: sourcekitdDocumentName),
+      URLQueryItem(name: Parameters.buildSettingsFrom, value: buildSettingsFrom.stringValue),
+    ]
+    return result
+  }
+
+  var uri: DocumentURI {
+    get throws {
+      try ReferenceDocumentURL.generatedInterface(self).uri
+    }
+  }
+
+  init(moduleName: String, groupName: String?, sourcekitdDocumentName: String, primaryFile: DocumentURI) {
+    self.moduleName = moduleName
+    self.groupName = groupName
+    self.sourcekitdDocumentName = sourcekitdDocumentName
+    self.buildSettingsFrom = primaryFile
+  }
+
+  init(queryItems: [URLQueryItem]) throws {
+    guard let moduleName = queryItems.last(where: { $0.name == Parameters.moduleName })?.value,
+      let sourcekitdDocumentName = queryItems.last(where: { $0.name == Parameters.sourcekitdDocumentName })?.value,
+      let primaryFile = queryItems.last(where: { $0.name == Parameters.buildSettingsFrom })?.value
+    else {
+      throw ReferenceDocumentURLError(description: "Invalid queryItems for generated interface reference document url")
+    }
+
+    self.moduleName = moduleName
+    self.groupName = queryItems.last(where: { $0.name == Parameters.groupName })?.value
+    self.sourcekitdDocumentName = sourcekitdDocumentName
+    self.buildSettingsFrom = try DocumentURI(string: primaryFile)
+  }
+}

Sources/SourceKitLSP/Swift/GeneratedInterfaceManager.swift

@@ -0,0 +1,233 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LanguageServerProtocol
+import SKLogging
+import SKUtilities
+import SourceKitD
+import SwiftExtensions
+
+/// When information about a generated interface is requested, this opens the generated interface in sourcekitd and
+/// caches the generated interface contents.
+///
+/// It keeps the generated interface open in sourcekitd until the corresponding reference document is closed in the
+/// editor. Additionally, it also keeps a few recently requested interfaces cached. This way we don't need to recompute
+/// the generated interface contents between the initial generated interface request to find a USR's position in the
+/// interface until the editor actually opens the reference document.
+actor GeneratedInterfaceManager {
+  private struct OpenGeneratedInterfaceDocumentDetails {
+    let url: GeneratedInterfaceDocumentURLData
+
+    /// The contents of the generated interface.
+    let snapshot: DocumentSnapshot
+
+    /// The number of `GeneratedInterfaceManager` that are actively working with the sourcekitd document. If this value
+    /// is 0, the generated interface may be closed in sourcekitd.
+    ///
+    /// Usually, this value is 1, while the reference document for this generated interface is open in the editor.
+    var refCount: Int
+  }
+
+  private weak var swiftLanguageService: SwiftLanguageService?
+
+  /// The number of generated interface documents that are not in editor but should still be cached.
+  private let cacheSize = 2
+
+  /// Details about the generated interfaces that are currently open in sourcekitd.
+  ///
+  /// Conceptually, this is a dictionary with `url` being the key. To prevent excessive memory usage we only keep
+  ///  `cacheSize` entries with a ref count of 0 in the array. Older entries are at the end of the list, newer entries
+  /// at the front.
+  private var openInterfaces: [OpenGeneratedInterfaceDocumentDetails] = []
+
+  init(swiftLanguageService: SwiftLanguageService) {
+    self.swiftLanguageService = swiftLanguageService
+  }
+
+  /// If there are more than `cacheSize` entries in `openInterfaces` that have a ref count of 0, close the oldest ones.
+  private func purgeCache() {
+    var documentsToClose: [String] = []
+    while openInterfaces.count(where: { $0.refCount == 0 }) > cacheSize,
+      let indexToPurge = openInterfaces.lastIndex(where: { $0.refCount == 0 })
+    {
+      documentsToClose.append(openInterfaces[indexToPurge].url.sourcekitdDocumentName)
+      openInterfaces.remove(at: indexToPurge)
+    }
+    if !documentsToClose.isEmpty, let swiftLanguageService {
+      Task {
+        let sourcekitd = swiftLanguageService.sourcekitd
+        for documentToClose in documentsToClose {
+          await orLog("Closing generated interface") {
+            _ = try await swiftLanguageService.sendSourcekitdRequest(
+              sourcekitd.dictionary([
+                sourcekitd.keys.request: sourcekitd.requests.editorClose,
+                sourcekitd.keys.name: documentToClose,
+                sourcekitd.keys.cancelBuilds: 0,
+              ]),
+              fileContents: nil
+            )
+          }
+        }
+      }
+    }
+  }
+
+  /// If we don't have the generated interface for the given `document` open in sourcekitd, open it, otherwise return
+  /// its details from the cache.
+  ///
+  /// If `incrementingRefCount` is `true`, then the document manager will keep the generated interface open in
+  /// sourcekitd, independent of the cache size. If `incrementingRefCount` is `true`, then `decrementRefCount` must be
+  /// called to allow the document to be closed again.
+  private func details(
+    for document: GeneratedInterfaceDocumentURLData,
+    incrementingRefCount: Bool
+  ) async throws -> OpenGeneratedInterfaceDocumentDetails {
+    func loadFromCache() -> OpenGeneratedInterfaceDocumentDetails? {
+      guard let cachedIndex = openInterfaces.firstIndex(where: { $0.url == document }) else {
+        return nil
+      }
+      if incrementingRefCount {
+        openInterfaces[cachedIndex].refCount += 1
+      }
+      return openInterfaces[cachedIndex]
+
+    }
+    if let cached = loadFromCache() {
+      return cached
+    }
+
+    guard let swiftLanguageService else {
+      // `SwiftLanguageService` has been destructed. We are tearing down the language server. Nothing left to do.
+      throw ResponseError.unknown("Connection to the editor closed")
+    }
+
+    let sourcekitd = swiftLanguageService.sourcekitd
+
+    let keys = sourcekitd.keys
+    let skreq = sourcekitd.dictionary([
+      keys.request: sourcekitd.requests.editorOpenInterface,
+      keys.moduleName: document.moduleName,
+      keys.groupName: document.groupName,
+      keys.name: document.sourcekitdDocumentName,
+      keys.synthesizedExtension: 1,
+      keys.compilerArgs: await swiftLanguageService.buildSettings(for: try document.uri, fallbackAfterTimeout: false)?
+        .compilerArgs as [SKDRequestValue]?,
+    ])
+
+    let dict = try await swiftLanguageService.sendSourcekitdRequest(skreq, fileContents: nil)
+
+    guard let contents: String = dict[keys.sourceText] else {
+      throw ResponseError.unknown("sourcekitd response is missing sourceText")
+    }
+
+    if let cached = loadFromCache() {
+      // Another request raced us to create the generated interface. Discard what we computed here and return the cached
+      // value.
+      await orLog("Closing generated interface created during race") {
+        _ = try await swiftLanguageService.sendSourcekitdRequest(
+          sourcekitd.dictionary([
+            keys.request: sourcekitd.requests.editorClose,
+            keys.name: document.sourcekitdDocumentName,
+            keys.cancelBuilds: 0,
+          ]),
+          fileContents: nil
+        )
+      }
+      return cached
+    }
+
+    let details = OpenGeneratedInterfaceDocumentDetails(
+      url: document,
+      snapshot: DocumentSnapshot(
+        uri: try document.uri,
+        language: .swift,
+        version: 0,
+        lineTable: LineTable(contents)
+      ),
+      refCount: incrementingRefCount ? 1 : 0
+    )
+    openInterfaces.insert(details, at: 0)
+    purgeCache()
+    return details
+  }
+
+  private func decrementRefCount(for document: GeneratedInterfaceDocumentURLData) {
+    guard let cachedIndex = openInterfaces.firstIndex(where: { $0.url == document }) else {
+      logger.fault(
+        "Generated interface document for \(document.moduleName) is not open anymore. Unbalanced retain and releases?"
+      )
+      return
+    }
+    if openInterfaces[cachedIndex].refCount == 0 {
+      logger.fault(
+        "Generated interface document for \(document.moduleName) is already 0. Unbalanced retain and releases?"
+      )
+      return
+    }
+    openInterfaces[cachedIndex].refCount -= 1
+    purgeCache()
+  }
+
+  func position(ofUsr usr: String, in document: GeneratedInterfaceDocumentURLData) async throws -> Position {
+    guard let swiftLanguageService else {
+      // `SwiftLanguageService` has been destructed. We are tearing down the language server. Nothing left to do.
+      throw ResponseError.unknown("Connection to the editor closed")
+    }
+
+    let details = try await details(for: document, incrementingRefCount: true)
+    defer {
+      decrementRefCount(for: document)
+    }
+
+    let sourcekitd = swiftLanguageService.sourcekitd
+    let keys = sourcekitd.keys
+    let skreq = sourcekitd.dictionary([
+      keys.request: sourcekitd.requests.editorFindUSR,
+      keys.sourceFile: document.sourcekitdDocumentName,
+      keys.usr: usr,
+    ])
+
+    let dict = try await swiftLanguageService.sendSourcekitdRequest(skreq, fileContents: details.snapshot.text)
+    guard let offset: Int = dict[keys.offset] else {
+      throw ResponseError.unknown("Missing key 'offset'")
+    }
+    return details.snapshot.positionOf(utf8Offset: offset)
+  }
+
+  func snapshot(of document: GeneratedInterfaceDocumentURLData) async throws -> DocumentSnapshot {
+    return try await details(for: document, incrementingRefCount: false).snapshot
+  }
+
+  func open(document: GeneratedInterfaceDocumentURLData) async throws {
+    _ = try await details(for: document, incrementingRefCount: true)
+  }
+
+  func close(document: GeneratedInterfaceDocumentURLData) async {
+    decrementRefCount(for: document)
+  }
+
+  func reopen(interfacesWithBuildSettingsFrom buildSettingsFile: DocumentURI) async {
+    for openInterface in openInterfaces {
+      guard openInterface.url.buildSettingsFrom == buildSettingsFile else {
+        continue
+      }
+      await orLog("Reopening generated interface") {
+        // `MessageHandlingDependencyTracker` ensures that we don't handle a request for the generated interface while
+        // it is being re-opened because `documentUpdate` and `documentRequest` use the `buildSettingsFile` to determine
+        // their dependencies.
+        await close(document: openInterface.url)
+        openInterfaces.removeAll(where: { $0.url == openInterface.url })
+        try await open(document: openInterface.url)
+      }
+    }
+  }
+}