Make Nested Macro Expansions to work (Tested upto 3rd level)

Author: lokesh-tr

Sources/LanguageServerProtocol/SupportTypes/DocumentURI.swift

@@ -76,10 +76,7 @@ public struct DocumentURI: Codable, Hashable, Sendable {
 
   public init(_ url: URL) {
     self.storage = url
-
-    if !self.storage.absoluteString.hasPrefix("@") {
-      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/Swift/CursorInfo.swift

@@ -153,7 +153,7 @@ extension SwiftLanguageService {
       keys.cancelOnSubsequentRequest: 0,
       keys.offset: offsetRange.lowerBound,
       keys.length: offsetRange.upperBound != offsetRange.lowerBound ? offsetRange.count : nil,
-      keys.sourceFile: snapshot.uri.actualFile.pseudoPath,
+      keys.sourceFile: snapshot.uri.actualFilePath,
       keys.compilerArgs: await self.buildSettings(for: uri)?.compilerArgs as [SKDRequestValue]?,
     ])
 

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

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

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -80,7 +80,7 @@ extension SwiftLanguageService {
               MacroExpansionReferenceDocumentURLData(
                 macroExpansionEditRange: macroEdit.range,
                 primaryFileURL: primaryFileURL,
-                sourceFileURL: expandMacroCommand.textDocument.uri.arbitrarySchemeURL,
+                parentReferenceDocumentURL: expandMacroCommand.textDocument.uri.arbitrarySchemeURL,
                 selectionRange: expandMacroCommand.positionRange,
                 bufferName: bufferName
               )
@@ -90,7 +90,7 @@ extension SwiftLanguageService {
               MacroExpansionReferenceDocumentURLData(
                 macroExpansionEditRange: macroEdit.range,
                 primaryFileURL: primaryFileURL,
-                sourceFileURL: nil,
+                parentReferenceDocumentURL: nil,
                 selectionRange: expandMacroCommand.positionRange,
                 bufferName: bufferName
               )
@@ -203,10 +203,10 @@ extension SwiftLanguageService {
 
   func expandMacro(macroExpansionURLData: MacroExpansionReferenceDocumentURLData) async throws -> String {
     var expandMacroCommand: ExpandMacroCommand
-    if let sourceFileURL = macroExpansionURLData.sourceFileURL {
+    if let parentReferenceDocumentURL = macroExpansionURLData.parentReferenceDocumentURL {
       expandMacroCommand = ExpandMacroCommand(
         positionRange: macroExpansionURLData.selectionRange,
-        textDocument: TextDocumentIdentifier(DocumentURI(sourceFileURL))
+        textDocument: TextDocumentIdentifier(DocumentURI(parentReferenceDocumentURL))  // usage
       )
     } else {
       expandMacroCommand = ExpandMacroCommand(

Sources/SourceKitLSP/Swift/MacroExpansionReferenceDocumentURLData.swift

@@ -15,32 +15,34 @@ import LanguageServerProtocol
 import RegexBuilder
 
 /// Represents url of macro expansion reference document as follows:
-/// `sourcekit-lsp://swift-macro-expansion/LaCb-LcCd.swift?primaryFilePath=&sourceFilePath=&fromLine=&fromColumn=&toLine=&toColumn=&bufferName=`
+/// `sourcekit-lsp://swift-macro-expansion/LaCb-LcCd.swift?primaryFilePath=&fromLine=&fromColumn=&toLine=&toColumn=&bufferName=&parentReferenceDocumentURL=`
 ///
 /// Here,
 ///  - `LaCb-LcCd.swift`, the `displayName`, represents where the macro will expand to or
 ///    replace in the source file (i.e. `macroExpansionEditRange`)
 ///  - `primaryFilePath` denoting the URL of the source file
 ///  - `fromLine`, `fromColumn`, `toLine`, `toColumn` represents the cursor's `selectionRange`
 ///  - `bufferName` denotes the buffer name of the specific macro expansion edit
+///  - `parentReferenceDocumentURL` the reference document url of the macro
+///    which expands to the current macro (optional)
 package struct MacroExpansionReferenceDocumentURLData {
   package static let documentType = "swift-macro-expansion"
 
   package var primaryFileURL: URL
-  package var sourceFileURL: URL?
+  package var parentReferenceDocumentURL: URL?
   package var selectionRange: Range<Position>
   package var bufferName: String
   package var macroExpansionEditRange: Range<Position>
 
   package init(
     macroExpansionEditRange: Range<Position>,
     primaryFileURL: URL,
-    sourceFileURL: URL?,
+    parentReferenceDocumentURL: URL?,
     selectionRange: Range<Position>,
     bufferName: String
   ) {
     self.primaryFileURL = primaryFileURL
-    self.sourceFileURL = sourceFileURL
+    self.parentReferenceDocumentURL = parentReferenceDocumentURL
     self.selectionRange = selectionRange
     self.bufferName = bufferName
     self.macroExpansionEditRange = macroExpansionEditRange
@@ -51,7 +53,7 @@ package struct MacroExpansionReferenceDocumentURLData {
   }
 
   package var queryItems: [URLQueryItem] {
-    if let sourceFileURL {
+    if let parentReferenceDocumentURL {  // usage
       [
         URLQueryItem(name: Parameters.fromLine, value: String(selectionRange.lowerBound.line)),
         URLQueryItem(name: Parameters.fromColumn, value: String(selectionRange.lowerBound.utf16index)),
@@ -60,8 +62,8 @@ package struct MacroExpansionReferenceDocumentURLData {
         URLQueryItem(name: Parameters.bufferName, value: bufferName),
         URLQueryItem(name: Parameters.primaryFilePath, value: primaryFileURL.path(percentEncoded: false)),
         URLQueryItem(
-          name: Parameters.sourceFilePath,
-          value: sourceFileURL.absoluteString.addingPercentEncoding(withAllowedCharacters: .alphanumerics)
+          name: Parameters.parentReferenceDocumentURL,
+          value: parentReferenceDocumentURL.absoluteString.addingPercentEncoding(withAllowedCharacters: .alphanumerics)
         ),
       ]
     } else {
@@ -93,62 +95,30 @@ package struct MacroExpansionReferenceDocumentURLData {
       )
     }
 
-    var sourceFileURL: URL? = nil
-    if let sourceFilePath = queryItems.last(where: { $0.name == Parameters.sourceFilePath })?.value?
-      .removingPercentEncoding
-    {
-      guard let url = URL(string: sourceFilePath) else {
+    var parentReferenceDocumentURL: URL? = nil
+    if let parentReferenceDocumentURLString = queryItems.last(where: {
+      $0.name == Parameters.parentReferenceDocumentURL
+    })?.value?
+    .removingPercentEncoding {
+      guard let url = URL(string: parentReferenceDocumentURLString) else {
         throw ReferenceDocumentURLError(
           description: "Unable to parse source file url"
         )
       }
 
-      sourceFileURL = url
+      parentReferenceDocumentURL = url
     }
 
     self.primaryFileURL = primaryFileURL
-    self.sourceFileURL = sourceFileURL
+    self.parentReferenceDocumentURL = parentReferenceDocumentURL
     self.selectionRange =
       Position(line: fromLine, utf16index: fromColumn)..<Position(line: toLine, utf16index: toColumn)
     self.bufferName = bufferName
     self.macroExpansionEditRange = try Self.parse(displayName: displayName)
   }
 
-  package var actualFile: DocumentURI {
-    get throws {
-      guard let uri = try? DocumentURI(string: bufferName) else {
-        throw ReferenceDocumentURLError(
-          description: "Unable to retrieve actual file uri of macro expansion reference document"
-        )
-      }
-
-      return uri
-    }
-  }
-
-  package var sourceFile: DocumentURI {
-    get throws {
-      var uri: DocumentURI?
-      if let sourceFileURL {
-        let referenceDocumentURL = try ReferenceDocumentURL(from: sourceFileURL)
-        guard case let .macroExpansion(urlData) = referenceDocumentURL else {
-          throw ReferenceDocumentURLError(
-            description: "Unable to retrieve buffer name from source file"
-          )
-        }
-        uri = try? DocumentURI(string: urlData.bufferName)
-      } else {
-        uri = try? DocumentURI(string: bufferName)
-      }
-
-      guard let uri else {
-        throw ReferenceDocumentURLError(
-          description: "Unable to retrieve source file uri of macro expansion reference document"
-        )
-      }
-
-      return uri
-    }
+  package var actualFilePath: String {
+    bufferName
   }
 
   package var primaryFile: DocumentURI {
@@ -157,7 +127,7 @@ package struct MacroExpansionReferenceDocumentURLData {
 
   private struct Parameters {
     static let primaryFilePath = "primaryFilePath"
-    static let sourceFilePath = "sourceFilePath"
+    static let parentReferenceDocumentURL = "parentReferenceDocumentURL"
     static let fromLine = "fromLine"
     static let fromColumn = "fromColumn"
     static let toLine = "toLine"
@@ -173,8 +143,8 @@ package struct MacroExpansionReferenceDocumentURLData {
     }
 
     var result = URL(string: urlWithoutEncoding)
-    if urlWithoutEncoding.contains(Parameters.sourceFilePath) {
-      let location = urlWithoutEncoding.firstRange(of: "&\(Parameters.sourceFilePath)=")
+    if urlWithoutEncoding.contains(Parameters.parentReferenceDocumentURL) {
+      let location = urlWithoutEncoding.firstRange(of: "&\(Parameters.parentReferenceDocumentURL)=")
       guard let location else {
         return nil
       }
@@ -184,7 +154,7 @@ package struct MacroExpansionReferenceDocumentURLData {
       let range = start..<end
 
       guard
-        let sourceFileURL = String(urlWithoutEncoding[range]).addingPercentEncoding(
+        let parentReferenceDocumentURL = String(urlWithoutEncoding[range]).addingPercentEncoding(
           withAllowedCharacters: .alphanumerics
         )
       else {
@@ -195,9 +165,10 @@ package struct MacroExpansionReferenceDocumentURLData {
       let firstEnd = location.lowerBound
       let firstRange = firstStart..<firstEnd
 
-      let urlExceptSourceFileURL = String(urlWithoutEncoding[firstRange])
+      let urlExceptParentReferenceDocumentURL = String(urlWithoutEncoding[firstRange])
 
-      let finalURLString = urlExceptSourceFileURL + "&\(Parameters.sourceFilePath)=" + sourceFileURL
+      let finalURLString =
+        urlExceptParentReferenceDocumentURL + "&\(Parameters.parentReferenceDocumentURL)=" + parentReferenceDocumentURL
       result = URL(string: finalURLString)
     }
 

Sources/SourceKitLSP/Swift/RefactoringResponse.swift

@@ -103,35 +103,63 @@ extension SwiftLanguageService {
     let uri = refactorCommand.textDocument.uri
     let referenceDocument = try? ReferenceDocumentURL(from: uri)
 
-    let snapshot = try self.documentManager.latestSnapshot(referenceDocument?.primaryFile ?? uri)
+    var snapshot: DocumentSnapshot
+    var opened = false
+
+    if let referenceDocument {
+      let primaryFileSnapshot = try self.documentManager.latestSnapshot(referenceDocument.primaryFile)
+
+      var newSnapshot: DocumentSnapshot
+      do {
+        newSnapshot = try self.documentManager.latestSnapshot(uri)
+      } catch {
+        newSnapshot = try await documentManager.open(
+          uri,
+          language: primaryFileSnapshot.language,
+          version: primaryFileSnapshot.version,
+          text: getReferenceDocument(GetReferenceDocumentRequest(uri: refactorCommand.textDocument.uri)).content
+        )
+        opened = true
+      }
+      snapshot = newSnapshot
+    } else {
+      snapshot = try self.documentManager.latestSnapshot(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]?
+
+    if opened {
+      try? documentManager.close(snapshot.uri)
+    }
 
     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.actualFilePath,
       // 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,
     ])
 
     if let referenceDocument {
-      skreq.set(keys.sourceFile, to: try referenceDocument.actualFile.pseudoPath)
       skreq.set(keys.primaryFile, to: referenceDocument.primaryFile.pseudoPath)
-    } else {
-      skreq.set(keys.sourceFile, to: uri.actualFile.pseudoPath)
     }
 
     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

@@ -89,28 +89,17 @@ package enum ReferenceDocumentURL {
     }
   }
 
-  /// The URI of the document which originally contains the contents of the reference document. This is used to
+  /// The file path of the document which originally contains the contents of the reference document. This is used to
   /// communicate with sourcekitd.
-  var actualFile: DocumentURI {
-    get throws {
-      switch self {
-      case let .macroExpansion(data):
-        return try data.actualFile
-      }
-    }
-  }
-
-  var sourceFile: DocumentURI {
-    get throws {
-      switch self {
-      case let .macroExpansion(data):
-        return try data.sourceFile
-      }
+  var actualFilePath: String {
+    switch self {
+    case let .macroExpansion(data):
+      data.actualFilePath
     }
   }
 
-  /// 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 URI of the document from which this reference document and its parent 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):
@@ -120,13 +109,11 @@ package enum ReferenceDocumentURL {
 }
 
 extension DocumentURI {
-  var actualFile: DocumentURI {
-    if let referenceDocument = try? ReferenceDocumentURL(from: self),
-      let actualFile = try? referenceDocument.actualFile
-    {
-      actualFile
+  var actualFilePath: String {
+    if let referenceDocument = try? ReferenceDocumentURL(from: self) {
+      referenceDocument.actualFilePath
     } else {
-      self
+      self.pseudoPath
     }
   }
 

Sources/SourceKitLSP/Swift/RelatedIdentifiers.swift

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

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -26,7 +26,7 @@ extension SwiftLanguageService {
 
     let skreq = sourcekitd.dictionary([
       keys.request: requests.semanticTokens,
-      keys.sourceFile: snapshot.uri.actualFile.pseudoPath,
+      keys.sourceFile: snapshot.uri.actualFilePath,
       keys.compilerArgs: buildSettings.compilerArgs as [SKDRequestValue],
     ])
 

Sources/SourceKitLSP/Swift/VariableTypeInfo.swift

@@ -87,7 +87,7 @@ extension SwiftLanguageService {
 
     let skreq = sourcekitd.dictionary([
       keys.request: requests.collectVariableType,
-      keys.sourceFile: snapshot.uri.actualFile.pseudoPath,
+      keys.sourceFile: snapshot.uri.actualFilePath,
       keys.compilerArgs: await self.buildSettings(for: uri)?.compilerArgs as [SKDRequestValue]?,
     ])