Extend copied file mapping to all LSP requests returning locations

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -962,6 +962,118 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     return locations.map { locationAdjustedForCopiedFiles($0) }
   }
 
+  private func uriAdjustedForCopiedFiles(_ uri: DocumentURI) -> DocumentURI {
+    guard let originalUri = cachedCopiedFileMap[uri] else {
+      return uri
+    }
+    return originalUri
+  }
+
+  package func workspaceEditAdjustedForCopiedFiles(_ workspaceEdit: WorkspaceEdit?) -> WorkspaceEdit? {
+    guard var edit = workspaceEdit else {
+      return nil
+    }
+    if let changes = edit.changes {
+      var newChanges: [DocumentURI: [TextEdit]] = [:]
+      for (uri, edits) in changes {
+        let newUri = self.uriAdjustedForCopiedFiles(uri)
+        newChanges[newUri, default: []] += edits
+      }
+      edit.changes = newChanges
+    }
+    if let documentChanges = edit.documentChanges {
+      edit.documentChanges = documentChanges.map { change in
+        switch change {
+        case .textDocumentEdit(var textEdit):
+          textEdit.textDocument.uri = self.uriAdjustedForCopiedFiles(textEdit.textDocument.uri)
+          return .textDocumentEdit(textEdit)
+        case .createFile(var create):
+          create.uri = self.uriAdjustedForCopiedFiles(create.uri)
+          return .createFile(create)
+        case .renameFile(var rename):
+          rename.oldUri = self.uriAdjustedForCopiedFiles(rename.oldUri)
+          rename.newUri = self.uriAdjustedForCopiedFiles(rename.newUri)
+          return .renameFile(rename)
+        case .deleteFile(var delete):
+          delete.uri = self.uriAdjustedForCopiedFiles(delete.uri)
+          return .deleteFile(delete)
+        }
+      }
+    }
+    return edit
+  }
+
+  package func locationsOrLocationLinksAdjustedForCopiedFiles(
+    _ response: LocationsOrLocationLinksResponse?
+  ) -> LocationsOrLocationLinksResponse? {
+    guard let response = response else {
+      return nil
+    }
+    switch response {
+    case .locations(let locations):
+      let remappedLocations = self.locationsAdjustedForCopiedFiles(locations)
+      return .locations(remappedLocations)
+    case .locationLinks(let locationLinks):
+      let remappedLinks = locationLinks.map { link -> LocationLink in
+        let adjustedTargetLocation = self.locationAdjustedForCopiedFiles(
+          Location(uri: link.targetUri, range: link.targetRange)
+        )
+        let adjustedTargetSelectionLocation = self.locationAdjustedForCopiedFiles(
+          Location(uri: link.targetUri, range: link.targetSelectionRange)
+        )
+        return LocationLink(
+          originSelectionRange: link.originSelectionRange,
+          targetUri: adjustedTargetLocation.uri,
+          targetRange: adjustedTargetLocation.range,
+          targetSelectionRange: adjustedTargetSelectionLocation.range
+        )
+      }
+      return .locationLinks(remappedLinks)
+    }
+  }
+
+  package func typeHierarchyItemAdjustedForCopiedFiles(_ item: TypeHierarchyItem) -> TypeHierarchyItem {
+    let adjustedLocation = self.locationAdjustedForCopiedFiles(Location(uri: item.uri, range: item.range))
+    let adjustedSelectionLocation = self.locationAdjustedForCopiedFiles(
+      Location(uri: item.uri, range: item.selectionRange)
+    )
+    return TypeHierarchyItem(
+      name: item.name,
+      kind: item.kind,
+      tags: item.tags,
+      detail: item.detail,
+      uri: adjustedLocation.uri,
+      range: adjustedLocation.range,
+      selectionRange: adjustedSelectionLocation.range,
+      data: item.data
+    )
+  }
+
+  package func callHierarchyItemAdjustedForCopiedFiles(_ item: CallHierarchyItem) -> CallHierarchyItem {
+    let adjustedLocation = self.locationAdjustedForCopiedFiles(Location(uri: item.uri, range: item.range))
+    let adjustedSelectionLocation = self.locationAdjustedForCopiedFiles(
+      Location(uri: item.uri, range: item.selectionRange)
+    )
+    return CallHierarchyItem(
+      name: item.name,
+      kind: item.kind,
+      tags: item.tags,
+      detail: item.detail,
+      uri: adjustedLocation.uri,
+      range: adjustedLocation.range,
+      selectionRange: adjustedSelectionLocation.range,
+      data: .dictionary([
+        "usr": item.data.flatMap { data in
+          if case let .dictionary(dict) = data {
+            return dict["usr"]
+          }
+          return nil
+        } ?? .null,
+        "uri": .string(adjustedLocation.uri.stringValue),
+      ])
+    )
+  }
+
   @discardableResult
   package func scheduleRecomputeCopyFileMap() -> Task<Void, Never> {
     let task = Task { [previousUpdateTask = copiedFileMapUpdateTask] in

Sources/ClangLanguageService/ClangLanguageService.swift

@@ -481,7 +481,11 @@ extension ClangLanguageService {
   }
 
   package func declaration(_ req: DeclarationRequest) async throws -> LocationsOrLocationLinksResponse? {
-    return try await forwardRequestToClangd(req)
+    let result = try await forwardRequestToClangd(req)
+    guard let workspace = self.workspace.value else {
+      return result
+    }
+    return await workspace.buildServerManager.locationsOrLocationLinksAdjustedForCopiedFiles(result)
   }
 
   package func completion(_ req: CompletionRequest) async throws -> CompletionList {
@@ -618,7 +622,11 @@ extension ClangLanguageService {
   }
 
   package func indexedRename(_ request: IndexedRenameRequest) async throws -> WorkspaceEdit? {
-    return try await forwardRequestToClangd(request)
+    let workspaceEdit = try await forwardRequestToClangd(request)
+    guard let workspace = self.workspace.value else {
+      return workspaceEdit
+    }
+    return await workspace.buildServerManager.workspaceEditAdjustedForCopiedFiles(workspaceEdit)
   }
 
   // MARK: - Other
@@ -634,7 +642,12 @@ extension ClangLanguageService {
       position: renameRequest.position
     )
     let symbolDetail = try await forwardRequestToClangd(symbolInfoRequest).only
-    return (try await edits ?? WorkspaceEdit(), symbolDetail?.usr)
+    let workspaceEdit = try await edits ?? WorkspaceEdit()
+    guard let workspace = self.workspace.value else {
+      return (workspaceEdit, symbolDetail?.usr)
+    }
+    let remappedEdit = await workspace.buildServerManager.workspaceEditAdjustedForCopiedFiles(workspaceEdit)
+    return (remappedEdit ?? WorkspaceEdit(), symbolDetail?.usr)
   }
 
   package func syntacticDocumentTests(