[SwiftLexicalLookup] Unqualified lookup caching

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,101 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 var ancestorResultsCache: [SyntaxIdentifier: [LookupResult]] = [:]
+  /// Cached results of `SequentialScopeSyntax.sequentialLookup` calls.
+  /// Identified by `SyntaxIdentifier`.
+  private var sequentialResultsCache: [SyntaxIdentifier: [LookupResult]] = [:]
+  /// Looked-up scope identifiers during cache accesses.
+  private var hits: Set<SyntaxIdentifier> = []
+
+  private let dropMod: Int
+  private var evictionCount = 0
+
+  /// Creates a new unqualified lookup cache.
+  /// `drop` parameter specifies how many eviction calls will be
+  /// ignored before evicting not-hit members of the cache.
+  ///
+  /// 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 init(drop: Int = 0) {
+    self.dropMod = drop + 1
+  }
+
+  /// 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 (as specified by `drop` in the initializer).
+  /// The dropping behavior can be disabled for this call with the `bypassDropCounter`
+  /// parameter, resulting in immediate eviction of entries without a hit.
+  public func evictEntriesWithoutHit(bypassDropCounter: Bool = false) {
+    if !bypassDropCounter {
+      evictionCount = (evictionCount + 1) % dropMod
+      guard evictionCount != 0 else { return }
+    }
+
+    for key in Set(ancestorResultsCache.keys).union(sequentialResultsCache.keys).subtracting(hits) {
+      ancestorResultsCache.removeValue(forKey: key)
+      sequentialResultsCache.removeValue(forKey: key)
+    }
+
+    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)
     }
   }
 }