From 37d0a6db5ae0116abb4fb6c809f73f472aac00c6 Mon Sep 17 00:00:00 2001 From: 2570165831 <2570165831@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:31:51 +0800 Subject: [PATCH 1/5] Let the model choose the variant of an anyOf Constrained generation always took the first variant of an anyOf, so a nullable union could never produce null and a union of object types always produced the first type (#289). The generator now resolves references and flattens nested unions, then samples once among the tokens that can start each variant and continues with the variants that can start with the sample. Object variants all start with `{`, so the key mask offers the keys of every live variant and each emitted key keeps the variants that declare it; a shared key whose values are string choices (enum case tags, constants) narrows by the chosen value. Variants that start alike but aren't all objects still fall back to the first of them, noted where it happens. Co-Authored-By: Claude Opus 5.5 --- .../Shared/StructuredGeneration.swift | 157 ++++++++++++++++-- .../AnyOfGenerationTests.swift | 133 +++++++++++++++ 2 files changed, 274 insertions(+), 16 deletions(-) create mode 100644 Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index 5bb0880d..da2a7a4e 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -398,15 +398,86 @@ struct ConstrainedJSONGenerator { guard !variants.isEmpty else { throw ConstrainedGenerationError.emptyAnyOf } - if variants.count == 1 { - return try await generateNode(variants[0]) + return try await generateAnyOf(variants) + } + } + + /// Lets the model choose a variant of an `anyOf`. + /// + /// Samples once, without decoding, among the tokens that can start each variant and + /// continues with the variants that can start with the sample, as + /// ``sampleWhetherToCloseEmptyArray(items:)`` does for an array's first item. + /// Object variants all start with `{`, so the keys the model emits tell them apart. + /// Other variants that start alike, like two strings, fall back to the first of them. + private mutating func generateAnyOf(_ variants: [GenerationSchema.Node]) async throws -> String { + let candidates = try flattenedVariants(variants, depth: 0) + guard let first = candidates.first else { + throw ConstrainedGenerationError.emptyAnyOf + } + guard candidates.count > 1 else { + return try await generateNode(first) + } + + let matching = try await variantsStartingWithSample(candidates) + let objects = matching.compactMap { variant -> GenerationSchema.ObjectNode? in + if case .object(let object) = variant { return object } + return nil + } + if objects.count > 1, objects.count == matching.count { + return try await generateObject(oneOf: objects) + } + // Known limitation: variants that start alike but aren't all objects (two strings, + // integer and number, two arrays) can't be told apart by their first token, so the + // first of them is generated. + return try await generateNode(matching[0]) + } + + /// Resolves references and flattens nested unions, keeping declaration order. + private func flattenedVariants(_ variants: [GenerationSchema.Node], depth: Int) throws -> [GenerationSchema.Node] { + guard depth < 64 else { return [] } + var flattened: [GenerationSchema.Node] = [] + for variant in variants { + switch variant { + case .ref(let typeName): + guard let referenced = schema.defs[typeName] else { + throw ConstrainedGenerationError.missingReference(typeName) + } + flattened += try flattenedVariants([referenced], depth: depth + 1) + case .anyOf(let nested): + flattened += try flattenedVariants(nested, depth: depth + 1) + default: + flattened.append(variant) } - // Choose the first variant to keep selection deterministic. - return try await generateNode(variants[0]) } + return flattened + } + + /// Samples among the tokens that can start each variant, without decoding the sample, + /// and returns the variants that can start with it. Returns every variant without + /// sampling when they all start alike. + private mutating func variantsStartingWithSample( + _ variants: [GenerationSchema.Node] + ) async throws -> [GenerationSchema.Node] { + let starts = try variants.map { try itemStartTokens(for: $0) } + guard starts.dropFirst().contains(where: { $0 != starts[0] }) else { + return variants + } + let allowed = starts.reduce(into: Set()) { $0.formUnion($1) } + guard !allowed.isEmpty else { return variants } + let token = try await backend.sample(from: allowed) + let matching = zip(variants, starts).filter { $0.1.contains(token) }.map(\.0) + return matching.isEmpty ? variants : matching } private mutating func generateObject(_ node: GenerationSchema.ObjectNode) async throws -> String { + try await generateObject(oneOf: [node]) + } + + /// Generates an object matching one of `variants`. + /// + /// With several variants, the mask offers the keys of every variant still in play, and + /// each emitted key keeps only the variants that declare it. + private mutating func generateObject(oneOf variants: [GenerationSchema.ObjectNode]) async throws -> String { // Object *key set* is model-driven under the JSON grammar. The previous // implementation pre-filtered optional properties with a hash of the field // name XOR the token budget, so each optional was always-on or always-off @@ -416,14 +487,20 @@ struct ConstrainedJSONGenerator { // `}` once every required property has been emitted. For schemas with no // required properties that makes `{}` reachable (schema-valid, and what the // model asked for) — a visible behaviour change for such schemas. - var remainingKeys = Set(node.properties.keys) - let required = node.required + var live = variants + var emittedKeys = Set() var output = try await emit("{") - var emittedAnyProperty = false - while !remainingKeys.isEmpty { - let missingRequired = required.intersection(remainingKeys) - let canClose = missingRequired.isEmpty + while true { + let remainingKeys = live.reduce(into: Set()) { $0.formUnion($1.properties.keys) } + .subtracting(emittedKeys) + if remainingKeys.isEmpty { break } + let missingRequired = live.reduce(into: Set()) { $0.formUnion($1.required) } + .intersection(remainingKeys) + let canClose = live.contains { variant in + variant.required.allSatisfy { emittedKeys.contains($0) || variant.properties[$0] == nil } + } + let emittedAnyProperty = !emittedKeys.isEmpty let budgetAllowsMoreOptionals = hasBudgetForOptionalStructure() let keysToOffer: [String] @@ -459,20 +536,68 @@ struct ConstrainedJSONGenerator { return output } - guard let key = propertyKey(fromPropertyStart: choice), - let valueNode = node.properties[key] - else { + guard let key = propertyKey(fromPropertyStart: choice) else { throw ConstrainedGenerationError.tokenizationFailed } - remainingKeys.remove(key) - output += try await generateNode(valueNode) - emittedAnyProperty = true + live = live.filter { $0.properties[key] != nil } + guard !live.isEmpty else { + throw ConstrainedGenerationError.tokenizationFailed + } + emittedKeys.insert(key) + output += try await generateValue(forKey: key, narrowing: &live) } output += try await emit("}") return output } + /// Generates the value of `key`, keeping only the variants whose schema for it the value matches. + /// + /// Variants that agree on the key's schema all stay. Variants that name themselves with + /// a string choice, like the `type` of an enum case, are told apart by the choice. + /// Otherwise the value's first token chooses, and the first matching variant stays. + private mutating func generateValue( + forKey key: String, + narrowing live: inout [GenerationSchema.ObjectNode] + ) async throws -> String { + let nodes = live.compactMap { $0.properties[key] } + guard let first = nodes.first else { + throw ConstrainedGenerationError.tokenizationFailed + } + if nodes.allSatisfy({ $0 == first }) { + return try await generateNode(first) + } + + let choices = try nodes.map { try stringChoices(of: $0) } + if choices.allSatisfy({ $0 != nil }) { + var all: [String] = [] + for choice in choices.compactMap({ $0 }).joined() where !all.contains(choice) { + all.append(choice) + } + var output = try await emit("\"") + let value = try await generateChoice(all) + output += value + output += try await emit("\"") + live = zip(live, choices).filter { $0.1?.contains(value) == true }.map(\.0) + return output + } + + // Same limitation as in `generateAnyOf(_:)`: among values that start alike, the first + // variant's schema is used. + let chosen = try await variantsStartingWithSample(nodes)[0] + live = zip(live, nodes).filter { $0.1 == chosen }.map(\.0) + return try await generateNode(chosen) + } + + /// The choices of a string enum without a pattern, through references. + private func stringChoices(of node: GenerationSchema.Node) throws -> [String]? { + let flattened = try flattenedVariants([node], depth: 0) + guard flattened.count == 1, case .string(let string) = flattened[0], string.pattern == nil, + let choices = string.enumChoices, !choices.isEmpty + else { return nil } + return choices + } + private mutating func generateArray(_ node: GenerationSchema.ArrayNode) async throws -> String { // Array *length* is model-driven under the JSON grammar — same family of fix as // model-driven optional object keys. After each element the mask permits both diff --git a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift new file mode 100644 index 00000000..799735f1 --- /dev/null +++ b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift @@ -0,0 +1,133 @@ +import Testing + +@testable import AnyLanguageModel + +/// Constrained generation of `anyOf` schemas: the model, not declaration order, picks the variant. +struct AnyOfGenerationTests { + private let quote = 0 + private let comma = 1 + private let rightBrace = 2 + private let colon = 4 + private let a = 8 + private let b = 9 + private let x = 10 + private let y = 11 + private let leftBrace = 15 + private let null = 70 + private let eos = 50 + + private func generate(_ schema: GenerationSchema, sampling queue: [Int]) async throws -> String { + let tokenToText: [Int: String] = [ + quote: "\"", comma: ",", rightBrace: "}", 3: "]", colon: ":", + a: "a", b: "b", x: "x", y: "y", leftBrace: "{", null: "null", eos: "", + ] + let textToTokens: [String: [Int]] = [ + "\"x\":": [quote, x, quote, colon], + "\"y\":": [quote, y, quote, colon], + ",\"x\":": [comma, quote, x, quote, colon], + ",\"y\":": [comma, quote, y, quote, colon], + ] + let backend = MockTokenBackend( + tokenToText: tokenToText, + textToTokens: textToTokens, + eosToken: eos, + endTokens: [eos], + maximumTokens: 64, + samplingQueue: queue + ) + var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema) + return try await generator.generate() + } + + private func letter(_ choices: [String]) -> GenerationSchema.Node { + .string(.init(enumChoices: choices)) + } + + private func object(_ properties: [String: GenerationSchema.Node]) -> GenerationSchema.Node { + .object(.init(description: nil, properties: properties, required: Set(properties.keys))) + } + + /// The shape of an enum with payloads: every case has the key `x`, whose value names the case. + private func discriminatedUnion() -> GenerationSchema { + GenerationSchema.primitive( + String.self, + node: .anyOf([ + object(["x": letter(["a"])]), + object(["x": letter(["b"]), "y": letter(["a"])]), + ]) + ) + } + + @Test func nullableStringCanBeNull() async throws { + let schema = GenerationSchema.primitive(String.self, node: .anyOf([letter(["a"]), .null])) + #expect(try await generate(schema, sampling: [null]) == "null") + } + + @Test func nullFirstUnionCanBeAString() async throws { + let schema = GenerationSchema.primitive(String.self, node: .anyOf([.null, letter(["a"])])) + #expect(try await generate(schema, sampling: [quote, a]) == "\"a\"") + } + + @Test func nullableObjectFromDynamicSchemaCanBeEitherBranch() async throws { + let person = DynamicGenerationSchema( + name: "Person", + properties: [.init(name: "x", schema: DynamicGenerationSchema(name: "Letter", anyOf: ["a"]))] + ) + let schema = try GenerationSchema( + root: DynamicGenerationSchema(name: "NullablePerson", anyOf: [person, .null]), + dependencies: [] + ) + #expect(try await generate(schema, sampling: [null]) == "null") + #expect(try await generate(schema, sampling: [leftBrace, quote, x, quote, colon, a]) == #"{"x":"a"}"#) + } + + @Test func objectUnionFollowsTheFirstKey() async throws { + let schema = GenerationSchema.primitive( + String.self, + node: .anyOf([object(["x": letter(["a"])]), object(["y": letter(["a"])])]) + ) + #expect(try await generate(schema, sampling: [quote, y, quote, colon, a]) == #"{"y":"a"}"#) + } + + @Test func discriminatedUnionFollowsTheSharedKeysValue() async throws { + let queue = [quote, x, quote, colon, b, comma, quote, y, quote, colon, a] + #expect(try await generate(discriminatedUnion(), sampling: queue) == #"{"x":"b","y":"a"}"#) + } + + @Test func discriminatedUnionAcceptsThePayloadKeyFirst() async throws { + // Only the second case declares `y`, so emitting it first already picks that case. + let queue = [quote, y, quote, colon, a, comma, quote, x, quote, colon, b] + #expect(try await generate(discriminatedUnion(), sampling: queue) == #"{"y":"a","x":"b"}"#) + } + + @Test func discriminatedUnionCanPickThePayloadlessCase() async throws { + #expect(try await generate(discriminatedUnion(), sampling: [quote, x, quote, colon, a]) == #"{"x":"a"}"#) + } + + @Test func constantGuidedDiscriminatorPicksTheCase() async throws { + // Since #304 a `.constant(_:)` guide is a one-choice string enum, so it narrows like any other choice. + let first = DynamicGenerationSchema( + name: "First", + properties: [.init(name: "x", schema: .init(type: String.self, guides: [.constant("a")]))] + ) + let second = DynamicGenerationSchema( + name: "Second", + properties: [ + .init(name: "x", schema: .init(type: String.self, guides: [.constant("b")])), + .init(name: "y", schema: .init(type: String.self, guides: [.constant("a")])), + ] + ) + let schema = try GenerationSchema( + root: DynamicGenerationSchema(name: "Union", anyOf: [first, second]), + dependencies: [] + ) + let queue = [quote, x, quote, colon, b, comma, quote, y, quote, colon, a] + #expect(try await generate(schema, sampling: queue) == #"{"x":"b","y":"a"}"#) + } + + @Test func variantsThatStartAlikeFallBackToTheFirst() async throws { + // Documents the known limitation: two strings both start with `"`, so the first is generated. + let schema = GenerationSchema.primitive(String.self, node: .anyOf([letter(["a"]), letter(["b"])])) + #expect(try await generate(schema, sampling: [a]) == "\"a\"") + } +} From 7f1a9d9db285f2b096da4610d02163edf68e29d3 Mon Sep 17 00:00:00 2001 From: 2570165831 <2570165831@users.noreply.github.com> Date: Sun, 11 Oct 2026 01:21:54 +1100 Subject: [PATCH 2/5] Ignore descriptions when comparing a shared key's variant schemas Node equality includes descriptions, so two object variants whose schemas for a shared key differ only in their descriptions were treated as different: the fallback chose the first and dropped the other, letting declaration order decide the union again. Compare the schemas with descriptions removed, both for the fast path and for narrowing after the sampled value. Co-Authored-By: Claude Opus 5.5 --- .../Shared/StructuredGeneration.swift | 32 +++++++++++++++++-- .../AnyOfGenerationTests.swift | 22 ++++++++++++- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index da2a7a4e..9fef5ba0 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -564,7 +564,10 @@ struct ConstrainedJSONGenerator { guard let first = nodes.first else { throw ConstrainedGenerationError.tokenizationFailed } - if nodes.allSatisfy({ $0 == first }) { + // Descriptions don't change what gets generated, so variants whose schemas for + // this key differ only in their descriptions all stay live. + let comparable = nodes.map(Self.withoutDescriptions) + if comparable.allSatisfy({ $0 == comparable[0] }) { return try await generateNode(first) } @@ -585,10 +588,35 @@ struct ConstrainedJSONGenerator { // Same limitation as in `generateAnyOf(_:)`: among values that start alike, the first // variant's schema is used. let chosen = try await variantsStartingWithSample(nodes)[0] - live = zip(live, nodes).filter { $0.1 == chosen }.map(\.0) + let chosenShape = Self.withoutDescriptions(chosen) + live = zip(live, comparable).filter { $0.1 == chosenShape }.map(\.0) return try await generateNode(chosen) } + /// `node` without any descriptions, so nodes that generate the same values compare equal. + private static func withoutDescriptions(_ node: GenerationSchema.Node) -> GenerationSchema.Node { + switch node { + case .object(var object): + object.description = nil + object.properties = object.properties.mapValues { withoutDescriptions($0) } + return .object(object) + case .array(var array): + array.description = nil + array.items = withoutDescriptions(array.items) + return .array(array) + case .string(var string): + string.description = nil + return .string(string) + case .number(var number): + number.description = nil + return .number(number) + case .anyOf(let variants): + return .anyOf(variants.map { withoutDescriptions($0) }) + case .boolean, .null, .ref: + return node + } + } + /// The choices of a string enum without a pattern, through references. private func stringChoices(of node: GenerationSchema.Node) throws -> [String]? { let flattened = try flattenedVariants([node], depth: 0) diff --git a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift index 799735f1..d5746724 100644 --- a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift @@ -14,14 +14,17 @@ struct AnyOfGenerationTests { private let y = 11 private let leftBrace = 15 private let null = 70 + private let one = 6 private let eos = 50 private func generate(_ schema: GenerationSchema, sampling queue: [Int]) async throws -> String { let tokenToText: [Int: String] = [ quote: "\"", comma: ",", rightBrace: "}", 3: "]", colon: ":", - a: "a", b: "b", x: "x", y: "y", leftBrace: "{", null: "null", eos: "", + a: "a", b: "b", x: "x", y: "y", leftBrace: "{", null: "null", one: "1", eos: "", ] let textToTokens: [String: [Int]] = [ + "\"": [quote], ",": [comma], "}": [rightBrace], ":": [colon], "{": [leftBrace], + "a": [a], "b": [b], "x": [x], "y": [y], "1": [one], "null": [null], "\"x\":": [quote, x, quote, colon], "\"y\":": [quote, y, quote, colon], ",\"x\":": [comma, quote, x, quote, colon], @@ -100,6 +103,23 @@ struct AnyOfGenerationTests { #expect(try await generate(discriminatedUnion(), sampling: queue) == #"{"y":"a","x":"b"}"#) } + @Test func sharedKeysThatDifferOnlyInDescriptionKeepBothVariants() async throws { + // Both variants declare `x` as the same integer except for its description, so + // emitting `x` must not decide the union; the later `y` picks the second variant. + let schema = GenerationSchema.primitive( + String.self, + node: .anyOf([ + object(["x": .number(.init(description: "first", integerOnly: true))]), + object([ + "x": .number(.init(description: "second", integerOnly: true)), + "y": letter(["a"]), + ]), + ]) + ) + let queue = [quote, x, quote, colon, one, comma, comma, quote, y, quote, colon, a] + #expect(try await generate(schema, sampling: queue) == #"{"x":1,"y":"a"}"#) + } + @Test func discriminatedUnionCanPickThePayloadlessCase() async throws { #expect(try await generate(discriminatedUnion(), sampling: [quote, x, quote, colon, a]) == #"{"x":"a"}"#) } From ca562b60c65845dfd49ca951740ec1f9dee2c96d Mon Sep 17 00:00:00 2001 From: 2570165831 <2570165831@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:50:40 +1100 Subject: [PATCH 3/5] Tell anyOf variants apart by nested objects and resolved references When a shared key's values are all objects, generate the value with the multi-variant object generator and keep the variants whose nested object still matches, instead of always keeping the first. Resolve references when comparing a shared key's schemas, stopping at a reference already being resolved so recursive definitions terminate. Keep the live variants when a generated choice matches none of the choices, as when a token decodes with a leading space, so the object does not close before the remaining required keys. Co-Authored-By: Claude Opus 5.5 --- .../Shared/StructuredGeneration.swift | 141 +++++++++++++----- .../AnyOfGenerationTests.swift | 86 ++++++++++- 2 files changed, 184 insertions(+), 43 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index 9fef5ba0..ecaca8a2 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -418,13 +418,13 @@ struct ConstrainedJSONGenerator { return try await generateNode(first) } - let matching = try await variantsStartingWithSample(candidates) + let matching = try await variantsStartingWithSample(candidates).map { candidates[$0] } let objects = matching.compactMap { variant -> GenerationSchema.ObjectNode? in if case .object(let object) = variant { return object } return nil } if objects.count > 1, objects.count == matching.count { - return try await generateObject(oneOf: objects) + return try await generateObject(oneOf: objects).output } // Known limitation: variants that start alike but aren't all objects (two strings, // integer and number, two arrays) can't be told apart by their first token, so the @@ -453,51 +453,57 @@ struct ConstrainedJSONGenerator { } /// Samples among the tokens that can start each variant, without decoding the sample, - /// and returns the variants that can start with it. Returns every variant without - /// sampling when they all start alike. + /// and returns the positions of the variants that can start with it. Returns every + /// position without sampling when the variants all start alike. private mutating func variantsStartingWithSample( _ variants: [GenerationSchema.Node] - ) async throws -> [GenerationSchema.Node] { + ) async throws -> [Int] { let starts = try variants.map { try itemStartTokens(for: $0) } guard starts.dropFirst().contains(where: { $0 != starts[0] }) else { - return variants + return Array(variants.indices) } let allowed = starts.reduce(into: Set()) { $0.formUnion($1) } - guard !allowed.isEmpty else { return variants } + guard !allowed.isEmpty else { return Array(variants.indices) } let token = try await backend.sample(from: allowed) - let matching = zip(variants, starts).filter { $0.1.contains(token) }.map(\.0) - return matching.isEmpty ? variants : matching + let matching = starts.indices.filter { starts[$0].contains(token) } + return matching.isEmpty ? Array(variants.indices) : matching } private mutating func generateObject(_ node: GenerationSchema.ObjectNode) async throws -> String { - try await generateObject(oneOf: [node]) + try await generateObject(oneOf: [node]).output } /// Generates an object matching one of `variants`. /// /// With several variants, the mask offers the keys of every variant still in play, and /// each emitted key keeps only the variants that declare it. - private mutating func generateObject(oneOf variants: [GenerationSchema.ObjectNode]) async throws -> String { + /// + /// - Returns: The object, and the positions in `variants` of the variants it still matches. + private mutating func generateObject( + oneOf variants: [GenerationSchema.ObjectNode] + ) async throws -> (output: String, live: [Int]) { // Object *key set* is model-driven under the JSON grammar. The previous // implementation pre-filtered optional properties with a hash of the field // name XOR the token budget, so each optional was always-on or always-off // for a given budget — decided before the model was consulted. // - // After each property the mask permits any not-yet-emitted key, and permits - // `}` once every required property has been emitted. For schemas with no - // required properties that makes `{}` reachable (schema-valid, and what the - // model asked for) — a visible behaviour change for such schemas. - var live = variants + // After each property the mask permits any not-yet-emitted key of a live + // variant, and permits `}` once some live variant has all of its required + // properties. For schemas with no required properties that makes `{}` + // reachable (schema-valid, and what the model asked for) — a visible + // behaviour change for such schemas. + var live = Array(variants.indices) var emittedKeys = Set() var output = try await emit("{") while true { - let remainingKeys = live.reduce(into: Set()) { $0.formUnion($1.properties.keys) } + let liveVariants = live.map { variants[$0] } + let remainingKeys = liveVariants.reduce(into: Set()) { $0.formUnion($1.properties.keys) } .subtracting(emittedKeys) if remainingKeys.isEmpty { break } - let missingRequired = live.reduce(into: Set()) { $0.formUnion($1.required) } + let missingRequired = liveVariants.reduce(into: Set()) { $0.formUnion($1.required) } .intersection(remainingKeys) - let canClose = live.contains { variant in + let canClose = liveVariants.contains { variant in variant.required.allSatisfy { emittedKeys.contains($0) || variant.properties[$0] == nil } } let emittedAnyProperty = !emittedKeys.isEmpty @@ -520,8 +526,8 @@ struct ConstrainedJSONGenerator { let prefix = emittedAnyProperty ? "," : "" return "\(prefix)\"\(key)\":" } - // Closing is legal once every required property has been emitted. The model - // may leave remaining optionals out; that is schema-valid JSON. + // Closing is legal once some live variant has all of its required properties. + // The model may leave remaining optionals out; that is schema-valid JSON. if canClose { candidates.append("}") } @@ -531,42 +537,45 @@ struct ConstrainedJSONGenerator { let choice = try await generateChoice(candidates) output += choice - // Model elected to close; remaining keys are optional and intentionally omitted. + // Model elected to close; the keys left out are optional for a live variant. if choice == "}" { - return output + return (output, live) } guard let key = propertyKey(fromPropertyStart: choice) else { throw ConstrainedGenerationError.tokenizationFailed } - live = live.filter { $0.properties[key] != nil } + live = live.filter { variants[$0].properties[key] != nil } guard !live.isEmpty else { throw ConstrainedGenerationError.tokenizationFailed } emittedKeys.insert(key) - output += try await generateValue(forKey: key, narrowing: &live) + output += try await generateValue(forKey: key, of: variants, narrowing: &live) } output += try await emit("}") - return output + return (output, live) } - /// Generates the value of `key`, keeping only the variants whose schema for it the value matches. + /// Generates the value of `key`, keeping in `live` only the variants whose schema for + /// it the value matches. /// /// Variants that agree on the key's schema all stay. Variants that name themselves with /// a string choice, like the `type` of an enum case, are told apart by the choice. - /// Otherwise the value's first token chooses, and the first matching variant stays. + /// Variants whose values are all objects are told apart by the keys of the nested + /// object. Otherwise the value's first token chooses, and the first matching variant stays. private mutating func generateValue( forKey key: String, - narrowing live: inout [GenerationSchema.ObjectNode] + of variants: [GenerationSchema.ObjectNode], + narrowing live: inout [Int] ) async throws -> String { - let nodes = live.compactMap { $0.properties[key] } - guard let first = nodes.first else { + let nodes = live.compactMap { variants[$0].properties[key] } + guard let first = nodes.first, nodes.count == live.count else { throw ConstrainedGenerationError.tokenizationFailed } // Descriptions don't change what gets generated, so variants whose schemas for // this key differ only in their descriptions all stay live. - let comparable = nodes.map(Self.withoutDescriptions) + let comparable = nodes.map { withoutDescriptions($0) } if comparable.allSatisfy({ $0 == comparable[0] }) { return try await generateNode(first) } @@ -581,28 +590,75 @@ struct ConstrainedJSONGenerator { let value = try await generateChoice(all) output += value output += try await emit("\"") - live = zip(live, choices).filter { $0.1?.contains(value) == true }.map(\.0) + // The value can differ from every choice when a token's text differs from the + // text it was tokenized from, like a piece with a leading space. Narrowing to + // no variant would close the object without the remaining required keys, so + // the live variants stay as they were. + let matching = zip(live, choices).filter { $0.1?.contains(value) == true }.map(\.0) + if !matching.isEmpty { + live = matching + } + return output + } + + let starting = try await variantsStartingWithSample(nodes) + if let objects = try distinctObjects(in: starting.map { nodes[$0] }) { + // Every value is an object and they all start with `{`, so the nested keys + // choose: the variants stay whose nested object the generated one still matches. + let (output, survivors) = try await generateObject(oneOf: objects.map(\.object)) + let kept = Set(survivors.flatMap { objects[$0].owners }) + live = kept.sorted().map { live[starting[$0]] } return output } // Same limitation as in `generateAnyOf(_:)`: among values that start alike, the first // variant's schema is used. - let chosen = try await variantsStartingWithSample(nodes)[0] - let chosenShape = Self.withoutDescriptions(chosen) + let chosen = nodes[starting[0]] + let chosenShape = comparable[starting[0]] live = zip(live, comparable).filter { $0.1 == chosenShape }.map(\.0) return try await generateNode(chosen) } + /// The distinct objects that `nodes` resolve to, each with the positions in `nodes` that + /// can be it, or `nil` when a node can be something other than an object. + private func distinctObjects( + in nodes: [GenerationSchema.Node] + ) throws -> [(object: GenerationSchema.ObjectNode, owners: [Int])]? { + var objects: [(object: GenerationSchema.ObjectNode, shape: GenerationSchema.Node, owners: [Int])] = [] + for (position, node) in nodes.enumerated() { + let flattened = try flattenedVariants([node], depth: 0) + guard !flattened.isEmpty else { return nil } + for variant in flattened { + guard case .object(let object) = variant else { return nil } + let shape = withoutDescriptions(variant) + if let index = objects.firstIndex(where: { $0.shape == shape }) { + if !objects[index].owners.contains(position) { + objects[index].owners.append(position) + } + } else { + objects.append((object, shape, [position])) + } + } + } + return objects.map { ($0.object, $0.owners) } + } + /// `node` without any descriptions, so nodes that generate the same values compare equal. - private static func withoutDescriptions(_ node: GenerationSchema.Node) -> GenerationSchema.Node { + /// + /// References are resolved, except one already being resolved further up, which stays a + /// reference so a recursive definition can't expand forever. + private func withoutDescriptions( + _ node: GenerationSchema.Node, + resolving: Set = [] + ) -> GenerationSchema.Node { switch node { case .object(var object): object.description = nil - object.properties = object.properties.mapValues { withoutDescriptions($0) } + object.properties = object.properties.mapValues { withoutDescriptions($0, resolving: resolving) } return .object(object) case .array(var array): array.description = nil - array.items = withoutDescriptions(array.items) + array.items = withoutDescriptions(array.items, resolving: resolving) return .array(array) case .string(var string): string.description = nil @@ -611,8 +667,13 @@ struct ConstrainedJSONGenerator { number.description = nil return .number(number) case .anyOf(let variants): - return .anyOf(variants.map { withoutDescriptions($0) }) - case .boolean, .null, .ref: + return .anyOf(variants.map { withoutDescriptions($0, resolving: resolving) }) + case .ref(let typeName): + guard !resolving.contains(typeName), let referenced = schema.defs[typeName] else { + return node + } + return withoutDescriptions(referenced, resolving: resolving.union([typeName])) + case .boolean, .null: return node } } diff --git a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift index d5746724..ae86e97b 100644 --- a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift @@ -1,3 +1,4 @@ +import Foundation import Testing @testable import AnyLanguageModel @@ -17,11 +18,17 @@ struct AnyOfGenerationTests { private let one = 6 private let eos = 50 - private func generate(_ schema: GenerationSchema, sampling queue: [Int]) async throws -> String { + /// - Parameter decodedText: Overrides the text a token decodes to, as for a piece with a + /// leading space that still tokenizes from the bare text. + private func generate( + _ schema: GenerationSchema, + sampling queue: [Int], + decodedText: [Int: String] = [:] + ) async throws -> String { let tokenToText: [Int: String] = [ quote: "\"", comma: ",", rightBrace: "}", 3: "]", colon: ":", a: "a", b: "b", x: "x", y: "y", leftBrace: "{", null: "null", one: "1", eos: "", - ] + ].merging(decodedText) { $1 } let textToTokens: [String: [Int]] = [ "\"": [quote], ",": [comma], "}": [rightBrace], ":": [colon], "{": [leftBrace], "a": [a], "b": [b], "x": [x], "y": [y], "1": [one], "null": [null], @@ -121,7 +128,80 @@ struct AnyOfGenerationTests { } @Test func discriminatedUnionCanPickThePayloadlessCase() async throws { - #expect(try await generate(discriminatedUnion(), sampling: [quote, x, quote, colon, a]) == #"{"x":"a"}"#) + // The payloadless case comes second, so declaration order alone can't pick it. + let schema = GenerationSchema.primitive( + String.self, + node: .anyOf([ + object(["x": letter(["b"]), "y": letter(["a"])]), + object(["x": letter(["a"])]), + ]) + ) + #expect(try await generate(schema, sampling: [quote, x, quote, colon, a]) == #"{"x":"a"}"#) + } + + @Test func choiceThatDecodesDifferentlyKeepsTheRemainingKeys() async throws { + // The token for "b" decodes as " b", so the value matches no choice exactly. The + // object must not close before `y`, which the case named "b" requires. + let queue = [quote, x, quote, colon, b, comma, quote, y, quote, colon, a] + let output = try await generate(discriminatedUnion(), sampling: queue, decodedText: [b: " b"]) + #expect(output == #"{"x":" b","y":"a"}"#) + } + + @Test func sharedKeyWithObjectValuesFollowsTheNestedKeys() async throws { + // Both values of `x` start with `{`, so the nested key picks the second variant, + // which then requires `y`. + let schema = GenerationSchema.primitive( + String.self, + node: .anyOf([ + object(["x": object(["x": letter(["a"])])]), + object(["x": object(["y": letter(["a"])]), "y": letter(["a"])]), + ]) + ) + let queue = [quote, x, quote, colon, quote, y, quote, colon, a, comma, quote, y, quote, colon, a] + #expect(try await generate(schema, sampling: queue) == #"{"x":{"y":"a"},"y":"a"}"#) + } + + @Test func sharedKeysWhoseReferencesDifferOnlyInDescriptionKeepBothVariants() async throws { + // `x` refers to `A` in one variant and `B` in the other, two integers that differ + // only in their descriptions, so emitting `x` must not decide the union. + let json = #""" + { + "$defs": { + "A": {"type": "integer", "description": "first"}, + "B": {"type": "integer", "description": "second"} + }, + "anyOf": [ + {"type": "object", "properties": {"x": {"$ref": "#/$defs/A"}}, "required": ["x"]}, + { + "type": "object", + "properties": {"x": {"$ref": "#/$defs/B"}, "y": {"type": "string", "enum": ["a"]}}, + "required": ["x", "y"] + } + ] + } + """# + let schema = try JSONDecoder().decode(GenerationSchema.self, from: Data(json.utf8)) + let queue = [quote, x, quote, colon, one, comma, comma, quote, y, quote, colon, a] + #expect(try await generate(schema, sampling: queue) == #"{"x":1,"y":"a"}"#) + } + + @Test func selfReferencingSharedKeysStillGenerate() async throws { + // `A` and `B` each refer to themselves, so comparing them must stop expanding. + let json = #""" + { + "$defs": { + "A": {"type": "object", "properties": {"x": {"$ref": "#/$defs/A"}}, "required": []}, + "B": {"type": "object", "properties": {"y": {"$ref": "#/$defs/B"}}, "required": []} + }, + "anyOf": [ + {"type": "object", "properties": {"x": {"$ref": "#/$defs/A"}}, "required": ["x"]}, + {"type": "object", "properties": {"x": {"$ref": "#/$defs/B"}}, "required": ["x"]} + ] + } + """# + let schema = try JSONDecoder().decode(GenerationSchema.self, from: Data(json.utf8)) + let queue = [quote, x, quote, colon, rightBrace] + #expect(try await generate(schema, sampling: queue) == #"{"x":{}}"#) } @Test func constantGuidedDiscriminatorPicksTheCase() async throws { From 02143b4150f52e4fbf9a2be4d7772f606e6a5172 Mon Sep 17 00:00:00 2001 From: 2570165831 <2570165831@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:48:17 +1100 Subject: [PATCH 4/5] Narrow a discriminator by the sampled choice, not its decoded text When the tokens sampled for a string choice decode to text that is not exactly one of the choices (a piece with a leading space, say), the live variants used to stay as they were. A variant that the choice ruled out could then let the object close before the chosen case's required keys. generateChoice now reports which candidate the sampled tokens spell, and the value path narrows by that candidate. Co-Authored-By: Claude Opus 5.5 --- .../Shared/StructuredGeneration.swift | 43 ++++++++++++------- .../AnyOfGenerationTests.swift | 9 ++++ 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index ecaca8a2..60f19494 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -244,6 +244,18 @@ struct ConstrainedJSONGenerator { } private mutating func generateChoice(_ candidates: [String]) async throws -> String { + try await generateIndexedChoice(candidates).text + } + + /// Generates one of `candidates`. + /// + /// - Returns: The text the sampled tokens decode to, and the position in `candidates` of + /// the choice they spell, or `nil` when the budget runs out before one is complete. + /// The text can differ from the choice when a token's text differs from the text it + /// was tokenized from, like a piece with a leading space; the position doesn't. + private mutating func generateIndexedChoice( + _ candidates: [String] + ) async throws -> (text: String, index: Int?) { guard !candidates.isEmpty else { throw ConstrainedGenerationError.tokenizationFailed } @@ -263,20 +275,20 @@ struct ConstrainedJSONGenerator { if !chosen.isEmpty { _ = try await emit(chosen) } - return chosen + return (chosen, candidates.firstIndex(of: chosen)) } - var prefixes = tokenized + var prefixes = Array(tokenized.indices) var emitted = "" var position = 0 while backend.remainingTokens > 0 { - if prefixes.contains(where: { $0.count == position }) { break } + if prefixes.contains(where: { tokenized[$0].count == position }) { break } let allowed = Set( - prefixes.compactMap { tokens -> Int? in - guard position < tokens.count else { return nil } - return tokens[position] + prefixes.compactMap { index -> Int? in + guard position < tokenized[index].count else { return nil } + return tokenized[index][position] } ) @@ -289,12 +301,12 @@ struct ConstrainedJSONGenerator { emittedText += text try await backend.decode(token) - prefixes = prefixes.filter { $0.count > position && $0[position] == token } + prefixes = prefixes.filter { tokenized[$0].count > position && tokenized[$0][position] == token } position += 1 if prefixes.isEmpty { break } } - return emitted + return (emitted, prefixes.first { tokenized[$0].count == position }) } private func maxNumberTokens(for node: GenerationSchema.NumberNode) -> Int { @@ -587,16 +599,15 @@ struct ConstrainedJSONGenerator { all.append(choice) } var output = try await emit("\"") - let value = try await generateChoice(all) + let (value, index) = try await generateIndexedChoice(all) output += value output += try await emit("\"") - // The value can differ from every choice when a token's text differs from the - // text it was tokenized from, like a piece with a leading space. Narrowing to - // no variant would close the object without the remaining required keys, so - // the live variants stay as they were. - let matching = zip(live, choices).filter { $0.1?.contains(value) == true }.map(\.0) - if !matching.isEmpty { - live = matching + // Narrow by the choice the sampled tokens spell, not by the text they decode to: + // the two differ when a token's text differs from the text it was tokenized + // from, like a piece with a leading space. Without a choice, the budget ran out + // and the live variants stay as they were. + if let index { + live = zip(live, choices).filter { $0.1?.contains(all[index]) == true }.map(\.0) } return output } diff --git a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift index ae86e97b..4a26e71f 100644 --- a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift @@ -147,6 +147,15 @@ struct AnyOfGenerationTests { #expect(output == #"{"x":" b","y":"a"}"#) } + @Test func choiceThatDecodesDifferentlyCannotCloseBeforeTheRemainingKeys() async throws { + // The sampled token is still the choice "b", so its case stays the only one live and + // `}` isn't offered before `y`. Scripting `}` there must be refused, not followed. + let queue = [quote, x, quote, colon, b, rightBrace] + await #expect(throws: ConstrainedGenerationError.self) { + try await generate(discriminatedUnion(), sampling: queue, decodedText: [b: " b"]) + } + } + @Test func sharedKeyWithObjectValuesFollowsTheNestedKeys() async throws { // Both values of `x` start with `{`, so the nested key picks the second variant, // which then requires `y`. From 829eb0763340db923e8f88844393c9ff6acfb8e9 Mon Sep 17 00:00:00 2001 From: 2570165831 <2570165831@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:37:12 +1100 Subject: [PATCH 5/5] Bound reference expansion in anyOf flattening and comparison A union that refers to the same definition twice, at each level, made `flattenedVariants` double its result per level: the depth guard capped the chain length, not the size. Each definition is now expanded once per flattening, which also drops the depth guard. `withoutDescriptions` had the same shape for definitions that refer to another one twice. It now resolves at most 256 references per call and leaves the rest as references, which can only make two schemas compare unequal sooner. Properties are visited in key order, so the cut-off is the same for equal schemas. Co-Authored-By: Claude Opus 5.5 --- .../Shared/StructuredGeneration.swift | 65 +++++++++++++++---- .../AnyOfGenerationTests.swift | 48 ++++++++++++++ 2 files changed, 99 insertions(+), 14 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift index 60f19494..8b35da22 100644 --- a/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift +++ b/Sources/AnyLanguageModel/Shared/StructuredGeneration.swift @@ -422,7 +422,7 @@ struct ConstrainedJSONGenerator { /// Object variants all start with `{`, so the keys the model emits tell them apart. /// Other variants that start alike, like two strings, fall back to the first of them. private mutating func generateAnyOf(_ variants: [GenerationSchema.Node]) async throws -> String { - let candidates = try flattenedVariants(variants, depth: 0) + let candidates = try flattenedVariants(variants) guard let first = candidates.first else { throw ConstrainedGenerationError.emptyAnyOf } @@ -445,8 +445,20 @@ struct ConstrainedJSONGenerator { } /// Resolves references and flattens nested unions, keeping declaration order. - private func flattenedVariants(_ variants: [GenerationSchema.Node], depth: Int) throws -> [GenerationSchema.Node] { - guard depth < 64 else { return [] } + /// + /// Each definition is expanded at most once: a later reference to it, including one + /// from inside itself, would only repeat variants already in the list. That keeps the + /// result linear in the number of definitions, where expanding every reference would + /// double the list at each union that refers to the same definition twice. + private func flattenedVariants(_ variants: [GenerationSchema.Node]) throws -> [GenerationSchema.Node] { + var expanded = Set() + return try flattenedVariants(variants, expanded: &expanded) + } + + private func flattenedVariants( + _ variants: [GenerationSchema.Node], + expanded: inout Set + ) throws -> [GenerationSchema.Node] { var flattened: [GenerationSchema.Node] = [] for variant in variants { switch variant { @@ -454,9 +466,10 @@ struct ConstrainedJSONGenerator { guard let referenced = schema.defs[typeName] else { throw ConstrainedGenerationError.missingReference(typeName) } - flattened += try flattenedVariants([referenced], depth: depth + 1) + guard expanded.insert(typeName).inserted else { continue } + flattened += try flattenedVariants([referenced], expanded: &expanded) case .anyOf(let nested): - flattened += try flattenedVariants(nested, depth: depth + 1) + flattened += try flattenedVariants(nested, expanded: &expanded) default: flattened.append(variant) } @@ -637,7 +650,7 @@ struct ConstrainedJSONGenerator { ) throws -> [(object: GenerationSchema.ObjectNode, owners: [Int])]? { var objects: [(object: GenerationSchema.ObjectNode, shape: GenerationSchema.Node, owners: [Int])] = [] for (position, node) in nodes.enumerated() { - let flattened = try flattenedVariants([node], depth: 0) + let flattened = try flattenedVariants([node]) guard !flattened.isEmpty else { return nil } for variant in flattened { guard case .object(let object) = variant else { return nil } @@ -654,22 +667,41 @@ struct ConstrainedJSONGenerator { return objects.map { ($0.object, $0.owners) } } + /// How many references one ``withoutDescriptions(_:)`` call resolves before it leaves the + /// rest as references. + private static var comparableReferenceBudget: Int { 256 } + /// `node` without any descriptions, so nodes that generate the same values compare equal. /// /// References are resolved, except one already being resolved further up, which stays a - /// reference so a recursive definition can't expand forever. + /// reference so a recursive definition can't expand forever. Resolution also stops after + /// ``comparableReferenceBudget`` references: a definition that refers to another one twice, + /// at each level, would otherwise expand exponentially. A reference left unresolved only + /// compares equal to the same reference, so two schemas then differ sooner, never later. + private func withoutDescriptions(_ node: GenerationSchema.Node) -> GenerationSchema.Node { + var budget = Self.comparableReferenceBudget + return withoutDescriptions(node, resolving: [], budget: &budget) + } + private func withoutDescriptions( _ node: GenerationSchema.Node, - resolving: Set = [] + resolving: Set, + budget: inout Int ) -> GenerationSchema.Node { switch node { case .object(var object): object.description = nil - object.properties = object.properties.mapValues { withoutDescriptions($0, resolving: resolving) } + for key in object.properties.keys.sorted() { + object.properties[key] = withoutDescriptions( + object.properties[key]!, + resolving: resolving, + budget: &budget + ) + } return .object(object) case .array(var array): array.description = nil - array.items = withoutDescriptions(array.items, resolving: resolving) + array.items = withoutDescriptions(array.items, resolving: resolving, budget: &budget) return .array(array) case .string(var string): string.description = nil @@ -678,12 +710,17 @@ struct ConstrainedJSONGenerator { number.description = nil return .number(number) case .anyOf(let variants): - return .anyOf(variants.map { withoutDescriptions($0, resolving: resolving) }) + var stripped: [GenerationSchema.Node] = [] + for variant in variants { + stripped.append(withoutDescriptions(variant, resolving: resolving, budget: &budget)) + } + return .anyOf(stripped) case .ref(let typeName): - guard !resolving.contains(typeName), let referenced = schema.defs[typeName] else { + guard budget > 0, !resolving.contains(typeName), let referenced = schema.defs[typeName] else { return node } - return withoutDescriptions(referenced, resolving: resolving.union([typeName])) + budget -= 1 + return withoutDescriptions(referenced, resolving: resolving.union([typeName]), budget: &budget) case .boolean, .null: return node } @@ -691,7 +728,7 @@ struct ConstrainedJSONGenerator { /// The choices of a string enum without a pattern, through references. private func stringChoices(of node: GenerationSchema.Node) throws -> [String]? { - let flattened = try flattenedVariants([node], depth: 0) + let flattened = try flattenedVariants([node]) guard flattened.count == 1, case .string(let string) = flattened[0], string.pattern == nil, let choices = string.enumChoices, !choices.isEmpty else { return nil } diff --git a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift index 4a26e71f..b71dfd70 100644 --- a/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift +++ b/Tests/AnyLanguageModelTests/AnyOfGenerationTests.swift @@ -213,6 +213,54 @@ struct AnyOfGenerationTests { #expect(try await generate(schema, sampling: queue) == #"{"x":{}}"#) } + @Test(.timeLimit(.minutes(1))) + func unionsThatReferToADefinitionTwiceFlattenOnce() async throws { + // Each `D` is a union of the next one, twice, so expanding every reference would + // list 2^64 variants. Expanding each definition once leaves the integer and string. + var defs = (0 ..< 64).map { level in + let next = ##"{"$ref": "#/$defs/D\##(level + 1)"}"## + return #""D\#(level)": {"anyOf": [\#(next), \#(next)]}"# + } + defs.append(#""D64": {"type": "integer"}"#) + let json = #""" + { + "$defs": {\#(defs.joined(separator: ","))}, + "anyOf": [{"$ref": "#/$defs/D0"}, {"type": "string"}] + } + """# + let schema = try JSONDecoder().decode(GenerationSchema.self, from: Data(json.utf8)) + let queue = [quote, a, quote] + #expect(try await generate(schema, sampling: queue) == #""a""#) + } + + @Test(.timeLimit(.minutes(1))) + func sharedKeysOverADefinitionUsedTwicePerLevelStillCompare() async throws { + // Each `N` refers to the next one twice, so resolving every reference to compare + // the two `x` schemas would build 2^64 nodes. + var defs = (0 ..< 64).map { level in + let next = ##"{"$ref": "#/$defs/N\##(level + 1)"}"## + let properties = #"{"x": \#(next), "y": \#(next)}"# + return #""N\#(level)": {"type": "object", "properties": \#(properties), "required": []}"# + } + defs.append(#""N64": {"type": "object", "properties": {}, "required": []}"#) + let json = #""" + { + "$defs": {\#(defs.joined(separator: ","))}, + "anyOf": [ + {"type": "object", "properties": {"x": {"$ref": "#/$defs/N0"}}, "required": ["x"]}, + { + "type": "object", + "properties": {"x": {"$ref": "#/$defs/N0"}, "y": {"type": "string"}}, + "required": ["x", "y"] + } + ] + } + """# + let schema = try JSONDecoder().decode(GenerationSchema.self, from: Data(json.utf8)) + let queue = [quote, x, quote, colon, rightBrace, rightBrace] + #expect(try await generate(schema, sampling: queue) == #"{"x":{}}"#) + } + @Test func constantGuidedDiscriminatorPicksTheCase() async throws { // Since #304 a `.constant(_:)` guide is a one-choice string enum, so it narrows like any other choice. let first = DynamicGenerationSchema(