feat: Support code action of generating opposite relation fields

packages/schema/src/language-server/constants.ts

@@ -12,3 +12,7 @@ export const SCALAR_TYPES = ['String', 'Int', 'Float', 'Decimal', 'BigInt', 'Boo
  * Name of standard library module
  */
 export const STD_LIB_MODULE_NAME = 'stdlib.zmodel';
+
+export enum IssueCodes {
+    MissingOppositeRelation = 'miss-opposite-relation',
+}

packages/schema/src/language-server/validator/datamodel-validator.ts

@@ -15,7 +15,7 @@ import {
 import { ValidationAcceptor } from 'langium';
 import pluralize from 'pluralize';
 import { analyzePolicies } from '../../utils/ast-utils';
-import { SCALAR_TYPES } from '../constants';
+import { IssueCodes, SCALAR_TYPES } from '../constants';
 import { AstValidator } from '../types';
 import { assignableToAttributeParam, validateDuplicatedDeclarations } from './utils';
 
@@ -297,7 +297,7 @@ export default class DataModelValidator implements AstValidator<DataModel> {
             accept(
                 'error',
                 `The relation field "${field.name}" on model "${field.$container.name}" is missing an opposite relation field on model "${oppositeModel.name}"`,
-                { node: field }
+                { node: field, code: IssueCodes.MissingOppositeRelation }
             );
             return;
         } else if (oppositeFields.length > 1) {

packages/schema/src/language-server/zmodel-code-action.ts

@@ -0,0 +1,138 @@
+import { DataModel, DataModelField, isDataModel } from '@zenstackhq/language/ast';
+import {
+    AstReflection,
+    CodeActionProvider,
+    findDeclarationNodeAtOffset,
+    getContainerOfType,
+    IndexManager,
+    LangiumDocument,
+    LangiumServices,
+    MaybePromise,
+} from 'langium';
+
+import {
+    CancellationToken,
+    CodeAction,
+    CodeActionKind,
+    CodeActionParams,
+    Command,
+    Diagnostic,
+} from 'vscode-languageserver';
+import { IssueCodes } from './constants';
+import { ZModelFormatter } from './zmodel-formatter';
+
+export class ZModelCodeActionProvider implements CodeActionProvider {
+    protected readonly reflection: AstReflection;
+    protected readonly indexManager: IndexManager;
+    protected readonly formatter: ZModelFormatter;
+
+    constructor(services: LangiumServices) {
+        this.reflection = services.shared.AstReflection;
+        this.indexManager = services.shared.workspace.IndexManager;
+        this.formatter = services.lsp.Formatter as ZModelFormatter;
+    }
+
+    getCodeActions(
+        document: LangiumDocument,
+        params: CodeActionParams,
+        cancelToken?: CancellationToken
+    ): MaybePromise<Array<Command | CodeAction> | undefined> {
+        const result: CodeAction[] = [];
+        const acceptor = (ca: CodeAction | undefined) => ca && result.push(ca);
+        for (const diagnostic of params.context.diagnostics) {
+            this.createCodeActions(diagnostic, document, acceptor);
+        }
+        return result;
+    }
+
+    private createCodeActions(
+        diagnostic: Diagnostic,
+        document: LangiumDocument,
+        accept: (ca: CodeAction | undefined) => void
+    ) {
+        switch (diagnostic.code) {
+            case IssueCodes.MissingOppositeRelation:
+                accept(this.fixMissingOppositeRelation(diagnostic, document));
+        }
+
+        return undefined;
+    }
+
+    private fixMissingOppositeRelation(diagnostic: Diagnostic, document: LangiumDocument): CodeAction | undefined {
+        const offset = document.textDocument.offsetAt(diagnostic.range.start);
+        const rootCst = document.parseResult.value.$cstNode;
+
+        if (rootCst) {
+            const cstNode = findDeclarationNodeAtOffset(rootCst, offset);
+
+            const astNode = cstNode?.element as DataModelField;
+
+            const oppositeModel = astNode.type.reference!.ref! as DataModel;
+
+            const lastField = oppositeModel.fields[oppositeModel.fields.length - 1];
+
+            const container = getContainerOfType(cstNode?.element, isDataModel) as DataModel;
+
+            const idField = container.fields.find((f) =>
+                f.attributes.find((attr) => attr.decl.ref?.name === '@id')
+            ) as DataModelField;
+
+            if (container && container.$cstNode && idField) {
+                // indent
+                let indent = '\t';
+                const formatOptions = this.formatter.getFormatOptions();
+                if (formatOptions?.insertSpaces) {
+                    indent = ' '.repeat(formatOptions.tabSize);
+                }
+                indent = indent.repeat(this.formatter.getIndent());
+
+                const typeName = container.name;
+                const fieldName = this.lowerCaseFirstLetter(typeName);
+
+                // might already exist
+                let referenceField = '';
+
+                const idFieldName = idField.name;
+                const referenceIdFieldName = fieldName + this.upperCaseFirstLetter(idFieldName);
+
+                if (!oppositeModel.fields.find((f) => f.name === referenceIdFieldName)) {
+                    referenceField = '\n' + indent + `${referenceIdFieldName} ${idField.type.type}`;
+                }
+
+                return {
+                    title: `Add opposite relation fields on ${oppositeModel.name}`,
+                    kind: CodeActionKind.QuickFix,
+                    diagnostics: [diagnostic],
+                    isPreferred: false,
+                    edit: {
+                        changes: {
+                            [document.textDocument.uri]: [
+                                {
+                                    range: {
+                                        start: lastField.$cstNode!.range.end,
+                                        end: lastField.$cstNode!.range.end,
+                                    },
+                                    newText:
+                                        '\n' +
+                                        indent +
+                                        `${fieldName} ${typeName} @relation(fields: [${referenceIdFieldName}], references: [${idFieldName}])` +
+                                        referenceField,
+                                },
+                            ],
+                        },
+                    },
+                };
+            }
+        }
+
+        return undefined;
+    }
+
+    private lowerCaseFirstLetter(str: string) {
+        return str.charAt(0).toLowerCase() + str.slice(1);
+    }
+
+    private upperCaseFirstLetter(str: string) {
+        return str.charAt(0).toUpperCase() + str.slice(1);
+    }
+}

packages/schema/src/language-server/zmodel-formatter.ts

@@ -1,13 +1,16 @@
-import { AbstractFormatter, AstNode, Formatting } from 'langium';
+import { AbstractFormatter, AstNode, Formatting, LangiumDocument } from 'langium';
 
 import * as ast from '@zenstackhq/language/ast';
+import { FormattingOptions, Range, TextEdit } from 'vscode-languageserver';
 
 export class ZModelFormatter extends AbstractFormatter {
+    private formatOptions?: FormattingOptions;
     protected format(node: AstNode): void {
         const formatter = this.getNodeFormatter(node);
         if (ast.isAbstractDeclaration(node)) {
             const bracesOpen = formatter.keyword('{');
             const bracesClose = formatter.keyword('}');
+            // this line decide the indent count return by this.getIndent()
             formatter.interior(bracesOpen, bracesClose).prepend(Formatting.indent());
             bracesOpen.prepend(Formatting.oneSpace());
             bracesClose.prepend(Formatting.newLine());
@@ -17,4 +20,21 @@ export class ZModelFormatter extends AbstractFormatter {
             nodes.prepend(Formatting.noIndent());
         }
     }
+
+    protected override doDocumentFormat(
+        document: LangiumDocument<AstNode>,
+        options: FormattingOptions,
+        range?: Range | undefined
+    ): TextEdit[] {
+        this.formatOptions = options;
+        return super.doDocumentFormat(document, options, range);
+    }
+
+    public getFormatOptions(): FormattingOptions | undefined {
+        return this.formatOptions;
+    }
+
+    public getIndent() {
+        return 1;
+    }
 }

packages/schema/src/language-server/zmodel-module.ts

@@ -20,6 +20,7 @@ import {
 import { TextDocuments } from 'vscode-languageserver';
 import { TextDocument } from 'vscode-languageserver-textdocument';
 import { ZModelValidationRegistry, ZModelValidator } from './validator/zmodel-validator';
+import { ZModelCodeActionProvider } from './zmodel-code-action';
 import { ZModelFormatter } from './zmodel-formatter';
 import { ZModelLinker } from './zmodel-linker';
 import { ZModelScopeComputation } from './zmodel-scope';
@@ -56,6 +57,7 @@ export const ZModelModule: Module<ZModelServices, PartialLangiumServices & ZMode
     },
     lsp: {
         Formatter: () => new ZModelFormatter(),
+        CodeActionProvider: (services) => new ZModelCodeActionProvider(services),
     },
 };