@typespec/mutator-framework 0.12.0

package/src/mutation-node/model.ts

@@ -0,0 +1,89 @@
+import type { Model, ModelProperty, Type } from "@typespec/compiler";
+import { MutationEdge } from "./mutation-edge.js";
+import { MutationNode } from "./mutation-node.js";
+
+export class ModelMutationNode extends MutationNode<Model> {
+  readonly kind = "Model";
+
+  traverse() {
+    if (this.sourceType.baseModel) {
+      const baseNode = this.subgraph.getNode(this.sourceType.baseModel);
+      this.connectToBase(baseNode);
+    }
+
+    for (const [propName, prop] of this.sourceType.properties) {
+      const propNode = this.subgraph.getNode(prop);
+      this.connectProperty(propNode, propName);
+    }
+
+    if (this.sourceType.indexer) {
+      const indexerNode = this.subgraph.getNode(this.sourceType.indexer.value);
+      this.connectIndexerValue(indexerNode);
+    }
+  }
+
+  connectToBase(baseNode: MutationNode<Model>) {
+    MutationEdge.create(this, baseNode, {
+      onTailMutation: () => {
+        this.mutatedType!.baseModel = baseNode.mutatedType;
+      },
+      onTailDeletion: () => {
+        this.mutatedType.baseModel = undefined;
+      },
+      onTailReplaced: (newTail) => {
+        if (newTail.mutatedType.kind !== "Model") {
+          throw new Error("Cannot replace base model with non-model type");
+        }
+        this.mutatedType.baseModel = newTail.mutatedType;
+      },
+    });
+  }
+
+  connectProperty(propNode: MutationNode<ModelProperty>, sourcePropName: string) {
+    MutationEdge.create(this, propNode, {
+      onTailMutation: () => {
+        this.mutatedType.properties.delete(sourcePropName);
+        this.mutatedType.properties.set(propNode.mutatedType.name, propNode.mutatedType);
+      },
+      onTailDeletion: () => {
+        this.mutatedType.properties.delete(sourcePropName);
+      },
+      onTailReplaced: (newTail) => {
+        if (newTail.mutatedType.kind !== "ModelProperty") {
+          throw new Error("Cannot replace model property with non-model property type");
+        }
+        this.mutatedType.properties.delete(sourcePropName);
+        this.mutatedType.properties.set(newTail.mutatedType.name, newTail.mutatedType);
+      },
+    });
+  }
+
+  connectIndexerValue(indexerNode: MutationNode<Type>) {
+    MutationEdge.create(this, indexerNode, {
+      onTailMutation: () => {
+        if (this.mutatedType.indexer) {
+          this.mutatedType.indexer = {
+            key: this.mutatedType.indexer.key,
+            value: indexerNode.mutatedType,
+          };
+        }
+      },
+      onTailDeletion: () => {
+        if (this.mutatedType.indexer) {
+          this.mutatedType.indexer = {
+            key: this.mutatedType.indexer.key,
+            value: this.$.intrinsic.any,
+          };
+        }
+      },
+      onTailReplaced: (newTail) => {
+        if (this.mutatedType.indexer) {
+          this.mutatedType.indexer = {
+            key: this.mutatedType.indexer.key,
+            value: newTail.mutatedType,
+          };
+        }
+      },
+    });
+  }
+}

package/src/mutation-node/mutation-edge.ts

@@ -0,0 +1,43 @@
+import type { Type } from "@typespec/compiler";
+import { MutationNode } from "./mutation-node.js";
+
+export interface MutationEdgeOptions {
+  onTailMutation: () => void;
+  onTailDeletion: () => void;
+  onTailReplaced: (newTail: MutationNode<Type>) => void;
+}
+
+export class MutationEdge<THead extends Type, TTail extends Type> {
+  public head: MutationNode<THead>;
+  public tail: MutationNode<TTail>;
+  #options: MutationEdgeOptions;
+
+  constructor(head: MutationNode<THead>, tail: MutationNode<TTail>, options: MutationEdgeOptions) {
+    this.head = head;
+    this.tail = tail;
+    this.#options = options;
+    this.tail.addInEdge(this);
+  }
+
+  static create(head: MutationNode<Type>, tail: MutationNode<Type>, options: MutationEdgeOptions) {
+    return new MutationEdge(head, tail, options);
+  }
+
+  tailMutated(): void {
+    this.head.mutate();
+    this.#options.onTailMutation();
+  }
+
+  tailDeleted() {
+    this.head.mutate();
+    this.#options.onTailDeletion();
+  }
+
+  tailReplaced(newTail: MutationNode<TTail>) {
+    this.head.mutate();
+    this.tail.deleteInEdge(this);
+    this.tail = newTail;
+    this.tail.addInEdge(this);
+    this.#options.onTailReplaced(newTail);
+  }
+}

package/src/mutation-node/mutation-node.test.ts

@@ -0,0 +1,94 @@
+import { t, type TesterInstance } from "@typespec/compiler/testing";
+import { beforeEach, expect, it } from "vitest";
+import { Tester } from "../../test/test-host.js";
+import { getSubgraph } from "../../test/utils.js";
+
+let runner: TesterInstance;
+beforeEach(async () => {
+  runner = await Tester.createInstance();
+});
+
+it("Subgraph#getNode returns the same node for the same type when called", async () => {
+  const { Foo, program } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        prop: string;
+      }
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode1 = subgraph.getNode(Foo);
+  const fooNode2 = subgraph.getNode(Foo);
+  expect(fooNode1 === fooNode2).toBe(true);
+});
+
+it("Creates the same node when constructing the subgraph and coming back to the same type", async () => {
+  const { Foo, Bar, Baz, program } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        prop: string; 
+      }
+
+      model ${t.model("Bar")} {
+        foo: Foo;
+      }
+
+      model ${t.model("Baz")} {
+        foo: Foo;
+      }
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  subgraph.getNode(Bar);
+  subgraph.getNode(Baz);
+
+  expect(fooNode.inEdges.size).toBe(2);
+});
+
+it("starts with the mutatedType and sourceType being the same", async () => {
+  const { Foo, program } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        prop: string;
+      }
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  expect(fooNode.isMutated).toBe(false);
+  expect(fooNode.sourceType === fooNode.mutatedType).toBe(true);
+});
+
+it("clones the source type when mutating and sets isMutated to true", async () => {
+  const { Foo, program } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        prop: string;
+      }
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  expect(fooNode.isMutated).toBe(false);
+  expect(fooNode.sourceType === fooNode.mutatedType).toBe(true);
+  fooNode.mutate();
+  expect(fooNode.isMutated).toBe(true);
+  expect(fooNode.sourceType === fooNode.mutatedType).toBe(false);
+  expect(fooNode.sourceType.name).toEqual(fooNode.mutatedType.name);
+});
+
+it("invokes whenMutated callbacks when mutating", async () => {
+  const { Foo, program } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        prop: Bar;
+      }
+
+      model ${t.model("Bar")} {
+        prop: string;
+      }
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  const barNode = subgraph.getNode(Foo);
+  let called = false;
+  fooNode.whenMutated((mutatedType) => {
+    called = true;
+    expect(mutatedType).toBe(fooNode.mutatedType);
+  });
+  expect(called).toBe(false);
+  barNode.mutate();
+  expect(called).toBe(true);
+});

package/src/mutation-node/mutation-node.ts

@@ -0,0 +1,110 @@
+import type { MemberType, Type } from "@typespec/compiler";
+import type { Typekit } from "@typespec/compiler/typekit";
+import { mutationNodeFor, type MutationNodeForType } from "./factory.js";
+import type { MutationEdge } from "./mutation-edge.js";
+import type { MutationSubgraph } from "./mutation-subgraph.js";
+
+let nextId = 0;
+
+export abstract class MutationNode<T extends Type> {
+  abstract readonly kind: string;
+
+  id = nextId++;
+  sourceType: T;
+  mutatedType: T;
+  isMutated: boolean = false;
+  isDeleted: boolean = false;
+  isReplaced: boolean = false;
+  replacementNode: MutationNodeForType<Type> | null = null;
+  inEdges: Set<MutationEdge<Type, Type>> = new Set();
+  subgraph: MutationSubgraph;
+  referenceType: MemberType | null = null;
+  $: Typekit;
+
+  #whenMutatedCallbacks: ((mutatedType: Type | null) => void)[] = [];
+
+  constructor(subgraph: MutationSubgraph, sourceNode: T) {
+    this.subgraph = subgraph;
+    this.$ = this.subgraph.engine.$;
+    this.sourceType = sourceNode;
+    this.mutatedType = sourceNode;
+  }
+
+  abstract traverse(): void;
+
+  addInEdge(edge: MutationEdge<Type, Type>) {
+    this.inEdges.add(edge);
+  }
+
+  deleteInEdge(edge: MutationEdge<Type, Type>) {
+    this.inEdges.delete(edge);
+  }
+
+  whenMutated(cb: (mutatedType: T | null) => void) {
+    this.#whenMutatedCallbacks.push(cb as any);
+  }
+
+  mutate(initializeMutation?: (type: T) => void) {
+    if (this.isMutated || this.isDeleted || this.isReplaced) {
+      return;
+    }
+
+    this.mutatedType = this.$.type.clone(this.sourceType);
+    this.isMutated = true;
+    initializeMutation?.(this.mutatedType);
+    for (const cb of this.#whenMutatedCallbacks) {
+      cb(this.mutatedType);
+    }
+
+    for (const edge of this.inEdges) {
+      edge.tailMutated();
+    }
+
+    this.$.type.finishType(this.mutatedType);
+  }
+
+  delete() {
+    if (this.isMutated || this.isDeleted || this.isReplaced) {
+      return;
+    }
+
+    this.isDeleted = true;
+
+    for (const cb of this.#whenMutatedCallbacks) {
+      cb(null);
+    }
+
+    this.mutatedType = this.$.intrinsic.never as T;
+
+    for (const edge of this.inEdges) {
+      edge.tailDeleted();
+    }
+  }
+
+  replace(newType: Type) {
+    if (this.isMutated || this.isDeleted || this.isReplaced) {
+      return this;
+    }
+
+    // We need to make a new node because different types need to handle edge mutations differently.
+
+    this.isReplaced = true;
+    this.replacementNode = mutationNodeFor(this.subgraph, newType);
+    this.replacementNode.traverse();
+    // we don't need to do the clone stuff with this node, but we mark it as
+    // mutated because we don't want to allow further mutations on it.
+    this.replacementNode.isMutated = true;
+
+    if (this.referenceType) {
+      this.subgraph.replaceReferenceNode(this.referenceType, this.replacementNode);
+    } else {
+      this.subgraph.replaceNode(this, this.replacementNode);
+    }
+
+    for (const edge of this.inEdges) {
+      edge.tailReplaced(this.replacementNode);
+    }
+
+    return this.replacementNode;
+  }
+}

package/src/mutation-node/mutation-subgraph.ts

@@ -0,0 +1,59 @@
+import type { MemberType, Type } from "@typespec/compiler";
+import type { MutationEngine } from "../mutation/mutation-engine.js";
+import { mutationNodeFor, type MutationNodeForType } from "./factory.js";
+import type { MutationNode } from "./mutation-node.js";
+
+/**
+ * A subgraph of mutation nodes such that there is one node per type in the graph.
+ */
+export class MutationSubgraph {
+  #seenNodes = new Map<Type, MutationNode<Type>>();
+  #seenReferenceNodes = new Map<MemberType, MutationNode<Type>>();
+
+  engine: MutationEngine<any>;
+
+  constructor(engine: MutationEngine<any>) {
+    this.engine = engine;
+  }
+
+  getNode<T extends Type>(type: T, memberReferences: MemberType[] = []): MutationNodeForType<T> {
+    if (memberReferences.length > 0) {
+      return this.getReferenceNode(memberReferences[0]!) as any;
+    }
+
+    if (this.#seenNodes.has(type)) {
+      return this.#seenNodes.get(type)! as MutationNodeForType<T>;
+    }
+
+    const node = mutationNodeFor(this, type);
+    this.#seenNodes.set(type, node);
+    node.traverse();
+
+    return node;
+  }
+
+  getReferenceNode(memberType: MemberType): MutationNode<Type> {
+    if (this.#seenReferenceNodes.has(memberType)) {
+      return this.#seenReferenceNodes.get(memberType)!;
+    }
+
+    let referencedType: Type = memberType;
+    while (referencedType.kind === "ModelProperty" || referencedType.kind === "UnionVariant") {
+      referencedType = referencedType.type;
+    }
+    const node = mutationNodeFor(this, referencedType);
+    node.referenceType = memberType;
+    this.#seenReferenceNodes.set(memberType, node);
+    node.traverse();
+
+    return node;
+  }
+
+  replaceNode(oldNode: MutationNode<Type>, newNode: MutationNode<Type>) {
+    this.#seenNodes.set(oldNode.sourceType, newNode);
+  }
+
+  replaceReferenceNode(referenceType: MemberType, newNode: MutationNode<Type>) {
+    this.#seenReferenceNodes.set(referenceType, newNode);
+  }
+}

package/src/mutation-node/operation.ts

@@ -0,0 +1,49 @@
+import type { Model, Operation, Type } from "@typespec/compiler";
+import { MutationEdge } from "./mutation-edge.js";
+import { MutationNode } from "./mutation-node.js";
+
+export class OperationMutationNode extends MutationNode<Operation> {
+  readonly kind = "Operation";
+
+  traverse() {
+    const parameterNode = this.subgraph.getNode(this.sourceType.parameters);
+    this.connectParameters(parameterNode);
+
+    const returnTypeNode = this.subgraph.getNode(this.sourceType.returnType);
+    this.connectReturnType(returnTypeNode);
+  }
+
+  connectParameters(baseNode: MutationNode<Model>) {
+    MutationEdge.create(this, baseNode, {
+      onTailMutation: () => {
+        this.mutatedType!.parameters = baseNode.mutatedType;
+      },
+      onTailDeletion: () => {
+        this.mutatedType.parameters = this.$.model.create({
+          name: "",
+          properties: {},
+        });
+      },
+      onTailReplaced: (newTail) => {
+        if (newTail.mutatedType.kind !== "Model") {
+          throw new Error("Cannot replace parameters with non-model type");
+        }
+        this.mutatedType.parameters = newTail.mutatedType;
+      },
+    });
+  }
+
+  connectReturnType(typeNode: MutationNode<Type>) {
+    MutationEdge.create(this, typeNode, {
+      onTailMutation: () => {
+        this.mutatedType!.returnType = typeNode.mutatedType;
+      },
+      onTailDeletion: () => {
+        this.mutatedType.returnType = this.$.intrinsic.void;
+      },
+      onTailReplaced: (newTail) => {
+        this.mutatedType.returnType = newTail.mutatedType;
+      },
+    });
+  }
+}

package/src/mutation-node/scalar.test.ts

@@ -0,0 +1,44 @@
+import { t, type TesterInstance } from "@typespec/compiler/testing";
+import { $ } from "@typespec/compiler/typekit";
+import { beforeEach, expect, it } from "vitest";
+import { Tester } from "../../test/test-host.js";
+import { getSubgraph } from "../../test/utils.js";
+let runner: TesterInstance;
+beforeEach(async () => {
+  runner = await Tester.createInstance();
+});
+it("handles mutation of base scalars", async () => {
+  const { program, Base, Derived } = await runner.compile(t.code`
+      scalar ${t.scalar("Base")};
+      scalar ${t.scalar("Derived")} extends Base;
+    `);
+  const subgraph = getSubgraph(program);
+  const baseNode = subgraph.getNode(Base);
+  const derivedNode = subgraph.getNode(Derived);
+
+  baseNode.mutate();
+  expect(baseNode.isMutated).toBe(true);
+  expect(derivedNode.isMutated).toBe(true);
+  expect(derivedNode.mutatedType.baseScalar === baseNode.mutatedType).toBeTruthy();
+});
+
+it("handles replacement of scalars", async () => {
+  const { program, Base, Derived } = await runner.compile(t.code`
+      scalar ${t.scalar("Base")};
+      scalar ${t.scalar("Derived")} extends Base;
+    `);
+  const subgraph = getSubgraph(program);
+  const baseNode = subgraph.getNode(Base);
+  const derivedNode = subgraph.getNode(Derived);
+
+  const replacedNode = baseNode.replace($(program).builtin.string);
+  expect(replacedNode.isMutated).toBe(true);
+  expect(baseNode.isReplaced).toBe(true);
+
+  // subgraph is updated
+  expect(replacedNode === subgraph.getNode(Base)).toBe(true);
+
+  // derived node is updated
+  expect(derivedNode.isMutated).toBe(true);
+  expect(derivedNode.mutatedType.baseScalar === replacedNode.sourceType).toBe(true);
+});

package/src/mutation-node/scalar.ts

@@ -0,0 +1,32 @@
+import type { Scalar } from "@typespec/compiler";
+import { MutationEdge } from "./mutation-edge.js";
+import { MutationNode } from "./mutation-node.js";
+
+export class ScalarMutationNode extends MutationNode<Scalar> {
+  readonly kind = "Scalar";
+
+  traverse() {
+    if (this.sourceType.baseScalar) {
+      const baseScalarNode = this.subgraph.getNode(this.sourceType.baseScalar);
+      this.connectBaseScalar(baseScalarNode);
+    }
+  }
+
+  connectBaseScalar(baseScalar: MutationNode<Scalar>) {
+    MutationEdge.create(this, baseScalar, {
+      onTailReplaced: (newTail) => {
+        if (!this.$.scalar.is(newTail.mutatedType)) {
+          throw new Error("Cannot replace base scalar with non-scalar type");
+        }
+
+        this.mutatedType.baseScalar = newTail.mutatedType;
+      },
+      onTailMutation: () => {
+        this.mutatedType.baseScalar = baseScalar.mutatedType;
+      },
+      onTailDeletion: () => {
+        this.mutatedType.baseScalar = undefined;
+      },
+    });
+  }
+}

package/src/mutation-node/tuple.test.ts

@@ -0,0 +1,30 @@
+import type { Tuple } from "@typespec/compiler";
+import { t, type TesterInstance } from "@typespec/compiler/testing";
+import { $ } from "@typespec/compiler/typekit";
+import { beforeEach, expect, it } from "vitest";
+import { Tester } from "../../test/test-host.js";
+import { getSubgraph } from "../../test/utils.js";
+let runner: TesterInstance;
+beforeEach(async () => {
+  runner = await Tester.createInstance();
+});
+
+it("handles mutation of element types", async () => {
+  const { program, Foo, prop, Bar } = await runner.compile(t.code`
+      model ${t.model("Foo")} {
+        ${t.modelProperty("prop")}: [Bar, string];
+      }
+      model ${t.model("Bar")} {}
+
+    `);
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  const propNode = subgraph.getNode(prop);
+  const barNode = subgraph.getNode(Bar);
+  barNode.mutate();
+  expect(barNode.isMutated).toBe(true);
+  expect(propNode.isMutated).toBe(true);
+  expect(fooNode.isMutated).toBe(true);
+  expect((propNode.mutatedType.type as Tuple).values[0] === barNode.mutatedType).toBeTruthy();
+  expect((propNode.mutatedType.type as Tuple).values[1] === $(program).builtin.string).toBeTruthy();
+});

package/src/mutation-node/tuple.ts

@@ -0,0 +1,35 @@
+import type { Tuple, Type } from "@typespec/compiler";
+import { MutationEdge } from "./mutation-edge.js";
+import { MutationNode } from "./mutation-node.js";
+
+export class TupleMutationNode extends MutationNode<Tuple> {
+  readonly kind = "Tuple";
+  #indexMap: number[] = [];
+
+  traverse() {
+    for (let i = 0; i < this.sourceType.values.length; i++) {
+      const elemType = this.sourceType.values[i];
+      const elemNode = this.subgraph.getNode(elemType);
+      this.#indexMap[i] = i;
+      this.connectElement(elemNode, i);
+    }
+  }
+
+  connectElement(elemNode: MutationNode<Type>, index: number) {
+    MutationEdge.create(this, elemNode, {
+      onTailMutation: () => {
+        this.mutatedType.values[this.#indexMap[index]] = elemNode.mutatedType;
+      },
+      onTailDeletion: () => {
+        const spliceIndex = this.#indexMap[index];
+        this.mutatedType.values.splice(spliceIndex, 1);
+        for (let i = spliceIndex + 1; i < this.#indexMap.length; i++) {
+          this.#indexMap[i]--;
+        }
+      },
+      onTailReplaced: (newTail) => {
+        this.mutatedType.values[this.#indexMap[index]] = newTail.mutatedType;
+      },
+    });
+  }
+}

package/src/mutation-node/union-variant.test.ts

@@ -0,0 +1,28 @@
+import { t, type TesterInstance } from "@typespec/compiler/testing";
+import { $ } from "@typespec/compiler/typekit";
+import { beforeEach, expect, it } from "vitest";
+import { Tester } from "../../test/test-host.js";
+import { getSubgraph } from "../../test/utils.js";
+let runner: TesterInstance;
+beforeEach(async () => {
+  runner = await Tester.createInstance();
+});
+
+it("handles mutation of variant types", async () => {
+  const { program, Foo, v1 } = await runner.compile(t.code`
+      union ${t.union("Foo")} {
+        ${t.unionVariant("v1")}: string;
+        v2: int32;
+      }
+    `);
+
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  const v1Node = subgraph.getNode(v1);
+  const stringNode = subgraph.getNode($(program).builtin.string);
+  stringNode.mutate();
+  expect(stringNode.isMutated).toBe(true);
+  expect(v1Node.isMutated).toBe(true);
+  expect(fooNode.isMutated).toBe(true);
+  expect(v1Node.mutatedType.type === stringNode.mutatedType).toBeTruthy();
+});

package/src/mutation-node/union-variant.ts

@@ -0,0 +1,26 @@
+import type { Type, UnionVariant } from "@typespec/compiler";
+import { MutationEdge } from "./mutation-edge.js";
+import { MutationNode } from "./mutation-node.js";
+
+export class UnionVariantMutationNode extends MutationNode<UnionVariant> {
+  readonly kind = "UnionVariant";
+
+  traverse() {
+    const typeNode = this.subgraph.getNode(this.sourceType.type);
+    this.connectType(typeNode);
+  }
+
+  connectType(typeNode: MutationNode<Type>) {
+    MutationEdge.create(this, typeNode, {
+      onTailMutation: () => {
+        this.mutatedType.type = typeNode.mutatedType;
+      },
+      onTailDeletion: () => {
+        this.mutatedType.type = this.$.intrinsic.any;
+      },
+      onTailReplaced: (newTail) => {
+        this.mutatedType.type = newTail.mutatedType;
+      },
+    });
+  }
+}

package/src/mutation-node/union.test.ts

@@ -0,0 +1,26 @@
+import { t, type TesterInstance } from "@typespec/compiler/testing";
+import { beforeEach, expect, it } from "vitest";
+import { Tester } from "../../test/test-host.js";
+import { getSubgraph } from "../../test/utils.js";
+let runner: TesterInstance;
+beforeEach(async () => {
+  runner = await Tester.createInstance();
+});
+
+it("handles mutation of variants", async () => {
+  const { program, Foo, v1 } = await runner.compile(t.code`
+      union ${t.union("Foo")} {
+        ${t.unionVariant("v1")}: string;
+        v2: int32;
+      }
+    `);
+
+  const subgraph = getSubgraph(program);
+  const fooNode = subgraph.getNode(Foo);
+  const v1Node = subgraph.getNode(v1);
+  v1Node.mutate((clone) => (clone.name = "v1Renamed"));
+  expect(v1Node.isMutated).toBe(true);
+  expect(fooNode.isMutated).toBe(true);
+  expect(fooNode.mutatedType.variants.get("v1") === undefined).toBeTruthy();
+  expect(fooNode.mutatedType.variants.get("v1Renamed") === v1Node.mutatedType).toBeTruthy();
+});