From c94d708bcf6c0f661ae8cacfe9649eaa7e04c502 Mon Sep 17 00:00:00 2001
From: Reda Al Sulais <u.yokozuna@gmail.com>
Date: Wed, 24 Apr 2024 10:38:22 +0300
Subject: [PATCH] update all refered states

---
 .../statemachine-code-actions.ts              | 90 +++++++++++++++----
 1 file changed, 71 insertions(+), 19 deletions(-)

diff --git a/examples/statemachine/src/language-server/statemachine-code-actions.ts b/examples/statemachine/src/language-server/statemachine-code-actions.ts
index 3a7158165..01d865648 100644
--- a/examples/statemachine/src/language-server/statemachine-code-actions.ts
+++ b/examples/statemachine/src/language-server/statemachine-code-actions.ts
@@ -16,9 +16,13 @@ import type {
     Command,
     CodeAction,
     Diagnostic,
+    TextEdit,
+    Range,
+    Position,
 } from 'vscode-languageserver';
 import { CodeActionKind } from 'vscode-languageserver';
 import { IssueCodes } from './statemachine-validator.js';
+import type { Statemachine } from './generated/ast.js';
 
 export class StatemachineCodeActionProvider implements CodeActionProvider {
     getCodeActions(
@@ -33,7 +37,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,14 +53,40 @@ 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;
+        if (init.ref && init.$refNode && init.ref.name === stateName) {
+            init.$refNode;
+            const { range } = init.$refNode;
+            const changeRange = this.getFirstLetterRange(range.start);
+            changes.push(this.updateChangesForState(changeRange, document));
+        }
+
+        states.forEach(({ transitions }) => {
+            transitions.forEach((transition) => {
+                const { state } = transition;
+                if (
+                    state.ref &&
+                    state.$refNode &&
+                    state.ref.name === stateName
+                ) {
+                    const { range } = state.$refNode;
+                    const changeRange = this.getFirstLetterRange(range.start);
+                    changes.push(
+                        this.updateChangesForState(changeRange, document)
+                    );
+                }
+            });
+        });
+
+        const range = this.getFirstLetterRange(diagnostic.range.start);
+        changes.push(this.updateChangesForState(range, document));
         return {
             title: 'First letter to upper case',
             kind: CodeActionKind.QuickFix,
@@ -60,20 +94,38 @@ export class StatemachineCodeActionProvider implements CodeActionProvider {
             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 updateChangesForState(
+        range: Range,
+        document: LangiumDocument
+    ): TextEdit {
+        const changeRange = this.getFirstLetterRange(range.start);
+        return {
+            range: changeRange,
+            newText: document.textDocument.getText(changeRange).toUpperCase(),
+        };
+    }
+
+    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 +136,7 @@ export class StatemachineCodeActionProvider implements CodeActionProvider {
                     [document.textDocument.uri]: [
                         {
                             range: diagnostic.range,
-                            newText: ''
+                            newText: '',
                         },
                     ],
                 },