Implement initial state insertions before existing states property (#467

)
statelyai · Feb 13, 2024 · 2e5a530 · 2e5a530
1 parent 3bf6f3a
commit 2e5a530
Show file tree

Hide file tree

Showing 8 changed files with 427 additions and 46 deletions.
diff --git a/new-packages/language-server/src/index.ts b/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,
     };
   });
 });

diff --git a/new-packages/language-server/src/protocol.ts b/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<
   {

diff --git a/new-packages/ts-project/__tests__/source-edits/set-initial-state.test.ts b/new-packages/ts-project/__tests__/source-edits/set-initial-state.test.ts
@@ -1,3 +1,4 @@
+import { outdent } from 'outdent';
 import { expect, test } from 'vitest';
 import { createTestProject, testdir, ts } from '../utils';
 
@@ -103,7 +104,7 @@ test("should override nested state's existing initial state", async () => {
   );
 });
 
-test.todo('should add initial state to root', async () => {
+test('should add initial state to root', async () => {
   const tmpPath = await testdir({
     'tsconfig.json': JSON.stringify({}),
     'index.ts': ts`
@@ -131,5 +132,148 @@ test.todo('should add initial state to root', async () => {
       initialState: 'bar',
     },
   );
-  expect(await project.applyTextEdits(textEdits)).toMatchInlineSnapshot();
+  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: {},
+      },
+    });",
+    }
+  `);
 });
diff --git a/new-packages/ts-project/__tests__/utils.ts b/new-packages/ts-project/__tests__/utils.ts
@@ -357,15 +357,24 @@ export async function createTestProject(
     applyTextEdits: async (edits: readonly TextEdit[]) => {
       const edited: Record<string, string> = {};
 
-      for (const edit of [...edits].sort(
-        (a, b) => b.range.start - a.range.start,
-      )) {
+      for (const edit of [...edits].sort((a, b) => {
+        const startA = 'range' in a ? a.range.start : a.position;
+        const startB = 'range' in b ? b.range.start : b.position;
+        return startB - startA;
+      })) {
+        const relativeFileName = path.relative(cwd, edit.fileName);
+        const source =
+          edited[relativeFileName] ??
+          program.getSourceFile(edit.fileName)!.text;
         switch (edit.type) {
+          case 'insert':
+            edited[relativeFileName] =
+              source.slice(0, edit.position) +
+              edit.newText +
+              source.slice(edit.position);
+            break;
           case 'replace':
-            const source =
-              edited[edit.fileName] ??
-              program.getSourceFile(edit.fileName)!.text;
-            edited[edit.fileName] =
+            edited[relativeFileName] =
               source.slice(0, edit.range.start) +
               edit.newText +
               source.slice(edit.range.end);