[DocC Live Preview] Support parameter and return type disambiguations

Sources/DocCDocumentation/DocCSymbolInformation.swift

@@ -11,36 +11,30 @@
 //===----------------------------------------------------------------------===//
 
 import Foundation
-import IndexStoreDB
-package import SemanticIndex
 @_spi(LinkCompletion) @preconcurrency import SwiftDocC
+import SwiftExtensions
 import SymbolKit
 
 package struct DocCSymbolInformation {
-  let components: [(name: String, information: LinkCompletionTools.SymbolInformation)]
+  struct Component {
+    let name: String
+    let information: LinkCompletionTools.SymbolInformation
 
-  /// Find the DocCSymbolLink for a given symbol USR.
-  ///
-  /// - Parameters:
-  ///   - usr: The symbol USR to find in the index.
-  ///   - index: The CheckedIndex to search within.
-  package init?(fromUSR usr: String, in index: CheckedIndex) {
-    guard let topLevelSymbolOccurrence = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
-      return nil
+    init(fromModuleName moduleName: String) {
+      self.name = moduleName
+      self.information = LinkCompletionTools.SymbolInformation(fromModuleName: moduleName)
     }
-    let moduleName = topLevelSymbolOccurrence.location.moduleName
-    var components = [topLevelSymbolOccurrence]
-    // Find any parent symbols
-    var symbolOccurrence: SymbolOccurrence = topLevelSymbolOccurrence
-    while let parentSymbolOccurrence = symbolOccurrence.parent(index) {
-      components.insert(parentSymbolOccurrence, at: 0)
-      symbolOccurrence = parentSymbolOccurrence
+
+    init(fromSymbol symbol: SymbolGraph.Symbol) {
+      self.name = symbol.pathComponents.last ?? symbol.names.title
+      self.information = LinkCompletionTools.SymbolInformation(symbol: symbol)
     }
-    self.components =
-      [(name: moduleName, LinkCompletionTools.SymbolInformation(fromModuleName: moduleName))]
-      + components.map {
-        (name: $0.symbol.name, information: LinkCompletionTools.SymbolInformation(fromSymbolOccurrence: $0))
-      }
+  }
+
+  let components: [Component]
+
+  init(components: [Component]) {
+    self.components = components
   }
 
   package func matches(_ link: DocCSymbolLink) -> Bool {
@@ -55,73 +49,11 @@ package struct DocCSymbolInformation {
 
 fileprivate typealias KindIdentifier = SymbolGraph.Symbol.KindIdentifier
 
-extension SymbolOccurrence {
-  var doccSymbolKind: String {
-    switch symbol.kind {
-    case .module:
-      KindIdentifier.module.identifier
-    case .namespace, .namespaceAlias:
-      KindIdentifier.namespace.identifier
-    case .macro:
-      KindIdentifier.macro.identifier
-    case .enum:
-      KindIdentifier.enum.identifier
-    case .struct:
-      KindIdentifier.struct.identifier
-    case .class:
-      KindIdentifier.class.identifier
-    case .protocol:
-      KindIdentifier.protocol.identifier
-    case .extension:
-      KindIdentifier.extension.identifier
-    case .union:
-      KindIdentifier.union.identifier
-    case .typealias:
-      KindIdentifier.typealias.identifier
-    case .function:
-      KindIdentifier.func.identifier
-    case .variable:
-      KindIdentifier.var.identifier
-    case .field:
-      KindIdentifier.property.identifier
-    case .enumConstant:
-      KindIdentifier.case.identifier
-    case .instanceMethod:
-      KindIdentifier.func.identifier
-    case .classMethod:
-      KindIdentifier.func.identifier
-    case .staticMethod:
-      KindIdentifier.func.identifier
-    case .instanceProperty:
-      KindIdentifier.property.identifier
-    case .classProperty, .staticProperty:
-      KindIdentifier.typeProperty.identifier
-    case .constructor:
-      KindIdentifier.`init`.identifier
-    case .destructor:
-      KindIdentifier.deinit.identifier
-    case .conversionFunction:
-      KindIdentifier.func.identifier
-    case .unknown, .using, .concept, .commentTag, .parameter:
-      "unknown"
-    }
-  }
-}
-
 extension LinkCompletionTools.SymbolInformation {
   init(fromModuleName moduleName: String) {
     self.init(
       kind: KindIdentifier.module.identifier,
       symbolIDHash: Self.hash(uniqueSymbolID: moduleName)
     )
   }
-
-  init(fromSymbolOccurrence occurrence: SymbolOccurrence) {
-    self.init(
-      kind: occurrence.doccSymbolKind,
-      symbolIDHash: Self.hash(uniqueSymbolID: occurrence.symbol.usr),
-      parameterTypes: nil,
-      returnTypes: nil
-    )
-  }
 }

Sources/DocCDocumentation/IndexStoreDB+Extensions.swift

@@ -10,65 +10,98 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Foundation
 package import IndexStoreDB
 import SKLogging
 import SemanticIndex
-@_spi(LinkCompletion) import SwiftDocC
+@preconcurrency @_spi(LinkCompletion) import SwiftDocC
+import SwiftExtensions
+import SymbolKit
 
 extension CheckedIndex {
   /// Find a `SymbolOccurrence` that is considered the primary definition of the symbol with the given `DocCSymbolLink`.
   ///
   /// If the `DocCSymbolLink` has an ambiguous definition, the most important role of this function is to deterministically return
   /// the same result every time.
   package func primaryDefinitionOrDeclarationOccurrence(
-    ofDocCSymbolLink symbolLink: DocCSymbolLink
-  ) -> SymbolOccurrence? {
-    var components = symbolLink.components
-    guard components.count > 0 else {
-      return nil
+    ofDocCSymbolLink symbolLink: DocCSymbolLink,
+    fetchSymbolGraph: @Sendable (SymbolLocation) async throws -> String?
+  ) async throws -> SymbolOccurrence? {
+    guard let topLevelSymbolName = symbolLink.components.last?.name else {
+      throw DocCCheckedIndexError.emptyDocCSymbolLink
     }
-    // Do a lookup to find the top level symbol
-    let topLevelSymbol = components.removeLast()
+    // Find all occurrences of the symbol by name alone
     var topLevelSymbolOccurrences: [SymbolOccurrence] = []
-    forEachCanonicalSymbolOccurrence(byName: topLevelSymbol.name) { symbolOccurrence in
+    forEachCanonicalSymbolOccurrence(byName: topLevelSymbolName) { symbolOccurrence in
       topLevelSymbolOccurrences.append(symbolOccurrence)
       return true  // continue
     }
-    topLevelSymbolOccurrences = topLevelSymbolOccurrences.filter {
-      let symbolInformation = LinkCompletionTools.SymbolInformation(fromSymbolOccurrence: $0)
-      return symbolInformation.matches(topLevelSymbol.disambiguation)
-    }
-    // Search each potential symbol's parents to find an exact match
-    let symbolOccurences = topLevelSymbolOccurrences.filter { topLevelSymbolOccurrence in
-      var components = components
-      var symbolOccurrence = topLevelSymbolOccurrence
-      while let nextComponent = components.popLast(), let parentSymbolOccurrence = symbolOccurrence.parent(self) {
-        let parentSymbolInformation = LinkCompletionTools.SymbolInformation(
-          fromSymbolOccurrence: parentSymbolOccurrence
-        )
-        guard parentSymbolOccurrence.symbol.name == nextComponent.name,
-          parentSymbolInformation.matches(nextComponent.disambiguation)
-        else {
-          return false
-        }
-        symbolOccurrence = parentSymbolOccurrence
+    // Determine which of the symbol occurrences actually matches the symbol link
+    var result: [SymbolOccurrence] = []
+    for occurrence in topLevelSymbolOccurrences {
+      let info = try await doccSymbolInformation(ofUSR: occurrence.symbol.usr, fetchSymbolGraph: fetchSymbolGraph)
+      if let info, info.matches(symbolLink) {
+        result.append(occurrence)
       }
-      // If we have exactly one component left, check to see if it's the module name
-      if components.count == 1 {
-        let lastComponent = components.removeLast()
-        guard lastComponent.name == topLevelSymbolOccurrence.location.moduleName else {
-          return false
-        }
+    }
+    // Ensure that this is deterministic by sorting the results
+    result.sort()
+    if result.count > 1 {
+      logger.debug("Multiple symbols found for DocC symbol link '\(symbolLink.linkString)'")
+    }
+    return result.first
+  }
+
+  /// Find the DocCSymbolLink for a given symbol USR.
+  ///
+  /// - Parameters:
+  ///   - usr: The symbol USR to find in the index.
+  ///   - fetchSymbolGraph: Callback that returns a SymbolGraph for a given SymbolLocation
+  package func doccSymbolInformation(
+    ofUSR usr: String,
+    fetchSymbolGraph: (SymbolLocation) async throws -> String?
+  ) async throws -> DocCSymbolInformation? {
+    guard let topLevelSymbolOccurrence = primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
+      return nil
+    }
+    let moduleName = topLevelSymbolOccurrence.location.moduleName
+    var symbols = [topLevelSymbolOccurrence]
+    // Find any parent symbols
+    var symbolOccurrence: SymbolOccurrence = topLevelSymbolOccurrence
+    while let parentSymbolOccurrence = symbolOccurrence.parent(self) {
+      symbols.insert(parentSymbolOccurrence, at: 0)
+      symbolOccurrence = parentSymbolOccurrence
+    }
+    // Fetch symbol information from the symbol graph
+    var components = [DocCSymbolInformation.Component(fromModuleName: moduleName)]
+    for symbolOccurence in symbols {
+      guard let rawSymbolGraph = try await fetchSymbolGraph(symbolOccurence.location) else {
+        throw DocCCheckedIndexError.noSymbolGraph(symbolOccurence.symbol.usr)
       }
-      guard components.isEmpty else {
-        return false
+      let symbolGraph = try JSONDecoder().decode(SymbolGraph.self, from: Data(rawSymbolGraph.utf8))
+      guard let symbol = symbolGraph.symbols[symbolOccurence.symbol.usr] else {
+        throw DocCCheckedIndexError.symbolNotFound(symbolOccurence.symbol.usr)
       }
-      return true
-    }.sorted()
-    if symbolOccurences.count > 1 {
-      logger.debug("Multiple symbols found for DocC symbol link '\(symbolLink.linkString)'")
+      components.append(DocCSymbolInformation.Component(fromSymbol: symbol))
+    }
+    return DocCSymbolInformation(components: components)
+  }
+}
+
+enum DocCCheckedIndexError: LocalizedError {
+  case emptyDocCSymbolLink
+  case noSymbolGraph(String)
+  case symbolNotFound(String)
+
+  var errorDescription: String? {
+    switch self {
+    case .emptyDocCSymbolLink:
+      "The provided DocCSymbolLink was empty and could not be resolved"
+    case .noSymbolGraph(let usr):
+      "Unable to locate symbol graph for \(usr)"
+    case .symbolNotFound(let usr):
+      "Symbol \(usr) was not found in its symbol graph"
     }
-    return symbolOccurences.first
   }
 }
 

Sources/SourceKitLSP/Documentation/DoccDocumentationHandler.swift

@@ -59,7 +59,29 @@ extension DocumentationLanguageService {
           throw ResponseError.requestFailed(doccDocumentationError: .indexNotAvailable)
         }
         guard let symbolLink = DocCSymbolLink(linkString: symbolName),
-          let symbolOccurrence = index.primaryDefinitionOrDeclarationOccurrence(ofDocCSymbolLink: symbolLink)
+          let symbolOccurrence = try await index.primaryDefinitionOrDeclarationOccurrence(
+            ofDocCSymbolLink: symbolLink,
+            fetchSymbolGraph: { location in
+              guard let symbolWorkspace = try await workspaceForDocument(uri: location.documentUri),
+                let languageService = try await languageService(for: location.documentUri, .swift, in: symbolWorkspace)
+                  as? SwiftLanguageService
+              else {
+                throw ResponseError.internalError("Unable to find Swift language service for \(location.documentUri)")
+              }
+              return try await languageService.withSnapshotFromDiskOpenedInSourcekitd(
+                uri: location.documentUri,
+                fallbackSettingsAfterTimeout: false
+              ) { (snapshot, compileCommand) in
+                let (_, _, symbolGraph) = try await languageService.cursorInfo(
+                  snapshot,
+                  compileCommand: compileCommand,
+                  Range(snapshot.position(of: location)),
+                  includeSymbolGraph: true
+                )
+                return symbolGraph
+              }
+            }
+          )
         else {
           throw ResponseError.requestFailed(doccDocumentationError: .symbolNotFound(symbolName))
         }

Sources/SourceKitLSP/Swift/DoccDocumentation.swift

@@ -14,6 +14,7 @@
 import BuildSystemIntegration
 import DocCDocumentation
 import Foundation
+import IndexStoreDB
 package import LanguageServerProtocol
 import SemanticIndex
 import SKLogging
@@ -73,7 +74,21 @@ extension SwiftLanguageService {
         workspace: workspace,
         documentationManager: documentationManager,
         catalogURL: catalogURL,
-        for: symbolUSR
+        for: symbolUSR,
+        fetchSymbolGraph: { symbolLocation in
+          try await withSnapshotFromDiskOpenedInSourcekitd(
+            uri: symbolLocation.documentUri,
+            fallbackSettingsAfterTimeout: false
+          ) { (snapshot, compileCommand) in
+            let (_, _, symbolGraph) = try await self.cursorInfo(
+              snapshot,
+              compileCommand: compileCommand,
+              Range(snapshot.position(of: symbolLocation)),
+              includeSymbolGraph: true
+            )
+            return symbolGraph
+          }
+        }
       )
     }
     return try await documentationManager.renderDocCDocumentation(
@@ -90,14 +105,18 @@ extension SwiftLanguageService {
     workspace: Workspace,
     documentationManager: DocCDocumentationManager,
     catalogURL: URL?,
-    for symbolUSR: String
+    for symbolUSR: String,
+    fetchSymbolGraph: @Sendable (SymbolLocation) async throws -> String?
   ) async throws -> String? {
     guard let catalogURL else {
       return nil
     }
     let catalogIndex = try await documentationManager.catalogIndex(for: catalogURL)
     guard let index = workspace.index(checkedFor: .deletedFiles),
-      let symbolInformation = DocCSymbolInformation(fromUSR: symbolUSR, in: index),
+      let symbolInformation = try await index.doccSymbolInformation(
+        ofUSR: symbolUSR,
+        fetchSymbolGraph: fetchSymbolGraph
+      ),
       let markupExtensionFileURL = catalogIndex.documentationExtension(for: symbolInformation)
     else {
       return nil
@@ -140,7 +159,10 @@ fileprivate struct DocumentableSymbol {
     } else if let functionDecl = node.as(FunctionDeclSyntax.self) {
       self = DocumentableSymbol(node: functionDecl, position: functionDecl.name.positionAfterSkippingLeadingTrivia)
     } else if let subscriptDecl = node.as(SubscriptDeclSyntax.self) {
-      self = DocumentableSymbol(node: subscriptDecl, position: subscriptDecl.positionAfterSkippingLeadingTrivia)
+      self = DocumentableSymbol(
+        node: subscriptDecl.subscriptKeyword,
+        position: subscriptDecl.subscriptKeyword.positionAfterSkippingLeadingTrivia
+      )
     } else if let variableDecl = node.as(VariableDeclSyntax.self) {
       guard let identifier = variableDecl.bindings.only?.pattern.as(IdentifierPatternSyntax.self) else {
         return nil