Add Semantic Functionality to Macro Expansion Reference Documents (including Nested Macro Expansion) 🚦

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1674,10 +1674,18 @@ extension SourceKitLSPServer {
   }
 
   func executeCommand(_ req: ExecuteCommandRequest) async throws -> LSPAny? {
-    guard let uri = req.textDocument?.uri else {
+    guard let reqURI = req.textDocument?.uri,
+      let uri =
+        if let primaryFileURI = (try? ReferenceDocumentURL(from: reqURI))?.primaryFile {
+          primaryFileURI
+        } else {
+          reqURI
+        }
+    else {
       logger.error("Attempted to perform executeCommand request without an URL")
       return nil
     }
+
     guard let workspace = await workspaceForDocument(uri: uri) else {
       throw ResponseError.workspaceNotOpen(uri)
     }

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -158,10 +158,10 @@ 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.actualFilePath,
+      keys.primaryFile: (try? ReferenceDocumentURL(from: snapshot.uri))?.primaryFile.pseudoPath,
       keys.compilerArgs: await self.buildSettings(for: uri)?.compilerArgs as [SKDRequestValue]?,
     ])
-
     appendAdditionalParameters?(skreq)
 
     let dict = try await sendSourcekitdRequest(skreq, fileContents: snapshot.text)

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.actualFilePath,
+      keys.primaryFile: (try? ReferenceDocumentURL(from: snapshot.uri))?.primaryFile.pseudoPath,
       keys.compilerArgs: compilerArgs as [SKDRequestValue],
     ])
 

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -44,6 +44,21 @@ struct MacroExpansion: RefactoringResponse {
   }
 }
 
+actor GeneratedMacroExpansionsStorage {
+  static var shared = GeneratedMacroExpansionsStorage()
+  private init() {}
+
+  private var generatedMacroExpansions: [String: String] = [:]
+
+  func addOrUpdateMacroExpansion(havingBufferName bufferName: String, withContents content: String) {
+    generatedMacroExpansions[bufferName] = content
+  }
+
+  func retrieveAndDeleteMacroExpansion(havingBufferName bufferName: String) -> String? {
+    return generatedMacroExpansions.removeValue(forKey: bufferName)
+  }
+}
+
 extension SwiftLanguageService {
   /// Handles the `ExpandMacroCommand`.
   ///
@@ -62,9 +77,9 @@ extension SwiftLanguageService {
       throw ResponseError.unknown("Connection to the editor closed")
     }
 
-    guard let primaryFileURL = expandMacroCommand.textDocument.uri.fileURL else {
-      throw ResponseError.unknown("Given URI is not a file URL")
-    }
+    let referenceDocumentURL = try? ReferenceDocumentURL(from: expandMacroCommand.textDocument.uri)
+    let primaryFileURL =
+      referenceDocumentURL?.primaryFile.arbitrarySchemeURL ?? expandMacroCommand.textDocument.uri.arbitrarySchemeURL
 
     let expansion = try await self.refactoring(expandMacroCommand)
 
@@ -74,6 +89,11 @@ extension SwiftLanguageService {
     var macroExpansionReferenceDocumentURLs: [ReferenceDocumentURL] = []
     for macroEdit in expansion.edits {
       if let bufferName = macroEdit.bufferName {
+        await GeneratedMacroExpansionsStorage.shared.addOrUpdateMacroExpansion(
+          havingBufferName: bufferName,
+          withContents: macroEdit.newText
+        )
+
         let macroExpansionReferenceDocumentURLData =
           ReferenceDocumentURL.macroExpansion(
             MacroExpansionReferenceDocumentURLData(
@@ -108,10 +128,24 @@ extension SwiftLanguageService {
       let expansionURIs = try macroExpansionReferenceDocumentURLs.map {
         return DocumentURI(try $0.url)
       }
+
       Task {
+        let (uri, position) =
+          if let referenceDocumentURL, case let .macroExpansion(referenceDocumentURLData) = referenceDocumentURL {
+            (
+              referenceDocumentURL.primaryFile,
+              referenceDocumentURLData.macroExpansionEditRange.lowerBound
+            )
+          } else {
+            (
+              expandMacroCommand.textDocument.uri,
+              expandMacroCommand.positionRange.lowerBound
+            )
+          }
+
         let req = PeekDocumentsRequest(
-          uri: expandMacroCommand.textDocument.uri,
-          position: expandMacroCommand.positionRange.lowerBound,
+          uri: uri,
+          position: position,
           locations: expansionURIs
         )
 
@@ -179,22 +213,15 @@ extension SwiftLanguageService {
     }
   }
 
-  func expandMacro(macroExpansionURLData: MacroExpansionReferenceDocumentURLData) async throws -> String {
-    let expandMacroCommand = ExpandMacroCommand(
-      positionRange: macroExpansionURLData.selectionRange,
-      textDocument: TextDocumentIdentifier(macroExpansionURLData.primaryFile)
-    )
-
-    let expansion = try await self.refactoring(expandMacroCommand)
-
+  func getMacroExpansion(macroExpansionURLData: MacroExpansionReferenceDocumentURLData) async throws -> String {
     guard
-      let macroExpansionEdit = expansion.edits.filter({
-        $0.bufferName == macroExpansionURLData.bufferName
-      }).only
+      let content = await GeneratedMacroExpansionsStorage.shared.retrieveAndDeleteMacroExpansion(
+        havingBufferName: macroExpansionURLData.bufferName
+      )
     else {
       throw ResponseError.unknown("Macro expansion edit doesn't exist")
     }
 
-    return macroExpansionEdit.newText
+    return content
   }
 }

Sources/SourceKitLSP/Swift/MacroExpansionReferenceDocumentURLData.swift

@@ -71,7 +71,7 @@ package struct MacroExpansionReferenceDocumentURLData {
 
     guard let primaryFileURL = URL(string: "file://\(primaryFilePath)") else {
       throw ReferenceDocumentURLError(
-        description: "Unable to parse source file url"
+        description: "Unable to parse primary file url"
       )
     }
 
@@ -82,6 +82,61 @@ package struct MacroExpansionReferenceDocumentURLData {
     self.macroExpansionEditRange = try Self.parse(displayName: displayName)
   }
 
+  /// The file path of the document that originally contains the contents of this reference document.
+  /// This is used since `sourcekitd` cannot understand reference document urls.
+  ///
+  /// For any `ReferenceDocumentURL.macroExpansion`, its `actualFilePath` will be its sourcekitd `bufferName`
+  ///
+  /// *Example:*
+  ///
+  /// User's source File:
+  /// URL: `file:///path/to/swift_file.swift`
+  /// ```swift
+  /// let a = 10
+  /// let b = 5
+  /// print(#stringify(a + b))
+  /// ```
+  ///
+  /// 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`
+  /// ```swift
+  /// (a + b, "a + b")
+  /// ```
+  ///
+  /// Here the `actualFilePath` of the reference document url is `@__swift_macro_..._Stringify_.swift`
+  ///
+  /// *NOTE*: In case of nested macro expansion reference documents, the `actualFilePath` will be their corresponding
+  /// `bufferName`s
+  package var actualFilePath: String {
+    bufferName
+  }
+
+  /// The URI of the document from which this reference document was derived.
+  /// This is used to determine the workspace and language service that is used to generate the reference document.
+  ///
+  /// *Example:*
+  ///
+  /// User's source File:
+  /// URL: `file://path/to/swift_file.swift`
+  /// ```swift
+  /// let a = 10
+  /// let b = 5
+  /// print(#stringify(a + b))
+  /// ```
+  ///
+  /// 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`
+  /// ```swift
+  /// (a + b, "a + b")
+  /// ```
+  ///
+  /// Here the `primaryFile` of the reference document url is a `DocumentURI`
+  /// with the following url: `file:///path/to/swift_file.swift`
+  ///
+  /// *NOTE*: In case of nested macro expansion reference documents, they all will have the same `primaryFile`
+  /// as that of the first macro expansion reference document i.e. `primaryFile` doesn't change.
   package var primaryFile: DocumentURI {
     DocumentURI(primaryFileURL)
   }

Sources/SourceKitLSP/Swift/OpenInterface.swift

@@ -44,12 +44,20 @@ extension SwiftLanguageService {
         symbol: symbol
       )
     } else {
+      let primaryDocument =
+        if let referenceDocument = try? ReferenceDocumentURL(from: document) {
+          referenceDocument.primaryFile
+        } else {
+          document
+        }
+
       let interfaceInfo = try await self.generatedInterfaceInfo(
-        document: document,
+        document: primaryDocument,
         moduleName: moduleName,
         groupName: groupName,
         interfaceURI: interfaceDocURI
       )
+
       try interfaceInfo.contents.write(to: interfaceFilePath, atomically: true, encoding: String.Encoding.utf8)
       let snapshot = DocumentSnapshot(
         uri: interfaceDocURI,

Sources/SourceKitLSP/Swift/RefactoringResponse.swift

@@ -102,28 +102,35 @@ extension SwiftLanguageService {
 
     let uri = refactorCommand.textDocument.uri
     let snapshot = try self.documentManager.latestSnapshot(uri)
+    let referenceDocument = try? ReferenceDocumentURL(from: snapshot.uri)
+
     let line = refactorCommand.positionRange.lowerBound.line
     let utf16Column = refactorCommand.positionRange.lowerBound.utf16index
     let utf8Column = snapshot.lineTable.utf8ColumnAt(line: line, utf16Column: utf16Column)
+    let length = snapshot.utf8OffsetRange(of: refactorCommand.positionRange).count
+    let buildSettings =
+      await self.buildSettings(for: referenceDocument?.primaryFile ?? snapshot.uri)?.compilerArgs as [SKDRequestValue]?
 
     let skreq = sourcekitd.dictionary([
       keys.request: self.requests.semanticRefactoring,
       // Preferred name for e.g. an extracted variable.
       // Empty string means sourcekitd chooses a name automatically.
       keys.name: "",
-      keys.sourceFile: uri.pseudoPath,
+      keys.sourceFile: uri.actualFilePath,
+      keys.primaryFile: referenceDocument?.primaryFile.pseudoPath,
       // LSP is zero based, but this request is 1 based.
       keys.line: line + 1,
       keys.column: utf8Column + 1,
-      keys.length: snapshot.utf8OffsetRange(of: refactorCommand.positionRange).count,
+      keys.length: length,
       keys.actionUID: self.sourcekitd.api.uid_get_from_cstr(refactorCommand.actionString)!,
-      keys.compilerArgs: await self.buildSettings(for: snapshot.uri)?.compilerArgs as [SKDRequestValue]?,
+      keys.compilerArgs: buildSettings,
     ])
 
     let dict = try await sendSourcekitdRequest(skreq, fileContents: snapshot.text)
     guard let refactor = T.Response(refactorCommand.title, dict, snapshot, self.keys) else {
       throw SemanticRefactoringError.noEditsNeeded(uri)
     }
+
     return refactor
   }
 }

Sources/SourceKitLSP/Swift/ReferenceDocumentURL.swift

@@ -63,7 +63,10 @@ package enum ReferenceDocumentURL {
 
     switch documentType {
     case MacroExpansionReferenceDocumentURLData.documentType:
-      guard let queryItems = URLComponents(string: url.absoluteString)?.queryItems else {
+      guard
+        let queryItems = URLComponents(string: url.absoluteString.removingPercentEncoding ?? url.absoluteString)?
+          .queryItems
+      else {
         throw ReferenceDocumentURLError(
           description: "No queryItems passed for macro expansion reference document: \(url)"
         )
@@ -85,8 +88,17 @@ package enum ReferenceDocumentURL {
     }
   }
 
-  /// The URI of the document from which this reference document was derived. This is used to determine the
-  /// workspace and language service that is used to generate the reference document.
+  /// The file path of the document that originally contains the contents of this reference document.
+  /// This is used since sourcekitd cannot understand reference document urls.
+  var actualFilePath: String {
+    switch self {
+    case let .macroExpansion(data):
+      data.actualFilePath
+    }
+  }
+
+  /// The URI of the source file from which this reference document was derived.
+  /// This is used to determine the workspace and language service that is used to generate the reference document.
   var primaryFile: DocumentURI {
     switch self {
     case let .macroExpansion(data):
@@ -95,6 +107,20 @@ package enum ReferenceDocumentURL {
   }
 }
 
+extension DocumentURI {
+  /// The file path of the document that originally contains the contents of the reference document
+  /// of this `DocumentURI`. This is used since sourcekitd cannot understand reference document urls.
+  ///
+  /// If the `DocumentURI` is not a reference document, this gives the `pseudoPath`
+  var actualFilePath: String {
+    if let referenceDocument = try? ReferenceDocumentURL(from: self) {
+      referenceDocument.actualFilePath
+    } else {
+      self.pseudoPath
+    }
+  }
+}
+
 package struct ReferenceDocumentURLError: Error, CustomStringConvertible {
   package var description: String
 

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.actualFilePath,
+      keys.primaryFile: (try? ReferenceDocumentURL(from: 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.actualFilePath,
+      keys.primaryFile: (try? ReferenceDocumentURL(from: snapshot.uri))?.primaryFile.pseudoPath,
       keys.compilerArgs: buildSettings.compilerArgs as [SKDRequestValue],
     ])