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

Author: lokesh-tr

Sources/LanguageServerProtocol/SupportTypes/DocumentURI.swift

@@ -76,7 +76,7 @@ public struct DocumentURI: Codable, Hashable, Sendable {
 
   public init(_ url: URL) {
     self.storage = url
-    assert(self.storage.scheme != nil, "Received invalid URI without a scheme '\(self.storage.absoluteString)'")
+    // assert(self.storage.scheme != nil, "Received invalid URI without a scheme '\(self.storage.absoluteString)'")
   }
 
   public init(filePath: String, isDirectory: Bool) {

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1674,10 +1674,16 @@ extension SourceKitLSPServer {
   }
 
   func executeCommand(_ req: ExecuteCommandRequest) async throws -> LSPAny? {
-    guard let uri = req.textDocument?.uri else {
+    guard var uri = req.textDocument?.uri else {
       logger.error("Attempted to perform executeCommand request without an URL")
       return nil
     }
+
+    let referenceDocumentURL = try? ReferenceDocumentURL(from: uri)
+    if let primaryFileURI = referenceDocumentURL?.primaryFile {
+      uri = primaryFileURI
+    }
+
     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")
-    }
+    var primaryFileURL = expandMacroCommand.textDocument.uri.arbitrarySchemeURL
+    let referenceDocumentURL = try? ReferenceDocumentURL(from: expandMacroCommand.textDocument.uri)
+    primaryFileURL = referenceDocumentURL?.primaryFile.arbitrarySchemeURL ?? primaryFileURL
 
     let expansion = try await self.refactoring(expandMacroCommand)
 
@@ -74,15 +89,32 @@ 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(
-              macroExpansionEditRange: macroEdit.range,
-              primaryFileURL: primaryFileURL,
-              selectionRange: expandMacroCommand.positionRange,
-              bufferName: bufferName
+          switch referenceDocumentURL {
+          case .macroExpansion(let parentReferenceDocument):
+            ReferenceDocumentURL.macroExpansion(
+              MacroExpansionReferenceDocumentURLData(
+                macroExpansionEditRange: macroEdit.range,
+                primaryFileURL: primaryFileURL,
+                selectionRange: expandMacroCommand.positionRange,
+                bufferName: bufferName
+              )
             )
-          )
+          case nil:
+            ReferenceDocumentURL.macroExpansion(
+              MacroExpansionReferenceDocumentURLData(
+                macroExpansionEditRange: macroEdit.range,
+                primaryFileURL: primaryFileURL,
+                selectionRange: expandMacroCommand.positionRange,
+                bufferName: bufferName
+              )
+            )
+          }
 
         macroExpansionReferenceDocumentURLs.append(macroExpansionReferenceDocumentURLData)
 
@@ -108,10 +140,19 @@ extension SwiftLanguageService {
       let expansionURIs = try macroExpansionReferenceDocumentURLs.map {
         return DocumentURI(try $0.url)
       }
+
       Task {
+        var uri = expandMacroCommand.textDocument.uri
+        var position = expandMacroCommand.positionRange.lowerBound
+
+        if let referenceDocumentURL, case let .macroExpansion(referenceDocumentURLData) = referenceDocumentURL {
+          uri = referenceDocumentURL.primaryFile
+          position = referenceDocumentURLData.macroExpansionEditRange.lowerBound
+        }
+
         let req = PeekDocumentsRequest(
-          uri: expandMacroCommand.textDocument.uri,
-          position: expandMacroCommand.positionRange.lowerBound,
+          uri: uri,
+          position: position,
           locations: expansionURIs
         )
 
@@ -180,21 +221,14 @@ 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)
-
     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

@@ -48,14 +48,16 @@ package struct MacroExpansionReferenceDocumentURLData {
   }
 
   package var queryItems: [URLQueryItem] {
-    [
-      URLQueryItem(name: Parameters.primaryFilePath, value: primaryFileURL.path(percentEncoded: false)),
+    var queryItems = [
       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.bufferName, value: bufferName),
+      URLQueryItem(name: Parameters.primaryFilePath, value: primaryFileURL.path(percentEncoded: false)),
     ]
+
+    return queryItems
   }
 
   package init(displayName: String, queryItems: [URLQueryItem]) throws {
@@ -71,7 +73,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 +84,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
+  /// 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

@@ -12,6 +12,7 @@
 
 import LanguageServerProtocol
 import SKLogging
+import SKSupport
 import SourceKitD
 
 protocol RefactoringResponse {
@@ -102,28 +103,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
+  /// 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` 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]?,
     ])