[SwiftLexicalLookup] Unqualified lookup caching

CMakeLists.txt

@@ -57,6 +57,10 @@ if(NOT DEFINED SWIFTSYNTAX_PACKAGE_NAME)
   set(SWIFTSYNTAX_PACKAGE_NAME "${SWIFT_MODULE_ABI_NAME_PREFIX}${PROJECT_NAME}")
 endif()
 
+add_compile_options(
+  "$<$<COMPILE_LANGUAGE:Swift>:SHELL:-Xfrontend -package-name -Xfrontend ${SWIFTSYNTAX_PACKAGE_NAME}>"
+)
+
 # Determine the module triple.
 if("${SWIFT_HOST_MODULE_TRIPLE}" STREQUAL "")
   set(module_triple_command "${CMAKE_Swift_COMPILER}" -print-target-info)

Sources/SwiftCompilerPluginMessageHandling/CMakeLists.txt

@@ -9,7 +9,6 @@
 add_swift_syntax_library(SwiftCompilerPluginMessageHandling
   CompilerPluginMessageHandler.swift
   Diagnostics.swift
-  LRUCache.swift
   Macros.swift
   PluginMacroExpansionContext.swift
   PluginMessageCompatibility.swift

Sources/SwiftLexicalLookup/CMakeLists.txt

@@ -8,6 +8,7 @@
 
 add_swift_syntax_library(SwiftLexicalLookup
   IdentifiableSyntax.swift
+  LookupCache.swift
   LookupName.swift
   LookupResult.swift
   SimpleLookupQueries.swift

Sources/SwiftLexicalLookup/LookupCache.swift

@@ -0,0 +1,109 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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 SwiftSyntax
+
+/// Unqualified lookup cache. Should be used when performing
+/// large sequences of adjacent lookups to maximise performance.
+public class LookupCache {
+  /// Cached results of `ScopeSyntax.lookupParent` calls.
+  /// Identified by `SyntaxIdentifier`.
+  private let ancestorResultsCache: LRUCache<SyntaxIdentifier, [LookupResult]>
+  /// Cached results of `SequentialScopeSyntax.sequentialLookup` calls.
+  /// Identified by `SyntaxIdentifier`.
+  private let sequentialResultsCache: LRUCache<SyntaxIdentifier, [LookupResult]>
+  /// Looked-up scope identifiers during cache accesses.
+  private var hits: Set<SyntaxIdentifier> = [] {
+    didSet {
+      if hits.count > capacity * 2 {
+        hits.removeAll()
+      }
+    }
+  }
+
+  private let capacity: Int
+  private var evictionCount = 0
+
+  /// Creates a new unqualified lookup cache.
+  /// `capacity` describes the maximum amount of entries in the cache.
+  /// The cache size is maintained according to the LRU (Least Recently Used) policy.
+  public init(capacity: Int) {
+    self.capacity = capacity
+    self.ancestorResultsCache = LRUCache(capacity: (capacity + 1) / 2)
+    self.sequentialResultsCache = LRUCache(capacity: capacity / 2)
+  }
+
+  /// Get cached ancestor results for the given `id`.
+  /// `nil` if there's no cache entry for the given `id`.
+  /// Adds `id` and ids of all ancestors to the cache `hits`.
+  func getCachedAncestorResults(id: SyntaxIdentifier) -> [LookupResult]? {
+    guard let results = ancestorResultsCache[id] else { return nil }
+    hits.formUnion(results.map(\.scope.id))
+    hits.insert(id)
+    return results
+  }
+
+  /// Set cached ancestor results for the given `id`.
+  /// Adds `id` to the cache `hits`.
+  func setCachedAncestorResults(id: SyntaxIdentifier, results: [LookupResult]) {
+    hits.insert(id)
+    ancestorResultsCache[id] = results
+  }
+
+  /// Get cached sequential lookup results for the given `id`.
+  /// `nil` if there's no cache entry for the given `id`.
+  /// Adds `id` to the cache `hits`.
+  func getCachedSequentialResults(id: SyntaxIdentifier) -> [LookupResult]? {
+    guard let results = sequentialResultsCache[id] else { return nil }
+    hits.insert(id)
+    return results
+  }
+
+  /// Set cached sequential lookup results for the given `id`.
+  /// Adds `id` to the cache `hits`.
+  func setCachedSequentialResults(id: SyntaxIdentifier, results: [LookupResult]) {
+    hits.insert(id)
+    sequentialResultsCache[id] = results
+  }
+
+  /// Removes all cached entries without a hit, unless it's prohibited
+  /// by the internal `drop` counter.
+  /// `drop` parameter specifies how many eviction calls will be
+  /// ignored before evicting not-hit members from subsequent lookups.
+  ///
+  /// Example cache eviction sequences (s - skip, e - evict):
+  /// - `drop = 0` - `e -> e -> e -> e -> e -> ...`
+  /// - `drop = 1` - `s -> e -> s -> s -> e -> ...`
+  /// - `drop = 3` - `s -> s -> s -> e -> s -> ...`
+  ///
+  /// - Note: `drop = 0` effectively maintains exactly one path of cached results to
+  /// the root in the cache (assuming we evict cache members after each lookup in a sequence of lookups).
+  /// Higher the `drop` value, more such paths can potentially be stored in the cache at any given moment.
+  /// Because of that, a higher `drop` value also translates to a higher number of cache-hits,
+  /// but it might not directly translate to better performance. Because of a larger memory footprint,
+  /// memory accesses could take longer, slowing down the eviction process. That's why the `drop` value
+  /// could be fine-tuned to maximize the performance given file size,
+  /// number of lookups, and amount of available memory.
+  public func evictEntriesWithoutHit(drop: Int = 0) {
+    evictionCount = (evictionCount + 1) % (drop + 1)
+    guard evictionCount != 0 else { return }
+
+    for key in ancestorResultsCache.keys where !hits.contains(key) {
+      ancestorResultsCache[key] = nil
+    }
+    for key in sequentialResultsCache.keys where !hits.contains(key) {
+      sequentialResultsCache[key] = nil
+    }
+
+    hits = []
+  }
+}

Sources/SwiftLexicalLookup/Scopes/CanInterleaveResultsLaterScopeSyntax.swift

@@ -20,6 +20,7 @@ protocol CanInterleaveResultsLaterScopeSyntax: ScopeSyntax {
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
     with config: LookupConfig,
+    cache: LookupCache?,
     resultsToInterleave: [LookupResult]
   ) -> [LookupResult]
 }

Sources/SwiftLexicalLookup/Scopes/FunctionScopeSyntax.swift

@@ -30,7 +30,8 @@ extension FunctionScopeSyntax {
   @_spi(Experimental) public func lookup(
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
-    with config: LookupConfig
+    with config: LookupConfig,
+    cache: LookupCache?
   ) -> [LookupResult] {
     var thisScopeResults: [LookupResult] = []
 
@@ -39,6 +40,7 @@ extension FunctionScopeSyntax {
         identifier,
         at: position,
         with: config,
+        cache: cache,
         propagateToParent: false
       )
     }
@@ -47,7 +49,8 @@ extension FunctionScopeSyntax {
       + lookupThroughGenericParameterScope(
         identifier,
         at: lookUpPosition,
-        with: config
+        with: config,
+        cache: cache
       )
   }
 }

Sources/SwiftLexicalLookup/Scopes/GenericParameterScopeSyntax.swift

@@ -40,18 +40,21 @@ protocol GenericParameterScopeSyntax: ScopeSyntax {}
   @_spi(Experimental) public func lookup(
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
-    with config: LookupConfig
+    with config: LookupConfig,
+    cache: LookupCache?
   ) -> [LookupResult] {
     return defaultLookupImplementation(
       identifier,
       at: lookUpPosition,
       with: config,
+      cache: cache,
       propagateToParent: false
     )
       + lookupBypassingParentResults(
         identifier,
         at: lookUpPosition,
-        with: config
+        with: config,
+        cache: cache
       )
   }
 
@@ -76,16 +79,22 @@ protocol GenericParameterScopeSyntax: ScopeSyntax {}
   private func lookupBypassingParentResults(
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
-    with config: LookupConfig
+    with config: LookupConfig,
+    cache: LookupCache?
   ) -> [LookupResult] {
     guard let parentScope else { return [] }
 
     if let parentScope = Syntax(parentScope).asProtocol(SyntaxProtocol.self)
       as? WithGenericParametersScopeSyntax
     {
-      return parentScope.returningLookupFromGenericParameterScope(identifier, at: lookUpPosition, with: config)
+      return parentScope.returningLookupFromGenericParameterScope(
+        identifier,
+        at: lookUpPosition,
+        with: config,
+        cache: cache
+      )
     } else {
-      return lookupInParent(identifier, at: lookUpPosition, with: config)
+      return lookupInParent(identifier, at: lookUpPosition, with: config, cache: cache)
     }
   }
 }

Sources/SwiftLexicalLookup/Scopes/IntroducingToSequentialParentScopeSyntax.swift

@@ -21,6 +21,7 @@ protocol IntroducingToSequentialParentScopeSyntax: ScopeSyntax {
   func lookupFromSequentialParent(
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
-    with config: LookupConfig
+    with config: LookupConfig,
+    cache: LookupCache?
   ) -> [LookupResult]
 }

Sources/SwiftLexicalLookup/Scopes/NominalTypeDeclSyntax.swift

@@ -34,16 +34,18 @@ extension NominalTypeDeclSyntax {
   @_spi(Experimental) public func returningLookupFromGenericParameterScope(
     _ identifier: Identifier?,
     at lookUpPosition: AbsolutePosition,
-    with config: LookupConfig
+    with config: LookupConfig,
+    cache: LookupCache?
   ) -> [LookupResult] {
     if let inheritanceClause, inheritanceClause.range.contains(lookUpPosition) {
-      return lookupInParent(identifier, at: lookUpPosition, with: config)
+      return lookupInParent(identifier, at: lookUpPosition, with: config, cache: cache)
     } else if let genericParameterClause, genericParameterClause.range.contains(lookUpPosition) {
-      return lookupInParent(identifier, at: lookUpPosition, with: config)
+      return lookupInParent(identifier, at: lookUpPosition, with: config, cache: cache)
     } else if name.range.contains(lookUpPosition) || genericWhereClause?.range.contains(lookUpPosition) ?? false {
-      return lookupInParent(identifier, at: lookUpPosition, with: config)
+      return lookupInParent(identifier, at: lookUpPosition, with: config, cache: cache)
     } else {
-      return [.lookForMembers(in: Syntax(self))] + lookupInParent(identifier, at: lookUpPosition, with: config)
+      return [.lookForMembers(in: Syntax(self))]
+        + lookupInParent(identifier, at: lookUpPosition, with: config, cache: cache)
     }
   }
 }