Better definitions & symbols

jayadamsmorgan · Apr 1, 2024 · abc92b7 · abc92b7
1 parent 571d807
commit abc92b7
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 92 deletions.
diff --git a/Sources/pkl-lsp/ASTHelper.swift b/Sources/pkl-lsp/ASTHelper.swift
@@ -123,21 +123,24 @@ public enum ASTHelper {
         }
     }
 
-    static func getPositionContext(module: ASTNode, position: Position) -> (ASTNode)? {
-        allEndNodes(node: module, importDepth: 0).first { node in
-            if node.range.positionRange.lowerBound.line > position.line {
-                return false
-            }
-            if node.range.positionRange.upperBound.line < position.line {
-                return false
-            }
-            if node.range.positionRange.lowerBound.character / 2 > position.character {
-                return false
-            }
-            if node.range.positionRange.upperBound.character / 2 < position.character {
-                return false
+    static func getPositionContext(module: ASTNode, position: Position) -> ASTNode? {
+        var contextNode: ASTNode?
+        var smallestRange: Int = Int.max
+
+        enumerate(node: module) { node in
+            let range = node.range.positionRange
+            if position.line >= range.lowerBound.line &&
+                (position.line > range.lowerBound.line || position.character >= range.lowerBound.character / 2) &&
+                position.line <= range.upperBound.line &&
+                (position.line < range.upperBound.line || position.character <= range.upperBound.character / 2) {
+
+                let rangeSize = (range.upperBound.line - range.lowerBound.line) * 1000 + (range.upperBound.character - range.lowerBound.character)
+                if rangeSize < smallestRange {
+                    smallestRange = rangeSize
+                    contextNode = node
+                }
             }
-            return true
         }
+        return contextNode
     }
 }
diff --git a/Sources/pkl-lsp/RequestHandlers/DefinitionHandler.swift b/Sources/pkl-lsp/RequestHandlers/DefinitionHandler.swift
@@ -20,42 +20,44 @@ public class DefinitionHandler {
 
         if let context = context as? PklStringLiteral {
             if context.type == .importString {
-                logger.debug("DefinitionHandler: Trying to find imported module.")
-                var relPath = context.value ?? ""
-                relPath.removeAll(where: { $0 == "\"" })
-                let modulePath = URL(fileURLWithPath: document.uri)
-                    .deletingLastPathComponent()
-                    .appendingPathComponent(relPath)
-                    .standardized
-                do {
-                    guard try modulePath.checkResourceIsReachable() else {
-                        logger.debug("DefinitionHandler: Module at path \(modulePath.absoluteString) is not reachable.")
-                        return nil
-                    }
-                    logger.debug("DefinitionHandler: Module at path \(modulePath.absoluteString) found.")
-                    return .optionA(Location(uri: modulePath.absoluteString, range: LSPRange.zero))
-                } catch {
-                    logger.debug("DefinitionHandler: Unable to check if module exists: \(error)")
-                    return nil
-                }
+                return await provideForModuleImport(path: context)
             }
-        }
-
-        guard let parent = context.parent else {
-            logger.debug("DefinitionHandler: Parent node is nil.")
             return nil
         }
-        logger.debug("DefinitionHandler: Parent node: \(String(describing: parent))")
-        if let parent = parent as? PklVariable {
-            logger.debug("DefinitionHandler: Searching for definition of variable \(parent.identifier?.value ?? "nil")")
-            guard let reference = parent.reference else {
-                logger.debug("DefinitionHandler: Reference is nil.")
+        if let context = context as? PklModuleImport {
+            return await provideForModuleImport(path: context.path)
+        }
+
+        if let context = context as? PklVariable {
+            logger.debug("DefinitionHandler: Searching for definition of variable \(context.identifier?.value ?? "nil")")
+            guard let reference = context.reference else {
+                logger.debug("DefinitionHandler: Variable reference is nil.")
                 return nil
             }
             return .optionA(Location(uri: reference.document.uri, range: reference.range.getLSPRange()))
         }
 
-        let range = context.range.getLSPRange()
-        return .optionA(Location(uri: document.uri, range: range))
+        return nil
+    }
+
+    private func provideForModuleImport(path: PklStringLiteral) async -> DefinitionResponse {
+            logger.debug("DefinitionHandler: Trying to find imported module.")
+            var relPath = path.value ?? ""
+            relPath.removeAll(where: { $0 == "\"" })
+            let modulePath = URL(fileURLWithPath: path.document.uri)
+                .deletingLastPathComponent()
+                .appendingPathComponent(relPath)
+                .standardized
+            do {
+                guard try modulePath.checkResourceIsReachable() else {
+                    logger.debug("DefinitionHandler: Module at path \(modulePath.absoluteString) is not reachable.")
+                    return nil
+                }
+                logger.debug("DefinitionHandler: Module at path \(modulePath.absoluteString) found.")
+                return .optionA(Location(uri: modulePath.absoluteString, range: LSPRange.zero))
+            } catch {
+                logger.debug("DefinitionHandler: Unable to check if module exists: \(error)")
+                return nil
+            }
     }
 }
diff --git a/Sources/pkl-lsp/RequestHandlers/DocumentSymbolsHandler.swift b/Sources/pkl-lsp/RequestHandlers/DocumentSymbolsHandler.swift
@@ -9,73 +9,55 @@ public class DocumentSymbolsHandler {
         self.logger = logger
     }
 
-    private func getSymbols(node: ASTNode) -> [DocumentSymbol] {
+    private func getSymbols(node: ASTNode) async -> [DocumentSymbol] {
         guard let children = node.children else {
             return []
         }
         var symbols: [DocumentSymbol] = []
         for child in children {
+            if let objectBody = child as? PklObjectBody {
+                symbols.append(contentsOf: await getSymbols(node: objectBody))
+            }
             if let classNode = child as? PklClassDeclaration {
-                let classSymbol = DocumentSymbol(
-                    name: classNode.classIdentifier?.value ?? "Unknown",
-                    detail: "class",
-                    kind: .class,
-                    range: classNode.range.getLSPRange(),
-                    selectionRange: classNode.range.getLSPRange(),
-                    children: getSymbols(node: classNode)
-                )
-                symbols.append(classSymbol)
+                symbols.append(await createDocumentSymbol(node: classNode, name: classNode.classIdentifier?.value, kind: .class))
             }
             if let functionNode = child as? PklFunctionDeclaration {
-                let functionSymbol = DocumentSymbol(
-                    name: functionNode.body?.identifier?.value ?? "Unknown",
-                    detail: "function",
-                    kind: .function,
-                    range: functionNode.range.getLSPRange(),
-                    selectionRange: functionNode.range.getLSPRange(),
-                    children: getSymbols(node: functionNode)
-                )
-                symbols.append(functionSymbol)
-            }
-            if let objectNode = child as? PklObjectBody {
-                let objectSymbol = DocumentSymbol(
-                    name: "Object",
-                    detail: "object",
-                    kind: .module,
-                    range: objectNode.range.getLSPRange(),
-                    selectionRange: objectNode.range.getLSPRange(),
-                    children: getSymbols(node: objectNode)
-                )
-                symbols.append(objectSymbol)
+                symbols.append(await createDocumentSymbol(node: functionNode, name: functionNode.body?.identifier?.value, kind: .function))
             }
             if let propertyNode = child as? PklObjectProperty {
-                let propertySymbol = DocumentSymbol(
-                    name: propertyNode.identifier?.value ?? "Unknown",
-                    detail: "property",
-                    kind: .property,
-                    range: propertyNode.range.getLSPRange(),
-                    selectionRange: propertyNode.range.getLSPRange(),
-                    children: getSymbols(node: propertyNode)
-                )
-                symbols.append(propertySymbol)
+                symbols.append(await createDocumentSymbol(node: propertyNode, name: propertyNode.identifier?.value, kind: .property))
             }
             if let propertyNode = child as? PklClassProperty {
-                let propertySymbol = DocumentSymbol(
-                    name: propertyNode.identifier?.value ?? "Unknown",
-                    detail: "property",
-                    kind: .property,
-                    range: propertyNode.range.getLSPRange(),
-                    selectionRange: propertyNode.range.getLSPRange(),
-                    children: getSymbols(node: propertyNode)
-                )
-                symbols.append(propertySymbol)
+                symbols.append(await createDocumentSymbol(node: propertyNode, name: propertyNode.identifier?.value, kind: .property))
             }
         }
         return symbols
     }
 
+    private func createDocumentSymbol<T: ASTNode>(node: T, name: String?, kind: SymbolKind) async -> DocumentSymbol {
+        DocumentSymbol(
+            name: name ?? "Unknown",
+            detail: name ?? "Unknown",
+            kind: kind,
+            range: node.range.getLSPRange(),
+            selectionRange: node.range.getLSPRange(),
+            children: await getSymbols(node: node)
+        )
+    }
+
     public func provide(document _: Document, module: ASTNode, params: DocumentSymbolParams) async -> DocumentSymbolResponse {
-        let symbols: [DocumentSymbol] = getSymbols(node: module)
+        if let moduleHeader = module.children?.first(where: { $0 is PklModuleHeader }) as? PklModuleHeader {
+            let symbols = DocumentSymbol(
+                name: moduleHeader.moduleClause?.name?.value ?? "Module",
+                detail: moduleHeader.moduleClause?.name?.value ?? "Module",
+                kind: .module,
+                range: module.range.getLSPRange(),
+                selectionRange: module.range.getLSPRange(),
+                children: await getSymbols(node: module)
+            )
+            return DocumentSymbolResponse(.optionA([symbols]))
+        }
+        let symbols: [DocumentSymbol] = await getSymbols(node: module)
         logger.debug("LSP DocumentSymbols: Found \(symbols.count) symbols in \(params.textDocument.uri).")
         return DocumentSymbolResponse(.optionA(symbols))
     }