Fully functionable DocComments

You can now see doc comments in completion descriptions!

Sources/pkl-lsp/AST/AST.swift

@@ -23,6 +23,8 @@ public class ASTNode: ASTEvaluation, IdentifiableNode {
     public var importDepth: Int
     public var document: Document
 
+    public var docComment: PklDocComment?
+
     public var parent: ASTNode?
 
     public var children: [ASTNode]? = nil
@@ -46,6 +48,10 @@ public struct ASTRange: Hashable {
     let positionRange: Range<Position>
     let byteRange: Range<UInt32>
 
+    var pointRange: Range<Point> {
+        positionRange.lowerBound.getPoint() ..< positionRange.upperBound.getPoint()
+    }
+
     public init(positionRange: Range<Position>, byteRange: Range<UInt32>) {
         self.positionRange = positionRange
         self.byteRange = byteRange

Sources/pkl-lsp/AST/DocComment.swift

@@ -1,7 +1,7 @@
 import Foundation
 import LanguageServerProtocol
 
-class DocComment: ASTNode {
+public class PklDocComment: ASTNode {
     var text: String
 
     init(text: String, range: ASTRange, importDepth: Int, document: Document) {

Sources/pkl-lsp/ASTHelper.swift

@@ -125,15 +125,15 @@ public enum ASTHelper {
 
     static func getPositionContext(module: ASTNode, position: Position) -> ASTNode? {
         var contextNode: ASTNode?
-        var smallestRange: Int = Int.max
+        var smallestRange = 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) {
-
+            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

Sources/pkl-lsp/DocumentProvider.swift

@@ -143,7 +143,7 @@ public actor DocumentProvider {
             logger.error("LSP Completion: Document \(params.textDocument.uri) is not available.")
             return nil
         }
-        return await completionHandler.provide(document: document, module: module, params: params)
+        return await completionHandler.provide(module: module, params: params)
     }
 
     public func provideRenaming(params: RenameParams) async -> RenameResponse {
@@ -207,7 +207,7 @@ public actor DocumentProvider {
             logger.error("LSP Definition: AST for \(params.textDocument.uri) is not available.")
             return nil
         }
-        return await definitionHandler.provide(document: document, module: module, params: params)
+        return await definitionHandler.provide(module: module, params: params)
     }
 
     public func provideDiagnostics(document: Document) async throws {

Sources/pkl-lsp/RequestHandlers/CompletionHandler.swift

@@ -9,47 +9,47 @@ public class CompletionHandler {
         self.logger = logger
     }
 
-    public func provide(document _: Document, module: ASTNode, params _: CompletionParams) async -> CompletionResponse {
+    public func provide(module: ASTNode, params _: CompletionParams) async -> CompletionResponse {
         var completions: [CompletionItem] = []
 
         ASTHelper.enumerate(node: module) { node in
             if let classObject = node as? PklClassDeclaration {
                 completions.append(CompletionItem(
                     label: classObject.classIdentifier?.value ?? "",
                     kind: .class,
-                    detail: "Pickle object"
+                    detail: node.docComment?.text ?? "Pickle object"
                 ))
                 return
             }
             if let function = node as? PklFunctionDeclaration {
                 completions.append(CompletionItem(
                     label: function.body?.identifier?.value ?? "",
                     kind: .function,
-                    detail: "Pickle function"
+                    detail: node.docComment?.text ?? "Pickle function"
                 ))
                 return
             }
             if let object = node as? PklObjectProperty {
                 completions.append(CompletionItem(
                     label: object.identifier?.value ?? "",
                     kind: .property,
-                    detail: "Pickle object property"
+                    detail: node.docComment?.text ?? "Pickle object property"
                 ))
                 return
             }
             if let objectEntry = node as? PklObjectEntry {
                 completions.append(CompletionItem(
                     label: objectEntry.strIdentifier?.value ?? "",
                     kind: .property,
-                    detail: "Pickle object entry"
+                    detail: node.docComment?.text ?? "Pickle object entry"
                 ))
                 return
             }
             if let classProperty = node as? PklClassProperty {
                 completions.append(CompletionItem(
                     label: classProperty.identifier?.value ?? "",
                     kind: .property,
-                    detail: "Pickle property"
+                    detail: node.docComment?.text ?? "Pickle property"
                 ))
                 return
             }

Sources/pkl-lsp/RequestHandlers/DefinitionHandler.swift

@@ -9,7 +9,7 @@ public class DefinitionHandler {
         self.logger = logger
     }
 
-    public func provide(document: Document, module: ASTNode, params: TextDocumentPositionParams) async -> DefinitionResponse {
+    public func provide(module: ASTNode, params: TextDocumentPositionParams) async -> DefinitionResponse {
         let positionContext = ASTHelper.getPositionContext(module: module, position: params.position)
 
         guard let context = positionContext else {
@@ -27,7 +27,7 @@ public class DefinitionHandler {
         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 {
@@ -41,23 +41,23 @@ public class DefinitionHandler {
     }
 
     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)")
+        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
+        }
     }
 }

Sources/pkl-lsp/RequestHandlers/DocumentSymbolsHandler.swift

@@ -16,44 +16,44 @@ public class DocumentSymbolsHandler {
         var symbols: [DocumentSymbol] = []
         for child in children {
             if let objectBody = child as? PklObjectBody {
-                symbols.append(contentsOf: await getSymbols(node: objectBody))
+                await symbols.append(contentsOf: getSymbols(node: objectBody))
             }
             if let classNode = child as? PklClassDeclaration {
-                symbols.append(await createDocumentSymbol(node: classNode, name: classNode.classIdentifier?.value, kind: .class))
+                await symbols.append(createDocumentSymbol(node: classNode, name: classNode.classIdentifier?.value, kind: .class))
             }
             if let functionNode = child as? PklFunctionDeclaration {
-                symbols.append(await createDocumentSymbol(node: functionNode, name: functionNode.body?.identifier?.value, kind: .function))
+                await symbols.append(createDocumentSymbol(node: functionNode, name: functionNode.body?.identifier?.value, kind: .function))
             }
             if let propertyNode = child as? PklObjectProperty {
-                symbols.append(await createDocumentSymbol(node: propertyNode, name: propertyNode.identifier?.value, kind: .property))
+                await symbols.append(createDocumentSymbol(node: propertyNode, name: propertyNode.identifier?.value, kind: .property))
             }
             if let propertyNode = child as? PklClassProperty {
-                symbols.append(await createDocumentSymbol(node: propertyNode, name: propertyNode.identifier?.value, kind: .property))
+                await symbols.append(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(
+    private func createDocumentSymbol(node: some ASTNode, name: String?, kind: SymbolKind) async -> DocumentSymbol {
+        await DocumentSymbol(
             name: name ?? "Unknown",
             detail: name ?? "Unknown",
             kind: kind,
             range: node.range.getLSPRange(),
             selectionRange: node.range.getLSPRange(),
-            children: await getSymbols(node: node)
+            children: getSymbols(node: node)
         )
     }
 
     public func provide(document _: Document, module: ASTNode, params: DocumentSymbolParams) async -> DocumentSymbolResponse {
         if let moduleHeader = module.children?.first(where: { $0 is PklModuleHeader }) as? PklModuleHeader {
-            let symbols = DocumentSymbol(
+            let symbols = await 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)
+                children: getSymbols(node: module)
             )
             return DocumentSymbolResponse(.optionA([symbols]))
         }

Sources/pkl-lsp/TreeSitterParser.swift

@@ -110,14 +110,17 @@ public class TreeSitterParser {
             return
         }
         let astRoot = await tsNodeToASTNode(node: rootNode, in: document, importDepth: importDepth, parent: nil)
+        astParsedTrees[document] = astRoot
+        await parseVariableReferences(document: document)
+        await parseImportModules(document: document)
+        if let astRoot {
+            await processAndAttachDocComments(node: astRoot)
+        }
         if logger.logLevel == .trace {
             if let astRoot {
                 listASTNodes(rootNode: astRoot)
             }
         }
-        astParsedTrees[document] = astRoot
-        await parseVariableReferences(document: document)
-        await parseImportModules(document: document)
     }
 
     private func parseVariableReferences(document: Document) async {
@@ -126,7 +129,7 @@ public class TreeSitterParser {
             return
         }
         guard let references = objectReferences[document] else {
-            logger.debug("No references found for document \(document.uri)")
+            logger.debug("No object references found for document \(document.uri)")
             return
         }
         for variable in variables {
@@ -146,6 +149,26 @@ public class TreeSitterParser {
         }
     }
 
+    private func processAndAttachDocComments(node: ASTNode) async {
+        guard let children = node.children else {
+            return
+        }
+        var docComment: PklDocComment?
+        for child in children {
+            if let child = child as? PklDocComment {
+                if docComment != nil {
+                    continue
+                }
+                docComment = child
+                continue
+            }
+            if let docComment {
+                child.docComment = docComment
+            }
+            docComment = nil
+        }
+    }
+
     private func parseImportModules(document: Document) async {
         guard let importModules = importModules[document] else {
             logger.debug("No import modules found for document \(document.uri)")
@@ -293,6 +316,34 @@ public class TreeSitterParser {
         return module
     }
 
+    private func buildDocComment(node: Node, in document: Document, importDepth: Int, parent: ASTNode?) async -> PklDocComment? {
+        var text = document.getTextInByteRange(node.byteRange)
+        if let previousSubling = node.previousSibling {
+            if previousSubling.symbol == PklTreeSitterSymbols.sym_docComment.rawValue {
+                // If the previous sibling is a docComment, we don't want to build another one
+                return nil
+            }
+        }
+        logger.debug("Starting building doc comment...")
+        var range = ASTRange(pointRange: node.pointRange, byteRange: node.byteRange)
+        var nextSibling = node.nextSibling
+        while let nextSiblingUnwrapped = nextSibling,
+              nextSiblingUnwrapped.symbol == PklTreeSitterSymbols.sym_lineComment.rawValue,
+              document.getTextInByteRange(nextSiblingUnwrapped.byteRange).starts(with: "///")
+        {
+            // If the next sibling is a docComment, we append it to the text and change range
+            text.append(contentsOf: "\n\(document.getTextInByteRange(nextSiblingUnwrapped.byteRange))")
+            range = ASTRange(
+                pointRange: range.pointRange.lowerBound ..< nextSiblingUnwrapped.pointRange.upperBound,
+                byteRange: range.byteRange.lowerBound ..< nextSiblingUnwrapped.byteRange.upperBound
+            )
+            nextSibling = nextSiblingUnwrapped.nextSibling
+        }
+        let comment = PklDocComment(text: text, range: range, importDepth: importDepth, document: document)
+        comment.parent = parent
+        return comment
+    }
+
     private func tsNodeToASTNode(node: Node, in document: Document, importDepth: Int, parent: ASTNode?) async -> ASTNode? {
         guard let tsSymbol = PklTreeSitterSymbols(node.symbol) else {
             logger.debug("Unable to parse node with symbol \(node.symbol)")
@@ -666,10 +717,15 @@ public class TreeSitterParser {
             logger.debug("Not implemented")
 
         case .sym_lineComment:
+            // Pkl tree-sitter recognizes docComment as a usual lineComment for some reason, so we do a workaround:
+            let text = document.getTextInByteRange(node.byteRange)
+            if text.starts(with: "///") {
+                return await buildDocComment(node: node, in: document, importDepth: importDepth, parent: parent)
+            }
             logger.debug("Line comment at \(node.pointRange).")
 
         case .sym_docComment:
-            logger.debug("Not implemented")
+            return await buildDocComment(node: node, in: document, importDepth: importDepth, parent: parent)
 
         case .sym_blockComment:
             logger.debug("Block comment at \(node.pointRange)")