Make DocumentManager.latestSnapshot throw if no snapshot exists for the URI

This simplifies most calls that would log an error + return an empty response if no document was found.

Author: ahoppen

Sources/SourceKitLSP/DocumentManager.swift

@@ -184,10 +184,10 @@ public final class DocumentManager {
     }
   }
 
-  public func latestSnapshot(_ uri: DocumentURI) -> DocumentSnapshot? {
-    return queue.sync {
+  public func latestSnapshot(_ uri: DocumentURI) throws -> DocumentSnapshot {
+    return try queue.sync {
       guard let document = documents[uri] else {
-        return nil
+        throw ResponseError.unknown("Failed to find snapshot for '\(uri)'")
       }
       return document.latestSnapshot
     }

Sources/SourceKitLSP/SourceKitServer.swift

@@ -456,7 +456,7 @@ public actor SourceKitServer {
       guard workspace.documentService[documentUri] === languageService else {
         continue
       }
-      guard let snapshot = self.documentManager.latestSnapshot(documentUri) else {
+      guard let snapshot = try? self.documentManager.latestSnapshot(documentUri) else {
         // The document has been closed since we retrieved its URI. We don't care about it anymore.
         continue
       }
@@ -1293,7 +1293,7 @@ extension SourceKitServer {
       let oldWorkspace = preChangeWorkspaces[docUri]
       let newWorkspace = await self.workspaceForDocument(uri: docUri)
       if newWorkspace !== oldWorkspace {
-        guard let snapshot = documentManager.latestSnapshot(docUri) else {
+        guard let snapshot = try? documentManager.latestSnapshot(docUri) else {
           continue
         }
         if let oldWorkspace = oldWorkspace {

Sources/SourceKitLSP/Swift/CodeCompletion.swift

@@ -18,10 +18,7 @@ import SourceKitD
 extension SwiftLanguageServer {
 
   public func completion(_ req: CompletionRequest) async throws -> CompletionList {
-    guard let snapshot = documentManager.latestSnapshot(req.textDocument.uri) else {
-      logger.error("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-      return CompletionList(isIncomplete: true, items: [])
-    }
+    let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
 
     guard let completionPos = adjustCompletionLocation(req.position, in: snapshot) else {
       logger.error("invalid completion position \(req.position, privacy: .public)")

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -50,10 +50,6 @@ struct CursorInfo {
 
 /// An error from a cursor info request.
 enum CursorInfoError: Error, Equatable {
-
-  /// The given URL is not a known document.
-  case unknownDocument(DocumentURI)
-
   /// The given range is not valid in the document snapshot.
   case invalidRange(Range<Position>)
 
@@ -64,8 +60,6 @@ enum CursorInfoError: Error, Equatable {
 extension CursorInfoError: CustomStringConvertible {
   var description: String {
     switch self {
-    case .unknownDocument(let url):
-      return "failed to find snapshot for url \(url)"
     case .invalidRange(let range):
       return "invalid range \(range)"
     case .responseError(let error):
@@ -90,9 +84,7 @@ extension SwiftLanguageServer {
     _ range: Range<Position>,
     additionalParameters appendAdditionalParameters: ((SKDRequestDictionary) -> Void)? = nil
   ) async throws -> CursorInfo? {
-    guard let snapshot = documentManager.latestSnapshot(uri) else {
-      throw CursorInfoError.unknownDocument(uri)
-    }
+    let snapshot = try documentManager.latestSnapshot(uri)
 
     guard let offsetRange = snapshot.utf8OffsetRange(of: range) else {
       throw CursorInfoError.invalidRange(range)

Sources/SourceKitLSP/Swift/DocumentSymbols.swift

@@ -16,10 +16,7 @@ import LSPLogging
 
 extension SwiftLanguageServer {
   public func documentSymbol(_ req: DocumentSymbolRequest) async throws -> DocumentSymbolResponse? {
-    guard let snapshot = self.documentManager.latestSnapshot(req.textDocument.uri) else {
-      logger.error("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-      return nil
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
 

Sources/SourceKitLSP/Swift/OpenInterface.swift

@@ -29,7 +29,7 @@ extension SwiftLanguageServer {
     let interfaceFilePath = self.generatedInterfacesPath.appendingPathComponent("\(name).swiftinterface")
     let interfaceDocURI = DocumentURI(interfaceFilePath)
     // has interface already been generated
-    if let snapshot = self.documentManager.latestSnapshot(interfaceDocURI) {
+    if let snapshot = try? self.documentManager.latestSnapshot(interfaceDocURI) {
       return await self.interfaceDetails(request: request, uri: interfaceDocURI, snapshot: snapshot, symbol: symbol)
     } else {
       // generate interface

Sources/SourceKitLSP/Swift/SemanticRefactoring.swift

@@ -84,10 +84,6 @@ struct SemanticRefactoring {
 
 /// An error from a semantic refactoring request.
 enum SemanticRefactoringError: Error {
-
-  /// The given URL is not a known document.
-  case unknownDocument(DocumentURI)
-
   /// The given position range is invalid.
   case invalidRange(Range<Position>)
 
@@ -104,8 +100,6 @@ enum SemanticRefactoringError: Error {
 extension SemanticRefactoringError: CustomStringConvertible {
   var description: String {
     switch self {
-    case .unknownDocument(let url):
-      return "failed to find snapshot for url \(url)"
     case .invalidRange(let range):
       return "failed to refactor due to invalid range: \(range)"
     case .failedToRetrieveOffset(let range):
@@ -134,9 +128,7 @@ extension SwiftLanguageServer {
     let keys = self.keys
 
     let uri = refactorCommand.textDocument.uri
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
-      throw SemanticRefactoringError.unknownDocument(uri)
-    }
+    let snapshot = try self.documentManager.latestSnapshot(uri)
     guard let offsetRange = snapshot.utf8OffsetRange(of: refactorCommand.positionRange) else {
       throw SemanticRefactoringError.failedToRetrieveOffset(refactorCommand.positionRange)
     }

Sources/SourceKitLSP/Swift/SemanticTokens.swift

@@ -71,12 +71,7 @@ extension SwiftLanguageServer {
   public func documentSemanticTokens(
     _ req: DocumentSemanticTokensRequest
   ) async throws -> DocumentSemanticTokensResponse? {
-    let uri = req.textDocument.uri
-
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
-      logger.error("failed to find snapshot for uri \(uri.forLogging)")
-      return DocumentSemanticTokensResponse(data: [])
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
     let tokens = await mergedAndSortedTokens(for: snapshot)
     let encodedTokens = tokens.lspEncoded
@@ -93,15 +88,9 @@ extension SwiftLanguageServer {
   public func documentSemanticTokensRange(
     _ req: DocumentSemanticTokensRangeRequest
   ) async throws -> DocumentSemanticTokensResponse? {
-    let uri = req.textDocument.uri
-    let range = req.range
-
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
-      logger.error("failed to find snapshot for uri \(uri.forLogging)")
-      return DocumentSemanticTokensResponse(data: [])
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
-    let tokens = await mergedAndSortedTokens(for: snapshot, in: range)
+    let tokens = await mergedAndSortedTokens(for: snapshot, in: req.range)
     let encodedTokens = tokens.lspEncoded
 
     return DocumentSemanticTokensResponse(data: encodedTokens)

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -293,7 +293,7 @@ extension SwiftLanguageServer {
 
   public func documentUpdatedBuildSettings(_ uri: DocumentURI) async {
     // We may not have a snapshot if this is called just before `openDocument`.
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
+    guard let snapshot = try? self.documentManager.latestSnapshot(uri) else {
       return
     }
 
@@ -303,7 +303,7 @@ extension SwiftLanguageServer {
   }
 
   public func documentDependenciesUpdated(_ uri: DocumentURI) async {
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
+    guard let snapshot = try? self.documentManager.latestSnapshot(uri) else {
       return
     }
 
@@ -530,10 +530,7 @@ extension SwiftLanguageServer {
   }
 
   public func documentColor(_ req: DocumentColorRequest) async throws -> [ColorInformation] {
-    guard let snapshot = self.documentManager.latestSnapshot(req.textDocument.uri) else {
-      logger.error("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-      return []
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
 
@@ -605,12 +602,7 @@ extension SwiftLanguageServer {
   }
 
   public func documentSymbolHighlight(_ req: DocumentHighlightRequest) async throws -> [DocumentHighlight]? {
-    let keys = self.keys
-
-    guard let snapshot = self.documentManager.latestSnapshot(req.textDocument.uri) else {
-      logger.error("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-      return nil
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
     guard let offset = snapshot.utf8Offset(of: req.position) else {
       logger.error("invalid position \(req.position, privacy: .public)")
@@ -658,11 +650,7 @@ extension SwiftLanguageServer {
 
   public func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]? {
     let foldingRangeCapabilities = capabilityRegistry.clientCapabilities.textDocument?.foldingRange
-    let uri = req.textDocument.uri
-    guard let snapshot = self.documentManager.latestSnapshot(uri) else {
-      logger.error("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-      return nil
-    }
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
 
     let sourceFile = await syntaxTreeManager.syntaxTree(for: snapshot)
 
@@ -1059,9 +1047,7 @@ extension SwiftLanguageServer {
   private func fullDocumentDiagnosticReport(
     _ req: DocumentDiagnosticsRequest
   ) async throws -> RelatedFullDocumentDiagnosticReport {
-    guard let snapshot = documentManager.latestSnapshot(req.textDocument.uri) else {
-      throw ResponseError.unknown("failed to find snapshot for url \(req.textDocument.uri.forLogging)")
-    }
+    let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
     guard let buildSettings = await self.buildSettings(for: req.textDocument.uri), !buildSettings.isFallback else {
       logger.log(
         "Producing syntactic diagnostics from the built-in swift-syntax because we have fallback arguments"
@@ -1087,7 +1073,7 @@ extension SwiftLanguageServer {
     let dict = try await self.sourcekitd.send(skreq)
 
     try Task.checkCancellation()
-    guard documentManager.latestSnapshot(req.textDocument.uri)?.id == snapshot.id else {
+    guard (try? documentManager.latestSnapshot(req.textDocument.uri).id) == snapshot.id else {
       // Check that the document wasn't modified while we were getting diagnostics. This could happen because we are
       // calling `fullDocumentDiagnosticReport` from `publishDiagnosticsIfNeeded` outside of `messageHandlingQueue`
       // and thus a concurrent edit is possible while we are waiting for the sourcekitd request to return a result.

Sources/SourceKitLSP/Swift/VariableTypeInfo.swift

@@ -76,14 +76,6 @@ struct VariableTypeInfo {
   }
 }
 
-enum VariableTypeInfoError: Error, Equatable {
-  /// The given URL is not a known document.
-  case unknownDocument(DocumentURI)
-
-  /// The underlying sourcekitd request failed with the given error.
-  case responseError(ResponseError)
-}
-
 extension SwiftLanguageServer {
   /// Provides typed variable declarations in a document.
   ///
@@ -93,11 +85,7 @@ extension SwiftLanguageServer {
     _ uri: DocumentURI,
     _ range: Range<Position>? = nil
   ) async throws -> [VariableTypeInfo] {
-    guard let snapshot = documentManager.latestSnapshot(uri) else {
-      throw VariableTypeInfoError.unknownDocument(uri)
-    }
-
-    let keys = self.keys
+    let snapshot = try documentManager.latestSnapshot(uri)
 
     let skreq = SKDRequestDictionary(sourcekitd: sourcekitd)
     skreq[keys.request] = requests.variable_type

Tests/SourceKitDTests/CrashRecoveryTests.swift

@@ -100,10 +100,6 @@ final class CrashRecoveryTests: XCTestCase {
     try await fulfillmentOfOrThrow([sourcekitdCrashed], timeout: 5)
     try await fulfillmentOfOrThrow([sourcekitdRestarted], timeout: 30)
 
-    // Check that we have syntactic functionality again
-
-    _ = try await testClient.send(FoldingRangeRequest(textDocument: TextDocumentIdentifier(uri)))
-
     // sourcekitd's semantic request timer is only started when the first semantic request comes in.
     // Send a hover request (which will fail) to trigger that timer.
     // Afterwards wait for semantic functionality to be restored.

Tests/SourceKitLSPTests/BuildSystemTests.swift

@@ -143,7 +143,7 @@ final class BuildSystemTests: XCTestCase {
 
     let diags = try await testClient.nextDiagnosticsNotification()
     XCTAssertEqual(diags.diagnostics.count, 1)
-    XCTAssertEqual(text, documentManager.latestSnapshot(doc)!.text)
+    XCTAssertEqual(text, try documentManager.latestSnapshot(doc).text)
 
     // Modify the build settings and inform the delegate.
     // This should trigger a new publish diagnostics and we should no longer have errors.
@@ -155,7 +155,7 @@ final class BuildSystemTests: XCTestCase {
     var receivedCorrectDiagnostic = false
     for _ in 0..<Int(defaultTimeout) {
       let refreshedDiags = try await testClient.nextDiagnosticsNotification(timeout: 1)
-      if refreshedDiags.diagnostics.count == 0, text == documentManager.latestSnapshot(doc)!.text {
+      if refreshedDiags.diagnostics.count == 0, try text == documentManager.latestSnapshot(doc).text {
         receivedCorrectDiagnostic = true
         break
       }
@@ -182,7 +182,7 @@ final class BuildSystemTests: XCTestCase {
     testClient.openDocument(text, uri: doc)
     let diags1 = try await testClient.nextDiagnosticsNotification()
     XCTAssertEqual(diags1.diagnostics.count, 1)
-    XCTAssertEqual(text, documentManager.latestSnapshot(doc)!.text)
+    XCTAssertEqual(text, try documentManager.latestSnapshot(doc).text)
 
     // Modify the build settings and inform the delegate.
     // This should trigger a new publish diagnostics and we should no longer have errors.
@@ -218,7 +218,7 @@ final class BuildSystemTests: XCTestCase {
     let openDiags = try await testClient.nextDiagnosticsNotification()
     // Expect diagnostics to be withheld.
     XCTAssertEqual(openDiags.diagnostics.count, 0)
-    XCTAssertEqual(text, documentManager.latestSnapshot(doc)!.text)
+    XCTAssertEqual(text, try documentManager.latestSnapshot(doc).text)
 
     // Modify the build settings and inform the delegate.
     // This should trigger a new publish diagnostics and we should see a diagnostic.
@@ -229,7 +229,7 @@ final class BuildSystemTests: XCTestCase {
 
     let refreshedDiags = try await testClient.nextDiagnosticsNotification()
     XCTAssertEqual(refreshedDiags.diagnostics.count, 1)
-    XCTAssertEqual(text, documentManager.latestSnapshot(doc)!.text)
+    XCTAssertEqual(text, try documentManager.latestSnapshot(doc).text)
   }
 
   func testSwiftDocumentFallbackWithholdsSemanticDiagnostics() async throws {
@@ -254,7 +254,7 @@ final class BuildSystemTests: XCTestCase {
     testClient.openDocument(text, uri: doc)
     let openDiags = try await testClient.nextDiagnosticsNotification()
     XCTAssertEqual(openDiags.diagnostics.count, 1)
-    XCTAssertEqual(text, documentManager.latestSnapshot(doc)!.text)
+    XCTAssertEqual(text, try documentManager.latestSnapshot(doc).text)
 
     // Swap from fallback settings to primary build system settings.
     buildSystem.buildSettingsByFile[doc] = primarySettings