Fix recursive union flattening

packages/quicktype-core/src/rewrites/FlattenUnions.ts

@@ -7,7 +7,7 @@ import { assert } from "../support/Support";
 import { IntersectionType, type Type, UnionType } from "../Type/Type";
 import type { StringTypeMapping } from "../Type/TypeBuilderUtils";
 import type { TypeGraph } from "../Type/TypeGraph";
-import { type TypeRef, derefTypeRef } from "../Type/TypeRef";
+import { type TypeRef, derefTypeRef, typeRefIndex } from "../Type/TypeRef";
 import { makeGroupsToFlatten } from "../Type/TypeUtils";
 import { UnifyUnionBuilder, unifyTypes } from "../UnifyClasses";
 
@@ -20,20 +20,74 @@ export function flattenUnions(
 ): [TypeGraph, boolean] {
     let needsRepeat = false;
 
+    // During flattening, recursive unions can appear again while their replacement
+    // is still being built.  For example, flattening `A | M` can later ask for
+    // `A | (A | M)` through an array item or object property.  Those are the same
+    // union by associativity/idempotence, so remember each in-progress flattened
+    // union by a normalized key of its non-union member refs.  Looking up that key
+    // lets recursive occurrences reuse the replacement's forwarding ref instead
+    // of allocating another finite unrolling of the same cycle.
+    const unionKeyToRef = new Map<string, TypeRef>();
+
+    function addUnionKeyAtoms(
+        t: Type,
+        atoms: Set<number>,
+        seen: Set<number>,
+    ): void {
+        const index = t.index;
+        if (seen.has(index)) return;
+        seen.add(index);
+
+        if (t instanceof UnionType) {
+            for (const m of t.members) {
+                addUnionKeyAtoms(m, atoms, seen);
+            }
+        } else {
+            atoms.add(typeRefIndex(t.typeRef));
+        }
+    }
+
+    function unionKeyForTypes(types: Iterable<Type>): string {
+        const atoms = new Set<number>();
+        const seen = new Set<number>();
+        for (const t of types) {
+            addUnionKeyAtoms(t, atoms, seen);
+        }
+
+        return Array.from(atoms)
+            .sort((a, b) => a - b)
+            .join(",");
+    }
+
     function replace(
         types: ReadonlySet<Type>,
         builder: GraphRewriteBuilder<Type>,
         forwardingRef: TypeRef,
     ): TypeRef {
-        const unionBuilder = new UnifyUnionBuilder(
-            builder,
-            makeObjectTypes,
-            true,
-            (trefs) => {
-                assert(
-                    trefs.length > 0,
-                    "Must have at least one type to build union",
-                );
+        unionKeyToRef.set(unionKeyForTypes(types), forwardingRef);
+
+        let unionBuilder: UnifyUnionBuilder;
+        const unifyTypeRefs = (trefs: TypeRef[]): TypeRef => {
+            assert(
+                trefs.length > 0,
+                "Must have at least one type to build union",
+            );
+
+            const maybeReconstituted = builder.lookupTypeRefs(
+                trefs,
+                undefined,
+                false,
+            );
+            if (maybeReconstituted !== undefined) {
+                return maybeReconstituted;
+            }
+
+            const typesToUnify = new Set(
+                trefs.map((tref) => derefTypeRef(tref, graph)),
+            );
+            if (
+                iterableSome(typesToUnify, (t) => t instanceof IntersectionType)
+            ) {
                 trefs = trefs.map((tref) =>
                     builder.reconstituteType(derefTypeRef(tref, graph)),
                 );
@@ -46,7 +100,35 @@ export function flattenUnions(
                     emptyTypeAttributes,
                     new Set(trefs),
                 );
-            },
+            }
+
+            const key = unionKeyForTypes(typesToUnify);
+            const maybeRef = unionKeyToRef.get(key);
+            if (maybeRef !== undefined) {
+                return maybeRef;
+            }
+
+            return builder.withForwardingRef(
+                undefined,
+                (nestedForwardingRef) => {
+                    unionKeyToRef.set(key, nestedForwardingRef);
+                    return unifyTypes(
+                        typesToUnify,
+                        emptyTypeAttributes,
+                        builder,
+                        unionBuilder,
+                        conflateNumbers,
+                        nestedForwardingRef,
+                    );
+                },
+            );
+        };
+
+        unionBuilder = new UnifyUnionBuilder(
+            builder,
+            makeObjectTypes,
+            true,
+            unifyTypeRefs,
         );
         return unifyTypes(
             types,

test/inputs/schema/recursive-union-flattening.1.json

@@ -0,0 +1,9 @@
+[
+  {
+    "x": [
+      {
+        "x": []
+      }
+    ]
+  }
+]

test/inputs/schema/recursive-union-flattening.schema

@@ -0,0 +1,17 @@
+{
+  "items": {
+    "anyOf": [
+      { "$ref": "#/definitions/a" },
+      { "additionalProperties": { "$ref": "#" } }
+    ]
+  },
+  "definitions": {
+    "a": {
+      "properties": {
+        "x": {
+          "items": { "$ref": "#/definitions/a" }
+        }
+      }
+    }
+  }
+}

test/languages.ts

@@ -567,6 +567,7 @@ export const CJSONLanguage: Language = {
         "any.schema",
         "direct-union.schema",
         "optional-any.schema",
+        "recursive-union-flattening.schema",
         "required-non-properties.schema",
         /* Class elements with invalid type are not checked (for the current implementation, can be added later, should abord parsing and return NULL) */
         ...skipsUntypedUnions,
@@ -629,6 +630,8 @@ export const CPlusPlusLanguage: Language = {
     skipSchema: [
         // uses too much memory
         "keyword-unions.schema",
+        // Recursive top-level unions produce aliases that can refer to later aliases.
+        "recursive-union-flattening.schema",
     ],
     rendererOptions: {},
     quickTestRendererOptions: [
@@ -1202,9 +1205,10 @@ export const KotlinLanguage: Language = {
         // Some weird name collision
         "keyword-enum.schema",
         "keyword-unions.schema",
-        // Klaxon does not support top-level primitives
+        // Klaxon does not support top-level primitives/unions
         "top-level-enum.schema",
         "top-level-primitive.schema",
+        "recursive-union-flattening.schema",
     ],
     skipMiscJSON: false,
     rendererOptions: {},
@@ -1286,9 +1290,10 @@ export const KotlinJacksonLanguage: Language = {
         // Some weird name collision
         "keyword-enum.schema",
         "keyword-unions.schema",
-        // Klaxon does not support top-level primitives
+        // Klaxon does not support top-level primitives/unions
         "top-level-enum.schema",
         "top-level-primitive.schema",
+        "recursive-union-flattening.schema",
     ],
     skipMiscJSON: false,
     rendererOptions: { framework: "jackson" },
@@ -1626,6 +1631,7 @@ export const TypeScriptZodLanguage: Language = {
         "multi-type-enum.schema",
         "keyword-unions.schema",
         "optional-any.schema",
+        "recursive-union-flattening.schema",
         "required.schema",
         "required-non-properties.schema",
     ],
@@ -1804,6 +1810,9 @@ export const ElixirLanguage: Language = {
         // The test incorrectly succeeds due to the emitter being permissive for unions that contain only primitives. A future enhancement
         // for the Elixir emitter could be a user-controlled 'strict' mode that pattern matches even on unions of only primitive types.
         "go-schema-pattern-properties.schema",
+
+        // The generated top-level type is not emitted as a TopLevel module the fixture can call.
+        "recursive-union-flattening.schema",
     ],
     rendererOptions: {},
     quickTestRendererOptions: [],