Show one state machine diagram per state set

client/media/walkthrough/state-machine.md

@@ -1,4 +1,4 @@
-The **State Machine** tab shows object states and method transitions declared with `@StateSet` and `@StateRefinement`.
+The **State Machine** tab shows object states and method transitions declared with `@StateSet` and `@StateRefinement`. A class with several `@StateSet` annotations has one diagram per state set, one below the other.
 Pan the diagram; use **Zoom In**, **Zoom Out**, **Reset Zoom**, and **Toggle Orientation** to inspect it.
 Use **Expand Conditions** or **Collapse Conditions** to show or hide preconditions and postconditions. **Copy Mermaid Source** copies the diagram definition.
-From a state error, choose **Show in state machine** to see the related states and call. The tab opens on the first state error of the file, keeps the one you choose, and follows the error at the cursor; step through several with the arrows, and use **View error** to go back to it.
+From a state error, choose **Show in state machine** to see the related states and call, on the diagrams of the state sets the call involves. The tab opens on the first state error of the file, keeps the one you choose, and follows the error at the cursor; step through several with the arrows, and use **View error** to go back to it.

client/src/services/codelens.ts

@@ -40,7 +40,7 @@ function createDiagnosticCodeLenses(diagnostic: LJDiagnostic, file: string): vsc
         command: "liquidjava.showView",
         arguments: [targetDiagnostic]
     })];
-    if (diagnostic.type === 'state-refinement-error' && diagnostic.stateMachine) {
+    if (diagnostic.type === 'state-refinement-error' && diagnostic.stateMachines.length > 0) {
         codeLenses.push(new vscode.CodeLens(range, {
             title: "Show in state machine",
             command: "liquidjava.showView",

client/src/services/events.ts

@@ -36,7 +36,7 @@ export function registerEvents(context: vscode.ExtensionContext) {
  */
 export async function onActiveFileChange(editor: vscode.TextEditor) {
     extension.file = normalizeFilePath(editor.document.uri.fsPath);
-    extension.stateMachine = undefined;
+    extension.stateMachines = undefined;
     extension.webview?.sendMessage({ type: "file", file: extension.file });
     await updateStateMachine(editor.document);
     handleContextUpdate(editor.selection);

client/src/services/state-machine.ts

@@ -4,14 +4,14 @@ import { LJStateMachine } from "../types/fsm";
 import { normalizeFilePath } from "../utils/utils";
 
 /**
- * Requests the state machine for the given document from the language server
+ * Requests the state machines for the given document from the language server, one per state set
  * @param document The text document
  */
 export async function updateStateMachine(document: vscode.TextDocument) {
     const file = normalizeFilePath(document.uri.fsPath);
-    const sm = await extension.client?.sendRequest<LJStateMachine | null>("liquidjava/fsm", { uri: document.uri.toString() });
+    const sms = await extension.client?.sendRequest<LJStateMachine[] | null>("liquidjava/fsm", { uri: document.uri.toString() });
     if (file !== extension.file) return;
 
-    extension.stateMachine = sm;
-    extension.webview?.sendMessage({ type: "fsm", sm });
+    extension.stateMachines = sms;
+    extension.webview?.sendMessage({ type: "fsm", sms });
 }

client/src/services/webview.ts

@@ -39,7 +39,7 @@ export function registerWebview(context: vscode.ExtensionContext) {
                 if (extension.file) extension.webview?.sendMessage({ type: "file", file: extension.file });
                 if (extension.diagnostics) extension.webview?.sendMessage({ type: "diagnostics", diagnostics: extension.diagnostics });
                 if (extension.context) extension.webview?.sendMessage({ type: "context", context: extension.context , errorAtCursor: extension.errorAtCursor });
-                if (extension.stateMachine !== undefined) extension.webview?.sendMessage({ type: "fsm", sm: extension.stateMachine });
+                if (extension.stateMachines !== undefined) extension.webview?.sendMessage({ type: "fsm", sms: extension.stateMachines });
                 if (extension.status) extension.webview?.sendMessage({ type: "status", status: extension.status });
                 if (pendingDiagnosticReveal) {
                     extension.webview?.sendMessage({ type: "revealDiagnostic", ...pendingDiagnosticReveal });

client/src/state.ts

@@ -27,7 +27,7 @@ export class ExtensionState {
     diagnostics?: LJDiagnostic[];
     readonly diagnosticsEmitter = new vscode.EventEmitter<LJDiagnostic[]>();
     readonly failureEmitter = new vscode.EventEmitter<void>();
-    stateMachine?: LJStateMachine | null;
+    stateMachines?: LJStateMachine[] | null;
     context?: LJContext;
     currentSelection?: Range;
     errorAtCursor?: RefinementMismatchError;

client/src/test/state-errors.test.ts

@@ -47,8 +47,8 @@ suite('Bundled LiquidJava state machine errors', () => {
             await ready;
 
             const machine = nextWebviewMessage('the refinements file state machine', event =>
-                event.direction === 'toWebview' && event.message.type === 'fsm' && event.message.sm
-                    ? event.message.sm as LJStateMachine : undefined);
+                event.direction === 'toWebview' && event.message.type === 'fsm' && event.message.sms?.length
+                    ? (event.message.sms as LJStateMachine[])[0] : undefined);
             const stateError = nextWebviewMessage('the state error diagnostics', event =>
                 event.direction === 'toWebview' && event.message.type === 'diagnostics'
                     ? (event.message.diagnostics as LJDiagnostic[]).find(isStateError) : undefined);
@@ -58,9 +58,9 @@ suite('Bundled LiquidJava state machine errors', () => {
 
             const [sm, error] = await Promise.all([machine, stateError]);
             assert.equal(sm.className, 'javax.swing.tree.DefaultMutableTreeNode');
-            assert.ok(error.stateMachine, 'the state error must carry its state machine');
-            assert.equal(error.stateMachine.className, sm.className, 'errors are matched to the machine by class');
-            assert.ok(error.stateMachine.errorContext, 'the state error must carry its error context');
+            assert.equal(error.stateMachines.length, 1, 'the state error must carry the state machine of each state set');
+            assert.equal(error.stateMachines[0].className, sm.className, 'errors are matched to the machine by class');
+            assert.ok(error.stateMachines[0].errorContext, 'the state error must carry its error context');
             assert.ok(error.position && sameFile(error.position.file ?? error.file, callerUri));
             assert.ok(error.declarationPosition?.file && sameFile(error.declarationPosition.file, refinementsUri));
 

client/src/types/diagnostics.ts

@@ -99,7 +99,8 @@ export type StateRefinementError = BaseDiagnostic & {
     found: VCSimplificationResult;
     customMessage: string;
     declarationPosition: SourcePosition | null;
-    stateMachine: LJStateMachine | null;
+    // one state machine per state set of the receiver's class
+    stateMachines: LJStateMachine[];
 }
 
 export type ArgumentMismatchError = BaseDiagnostic & {

client/src/types/fsm.ts

@@ -1,11 +1,15 @@
 // Type definitions used for representing finite state machines
 
+// The state machine of one state set of a class
 export type LJStateMachine = {
     className: string;
+    // index of the state set among the class's @StateSet annotations, in declaration order
+    stateSet?: number;
     initialTransitions: { to: string; toCondition?: string | null; constructorSignature?: string | null }[];
     states: string[];
     transitions: { from: string; to: string; label: string; fromCondition?: string | null; toCondition?: string | null }[];
-    errorContext: LJStateMachineErrorContext | null;
+    // the state error on this state set, null if the call does not involve it
+    errorContext?: LJStateMachineErrorContext | null;
 };
 
 // The verifier's states of a rejected call, on the diagram's states

client/src/webview/diagram-focus.spec.ts

@@ -1,7 +1,7 @@
 import { afterEach, describe, expect, it, vi } from 'vitest';
 import type { LJStateMachine, LJStateMachineErrorContext } from '../types/fsm';
 import { createMermaidDiagram, renderMermaidDiagram } from './diagram';
-import { getRelatedStates, getStateErrorFocus } from './diagram-focus';
+import { getRelatedStates, getStateErrorFocus, getStateErrorFocuses, orderStateMachines } from './diagram-focus';
 
 // DefaultMutableTreeNode from the state error focus issue:
 // root.add(...); root.removeAllChildren(); root.getFirstChild(); // state error
@@ -250,3 +250,19 @@ describe('rendered state error focus', () => {
         expect(dimmed('g.label.lj-dimmed')).toEqual(['edge0', 'edge1', 'edge12']);
     });
 });
+
+describe('state errors on several state sets', () => {
+    const tiling: LJStateMachine = { ...treeNode(null), className: 'ImageWriteParam', stateSet: 0 };
+    const compression: LJStateMachine = { ...treeNode(), className: 'ImageWriteParam', stateSet: 1 };
+    const other: LJStateMachine = { ...treeNode(null), className: 'ImageWriteParam', stateSet: 2 };
+
+    it('puts the state sets with the error first, keeping the declaration order otherwise', () => {
+        expect(orderStateMachines([tiling, compression, other])).toEqual([compression, tiling, other]);
+        expect(orderStateMachines([tiling, other])).toEqual([tiling, other]);
+    });
+
+    it('focuses on the error of each state set that has it', () => {
+        expect(getStateErrorFocuses([tiling, compression, other])).toEqual([getStateErrorFocus(compression)]);
+        expect(getStateErrorFocuses([tiling, other])).toEqual([]);
+    });
+});

client/src/webview/diagram-focus.ts

@@ -54,10 +54,28 @@ export function getRelatedStates(focus: StateErrorFocus): string[] {
 
 /**
  * Gets the text shown for the rejected call, e.g. root.getFirstChild()
- * @param focus
+ * @param call the called method and receiver of the state error
  * @returns the call text, or null when the called method is unknown
  */
-export function getRejectedCallText(focus: StateErrorFocus): string | null {
-    if (!focus.calledMethod) return null;
-    return `${focus.receiver ? `${focus.receiver}.` : ''}${focus.calledMethod}()`;
+export function getRejectedCallText(call: Pick<StateErrorFocus, 'calledMethod' | 'receiver'>): string | null {
+    if (!call.calledMethod) return null;
+    return `${call.receiver ? `${call.receiver}.` : ''}${call.calledMethod}()`;
+}
+
+/**
+ * Orders the state machines of a class so that those of the state sets with the state error come first
+ * @param sms the state machines, one per state set, in declaration order
+ * @returns the state machines with the state error, then the others, each in declaration order
+ */
+export function orderStateMachines(sms: LJStateMachine[]): LJStateMachine[] {
+    return [...sms.filter(sm => sm.errorContext), ...sms.filter(sm => !sm.errorContext)];
+}
+
+/**
+ * Gets the state error to focus on in each state machine that has one
+ * @param sms the state machines, one per state set
+ * @returns the state error focus of each state set with the error
+ */
+export function getStateErrorFocuses(sms: LJStateMachine[]): StateErrorFocus[] {
+    return sms.map(getStateErrorFocus).filter((focus): focus is StateErrorFocus => !!focus);
 }

client/src/webview/script.spec.ts

@@ -63,7 +63,7 @@ function stateError(title: string, sm: LJStateMachine, line = position.lineStart
         title, message: `${title} message`, hints: [], file, position: errorPosition,
         category: 'error', type: 'state-refinement-error', translationTable: {}, expected: 'Open(counter)',
         found: { implication: null as never, origin: null, simplification: 'Closed(counter)' },
-        customMessage: '', declarationPosition: null, stateMachine: sm,
+        customMessage: '', declarationPosition: null, stateMachines: [sm],
     };
 }
 
@@ -142,7 +142,7 @@ describe('webview script', () => {
     it('switches between diagnostics, context, and state machine tabs', async () => {
         showDiagnostics([diagnostic(currentFile, 'Current error')]);
         receive({ type: 'context', context: context() });
-        receive({ type: 'fsm', sm: stateMachine });
+        receive({ type: 'fsm', sms: [stateMachine] });
         postMessage.mockClear();
 
         click('[data-tab="context"]');
@@ -194,7 +194,7 @@ describe('webview script', () => {
 
         // the tab keeps the error after visiting another tab
         click('[data-tab="diagnostics"]');
-        receive({ type: 'fsm', sm: stateMachine });
+        receive({ type: 'fsm', sms: [stateMachine] });
         click('[data-tab="fsm"]');
         expect(root.querySelector('#diagram-focus-btn')?.textContent?.trim()).toBe('Show all');
         expect(root.querySelector('.mermaid')?.textContent).toContain('class Broken ljStateDimmed');
@@ -216,11 +216,35 @@ describe('webview script', () => {
         await vi.waitFor(() => expect(transform()).toBe('matrix(1, 0, 0, 1, 0, 0)'));
     });
 
+    it('shows the diagram of the state set with the error first and copies every diagram', async () => {
+        const writeText = vi.fn().mockResolvedValue(undefined);
+        vi.stubGlobal('navigator', { clipboard: { writeText } });
+        const other: LJStateMachine = {
+            ...stateMachine, className: 'Counter', stateSet: 0, states: ['Idle', 'Busy'],
+            transitions: [{ from: 'Idle', to: 'Busy', label: 'start()' }],
+        };
+        const withError = { ...errorMachine('Counter', 'close'), stateSet: 1 };
+        showDiagnostics([{ ...stateError('State refinement error', withError), stateMachines: [other, withError] } as LJDiagnostic]);
+        click('.diagnostic-state-machine-btn');
+
+        await vi.waitFor(() => expect(runMermaid).toHaveBeenCalledTimes(1));
+        expect(runMermaid.mock.calls[0][0].nodes).toHaveLength(2);
+        expect(Array.from(root.querySelectorAll('.diagram-state-set-title'), node => node.textContent))
+            .toEqual(['Open \u00B7 Closed \u00B7 Broken', 'Idle \u00B7 Busy']);
+        expect(root.querySelector('.diagram-state-set:last-child')?.classList.contains('diagram-state-set-dimmed')).toBe(true);
+
+        click('#copy-diagram-btn');
+        const copied: string = writeText.mock.calls[0][0];
+        expect(copied.match(/stateDiagram-v2/g)).toHaveLength(2);
+        expect(copied).toContain('Idle --> Busy');
+        expect(copied).not.toMatch(/ljState|lj-dimmed|lj-state-error-method/);
+    });
+
     it('copies the Mermaid source without the error overlay', async () => {
         const writeText = vi.fn().mockResolvedValue(undefined);
         vi.stubGlobal('navigator', { clipboard: { writeText } });
         showDiagnostics([]);
-        receive({ type: 'fsm', sm: { ...errorMachine('Counter', 'close'), errorContext: null } });
+        receive({ type: 'fsm', sms: [{ ...errorMachine('Counter', 'close'), errorContext: null }] });
         click('[data-tab="fsm"]');
         click('#copy-diagram-btn');
         const plain = writeText.mock.calls[0][0];
@@ -296,7 +320,7 @@ describe('webview script', () => {
             stateError('Second state error', errorMachine('Second', 'reset'), 9),
             stateError('First state error', errorMachine('First', 'close'), 4),
         ]);
-        receive({ type: 'fsm', sm: stateMachine });
+        receive({ type: 'fsm', sms: [stateMachine] });
         click('[data-tab="fsm"]');
 
         expect(root.querySelector('.diagram-title')?.textContent).toBe('First');