Complex State Transitions in DFA

client/src/types/fsm.ts

@@ -2,7 +2,7 @@
 
 export type LJStateMachine = {
     className: string;
-    initialStates: string[];
+    initialTransitions: { to: string; postCond?: string | null }[];
     states: string[];
-    transitions: { from: string; to: string; label: string }[];
+    transitions: { from: string; to: string; label: string; preCond?: string | null; postCond?: string | null }[];
 };

client/src/webview/diagram.ts

@@ -21,7 +21,7 @@ let startY = 0;
  * @param sm 
  * @returns Mermaid diagram string
  */
-export function createMermaidDiagram(sm: LJStateMachine | undefined, orientation: "LR" | "TB"): string {
+export function createMermaidDiagram(sm: LJStateMachine | undefined, orientation: "LR" | "TB", showConditions = false): string {
     if (!sm) return '';
 
     const lines: string[] = [];
@@ -33,17 +33,19 @@ export function createMermaidDiagram(sm: LJStateMachine | undefined, orientation
     lines.push('stateDiagram-v2');
     lines.push(`    direction ${orientation}`);
 
-    // initial states
-    sm.initialStates.forEach(state => {
-        lines.push(`    [*] --> ${state}`);
+    // initial transitions
+    sm.initialTransitions.forEach(transition => {
+        const label = getInitialTransitionLabel(transition.postCond, showConditions);
+        lines.push(`    [*] --> ${transition.to}${label ? ` : ${label}` : ''}`);
     });
 
     // group transitions by from/to states and merge labels
     const transitionMap = new Map<string, string[]>();
     sm.transitions.forEach(transition => {
+        const label = getTransitionLabel(transition.label, transition.preCond, transition.postCond, showConditions);
         const key = `${transition.from}|${transition.to}`;
         if (!transitionMap.has(key)) transitionMap.set(key, []);
-        transitionMap.get(key)?.push(transition.label);
+        transitionMap.get(key)?.push(label);
     });
 
     // add transitions
@@ -56,6 +58,33 @@ export function createMermaidDiagram(sm: LJStateMachine | undefined, orientation
     return lines.join('\n');
 }
 
+function getTransitionLabel(label: string, preCond?: string | null, postCond?: string | null, showConditions = false): string {
+    if (!showConditions) {
+        return escapeMermaidLabel(label);
+    }
+
+    return [
+        getConditionLabel('pre', preCond),
+        escapeMermaidLabel(label),
+        getConditionLabel('post', postCond)
+    ].filter(Boolean).join('<br/>');
+}
+
+function getInitialTransitionLabel(postCond?: string | null, showConditions = false): string {
+    return showConditions ? getConditionLabel('post', postCond) : '';
+}
+
+function getConditionLabel(kind: 'pre' | 'post', cond?: string | null): string {
+    if (!cond) {
+        return '';
+    }
+    return `<span class="state-cond state-cond-${kind}">${escapeMermaidLabel(cond)}</span>`;
+}
+
+function escapeMermaidLabel(label: string): string {
+    return label.replace(/&/g, '&amp;').replace(/"/g, '\\"').replace(/</g, '&lt;').replace(/>/g, '&gt;');
+}
+
 /**
  * Renders Mermaid diagrams in the document
  * @param document The document object

client/src/webview/script.ts

@@ -34,6 +34,7 @@ export function getScript(vscode: VSCodeApi, document: Document, window: Window)
     let errorAtCursor: RefinementMismatchError;
     let selectedTab: NavTab = 'diagnostics';
     let diagramOrientation: "LR" | "TB" = "TB";
+    let showDiagramConditions = false;
     let currentDiagram: string = '';
     let revealTimeout: ReturnType<typeof setTimeout> | undefined;
     const contextSectionState: ContextSectionState = {
@@ -147,6 +148,15 @@ export function getScript(vscode: VSCodeApi, document: Document, window: Window)
             return;
         }
 
+        // toggle diagram conditions
+        if (target.id === 'diagram-conditions-btn') {
+            e.stopPropagation();
+            if ((target as HTMLButtonElement).disabled) return;
+            showDiagramConditions = !showDiagramConditions;
+            updateView();
+            return;
+        }
+
         // zoom in
         if (target.id === 'zoom-in-btn') {
             e.stopPropagation();
@@ -260,6 +270,7 @@ export function getScript(vscode: VSCodeApi, document: Document, window: Window)
                 break;
             case 'fsm':
                 stateMachine = msg.sm as LJStateMachine;
+                showDiagramConditions = false;
                 if (selectedTab === 'fsm') updateView();
                 break;
             case 'context':
@@ -282,9 +293,9 @@ export function getScript(vscode: VSCodeApi, document: Document, window: Window)
                 root.innerHTML = renderDiagnosticsView(diagnostics, showAllDiagnostics, currentFile, expandedErrors);
                 break;
             case 'fsm': {
-                const diagram = createMermaidDiagram(stateMachine, diagramOrientation);
+                const diagram = createMermaidDiagram(stateMachine, diagramOrientation, showDiagramConditions);
                 currentDiagram = diagram;
-                root.innerHTML = renderStateMachineView(stateMachine, diagram, diagramOrientation);
+                root.innerHTML = renderStateMachineView(stateMachine, diagram, diagramOrientation, showDiagramConditions);
                 if (stateMachine) renderMermaidDiagram(document, window);
                 break;
             }

client/src/webview/styles.ts

@@ -38,6 +38,8 @@ export function getStyles(): string {
             --lj-token-boolean: var(--vscode-debugTokenExpression-boolean, var(--vscode-symbolIcon-booleanForeground, var(--vscode-editor-foreground)));
             --lj-token-identifier: var(--vscode-debugTokenExpression-name, var(--vscode-editor-foreground));
             --lj-token-punctuation: var(--vscode-editor-foreground);
+            --lj-state-cond-pre: var(--lj-token-function);
+            --lj-state-cond-post: var(--lj-token-type);
         }
         body.vscode-light {
             --lj-token-keyword: #0000FF;
@@ -619,9 +621,12 @@ export function getStyles(): string {
             background: none;
             opacity: 1;
         }
+        .diagram-control-btn.active {
+            opacity: 1;
+        }
         .diagram-control-btn:disabled {
             cursor: default;
-            opacity: 0.8;
+            opacity: 0.35;
         }
         .mermaid .statediagramTitleText {
             font-size: 30px!important;
@@ -650,6 +655,18 @@ export function getStyles(): string {
             color: var(--vscode-foreground) !important;
             background: var(--vscode-editor-background) !important;
         }
+        .diagram-container .mermaid .edgeLabel .state-cond {
+            color: var(--vscode-descriptionForeground) !important;
+            display: inline-block !important;
+            font-size: 0.82em !important;
+            line-height: 1.2 !important;
+        }
+        .diagram-container .mermaid .edgeLabel .state-cond-pre {
+            color: var(--lj-state-cond-pre) !important;
+        }
+        .diagram-container .mermaid .edgeLabel .state-cond-post {
+            color: var(--lj-state-cond-post) !important;
+        }
         .diagram-container .mermaid svg rect,
         .diagram-container .mermaid svg circle,
         .diagram-container .mermaid svg ellipse,

client/src/webview/views/fsm/fsm.ts

@@ -1,7 +1,15 @@
 import type { LJStateMachine } from "../../../types/fsm";
 import { renderMainHeader } from "../sections";
 
-export function renderStateMachineView(sm: LJStateMachine | undefined, diagram: string, orientation: "LR" | "TB"): string {
+export function renderStateMachineView(sm: LJStateMachine | undefined, diagram: string, orientation: "LR" | "TB", showConditions: boolean): string {
+    const initialStateNames = sm ? [...new Set(sm.initialTransitions.map(transition => transition.to))] : [];
+    const hasConditionExpansions = sm
+        ? sm.initialTransitions.some(transition => !!transition.postCond)
+            || sm.transitions.some(transition => !!transition.preCond || !!transition.postCond)
+        : false;
+    const conditionToggleLabel = showConditions ? 'Collapse Conditions' : 'Expand Conditions';
+    const conditionToggleIcon = showConditions ? '⊟' : '⊞';
+
     return /*html*/`
         <div>
             ${renderMainHeader("", 'fsm')}
@@ -13,6 +21,7 @@ export function renderStateMachineView(sm: LJStateMachine | undefined, diagram:
                             <button id="zoom-out-btn" class="diagram-control-btn" title="Zoom Out">-</button>
                             <button id="zoom-reset-btn" class="diagram-control-btn" title="Reset Zoom">⟲</button>
                             <button id="diagram-orientation-btn" class="diagram-control-btn" title="Rotate Diagram">${orientation === "TB" ? "↓" : "→"}</button>
+                            <button id="diagram-conditions-btn" class="diagram-control-btn${showConditions ? ' active' : ''}" title="${conditionToggleLabel}" aria-label="${conditionToggleLabel}" aria-pressed="${showConditions ? 'true' : 'false'}" ${hasConditionExpansions ? '' : 'disabled'}>${conditionToggleIcon}</button>
                             <button id="copy-diagram-btn" class="diagram-control-btn" title="Copy Mermaid Source">⎘</button>
                         </div>
                         <div id="diagram-wrapper" class="diagram-wrapper">
@@ -21,9 +30,9 @@ export function renderStateMachineView(sm: LJStateMachine | undefined, diagram:
                     </div>
                     <div>
                         <p><strong>States:</strong> ${sm.states.join(', ')}</p>
-                        <p><strong>Initial state${sm.initialStates.length > 1 ? 's' : ''}:</strong> ${sm.initialStates.join(', ')}</p>
+                        <p><strong>Initial state${initialStateNames.length > 1 ? 's' : ''}:</strong> ${initialStateNames.join(', ')}</p>
                         <p><strong>Number of states:</strong> ${sm.states.length}</p>
-                        <p><strong>Number of transitions:</strong> ${sm.transitions.length + 1}</p>
+                        <p><strong>Number of transitions:</strong> ${sm.transitions.length + sm.initialTransitions.length}</p>
                     </div>
                 </div>`
             : 'No state machine available for the current file'}

server/src/main/java/fsm/StateMachine.java

@@ -7,7 +7,7 @@
  */
 public record StateMachine(
     String className,
-    List<String> initialStates,
     List<String> states,
-    List<StateMachineTransition> transitions
+    List<StateMachineTransition> transitions,
+    List<StateMachineInitialTransition> initialTransitions
 ) { }

server/src/main/java/fsm/StateMachineInitialTransition.java

@@ -0,0 +1,11 @@
+package fsm;
+
+/**
+ * Represents an initial transition in a state machine
+ */
+public record StateMachineInitialTransition(String to, String postCond) {
+
+    public StateMachineInitialTransition(String to) {
+        this(to, null);
+    }
+}