Implement initial state edits

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

@@ -101,11 +101,26 @@ connection.onRequest(applyPatches, async ({ uri, machineIndex, patches }) => {
     patches,
   });
 
-  return edits.map(({ fileName, range, ...rest }) => {
+  return edits.map((edit) => {
+    if (edit.type === 'replace') {
+      return {
+        type: 'replace' as const,
+        uri: server.env.fileNameToUri(edit.fileName),
+        range: xstateProject.getLinesAndCharactersRange(
+          edit.fileName,
+          edit.range,
+        ),
+        newText: edit.newText,
+      };
+    }
     return {
-      ...rest,
-      uri: server.env.fileNameToUri(fileName),
-      range: xstateProject.getLinesAndCharactersRange(fileName, range),
+      type: edit.type,
+      uri: server.env.fileNameToUri(edit.fileName),
+      position: xstateProject.getLineAndCharacterOfPosition(
+        edit.fileName,
+        edit.position,
+      ),
+      newText: edit.newText,
     };
   });
 });

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

@@ -1,21 +1,24 @@
 import type {
   ExtractorDigraphDef,
+  LineAndCharacterPosition,
   LinesAndCharactersRange,
   Patch,
 } from '@xstate/ts-project';
 import * as vscode from 'vscode-languageserver-protocol';
 
-type DistributiveOmit<T, K extends PropertyKey> = T extends unknown
-  ? Omit<T, K>
-  : never;
-
-type TextEdit = DistributiveOmit<
-  import('@xstate/ts-project').TextEdit,
-  'fileName' | 'range'
-> & {
-  uri: string;
-  range: LinesAndCharactersRange;
-};
+type TextEdit =
+  | {
+      type: 'insert';
+      uri: string;
+      position: LineAndCharacterPosition;
+      newText: string;
+    }
+  | {
+      type: 'replace';
+      uri: string;
+      range: LinesAndCharactersRange;
+      newText: string;
+    };
 
 export const getMachineAtIndex = new vscode.RequestType<
   {

new-packages/ts-project/__tests__/source-edits/rename-state.test.ts

@@ -456,3 +456,47 @@ test('should rename a square bracket state name to an identifier', async () => {
     }
   `);
 });
+
+test('should adjust initial state of the parent when renaming a state', async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    'index.ts': ts`
+      import { createMachine } from "xstate";
+
+      createMachine({
+        initial: "foo",
+        states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'rename_state',
+      path: ['foo'],
+      name: 'NEW_NAME',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(`
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+      initial: "NEW_NAME",
+      states: {
+        NEW_NAME: {},
+        bar: {},
+      },
+    });",
+    }
+  `);
+});

new-packages/ts-project/__tests__/source-edits/set-initial-state.test.ts

@@ -0,0 +1,279 @@
+import { outdent } from 'outdent';
+import { expect, test } from 'vitest';
+import { createTestProject, testdir, ts } from '../utils';
+
+test("should override root's existing initial state", async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    'index.ts': ts`
+      import { createMachine } from "xstate";
+
+      createMachine({
+        initial: "foo",
+        states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: [],
+      initialState: 'bar',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(
+    `
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+      initial: "bar",
+      states: {
+        foo: {},
+        bar: {},
+      },
+    });",
+    }
+  `,
+  );
+});
+
+test("should override nested state's existing initial state", async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    'index.ts': ts`
+      import { createMachine } from "xstate";
+
+      createMachine({
+        initial: "foo",
+        states: {
+          foo: {
+            initial: "bar",
+            states: {
+              bar: {},
+              baz: {},
+            },
+          },
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: ['foo'],
+      initialState: 'baz',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(
+    `
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+      initial: "foo",
+      states: {
+        foo: {
+          initial: "baz",
+          states: {
+            bar: {},
+            baz: {},
+          },
+        },
+      },
+    });",
+    }
+  `,
+  );
+});
+
+test('should add initial state to root', async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    'index.ts': ts`
+      import { createMachine } from "xstate";
+
+      createMachine({
+        states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: [],
+      initialState: 'bar',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(`
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+      initial: "bar",
+      states: {
+        foo: {},
+        bar: {},
+      },
+    });",
+    }
+  `);
+});
+
+test("should add initial state before `states` property's comments", async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    'index.ts': ts`
+      import { createMachine } from "xstate";
+
+      createMachine({
+        // comment
+        states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: [],
+      initialState: 'bar',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(`
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+      initial: "bar",
+      // comment
+      states: {
+        foo: {},
+        bar: {},
+      },
+    });",
+    }
+  `);
+});
+
+test('should successfully add initial state before `states` property with no leading whitespace whatsoever', async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    // ignore Prettier here by using outdent
+    'index.ts': outdent`
+      import { createMachine } from "xstate";
+
+      createMachine({states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: [],
+      initialState: 'bar',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(`
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({initial: "bar",
+    states: {
+        foo: {},
+        bar: {},
+      },
+    });",
+    }
+  `);
+});
+
+test('should add initial state using `states` property indentation', async () => {
+  const tmpPath = await testdir({
+    'tsconfig.json': JSON.stringify({}),
+    // ignore Prettier here by using outdent
+    'index.ts': outdent`
+      import { createMachine } from "xstate";
+
+      createMachine({
+                    states: {
+          foo: {},
+          bar: {},
+        },
+      });
+    `,
+  });
+
+  const project = await createTestProject(tmpPath);
+
+  const textEdits = project.editDigraph(
+    {
+      fileName: 'index.ts',
+      machineIndex: 0,
+    },
+    {
+      type: 'set_initial_state',
+      path: [],
+      initialState: 'bar',
+    },
+  );
+  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot(`
+    {
+      "index.ts": "import { createMachine } from "xstate";
+
+    createMachine({
+                  initial: "bar",
+                  states: {
+        foo: {},
+        bar: {},
+      },
+    });",
+    }
+  `);
+});