Add DiagnosticReportManager to manage the cache for diagnostis

Author: joehsieh

Sources/SourceKitLSP/CMakeLists.txt

@@ -20,6 +20,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/CommentXML.swift
   Swift/CursorInfo.swift
   Swift/Diagnostic.swift
+  Swift/DiagnosticReportManager.swift
   Swift/DocumentSymbols.swift
   Swift/EditorPlaceholder.swift
   Swift/OpenInterface.swift

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -0,0 +1,98 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LSPLogging
+import LanguageServerProtocol
+
+protocol DiagnosticReportManagerDelegate: AnyObject {
+  /// Notify the delegate that compute a report with a given snapshot
+  func reportWithSnapshotReceived(_ snapshot: DocumentSnapshot, compilerArgs: [String]) async throws
+    -> RelatedFullDocumentDiagnosticReport
+
+  /// Notify the delegate that compute a report with a given snapshot for fallback
+  func fallbackReportWithSnapshotReceived(_ snapshot: DocumentSnapshot) async throws
+    -> RelatedFullDocumentDiagnosticReport
+}
+
+actor DiagnosticReportManager {
+  /// The cache that stores reports for snapshot ids
+  ///
+  /// 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 reportCache:
+    [(
+      snapshotID: DocumentSnapshot.ID,
+      report: RelatedFullDocumentDiagnosticReport
+    )] = []
+
+  /// The number of reports to keep
+  ///
+  /// - Note: This has been chosen without scientific measurements.
+  private let cacheSize = 5
+
+  /// Delegate to handle any DiagnosticReportManager events.
+  private weak var delegate: DiagnosticReportManagerDelegate?
+
+  func setDelegate(_ delegate: DiagnosticReportManagerDelegate?) async {
+    self.delegate = delegate
+  }
+
+  func diagnosticReport(for snapshot: DocumentSnapshot, buildSettings: SwiftCompileCommand?, useCache: Bool = false)
+    async throws -> RelatedFullDocumentDiagnosticReport
+  {
+    guard let delegate = self.delegate else {
+      throw ResponseError.unknown("DiagnosticReportManagerDelegate should not be nil")
+    }
+    if useCache, let report = report(for: snapshot.id) {
+      return report
+    }
+    guard let buildSettings = buildSettings, !buildSettings.isFallback else {
+      logger.log(
+        "Producing syntactic diagnostics from the built-in swift-syntax because we have fallback arguments"
+      )
+      // If we don't have build settings or we only have fallback build settings,
+      // sourcekitd won't be able to give us accurate semantic diagnostics.
+      // Fall back to providing syntactic diagnostics from the built-in
+      // swift-syntax. That's the best we can do for now.
+      let report = try await delegate.fallbackReportWithSnapshotReceived(snapshot)
+      setReport(for: snapshot.id, report: report)
+      return report
+    }
+
+    let report = try await delegate.reportWithSnapshotReceived(snapshot, compilerArgs: buildSettings.compilerArgs)
+    setReport(for: snapshot.id, report: report)
+    return report
+  }
+
+  /// The report for the given document snapshot.
+  private func report(for snapshotID: DocumentSnapshot.ID) -> RelatedFullDocumentDiagnosticReport? {
+    return reportCache.first(where: { $0.snapshotID == snapshotID })?.report
+  }
+
+  /// Set the report for the given document snapshot.
+  ///
+  /// If we are already storing `cacheSize` many reports, the oldest one
+  /// will get discarded.
+  private func setReport(for snapshotID: DocumentSnapshot.ID, report: RelatedFullDocumentDiagnosticReport) {
+    reportCache.insert((snapshotID, report), at: 0)
+
+    // Remove any reports for old versions of this document.
+    reportCache.removeAll(where: { $0.snapshotID < snapshotID })
+
+    // If we still have more than `cacheSize` reports, delete the ones that
+    // were produced last. We can always re-request them on-demand.
+    while reportCache.count > cacheSize {
+      reportCache.removeLast()
+    }
+  }
+}

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -139,6 +139,8 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
 
   private var stateChangeHandlers: [(_ oldState: LanguageServerState, _ newState: LanguageServerState) -> Void] = []
 
+  private let diagnosticReportManager: DiagnosticReportManager
+
   /// Creates a language server for the given client using the sourcekitd dylib specified in `toolchain`.
   /// `reopenDocuments` is a closure that will be called if sourcekitd crashes and the `SwiftLanguageServer` asks its parent server to reopen all of its documents.
   /// Returns `nil` if `sourcektid` couldn't be found.
@@ -157,6 +159,7 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
     self.state = .connected
     self.generatedInterfacesPath = options.generatedInterfacesPath.asURL
     try FileManager.default.createDirectory(at: generatedInterfacesPath, withIntermediateDirectories: true)
+    self.diagnosticReportManager = DiagnosticReportManager()
   }
 
   /// - Important: For testing only
@@ -196,7 +199,8 @@ public actor SwiftLanguageServer: ToolchainLanguageServer {
 
 extension SwiftLanguageServer {
 
-  public func initialize(_ initialize: InitializeRequest) throws -> InitializeResult {
+  public func initialize(_ initialize: InitializeRequest) async throws -> InitializeResult {
+    await diagnosticReportManager.setDelegate(self)
     sourcekitd.addNotificationHandler(self)
 
     return InitializeResult(
@@ -396,8 +400,11 @@ extension SwiftLanguageServer {
         return
       }
       do {
-        let diagnosticReport = try await self.fullDocumentDiagnosticReport(
-          DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(document))
+        let snapshot = try await documentManager.latestSnapshot(document)
+        let buildSettings = await self.buildSettings(for: document)
+        let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
+          for: snapshot,
+          buildSettings: buildSettings
         )
 
         await sourceKitServer.sendNotificationToClient(
@@ -967,10 +974,12 @@ extension SwiftLanguageServer {
   }
 
   func retrieveQuickFixCodeActions(_ params: CodeActionRequest) async throws -> [CodeAction] {
-    let diagnosticReport = try await self.fullDocumentDiagnosticReport(
-      DocumentDiagnosticsRequest(
-        textDocument: params.textDocument
-      )
+    let snapshot = try documentManager.latestSnapshot(params.textDocument.uri)
+    let buildSettings = await self.buildSettings(for: params.textDocument.uri)
+    let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
+      for: snapshot,
+      buildSettings: buildSettings,
+      useCache: true
     )
 
     let codeActions = diagnosticReport.items.flatMap { (diag) -> [CodeAction] in
@@ -1061,55 +1070,13 @@ extension SwiftLanguageServer {
   }
 
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {
-    return try await .full(fullDocumentDiagnosticReport(req))
-  }
-
-  private func fullDocumentDiagnosticReport(
-    _ req: DocumentDiagnosticsRequest
-  ) async throws -> RelatedFullDocumentDiagnosticReport {
     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"
-      )
-      // If we don't have build settings or we only have fallback build settings,
-      // sourcekitd won't be able to give us accurate semantic diagnostics.
-      // Fall back to providing syntactic diagnostics from the built-in
-      // swift-syntax. That's the best we can do for now.
-      return try await syntacticDiagnosticFromBuiltInSwiftSyntax(for: snapshot)
-    }
-
-    try Task.checkCancellation()
-
-    let keys = self.keys
-
-    let skreq = SKDRequestDictionary(sourcekitd: self.sourcekitd)
-    skreq[keys.request] = requests.diagnostics
-    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
-
-    // FIXME: SourceKit should probably cache this for us.
-    skreq[keys.compilerargs] = buildSettings.compilerArgs
-
-    let dict = try await self.sourcekitd.send(skreq, fileContents: snapshot.text)
-
-    try Task.checkCancellation()
-    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.
-      throw ResponseError.unknown("Document was modified while loading document")
-    }
-
-    let supportsCodeDescription = capabilityRegistry.clientHasDiagnosticsCodeDescriptionSupport
-    var diagnostics: [Diagnostic] = []
-    dict[keys.diagnostics]?.forEach { _, diag in
-      if let diag = Diagnostic(diag, in: snapshot, useEducationalNoteAsCode: supportsCodeDescription) {
-        diagnostics.append(diag)
-      }
-      return true
-    }
-
-    return RelatedFullDocumentDiagnosticReport(items: diagnostics)
+    let buildSettings = await self.buildSettings(for: req.textDocument.uri)
+    let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(
+      for: snapshot,
+      buildSettings: buildSettings
+    )
+    return .full(diagnosticReport)
   }
 
   public func executeCommand(_ req: ExecuteCommandRequest) async throws -> LSPAny? {
@@ -1331,3 +1298,51 @@ extension sourcekitd_uid_t {
     }
   }
 }
+
+extension SwiftLanguageServer: DiagnosticReportManagerDelegate {
+  func reportWithSnapshotReceived(_ snapshot: DocumentSnapshot, compilerArgs: [String]) async throws
+    -> LanguageServerProtocol.RelatedFullDocumentDiagnosticReport
+  {
+    try Task.checkCancellation()
+
+    let keys = self.keys
+
+    let skreq = SKDRequestDictionary(sourcekitd: self.sourcekitd)
+    skreq[keys.request] = requests.diagnostics
+    skreq[keys.sourcefile] = snapshot.uri.pseudoPath
+
+    // FIXME: SourceKit should probably cache this for us.
+    skreq[keys.compilerargs] = compilerArgs
+
+    let dict = try await self.sourcekitd.send(skreq, fileContents: snapshot.text)
+
+    try Task.checkCancellation()
+    guard (try? documentManager.latestSnapshot(snapshot.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.
+      throw ResponseError.unknown("Document was modified while loading document")
+    }
+
+    let supportsCodeDescription = capabilityRegistry.clientHasDiagnosticsCodeDescriptionSupport
+    var diagnostics: [Diagnostic] = []
+    dict[keys.diagnostics]?.forEach { _, diag in
+      if let diag = Diagnostic(diag, in: snapshot, useEducationalNoteAsCode: supportsCodeDescription) {
+        diagnostics.append(diag)
+      }
+      return true
+    }
+
+    return RelatedFullDocumentDiagnosticReport(items: diagnostics)
+  }
+
+  func fallbackReportWithSnapshotReceived(_ snapshot: DocumentSnapshot) async throws
+    -> LanguageServerProtocol.RelatedFullDocumentDiagnosticReport
+  {
+    // If we don't have build settings or we only have fallback build settings,
+    // sourcekitd won't be able to give us accurate semantic diagnostics.
+    // Fall back to providing syntactic diagnostics from the built-in
+    // swift-syntax. That's the best we can do for now.
+    return try await syntacticDiagnosticFromBuiltInSwiftSyntax(for: snapshot)
+  }
+}