Add support for semantic functionality in macro expansion reference documents

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -147,7 +147,7 @@ extension SwiftLanguageService {
     additionalParameters appendAdditionalParameters: ((SKDRequestDictionary) -> Void)? = nil
   ) async throws -> (cursorInfo: [CursorInfo], refactorActions: [SemanticRefactorCommand]) {
     let documentManager = try self.documentManager
-    let snapshot = try documentManager.latestSnapshot(uri)
+    let snapshot = try await self.latestSnapshot(for: uri)
 
     let offsetRange = snapshot.utf8OffsetRange(of: range)
 
@@ -158,7 +158,8 @@ extension SwiftLanguageService {
       keys.cancelOnSubsequentRequest: 0,
       keys.offset: offsetRange.lowerBound,
       keys.length: offsetRange.upperBound != offsetRange.lowerBound ? offsetRange.count : nil,
-      keys.sourceFile: snapshot.uri.pseudoPath,
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
       keys.compilerArgs: await self.buildSettings(for: uri)?.compilerArgs as [SKDRequestValue]?,
     ])
 

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -104,7 +104,8 @@ actor DiagnosticReportManager {
 
     let skreq = sourcekitd.dictionary([
       keys.request: requests.diagnostics,
-      keys.sourceFile: snapshot.uri.pseudoPath,
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
       keys.compilerArgs: compilerArgs as [SKDRequestValue],
     ])
 

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -68,7 +68,7 @@ actor MacroExpansionManager {
   ) async throws -> String {
     let expansions = try await macroExpansions(
       in: macroExpansionURLData.parent,
-      at: macroExpansionURLData.selectionRange
+      at: macroExpansionURLData.parentSelectionRange
     )
     guard let expansion = expansions.filter({ $0.bufferName == macroExpansionURLData.bufferName }).only else {
       throw ResponseError.unknown("Failed to find macro expansion for \(macroExpansionURLData.bufferName).")
@@ -195,7 +195,7 @@ extension SwiftLanguageService {
             MacroExpansionReferenceDocumentURLData(
               macroExpansionEditRange: macroEdit.range,
               parent: expandMacroCommand.textDocument.uri,
-              selectionRange: expandMacroCommand.positionRange,
+              parentSelectionRange: expandMacroCommand.positionRange,
               bufferName: bufferName
             )
           )
@@ -224,10 +224,21 @@ extension SwiftLanguageService {
       let expansionURIs = try macroExpansionReferenceDocumentURLs.map {
         return DocumentURI(try $0.url)
       }
+
+      let uri = expandMacroCommand.textDocument.uri.primaryFile ?? expandMacroCommand.textDocument.uri
+
+      let position =
+        switch try? ReferenceDocumentURL(from: expandMacroCommand.textDocument.uri) {
+        case .macroExpansion(let data):
+          data.primaryFileSelectionRange.lowerBound
+        case nil:
+          expandMacroCommand.positionRange.lowerBound
+        }
+
       Task {
         let req = PeekDocumentsRequest(
-          uri: expandMacroCommand.textDocument.uri,
-          position: expandMacroCommand.positionRange.lowerBound,
+          uri: uri,
+          position: position,
           locations: expansionURIs
         )
 

Sources/SourceKitLSP/Swift/MacroExpansionReferenceDocumentURLData.swift

@@ -15,15 +15,16 @@ import LanguageServerProtocol
 import RegexBuilder
 
 /// Represents url of macro expansion reference document as follows:
-/// `sourcekit-lsp://swift-macro-expansion/LaCb-LcCd.swift?parent=&fromLine=&fromColumn=&toLine=&toColumn=&bufferName=`
+/// `sourcekit-lsp://swift-macro-expansion/LaCb-LcCd.swift?fromLine=&fromColumn=&toLine=&toColumn=&bufferName=&parent=`
 ///
 /// Here,
 ///  - `LaCb-LcCd.swift`, the `displayName`, represents where the macro will expand to or
 ///    replace in the source file (i.e. `macroExpansionEditRange`)
+///  - `fromLine`, `fromColumn`, `toLine`, `toColumn` represents the cursor's selection range in
+///    its `parent` (i.e. `parentSelectionRange`)
+///  - `bufferName` denotes the buffer name of the specific macro expansion edit
 ///  - `parent` denoting the URI of the document from which the macro was expanded. For a first-level macro expansion,
 ///    this is a file URI. For nested macro expansions, this is a `sourcekit-lsp://swift-macro-expansion` URL.
-///  - `fromLine`, `fromColumn`, `toLine`, `toColumn` represents the cursor's `selectionRange`
-///  - `bufferName` denotes the buffer name of the specific macro expansion edit
 package struct MacroExpansionReferenceDocumentURLData {
   package static let documentType = "swift-macro-expansion"
 
@@ -33,7 +34,7 @@ package struct MacroExpansionReferenceDocumentURLData {
   package var parent: DocumentURI
 
   /// The range that was selected in `parent` when the macro was expanded.
-  package var selectionRange: Range<Position>
+  package var parentSelectionRange: Range<Position>
 
   /// ## Example
   ///
@@ -47,7 +48,7 @@ package struct MacroExpansionReferenceDocumentURLData {
   ///
   /// Generated content of reference document url:
   /// URL:
-  /// `sourcekit-lsp://swift-macro-expansion/L3C7-L3C23.swift?primaryFilePath=/path/to/swift_file.swift&fromLine=3&fromColumn=8&toLine=3&toColumn=8&bufferName=@__swift_macro_..._Stringify_.swift`
+  /// `sourcekit-lsp://swift-macro-expansion/L3C7-L3C23.swift?fromLine=3&fromColumn=8&toLine=3&toColumn=8&bufferName=@__swift_macro_..._Stringify_.swift&parent=/path/to/swift_file.swift`
   /// ```swift
   /// (a + b, "a + b")
   /// ```
@@ -64,12 +65,12 @@ package struct MacroExpansionReferenceDocumentURLData {
   package init(
     macroExpansionEditRange: Range<Position>,
     parent: DocumentURI,
-    selectionRange: Range<Position>,
+    parentSelectionRange: Range<Position>,
     bufferName: String
   ) {
     self.macroExpansionEditRange = macroExpansionEditRange
     self.parent = parent
-    self.selectionRange = selectionRange
+    self.parentSelectionRange = parentSelectionRange
     self.bufferName = bufferName
   }
 
@@ -78,13 +79,17 @@ package struct MacroExpansionReferenceDocumentURLData {
   }
 
   package var queryItems: [URLQueryItem] {
-    return [
-      URLQueryItem(name: Parameters.parent, value: parent.stringValue),
-      URLQueryItem(name: Parameters.fromLine, value: String(selectionRange.lowerBound.line)),
-      URLQueryItem(name: Parameters.fromColumn, value: String(selectionRange.lowerBound.utf16index)),
-      URLQueryItem(name: Parameters.toLine, value: String(selectionRange.upperBound.line)),
-      URLQueryItem(name: Parameters.toColumn, value: String(selectionRange.upperBound.utf16index)),
+    [
+      URLQueryItem(name: Parameters.fromLine, value: String(parentSelectionRange.lowerBound.line)),
+      URLQueryItem(name: Parameters.fromColumn, value: String(parentSelectionRange.lowerBound.utf16index)),
+      URLQueryItem(name: Parameters.toLine, value: String(parentSelectionRange.upperBound.line)),
+      URLQueryItem(name: Parameters.toColumn, value: String(parentSelectionRange.upperBound.utf16index)),
       URLQueryItem(name: Parameters.bufferName, value: bufferName),
+
+      // *Note*: Having `parent` as the last parameter will ensure that the url's parameters aren't mistaken to be its
+      // `parent`'s parameters in certain environments where percent encoding gets removed or added
+      // unnecessarily (for example: VS Code)
+      URLQueryItem(name: Parameters.parent, value: parent.stringValue),
     ]
   }
 
@@ -100,7 +105,7 @@ package struct MacroExpansionReferenceDocumentURLData {
     }
 
     self.parent = try DocumentURI(string: parent)
-    self.selectionRange =
+    self.parentSelectionRange =
       Position(line: fromLine, utf16index: fromColumn)..<Position(line: toLine, utf16index: toColumn)
     self.bufferName = bufferName
     self.macroExpansionEditRange = try Self.parse(displayName: displayName)
@@ -121,7 +126,7 @@ package struct MacroExpansionReferenceDocumentURLData {
   ///
   /// Generated content of reference document url:
   /// URL:
-  /// `sourcekit-lsp://swift-macro-expansion/L3C7-L3C23.swift?primaryFilePath=/path/to/swift_file.swift&fromLine=3&fromColumn=8&toLine=3&toColumn=8&bufferName=@__swift_macro_..._Stringify_.swift`
+  /// `sourcekit-lsp://swift-macro-expansion/L3C7-L3C23.swift?fromLine=3&fromColumn=8&toLine=3&toColumn=8&bufferName=@__swift_macro_..._Stringify_.swift&parent=/path/to/swift_file.swift`
   /// ```swift
   /// (a + b, "a + b")
   /// ```
@@ -134,9 +139,18 @@ package struct MacroExpansionReferenceDocumentURLData {
   package var primaryFile: DocumentURI {
     switch try? ReferenceDocumentURL(from: parent) {
     case .macroExpansion(let data):
-      return data.primaryFile
+      data.primaryFile
+    case nil:
+      parent
+    }
+  }
+
+  package var primaryFileSelectionRange: Range<Position> {
+    switch try? ReferenceDocumentURL(from: parent) {
+    case .macroExpansion(let data):
+      data.primaryFileSelectionRange
     case nil:
-      return parent
+      self.parentSelectionRange
     }
   }
 

Sources/SourceKitLSP/Swift/OpenInterface.swift

@@ -37,7 +37,7 @@ extension SwiftLanguageService {
     let interfaceFilePath = self.generatedInterfacesPath.appendingPathComponent("\(name).swiftinterface")
     let interfaceDocURI = DocumentURI(interfaceFilePath)
     // has interface already been generated
-    if let snapshot = try? self.documentManager.latestSnapshot(interfaceDocURI) {
+    if let snapshot = try? await self.latestSnapshot(for: interfaceDocURI) {
       return await self.generatedInterfaceDetails(
         uri: interfaceDocURI,
         snapshot: snapshot,
@@ -111,7 +111,8 @@ extension SwiftLanguageService {
       let keys = self.keys
       let skreq = sourcekitd.dictionary([
         keys.request: requests.editorFindUSR,
-        keys.sourceFile: uri.pseudoPath,
+        keys.sourceFile: uri.sourcekitdSourceFile,
+        keys.primaryFile: uri.primaryFile?.pseudoPath,
         keys.usr: symbol,
       ])
 

Sources/SourceKitLSP/Swift/RelatedIdentifiers.swift

@@ -69,7 +69,8 @@ extension SwiftLanguageService {
       keys.request: requests.relatedIdents,
       keys.cancelOnSubsequentRequest: 0,
       keys.offset: snapshot.utf8Offset(of: position),
-      keys.sourceFile: snapshot.uri.pseudoPath,
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
       keys.includeNonEditableBaseNames: includeNonEditableBaseNames ? 1 : 0,
       keys.compilerArgs: await self.buildSettings(for: snapshot.uri)?.compilerArgs as [SKDRequestValue]?,
     ])

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -26,7 +26,8 @@ extension SwiftLanguageService {
 
     let skreq = sourcekitd.dictionary([
       keys.request: requests.semanticTokens,
-      keys.sourceFile: snapshot.uri.pseudoPath,
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
       keys.compilerArgs: buildSettings.compilerArgs as [SKDRequestValue],
     ])
 
@@ -84,7 +85,7 @@ extension SwiftLanguageService {
   package func documentSemanticTokens(
     _ req: DocumentSemanticTokensRequest
   ) async throws -> DocumentSemanticTokensResponse? {
-    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
+    let snapshot = try await self.latestSnapshot(for: req.textDocument.uri)
 
     let tokens = try await mergedAndSortedTokens(for: snapshot)
     let encodedTokens = tokens.lspEncoded

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -529,13 +529,13 @@ extension SwiftLanguageService {
         return
       }
       do {
-        let snapshot = try await documentManager.latestSnapshot(document)
+        let snapshot = try await self.latestSnapshot(for: document)
         let buildSettings = await self.buildSettings(for: document)
         let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
           for: snapshot,
           buildSettings: buildSettings
         )
-        let latestSnapshotID = try? await documentManager.latestSnapshot(snapshot.uri).id
+        let latestSnapshotID = try? await self.latestSnapshot(for: snapshot.uri).id
         if latestSnapshotID != snapshot.id {
           // Check that the document wasn't modified while we were getting diagnostics. This could happen because we are
           // calling `publishDiagnosticsIfNeeded` outside of `messageHandlingQueue` and thus a concurrent edit is
@@ -766,7 +766,7 @@ extension SwiftLanguageService {
   }
 
   package func documentSymbolHighlight(_ req: DocumentHighlightRequest) async throws -> [DocumentHighlight]? {
-    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
+    let snapshot = try await self.latestSnapshot(for: req.textDocument.uri)
 
     let relatedIdentifiers = try await self.relatedIdentifiers(
       at: req.position,
@@ -867,7 +867,7 @@ extension SwiftLanguageService {
   }
 
   func retrieveQuickFixCodeActions(_ params: CodeActionRequest) async throws -> [CodeAction] {
-    let snapshot = try documentManager.latestSnapshot(params.textDocument.uri)
+    let snapshot = try await self.latestSnapshot(for: params.textDocument.uri)
     let buildSettings = await self.buildSettings(for: params.textDocument.uri)
     let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
       for: snapshot,
@@ -960,7 +960,7 @@ extension SwiftLanguageService {
       await semanticIndexManager?.prepareFileForEditorFunctionality(
         req.textDocument.uri.primaryFile ?? req.textDocument.uri
       )
-      let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
+      let snapshot = try await self.latestSnapshot(for: req.textDocument.uri)
       let buildSettings = await self.buildSettings(for: req.textDocument.uri)
       let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
         for: snapshot,

Sources/SourceKitLSP/Swift/VariableTypeInfo.swift

@@ -83,11 +83,12 @@ extension SwiftLanguageService {
     _ uri: DocumentURI,
     _ range: Range<Position>? = nil
   ) async throws -> [VariableTypeInfo] {
-    let snapshot = try documentManager.latestSnapshot(uri)
+    let snapshot = try await self.latestSnapshot(for: uri)
 
     let skreq = sourcekitd.dictionary([
       keys.request: requests.collectVariableType,
-      keys.sourceFile: snapshot.uri.pseudoPath,
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
       keys.compilerArgs: await self.buildSettings(for: uri)?.compilerArgs as [SKDRequestValue]?,
     ])