Merge pull request #1095 from ahoppen/ahoppen/sort-index-results

Sort results returned by the index

Author: ahoppen

Sources/LanguageServerProtocol/SupportTypes/Location.swift

@@ -13,7 +13,16 @@
 /// Range within a particular document.
 ///
 /// For a location where the document is implied, use `Position` or `Range<Position>`.
-public struct Location: ResponseType, Hashable, Codable, CustomDebugStringConvertible {
+public struct Location: ResponseType, Hashable, Codable, CustomDebugStringConvertible, Comparable {
+  public static func < (lhs: Location, rhs: Location) -> Bool {
+    if lhs.uri != rhs.uri {
+      return lhs.uri.stringValue < rhs.uri.stringValue
+    }
+    if lhs.range.lowerBound != rhs.range.lowerBound {
+      return lhs.range.lowerBound < rhs.range.lowerBound
+    }
+    return lhs.range.upperBound < rhs.range.upperBound
+  }
 
   public var uri: DocumentURI
 

Sources/SourceKitLSP/SourceKitServer.swift

@@ -1607,7 +1607,7 @@ extension SourceKitServer {
 
   /// Find all symbols in the workspace that include a string in their name.
   /// - returns: An array of SymbolOccurrences that match the string.
-  func findWorkspaceSymbols(matching: String) -> [SymbolOccurrence] {
+  func findWorkspaceSymbols(matching: String) throws -> [SymbolOccurrence] {
     // Ignore short queries since they are:
     // - noisy and slow, since they can match many symbols
     // - normally unintentional, triggered when the user types slowly or if the editor doesn't
@@ -1624,25 +1624,24 @@ extension SourceKitServer {
         subsequence: true,
         ignoreCase: true
       ) { symbol in
+        if Task.isCancelled {
+          return false
+        }
         guard !symbol.location.isSystem && !symbol.roles.contains(.accessorOf) else {
           return true
         }
         symbolOccurrenceResults.append(symbol)
-        // FIXME: Once we have cancellation support, we should fetch all results and take the top
-        // `maxWorkspaceSymbolResults` symbols but bail if cancelled.
-        //
-        // Until then, take the first `maxWorkspaceSymbolResults` symbols to limit the impact of
-        // queries which match many symbols.
-        return symbolOccurrenceResults.count < maxWorkspaceSymbolResults
+        return true
       }
+      try Task.checkCancellation()
     }
-    return symbolOccurrenceResults
+    return symbolOccurrenceResults.sorted()
   }
 
   /// Handle a workspace/symbol request, returning the SymbolInformation.
   /// - returns: An array with SymbolInformation for each matching symbol in the workspace.
   func workspaceSymbols(_ req: WorkspaceSymbolsRequest) async throws -> [WorkspaceSymbolItem]? {
-    let symbols = findWorkspaceSymbols(matching: req.query).map(WorkspaceSymbolItem.init)
+    let symbols = try findWorkspaceSymbols(matching: req.query).map(WorkspaceSymbolItem.init)
     return symbols
   }
 
@@ -1870,6 +1869,7 @@ extension SourceKitServer {
       }
       return indexToLSPLocation(occurrence.location)
     }
+    .sorted()
   }
 
   /// Returns the result of a `DefinitionRequest` by running a `SymbolInfoRequest`, inspecting
@@ -1962,18 +1962,19 @@ extension SourceKitServer {
         position: req.position
       )
     )
-    guard let symbol = symbols.first,
-      let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
-    else {
+    guard let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index else {
       return nil
     }
-    guard let usr = symbol.usr else { return nil }
-    var occurrences = index.occurrences(ofUSR: usr, roles: .baseOf)
-    if occurrences.isEmpty {
-      occurrences = index.occurrences(relatedToUSR: usr, roles: .overrideOf)
-    }
+    let locations = symbols.flatMap { (symbol) -> [Location] in
+      guard let usr = symbol.usr else { return [] }
+      var occurrences = index.occurrences(ofUSR: usr, roles: .baseOf)
+      if occurrences.isEmpty {
+        occurrences = index.occurrences(relatedToUSR: usr, roles: .overrideOf)
+      }
 
-    return .locations(occurrences.compactMap { indexToLSPLocation($0.location) })
+      return occurrences.compactMap { indexToLSPLocation($0.location) }
+    }
+    return .locations(locations.sorted())
   }
 
   func references(
@@ -1987,17 +1988,19 @@ extension SourceKitServer {
         position: req.position
       )
     )
-    guard let symbol = symbols.first, let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
-    else {
+    guard let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index else {
       return []
     }
-    guard let usr = symbol.usr else { return [] }
-    logger.info("performing indexed jump-to-def with usr \(usr)")
-    var roles: SymbolRole = [.reference]
-    if req.context.includeDeclaration {
-      roles.formUnion([.declaration, .definition])
+    let locations = symbols.flatMap { (symbol) -> [Location] in
+      guard let usr = symbol.usr else { return [] }
+      logger.info("Finding references for USR \(usr)")
+      var roles: SymbolRole = [.reference]
+      if req.context.includeDeclaration {
+        roles.formUnion([.declaration, .definition])
+      }
+      return index.occurrences(ofUSR: usr, roles: roles).compactMap { indexToLSPLocation($0.location) }
     }
-    return index.occurrences(ofUSR: usr, roles: roles).compactMap { indexToLSPLocation($0.location) }
+    return locations.sorted()
   }
 
   private func indexToLSPCallHierarchyItem(
@@ -2032,29 +2035,31 @@ extension SourceKitServer {
         position: req.position
       )
     )
-    guard let symbol = symbols.first, let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
-    else {
+    guard let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index else {
       return nil
     }
     // For call hierarchy preparation we only locate the definition
-    guard let usr = symbol.usr else { return nil }
+    let usrs = symbols.compactMap(\.usr)
 
     // Only return a single call hierarchy item. Returning multiple doesn't make sense because they will all have the
     // same USR (because we query them by USR) and will thus expand to the exact same call hierarchy.
-    // Also, VS Code doesn't seem to like multiple call hierarchy items being returned and fails to display any results
-    // if they are, failing with `Cannot read properties of undefined (reading 'map')`.
-    guard let definition = index.definitionOrDeclarationOccurrences(ofUSR: usr).first else {
-      return nil
-    }
-    guard let location = indexToLSPLocation(definition.location) else {
-      return nil
-    }
-    let callHierarchyItem = self.indexToLSPCallHierarchyItem(
-      symbol: definition.symbol,
-      moduleName: definition.location.moduleName,
-      location: location
-    )
-    return [callHierarchyItem]
+    var callHierarchyItems = usrs.compactMap { (usr) -> CallHierarchyItem? in
+      guard let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: usr) else {
+        return nil
+      }
+      guard let location = indexToLSPLocation(definition.location) else {
+        return nil
+      }
+      return self.indexToLSPCallHierarchyItem(
+        symbol: definition.symbol,
+        moduleName: definition.location.moduleName,
+        location: location
+      )
+    }.sorted(by: { Location(uri: $0.uri, range: $0.range) < Location(uri: $1.uri, range: $1.range) })
+
+    // Ideally, we should show multiple symbols. But VS Code fails to display call hierarchies with multiple root items,
+    // failing with `Cannot read properties of undefined (reading 'map')`. Pick the first one.
+    return Array(callHierarchyItems.prefix(1))
   }
 
   /// Extracts our implementation-specific data about a call hierarchy
@@ -2090,7 +2095,7 @@ extension SourceKitServer {
       return occurrence.relations.filter { $0.symbol.kind.isCallable }
         .map { related in
           // Resolve the caller's definition to find its location
-          let definition = index.occurrences(ofUSR: related.symbol.usr, roles: [.definition, .declaration]).first
+          let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: related.symbol.usr)
           let definitionSymbolLocation = definition?.location
           let definitionLocation = definitionSymbolLocation.flatMap(indexToLSPLocation)
 
@@ -2104,7 +2109,7 @@ extension SourceKitServer {
           )
         }
     }
-    return calls
+    return calls.sorted(by: { $0.from.name < $1.from.name })
   }
 
   func outgoingCalls(_ req: CallHierarchyOutgoingCallsRequest) async throws -> [CallHierarchyOutgoingCall]? {
@@ -2121,7 +2126,7 @@ extension SourceKitServer {
       }
 
       // Resolve the callee's definition to find its location
-      let definition = index.occurrences(ofUSR: occurrence.symbol.usr, roles: [.definition, .declaration]).first
+      let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: occurrence.symbol.usr)
       let definitionSymbolLocation = definition?.location
       let definitionLocation = definitionSymbolLocation.flatMap(indexToLSPLocation)
 
@@ -2134,7 +2139,7 @@ extension SourceKitServer {
         fromRanges: [location.range]
       )
     }
-    return calls
+    return calls.sorted(by: { $0.to.name < $1.to.name })
   }
 
   private func indexToLSPTypeHierarchyItem(
@@ -2153,7 +2158,7 @@ extension SourceKitServer {
       if conformances.isEmpty {
         name = symbol.name
       } else {
-        name = "\(symbol.name): \(conformances.map(\.symbol.name).joined(separator: ", "))"
+        name = "\(symbol.name): \(conformances.map(\.symbol.name).sorted().joined(separator: ", "))"
       }
       // Add the file name and line to the detail string
       if let url = location.uri.fileURL,
@@ -2197,25 +2202,42 @@ extension SourceKitServer {
         position: req.position
       )
     )
-    guard let symbol = symbols.first else {
+    guard !symbols.isEmpty else {
       return nil
     }
     guard let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index else {
       return nil
     }
-    guard let usr = symbol.usr else { return nil }
-    return index.occurrences(ofUSR: usr, roles: [.definition, .declaration])
-      .compactMap { info -> TypeHierarchyItem? in
-        guard let location = indexToLSPLocation(info.location) else {
-          return nil
+    let usrs =
+      symbols
+      .filter {
+        // Only include references to type. For example, we don't want to find the type hierarchy of a constructor when
+        // starting the type hierarchy on `Foo()``.
+        switch $0.kind {
+        case .class, .enum, .interface, .struct: return true
+        default: return false
         }
-        return self.indexToLSPTypeHierarchyItem(
-          symbol: info.symbol,
-          moduleName: info.location.moduleName,
-          location: location,
-          index: index
-        )
       }
+      .compactMap(\.usr)
+    let typeHierarchyItems = usrs.compactMap { (usr) -> TypeHierarchyItem? in
+      guard
+        let info = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: usr),
+        let location = indexToLSPLocation(info.location)
+      else {
+        return nil
+      }
+      return self.indexToLSPTypeHierarchyItem(
+        symbol: info.symbol,
+        moduleName: info.location.moduleName,
+        location: location,
+        index: index
+      )
+    }
+    .sorted(by: { $0.name < $1.name })
+
+    // Ideally, we should show multiple symbols. But VS Code fails to display type hierarchies with multiple root items,
+    // failing with `Cannot read properties of undefined (reading 'map')`. Pick the first one.
+    return Array(typeHierarchyItems.prefix(1))
   }
 
   /// Extracts our implementation-specific data about a type hierarchy
@@ -2249,7 +2271,12 @@ extension SourceKitServer {
     // Resolve retroactive conformances via the extensions
     let extensions = index.occurrences(ofUSR: data.usr, roles: .extendedBy)
     let retroactiveConformanceOccurs = extensions.flatMap { occurrence -> [SymbolOccurrence] in
-      guard let related = occurrence.relations.first else {
+      if occurrence.relations.count > 1 {
+        // When the occurrence has an `extendedBy` relation, it's an extension declaration. An extension can only extend
+        // a single type, so there can only be a single relation here.
+        logger.fault("Expected at most extendedBy relation but got \(occurrence.relations.count)")
+      }
+      guard let related = occurrence.relations.sorted().first else {
         return []
       }
       return index.occurrences(relatedToUSR: related.symbol.usr, roles: .baseOf)
@@ -2263,7 +2290,7 @@ extension SourceKitServer {
       }
 
       // Resolve the supertype's definition to find its location
-      let definition = index.occurrences(ofUSR: occurrence.symbol.usr, roles: [.definition, .declaration]).first
+      let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: occurrence.symbol.usr)
       let definitionSymbolLocation = definition?.location
       let definitionLocation = definitionSymbolLocation.flatMap(indexToLSPLocation)
 
@@ -2274,7 +2301,7 @@ extension SourceKitServer {
         index: index
       )
     }
-    return types
+    return types.sorted(by: { $0.name < $1.name })
   }
 
   func subtypes(_ req: TypeHierarchySubtypesRequest) async throws -> [TypeHierarchyItem]? {
@@ -2289,14 +2316,19 @@ extension SourceKitServer {
 
     // Convert occurrences to type hierarchy items
     let types = occurs.compactMap { occurrence -> TypeHierarchyItem? in
-      guard let location = indexToLSPLocation(occurrence.location),
-        let related = occurrence.relations.first
+      if occurrence.relations.count > 1 {
+        // An occurrence with a `baseOf` or `extendedBy` relation is an occurrence inside an inheritance clause.
+        // Such an occurrence can only be the source of a single type, namely the one that the inheritance clause belongs
+        // to.
+        logger.fault("Expected at most extendedBy or baseOf relation but got \(occurrence.relations.count)")
+      }
+      guard let related = occurrence.relations.sorted().first, let location = indexToLSPLocation(occurrence.location)
       else {
         return nil
       }
 
       // Resolve the subtype's definition to find its location
-      let definition = index.occurrences(ofUSR: related.symbol.usr, roles: [.definition, .declaration]).first
+      let definition = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: related.symbol.usr)
       let definitionSymbolLocation = definition.map(\.location)
       let definitionLocation = definitionSymbolLocation.flatMap(indexToLSPLocation)
 
@@ -2307,7 +2339,7 @@ extension SourceKitServer {
         index: index
       )
     }
-    return types
+    return types.sorted { $0.name < $1.name }
   }
 
   func pollIndex(_ req: PollIndexRequest) async throws -> VoidResponse {
@@ -2348,6 +2380,14 @@ fileprivate extension IndexStoreDB {
     }
     return occurrences(ofUSR: usr, roles: [.declaration])
   }
+
+  /// Find a `SymbolOccurrence` that is considered the primary definition of the symbol with the given USR.
+  ///
+  /// If the USR has an ambiguous definition, the most important role of this function is to deterministically return
+  /// the same result every time.
+  func primaryDefinitionOrDeclarationOccurrence(ofUSR usr: String) -> SymbolOccurrence? {
+    return definitionOrDeclarationOccurrences(ofUSR: usr).sorted().first
+  }
 }
 
 extension IndexSymbolKind {
@@ -2394,7 +2434,11 @@ extension IndexSymbolKind {
 extension SymbolOccurrence {
   /// Get the name of the symbol that is a parent of this symbol, if one exists
   func getContainerName() -> String? {
-    return relations.first(where: { $0.roles.contains(.childOf) })?.symbol.name
+    let containers = 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")
+    }
+    return containers.sorted().first?.symbol.name
   }
 }