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.

.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,16 +91,16 @@ 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;
 
@@ -102,6 +108,11 @@ it("removes model reference from namespace", async () => {
   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 () => {