@kubb/ast 5.0.0-beta.3 → 5.0.0-beta.31

package/src/visitor.ts

@@ -1,7 +1,19 @@
 import type { VisitorDepth } from './constants.ts'
 import { visitorDepths, WALK_CONCURRENCY } from './constants.ts'
 import { createParameter, createProperty } from './factory.ts'
-import type { InputNode, Node, OperationNode, OutputNode, ParameterNode, PropertyNode, ResponseNode, SchemaNode } from './nodes/index.ts'
+import type {
+  ContentNode,
+  InputNode,
+  Node,
+  NodeKind,
+  OperationNode,
+  OutputNode,
+  ParameterNode,
+  PropertyNode,
+  RequestBodyNode,
+  ResponseNode,
+  SchemaNode,
+} from './nodes/index.ts'
 
 /**
  * Creates a small async concurrency limiter.
@@ -54,7 +66,9 @@ type ParentNodeMap = [
   [InputNode, undefined],
   [OutputNode, undefined],
   [OperationNode, InputNode],
-  [SchemaNode, InputNode | OperationNode | SchemaNode | PropertyNode | ParameterNode | ResponseNode],
+  [RequestBodyNode, OperationNode],
+  [ContentNode, RequestBodyNode | ResponseNode],
+  [SchemaNode, InputNode | ContentNode | SchemaNode | PropertyNode | ParameterNode],
   [PropertyNode, SchemaNode],
   [ParameterNode, OperationNode],
   [ResponseNode, OperationNode],
@@ -115,25 +129,40 @@ export type VisitorContext<T extends Node = Node> = {
 }
 
 /**
- * Synchronous visitor used by `transform`.
+ * Synchronous visitor consumed by `transform`. Each optional callback runs
+ * for the matching node type. Return a new node to replace it, or `undefined`
+ * to leave it untouched.
  *
- * @example
+ * Plugins typically expose `transformer` so users can supply a `Visitor` that
+ * rewrites operation IDs, drops descriptions, or otherwise tweaks the AST
+ * before printing.
+ *
+ * @example Prefix every operationId
  * ```ts
  * const visitor: Visitor = {
  *   operation(node) {
- *     return { ...node, operationId: `x_${node.operationId}` }
+ *     return { ...node, operationId: `api_${node.operationId}` }
+ *   },
+ * }
+ * ```
+ *
+ * @example Strip schema descriptions
+ * ```ts
+ * const visitor: Visitor = {
+ *   schema(node) {
+ *     return { ...node, description: undefined }
  *   },
  * }
  * ```
  */
 export type Visitor = {
-  input?(node: InputNode, context: VisitorContext<InputNode>): void | InputNode
-  output?(node: OutputNode, context: VisitorContext<OutputNode>): void | OutputNode
-  operation?(node: OperationNode, context: VisitorContext<OperationNode>): void | OperationNode
-  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): void | SchemaNode
-  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): void | PropertyNode
-  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): void | ParameterNode
-  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): void | ResponseNode
+  input?(node: InputNode, context: VisitorContext<InputNode>): undefined | null | InputNode
+  output?(node: OutputNode, context: VisitorContext<OutputNode>): undefined | null | OutputNode
+  operation?(node: OperationNode, context: VisitorContext<OperationNode>): undefined | null | OperationNode
+  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): undefined | null | SchemaNode
+  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): undefined | null | PropertyNode
+  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): undefined | null | ParameterNode
+  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): undefined | null | ResponseNode
 }
 
 /**
@@ -154,13 +183,13 @@ type MaybePromise<T> = T | Promise<T>
  * ```
  */
 export type AsyncVisitor = {
-  input?(node: InputNode, context: VisitorContext<InputNode>): MaybePromise<void | InputNode>
-  output?(node: OutputNode, context: VisitorContext<OutputNode>): MaybePromise<void | OutputNode>
-  operation?(node: OperationNode, context: VisitorContext<OperationNode>): MaybePromise<void | OperationNode>
-  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): MaybePromise<void | SchemaNode>
-  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): MaybePromise<void | PropertyNode>
-  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): MaybePromise<void | ParameterNode>
-  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): MaybePromise<void | ResponseNode>
+  input?(node: InputNode, context: VisitorContext<InputNode>): MaybePromise<undefined | null | InputNode>
+  output?(node: OutputNode, context: VisitorContext<OutputNode>): MaybePromise<undefined | null | OutputNode>
+  operation?(node: OperationNode, context: VisitorContext<OperationNode>): MaybePromise<undefined | null | OperationNode>
+  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): MaybePromise<undefined | null | SchemaNode>
+  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): MaybePromise<undefined | null | PropertyNode>
+  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): MaybePromise<undefined | null | ParameterNode>
+  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): MaybePromise<undefined | null | ResponseNode>
 }
 
 /**
@@ -176,13 +205,13 @@ export type AsyncVisitor = {
  * ```
  */
 export type CollectVisitor<T> = {
-  input?(node: InputNode, context: VisitorContext<InputNode>): T | undefined
-  output?(node: OutputNode, context: VisitorContext<OutputNode>): T | undefined
-  operation?(node: OperationNode, context: VisitorContext<OperationNode>): T | undefined
-  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): T | undefined
-  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): T | undefined
-  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): T | undefined
-  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): T | undefined
+  input?(node: InputNode, context: VisitorContext<InputNode>): T | null | undefined
+  output?(node: OutputNode, context: VisitorContext<OutputNode>): T | null | undefined
+  operation?(node: OperationNode, context: VisitorContext<OperationNode>): T | null | undefined
+  schema?(node: SchemaNode, context: VisitorContext<SchemaNode>): T | null | undefined
+  property?(node: PropertyNode, context: VisitorContext<PropertyNode>): T | null | undefined
+  parameter?(node: ParameterNode, context: VisitorContext<ParameterNode>): T | null | undefined
+  response?(node: ResponseNode, context: VisitorContext<ResponseNode>): T | null | undefined
 }
 
 /**
@@ -253,60 +282,102 @@ export type CollectOptions<T> = CollectVisitor<T> & {
 }
 
 /**
- * Returns the immediate traversable children of `node`.
+ * Child node fields per node kind, in traversal order (Babel's `VISITOR_KEYS`).
  *
- * For `Schema` nodes, children (`properties`, `items`, `members`, and non-boolean
- * `additionalProperties`) are only included
- * when `recurse` is `true`; shallow mode skips them.
+ * Each listed property holds a child node, an array of child nodes, or — for
+ * `additionalProperties` — a node or the literal `true` (skipped). Every value
+ * in a child slot is a node, so one table drives both `getChildren` and `transform`.
+ */
+const VISITOR_KEYS = {
+  Input: ['schemas', 'operations'],
+  Operation: ['parameters', 'requestBody', 'responses'],
+  RequestBody: ['content'],
+  Content: ['schema'],
+  Response: ['content'],
+  Schema: ['properties', 'items', 'members', 'additionalProperties'],
+  Property: ['schema'],
+  Parameter: ['schema'],
+} as const satisfies Partial<Record<NodeKind, ReadonlyArray<string>>>
+
+const visitorKeysByKind = VISITOR_KEYS as Record<string, ReadonlyArray<string> | undefined>
+
+/**
+ * Returns `true` when `value` is an AST node (an object carrying a `kind`).
+ */
+function isNode(value: unknown): value is Node {
+  return typeof value === 'object' && value !== null && 'kind' in value
+}
+
+/**
+ * Returns the immediate traversable children of `node` based on {@link VISITOR_KEYS}.
+ *
+ * `Schema` children are only included when `recurse` is `true`; shallow mode skips them.
  *
  * @example
  * ```ts
  * const children = getChildren(operationNode, true)
- * // returns parameters, requestBody schema (if present), and responses
+ * // returns parameters, the request body, and responses
  * ```
  */
-function getChildren(node: Node, recurse: boolean): Array<Node> {
-  switch (node.kind) {
-    case 'Input':
-      return [...node.schemas, ...node.operations]
-    case 'Output':
-      return []
-    case 'Operation':
-      return [...node.parameters, ...(node.requestBody?.content?.flatMap((c) => (c.schema ? [c.schema] : [])) ?? []), ...node.responses]
-    case 'Schema': {
-      const children: Array<Node> = []
-
-      if (!recurse) return []
-
-      if ('properties' in node && node.properties.length > 0) children.push(...node.properties)
-      if ('items' in node && node.items) children.push(...node.items)
-      if ('members' in node && node.members) children.push(...node.members)
-      if ('additionalProperties' in node && node.additionalProperties && node.additionalProperties !== true) children.push(node.additionalProperties)
-
-      return children
+function* getChildren(node: Node, recurse: boolean): Generator<Node, void, undefined> {
+  if (node.kind === 'Schema' && !recurse) return
+
+  const keys = visitorKeysByKind[node.kind]
+  if (!keys) return
+
+  const record = node as unknown as Record<string, unknown>
+  for (const key of keys) {
+    const value = record[key]
+    if (Array.isArray(value)) {
+      for (const item of value) if (isNode(item)) yield item
+    } else if (isNode(value)) {
+      yield value
     }
-    case 'Property':
-      return [node.schema]
-    case 'Parameter':
-      return [node.schema]
-    case 'Response':
-      return node.schema ? [node.schema] : []
-    case 'FunctionParameter':
-    case 'ParameterGroup':
-    case 'FunctionParameters':
-    case 'Type':
-      return []
-    default:
-      return []
   }
 }
 
 /**
- * Depth-first traversal for side effects. Visitor return values are ignored.
- * Sibling nodes at each level are visited concurrently up to `options.concurrency`
- * (default: `WALK_CONCURRENCY`).
+ * Maps a node `kind` to the matching visitor callback name. Only the seven
+ * traversable node kinds have an entry; every other kind resolves to
+ * `undefined` and is skipped.
+ */
+const VISITOR_KEY_BY_KIND: Partial<Record<NodeKind, keyof Visitor>> = {
+  Input: 'input',
+  Output: 'output',
+  Operation: 'operation',
+  Schema: 'schema',
+  Property: 'property',
+  Parameter: 'parameter',
+  Response: 'response',
+}
+
+/**
+ * Invokes the visitor callback that matches `node.kind`, passing the traversal
+ * context. Returns the callback's result (a replacement node, a collected
+ * value, or `undefined` when no callback is registered for the kind).
  *
- * @example
+ * Shared by `walk`, `transform`, and `collectLazy` so node-kind dispatch lives
+ * in one place. `TResult` is the caller's expected return: the same node type
+ * for `transform`, the collected value type for `collectLazy`, ignored for `walk`.
+ */
+function applyVisitor<TResult>(node: Node, visitor: Visitor | AsyncVisitor | CollectVisitor<unknown>, parent: Node | undefined): TResult | null | undefined {
+  const key = VISITOR_KEY_BY_KIND[node.kind]
+  if (!key) return undefined
+
+  const fn = visitor[key] as ((node: Node, context: VisitorContext) => TResult | null | undefined) | undefined
+
+  return fn?.(node, { parent })
+}
+
+/**
+ * Async depth-first traversal for side effects. Visitor return values are
+ * ignored. Use `transform` when you want to rewrite nodes.
+ *
+ * Sibling nodes at each depth run concurrently up to `options.concurrency`
+ * (defaults to `WALK_CONCURRENCY`). Higher values overlap I/O-bound visitor
+ * work; lower values reduce memory pressure.
+ *
+ * @example Log every operation
  * ```ts
  * await walk(root, {
  *   operation(node) {
@@ -315,10 +386,9 @@ function getChildren(node: Node, recurse: boolean): Array<Node> {
  * })
  * ```
  *
- * @example
+ * @example Only visit the root node
  * ```ts
- * // Visit only the current node
- * await walk(root, { depth: 'shallow', root: () => {} })
+ * await walk(root, { depth: 'shallow', input: () => {} })
  * ```
  */
 export async function walk(node: Node, options: WalkOptions): Promise<void> {
@@ -329,41 +399,7 @@ export async function walk(node: Node, options: WalkOptions): Promise<void> {
 }
 
 async function _walk(node: Node, visitor: AsyncVisitor, recurse: boolean, limit: LimitFn, parent: Node | undefined): Promise<void> {
-  switch (node.kind) {
-    case 'Input':
-      await limit(() => visitor.input?.(node, { parent: parent as ParentOf<InputNode> }))
-      break
-    case 'Output':
-      await limit(() => visitor.output?.(node, { parent: parent as ParentOf<OutputNode> }))
-      break
-    case 'Operation':
-      await limit(() =>
-        visitor.operation?.(node, {
-          parent: parent as ParentOf<OperationNode>,
-        }),
-      )
-      break
-    case 'Schema':
-      await limit(() => visitor.schema?.(node, { parent: parent as ParentOf<SchemaNode> }))
-      break
-    case 'Property':
-      await limit(() => visitor.property?.(node, { parent: parent as ParentOf<PropertyNode> }))
-      break
-    case 'Parameter':
-      await limit(() =>
-        visitor.parameter?.(node, {
-          parent: parent as ParentOf<ParameterNode>,
-        }),
-      )
-      break
-    case 'Response':
-      await limit(() => visitor.response?.(node, { parent: parent as ParentOf<ResponseNode> }))
-      break
-    case 'FunctionParameter':
-    case 'ParameterGroup':
-    case 'FunctionParameters':
-      break
-  }
+  await limit(() => applyVisitor(node, visitor, parent))
 
   const children = getChildren(node, recurse)
   for (const child of children) {
@@ -372,12 +408,13 @@ async function _walk(node: Node, visitor: AsyncVisitor, recurse: boolean, limit:
 }
 
 /**
- * Runs a depth-first immutable transform.
+ * Synchronous depth-first transform. Each visitor callback gets a chance to
+ * return a replacement node; `undefined` keeps the original.
  *
- * If a visitor returns a node, it replaces the current node.
- * If it returns `undefined`, the current node stays the same.
+ * The transform is immutable. The original tree is not mutated; a new tree
+ * is returned. Use `depth: 'shallow'` to skip recursion into children.
  *
- * @example
+ * @example Prefix every operationId
  * ```ts
  * const next = transform(root, {
  *   operation(node) {
@@ -386,10 +423,12 @@ async function _walk(node: Node, visitor: AsyncVisitor, recurse: boolean, limit:
  * })
  * ```
  *
- * @example
+ * @example Replace only the root node
  * ```ts
- * // Shallow mode: only transform the input node
- * const next = transform(root, { depth: 'shallow', root: (node) => node })
+ * const next = transform(root, {
+ *   depth: 'shallow',
+ *   input: (node) => ({ ...node, meta: { ...node.meta, title: 'Rewritten' } }),
+ * })
  * ```
  */
 export function transform(node: InputNode, options: TransformOptions): InputNode
@@ -404,189 +443,104 @@ export function transform(node: Node, options: TransformOptions): Node {
   const { depth, parent, ...visitor } = options
   const recurse = (depth ?? visitorDepths.deep) === visitorDepths.deep
 
-  switch (node.kind) {
-    case 'Input': {
-      let input = node
-      const replaced = visitor.input?.(input, {
-        parent: parent as ParentOf<InputNode>,
-      })
-      if (replaced) input = replaced
+  const visited = applyVisitor<Node>(node, visitor, parent) ?? node
+  const rebuilt = transformChildren(visited, options, recurse)
 
-      return {
-        ...input,
-        schemas: input.schemas.map((s) => transform(s, { ...options, parent: input })),
-        operations: input.operations.map((op) => transform(op, { ...options, parent: input })),
-      }
-    }
-    case 'Output': {
-      let output = node
-      const replaced = visitor.output?.(output, {
-        parent: parent as ParentOf<OutputNode>,
-      })
-      if (replaced) output = replaced
+  // Structural sharing: when the visitor and child rebuild both left this node
+  // untouched, return the original reference so callers can detect "nothing
+  // changed" by identity and ancestors can avoid reallocating.
+  if (rebuilt === node) return node
 
-      return output
-    }
-    case 'Operation': {
-      let op = node
-      const replaced = visitor.operation?.(op, {
-        parent: parent as ParentOf<OperationNode>,
-      })
-      if (replaced) op = replaced
-
-      return {
-        ...op,
-        parameters: op.parameters.map((p) => transform(p, { ...options, parent: op })),
-        requestBody: op.requestBody
-          ? {
-              ...op.requestBody,
-              content: op.requestBody.content?.map((c) => ({
-                ...c,
-                schema: c.schema ? transform(c.schema, { ...options, parent: op }) : undefined,
-              })),
-            }
-          : undefined,
-        responses: op.responses.map((r) => transform(r, { ...options, parent: op })),
-      }
-    }
-    case 'Schema': {
-      let schema = node
-      const replaced = visitor.schema?.(schema, {
-        parent: parent as ParentOf<SchemaNode>,
-      })
-      if (replaced) schema = replaced
-
-      const childOptions = { ...options, parent: schema }
-
-      return {
-        ...schema,
-        ...('properties' in schema && recurse
-          ? {
-              properties: schema.properties.map((p) => transform(p, childOptions)),
-            }
-          : {}),
-        ...('items' in schema && recurse ? { items: schema.items?.map((i) => transform(i, childOptions)) } : {}),
-        ...('members' in schema && recurse ? { members: schema.members?.map((m) => transform(m, childOptions)) } : {}),
-        ...('additionalProperties' in schema && recurse && schema.additionalProperties && schema.additionalProperties !== true
-          ? {
-              additionalProperties: transform(schema.additionalProperties, childOptions),
-            }
-          : {}),
-      } as SchemaNode
-    }
-    case 'Property': {
-      let prop = node
-      const replaced = visitor.property?.(prop, {
-        parent: parent as ParentOf<PropertyNode>,
-      })
-      if (replaced) prop = replaced
+  const finalize = nodeFinalizers[rebuilt.kind]
+  return finalize ? finalize(rebuilt) : rebuilt
+}
 
-      return createProperty({
-        ...prop,
-        schema: transform(prop.schema, { ...options, parent: prop }),
-      })
-    }
-    case 'Parameter': {
-      let param = node
-      const replaced = visitor.parameter?.(param, {
-        parent: parent as ParentOf<ParameterNode>,
-      })
-      if (replaced) param = replaced
+/**
+ * Per-kind builders rerun after children are rebuilt. `Property`/`Parameter`
+ * resync schema optionality against their `required` flag once the schema may
+ * have changed.
+ */
+const nodeFinalizers: Partial<Record<NodeKind, (node: Node) => Node>> = {
+  Property: (node) => createProperty(node as PropertyNode),
+  Parameter: (node) => createParameter(node as ParameterNode),
+}
 
-      return createParameter({
-        ...param,
-        schema: transform(param.schema, { ...options, parent: param }),
-      })
-    }
-    case 'Response': {
-      let response = node
-      const replaced = visitor.response?.(response, {
-        parent: parent as ParentOf<ResponseNode>,
+/**
+ * Immutably rebuilds a node's children using {@link VISITOR_KEYS}, transforming
+ * each child node and leaving non-node values (e.g. `additionalProperties: true`) intact.
+ * `Schema` children are skipped in shallow mode.
+ */
+function transformChildren(node: Node, options: TransformOptions, recurse: boolean): Node {
+  if (node.kind === 'Schema' && !recurse) return node
+
+  const keys = visitorKeysByKind[node.kind]
+  if (!keys) return node
+
+  const record = node as unknown as Record<string, unknown>
+  const childOptions = { ...options, parent: node }
+  let updates: Record<string, unknown> | undefined
+
+  for (const key of keys) {
+    if (!(key in record)) continue
+    const value = record[key]
+    if (Array.isArray(value)) {
+      let changed = false
+      const mapped = value.map((item) => {
+        if (!isNode(item)) return item
+        const next = transform(item, childOptions)
+        if (next !== item) changed = true
+        return next
       })
-      if (replaced) response = replaced
-
-      return {
-        ...response,
-        schema: transform(response.schema, { ...options, parent: response }),
-      }
+      if (changed) (updates ??= {})[key] = mapped
+    } else if (isNode(value)) {
+      const next = transform(value, childOptions)
+      if (next !== value) (updates ??= {})[key] = next
     }
-    case 'FunctionParameter':
-    case 'ParameterGroup':
-    case 'FunctionParameters':
-    case 'Type':
-      return node
-    default:
-      return node
   }
+
+  return updates ? ({ ...node, ...updates } as Node) : node
 }
 /**
- * Runs a depth-first synchronous collection pass.
- *
- * Non-`undefined` values returned by visitor callbacks are appended to the result.
+ * Lazy depth-first collection pass. Yields every non-null value returned by
+ * the visitor callbacks. Use `collect` for the eager array form.
  *
- * @example
+ * @example Collect every operationId
  * ```ts
- * const ids = collect(root, {
+ * const ids: string[] = []
+ * for (const id of collectLazy<string>(root, {
  *   operation(node) {
  *     return node.operationId
  *   },
- * })
- * ```
- *
- * @example
- * ```ts
- * // Collect from only the current node
- * const values = collect(root, { depth: 'shallow', root: () => 'root' })
+ * })) {
+ *   ids.push(id)
+ * }
  * ```
  */
-export function collect<T>(node: Node, options: CollectOptions<T>): Array<T> {
+export function* collectLazy<T>(node: Node, options: CollectOptions<T>): Generator<T, void, undefined> {
   const { depth, parent, ...visitor } = options
   const recurse = (depth ?? visitorDepths.deep) === visitorDepths.deep
-  const results: Array<T> = []
-
-  let v: T | undefined
-  switch (node.kind) {
-    case 'Input':
-      v = visitor.input?.(node, { parent: parent as ParentOf<InputNode> })
-      break
-    case 'Output':
-      v = visitor.output?.(node, { parent: parent as ParentOf<OutputNode> })
-      break
-    case 'Operation':
-      v = visitor.operation?.(node, {
-        parent: parent as ParentOf<OperationNode>,
-      })
-      break
-    case 'Schema':
-      v = visitor.schema?.(node, { parent: parent as ParentOf<SchemaNode> })
-      break
-    case 'Property':
-      v = visitor.property?.(node, {
-        parent: parent as ParentOf<PropertyNode>,
-      })
-      break
-    case 'Parameter':
-      v = visitor.parameter?.(node, {
-        parent: parent as ParentOf<ParameterNode>,
-      })
-      break
-    case 'Response':
-      v = visitor.response?.(node, {
-        parent: parent as ParentOf<ResponseNode>,
-      })
-      break
-    case 'FunctionParameter':
-    case 'ParameterGroup':
-    case 'FunctionParameters':
-      break
-  }
-  if (v !== undefined) results.push(v)
+
+  const v = applyVisitor<T>(node, visitor, parent)
+  if (v != null) yield v
 
   for (const child of getChildren(node, recurse)) {
-    for (const item of collect(child, { ...options, parent: node })) {
-      results.push(item)
-    }
+    yield* collectLazy(child, { ...options, parent: node })
   }
+}
 
-  return results
+/**
+ * Eager depth-first collection pass. Returns an array of every non-null value
+ * the visitor callbacks return.
+ *
+ * @example Collect every operationId
+ * ```ts
+ * const ids = collect<string>(root, {
+ *   operation(node) {
+ *     return node.operationId
+ *   },
+ * })
+ * ```
+ */
+export function collect<T>(node: Node, options: CollectOptions<T>): Array<T> {
+  return Array.from(collectLazy(node, options))
 }