Support expansion of nested macros

The basic idea is that a `sourcekit-lsp://swift-macro-expansion` URL should have sufficient information to reconstruct the contents of that macro buffer without relying on any state in SourceKit-LSP. The benefit of not having any cross-request state in SourceKit-LSP is that an editor might can send the `workspace/getReferenceDocument` request at any time and it will succeed independent of the previous requests. Furthermore, we can always get the contents of the macro expansion to form a `DocumentSnapshot`, which can be used to provide semantic functionality inside macro expansion buffers.

To do that, the `sourcekit-lsp:` URL scheme was changed to have a parent instead of a `primary`, which is the URI of the document that the buffer was expanded from. For nested macro expansions, this will be a `sourcekit-lsp://swift-macro-expansion` URL itself.

With that parent, we can reconstruct the macro expansion chain all the way from the primary source file. To avoid sending the same expand macro request to sourcekitd all the time, we introduce `MacroExpansionManager`, which caches the last 10 macro expansions.

`SwiftLanguageService` now has a `latestSnapshot` method that returns the contents of the reference document when asked for a reference document URL and only consults the document manager for other URIs. To support semantic functionality in macro expansion buffers, we need to call that `latestSnapshot` method so we have a document snapshot of the macro expansion buffer for position conversions and pass the following to the sourcekitd requests.
```
keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
```

We should consider if there’s a way to make the `latestSnapshot` method on `documentManager` less accessible so that the method which also returns snapshots for reference documents is the one being used by default.

Co-Authored-By: Lokesh T R <[email protected]>

Author: ahoppen

Sources/SourceKitLSP/CMakeLists.txt

@@ -51,7 +51,6 @@ target_sources(SourceKitLSP PRIVATE
   Swift/OpenInterface.swift
   Swift/RefactoringResponse.swift
   Swift/RefactoringEdit.swift
-  Swift/RefactorCommand.swift
   Swift/ReferenceDocumentURL.swift
   Swift/RelatedIdentifiers.swift
   Swift/RewriteSourceKitPlaceholders.swift

Sources/SourceKitLSP/MessageHandlingDependencyTracker.swift

@@ -179,6 +179,8 @@ enum MessageHandlingDependencyTracker: DependencyTracker {
       } else {
         self = .freestanding
       }
+    case let request as GetReferenceDocumentRequest:
+      self = .documentRequest(request.uri)
     case is InitializeRequest:
       self = .globalConfigurationChange
     case is InlayHintRefreshRequest:

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -290,6 +290,7 @@ package actor SourceKitLSPServer {
   }
 
   package func workspaceForDocument(uri: DocumentURI) async -> Workspace? {
+    let uri = uri.primaryFile ?? uri
     if let cachedWorkspace = self.uriToWorkspaceCache[uri]?.value {
       return cachedWorkspace
     }
@@ -1693,8 +1694,7 @@ extension SourceKitLSPServer {
   }
 
   func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
-    let referenceDocumentURL = try ReferenceDocumentURL(from: req.uri)
-    let primaryFileURI = referenceDocumentURL.primaryFile
+    let primaryFileURI = try ReferenceDocumentURL(from: req.uri).primaryFile
 
     guard let workspace = await workspaceForDocument(uri: primaryFileURI) else {
       throw ResponseError.workspaceNotOpen(primaryFileURI)

Sources/SourceKitLSP/Swift/ExpandMacroCommand.swift

@@ -13,9 +13,7 @@
 import LanguageServerProtocol
 import SourceKitD
 
-package struct ExpandMacroCommand: RefactorCommand {
-  typealias Response = MacroExpansion
-
+package struct ExpandMacroCommand: SwiftCommand {
   package static let identifier: String = "expand.macro.command"
 
   /// The name of this refactoring action.

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -14,6 +14,8 @@ import Crypto
 import Foundation
 import LanguageServerProtocol
 import SKLogging
+import SKOptions
+import SKSupport
 import SourceKitD
 
 /// Detailed information about the result of a macro expansion operation.
@@ -44,6 +46,141 @@ struct MacroExpansion: RefactoringResponse {
   }
 }
 
+/// Caches the contents of macro expansions that were recently requested by the user.
+actor MacroExpansionManager {
+  private struct CacheEntry {
+    // Key
+    let snapshotID: DocumentSnapshot.ID
+    let range: Range<Position>
+    let buildSettings: SwiftCompileCommand?
+
+    // Value
+    let value: [RefactoringEdit]
+
+    fileprivate init(
+      snapshot: DocumentSnapshot,
+      range: Range<Position>,
+      buildSettings: SwiftCompileCommand?,
+      value: [RefactoringEdit]
+    ) {
+      self.snapshotID = snapshot.id
+      self.range = range
+      self.buildSettings = buildSettings
+      self.value = value
+    }
+  }
+
+  init(swiftLanguageService: SwiftLanguageService?) {
+    self.swiftLanguageService = swiftLanguageService
+  }
+
+  private weak var swiftLanguageService: SwiftLanguageService?
+
+  /// The number of macro expansions to cache.
+  ///
+  /// - Note: This should be bigger than the maximum expansion depth of macros a user might do to avoid re-generating
+  ///   all parent macros to a nested macro expansion's buffer. 10 seems to be big enough for that because it's
+  ///   unlikely that a macro will expand to more than 10 levels.
+  private let cacheSize = 10
+
+  /// The cache that stores reportTasks for a combination of uri, range and build settings.
+  ///
+  /// Conceptually, this is a dictionary. To prevent excessive memory usage we
+  /// only keep `cacheSize` entries within the array. Older entries are at the
+  /// end of the list, newer entries at the front.
+  private var cache: [CacheEntry] = []
+
+  /// Return the text of the macro expansion referenced by `macroExpansionURLData`.
+  func macroExpansion(
+    for macroExpansionURLData: MacroExpansionReferenceDocumentURLData
+  ) async throws -> String {
+    let expansions = try await macroExpansions(
+      in: macroExpansionURLData.parent,
+      at: macroExpansionURLData.selectionRange
+    )
+    guard let expansion = expansions.filter({ $0.bufferName == macroExpansionURLData.bufferName }).only else {
+      throw ResponseError.unknown("Failed to find macro expansion for \(macroExpansionURLData.bufferName).")
+    }
+    return expansion.newText
+  }
+
+  func macroExpansions(
+    in uri: DocumentURI,
+    at range: Range<Position>
+  ) async throws -> [RefactoringEdit] {
+    guard let swiftLanguageService else {
+      // `SwiftLanguageService` has been destructed. We are tearing down the language server. Nothing left to do.
+      throw ResponseError.unknown("Connection to the editor closed")
+    }
+
+    let snapshot = try await swiftLanguageService.latestSnapshot(for: uri)
+    let buildSettings = await swiftLanguageService.buildSettings(for: uri)
+
+    if let cacheEntry = cache.first(where: {
+      $0.snapshotID == snapshot.id && $0.range == range && $0.buildSettings == buildSettings
+    }) {
+      return cacheEntry.value
+    }
+    let macroExpansions = try await macroExpansionsImpl(in: snapshot, at: range, buildSettings: buildSettings)
+    cache.insert(
+      CacheEntry(snapshot: snapshot, range: range, buildSettings: buildSettings, value: macroExpansions),
+      at: 0
+    )
+
+    while cache.count > cacheSize {
+      cache.removeLast()
+    }
+
+    return macroExpansions
+  }
+
+  private func macroExpansionsImpl(
+    in snapshot: DocumentSnapshot,
+    at range: Range<Position>,
+    buildSettings: SwiftCompileCommand?
+  ) async throws -> [RefactoringEdit] {
+    guard let swiftLanguageService else {
+      // `SwiftLanguageService` has been destructed. We are tearing down the language server. Nothing left to do.
+      throw ResponseError.unknown("Connection to the editor closed")
+    }
+    let keys = swiftLanguageService.keys
+
+    let line = range.lowerBound.line
+    let utf16Column = range.lowerBound.utf16index
+    let utf8Column = snapshot.lineTable.utf8ColumnAt(line: line, utf16Column: utf16Column)
+    let length = snapshot.utf8OffsetRange(of: range).count
+
+    let skreq = swiftLanguageService.sourcekitd.dictionary([
+      keys.request: swiftLanguageService.requests.semanticRefactoring,
+      // Preferred name for e.g. an extracted variable.
+      // Empty string means sourcekitd chooses a name automatically.
+      keys.name: "",
+      keys.sourceFile: snapshot.uri.sourcekitdSourceFile,
+      keys.primaryFile: snapshot.uri.primaryFile?.pseudoPath,
+      // LSP is zero based, but this request is 1 based.
+      keys.line: line + 1,
+      keys.column: utf8Column + 1,
+      keys.length: length,
+      keys.actionUID: swiftLanguageService.sourcekitd.api.uid_get_from_cstr("source.refactoring.kind.expand.macro")!,
+      keys.compilerArgs: buildSettings?.compilerArgs as [SKDRequestValue]?,
+    ])
+
+    let dict = try await swiftLanguageService.sendSourcekitdRequest(
+      skreq,
+      fileContents: snapshot.text
+    )
+    guard let expansions = [RefactoringEdit](dict, snapshot, keys) else {
+      throw SemanticRefactoringError.noEditsNeeded(snapshot.uri)
+    }
+    return expansions
+  }
+
+  /// Remove all cached macro expansions for the given primary file, eg. because the macro's plugin might have changed.
+  func purge(primaryFile: DocumentURI) {
+    cache.removeAll { $0.snapshotID.uri.primaryFile ?? $0.snapshotID.uri == primaryFile }
+  }
+}
+
 extension SwiftLanguageService {
   /// Handles the `ExpandMacroCommand`.
   ///
@@ -62,23 +199,30 @@ 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 primaryFileDisplayName =
+      switch try? ReferenceDocumentURL(from: expandMacroCommand.textDocument.uri) {
+      case .macroExpansion(let data):
+        data.bufferName
+      case nil:
+        expandMacroCommand.textDocument.uri.fileURL?.lastPathComponent ?? expandMacroCommand.textDocument.uri.pseudoPath
+      }
 
-    let expansion = try await self.refactoring(expandMacroCommand)
+    let expansions = try await macroExpansionManager.macroExpansions(
+      in: expandMacroCommand.textDocument.uri,
+      at: expandMacroCommand.positionRange
+    )
 
     var completeExpansionFileContent = ""
     var completeExpansionDirectoryName = ""
 
     var macroExpansionReferenceDocumentURLs: [ReferenceDocumentURL] = []
-    for macroEdit in expansion.edits {
+    for macroEdit in expansions {
       if let bufferName = macroEdit.bufferName {
         let macroExpansionReferenceDocumentURLData =
           ReferenceDocumentURL.macroExpansion(
             MacroExpansionReferenceDocumentURLData(
               macroExpansionEditRange: macroEdit.range,
-              primaryFileURL: primaryFileURL,
+              parent: expandMacroCommand.textDocument.uri,
               selectionRange: expandMacroCommand.positionRange,
               bufferName: bufferName
             )
@@ -90,7 +234,7 @@ extension SwiftLanguageService {
 
         let editContent =
           """
-          // \(primaryFileURL.lastPathComponent) @ \(macroEdit.range.lowerBound.line + 1):\(macroEdit.range.lowerBound.utf16index + 1) - \(macroEdit.range.upperBound.line + 1):\(macroEdit.range.upperBound.utf16index + 1)
+          // \(primaryFileDisplayName) @ \(macroEdit.range.lowerBound.line + 1):\(macroEdit.range.lowerBound.utf16index + 1) - \(macroEdit.range.upperBound.line + 1):\(macroEdit.range.upperBound.utf16index + 1)
           \(macroEdit.newText)
 
 
@@ -154,7 +298,7 @@ extension SwiftLanguageService {
       }
 
       completeExpansionFilePath =
-        completeExpansionFilePath.appendingPathComponent(primaryFileURL.lastPathComponent)
+        completeExpansionFilePath.appendingPathComponent(primaryFileDisplayName)
       do {
         try completeExpansionFileContent.write(to: completeExpansionFilePath, atomically: true, encoding: .utf8)
       } catch {
@@ -178,23 +322,4 @@ 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
-    else {
-      throw ResponseError.unknown("Macro expansion edit doesn't exist")
-    }
-
-    return macroExpansionEdit.newText
-  }
 }