Propagate Namespace mutation to its types (#4937)

- Expose mutateSubgraphWithNamespace which allows Namespace mutations - mutateSubgraph doesn't allow Namespace mutators. We want to keep namespace mutator hidden as it is a more advanced use case.

Author: joheredi

.chronus/changes/feature-mutate-namespaces-cascade-2024-9-31-21-41-8.md

@@ -0,0 +1,7 @@
+---
+changeKind: feature
+packages:
+ - "@typespec/compiler"
+---
+
+Add mutateSubgraphWithNamespace as a separate API

packages/compiler/src/experimental/index.ts

@@ -6,7 +6,9 @@ export {
  MutatorFn as unsafe_MutatorFn,
  MutatorRecord as unsafe_MutatorRecord,
  MutatorReplaceFn as unsafe_MutatorReplaceFn,
+ MutatorWithNamespace as unsafe_MutatorWithNamespace,
  mutateSubgraph as unsafe_mutateSubgraph,
+ mutateSubgraphWithNamespace as unsafe_mutateSubgraphWithNamespace,
 } from "./mutators.js";
 export { Realm as unsafe_Realm } from "./realm.js";
 export { unsafe_useStateMap, unsafe_useStateSet } from "./state-accessor.js";

packages/compiler/src/experimental/mutators.ts

@@ -73,9 +73,15 @@ export interface Mutator {
  ScalarConstructor?: MutatorRecord<ScalarConstructor>;
  StringTemplate?: MutatorRecord<StringTemplate>;
  StringTemplateSpan?: MutatorRecord<StringTemplateSpan>;
- Namespace?: MutatorRecord<Namespace>;
 }
 
+/**
+ * @experimental - This is a type that extends Mutator with a Namespace property.
+ */
+export type MutatorWithNamespace = Mutator & {
+ Namespace: MutatorRecord<Namespace>;
+};
+
 /** @experimental */
 export enum MutatorFlow {
  MutateAndRecurse = 0,
@@ -93,7 +99,10 @@ export type MutableType = Exclude<
  | FunctionParameter
  | ObjectType
  | Projection
+ | Namespace
 >;
+/** @experimental */
+export type MutableTypeWithNamespace = MutableType | Namespace;
 const typeId = CustomKeyMap.objectKeyer();
 const mutatorId = CustomKeyMap.objectKeyer();
 const seen = new CustomKeyMap<[MutableType, Set<Mutator> | Mutator[]], Type>(([type, mutators]) => {
@@ -103,6 +112,21 @@ const seen = new CustomKeyMap<[MutableType, Set<Mutator> | Mutator[]], Type>(([t
  return key;
 });
 
+/**
+ * Mutate the type graph with some namespace mutation.
+ * **Warning** this will most likely end up mutating the entire TypeGraph
+ * as every type relate to namespace in some way or another
+ * causing parent navigation which in turn would mutate everything in that namespace.
+ * @experimental
+ */
+export function mutateSubgraphWithNamespace<T extends MutableTypeWithNamespace>(
+ program: Program,
+ mutators: MutatorWithNamespace[],
+ type: T,
+): { realm: Realm | null; type: MutableTypeWithNamespace } {
+ return mutateSubgraph(program, mutators, type as any);
+}
+
 /** @experimental */
 export function mutateSubgraph<T extends MutableType>(
  program: Program,

packages/compiler/src/experimental/typekit/kits/type.ts

@@ -1,5 +1,5 @@
-import type { Enum, Model, Type } from "../../../core/types.js";
-import { defineKit } from "../define-kit.js";
+import { type Namespace, type Type } from "../../../core/types.js";
+import { $, defineKit } from "../define-kit.js";
 import { copyMap } from "../utils.js";
 
 /** @experimental */
@@ -70,20 +70,33 @@ defineKit<BaseTypeKit>({
  clone = this.program.checker.createType({
  ...type,
  decorators: [...type.decorators],
- decoratorDeclarations: new Map(type.decoratorDeclarations),
- models: new Map<string, Model>(type.models),
- enums: new Map<string, Enum>(type.enums),
- functionDeclarations: new Map(type.functionDeclarations),
  instantiationParameters: type.instantiationParameters
  ? [...type.instantiationParameters]
  : undefined,
- interfaces: new Map(type.interfaces),
- namespaces: new Map(type.namespaces),
- operations: new Map(type.operations),
  projections: [...type.projections],
- scalars: new Map(type.scalars),
- unions: new Map(type.unions),
  });
+ const clonedNamespace = clone as Namespace;
+ clonedNamespace.decoratorDeclarations = cloneTypeCollection(type.decoratorDeclarations, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.models = cloneTypeCollection(type.models, { namespace: clonedNamespace });
+ clonedNamespace.enums = cloneTypeCollection(type.enums, { namespace: clonedNamespace });
+ clonedNamespace.functionDeclarations = cloneTypeCollection(type.functionDeclarations, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.interfaces = cloneTypeCollection(type.interfaces, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.namespaces = cloneTypeCollection(type.namespaces, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.operations = cloneTypeCollection(type.operations, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.scalars = cloneTypeCollection(type.scalars, {
+ namespace: clonedNamespace,
+ });
+ clonedNamespace.unions = cloneTypeCollection(type.unions, { namespace: clonedNamespace });
  break;
  default:
  clone = this.program.checker.createType({
@@ -97,3 +110,18 @@ defineKit<BaseTypeKit>({
  },
  },
 });
+
+function cloneTypeCollection<T extends Type>(
+ collection: Map<string, T>,
+ options: { namespace?: Namespace } = {},
+): Map<string, T> {
+ const cloneCollection = new Map<string, T>();
+ for (const [key, type] of collection) {
+ const clone = $.type.clone(type);
+ if ("namespace" in clone && options.namespace) {
+ clone.namespace = options.namespace;
+ }
+ cloneCollection.set(key, clone);
+ }
+ return cloneCollection;
+}

packages/compiler/test/experimental/mutator.test.ts

@@ -1,5 +1,11 @@
 import { beforeEach, expect, it } from "vitest";
-import { mutateSubgraph, Mutator, MutatorFlow } from "../../src/experimental/mutators.js";
+import {
+ mutateSubgraph,
+ mutateSubgraphWithNamespace,
+ Mutator,
+ MutatorFlow,
+ MutatorWithNamespace,
+} from "../../src/experimental/mutators.js";
 import { Model, Namespace } from "../../src/index.js";
 import { createTestHost } from "../../src/testing/test-host.js";
 import { createTestWrapper } from "../../src/testing/test-utils.js";
@@ -85,23 +91,28 @@ it("removes model reference from namespace", async () => {
  `;
 
  const { Foo } = (await runner.compile(code)) as { Foo: Namespace; Bar: Model; Baz: Model };
- const mutator: Mutator = {
+ const mutator: MutatorWithNamespace = {
  name: "test",
  Namespace: {
- mutate: (ns, clone, p, realm) => {
+ mutate: (_ns, clone) => {
  clone.models.delete("Bar");
  },
  },
  };
 
- const { type } = mutateSubgraph(runner.program, [mutator], Foo);
+ const { type } = mutateSubgraphWithNamespace(runner.program, [mutator], Foo);
 
  const mutatedNs = type as Namespace;
 
  //Original namespace should have Bar model
  expect(Foo.models.has("Bar")).toBeTruthy();
  // Mutated namespace should not have Bar model
  expect(mutatedNs.models.has("Bar")).toBeFalsy();
+ // Mutated namespace is propagated to the models
+ expect(mutatedNs.models.get("Baz")!.namespace?.models.get("Bar")).toBeUndefined();
+ // Original should be unchanged
+ expect(Foo.models.get("Baz")!.namespace?.models.get("Bar")).toBeDefined();
+ expect(Foo.models.get("Baz")!.namespace).toBe(Foo);
 });
 
 it("do not recurse the model", async () => {