From 0b2ca89903d7557ec8cf05ec7acbc7bfe4f908fb Mon Sep 17 00:00:00 2001 From: Alex Hoppen Date: Thu, 29 Feb 2024 13:05:52 -0800 Subject: [PATCH] Ensure we handle multiple results returned by the index correctly This replaces all uses of `.first` on index results in `SourceKitServer` by code that explicitly handles multiple results. --- Sources/SourceKitLSP/SourceKitServer.swift | 153 +++++++++++------- .../Swift/SwiftLanguageServer.swift | 35 ++-- 2 files changed, 115 insertions(+), 73 deletions(-) diff --git a/Sources/SourceKitLSP/SourceKitServer.swift b/Sources/SourceKitLSP/SourceKitServer.swift index c3805690..1601677b 100644 --- a/Sources/SourceKitLSP/SourceKitServer.swift +++ b/Sources/SourceKitLSP/SourceKitServer.swift @@ -1961,18 +1961,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) }.sorted()) + return occurrences.compactMap { indexToLSPLocation($0.location) } + } + return .locations(locations.sorted()) } func references( @@ -1986,17 +1987,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) }.sorted() + return locations.sorted() } private func indexToLSPCallHierarchyItem( @@ -2031,29 +2034,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 @@ -2089,7 +2094,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) @@ -2120,7 +2125,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) @@ -2196,26 +2201,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 - ) } - .sorted(by: { $0.name < $1.name }) + .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 +2270,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 +2289,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) @@ -2289,14 +2315,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) @@ -2348,6 +2379,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 +2433,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 } } diff --git a/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift b/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift index f28a3934..fa3836e6 100644 --- a/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift +++ b/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift @@ -1091,38 +1091,37 @@ extension sourcekitd_api_uid_t { func asSymbolKind(_ vals: sourcekitd_api_values) -> SymbolKind? { switch self { - case vals.declClass: + case vals.declClass, vals.refClass, vals.declActor, vals.refActor: return .class - case vals.declMethodInstance, - vals.declMethodStatic, - vals.declMethodClass: + case vals.declMethodInstance, vals.refMethodInstance, + vals.declMethodStatic, vals.refMethodStatic, + vals.declMethodClass, vals.refMethodClass: return .method - case vals.declVarInstance, - vals.declVarStatic, - vals.declVarClass: + case vals.declVarInstance, vals.refVarInstance, + vals.declVarStatic, vals.refVarStatic, + vals.declVarClass, vals.refVarClass: return .property - case vals.declEnum: + case vals.declEnum, vals.refEnum: return .enum - case vals.declEnumElement: + case vals.declEnumElement, vals.refEnumElement: return .enumMember - case vals.declProtocol: + case vals.declProtocol, vals.refProtocol: return .interface - case vals.declFunctionFree: + case vals.declFunctionFree, vals.refFunctionFree: return .function - case vals.declVarGlobal, - vals.declVarLocal: + case vals.declVarGlobal, vals.refVarGlobal, + vals.declVarLocal, vals.refVarLocal: return .variable - case vals.declStruct: + case vals.declStruct, vals.refStruct: return .struct - case vals.declGenericTypeParam: + case vals.declGenericTypeParam, vals.refGenericTypeParam: return .typeParameter case vals.declExtension: - // There are no extensions in LSP, so I return something vaguely similar + // There are no extensions in LSP, so we return something vaguely similar return .namespace case vals.refModule: return .module - case vals.refConstructor, - vals.declConstructor: + case vals.declConstructor, vals.refConstructor: return .constructor default: return nil