Merge pull request #38 from RumbleDB/object-completion

feat: enrich completion intent with object lookup support

Author: CaiJimmy

.changeset/tasty-mails-grin.md

@@ -0,0 +1,5 @@
+---
+"jsoniq-language-server": minor
+---
+
+feat: enrich completion intent with object lookup support

packages/language-server/src/completion.ts

@@ -29,11 +29,23 @@ import {
 } from "./assets/function-docs.js";
 import { collectCompletionIntent } from "./parser/index.js";
 import { getDocumentText } from "./parser/utils.js";
-import { formatSequenceType } from "./static-typecheck/types.js";
+import {
+    formatSequenceType,
+    formatTypeDefinition,
+    type ObjectTypeDefinition,
+} from "./static-typecheck/types.js";
+import { getTypeAtPositionFromSource } from "./type-at-position/service.js";
 
 const VARIABLE_PREFIX_PATTERN = /\$[A-Za-z0-9_.:-]*$/;
+const OBJECT_FIELD_PREFIX_PATTERN = /[A-Za-z_][A-Za-z0-9_:-]*$/;
 const GENERIC_BUILTIN_PARAMETER_PREFIX = "$arg";
 
+interface DotCompletionContext {
+    dotOffset: number;
+    fieldPrefix: string;
+    syntheticSource: string;
+}
+
 export function findCompletions(document: TextDocument, position: Position): CompletionItem[] {
     const source = getDocumentText(document);
     const cursorOffset = document.offsetAt(position);
@@ -103,6 +115,85 @@ export function findCompletions(document: TextDocument, position: Position): Com
     ]);
 }
 
+export async function findCompletionsWithTypeInfo(
+    document: TextDocument,
+    position: Position,
+): Promise<CompletionItem[]> {
+    const source = getDocumentText(document);
+    const cursorOffset = document.offsetAt(position);
+    const intent = collectCompletionIntent(document, cursorOffset);
+    const dotContext = getDotCompletionContext(source, cursorOffset);
+
+    if (
+        dotContext !== null &&
+        intent?.allowObjectLookup &&
+        intent.objectLookupDotOffset === dotContext.dotOffset
+    ) {
+        /// Return only dot completions because it's more relevant
+        return findDotCompletions(document, dotContext);
+    }
+
+    return findCompletions(document, position);
+}
+
+async function findDotCompletions(
+    document: TextDocument,
+    context: DotCompletionContext,
+): Promise<CompletionItem[]> {
+    const result = await getTypeAtPositionFromSource(
+        document.uri,
+        context.syntheticSource,
+        document.positionAt(context.dotOffset),
+    );
+    const objectType = result.sequenceType?.itemType;
+
+    if (objectType?.kind !== "object") {
+        return [];
+    }
+
+    const replaceRange = {
+        start: document.positionAt(context.dotOffset + 1),
+        end: document.positionAt(context.dotOffset + 1 + context.fieldPrefix.length),
+    };
+
+    return withSortText(
+        objectFieldCompletions(objectType, context.fieldPrefix).map(([fieldName, fieldType]) => ({
+            label: fieldName,
+            kind: CompletionItemKind.Field,
+            detail: formatTypeDefinition(fieldType),
+            textEdit: TextEdit.replace(replaceRange, fieldName),
+        })),
+    );
+}
+
+function objectFieldCompletions(
+    objectType: ObjectTypeDefinition,
+    fieldPrefix: string,
+): Array<[string, ObjectTypeDefinition["fields"][string]]> {
+    return Object.entries(objectType.fields).filter(([fieldName]) =>
+        fieldName.startsWith(fieldPrefix),
+    );
+}
+
+function getDotCompletionContext(
+    source: string,
+    cursorOffset: number,
+): DotCompletionContext | null {
+    const prefix = source.slice(0, cursorOffset);
+    const fieldPrefix = prefix.match(OBJECT_FIELD_PREFIX_PATTERN)?.[0] ?? "";
+    const dotOffset = cursorOffset - fieldPrefix.length - 1;
+
+    if (dotOffset < 0 || source[dotOffset] !== ".") {
+        return null;
+    }
+
+    return {
+        dotOffset,
+        fieldPrefix,
+        syntheticSource: source.slice(0, dotOffset) + source.slice(cursorOffset),
+    };
+}
+
 function toCompletionItem(declaration: BaseDefinition): CompletionItem {
     const name = definitionNameToString(declaration);
     if (isSourceFunctionDefinition(declaration)) {

packages/language-server/src/main.ts

@@ -8,7 +8,7 @@ import {
     type InitializeResult,
 } from "vscode-languageserver/node";
 
-import { findCompletions } from "./completion.js";
+import { findCompletionsWithTypeInfo } from "./completion.js";
 import { config, InitializationOptions, mergeConfiguration } from "./config.js";
 import { findDefinitionLocation } from "./definitions.js";
 import { findHover } from "./hover.js";
@@ -98,7 +98,7 @@ connection.onInitialize((params: InitializeParams): InitializeResult => {
                 triggerCharacters: ["(", ","],
             },
             completionProvider: {
-                triggerCharacters: ["$"],
+                triggerCharacters: ["$", "."],
             },
             renameProvider: {
                 prepareProvider: true,
@@ -202,7 +202,7 @@ connection.onCompletion((params) => {
         return [];
     }
 
-    return findCompletions(document, params.position);
+    return findCompletionsWithTypeInfo(document, params.position);
 });
 
 connection.languages.semanticTokens.on((params) => {

packages/language-server/src/parser/adapters/jsoniq/completion-data.ts

@@ -32,6 +32,7 @@ export const PREFERRED_COMPLETION_RULES = new Set([
     JsoniqParser.RULE_varRef,
     JsoniqParser.RULE_qname,
     JsoniqParser.RULE_functionCall,
+    JsoniqParser.RULE_objectLookup,
 ]);
 
 export const KEYWORD_COMPLETIONS: LanguageKeywordCompletion[] = [

packages/language-server/src/parser/adapters/jsoniq/index.ts

@@ -9,6 +9,7 @@ import {
     PREFERRED_COMPLETION_RULES,
 } from "./completion-data.js";
 import { JsoniqTokenContextAnalyzer } from "./completion-token-context.js";
+import { JsoniqLexer } from "./grammar/JsoniqLexer.js";
 import { JsoniqParser } from "./grammar/JsoniqParser.js";
 import { parseJsoniq } from "./parse.js";
 
@@ -23,6 +24,8 @@ export const jsoniqParserAdapter: ParserAdapter = {
             preferredRules: PREFERRED_COMPLETION_RULES,
             languageKeywords: KEYWORD_COMPLETIONS,
             isFunctionCallRule: (ruleIndex) => ruleIndex === JsoniqParser.RULE_functionCall,
+            isObjectLookupRule: (ruleIndex) => ruleIndex === JsoniqParser.RULE_objectLookup,
+            isObjectLookupDotToken: (tokenType) => tokenType === JsoniqLexer.DOT,
             isVariableReferenceRule: (ruleIndex) => ruleIndex === JsoniqParser.RULE_varRef,
             isTypeReferenceRule: (ruleIndex) =>
                 ruleIndex === JsoniqParser.RULE_sequenceType ||

packages/language-server/src/parser/adapters/xquery/index.ts

@@ -23,6 +23,8 @@ export const xqueryParserAdapter: ParserAdapter = {
             preferredRules: PREFERRED_COMPLETION_RULES,
             languageKeywords: KEYWORD_COMPLETIONS,
             isFunctionCallRule: (ruleIndex) => ruleIndex === XQueryParser.RULE_functionCall,
+            isObjectLookupRule: () => false,
+            isObjectLookupDotToken: () => false,
             isVariableReferenceRule: (ruleIndex) => ruleIndex === XQueryParser.RULE_varRef,
             isTypeReferenceRule: (ruleIndex) =>
                 ruleIndex === XQueryParser.RULE_sequenceType ||

packages/language-server/src/parser/completion.ts

@@ -20,6 +20,8 @@ type CompletionOptions<T extends TokenContextAnalyzer> = {
     preferredRules: Set<number>;
     languageKeywords: LanguageKeywordCompletion[];
     isFunctionCallRule(ruleIndex: number): boolean;
+    isObjectLookupRule(ruleIndex: number): boolean;
+    isObjectLookupDotToken(tokenType: number): boolean;
     isVariableReferenceRule(ruleIndex: number): boolean;
     isTypeReferenceRule(ruleIndex: number): boolean;
     tokenName(tokenType: number): string | number;
@@ -46,6 +48,15 @@ export function getCompletionIntent<T extends TokenContextAnalyzer>(
     const allowVariables =
         candidates.tokenContext.allowReferences &&
         hasCandidateRule(candidates, options.isVariableReferenceRule);
+    const objectLookupDotOffset = findObjectLookupDotOffset(
+        parsed.tokens,
+        cursorOffset,
+        options.isObjectLookupDotToken,
+    );
+    const allowObjectLookup =
+        candidates.tokenContext.allowReferences &&
+        hasCandidateRule(candidates, options.isObjectLookupRule) &&
+        objectLookupDotOffset !== undefined;
     const allowTypes =
         candidates.tokenContext.allowTypeReferences ||
         (candidates.tokenContext.allowReferences &&
@@ -58,8 +69,10 @@ export function getCompletionIntent<T extends TokenContextAnalyzer>(
     logger.debug("Completion candidates:", {
         allowFunctions,
         allowVariables,
+        allowObjectLookup,
         allowTypes,
         allowVariableDeclarations,
+        objectLookupDotOffset,
         keywords,
         expectedTokens,
         expectedRules,
@@ -70,6 +83,8 @@ export function getCompletionIntent<T extends TokenContextAnalyzer>(
         allowVariableReferences: allowVariables,
         allowVariableDeclarations,
         allowFunctions,
+        allowObjectLookup,
+        ...(objectLookupDotOffset === undefined ? {} : { objectLookupDotOffset }),
         allowTypes,
         keywords,
     };
@@ -86,6 +101,28 @@ function hasCandidateToken(candidates: CompletionCandidates, tokenType: number):
     return candidates.tokenTypes.has(tokenType);
 }
 
+/// This is used to strip out object lookup dot tokens that are not actually part of an object lookup expression, e.g. in the following example:
+/// ```
+/// let x = 1;
+/// x. // <- cursor here
+/// ```
+/// We send only the query text up to the dot to the server so it does not give errors
+function findObjectLookupDotOffset(
+    tokens: Token[],
+    cursorOffset: number,
+    isObjectLookupDotToken: (tokenType: number) => boolean,
+): number | undefined {
+    return tokens
+        .filter(
+            (token) =>
+                token.type !== Token.EOF &&
+                isObjectLookupDotToken(token.type) &&
+                token.start < cursorOffset &&
+                token.stop < cursorOffset,
+        )
+        .at(-1)?.start;
+}
+
 function keywordCompletions(
     candidates: CompletionCandidates,
     languageKeywords: LanguageKeywordCompletion[],

packages/language-server/src/parser/types/completion.ts

@@ -8,6 +8,8 @@ export interface CompletionIntent {
     allowVariableDeclarations: boolean;
     allowFunctions: boolean;
     allowTypes: boolean;
+    allowObjectLookup: boolean;
+    objectLookupDotOffset?: number;
     keywords: KeywordCompletion[];
 }
 

packages/language-server/src/type-at-position/service.ts

@@ -17,6 +17,14 @@ const EMPTY_RESULT: TypeAtPositionWireResult = {};
 export async function getTypeAtPosition(
     document: TextDocument,
     position: Position,
+): Promise<TypeAtPositionWireResult> {
+    return getTypeAtPositionFromSource(document.uri, getDocumentText(document), position);
+}
+
+export async function getTypeAtPositionFromSource(
+    documentUri: string,
+    source: string,
+    position: Position,
 ): Promise<TypeAtPositionWireResult> {
     const client = getWrapperClient();
     if (!client.isUsable()) {
@@ -26,15 +34,15 @@ export async function getTypeAtPosition(
     try {
         const response = await client.sendRequest<TypeAtPositionRequestSpec>({
             requestType: REQUEST_TYPE_TYPE_AT_POSITION,
-            body: Buffer.from(getDocumentText(document), "utf8").toString("base64"),
-            documentUri: document.uri,
+            body: Buffer.from(source, "utf8").toString("base64"),
+            documentUri,
             position,
         });
 
         return response.body;
     } catch (error) {
         logger.warn(
-            `Type-at-position unavailable for ${document.uri}: ${error instanceof Error ? error.message : String(error)}`,
+            `Type-at-position unavailable for ${documentUri}: ${error instanceof Error ? error.message : String(error)}`,
         );
         return EMPTY_RESULT;
     }

packages/language-server/tests/completion.test.ts

@@ -1,4 +1,4 @@
-import { findCompletions } from "server/completion.js";
+import { findCompletions, findCompletionsWithTypeInfo } from "server/completion.js";
 import { describe, expect, it } from "vitest";
 import { type CompletionItem, type Position } from "vscode-languageserver";
 import { type TextDocument } from "vscode-languageserver-textdocument";
@@ -278,6 +278,51 @@ describe("JSONiq completion", () => {
         expect(labelsAtCursor).toContain("declare function");
         expect(labelsAtCursor).toContain("declare variable");
     });
+
+    it("suggests object fields after the dot operator using real type inference", async () => {
+        const document = testDocument("completion-dot-object", [
+            "declare variable $a := {",
+            '  "name": "Ada",',
+            '  "age": 42',
+            "};",
+            "",
+            "$a.",
+        ]);
+
+        const items = await findCompletionsWithTypeInfo(document, { line: 5, character: 3 });
+
+        expect(labels(items)).toContain("age");
+        expect(labels(items)).toContain("name");
+        expect(items.find((item) => item.label === "name")?.textEdit).toEqual({
+            range: {
+                start: { line: 5, character: 3 },
+                end: { line: 5, character: 3 },
+            },
+            newText: "name",
+        });
+    }, 45_000);
+
+    it("filters real object field completions by the prefix typed after dot", async () => {
+        const document = testDocument("completion-dot-object-prefix", [
+            "declare variable $a := {",
+            '  "name": "Ada",',
+            '  "age": 42',
+            "};",
+            "",
+            "$a.na",
+        ]);
+
+        const items = await findCompletionsWithTypeInfo(document, { line: 5, character: 5 });
+
+        expect(labels(items)).toContain("name");
+        expect(items.find((item) => item.label === "name")?.textEdit).toEqual({
+            range: {
+                start: { line: 5, character: 3 },
+                end: { line: 5, character: 5 },
+            },
+            newText: "name",
+        });
+    }, 45_000);
 });
 
 describe("XQuery completion", () => {