Support expansion of nested macros

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
-  }
 }