wire up things for apply patches

new-packages/language-server/src/index.ts

@@ -9,7 +9,7 @@ import {
   Provide,
   create as createTypeScriptService,
 } from 'volar-service-typescript';
-import { getMachineAtIndex } from './protocol';
+import { applyPatches, getMachineAtIndex } from './protocol';
 
 const projectCache = new WeakMap<Program, XStateProject>();
 
@@ -81,6 +81,25 @@ connection.onRequest(getMachineAtIndex, async ({ uri, machineIndex }) => {
   return digraph;
 });
 
+connection.onRequest(applyPatches, async ({ uri, machineIndex, patches }) => {
+  const xstateProject = await getXStateProject(uri);
+
+  if (!xstateProject) {
+    return [];
+  }
+
+  const edits = xstateProject.applyPatches({
+    fileName: server.env.uriToFileName(uri),
+    machineIndex,
+    patches,
+  });
+
+  return edits.map(({ fileName, ...rest }) => ({
+    ...rest,
+    uri: server.env.fileNameToUri(fileName),
+  }));
+});
+
 connection.onInitialized(() => {
   server.initialized();
 });

new-packages/language-server/src/protocol.ts

@@ -1,6 +1,10 @@
-import { ExtractorDigraphDef } from '@xstate/ts-project';
+import type { ExtractorDigraphDef, Patch, TextEdit } from '@xstate/ts-project';
 import * as vscode from 'vscode-languageserver-protocol';
 
+type DistributiveOmit<T, K extends PropertyKey> = T extends unknown
+  ? Omit<T, K>
+  : never;
+
 export const getMachineAtIndex = new vscode.RequestType<
   {
     uri: string;
@@ -9,3 +13,13 @@ export const getMachineAtIndex = new vscode.RequestType<
   ExtractorDigraphDef | undefined,
   never
 >('stately-xstate/get-machine-at-index');
+
+export const applyPatches = new vscode.RequestType<
+  {
+    uri: string;
+    machineIndex: number;
+    patches: Patch[];
+  },
+  (DistributiveOmit<TextEdit, 'fileName'> & { uri: string })[],
+  never
+>('stately-xstate/apply-patches');

new-packages/ts-project/package.json

@@ -18,7 +18,9 @@
     "./package.json": "./package.json"
   },
   "scripts": {},
-  "dependencies": {},
+  "dependencies": {
+    "immer": "^10.0.3"
+  },
   "devDependencies": {
     "@types/fs-extra": "^11.0.4",
     "fs-extra": "^11.2.0",

new-packages/ts-project/src/index.ts

@@ -1,13 +1,15 @@
+import type { Patch } from 'immer';
 import type { CallExpression, Program, SourceFile } from 'typescript';
 import { extractState } from './state';
 import type {
   ExtractionContext,
   ExtractionError,
   ExtractorDigraphDef,
+  Range,
+  TextEdit,
   TreeNode,
   XStateVersion,
 } from './types';
-export { ExtractorDigraphDef };
 
 function findCreateMachineCalls(
   ts: typeof import('typescript'),
@@ -139,22 +141,12 @@ export interface TSProjectOptions {
   xstateVersion?: XStateVersion | undefined;
 }
 
-interface Position {
-  line: number;
-  character: number;
-}
-
-interface Range {
-  start: Position;
-  end: Position;
-}
-
 export function createProject(
   ts: typeof import('typescript'),
   tsProgram: Program,
   { xstateVersion = '5' }: TSProjectOptions = {},
 ) {
-  return {
+  const api = {
     findMachines: (fileName: string): Range[] => {
       const sourceFile = tsProgram.getSourceFile(fileName);
       if (!sourceFile) {
@@ -197,7 +189,24 @@ export function createProject(
         return extractMachineConfig(ctx, ts, call);
       });
     },
+    // TODO: consider exposing an object representing a file or a machine to the caller of the `extractMachines` and add a similar-ish method there
+    // for now we are just doing through the full extraction process again which is wasteful
+    applyPatches({
+      fileName,
+      machineIndex,
+      patches,
+    }: {
+      fileName: string;
+      machineIndex: number;
+      patches: Patch[];
+    }): TextEdit[] {
+      const extractedMachine = api.extractMachines(fileName)[machineIndex];
+      console.log(extractedMachine, patches);
+      return [];
+    },
   };
+  return api;
 }
 
 export type XStateProject = ReturnType<typeof createProject>;
+export { ExtractorDigraphDef, Patch, TextEdit };

new-packages/ts-project/src/types.ts

@@ -205,3 +205,22 @@ export type JsonValue =
   | JsonObject
   | JsonValue[];
 export type JsonObject = { [key: string]: JsonValue };
+
+export interface Position {
+  line: number;
+  character: number;
+}
+
+export interface Range {
+  start: Position;
+  end: Position;
+}
+
+interface ReplaceTextEdit {
+  type: 'replace';
+  fileName: string;
+  range: Range;
+  newText: string;
+}
+
+export type TextEdit = ReplaceTextEdit;

new-packages/vscode-xstate/fixtures/basic/src/index.ts

@@ -6,8 +6,10 @@ createMachine({
     a: {
       on: {
         NEXT: 'b',
+        kappa: 'c',
       },
     },
     b: {},
+    c: {},
   },
 });

new-packages/vscode-xstate/src/languageClient.ts

@@ -1,7 +1,10 @@
 import { RequestType, getTsdk } from '@volar/vscode';
 import { LanguageClient, TransportKind } from '@volar/vscode/node.js';
-import { getMachineAtIndex } from '@xstate/language-server/protocol';
-import type { ExtractorDigraphDef } from '@xstate/ts-project';
+import {
+  applyPatches,
+  getMachineAtIndex,
+} from '@xstate/language-server/protocol';
+import type { ExtractorDigraphDef, Patch } from '@xstate/ts-project';
 import * as vscode from 'vscode';
 import {
   ActorRef,
@@ -20,10 +23,21 @@ type LanguageClientInput = {
   tsdk: Awaited<ReturnType<typeof getTsdk>>['tsdk'];
   isServerEnabled: boolean;
 };
+
+export type OpenMachineData = {
+  uri: string;
+  machineIndex: number;
+};
+
 export type LanguageClientEvent =
   | { type: 'SERVER_ENABLED_CHANGE' }
   | { type: 'WEBVIEW_CLOSED' }
-  | { type: 'OPEN_MACHINE'; digraph: ExtractorDigraphDef };
+  | ({ type: 'OPEN_MACHINE'; digraph: ExtractorDigraphDef } & OpenMachineData)
+  | ({
+      type: 'APPLY_PATCHES';
+      patches: Patch[];
+      reason: 'undo' | 'redo' | undefined;
+    } & OpenMachineData);
 
 export const languageClientMachine = setup({
   types: {} as {
@@ -41,11 +55,39 @@ export const languageClientMachine = setup({
     stopLanguageClient: ({ context }) => context.languageClient.stop(),
     updateDigraph: sendTo('webview', ({ event }) => {
       assertEvent(event, 'OPEN_MACHINE');
-      return {
-        type: 'UPDATE_DIGRAPH',
-        digraph: event.digraph,
-      };
+      return event;
     }),
+    applyPatches: ({ context, event }) => {
+      assertEvent(event, 'APPLY_PATCHES');
+      const { type, ...params } = event;
+      context.languageClient
+        .sendRequest(applyPatches, params)
+        .then((textEdits) => {
+          const workspaceEdit = new vscode.WorkspaceEdit();
+          for (const textEdit of textEdits) {
+            switch (textEdit.type) {
+              case 'replace': {
+                workspaceEdit.replace(
+                  vscode.Uri.parse(params.uri),
+                  new vscode.Range(
+                    new vscode.Position(
+                      textEdit.range.start.line,
+                      textEdit.range.start.character,
+                    ),
+                    new vscode.Position(
+                      textEdit.range.end.line,
+                      textEdit.range.end.character,
+                    ),
+                  ),
+                  textEdit.newText,
+                );
+                break;
+              }
+            }
+          }
+          return vscode.workspace.applyEdit(workspaceEdit);
+        });
+    },
   },
   actors: {
     listenOnServerEnabledChange: fromCallback(
@@ -93,6 +135,8 @@ export const languageClientMachine = setup({
               }
               parent.send({
                 type: 'OPEN_MACHINE',
+                uri,
+                machineIndex,
                 digraph,
               });
             });
@@ -202,10 +246,11 @@ export const languageClientMachine = setup({
                 id: 'webview',
                 input: ({ context, event, self }) => {
                   assertEvent(event, 'OPEN_MACHINE');
+                  const { type, ...params } = event;
                   return {
                     extensionContext: context.extensionContext,
                     parent: self,
-                    digraph: event.digraph,
+                    ...params,
                   };
                 },
               },
@@ -214,6 +259,9 @@ export const languageClientMachine = setup({
                 OPEN_MACHINE: {
                   actions: 'updateDigraph',
                 },
+                APPLY_PATCHES: {
+                  actions: 'applyPatches',
+                },
               },
             },
           },

new-packages/vscode-xstate/src/webview.ts

@@ -1,7 +1,7 @@
 import type { ExtractorDigraphDef } from '@xstate/ts-project';
 import * as vscode from 'vscode';
-import { ActorRef, Snapshot, fromCallback } from 'xstate';
-import { LanguageClientEvent } from './languageClient';
+import { ActorRef, ExtractEvent, Snapshot, fromCallback } from 'xstate';
+import { LanguageClientEvent, OpenMachineData } from './languageClient';
 
 // this should not be an async function, it should return `webviewPanel` synchronously
 // to allow the consumer to start listening to events from the webview immediately
@@ -26,11 +26,9 @@ function createWebviewPanel() {
 async function getWebviewHtml(
   extensionContext: vscode.ExtensionContext,
   webviewPanel: vscode.WebviewPanel,
-  {
-    digraph,
-  }: {
+  params: {
     digraph: ExtractorDigraphDef;
-  },
+  } & OpenMachineData,
 ) {
   const bundledEditorRootUri = vscode.Uri.joinPath(
     vscode.Uri.file(extensionContext.extensionPath),
@@ -58,27 +56,27 @@ async function getWebviewHtml(
             : 'light'
           : theme,
       distinctId: `vscode:${vscode.env.machineId}`,
-      digraph,
+      ...params,
     },
   )}</script>`;
 
   return htmlContent.replace('<head>', `<head>${baseTag}${initialDataScript}`);
 }
 
 export const webviewLogic = fromCallback<
-  { type: 'UPDATE_DIGRAPH'; digraph: ExtractorDigraphDef },
+  ExtractEvent<LanguageClientEvent, 'OPEN_MACHINE'>,
   {
     extensionContext: vscode.ExtensionContext;
     parent: ActorRef<Snapshot<unknown>, LanguageClientEvent>;
     digraph: ExtractorDigraphDef;
-  }
->(({ input: { extensionContext, parent, digraph }, receive }) => {
+  } & OpenMachineData
+>(({ input: { extensionContext, parent, ...initialParams }, receive }) => {
   let canceled = false;
   const webviewPanel = createWebviewPanel();
 
   receive((event) => {
     switch (event.type) {
-      case 'UPDATE_DIGRAPH':
+      case 'OPEN_MACHINE':
         webviewPanel.reveal(vscode.ViewColumn.Beside);
         webviewPanel.webview.postMessage(event);
         return;
@@ -88,18 +86,21 @@ export const webviewLogic = fromCallback<
   });
 
   const disposable = vscode.Disposable.from(
-    webviewPanel.webview.onDidReceiveMessage((event: LanguageClientEvent) =>
-      parent.send(event),
+    webviewPanel.webview.onDidReceiveMessage(
+      (event: ExtractEvent<LanguageClientEvent, 'APPLY_PATCHES'>) =>
+        parent.send(event),
     ),
     webviewPanel.onDidDispose(() => {
       parent.send({ type: 'WEBVIEW_CLOSED' });
     }),
   );
 
   (async () => {
-    const html = await getWebviewHtml(extensionContext, webviewPanel, {
-      digraph,
-    });
+    const html = await getWebviewHtml(
+      extensionContext,
+      webviewPanel,
+      initialParams,
+    );
     if (canceled) {
       return;
     }

yarn.lock

@@ -4811,6 +4811,11 @@ ignore@^5.2.4:
   resolved "https://registry.yarnpkg.com/ignore/-/ignore-5.3.0.tgz#67418ae40d34d6999c95ff56016759c718c82f78"
   integrity sha512-g7dmpshy+gD7mh88OC9NwSGTKoc3kyLAZQRU1mt53Aw/vnvfXnbC+F/7F7QoYVKbV+KNvJx8wArewKy1vXMtlg==
 
+immer@^10.0.3:
+  version "10.0.3"
+  resolved "https://registry.yarnpkg.com/immer/-/immer-10.0.3.tgz#a8de42065e964aa3edf6afc282dfc7f7f34ae3c9"
+  integrity sha512-pwupu3eWfouuaowscykeckFmVTpqbzW+rXFCX8rQLkZzM9ftBmU/++Ra+o+L27mz03zJTlyV4UUr+fdKNffo4A==
+
 import-fresh@^3.2.1:
   version "3.3.0"
   resolved "https://registry.yarnpkg.com/import-fresh/-/import-fresh-3.3.0.tgz#37162c25fcb9ebaa2e6e53d5b4d88ce17d9e0c2b"