add LRUCache struct to SKUtilities

Author: matthewbastien

Sources/DocCDocumentation/DocCCatalogIndexManager.swift

@@ -12,22 +12,29 @@
 
 package import Foundation
 import SKLogging
+import SKUtilities
 @_spi(LinkCompletion) @preconcurrency import SwiftDocC
 
 final actor DocCCatalogIndexManager {
   private let server: DocCServer
-  private var catalogToIndexMap: [URL: Result<DocCCatalogIndex, DocCIndexError>] = [:]
+
+  /// The cache of DocCCatalogIndex for a given SwiftDocC catalog URL
+  ///
+  /// - Note: The capacity has been chosen without scientific measurements. The
+  ///   feeling is that switching between SwiftDocC catalogs is rare and 5 catalog
+  ///   indexes won't take up much memory.
+  private var indexCache = LRUCache<URL, Result<DocCCatalogIndex, DocCIndexError>>(capacity: 5)
 
   init(server: DocCServer) {
     self.server = server
   }
 
   func invalidate(_ url: URL) {
-    catalogToIndexMap.removeValue(forKey: url)
+    indexCache.removeValue(forKey: url)
   }
 
   func index(for catalogURL: URL) async throws(DocCIndexError) -> DocCCatalogIndex {
-    if let existingCatalog = catalogToIndexMap[catalogURL] {
+    if let existingCatalog = indexCache[catalogURL] {
       return try existingCatalog.get()
     }
     do {
@@ -49,15 +56,15 @@ final actor DocCCatalogIndexManager {
       }
       let renderReferenceStore = try JSONDecoder().decode(RenderReferenceStore.self, from: renderReferenceStoreData)
       let catalogIndex = DocCCatalogIndex(from: renderReferenceStore)
-      catalogToIndexMap[catalogURL] = .success(catalogIndex)
+      indexCache[catalogURL] = .success(catalogIndex)
       return catalogIndex
     } catch {
       // Don't cache cancellation errors
       guard !(error is CancellationError) else {
         throw .cancelled
       }
       let internalError = error as? DocCIndexError ?? DocCIndexError.internalError(error)
-      catalogToIndexMap[catalogURL] = .failure(internalError)
+      indexCache[catalogURL] = .failure(internalError)
       throw internalError
     }
   }

Sources/SKUtilities/LRUCache.swift

@@ -0,0 +1,140 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 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
+//
+//===----------------------------------------------------------------------===//
+
+/// A cache that stores key-value pairs up to a given capacity.
+///
+/// The least recently used key-value pair is removed when the cache exceeds its capacity.
+package struct LRUCache<Key: Hashable, Value> {
+  private struct Priority {
+    var next: Key?
+    var previous: Key?
+
+    init(next: Key? = nil, previous: Key? = nil) {
+      self.next = next
+      self.previous = previous
+    }
+  }
+
+  // The hash map for accessing cached key-value pairs.
+  private var cache: [Key: Value]
+
+  // Doubly linked list of priorities keeping track of the first and last entries.
+  private var priorities: [Key: Priority]
+  private var firstPriority: Key? = nil
+  private var lastPriority: Key? = nil
+
+  /// The maximum number of key-value pairs that can be stored in the cache.
+  package let capacity: Int
+
+  /// The number of key-value pairs within the cache.
+  package var count: Int { cache.count }
+
+  /// A collection containing just the keys of the cache.
+  ///
+  /// - Note: Keys will **not** be in the same order that they were added to the cache.
+  package var keys: any Collection<Key> { cache.keys }
+
+  /// A collection containing just the values of the cache.
+  ///
+  /// - Note: Values will **not** be in the same order that they were added to the cache.
+  package var values: any Collection<Value> { cache.values }
+
+  package init(capacity: Int) {
+    assert(capacity > 0, "LRUCache capacity must be greater than 0")
+    self.capacity = capacity
+    self.cache = Dictionary(minimumCapacity: capacity)
+    self.priorities = Dictionary(minimumCapacity: capacity)
+  }
+
+  /// Adds the given key as the first priority in the doubly linked list of priorities.
+  private mutating func addPriority(forKey key: Key) {
+    // Make sure the key doesn't already exist in the list
+    removePriority(forKey: key)
+
+    guard let currentFirstPriority = firstPriority else {
+      firstPriority = key
+      lastPriority = key
+      priorities[key] = Priority()
+      return
+    }
+    priorities[key] = Priority(next: currentFirstPriority)
+    priorities[currentFirstPriority]?.previous = key
+    firstPriority = key
+  }
+
+  /// Removes the given key from the doubly linked list of priorities.
+  private mutating func removePriority(forKey key: Key) {
+    guard let priority = priorities.removeValue(forKey: key) else {
+      return
+    }
+    // Update the first and last priorities
+    if firstPriority == key {
+      firstPriority = priority.next
+    }
+    if lastPriority == key {
+      lastPriority = priority.previous
+    }
+    // Update the previous and next keys in the priority list
+    if let previousPriority = priority.previous {
+      priorities[previousPriority]?.next = priority.next
+    }
+    if let nextPriority = priority.next {
+      priorities[nextPriority]?.previous = priority.previous
+    }
+  }
+
+  /// Removes all key-value pairs from the cache.
+  package mutating func removeAll() {
+    cache.removeAll()
+    priorities.removeAll()
+    firstPriority = nil
+    lastPriority = nil
+  }
+
+  /// Removes all the elements that satisfy the given predicate.
+  package mutating func removeAll(where shouldBeRemoved: (_: ((key: Key, value: Value)) throws -> Bool)) rethrows {
+    cache = try cache.filter { entry in
+      guard try shouldBeRemoved(entry) else {
+        return true
+      }
+      removePriority(forKey: entry.key)
+      return false
+    }
+  }
+
+  /// Removes the given key and its associated value from the cache.
+  ///
+  /// Returns the value that was associated with the key.
+  @discardableResult
+  package mutating func removeValue(forKey key: Key) -> Value? {
+    removePriority(forKey: key)
+    return cache.removeValue(forKey: key)
+  }
+
+  package subscript(key: Key) -> Value? {
+    mutating _read {
+      addPriority(forKey: key)
+      yield cache[key]
+    }
+    set {
+      guard let newValue else {
+        removeValue(forKey: key)
+        return
+      }
+      cache[key] = newValue
+      addPriority(forKey: key)
+      if cache.count > capacity, let lastPriority {
+        removeValue(forKey: lastPriority)
+      }
+    }
+  }
+}

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -14,6 +14,7 @@ import LanguageServerProtocol
 import LanguageServerProtocolExtensions
 import SKLogging
 import SKOptions
+import SKUtilities
 import SourceKitD
 import SwiftDiagnostics
 import SwiftExtensions
@@ -25,6 +26,11 @@ actor DiagnosticReportManager {
     (report: RelatedFullDocumentDiagnosticReport, cachable: Bool)
   >
 
+  private struct CacheKey: Hashable {
+    let snapshotID: DocumentSnapshot.ID
+    let buildSettings: SwiftCompileCommand?
+  }
+
   private let sourcekitd: SourceKitD
   private let options: SourceKitLSPOptions
   private let syntaxTreeManager: SyntaxTreeManager
@@ -36,20 +42,8 @@ actor DiagnosticReportManager {
 
   /// The cache that stores reportTasks for snapshot id and buildSettings
   ///
-  /// 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 reportTaskCache:
-    [(
-      snapshotID: DocumentSnapshot.ID,
-      buildSettings: SwiftCompileCommand?,
-      reportTask: ReportTask
-    )] = []
-
-  /// The number of reportTasks to keep
-  ///
-  /// - Note: This has been chosen without scientific measurements.
-  private let cacheSize = 5
+  /// - Note: The capacity has been chosen without scientific measurements.
+  private var reportTaskCache = LRUCache<CacheKey, ReportTask>(capacity: 5)
 
   init(
     sourcekitd: SourceKitD,
@@ -101,7 +95,7 @@ actor DiagnosticReportManager {
   }
 
   func removeItemsFromCache(with uri: DocumentURI) async {
-    reportTaskCache.removeAll(where: { $0.snapshotID.uri == uri })
+    reportTaskCache.removeAll(where: { $0.key.snapshotID.uri == uri })
   }
 
   private func requestReport(
@@ -188,28 +182,18 @@ actor DiagnosticReportManager {
     for snapshotID: DocumentSnapshot.ID,
     buildSettings: SwiftCompileCommand?
   ) -> ReportTask? {
-    return reportTaskCache.first(where: { $0.snapshotID == snapshotID && $0.buildSettings == buildSettings })?
-      .reportTask
+    let key = CacheKey(snapshotID: snapshotID, buildSettings: buildSettings)
+    return reportTaskCache[key]
   }
 
   /// Set the reportTask for the given document snapshot and buildSettings.
-  ///
-  /// If we are already storing `cacheSize` many reports, the oldest one
-  /// will get discarded.
   private func setReportTask(
     for snapshotID: DocumentSnapshot.ID,
     buildSettings: SwiftCompileCommand?,
     reportTask: ReportTask
   ) {
     // Remove any reportTasks for old versions of this document.
-    reportTaskCache.removeAll(where: { $0.snapshotID <= snapshotID })
-
-    reportTaskCache.insert((snapshotID, buildSettings, reportTask), at: 0)
-
-    // If we still have more than `cacheSize` reportTasks, delete the ones that
-    // were produced last. We can always re-request them on-demand.
-    while reportTaskCache.count > cacheSize {
-      reportTaskCache.removeLast()
-    }
+    reportTaskCache.removeAll(where: { $0.key.snapshotID <= snapshotID })
+    reportTaskCache[CacheKey(snapshotID: snapshotID, buildSettings: buildSettings)] = reportTask
   }
 }

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -22,26 +22,10 @@ import SwiftExtensions
 
 /// Caches the contents of macro expansions that were recently requested by the user.
 actor MacroExpansionManager {
-  private struct CacheEntry {
-    // Key
+  private struct CacheKey: Hashable {
     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?) {
@@ -50,19 +34,12 @@ actor MacroExpansionManager {
 
   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] = []
+  /// - Note: The capacity of this cache 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 var cache = LRUCache<CacheKey, [RefactoringEdit]>(capacity: 10)
 
   /// Return the text of the macro expansion referenced by `macroExpansionURLData`.
   func macroExpansion(
@@ -90,20 +67,12 @@ actor MacroExpansionManager {
     let snapshot = try await swiftLanguageService.latestSnapshot(for: uri)
     let compileCommand = await swiftLanguageService.compileCommand(for: uri, fallbackAfterTimeout: false)
 
-    if let cacheEntry = cache.first(where: {
-      $0.snapshotID == snapshot.id && $0.range == range && $0.buildSettings == compileCommand
-    }) {
-      return cacheEntry.value
+    let cacheKey = CacheKey(snapshotID: snapshot.id, range: range, buildSettings: compileCommand)
+    if let valueFromCache = cache[cacheKey] {
+      return valueFromCache
     }
     let macroExpansions = try await macroExpansionsImpl(in: snapshot, at: range, buildSettings: compileCommand)
-    cache.insert(
-      CacheEntry(snapshot: snapshot, range: range, buildSettings: compileCommand, value: macroExpansions),
-      at: 0
-    )
-
-    while cache.count > cacheSize {
-      cache.removeLast()
-    }
+    cache[cacheKey] = macroExpansions
 
     return macroExpansions
   }
@@ -151,7 +120,9 @@ actor MacroExpansionManager {
 
   /// 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 }
+    cache.removeAll {
+      $0.key.snapshotID.uri.primaryFile ?? $0.key.snapshotID.uri == primaryFile
+    }
   }
 }
 

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -73,7 +73,7 @@ fileprivate func diagnosticsEnabled(for document: DocumentURI) -> Bool {
 }
 
 /// A swift compiler command derived from a `FileBuildSettingsChange`.
-package struct SwiftCompileCommand: Sendable, Equatable {
+package struct SwiftCompileCommand: Sendable, Equatable, Hashable {
 
   /// The compiler arguments, including working directory. This is required since sourcekitd only
   /// accepts the working directory via the compiler arguments.