Make the folding range request return the request result as an async return value

Author: ahoppen

Sources/SourceKitLSP/Clang/ClangLanguageServer.swift

@@ -320,6 +320,27 @@ actor ClangLanguageServerShim: ToolchainLanguageServer, MessageHandler {
     }
   }
 
+  /// Forward the given request to `clangd`.
+  ///
+  /// This method calls `readyToHandleNextRequest` once the request has been
+  /// transmitted to `clangd` and another request can be safely transmitted to
+  /// `clangd` while guaranteeing ordering.
+  ///
+  /// The response of the request is  returned asynchronously as the return value.
+  func forwardRequestToClangd<R: RequestType>(_ request: R) async throws -> R.Response {
+    try await withCheckedThrowingContinuation { continuation in
+      _ = clangd.send(request, queue: clangdCommunicationQueue) { result in
+        switch result {
+        case .success(let response):
+          continuation.resume(returning: response)
+        case .failure(let error):
+          continuation.resume(throwing: error)
+        }
+      }
+    }
+    // FIXME: (async) Cancellation
+  }
+
   func _crash() {
     // Since `clangd` doesn't have a method to crash it, kill it.
 #if os(Windows)
@@ -553,11 +574,11 @@ extension ClangLanguageServerShim {
     forwardRequestToClangd(req)
   }
 
-  func foldingRange(_ req: Request<FoldingRangeRequest>) {
+  func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]? {
     if self.capabilities?.foldingRangeProvider?.isSupported == true {
-      forwardRequestToClangd(req)
+      return try await forwardRequestToClangd(req)
     } else {
-      req.reply(.success(nil))
+      return nil
     }
   }
 

Sources/SourceKitLSP/SourceKitServer.swift

@@ -269,6 +269,32 @@ public actor SourceKitServer {
     await requestHandler(request, workspace, languageService)
   }
 
+  private func withLanguageServiceAndWorkspace<RequestType: TextDocumentRequest>(
+    for request: Request<RequestType>,
+    requestHandler: @escaping (RequestType, Workspace, ToolchainLanguageServer) async throws -> RequestType.Response,
+    fallback: RequestType.Response
+  ) async {
+    let doc = request.params.textDocument.uri
+    guard let workspace = await self.workspaceForDocument(uri: doc) else {
+      return request.reply(.failure(.workspaceNotOpen(doc)))
+    }
+
+    guard let languageService = workspace.documentService[doc] else {
+      return request.reply(fallback)
+    }
+
+    do {
+      let result = try await requestHandler(request.params, workspace, languageService)
+      request.reply(.success(result))
+    } catch let error as ResponseError {
+      request.reply(.failure(error))
+    } catch is CancellationError {
+      request.reply(.failure(.cancelled))
+    } catch {
+      request.reply(.failure(.unknown("Unknown error: \(error)")))
+    }
+  }
+
 
   /// Send the given notification to the editor.
   public func sendNotificationToClient(_ notification: some NotificationType) {
@@ -1253,10 +1279,11 @@ extension SourceKitServer {
   }
 
   func foldingRange(
-    _ req: Request<FoldingRangeRequest>,
+    _ req: FoldingRangeRequest,
     workspace: Workspace,
-    languageService: ToolchainLanguageServer) async {
-    await languageService.foldingRange(req)
+    languageService: ToolchainLanguageServer
+  ) async throws -> [FoldingRange]? {
+    return try await languageService.foldingRange(req)
   }
 
   func documentSymbol(
@@ -1322,13 +1349,6 @@ extension SourceKitServer {
       return
     }
 
-    await self.fowardExecuteCommand(req, languageService: languageService)
-  }
-
-  func fowardExecuteCommand(
-    _ req: Request<ExecuteCommandRequest>,
-    languageService: ToolchainLanguageServer
-  ) async {
     let params = req.params
     let executeCommand = ExecuteCommandRequest(command: params.command,
                                                arguments: params.argumentsWithoutSourceKitMetadata)

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -984,13 +984,12 @@ extension SwiftLanguageServer {
     _ = handle
   }
 
-  public func foldingRange(_ req: Request<FoldingRangeRequest>) async {
+  public func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]? {
     let foldingRangeCapabilities = capabilityRegistry.clientCapabilities.textDocument?.foldingRange
-    let uri = req.params.textDocument.uri
+    let uri = req.textDocument.uri
     guard var snapshot = self.documentManager.latestSnapshot(uri) else {
-      log("failed to find snapshot for url \(req.params.textDocument.uri)")
-      req.reply(nil)
-      return
+      log("failed to find snapshot for url \(req.textDocument.uri)")
+      return nil
     }
 
     // FIXME: (async) We might not have computed the syntax tree yet. Wait until we have a syntax tree.
@@ -1000,16 +999,14 @@ extension SwiftLanguageServer {
       if let newSnapshot = documentManager.latestSnapshot(uri) {
         snapshot = newSnapshot
       } else {
-        log("failed to find snapshot for url \(req.params.textDocument.uri)")
-        req.reply(nil)
-        return
+        log("failed to find snapshot for url \(req.textDocument.uri)")
+        return nil
       }
     }
 
     guard let sourceFile = snapshot.tokens.syntaxTree else {
-      log("no lexical structure available for url \(req.params.textDocument.uri)")
-      req.reply(nil)
-      return
+      log("no lexical structure available for url \(req.textDocument.uri)")
+      return nil
     }
 
     final class FoldingRangeFinder: SyntaxVisitor {
@@ -1205,10 +1202,11 @@ extension SwiftLanguageServer {
       }
     }
 
+    try Task.checkCancellation()
+
     // If the limit is less than one, do nothing.
     if let limit = foldingRangeCapabilities?.rangeLimit, limit <= 0 {
-      req.reply([])
-      return
+      return []
     }
 
     let rangeFinder = FoldingRangeFinder(
@@ -1218,7 +1216,7 @@ extension SwiftLanguageServer {
     rangeFinder.walk(sourceFile)
     let ranges = rangeFinder.finalize()
 
-    req.reply(ranges.sorted())
+    return ranges.sorted()
   }
 
   public func codeAction(_ req: Request<CodeActionRequest>) async {

Sources/SourceKitLSP/ToolchainLanguageServer.swift

@@ -88,7 +88,7 @@ public protocol ToolchainLanguageServer: AnyObject {
   func declaration(_ request: Request<DeclarationRequest>) async -> Bool
 
   func documentSymbolHighlight(_ req: Request<DocumentHighlightRequest>) async
-  func foldingRange(_ req: Request<FoldingRangeRequest>) async
+  func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]?
   func documentSymbol(_ req: Request<DocumentSymbolRequest>) async
   func documentColor(_ req: Request<DocumentColorRequest>) async
   func documentSemanticTokens(_ req: Request<DocumentSemanticTokensRequest>) async