Replace state maps with decorators for mutation engine metadata

Author: FionaBronwen

packages/graphql/lib/main.tsp

@@ -1,4 +1,6 @@
 import "./interface.tsp";
+import "./nullable.tsp";
+import "./one-of.tsp";
 import "./operation-fields.tsp";
 import "./operation-kind.tsp";
 import "./scalars.tsp";

packages/graphql/lib/nullable.tsp

@@ -0,0 +1,22 @@
+import "../dist/src/lib/nullable.js";
+
+using TypeSpec.Reflection;
+
+namespace TypeSpec.GraphQL;
+
+/**
+ * Mark a field, operation, or type as nullable in the emitted GraphQL schema.
+ *
+ * Applied automatically by the mutation engine when it strips `| null` from
+ * union types. The decorator's presence on the type's `decorators` array is
+ * the signal — the implementation is a no-op.
+ */
+extern dec nullable(target: ModelProperty | Operation | Union | Model);
+
+/**
+ * Mark a field or operation as having nullable array elements in the emitted GraphQL schema.
+ *
+ * Applied automatically by the mutation engine when it detects `Array<T | null>`
+ * patterns. Causes the emitter to emit `[T]` instead of `[T!]`.
+ */
+extern dec nullableElements(target: ModelProperty | Operation);

packages/graphql/lib/one-of.tsp

@@ -0,0 +1,16 @@
+import "../dist/src/lib/one-of.js";
+
+using TypeSpec.Reflection;
+
+namespace TypeSpec.GraphQL;
+
+/**
+ * Mark a model as a `@oneOf` input object in the emitted GraphQL schema.
+ *
+ * This decorator is applied automatically by the mutation engine when it converts
+ * a union type in input context to a synthetic input object (since GraphQL unions
+ * are output-only). The emitter uses this to emit the `@oneOf` directive.
+ *
+ * @see https://spec.graphql.org/September2025/#sec-OneOf-Input-Objects
+ */
+extern dec oneOf(target: Model);

packages/graphql/src/lib.ts

@@ -174,15 +174,6 @@ export const libDef = {
     interface: { description: "State for the @Interface decorator." },
     schema: { description: "State for the @schema decorator." },
     specifiedBy: { description: "State for the @specifiedBy decorator." },
-    oneOf: { description: "State for tracking @oneOf input objects created from input unions." },
-    nullable: {
-      description:
-        "State for tracking types and properties marked nullable after null-variant stripping by the mutation engine.",
-    },
-    nullableElements: {
-      description:
-        "State for tracking properties whose array element type was originally T | null before mutation.",
-    },
   },
 } as const;
 

packages/graphql/src/lib/nullable.ts

@@ -1,8 +1,38 @@
-import type { Program, Type } from "@typespec/compiler";
-import { useStateSet } from "@typespec/compiler/utils";
-import { GraphQLKeys } from "../lib.js";
+import type {
+  DecoratedType,
+  DecoratorContext,
+  DecoratorFunction,
+  Model,
+  ModelProperty,
+  Operation,
+  Type,
+  Union,
+} from "@typespec/compiler";
+import { NAMESPACE } from "../lib.js";
 
-const [getNullableState, setNullableState] = useStateSet<Type>(GraphQLKeys.nullable);
+// This will set the namespace for decorators implemented in this file
+export const namespace = NAMESPACE;
+
+/**
+ * Decorator implementation for `@nullable`.
+ *
+ * No-op — the decorator's presence on the type's `decorators` array is the
+ * signal. No additional state storage is needed.
+ */
+export const $nullable: DecoratorFunction = (
+  _context: DecoratorContext,
+  _target: ModelProperty | Operation | Union | Model,
+) => {};
+
+/**
+ * Decorator implementation for `@nullableElements`.
+ *
+ * No-op — presence on the decorators array is the signal.
+ */
+export const $nullableElements: DecoratorFunction = (
+  _context: DecoratorContext,
+  _target: ModelProperty | Operation,
+) => {};
 
 /**
  * Check whether a type was marked nullable after null-variant stripping.
@@ -12,30 +42,39 @@ const [getNullableState, setNullableState] = useStateSet<Type>(GraphQLKeys.nulla
  * - **Operation**: return type `T | null`
  * - **Union**: named unions like `Cat | Dog | null` (safe — new unique object)
  */
-export function isNullable(program: Program, type: Type): boolean {
-  return getNullableState(program, type);
+export function isNullable(type: Type): boolean {
+  if (!isDecoratedType(type)) return false;
+  return type.decorators.some((d) => d.decorator === $nullable);
 }
 
-/** Mark a type, property, or operation as nullable. */
-export function setNullable(program: Program, type: Type): void {
-  setNullableState(program, type);
+/**
+ * Mark a type, property, or operation as nullable.
+ * Called by the mutation engine when null variants are stripped.
+ */
+export function setNullable(type: Type): void {
+  if (!isDecoratedType(type)) return;
+  if (type.decorators.some((d) => d.decorator === $nullable)) return;
+  type.decorators.push({ decorator: $nullable, args: [] });
 }
 
-const [getNullableElementsState, setNullableElementsState] = useStateSet<Type>(
-  GraphQLKeys.nullableElements,
-);
-
 /**
  * Check whether a property's array elements were originally `T | null`.
  *
  * For `(string | null)[]`, marks the ModelProperty so components emit
  * `[String]` instead of `[String!]`.
  */
-export function hasNullableElements(program: Program, type: Type): boolean {
-  return getNullableElementsState(program, type);
+export function hasNullableElements(type: Type): boolean {
+  if (!isDecoratedType(type)) return false;
+  return type.decorators.some((d) => d.decorator === $nullableElements);
 }
 
 /** Mark a property as having nullable array elements. */
-export function setNullableElements(program: Program, type: Type): void {
-  setNullableElementsState(program, type);
+export function setNullableElements(type: Type): void {
+  if (!isDecoratedType(type)) return;
+  if (type.decorators.some((d) => d.decorator === $nullableElements)) return;
+  type.decorators.push({ decorator: $nullableElements, args: [] });
+}
+
+function isDecoratedType(type: Type): type is Type & DecoratedType {
+  return "decorators" in type;
 }

packages/graphql/src/lib/one-of.ts

@@ -1,22 +1,34 @@
-import type { Model, Program } from "@typespec/compiler";
-import { useStateSet } from "@typespec/compiler/utils";
-import { GraphQLKeys } from "../lib.js";
+import type { DecoratorContext, DecoratorFunction, Model } from "@typespec/compiler";
+import { NAMESPACE } from "../lib.js";
 
-const [getOneOfState, setOneOfState] = useStateSet<Model>(GraphQLKeys.oneOf);
+// This will set the namespace for decorators implemented in this file
+export const namespace = NAMESPACE;
+
+/**
+ * Decorator implementation for `@oneOf`.
+ *
+ * No-op — the decorator's presence on the type's `decorators` array is the
+ * signal. No additional state storage is needed.
+ */
+export const $oneOf: DecoratorFunction = (
+  _context: DecoratorContext,
+  _target: Model,
+) => {};
 
 /**
  * Check if a model has been marked as a @oneOf input object.
  * These are synthetic models created by the union mutation when a union
  * is used in input context — GraphQL unions are output-only, so input
  * unions become @oneOf input objects.
  */
-export function isOneOf(program: Program, model: Model): boolean {
-  return getOneOfState(program, model);
+export function isOneOf(model: Model): boolean {
+  return model.decorators.some((d) => d.decorator === $oneOf);
 }
 
 /**
  * Mark a model as a @oneOf input object.
  */
-export function setOneOf(program: Program, model: Model): void {
-  setOneOfState(program, model);
+export function setOneOf(model: Model): void {
+  if (model.decorators.some((d) => d.decorator === $oneOf)) return;
+  model.decorators.push({ decorator: $oneOf, args: [] });
 }

packages/graphql/src/mutation-engine/mutations/model-property.ts

@@ -52,10 +52,10 @@ export class GraphQLModelPropertyMutation extends SimpleModelPropertyMutation<Si
     super.mutate();
 
     if (isInlineNullable) {
-      setNullable(this.engine.$.program, this.mutatedType);
+      setNullable(this.mutatedType);
     }
     if (isArrayWithNullableElements) {
-      setNullableElements(this.engine.$.program, this.mutatedType);
+      setNullableElements(this.mutatedType);
     }
   }
 }

packages/graphql/src/mutation-engine/mutations/operation.ts

@@ -1,13 +1,17 @@
-import type { MemberType, Operation } from "@typespec/compiler";
+import { isArrayModelType, type MemberType, type Operation } from "@typespec/compiler";
 import {
   SimpleOperationMutation,
   type MutationInfo,
   type SimpleMutationEngine,
   type SimpleMutationOptions,
   type SimpleMutations,
 } from "@typespec/mutator-framework";
-import { setNullable } from "../../lib/nullable.js";
-import { isNullableUnion, sanitizeNameForGraphQL } from "../../lib/type-utils.js";
+import { setNullable, setNullableElements } from "../../lib/nullable.js";
+import {
+  isNullableUnion,
+  sanitizeNameForGraphQL,
+  unwrapNullableUnion,
+} from "../../lib/type-utils.js";
 import { GraphQLMutationOptions, GraphQLTypeContext } from "../options.js";
 
 /** GraphQL-specific Operation mutation. */
@@ -44,15 +48,29 @@ export class GraphQLOperationMutation extends SimpleOperationMutation<SimpleMuta
 
   mutate() {
     // Snapshot return-type nullability before mutation replaces it.
-    const hasNullableReturn = isNullableUnion(this.sourceType.returnType);
+    const returnType = this.sourceType.returnType;
+    const hasNullableReturn = isNullableUnion(returnType);
+
+    // For element nullability, look through an outer `| null` wrapper to find the array.
+    // e.g. `(string | null)[] | null` → unwrap outer null → check array elements.
+    const innerReturnType =
+      returnType.kind === "Union" ? (unwrapNullableUnion(returnType) ?? returnType) : returnType;
+
+    const hasNullableElements =
+      innerReturnType.kind === "Model" &&
+      isArrayModelType(innerReturnType) &&
+      isNullableUnion(innerReturnType.indexer.value);
 
     this.mutationNode.mutate((operation) => {
       operation.name = sanitizeNameForGraphQL(operation.name);
     });
     super.mutate();
 
     if (hasNullableReturn) {
-      setNullable(this.engine.$.program, this.mutatedType);
+      setNullable(this.mutatedType);
+    }
+    if (hasNullableElements) {
+      setNullableElements(this.mutatedType);
     }
   }
 }

packages/graphql/src/mutation-engine/mutations/union.ts

@@ -142,7 +142,7 @@ export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, Mu
     }
 
     if (hasNull) {
-      setNullable(program, this.mutatedType);
+      setNullable(this.mutatedType);
     }
 
     // GraphQL unions can only contain object types — wrap scalars in synthetic models
@@ -207,10 +207,10 @@ export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, Mu
       properties,
     });
 
-    setOneOf(program, oneOfModel);
+    setOneOf(oneOfModel);
 
     if (hasNull) {
-      setNullable(program, oneOfModel);
+      setNullable(oneOfModel);
     }
 
     this.#mutationNode.replace(oneOfModel);