Apply code review fixes for copied file handling

- Remove async from workspaceEditAdjustedForCopiedFiles
- Refactor to use uriAdjustedForCopiedFiles helper
- Update dictionary update logic with +=
- Adjust LocationLink creation to use adjusted ranges
- Ensure selectionRange adjustment in prepareCallHierarchy
- Provide default WorkspaceEdit in ClangLanguageService
- Revert asyncMap to map and remove await in SourceKitLSPServer
- Chain workspace and index retrieval in incomingCalls
- Use indexToLSPCallHierarchyItem and shared helper for CallHierarchyItem
- Fix indentation and remove duplicated detail setting
- Use shared helper for TypeHierarchyItem
- Remove .sort() from expected array in tests
- Enable testFindImplementationInCopiedHeader
- Add await for actor-isolated BuildServerManager calls

Author: loveucifer

Sources/BuildServerIntegration/BuildServerManager.swift

@@ -962,17 +962,22 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     return locations.map { locationAdjustedForCopiedFiles($0) }
   }
 
-  package func workspaceEditAdjustedForCopiedFiles(_ workspaceEdit: WorkspaceEdit?) async -> WorkspaceEdit? {
+  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.locationAdjustedForCopiedFiles(
-          Location(uri: uri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-        ).uri
-        newChanges[newUri, default: []].append(contentsOf: edits)
+        let newUri = self.uriAdjustedForCopiedFiles(uri)
+        newChanges[newUri, default: []] += edits
       }
       edit.changes = newChanges
     }
@@ -981,32 +986,17 @@ package actor BuildServerManager: QueueBasedMessageHandler {
       for change in documentChanges {
         switch change {
         case .textDocumentEdit(var textEdit):
-          let newUri = self.locationAdjustedForCopiedFiles(
-            Location(uri: textEdit.textDocument.uri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-          ).uri
-          textEdit.textDocument.uri = newUri
+          textEdit.textDocument.uri = self.uriAdjustedForCopiedFiles(textEdit.textDocument.uri)
           newDocumentChanges.append(.textDocumentEdit(textEdit))
         case .createFile(var create):
-          let newUri = self.locationAdjustedForCopiedFiles(
-            Location(uri: create.uri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-          ).uri
-          create.uri = newUri
+          create.uri = self.uriAdjustedForCopiedFiles(create.uri)
           newDocumentChanges.append(.createFile(create))
         case .renameFile(var rename):
-          let newOldUri = self.locationAdjustedForCopiedFiles(
-            Location(uri: rename.oldUri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-          ).uri
-          let newNewUri = self.locationAdjustedForCopiedFiles(
-            Location(uri: rename.newUri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-          ).uri
-          rename.oldUri = newOldUri
-          rename.newUri = newNewUri
+          rename.oldUri = self.uriAdjustedForCopiedFiles(rename.oldUri)
+          rename.newUri = self.uriAdjustedForCopiedFiles(rename.newUri)
           newDocumentChanges.append(.renameFile(rename))
         case .deleteFile(var delete):
-          let newUri = self.locationAdjustedForCopiedFiles(
-            Location(uri: delete.uri, range: Position(line: 0, utf16index: 0)..<Position(line: 0, utf16index: 0))
-          ).uri
-          delete.uri = newUri
+          delete.uri = self.uriAdjustedForCopiedFiles(delete.uri)
           newDocumentChanges.append(.deleteFile(delete))
         }
       }
@@ -1015,39 +1005,41 @@ package actor BuildServerManager: QueueBasedMessageHandler {
     return edit
   }
 
-  package func locationsOrLocationLinksAdjustedForCopiedFiles(_ response: LocationsOrLocationLinksResponse?) async -> LocationsOrLocationLinksResponse? {
+  package func locationsOrLocationLinksAdjustedForCopiedFiles(_ response: LocationsOrLocationLinksResponse?) -> LocationsOrLocationLinksResponse? {
     guard let response = response else {
       return nil
     }
     switch response {
     case .locations(let locations):
-      let remappedLocations = await self.locationsAdjustedForCopiedFiles(locations)
+      let remappedLocations = self.locationsAdjustedForCopiedFiles(locations)
       return .locations(remappedLocations)
     case .locationLinks(let locationLinks):
       var remappedLinks: [LocationLink] = []
       for link in locationLinks {
-        let newUri = self.locationAdjustedForCopiedFiles(Location(uri: link.targetUri, range: link.targetRange)).uri
+        let adjustedTargetLocation = self.locationAdjustedForCopiedFiles(Location(uri: link.targetUri, range: link.targetRange))
+        let adjustedTargetSelectionLocation = self.locationAdjustedForCopiedFiles(Location(uri: link.targetUri, range: link.targetSelectionRange))
         remappedLinks.append(LocationLink(
           originSelectionRange: link.originSelectionRange,
-          targetUri: newUri,
-          targetRange: link.targetRange,
-          targetSelectionRange: link.targetSelectionRange
+          targetUri: adjustedTargetLocation.uri,
+          targetRange: adjustedTargetLocation.range,
+          targetSelectionRange: adjustedTargetSelectionLocation.range
         ))
       }
       return .locationLinks(remappedLinks)
     }
   }
 
-  package func typeHierarchyItemAdjustedForCopiedFiles(_ item: TypeHierarchyItem) async -> TypeHierarchyItem {
-    let adjustedLocation = await self.locationAdjustedForCopiedFiles(Location(uri: item.uri, range: item.range))
+  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: item.selectionRange,
+      selectionRange: adjustedSelectionLocation.range,
       data: item.data
     )
   }

Sources/ClangLanguageService/ClangLanguageService.swift

@@ -642,9 +642,9 @@ extension ClangLanguageService {
       position: renameRequest.position
     )
     let symbolDetail = try await forwardRequestToClangd(symbolInfoRequest).only
-    let workspaceEdit = try await edits
+    let workspaceEdit = try await edits ?? WorkspaceEdit()
     guard let workspace = self.workspace.value else {
-      return (workspaceEdit ?? WorkspaceEdit(), symbolDetail?.usr)
+      return (workspaceEdit, symbolDetail?.usr)
     }
     let remappedEdit = await workspace.buildServerManager.workspaceEditAdjustedForCopiedFiles(workspaceEdit)
     return (remappedEdit ?? WorkspaceEdit(), symbolDetail?.usr)

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -2260,6 +2260,36 @@ extension SourceKitLSPServer {
     )
   }
 
+  private func callHierarchyItemAdjustedForCopiedFiles(
+    _ item: CallHierarchyItem,
+    workspace: Workspace
+  ) async -> CallHierarchyItem {
+    let adjustedLocation = await workspace.buildServerManager.locationAdjustedForCopiedFiles(
+      Location(uri: item.uri, range: item.range)
+    )
+    let adjustedSelectionLocation = await workspace.buildServerManager.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),
+      ])
+    )
+  }
+
   func prepareCallHierarchy(
     _ req: CallHierarchyPrepareRequest,
     workspace: Workspace,
@@ -2284,31 +2314,11 @@ extension SourceKitLSPServer {
       guard let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
         continue
       }
-      let location = indexToLSPLocation(definition.location)
-      guard let originalLocation = location else {
+      guard let item = indexToLSPCallHierarchyItem(definition: definition, index: index) else {
         continue
       }
-      let remappedLocation = await workspace.buildServerManager.locationAdjustedForCopiedFiles(originalLocation)
-
-      // Create a new CallHierarchyItem with the remapped location (similar to how we handle TypeHierarchyItem)
-      let name = index.fullyQualifiedName(of: definition)
-      let symbol = definition.symbol
-
-      let item = CallHierarchyItem(
-        name: name,
-        kind: symbol.kind.asLspSymbolKind(),
-        tags: nil,
-        detail: nil,
-        uri: remappedLocation.uri,
-        range: remappedLocation.range,
-        selectionRange: remappedLocation.range,
-        // We encode usr and uri for incoming/outgoing call lookups in the implementation-specific data field
-        data: .dictionary([
-          "usr": .string(symbol.usr),
-          "uri": .string(remappedLocation.uri.stringValue),
-        ])
-      )
-      callHierarchyItems.append(item)
+      let adjustedItem = await callHierarchyItemAdjustedForCopiedFiles(item, workspace: workspace)
+      callHierarchyItems.append(adjustedItem)
     }
     callHierarchyItems.sort(by: { Location(uri: $0.uri, range: $0.range) < Location(uri: $1.uri, range: $1.range) })
 
@@ -2358,7 +2368,7 @@ extension SourceKitLSPServer {
     var callersToCalls: [Symbol: [SymbolOccurrence]] = [:]
 
     for call in callOccurrences {
-      // Callers are all `calledBy` relations of a call to a USR in `callableUSrs`, ie. all the functions that contain a
+      // Callers are all `containedBy` relations of a call to a USR in `callableUSrs`, ie. all the functions that contain a
       // call to a USR in callableUSRs. In practice, this should always be a single item.
       let callers = call.relations.filter { $0.roles.contains(.containedBy) }.map(\.symbol)
       for caller in callers {
@@ -2506,32 +2516,22 @@ extension SourceKitLSPServer {
     }
 
     let symbol = definition.symbol
-      switch symbol.kind {
-      case .extension:
-        // Query the conformance added by this extension
-        let conformances = index.occurrences(relatedToUSR: symbol.usr, roles: .baseOf)
-        if conformances.isEmpty {
-          name = symbol.name
-        } else {
-          name = "\(symbol.name): \(conformances.map(\.symbol.name).sorted().joined(separator: ", "))"
-        }
-        // Add the file name and line to the detail string
-        if let url = remappedLocation.uri.fileURL,
-          let basename = (try? AbsolutePath(validating: url.filePath))?.basename
-        {
-          detail = "Extension at \(basename):\(remappedLocation.range.lowerBound.line + 1)"
-        } else if !definition.location.moduleName.isEmpty {
-          detail = "Extension in \(definition.location.moduleName)"
-        } else {
-          detail = "Extension"
-        }
+    switch symbol.kind {
+    case .extension:
+      // Query the conformance added by this extension
+      let conformances = index.occurrences(relatedToUSR: symbol.usr, roles: .baseOf)
+      if conformances.isEmpty {
+        name = symbol.name
+      } else {
+        name = "\(symbol.name): \(conformances.map(\.symbol.name).sorted().joined(separator: ", "))"
+      }
       // Add the file name and line to the detail string
       if let url = location.uri.fileURL,
         let basename = (try? AbsolutePath(validating: url.filePath))?.basename
       {
         detail = "Extension at \(basename):\(location.range.lowerBound.line + 1)"
-      } else if let moduleName = moduleName {
-        detail = "Extension in \(moduleName)"
+      } else if !definition.location.moduleName.isEmpty {
+        detail = "Extension in \(definition.location.moduleName)"
       } else {
         detail = "Extension"
       }
@@ -2602,55 +2602,16 @@ extension SourceKitLSPServer {
       }
 
       let location = indexToLSPLocation(info.location)
-      guard let originalLocation = location else {
+      guard location != nil else {
         continue
       }
-      let remappedLocation = await workspace.buildServerManager.locationAdjustedForCopiedFiles(originalLocation)
-
-      // Create a new TypeHierarchyItem with the original location, then adjust for copied files
-      let name: String
-      let detail: String?
-      let symbol = info.symbol
-      switch symbol.kind {
-      case .extension:
-        // Query the conformance added by this extension
-        let conformances = index.occurrences(relatedToUSR: symbol.usr, roles: .baseOf)
-        if conformances.isEmpty {
-          name = symbol.name
-        } else {
-          name = "\(symbol.name): \(conformances.map(\.symbol.name).sorted().joined(separator: ", "))"
-        }
-        // Add the file name and line to the detail string
-        if let url = remappedLocation.uri.fileURL,
-          let basename = (try? AbsolutePath(validating: url.filePath))?.basename
-        {
-          detail = "Extension at \(basename):\(remappedLocation.range.lowerBound.line + 1)"
-        } else if !info.location.moduleName.isEmpty {
-          detail = "Extension in \(info.location.moduleName)"
-        } else {
-          detail = "Extension"
-        }
-      default:
-        name = index.fullyQualifiedName(of: info)
-        detail = info.location.moduleName
-      }
-
-       let item = TypeHierarchyItem(
-         name: name,
-         kind: symbol.kind.asLspSymbolKind(),
-         tags: nil,
-         detail: detail,
-         uri: originalLocation.uri,
-         range: originalLocation.range,
-         selectionRange: originalLocation.range,
-         // We encode usr and uri for incoming/outgoing type lookups in the implementation-specific data field
-         data: .dictionary([
-           "usr": .string(symbol.usr),
-           "uri": .string(originalLocation.uri.stringValue),
-         ])
-       )
-       let remappedItem = await workspace.buildServerManager.typeHierarchyItemAdjustedForCopiedFiles(item)
-       typeHierarchyItems.append(remappedItem)
+      
+      let moduleName = info.location.moduleName.isEmpty ? nil : info.location.moduleName
+      guard let item = indexToLSPTypeHierarchyItem(definition: info, moduleName: moduleName, index: index) else {
+        continue
+      }
+      let remappedItem = await workspace.buildServerManager.typeHierarchyItemAdjustedForCopiedFiles(item)
+      typeHierarchyItems.append(remappedItem)
     }
     typeHierarchyItems.sort(by: { $0.name < $1.name })
 
@@ -2754,10 +2715,10 @@ extension SourceKitLSPServer {
           name = "\(symbol.name): \(conformances.map(\.symbol.name).sorted().joined(separator: ", "))"
         }
         // Add the file name and line to the detail string
-        if let url = adjustedLocation.uri.fileURL,
+        if let url = remappedLocation.uri.fileURL,
           let basename = (try? AbsolutePath(validating: url.filePath))?.basename
         {
-          detail = "Extension at \(basename):\(adjustedLocation.range.lowerBound.line + 1)"
+          detail = "Extension at \(basename):\(remappedLocation.range.lowerBound.line + 1)"
         } else if !definition.location.moduleName.isEmpty {
           detail = "Extension in \(definition.location.moduleName)"
         } else {

Tests/SourceKitLSPTests/CopiedHeaderTests.swift

@@ -120,16 +120,14 @@ class CopiedHeaderTests: SourceKitLSPTestCase {
         context: ReferencesContext(includeDeclaration: true)
       )
     )
-    response.sort()
     var expected = [
       try project.location(from: "1️⃣", to: "1️⃣", in: "Test.h"),
       try project.location(from: "2️⃣", to: "2️⃣", in: "Test.c"),
     ]
-    expected.sort()
     XCTAssertEqual(response, expected)
   }
 
-  func DISABLED_testFindImplementationInCopiedHeader() async throws {
+  func testFindImplementationInCopiedHeader() async throws {
     let project = try await CustomBuildServerTestProject(
       files: [
         "Test.h": """