Implement initial state edits

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,135 @@
+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.todo('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();
+});

new-packages/ts-project/__tests__/utils.ts

@@ -19,6 +19,15 @@ export const js = outdent;
 export const ts = outdent;
 export const tsx = outdent;
 
+function shuffle<T>(arr: T[]): T[] {
+  const copy = [...arr];
+  for (let i = copy.length - 1; i > 0; i--) {
+    const j = Math.floor(Math.random() * (i + 1));
+    [copy[i], copy[j]] = [copy[j], copy[i]];
+  }
+  return copy;
+}
+
 type Fixture = {
   [key: string]: string | { kind: 'symlink'; path: string };
 };
@@ -141,7 +150,7 @@ type MachineEdit =
   | {
       type: 'set_initial_state';
       path: string[];
-      initialState: string | null;
+      initialState: string | undefined;
     }
   | { type: 'set_state_id'; path: string[]; id: string | null }
   | {
@@ -293,12 +302,31 @@ export async function createTestProject(
           case 'add_state':
           case 'remove_state':
             throw new Error(`Not implemented: ${edit.type}`);
-          case 'rename_state':
+          case 'rename_state': {
             const node = findNodeByStatePath(digraphDraft, edit.path);
+            const oldName = node.data.key;
+
             node.data.key = edit.name;
+
+            const parentNode = findNodeByStatePath(
+              digraphDraft,
+              edit.path.slice(0, -1),
+            );
+
+            // TODO: it would be great if `.initial` could be a uniqueId and not a resolved value
+            // that would have to be changed in the Studio and adjusted in this package
+            if (parentNode.data.initial === oldName) {
+              parentNode.data.initial = edit.name;
+            }
             break;
+          }
           case 'reparent_state':
-          case 'set_initial_state':
+            throw new Error(`Not implemented: ${edit.type}`);
+          case 'set_initial_state': {
+            const node = findNodeByStatePath(digraphDraft, edit.path);
+            node.data.initial = edit.initialState;
+            break;
+          }
           case 'set_state_id':
           case 'set_state_type':
           case 'add_transition':
@@ -319,13 +347,18 @@ export async function createTestProject(
             throw new Error(`Not implemented: ${edit.type}`);
         }
       });
-      return project.applyPatches({ fileName, machineIndex, patches });
+      return project.applyPatches({
+        fileName,
+        machineIndex,
+        // shuffle patches to make sure the order doesn't matter
+        patches: shuffle(patches),
+      });
     },
     applyTextEdits: async (edits: readonly TextEdit[]) => {
       const edited: Record<string, string> = {};
 
       for (const edit of [...edits].sort(
-        (a, b) => a.range.start - b.range.start,
+        (a, b) => b.range.start - a.range.start,
       )) {
         switch (edit.type) {
           case 'replace':

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

@@ -5,6 +5,7 @@ import type {
   PropertyAssignment,
   SourceFile,
 } from 'typescript';
+import { safeStringLikeLiteralText } from './safeStringLikeLiteralText';
 import { extractState } from './state';
 import type {
   ExtractionContext,
@@ -20,6 +21,7 @@ import type {
 import {
   assert,
   findNodeByAstPath,
+  findProperty,
   getPreferredQuoteCharCode,
   isValidIdentifier,
   safePropertyNameString,
@@ -284,6 +286,39 @@ function createProjectMachine({
                           getPreferredQuoteCharCode(host.ts, sourceFile),
                         ),
                   });
+                  break;
+                }
+                if (patch.path[2] === 'data' && patch.path[3] === 'initial') {
+                  if (typeof patch.value === undefined) {
+                    // removing initial states is not supported in the Studio
+                    // but a patch like this can likely still be received when the last child of a state gets removed
+                    break;
+                  }
+                  const node = findNodeByAstPath(
+                    host.ts,
+                    createMachineCall,
+                    currentState.astPaths.nodes[nodeId],
+                  );
+                  assert(host.ts.isObjectLiteralExpression(node));
+                  const prop = findProperty(
+                    undefined,
+                    host.ts,
+                    node,
+                    'initial',
+                  );
+                  assert(prop && host.ts.isPropertyAssignment(prop));
+                  edits.push({
+                    type: 'replace',
+                    fileName,
+                    range: {
+                      start: prop.initializer.getStart(),
+                      end: prop.initializer.getEnd(),
+                    },
+                    newText: safeStringLikeLiteralText(
+                      patch.value,
+                      getPreferredQuoteCharCode(host.ts, sourceFile),
+                    ),
+                  });
                 }
             }
             break;

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

@@ -42,7 +42,7 @@ export const uniqueId = () => {
 };
 
 function getLiteralText(
-  ctx: ExtractionContext,
+  ctx: ExtractionContext | undefined,
   ts: typeof import('typescript'),
   node: Expression,
 ) {
@@ -54,7 +54,7 @@ function getLiteralText(
     // for big ints this loses precision or might even return `'Infinity'`
     const text = node.getText();
     if (text !== node.text) {
-      ctx.errors.push({
+      ctx?.errors.push({
         type: 'property_key_no_roundtrip',
       });
     }
@@ -63,7 +63,7 @@ function getLiteralText(
 }
 
 export function getPropertyKey(
-  ctx: ExtractionContext,
+  ctx: ExtractionContext | undefined,
   ts: typeof import('typescript'),
   prop: PropertyAssignment,
 ) {
@@ -78,14 +78,14 @@ export function getPropertyKey(
     if (typeof text === 'string') {
       return text;
     }
-    ctx.errors.push({
+    ctx?.errors.push({
       type: 'property_key_unhandled',
       propertyKind: 'computed',
     });
     return;
   }
   if (ts.isPrivateIdentifier(prop.name)) {
-    ctx.errors.push({
+    ctx?.errors.push({
       type: 'property_key_unhandled',
       propertyKind: 'private',
     });
@@ -179,7 +179,7 @@ export function everyDefined<T>(arr: T[]): arr is NonNullable<T>[] {
 }
 
 export function findProperty(
-  ctx: ExtractionContext,
+  ctx: ExtractionContext | undefined,
   ts: typeof import('typescript'),
   obj: ObjectLiteralExpression,
   key: string,