update all refered states

eclipse-langium · Apr 24, 2024 · a990392 · a990392
1 parent 45a1875
commit a990392
Showing 1 changed file with 74 additions and 19 deletions.
diff --git a/examples/statemachine/src/language-server/statemachine-code-actions.ts b/examples/statemachine/src/language-server/statemachine-code-actions.ts
@@ -9,16 +9,21 @@ import type {
     DiagnosticData,
     LangiumDocument,
     MaybePromise,
+    Reference,
 } from 'langium';
 import type { CodeActionProvider } from 'langium/lsp';
 import type {
     CodeActionParams,
     Command,
     CodeAction,
     Diagnostic,
+    TextEdit,
+    Range,
+    Position,
 } from 'vscode-languageserver';
 import { CodeActionKind } from 'vscode-languageserver';
 import { IssueCodes } from './statemachine-validator.js';
+import type { State, Statemachine } from './generated/ast.js';
 
 export class StatemachineCodeActionProvider implements CodeActionProvider {
     getCodeActions(
@@ -33,7 +38,11 @@ export class StatemachineCodeActionProvider implements CodeActionProvider {
         return result;
     }
 
-    private createCodeActions(diagnostic: Diagnostic, document: LangiumDocument, accept: (ca: CodeAction | undefined) => void): void {
+    private createCodeActions(
+        diagnostic: Diagnostic,
+        document: LangiumDocument,
+        accept: (ca: CodeAction | undefined) => void
+    ): void {
         switch ((diagnostic.data as DiagnosticData)?.code) {
             case IssueCodes.StateNameUppercase:
                 accept(this.makeUpperCase(diagnostic, document));
@@ -45,35 +54,81 @@ export class StatemachineCodeActionProvider implements CodeActionProvider {
         return undefined;
     }
 
-    private makeUpperCase(diagnostic: Diagnostic, document: LangiumDocument): CodeAction {
-        const range = {
-            start: diagnostic.range.start,
-            end: {
-                line: diagnostic.range.start.line,
-                character: diagnostic.range.start.character + 1,
-            },
-        };
+    private makeUpperCase(
+        diagnostic: Diagnostic,
+        document: LangiumDocument
+    ): CodeAction {
+        const changes: TextEdit[] = [];
+
+        const stateName = document.textDocument.getText(diagnostic.range);
+        const { init, states } = document.parseResult.value as Statemachine;
+        this.updateChangesForReferencedState(init, stateName, document, changes);
+
+        states.forEach(({ transitions }) => {
+            transitions.forEach(({ state }) => {
+                this.updateChangesForReferencedState(
+                    state,
+                    stateName,
+                    document,
+                    changes
+                );
+            });
+        });
+
+        const range = this.getFirstLetterRange(diagnostic.range.start);
+        changes.push(this.createTextEditForState(range, document));
         return {
             title: 'First letter to upper case',
             kind: CodeActionKind.QuickFix,
             diagnostics: [diagnostic],
             isPreferred: true,
             edit: {
                 changes: {
-                    [document.textDocument.uri]: [
-                        {
-                            range,
-                            newText: document.textDocument
-                                .getText(range)
-                                .toUpperCase(),
-                        },
-                    ],
+                    [document.textDocument.uri]: changes,
                 },
             },
         };
     }
 
-    private removeUnusedSymbol(diagnostic: Diagnostic, document: LangiumDocument): CodeAction {
+    private createTextEditForState(
+        range: Range,
+        document: LangiumDocument
+    ): TextEdit {
+        const changeRange = this.getFirstLetterRange(range.start);
+        return {
+            range: changeRange,
+            newText: document.textDocument.getText(changeRange).toUpperCase(),
+        };
+    }
+
+    private updateChangesForReferencedState(
+        state: Reference<State>,
+        name: string,
+        document: LangiumDocument,
+        changes: TextEdit[]
+    ): void {
+        if (state.$refNode && state.ref && state.ref.name === name) {
+            const { range } = state.$refNode;
+            const changeRange = this.getFirstLetterRange(range.start);
+            changes.push(this.createTextEditForState(changeRange, document));
+        }
+    }
+
+    private getFirstLetterRange(position: Position): Range {
+        const range: Range = {
+            start: position,
+            end: {
+                line: position.line,
+                character: position.character + 1,
+            },
+        };
+        return range;
+    }
+
+    private removeUnusedSymbol(
+        diagnostic: Diagnostic,
+        document: LangiumDocument
+    ): CodeAction {
         return {
             title: 'Remove unsed symbol',
             kind: CodeActionKind.QuickFix,
@@ -84,7 +139,7 @@ export class StatemachineCodeActionProvider implements CodeActionProvider {
                     [document.textDocument.uri]: [
                         {
                             range: diagnostic.range,
-                            newText: ''
+                            newText: '',
                         },
                     ],
                 },