Cache container names in CheckedIndex

Sources/SemanticIndex/CheckedIndex.swift

@@ -60,6 +60,32 @@ package final class CheckedIndex {
   private var checker: IndexOutOfDateChecker
   private let index: IndexStoreDB
 
+  /// Maps the USR of a symbol to its name and the name of all its containers, from outermost to innermost.
+  ///
+  /// It is important that we cache this because we might find a lot of symbols in the same container for eg. workspace
+  /// symbols (eg. consider many symbols in the same C++ namespace). If we didn't cache this value, then we would need
+  /// to perform a `primaryDefinitionOrDeclarationOccurrence` lookup for all of these containers, which is expensive.
+  ///
+  /// Since we don't expect `CheckedIndex` to be outlive a single request it is acceptable to cache these results
+  /// without having any invalidation logic (similar to how we don't invalide results cached in
+  /// `IndexOutOfDateChecker`).
+  ///
+  /// ### Examples
+  /// If we have
+  /// ```swift
+  /// struct Foo {}
+  /// ``` then
+  /// `containerNamesCache[<usr of Foo>]` will be `["Foo"]`.
+  ///
+  /// If we have
+  /// ```swift
+  /// struct Bar {
+  ///   struct Foo {}
+  /// }
+  /// ```, then
+  /// `containerNamesCache[<usr of Foo>]` will be `["Bar", "Foo"]`.
+  private var containerNamesCache: [String: [String]] = [:]
+
   fileprivate init(index: IndexStoreDB, checkLevel: IndexCheckLevel) {
     self.index = index
     self.checker = IndexOutOfDateChecker(checkLevel: checkLevel)
@@ -183,6 +209,84 @@ package final class CheckedIndex {
     }
     return result
   }
+
+  /// The names of all containers the symbol is contained in, from outermost to innermost.
+  ///
+  /// ### Examples
+  /// In the following, the container names of `test` are `["Foo"]`.
+  /// ```swift
+  /// struct Foo {
+  ///   func test() {}
+  /// }
+  /// ```
+  ///
+  /// In the following, the container names of `test` are `["Bar", "Foo"]`.
+  /// ```swift
+  /// struct Bar {
+  ///   struct Foo {
+  ///     func test() {}
+  ///   }
+  /// }
+  /// ```
+  package func containerNames(of symbol: SymbolOccurrence) -> [String] {
+    // The container name of accessors is the container of the surrounding variable.
+    let accessorOf = symbol.relations.filter { $0.roles.contains(.accessorOf) }
+    if let primaryVariable = accessorOf.sorted().first {
+      if accessorOf.count > 1 {
+        logger.fault("Expected an occurrence to an accessor of at most one symbol, not multiple")
+      }
+      if let primaryVariable = primaryDefinitionOrDeclarationOccurrence(ofUSR: primaryVariable.symbol.usr) {
+        return containerNames(of: primaryVariable)
+      }
+    }
+
+    let containers = symbol.relations.filter { $0.roles.contains(.childOf) }
+    if containers.count > 1 {
+      logger.fault("Expected an occurrence to a child of at most one symbol, not multiple")
+    }
+    let container = containers.filter {
+      switch $0.symbol.kind {
+      case .module, .namespace, .enum, .struct, .class, .protocol, .extension, .union:
+        return true
+      case .unknown, .namespaceAlias, .macro, .typealias, .function, .variable, .field, .enumConstant,
+        .instanceMethod, .classMethod, .staticMethod, .instanceProperty, .classProperty, .staticProperty, .constructor,
+        .destructor, .conversionFunction, .parameter, .using, .concept, .commentTag:
+        return false
+      }
+    }.sorted().first
+
+    guard var containerSymbol = container?.symbol else {
+      return []
+    }
+    if let cached = containerNamesCache[containerSymbol.usr] {
+      return cached
+    }
+
+    if containerSymbol.kind == .extension,
+      let extendedSymbol = self.occurrences(relatedToUSR: containerSymbol.usr, roles: .extendedBy).first?.symbol
+    {
+      containerSymbol = extendedSymbol
+    }
+    let result: [String]
+
+    // Use `forEachSymbolOccurrence` instead of `primaryDefinitionOrDeclarationOccurrence` to get a symbol occurrence
+    // for the container because it can be significantly faster: Eg. when searching for a C++ namespace (such as `llvm`),
+    // it may be declared in many files. Finding the canonical definition means that we would need to scan through all
+    // of these files. But we expect all all of these declarations to have the same parent container names and we don't
+    // care about locations here.
+    var containerDefinition: SymbolOccurrence?
+    forEachSymbolOccurrence(byUSR: containerSymbol.usr, roles: [.definition, .declaration]) { occurrence in
+      containerDefinition = occurrence
+      return false  // stop iteration
+    }
+    if let containerDefinition {
+      result = self.containerNames(of: containerDefinition) + [containerSymbol.name]
+    } else {
+      result = [containerSymbol.name]
+    }
+    containerNamesCache[containerSymbol.usr] = result
+    return result
+  }
 }
 
 /// A wrapper around `IndexStoreDB` that allows the retrieval of a `CheckedIndex` with a specified check level or the

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -2395,48 +2395,11 @@ private let maxWorkspaceSymbolResults = 4096
 package typealias Diagnostic = LanguageServerProtocol.Diagnostic
 
 fileprivate extension CheckedIndex {
-  func containerNames(of symbol: SymbolOccurrence) -> [String] {
-    // The container name of accessors is the container of the surrounding variable.
-    let accessorOf = symbol.relations.filter { $0.roles.contains(.accessorOf) }
-    if let primaryVariable = accessorOf.sorted().first {
-      if accessorOf.count > 1 {
-        logger.fault("Expected an occurrence to an accessor of at most one symbol, not multiple")
-      }
-      if let primaryVariable = primaryDefinitionOrDeclarationOccurrence(ofUSR: primaryVariable.symbol.usr) {
-        return containerNames(of: primaryVariable)
-      }
-    }
-
-    let containers = symbol.relations.filter { $0.roles.contains(.childOf) }
-    if containers.count > 1 {
-      logger.fault("Expected an occurrence to a child of at most one symbol, not multiple")
-    }
-    let container = containers.filter {
-      switch $0.symbol.kind {
-      case .module, .namespace, .enum, .struct, .class, .protocol, .extension, .union:
-        return true
-      case .unknown, .namespaceAlias, .macro, .typealias, .function, .variable, .field, .enumConstant,
-        .instanceMethod, .classMethod, .staticMethod, .instanceProperty, .classProperty, .staticProperty, .constructor,
-        .destructor, .conversionFunction, .parameter, .using, .concept, .commentTag:
-        return false
-      }
-    }.sorted().first
-
-    if let container {
-      if let containerDefinition = primaryDefinitionOrDeclarationOccurrence(ofUSR: container.symbol.usr) {
-        return self.containerNames(of: containerDefinition) + [container.symbol.name]
-      }
-      return [container.symbol.name]
-    } else {
-      return []
-    }
-  }
-
   /// Take the name of containers into account to form a fully-qualified name for the given symbol.
   /// This means that we will form names of nested types and type-qualify methods.
   func fullyQualifiedName(of symbolOccurrence: SymbolOccurrence) -> String {
     let symbol = symbolOccurrence.symbol
-    let containerNames = containerNames(of: symbolOccurrence)
+    let containerNames = self.containerNames(of: symbolOccurrence)
     guard let containerName = containerNames.last else {
       // No containers, so nothing to do.
       return symbol.name

Tests/SourceKitLSPTests/WorkspaceSymbolsTests.swift

@@ -95,4 +95,32 @@ class WorkspaceSymbolsTests: XCTestCase {
       ]
     )
   }
+
+  func testContainerNameOfFunctionInExtension() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      struct Foo {
+        struct Bar {}
+      }
+
+      extension Foo.Bar {
+        func 1️⃣barMethod() {}
+      }
+      """
+    )
+    let response = try await project.testClient.send(WorkspaceSymbolsRequest(query: "barMethod"))
+    XCTAssertEqual(
+      response,
+      [
+        .symbolInformation(
+          SymbolInformation(
+            name: "barMethod()",
+            kind: .method,
+            location: Location(uri: project.fileURI, range: Range(project.positions["1️⃣"])),
+            containerName: "Foo.Bar"
+          )
+        )
+      ]
+    )
+  }
 }