From 4284048d1aca2799e543a7400b3f07519befb515 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:24:31 +0700 Subject: [PATCH 1/7] The JSON writer escapes a UTF-8 sequence cut off at a string's end as U+FFFD instead of slicing past it. --- zig/src/server/json_write.zig | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/zig/src/server/json_write.zig b/zig/src/server/json_write.zig index 94ce0ccc62..d880ad126c 100644 --- a/zig/src/server/json_write.zig +++ b/zig/src/server/json_write.zig @@ -68,8 +68,10 @@ pub fn writeString(w: *Writer, s: []const u8, o: Options) Writer.Error!void { try w.print("\\u{x:0>4}", .{c}); i += 1; } else { - const n = std.unicode.utf8ByteSequenceLength(c) catch 1; - const code = decodeWtf8(s[i..@min(s.len, i + n)]); + // a sequence cut short at the string's end is U+FFFD and consumes only the bytes that are there + const full = std.unicode.utf8ByteSequenceLength(c) catch 1; + const n = @min(full, s.len - i); + const code = if (n < full) 0xfffd else decodeWtf8(s[i .. i + n]); if (code >= 0x10000) { const v = code - 0x10000; try w.print("\\u{x:0>4}\\u{x:0>4}", .{ 0xd800 + (v >> 10), 0xdc00 + (v & 0x3ff) }); @@ -171,3 +173,14 @@ test "dumps" { try std.testing.expectEqualStrings("{\"a\": [1, 2.5, \"\\u00e9\\n\\ud83d\\ude00\\u007f\"], \"b\": {}, \"c\": 0}", try json.stringify(a, r.ok, .{})); try std.testing.expectEqualStrings("{\"a\":[1,2.5,\"\xc3\xa9\\n\xf0\x9f\x98\x80\x7f\"],\"b\":{},\"c\":0}", try json.stringify(a, r.ok, .{ .ascii = false, .compact = true })); } + +test "a UTF-8 sequence cut short at the string's end is U+FFFD, not a slice past it" { + var buf: [64]u8 = undefined; + for ([_]struct { []const u8, []const u8 }{ + .{ "ab\xE2", "\"ab\\ufffd\"" }, .{ "ab\xE2\x82", "\"ab\\ufffd\"" }, .{ "\xF0\x9F\x98", "\"\\ufffd\"" }, .{ "x\xE2\x82\xAC", "\"x\\u20ac\"" }, + }) |c| { + var w: Writer = .fixed(&buf); + try writeString(&w, c[0], .{}); + try std.testing.expectEqualStrings(c[1], w.buffered()); + } +} From 986c996041e25e8b4a0341a8b54524a9d2248c12 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:28:21 +0700 Subject: [PATCH 2/7] CUDA keyed graph cache and device arena, both off until a family installs them. graph_cache.Cache captures one executable graph per caller-named key and replays it when the key comes back; ranks agree on every capture, eviction and fallback through an Agree hook, and a memory floor turns misses eager. arena.Arena sub-allocates DeviceBuffer.alloc from a few large chunks (VMM or plain cuMemAlloc) through memory.setHook; borrowed spans and pointers it does not own are freed as before. --- zig/src/cuda/abi.zig | 2 + zig/src/cuda/arena.zig | 331 +++++++++++++++++++++++++++++ zig/src/cuda/graph_cache.zig | 392 +++++++++++++++++++++++++++++++++++ zig/src/cuda/memory.zig | 52 ++++- zig/src/cuda/root.zig | 5 + 5 files changed, 779 insertions(+), 3 deletions(-) create mode 100644 zig/src/cuda/arena.zig create mode 100644 zig/src/cuda/graph_cache.zig diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index a608fd6114..95c9291dd0 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -39,6 +39,8 @@ pub const DeviceAttribute = enum(c_int) { compute_capability_minor = 76, max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, + virtual_memory_management_supported = 102, + gpu_direct_rdma_with_cuda_vmm_supported = 126, }; pub const FunctionAttribute = enum(c_int) { diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..2c80b7581b --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,331 @@ +//! A device sub-allocator under DeviceBuffer.alloc: buffers carved from a few large chunks, not one cuMemAlloc each. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); + +pub const Error = error{ DriverUnavailable, MissingSymbol, Unsupported, CudaFailed, OutOfMemory }; + +/// Where chunks come from: `chunk` maps `len` bytes (a multiple of `granularity`), `release` gives a chunk back. +pub const Backend = struct { + ctx: *anyopaque, + granularity: usize, + chunk: *const fn (ctx: *anyopaque, len: usize) ?u64, + release: *const fn (ctx: *anyopaque, base: u64, len: usize) void, + name: []const u8, +}; + +pub const Options = struct { + small_max: usize = 4 << 20, // requests up to this many bytes share chunks + small_chunk: usize = 64 << 20, // a shared chunk's bytes, rounded up to the granularity + alignment: usize = 512, // every range starts at a multiple of this +}; + +pub const Stats = struct { small: u64 = 0, large: u64 = 0, reused: u64 = 0, fallback: u64 = 0, chunk_bytes: u64 = 0 }; + +/// Small requests bump 512-byte aligned through shared chunks and are reused by size; large ones get their own chunk. +pub const Arena = struct { + gpa: std.mem.Allocator, + be: Backend, + o: Options, + mutex: std.atomic.Mutex = .unlocked, + small: std.ArrayList([2]u64) = .empty, // shared chunks (base, len) + cursor: u64 = 0, // the current shared chunk's next free byte + limit: u64 = 0, + free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size + live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size + large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { + return .{ .gpa = gpa, .be = be, .o = o }; + } + + pub fn deinit(a: *Arena) void { + var it = a.large.iterator(); + while (it.next()) |e| a.be.release(a.be.ctx, e.key_ptr.*, e.value_ptr.*); + for (a.small.items) |c| a.be.release(a.be.ctx, c[0], c[1]); + var f = a.free_small.valueIterator(); + while (f.next()) |l| l.deinit(a.gpa); + a.free_small.deinit(a.gpa); + a.live.deinit(a.gpa); + a.large.deinit(a.gpa); + a.small.deinit(a.gpa); + a.* = undefined; + } + + fn lock(a: *Arena) void { + while (!a.mutex.tryLock()) std.Thread.yield() catch {}; + } + + fn roundUp(x: usize, m: usize) usize { + return std.mem.alignForward(usize, x, m); + } + + /// An address for `len` bytes, or null (the caller then allocates as before). + pub fn alloc(a: *Arena, len: usize) ?u64 { + if (len == 0) return null; + a.lock(); + defer a.mutex.unlock(); + return a.allocLocked(len) catch null; + } + + fn allocLocked(a: *Arena, len: usize) !?u64 { + if (len > a.o.small_max) { + const n = roundUp(len, a.be.granularity); + const p = a.be.chunk(a.be.ctx, n) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, p, n); + try a.large.put(a.gpa, p, n); + a.stats.large += 1; + a.stats.chunk_bytes += n; + return p; + } + const n: u64 = roundUp(len, a.o.alignment); + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + try a.live.put(a.gpa, p, n); + a.stats.reused += 1; + return p; + }; + if (a.limit - a.cursor < n) { + const cl = roundUp(@max(a.o.small_chunk, n), a.be.granularity); + const base = a.be.chunk(a.be.ctx, cl) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, base, cl); + try a.small.append(a.gpa, .{ base, cl }); + a.cursor = base; + a.limit = base + cl; + a.stats.chunk_bytes += cl; + } + const p = a.cursor; + try a.live.put(a.gpa, p, n); + a.cursor += n; + a.stats.small += 1; + return p; + } + + /// True when `p` was the arena's (and is now free); false: not ours, the caller frees it as before. + pub fn free(a: *Arena, p: u64) bool { + a.lock(); + defer a.mutex.unlock(); + if (a.large.fetchRemove(p)) |e| { + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours + if (!g.found_existing) g.value_ptr.* = .empty; + g.value_ptr.append(a.gpa, p) catch {}; + return true; + } + + // -- memory.zig's hook ---------------------------------------------------------------------------------------- + + fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.alloc(len); + } + + fn hookFree(ctx: *anyopaque, p: u64) bool { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.free(p); + } + + /// Every later DeviceBuffer.alloc comes from the arena, which must outlive its buffers; only addresses change. + pub fn install(a: *Arena) void { + memory.setHook(.{ .ctx = a, .alloc = hookAlloc, .free = hookFree }); + } +}; + +// -- backends --------------------------------------------------------------------------------------------------------- + +/// One cuMemAlloc a chunk (the driver's own 2 MiB-aligned allocations). +pub const Plain = struct { + d: *const Driver, + + pub fn backend(p: *Plain) Backend { + return .{ .ctx = p, .granularity = 2 << 20, .chunk = chunk, .release = release, .name = "plain" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const p: *Plain = @ptrCast(@alignCast(ctx)); + var q: abi.DevicePtr = 0; + if (p.d.api.cuMemAlloc_v2(&q, len) != abi.success) return null; + return q; + } + + fn release(ctx: *anyopaque, base: u64, _: usize) void { + const p: *Plain = @ptrCast(@alignCast(ctx)); + _ = p.d.api.cuMemFree_v2(base); + } +}; + +/// CUDA's virtual memory management: one reserved virtual range, chunks of the recommended granularity mapped in order. +pub const Vmm = struct { + lib: std.DynLib, + device: abi.Device, + prop: Prop, + granularity: usize, + va: u64 = 0, + va_len: usize = 0, + next: u64 = 0, + api: Api, + + const Api = struct { + reserve: *const fn (*u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + addressFree: *const fn (u64, usize) callconv(.c) abi.Result, + create: *const fn (*u64, usize, *const Prop, c_ulonglong) callconv(.c) abi.Result, + release: *const fn (u64) callconv(.c) abi.Result, + map: *const fn (u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + unmap: *const fn (u64, usize) callconv(.c) abi.Result, + setAccess: *const fn (u64, usize, *const AccessDesc, usize) callconv(.c) abi.Result, + granularity: *const fn (*usize, *const Prop, c_int) callconv(.c) abi.Result, + }; + + pub const Prop = extern struct { // CUmemAllocationProp + type: c_int = 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: c_int = 0, + loc_type: c_int = 1, // CU_MEM_LOCATION_TYPE_DEVICE + loc_id: c_int = 0, + win32: ?*anyopaque = null, + compression: u8 = 0, + rdma: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = .{ 0, 0, 0, 0 }, + }; + const AccessDesc = extern struct { loc_type: c_int = 1, loc_id: c_int, flags: c_int = 3 }; // CUmemAccessDesc, read-write + + comptime { + std.debug.assert(@sizeOf(Prop) == 32 and @offsetOf(Prop, "compression") == 24 and @sizeOf(AccessDesc) == 12); + } + + /// `va_bytes` of virtual addresses reserved (cheap: no memory behind them until a chunk maps). + pub fn open(d: *const Driver, device: abi.Device, va_bytes: usize) Error!Vmm { + var vmm_ok: c_int = 0; + if (d.api.cuDeviceGetAttribute(&vmm_ok, .virtual_memory_management_supported, device) != abi.success or vmm_ok == 0) return error.Unsupported; // VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + var api: Api = undefined; + const names = .{ .{ "reserve", "cuMemAddressReserve" }, .{ "addressFree", "cuMemAddressFree" }, .{ "create", "cuMemCreate" }, .{ "release", "cuMemRelease" }, .{ "map", "cuMemMap" }, .{ "unmap", "cuMemUnmap" }, .{ "setAccess", "cuMemSetAccess" }, .{ "granularity", "cuMemGetAllocationGranularity" } }; + inline for (names) |n| @field(api, n[0]) = lib.lookup(@FieldType(Api, n[0]), n[1]) orelse return error.MissingSymbol; + var prop: Prop = .{ .loc_id = device }; + var rdma: c_int = 0; // GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED: keep the memory registrable as cuMemAlloc's is + if (d.api.cuDeviceGetAttribute(&rdma, .gpu_direct_rdma_with_cuda_vmm_supported, device) == abi.success and rdma != 0) prop.rdma = 1; + var g: usize = 0; + if (api.granularity(&g, &prop, 1) != abi.success or g == 0) return error.Unsupported; // RECOMMENDED + g = @max(g, 2 << 20); + const len = std.mem.alignForward(usize, va_bytes, g); + var va: u64 = 0; + if (api.reserve(&va, len, g, 0, 0) != abi.success) return error.CudaFailed; + return .{ .lib = lib, .device = device, .prop = prop, .granularity = g, .va = va, .va_len = len, .next = va, .api = api }; + } + + pub fn close(v: *Vmm) void { + _ = v.api.addressFree(v.va, v.va_len); + v.lib.close(); + } + + pub fn backend(v: *Vmm) Backend { + return .{ .ctx = v, .granularity = v.granularity, .chunk = chunk, .release = release, .name = "vmm" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + if (v.va + v.va_len - v.next < len) return null; + var h: u64 = 0; + if (v.api.create(&h, len, &v.prop, 0) != abi.success) return null; + defer _ = v.api.release(h); // the mapping holds the memory until it is unmapped + const at = v.next; + if (v.api.map(at, len, 0, h, 0) != abi.success) return null; + const desc: AccessDesc = .{ .loc_id = v.device }; + if (v.api.setAccess(at, len, &desc, 1) != abi.success) { + _ = v.api.unmap(at, len); + return null; + } + v.next += len; + return at; + } + + fn release(ctx: *anyopaque, base: u64, len: usize) void { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + _ = v.api.unmap(base, len); // the virtual range is not reused: the reservation is sized for the process + } +}; + +// -- tests (a fake backend: addresses only) --------------------------------------------------------------------------- + +const Fake = struct { + next: u64 = 1 << 40, + chunks: u32 = 0, + released: u32 = 0, + fail: bool = false, + + fn backend(f: *Fake, g: usize) Backend { + return .{ .ctx = f, .granularity = g, .chunk = chunk, .release = release, .name = "fake" }; + } + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail) return null; + const p = f.next; + f.next += len; + f.chunks += 1; + return p; + } + fn release(ctx: *anyopaque, _: u64, _: usize) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.released += 1; + } +}; + +test "arena: small requests pack 512-aligned into shared chunks, are reused by size, large ones get their own" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + const p1 = a.alloc(20480).?; // a norm weight + const p2 = a.alloc(983040).?; // an mHC fn16 + const p3 = a.alloc(100).?; + try t.expectEqual(p1 + 20480, p2); + try t.expectEqual(p2 + 983040, p3); + for ([_]u64{ p1, p2, p3 }) |p| try t.expectEqual(@as(u64, 0), p % 512); + try t.expectEqual(@as(u32, 1), f.chunks); + // a freed range comes back for the same rounded size + try t.expect(a.free(p3)); + try t.expectEqual(p3, a.alloc(17).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + // three more 1 MiB ranges fit the 4 MiB chunk, the fourth opens a new shared chunk + for (0..3) |_| _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 1), f.chunks); + _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 2), f.chunks); + // large: its own chunk, rounded to the granularity, released when freed + const big = a.alloc(5 << 20).?; + try t.expectEqual(@as(u32, 3), f.chunks); + try t.expectEqual(@as(u64, 6 << 20), a.large.get(big).?); + try t.expect(a.free(big)); + try t.expectEqual(@as(u32, 1), f.released); + // not ours + try t.expect(!a.free(12345)); + // a backend that cannot map: null, the caller allocates as before + f.fail = true; + try t.expectEqual(@as(?u64, null), a.alloc(64 << 20)); + try t.expectEqual(@as(u64, 1), a.stats.fallback); +} + +test "arena: under DeviceBuffer.alloc through memory.setHook" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{}); + defer a.deinit(); + a.install(); + defer memory.setHook(null); + const h = memory.currentHook().?; + const p = h.alloc(h.ctx, 4096).?; + try t.expect(h.free(h.ctx, p)); + try t.expect(!h.free(h.ctx, 7)); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..a2fb8e18bb --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,392 @@ +//! Keyed CUDA graphs for any family: one captured executable per key, launched again when the key comes back. + +const std = @import("std"); +const graph = @import("graph.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +/// One step's launches on `stream`; it must not synchronize, allocate device memory or read the device on the host. +pub const Body = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, +}; + +/// Every rank's verdict AND-ed, so ranks capture, evict and fall back at the same steps; never called in a capture. +pub const Agree = struct { + ctx: *anyopaque, + all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, +}; + +/// Capture a body into an executable graph, launch it, free it: a fake engine runs the policy in host tests. +pub const Engine = struct { + ctx: *anyopaque, + capture: *const fn (ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque, + launch: *const fn (ctx: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void, + free: *const fn (ctx: *anyopaque, exec: *anyopaque) void, +}; + +pub const Settings = struct { + on: bool = true, // off: every step runs eagerly + max: u32 = 48, // graphs held at most + min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing + hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory +}; + +pub const Outcome = enum { replayed, captured, eager }; + +pub fn Cache(comptime Key: type) type { + return struct { + const Self = @This(); + const Entry = struct { exec: *anyopaque, used: u64 }; + + gpa: std.mem.Allocator, + engine: Engine, + settings: Settings, + agree: ?Agree = null, + room: ?*const fn () bool = null, // false when this rank is below its memory floor; null: always room + entries: std.AutoArrayHashMapUnmanaged(Key, Entry) = .empty, + cap: u32, + tick: u64 = 0, + fingerprint: ?u64 = null, + failed: bool = false, // a capture failed on some rank: eager from here on + stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + return .{ .gpa = gpa, .engine = engine, .settings = settings, .cap = settings.max }; + } + + pub fn deinit(c: *Self) void { + c.dropAll(); + c.entries.deinit(c.gpa); + } + + pub fn count(c: *const Self) usize { + return c.entries.count(); + } + + pub fn has(c: *const Self, key: Key) bool { + return c.entries.contains(key); + } + + /// Every graph freed (a fingerprint change, a shutdown). + pub fn dropAll(c: *Self) void { + for (c.entries.values()) |e| c.engine.free(c.engine.ctx, e.exec); + c.stats.dropped += c.entries.count(); + c.entries.clearRetainingCapacity(); + } + + fn agreed(c: *Self, ok: bool) !bool { + const a = c.agree orelse return ok; + return a.all(a.ctx, ok); + } + + fn evictOldest(c: *Self) void { + var at: ?usize = null; + for (c.entries.values(), 0..) |e, i| { + if (at == null or e.used < c.entries.values()[at.?].used) at = i; + } + const i = at orelse return; + c.engine.free(c.engine.ctx, c.entries.values()[i].exec); + c.entries.swapRemoveAt(i); + c.stats.evicted += 1; + } + + fn eager(c: *Self, stream: Stream, body: Body) !Outcome { + try body.run(body.ctx, stream); + c.stats.eager += 1; + return .eager; + } + + /// One step: the key's graph launched, captured first on a miss; a new `fingerprint` drops every graph first. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed) return c.eager(stream, body); + if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); + c.fingerprint = fingerprint; + c.tick += 1; + if (c.entries.getPtr(key)) |e| { + e.used = c.tick; + try c.engine.launch(c.engine.ctx, e.exec, stream); + c.stats.replayed += 1; + return .replayed; + } + // a miss: room on every rank, else shed the least recent quarter and lower the cap (once below it, eager) + const room_here = if (c.room) |f| f() else true; + if (!try c.agreed(room_here and c.cap > 0)) { + if (c.settings.hold) { + if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); + c.stats.floor_hits += 1; + return c.eager(stream, body); + } + if (c.entries.count() > c.settings.min_keep) { + const drop = @max(1, c.entries.count() / 4); + for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); + } + const was = c.cap; + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); + c.stats.floor_hits += 1; + if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); + return c.eager(stream, body); + } + if (c.cap < c.settings.max) { + if (!quiet) std.log.info("graph cache: room again on every rank: cap {d} -> {d}", .{ c.cap, c.settings.max }); + c.cap = c.settings.max; + } + while (c.entries.count() >= c.cap) c.evictOldest(); + const exec = c.engine.capture(c.engine.ctx, stream, body) catch |err| blk: { + // the only trace of why graphs went off (the step then runs eagerly, which may fail the same way) + if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); + break :blk null; + }; + if (!try c.agreed(exec != null)) { + if (exec) |x| c.engine.free(c.engine.ctx, x); + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + try c.entries.put(c.gpa, key, .{ .exec = exec.?, .used = c.tick }); + try c.engine.launch(c.engine.ctx, exec.?, stream); + c.stats.captured += 1; + return .captured; + } + }; +} + +/// The driver's engine: thread-local stream capture (other threads' CUDA calls stay out), instantiate, upload, launch. +pub const CudaEngine = struct { + d: *const Driver, + + pub fn engine(e: *CudaEngine) Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + + const Held = struct { exec: graph.Exec }; + + fn capture(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const e: *CudaEngine = @ptrCast(@alignCast(ctx)); + try graph.beginCapture(stream, .thread_local); + body.run(body.ctx, stream) catch |err| { + // end the capture so the stream is usable again; the partial graph is dropped + if (graph.endCapture(stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try graph.endCapture(stream); + defer g.deinit(); + const h = try std.heap.c_allocator.create(Held); + errdefer std.heap.c_allocator.destroy(h); + h.exec = try g.instantiate(); + errdefer h.exec.deinit(); + try h.exec.upload(stream); + _ = e; + return h; + } + + fn launch(_: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void { + const h: *Held = @ptrCast(@alignCast(exec)); + try h.exec.launchOn(stream); + } + + fn free(_: *anyopaque, exec: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(exec)); + h.exec.deinit(); + std.heap.c_allocator.destroy(h); + } +}; + +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares. +pub fn fingerprintOf(words: []const u64) u64 { + var h: u64 = 0xcbf29ce484222325; + for (words) |w| { + var x = w; + for (0..8) |_| { + h = (h ^ (x & 0xff)) *% 0x100000001b3; + x >>= 8; + } + } + return h; +} + +// -- host tests: the policy with a fake engine -------------------------------------------------------------------- + +const Fake = struct { + captures: u32 = 0, + launches: u32 = 0, + live: u32 = 0, + fail_capture: bool = false, + slots: [64]u8 = undefined, + + fn engine(f: *Fake) Engine { + return .{ .ctx = f, .capture = cap, .launch = lau, .free = fre }; + } + fn cap(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail_capture) return error.CaptureFailed; + _ = stream; + _ = body; // a capture records the body; nothing runs + f.captures += 1; + f.live += 1; + return &f.slots[f.captures % f.slots.len]; + } + fn lau(ctx: *anyopaque, _: *anyopaque, _: Stream) anyerror!void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.launches += 1; + } + fn fre(ctx: *anyopaque, _: *anyopaque) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.live -= 1; + } +}; + +const Ran = struct { + n: u32 = 0, + fn body(r: *Ran) Body { + return .{ .ctx = r, .run = go }; + } + fn go(ctx: *anyopaque, _: Stream) anyerror!void { + const r: *Ran = @ptrCast(@alignCast(ctx)); + r.n += 1; + } +}; + +const no_stream: Stream = .{ .d = undefined, .handle = null }; + +test "graph cache: capture on first use, replay after, LRU at the cap" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 2 }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.replayed, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.captured, try c.run(2, no_stream, 7, r.body())); + _ = try c.run(1, no_stream, 7, r.body()); // 1 is now the most recent + try std.testing.expectEqual(Outcome.captured, try c.run(3, no_stream, 7, r.body())); // evicts 2 + try std.testing.expect(c.has(1) and c.has(3) and !c.has(2)); + try std.testing.expectEqual(@as(u32, 0), r.n); // the body never ran eagerly + try std.testing.expectEqual(@as(u32, 2), f.live); + try std.testing.expectEqual(@as(u32, 5), f.launches); +} + +test "graph cache: a new fingerprint drops every graph" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + defer c.deinit(); + _ = try c.run(1, no_stream, 7, r.body()); + _ = try c.run(2, no_stream, 7, r.body()); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 8, r.body())); + try std.testing.expectEqual(@as(usize, 1), c.count()); + try std.testing.expectEqual(@as(u32, 1), f.live); +} + +const Peer = struct { + theirs: []const bool, // the other rank's verdicts, in call order + at: usize = 0, + fn agree(p: *Peer) Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const t = p.theirs[p.at]; + p.at += 1; + return ok and t; + } +}; + +test "graph cache: a capture failing on another rank turns graphs off here too" { + var f: Fake = .{}; + var r: Ran = .{}; + // per miss: room, then the capture's verdict; the peer's capture fails on the second key + var peer: Peer = .{ .theirs = &.{ true, true, true, false } }; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + c.agree = peer.agree(); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(2, no_stream, 7, r.body())); + try std.testing.expect(c.failed); + try std.testing.expectEqual(@as(u32, 0), f.live); // this rank's good graphs are freed as well + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 2), r.n); +} + +var room_flag = true; +fn roomFn() bool { + return room_flag; +} + +test "graph cache: below the memory floor the least recent quarter goes, the cap follows, the step runs eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 6), c.count()); + try std.testing.expectEqual(@as(u32, 6), c.cap); + try std.testing.expect(!c.has(0) and !c.has(1) and c.has(7)); + room_flag = true; + // room again: the cap is back to max, the miss captures without evicting + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 7), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + try std.testing.expect(c.has(7) and c.has(2)); + try std.testing.expectEqual(@as(u64, 1), c.stats.floor_hits); +} + +test "graph cache: a dip under the floor does not leave later keys churning once memory is back" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 16, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..12) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; // a long prompt's workspace held + _ = try c.run(200, no_stream, 7, r.body()); + room_flag = true; // freed: the decode after it meets 12 new keys (a new context bucket), then replays them + const before = c.stats.evicted; + for (0..12) |k| _ = try c.run(@intCast(300 + k), no_stream, 7, r.body()); + for (0..12) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(300 + k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 16), c.cap); + // only what the cap of 16 forces: 9 left after the dip + 12 new = 21 -> 5 evictions, none during the replays + try std.testing.expectEqual(before + 5, c.stats.evicted); +} + +test "graph cache: hold - below the floor a miss runs eagerly, nothing is evicted, the held keys replay" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2, .hold = true }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(101, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 8), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + for (0..8) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u64, 0), c.stats.evicted); + try std.testing.expectEqual(@as(u64, 2), c.stats.floor_hits); + room_flag = true; + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 9), c.count()); +} + +test "graph cache: off runs every step eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .on = false }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 2c469bae2c..6d802980c7 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -49,16 +49,37 @@ fn lockCounts() void { while (!counts_mutex.tryLock()) std.Thread.yield() catch {}; } +/// A sub-allocator under DeviceBuffer (arena.zig): `alloc` gives an address or null for cuMemAlloc, `free` owns it. +pub const Hook = struct { + ctx: *anyopaque, + alloc: *const fn (ctx: *anyopaque, len: usize) ?abi.DevicePtr, + free: *const fn (ctx: *anyopaque, p: abi.DevicePtr) bool, // true when the address was the hook's +}; + +var alloc_hook: ?Hook = null; + +/// Installs, or with null removes, the sub-allocator: once at boot, before the buffers it should hold. +pub fn setHook(h: ?Hook) void { + alloc_hook = h; +} + +pub fn currentHook() ?Hook { + return alloc_hook; +} + pub const DeviceBuffer = struct { d: *const Driver, ptr: abi.DevicePtr, len: usize, borrowed: bool = false, // a span of memory another value owns (a carveout's): free leaves it alone - /// `len` bytes, 256-byte aligned by the driver; zero bytes allocate nothing and hold address 0. + /// `len` bytes, 256-byte aligned (512 from an installed arena); zero bytes allocate nothing and hold address 0. pub fn alloc(d: *const Driver, len: usize) Error!DeviceBuffer { var p: abi.DevicePtr = 0; - if (len > 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + if (len > 0) { + if (alloc_hook) |h| p = h.alloc(h.ctx, len) orelse 0; + if (p == 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + } held(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -73,7 +94,10 @@ pub const DeviceBuffer = struct { pub fn free(self: *DeviceBuffer) void { if (!self.borrowed) { - if (self.ptr != 0) _ = self.d.api.cuMemFree_v2(self.ptr); + if (self.ptr != 0) { + const mine = if (alloc_hook) |h| h.free(h.ctx, self.ptr) else false; + if (!mine) _ = self.d.api.cuMemFree_v2(self.ptr); + } freed(&device_bytes, self.len); } self.* = undefined; @@ -209,6 +233,28 @@ test "freeing a borrowed span leaves the driver and the counts alone" { try std.testing.expectEqual(before.device, usage(false).device); } +test "a buffer the hook owns is freed by the hook, never by the driver" { + const Owner = struct { + var frees: u32 = 0; + fn alloc(_: *anyopaque, _: usize) ?abi.DevicePtr { + return null; + } + fn free(_: *anyopaque, p: abi.DevicePtr) bool { + frees += 1; + return p == 0x2000; + } + }; + var ctx: u8 = 0; + setHook(.{ .ctx = &ctx, .alloc = Owner.alloc, .free = Owner.free }); + defer setHook(null); + held(&device_bytes, 4096); + var b: DeviceBuffer = .{ .d = undefined, .ptr = 0x2000, .len = 4096 }; + b.free(); + var c: DeviceBuffer = .{ .d = undefined, .ptr = 0x3000, .len = 64, .borrowed = true }; + c.free(); + try std.testing.expectEqual(@as(u32, 1), Owner.frees); +} + test "peak reset preserves an allocation between its read and store" { const before = usage(true); held(&device_bytes, 100); diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index ba8fb572ad..f72567af1a 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -6,6 +6,7 @@ pub const Error = @import("driver.zig").Error; pub const Context = @import("context.zig").Context; pub const DeviceBuffer = @import("memory.zig").DeviceBuffer; pub const HostBuffer = @import("memory.zig").HostBuffer; +pub const memory = @import("memory.zig"); pub const usage = @import("memory.zig").usage; pub const Usage = @import("memory.zig").Usage; pub const carveout = @import("carveout.zig"); @@ -31,6 +32,8 @@ pub const segments = @import("segments.zig"); pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const experts = @import("experts.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); test { _ = @import("memory.zig"); @@ -45,4 +48,6 @@ test { _ = @import("fp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); + _ = graph_cache; + _ = arena; } From 5fa8ad0fd9c53131c599c39f8885117d5a582fda Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:18 +0700 Subject: [PATCH 3/7] A failed CUDA kernel launch names its kernel and launch configuration; CUresult names fall back to a table when the driver cannot answer. cuFuncGetName is optional (CUDA 12.3+), so an older driver still opens. During the exit's teardown cuGetErrorName fails too, which printed CUDA_ERROR_DEINITIALIZED as "CUDA_ERROR_UNKNOWN_CODE (4)"; code 4 now says the process is exiting. --- zig/src/cuda/driver.zig | 67 ++++++++++++++++++++++++++++++++++++++--- zig/src/cuda/launch.zig | 36 ++++++++++++++++++++-- 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..88ee9050dc 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -8,6 +8,10 @@ pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDevi pub const Driver = struct { lib: std.DynLib, api: abi.Api, + /// cuFuncGetName (CUDA 12.3+; optional: an older driver opens without it): a failed launch names its kernel + func_name: ?FuncGetName = null, + + pub const FuncGetName = *const fn (*?[*:0]const u8, abi.Function) callconv(.c) abi.Result; pub fn open() Error!Driver { return openPath("libcuda.so.1"); @@ -25,7 +29,7 @@ pub const Driver = struct { return error.MissingSymbol; }; } - const d: Driver = .{ .lib = lib, .api = api }; + const d: Driver = .{ .lib = lib, .api = api, .func_name = lib.lookup(FuncGetName, "cuFuncGetName") }; try d.check(api.cuInit(0), "cuInit"); return d; } @@ -38,18 +42,33 @@ pub const Driver = struct { pub fn check(self: *const Driver, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; if (res == abi.error_not_ready) return error.NotReady; - std.log.err("{s}: {s} ({d}) {s}", .{ what, self.errorName(res), res, self.errorText(res) }); + std.log.err("{s}: {s} ({d}) {s}{s}", .{ what, self.errorName(res), res, self.errorText(res), shutdownNote(res) }); + return errorOf(res); + } + + /// The error a failed call returns (`check` without its log line, for callers that log more). + pub fn errorOf(res: abi.Result) Error { return switch (res) { 2 => error.OutOfDeviceMemory, abi.error_not_found => error.NotFound, + abi.error_not_ready => error.NotReady, else => error.CudaFailed, }; } + /// The driver's name for the code, else knownName's: during the exit's teardown cuGetErrorName fails too. pub fn errorName(self: *const Driver, res: abi.Result) []const u8 { var s: ?[*:0]const u8 = null; - if (self.api.cuGetErrorName(res, &s) != abi.success) return "CUDA_ERROR_UNKNOWN_CODE"; - return if (s) |p| std.mem.span(p) else "CUDA_ERROR_UNKNOWN_CODE"; + if (self.api.cuGetErrorName(res, &s) == abi.success) if (s) |p| return std.mem.span(p); + return knownName(res) orelse "CUDA_ERROR_UNKNOWN_CODE"; + } + + /// A kernel's symbol name (cuFuncGetName), or "?" when the driver cannot say. + pub fn functionName(self: *const Driver, f: abi.Function) []const u8 { + const get = self.func_name orelse return "?"; + var s: ?[*:0]const u8 = null; + if (get(&s, f) != abi.success) return "?"; + return if (s) |p| std.mem.span(p) else "?"; } pub fn errorText(self: *const Driver, res: abi.Result) []const u8 { @@ -82,3 +101,43 @@ pub const Driver = struct { return @intCast(10 * major + minor); } }; + +/// CUresult names (cuda.h) for when cuGetErrorName cannot answer. +pub fn knownName(res: abi.Result) ?[]const u8 { + return switch (res) { + 1 => "CUDA_ERROR_INVALID_VALUE", + 2 => "CUDA_ERROR_OUT_OF_MEMORY", + 3 => "CUDA_ERROR_NOT_INITIALIZED", + 4 => "CUDA_ERROR_DEINITIALIZED", + 101 => "CUDA_ERROR_INVALID_DEVICE", + 200 => "CUDA_ERROR_INVALID_IMAGE", + 201 => "CUDA_ERROR_INVALID_CONTEXT", + 209 => "CUDA_ERROR_NO_BINARY_FOR_GPU", + 400 => "CUDA_ERROR_INVALID_HANDLE", + 500 => "CUDA_ERROR_NOT_FOUND", + 600 => "CUDA_ERROR_NOT_READY", + 700 => "CUDA_ERROR_ILLEGAL_ADDRESS", + 701 => "CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", + 702 => "CUDA_ERROR_LAUNCH_TIMEOUT", + 709 => "CUDA_ERROR_CONTEXT_IS_DESTROYED", + 710 => "CUDA_ERROR_ASSERT", + 714 => "CUDA_ERROR_HARDWARE_STACK_ERROR", + 715 => "CUDA_ERROR_ILLEGAL_INSTRUCTION", + 716 => "CUDA_ERROR_MISALIGNED_ADDRESS", + 719 => "CUDA_ERROR_LAUNCH_FAILED", + 999 => "CUDA_ERROR_UNKNOWN", + else => null, + }; +} + +/// CUDA_ERROR_DEINITIALIZED: the process is exiting (a stop signal, then exit) while a thread still calls the driver. +pub fn shutdownNote(res: abi.Result) []const u8 { + return if (res == 4) " (the CUDA driver is shutting down: the process is exiting)" else ""; +} + +test "CUresult names without the driver: DEINITIALIZED is 4, unknown codes stay unknown" { + try std.testing.expectEqualStrings("CUDA_ERROR_DEINITIALIZED", knownName(4).?); + try std.testing.expectEqualStrings("CUDA_ERROR_ILLEGAL_ADDRESS", knownName(700).?); + try std.testing.expect(knownName(12345) == null); + try std.testing.expect(shutdownNote(4).len > 0 and shutdownNote(1).len == 0); +} diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index 2472042638..c80684f291 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -4,7 +4,8 @@ const std = @import("std"); const abi = @import("abi.zig"); const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; -const Error = @import("driver.zig").Error; +const driver = @import("driver.zig"); +const Error = driver.Error; pub const Dim3 = abi.Dim3; @@ -80,7 +81,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void const g = cfg.grid; const b = cfg.block; if (cfg.plain()) { - return d.check(d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); + return launched(f, cfg, d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); } var attrs: [4]abi.LaunchAttribute = undefined; const n = fillAttributes(cfg, &attrs); @@ -96,7 +97,36 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void .attrs = &attrs, .num_attrs = @intCast(n), }; - try d.check(d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); + try launched(f, cfg, d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); +} + +/// A launch's result; a failure names the kernel and its configuration with the code. +fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { + if (res == abi.success) return; + const d = f.d; + var buf: [160]u8 = undefined; + std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) }); + return driver.Driver.errorOf(res); +} + +/// "kernel , grid XxYxZ, block XxYxZ, shared N B" (+ cluster / PDL / cooperative when set). +pub fn describe(buf: []u8, name: []const u8, cfg: Config) []const u8 { + const g = cfg.grid; + const b = cfg.block; + var w = std.Io.Writer.fixed(buf); + w.print("kernel {s}, grid {d}x{d}x{d}, block {d}x{d}x{d}, shared {d} B", .{ name, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared }) catch return w.buffered(); + if (cfg.cluster) |c| w.print(", cluster {d}x{d}x{d}", .{ c.x, c.y, c.z }) catch return w.buffered(); + if (cfg.pdl) w.writeAll(", pdl") catch return w.buffered(); + if (cfg.cooperative) w.writeAll(", cooperative") catch return w.buffered(); + return w.buffered(); +} + +test "a failed launch's description: the kernel, its grid, block and shared bytes, bounded" { + var buf: [160]u8 = undefined; + const cfg: Config = .{ .grid = .{ .x = 70000, .y = 2, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, .shared = 49152, .pdl = true }; + try std.testing.expectEqualStrings("kernel _stream_pf, grid 70000x2x1, block 128x1x1, shared 49152 B, pdl", describe(&buf, "_stream_pf", cfg)); + var small: [24]u8 = undefined; + try std.testing.expect(describe(&small, "a_very_long_kernel_symbol_name", cfg).len <= small.len); } fn attribute(id: abi.LaunchAttributeId) abi.LaunchAttribute { From 0824e393b27ba8ebce8663762686d8bb2730f3da Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:38:24 +0700 Subject: [PATCH 4/7] Tensor parallelism across CUDA processes: bootstrap, NCCL, mailbox and RoCE collectives, plan link, fail-fast. zig/src/fabric/tp brings a rank up from TF_TP_WORLD, TF_TP_RANK, TF_TP_MASTER and TF_TP_PORT, agrees the settings digest over every rank, and gives a family all-gather, all-reduce and broadcast over NCCL (TF_COMM_BACKEND=nccl, the default), a pinned host mailbox for ranks on one host, or RoCE verbs between two hosts (TF_TP_ROCE_*), with an NCCL fallback when RoCE setup fails. Rank 0 sends each round's plan to the others over its own link (TF_TP_PLAN_*). A failure or a vanished rank ends every rank shortly after (TF_TP_FAILFAST, TF_TP_FAILFAST_GRACE_S). No family uses it yet. zig build test-tp runs the host tests; zig build tf-tp-test builds the two-GPU runner (-Dtp-nvcc= or -Dtp-fatbin= for the mailbox kernel). The RoCE all-gather adapts b12x's RoCEnante protocol. --- THIRD_PARTY_NOTICES.md | 6 + build.zig | 2 + zig/build/tp.zig | 59 ++++ zig/src/cuda/nccl.zig | 7 + zig/src/fabric/tp/bootstrap.zig | 262 +++++++++++++++ zig/src/fabric/tp/collective.zig | 136 ++++++++ zig/src/fabric/tp/config.zig | 245 ++++++++++++++ zig/src/fabric/tp/fate.zig | 261 +++++++++++++++ zig/src/fabric/tp/fate_test.zig | 146 ++++++++ zig/src/fabric/tp/gpu_bench.zig | 82 +++++ zig/src/fabric/tp/gpu_checks.zig | 188 +++++++++++ zig/src/fabric/tp/gpu_main.zig | 92 +++++ zig/src/fabric/tp/host.zig | 293 ++++++++++++++++ zig/src/fabric/tp/host_test.zig | 154 +++++++++ zig/src/fabric/tp/hybrid.zig | 167 ++++++++++ zig/src/fabric/tp/log.zig | 15 + zig/src/fabric/tp/mailbox.cu | 556 +++++++++++++++++++++++++++++++ zig/src/fabric/tp/mailbox.zig | 215 ++++++++++++ zig/src/fabric/tp/nccl.zig | 163 +++++++++ zig/src/fabric/tp/planlink.zig | 467 ++++++++++++++++++++++++++ zig/src/fabric/tp/reference.zig | 121 +++++++ zig/src/fabric/tp/roce.zig | 458 +++++++++++++++++++++++++ zig/src/fabric/tp/roce_host.zig | 58 ++++ zig/src/fabric/tp/roce_test.zig | 116 +++++++ zig/src/fabric/tp/roce_verbs.zig | 297 +++++++++++++++++ zig/src/fabric/tp/run_hosts.sh | 24 ++ zig/src/fabric/tp/run_pair.sh | 47 +++ zig/src/fabric/tp/sock.zig | 169 ++++++++++ zig/src/fabric/tp/tp.zig | 212 ++++++++++++ zig/src/fabric/verbs.zig | 2 +- 30 files changed, 5019 insertions(+), 1 deletion(-) create mode 100644 zig/build/tp.zig create mode 100644 zig/src/fabric/tp/bootstrap.zig create mode 100644 zig/src/fabric/tp/collective.zig create mode 100644 zig/src/fabric/tp/config.zig create mode 100644 zig/src/fabric/tp/fate.zig create mode 100644 zig/src/fabric/tp/fate_test.zig create mode 100644 zig/src/fabric/tp/gpu_bench.zig create mode 100644 zig/src/fabric/tp/gpu_checks.zig create mode 100644 zig/src/fabric/tp/gpu_main.zig create mode 100644 zig/src/fabric/tp/host.zig create mode 100644 zig/src/fabric/tp/host_test.zig create mode 100644 zig/src/fabric/tp/hybrid.zig create mode 100644 zig/src/fabric/tp/log.zig create mode 100644 zig/src/fabric/tp/mailbox.cu create mode 100644 zig/src/fabric/tp/mailbox.zig create mode 100644 zig/src/fabric/tp/nccl.zig create mode 100644 zig/src/fabric/tp/planlink.zig create mode 100644 zig/src/fabric/tp/reference.zig create mode 100644 zig/src/fabric/tp/roce.zig create mode 100644 zig/src/fabric/tp/roce_host.zig create mode 100644 zig/src/fabric/tp/roce_test.zig create mode 100644 zig/src/fabric/tp/roce_verbs.zig create mode 100644 zig/src/fabric/tp/run_hosts.sh create mode 100644 zig/src/fabric/tp/run_pair.sh create mode 100644 zig/src/fabric/tp/sock.zig create mode 100644 zig/src/fabric/tp/tp.zig diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index f63b0ac17a..d02e4cea6d 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -76,6 +76,12 @@ checked bit for bit against ExLlamaV3's dequantization. Flash Next's optional int8 and int4 KV caches (`families/qwen4_exp/cuda/kvcache.py`) follow the cache quantization scheme of [ExLlamaV3](https://github.com/turboderp-org/exllamav3) `-cq 8` and `-cq 4` (MIT License, Copyright (c) 2025 Turboderp, text below): groups of 32, one fp16 absmax scale per group, the group rotated by a 32-point Hadamard, midpoint-grid codes, `compand_a == 0`. 8-bit stores each code as a signed int8 (`q - 128`). 4-bit stores two unsigned codes per byte, low nibble first (the same bits as ExLlamaV3's little-endian packing, a uint8 tensor rather than their uint32 words). Their dequantizer folds another `1/sqrt(32)` into the scale and applies the unnormalized butterfly on the way out; this cache applies the normalized H32 to the query and to the merged output instead, and leaves the stored codes rotated. Scales match their quantizer bit for bit. Reconstructed values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit for bit. The quantizer and the attention dequant are written for TensorFold and checked against an independent reference of that arithmetic. +The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts +the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante +(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by +sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, +Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. + ## Vendored code and weights `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified `dflash/model_mlx.py` from diff --git a/build.zig b/build.zig index 9515dc3618..30659f7b7b 100644 --- a/build.zig +++ b/build.zig @@ -4,6 +4,7 @@ const cuda_build = @import("zig/build/cuda.zig"); const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); +const tp_build = @import("zig/build/tp.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -172,6 +173,7 @@ pub fn build(b: *std.Build) void { const warning_run = b.addRunArtifact(warning); test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); + tp_build.add(b, target, optimize, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..f5dacc46e3 --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,59 @@ +//! Tensor parallelism's build: host tests on any OS, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. + +const std = @import("std"); + +/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true })); + return m; +} + +/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else. +fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath { + const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" }); + const out = run.addOutputFileArg("tp_mailbox.fatbin"); + run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu")); + return out; +} + +/// `pub const mailbox: []const u8`, the fatbin's bytes or empty. +fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { + const wf = b.addWriteFiles(); + if (image) |file| { + _ = wf.addCopyFile(file, "tp_mailbox.fatbin"); + // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") }); + } + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") }); +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { + const host = b.graph.host; + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); + test_step.dependOn(&tests.step); + if (target.result.os.tag != .linux) return; + const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); + const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); + const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null; + const cuda = runtime(b, target, optimize); + const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + runner.addImport("cuda", cuda); + runner.addImport("tp", module(b, target, optimize, cuda)); + runner.addImport("tp_kernel", kernelModule(b, image)); + const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + b.installArtifact(exe); + b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..85f88bace8 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -12,6 +12,8 @@ pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; +pub const in_progress: Result = 7; // ncclInProgress: a non-blocking communicator's call is still running + const R = Result; const D = abi.DevicePtr; @@ -25,6 +27,11 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + /// Frees the communicator after ending its pending work; safe from another thread while a kernel waits. + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *Result) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..1486e3a8f7 --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,262 @@ +//! Rank discovery over TCP (a star on rank 0): small host blobs all-gathered, settings agreed, sockets kept for fate. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; + +pub const max_ranks = 64; +/// Largest blob a rank contributes to one `allGather`. +pub const max_blob = 64 * 1024; + +const magic: u32 = 0x31425054; // "TPB1" +const version: u32 = 2; + +pub const Error = error{ Protocol, Mismatch, TooLarge } || sock.Error; + +/// `role`: 0 the control connection (bootstrap, then the fate channel), 1 the plan link. +const Hello = extern struct { magic: u32, version: u32, world: u32, rank: u32, role: u32 }; + +pub const Bootstrap = struct { + rank: u32, + world: u32, + /// Rank 0: the socket of rank i at index i (index 0 unused); other ranks: rank 0's socket at index 0. + fds: [max_ranks]sock.Fd = @splat(-1), + /// The plan link's sockets, laid out as `fds`. + plan_fds: [max_ranks]sock.Fd = @splat(-1), + /// Each call's limit for the peers' bytes (the connect timeout). + timeout_ns: u64, + + fn deadline(b: *const Bootstrap) u64 { + return sock.nowNs() + b.timeout_ns; + } + + /// Rank 0 listens on `link_port` and accepts every other rank; the others connect to `master` (connect timeout). + pub fn open(cfg: Config) Error!Bootstrap { + if (cfg.world > max_ranks or cfg.rank >= cfg.world) return error.Protocol; + var b: Bootstrap = .{ .rank = cfg.rank, .world = cfg.world, .timeout_ns = cfg.connect_timeout_ns }; + errdefer b.close(); + if (cfg.world == 1) return b; + if (cfg.rank == 0) { + const l = try sock.listen(cfg.link_port); + defer sock.close(l); + try b.acceptAll(l); + } else { + const until = b.deadline(); + for ([_]u32{ 0, 1 }) |role| { + const fd = try sock.connect(cfg.master, cfg.link_port, until); + if (role == 0) b.fds[0] = fd else b.plan_fds[0] = fd; + const h: Hello = .{ .magic = magic, .version = version, .world = cfg.world, .rank = cfg.rank, .role = role }; + try sock.sendAll(fd, std.mem.asBytes(&h)); + } + } + return b; + } + + /// Rank 0 over an already listening socket (tests: an ephemeral port). + pub fn acceptOn(l: sock.Fd, world: u32, timeout_ns: u64) Error!Bootstrap { + var b: Bootstrap = .{ .rank = 0, .world = world, .timeout_ns = timeout_ns }; + errdefer b.close(); + try b.acceptAll(l); + return b; + } + + fn acceptAll(b: *Bootstrap, l: sock.Fd) Error!void { + // every other rank's control and plan connections, in any order + var joined: u32 = 0; + const until = b.deadline(); + while (joined < 2 * (b.world - 1)) { + const fd = try sock.accept(l, until); + var h: Hello = undefined; + sock.recvExact(fd, std.mem.asBytes(&h), until) catch { + sock.close(fd); + continue; + }; + if (h.magic != magic or h.version != version) { + log.warn("tp bootstrap: a connection that is not a rank (magic {x}); ignored", .{h.magic}); + sock.close(fd); + continue; + } + const table = if (h.role == 0) &b.fds else &b.plan_fds; + if (h.world != b.world or h.rank == 0 or h.rank >= b.world or h.role > 1 or table[h.rank] >= 0) { + log.warn("tp bootstrap: rank {d} of {d} joined a world of {d}", .{ h.rank, h.world, b.world }); + sock.close(fd); + return error.Mismatch; + } + table[h.rank] = fd; + joined += 1; + } + } + + /// Takes the control sockets (the fate channel keeps them); `close` then leaves them open. + pub fn release(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.fds; + b.fds = @splat(-1); + return out; + } + + /// Takes the plan link's sockets. + pub fn releasePlan(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.plan_fds; + b.plan_fds = @splat(-1); + return out; + } + + pub fn close(b: *Bootstrap) void { + for ([_]*[max_ranks]sock.Fd{ &b.fds, &b.plan_fds }) |table| for (table) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } + + /// Every rank's `mine` in rank order into `out`, returned as slices of `out`; every rank calls it equally often. + pub fn allGather(b: *Bootstrap, mine: []const u8, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + if (mine.len > max_blob) return error.TooLarge; + std.debug.assert(blobs.len >= b.world); + if (b.world == 1) { + if (out.len < mine.len) return error.TooLarge; + @memcpy(out[0..mine.len], mine); + blobs[0] = out[0..mine.len]; + return blobs[0..1]; + } + if (b.rank != 0) { + try sendFrame(b.fds[0], mine); + return b.readAll(b.fds[0], out, blobs); + } + // rank 0: everyone's blob, then the whole table to everyone + var used: usize = 0; + for (0..b.world) |i| { + const len: u32 = if (i == 0) @intCast(mine.len) else blk: { + var l: u32 = undefined; + try sock.recvExact(b.fds[i], std.mem.asBytes(&l), b.deadline()); + break :blk l; + }; + if (len > max_blob or used + len > out.len) return error.TooLarge; + if (i == 0) @memcpy(out[used..][0..len], mine) else try sock.recvExact(b.fds[i], out[used..][0..len], b.deadline()); + blobs[i] = out[used..][0..len]; + used += len; + } + for (1..b.world) |i| { + for (blobs[0..b.world]) |blob| try sendFrame(b.fds[i], blob); + } + return blobs[0..b.world]; + } + + fn readAll(b: *Bootstrap, fd: sock.Fd, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + var used: usize = 0; + for (0..b.world) |i| { + var l: u32 = undefined; + try sock.recvExact(fd, std.mem.asBytes(&l), b.deadline()); + if (l > max_blob or used + l > out.len) return error.TooLarge; + try sock.recvExact(fd, out[used..][0..l], b.deadline()); + blobs[i] = out[used..][0..l]; + used += l; + } + return blobs[0..b.world]; + } + + /// Rank 0's `bytes` on every rank (into `out` on the others; `bytes` is ignored there). + pub fn broadcast(b: *Bootstrap, bytes: []const u8, out: []u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [4096]u8 = undefined; + if (out.len * b.world > scratch.len) return error.TooLarge; + const all = try b.allGather(if (b.rank == 0) bytes else &.{}, &scratch, &blobs); + if (all[0].len != out.len) return error.Protocol; + @memcpy(out, all[0]); + } + + /// Every rank's `bytes` must be equal (settings that change bytes on the wire); `what` names them in the log. + pub fn agree(b: *Bootstrap, bytes: []const u8, what: []const u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [8192]u8 = undefined; + const all = try b.allGather(bytes, &scratch, &blobs); + for (all, 0..) |blob, i| if (!std.mem.eql(u8, blob, all[0])) { + log.warn("tp bootstrap: rank {d}'s {s} differs from rank 0's; both ranks refuse to start", .{ i, what }); + return error.Mismatch; + }; + } +}; + +fn sendFrame(fd: sock.Fd, bytes: []const u8) sock.Error!void { + const l: u32 = @intCast(bytes.len); + try sock.sendAll(fd, std.mem.asBytes(&l)); + try sock.sendAll(fd, bytes); +} + +/// An ephemeral listening port for tests. +pub fn testListener() !struct { fd: sock.Fd, port: u16 } { + const l = try sock.listen(0); + var addr: std.posix.sockaddr.in = undefined; + var len: std.posix.socklen_t = @sizeOf(std.posix.sockaddr.in); + _ = std.c.getsockname(l, @ptrCast(&addr), &len); + return .{ .fd = l, .port = std.mem.bigToNative(u16, addr.port) }; +} + +fn testRank(cfg: Config, out: *[3][32]u8, err: *?anyerror) void { + var b = Bootstrap.open(cfg) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var id: [32]u8 = undefined; + b.broadcast(&@as([32]u8, @splat(0xAB)), &id) catch |e| { + err.* = e; + return; + }; + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const mine: [3]u8 = @splat(@intCast(cfg.rank)); + const all = b.allGather(mine[0 .. cfg.rank + 1], &buf, &blobs) catch |e| { + err.* = e; + return; + }; + for (all, 0..) |blob, i| @memcpy(out[i][0..blob.len], blob); + b.agree("same", "settings") catch |e| { + err.* = e; + }; +} + +test "three ranks gather in rank order and agree" { + const l = try testListener(); + defer sock.close(l.fd); + var outs: [3][3][32]u8 = undefined; + var errs: [3]?anyerror = @splat(null); + var threads: [2]std.Thread = undefined; + for (0..2) |i| threads[i] = try std.Thread.spawn(.{}, testRank, .{ Config{ .world = 3, .rank = @intCast(i + 1), .link_port = l.port, .connect_timeout_ns = 5 * std.time.ns_per_s }, &outs[i + 1], &errs[i + 1] }); + var b = try Bootstrap.acceptOn(l.fd, 3, 5 * std.time.ns_per_s); + defer b.close(); + var id: [32]u8 = undefined; + try b.broadcast(&@as([32]u8, @splat(0xAB)), &id); + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const all = try b.allGather(&.{0}, &buf, &blobs); + try b.agree("same", "settings"); + for (threads) |t| t.join(); + for (errs) |e| try std.testing.expect(e == null); + try std.testing.expectEqual(@as(usize, 3), all.len); + try std.testing.expectEqualSlices(u8, &.{ 2, 2, 2 }, all[2]); + try std.testing.expectEqualSlices(u8, &.{ 1, 1 }, outs[2][1][0..2]); +} + +test "settings that differ stop every rank" { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, err: *?anyerror) void { + var b = Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = port, .connect_timeout_ns = 5 * std.time.ns_per_s }) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree("knobs-b", "settings") catch |e| { + err.* = e; + }; + } + }; + var err: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, &err }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + try std.testing.expectError(error.Mismatch, b.agree("knobs-a", "settings")); + t.join(); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); +} diff --git a/zig/src/fabric/tp/collective.zig b/zig/src/fabric/tp/collective.zig new file mode 100644 index 0000000000..d2334dc74a --- /dev/null +++ b/zig/src/fabric/tp/collective.zig @@ -0,0 +1,136 @@ +//! The tensor-parallel communicator every backend implements: stream-ordered, graph-capturable, rank-order sums. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DevicePtr = cuda.abi.DevicePtr; +pub const Stream = cuda.abi.Stream; + +pub const Error = error{ Aborted, Unsupported, Invalid, PeerFailed, Timeout, BackendFailed }; + +/// Element types the collectives move (`reduce` takes the float ones and the 32/64-bit integers). +pub const DType = enum(u8) { + u8, + i32, + i64, + f16, + bf16, + f32, + f64, + + pub fn size(t: DType) usize { + return switch (t) { + .u8 => 1, + .f16, .bf16 => 2, + .i32, .f32 => 4, + .i64, .f64 => 8, + }; + } +}; + +pub const Op = enum(u8) { sum, max, min }; + +/// Which implementation is behind a `Collective` (logs, bench labels). +pub const Kind = enum { nccl, mailbox, hybrid, roce, host }; + +pub const Collective = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + kind: Kind, + rank: *const fn (ptr: *anyopaque) u32, + world: *const fn (ptr: *anyopaque) u32, + /// recv[count] <- the element-wise reduction of every rank's send[count], added in rank order. + all_reduce: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void, + /// recv[world * count] <- every rank's send[count], rank 0's first (NCCL's layout). + all_gather: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, stream: Stream) Error!void, + /// Two ranks trade: this rank's send[count] lands in `peer`'s recv[count] and the peer's in ours (one group). + exchange: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + send: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + recv: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + /// Host: returns once every rank has called it and this rank's earlier collectives have finished. + barrier: *const fn (ptr: *anyopaque) Error!void, + /// Host: a failure the transport recorded since the last check (a peer that timed out, an NCCL async error). + check: *const fn (ptr: *anyopaque) Error!void, + /// Any thread: end pending and later work at once; safe while another thread is inside a call. + abort: *const fn (ptr: *anyopaque) void, + /// recv[world][stride]: rank r's first lens[r] bytes (device int32 [world]) land at recv + r * stride. + all_gather_v: ?*const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void = null, + /// Whether `all_gather_v` moves only the lengths for shards of `bytes` here (rank-local). Null: never. + var_fits: ?*const fn (ptr: *anyopaque, bytes: usize) bool = null, + /// Host: whether every rank's transport takes device lengths, agreed once per communicator. Null: false. + var_agreed: ?*const fn (ptr: *anyopaque) Error!bool = null, + }; + + pub fn kind(c: Collective) Kind { + return c.vtable.kind; + } + + pub fn rank(c: Collective) u32 { + return c.vtable.rank(c.ptr); + } + + pub fn world(c: Collective) u32 { + return c.vtable.world(c.ptr); + } + + pub fn allReduce(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void { + return c.vtable.all_reduce(c.ptr, src, dst, count, t, op, stream); + } + + pub fn allGather(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: Stream) Error!void { + return c.vtable.all_gather(c.ptr, src, dst, count, t, stream); + } + + /// `allGather` of `stride` elements a rank where only rank r's first lens[r] bytes matter (`VTable.all_gather_v`). + pub fn allGatherV(c: Collective, src: DevicePtr, dst: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void { + if (c.vtable.all_gather_v) |f| return f(c.ptr, src, dst, stride, lens, t, stream); + return c.vtable.all_gather(c.ptr, src, dst, stride, t, stream); + } + + /// Whether `allGatherV` of `bytes`-byte strides moves only the lengths here (callers agree first: `varAgreed`). + pub fn varFits(c: Collective, bytes: usize) bool { + const f = c.vtable.var_fits orelse return false; + return f(c.ptr, bytes); + } + + /// Every rank's transport takes device lengths: one agreement a communicator, called alike outside any graph. + pub fn varAgreed(c: Collective) Error!bool { + const f = c.vtable.var_agreed orelse return false; + return f(c.ptr); + } + + pub fn exchange(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.exchange(c.ptr, src, dst, count, t, peer, stream); + } + + pub fn send(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.send(c.ptr, buf, count, t, peer, stream); + } + + pub fn recv(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.recv(c.ptr, buf, count, t, peer, stream); + } + + pub fn barrier(c: Collective) Error!void { + return c.vtable.barrier(c.ptr); + } + + pub fn check(c: Collective) Error!void { + return c.vtable.check(c.ptr); + } + + pub fn abort(c: Collective) void { + c.vtable.abort(c.ptr); + } +}; + +/// Refuses a peer that is this rank or out of range. +pub fn checkPeer(rank: u32, world: u32, peer: u32) Error!void { + if (peer == rank or peer >= world) return error.Invalid; +} + +test "dtype sizes" { + try std.testing.expectEqual(@as(usize, 2), DType.bf16.size()); + try std.testing.expectEqual(@as(usize, 8), DType.i64.size()); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig new file mode 100644 index 0000000000..13647f5921 --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,245 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); + +pub const Backend = enum { + /// NCCL for everything (the default). + nccl, + /// Exchanges of at most `mailbox_max_bytes` a rank through the one-shot mailbox kernel, NCCL above that. + mailbox, + /// The same split between two hosts: the small exchanges over RoCE (doorbell + RDMA proxy), NCCL above. + roce, +}; + +/// What `roce` does when its setup or load-time probe fails on any rank (TF_TP_ROCE_FALLBACK). +pub const Fallback = enum { nccl, @"error" }; + +pub const Error = error{BadSetting}; + +/// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). +pub const PlanPin = union(enum) { auto, cpu: u32 }; + +pub const Config = struct { + world: u32 = 2, + rank: u32 = 0, + /// Rank 0's address as the other ranks reach it. + master: []const u8 = "localhost", + /// The rendezvous port (TF_TP_PORT); the link listens on `link_port`. + port: u16 = 29500, + /// Rank 0 listens here: TF_TP_PLAN_PORT, else `port` + 7. + link_port: u16 = 29507, + /// How long rank 0 waits for the others and they keep retrying the connect (TF_TP_PLAN_TIMEOUT). + connect_timeout_ns: u64 = 300 * std.time.ns_per_s, + /// CUDA device ordinal this rank uses. + device: u32 = 0, + backend: Backend = .nccl, + mailbox_max_bytes: usize = 256 * 1024, + /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). + mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, + failfast: bool = true, + /// After a failure: time for the HTTP threads to answer before exit 70. + grace_ns: u64 = std.time.ns_per_s, + /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). + report_ns: u64 = std.time.ns_per_s / 2, + nccl_lib: []const u8 = "libnccl.so.2", + /// `name[:gid],...` RDMA devices for `roce` (null: every active Ethernet port with an IPv4 RoCE v2 GID). + roce_hca: ?[]const u8 = null, + /// Stripe over at most this many HCAs (auto detection only). + roce_hcas: u8 = 2, + roce_tc: u8 = 0, + /// TF_TP_ROCE_FAST as the RoCE kernel's `opts` word: bit 0 on, bits 8-15 pollers, bits 16-31 stagger ns (local). + roce_fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest RoCE grid, a power of two 1..16; launch shape only, the bytes are the same. + roce_blocks: u32 = 16, + /// Pin the proxy thread (it busy-spins one core between exchanges). + roce_cpu: ?u32 = null, + /// `shm` runs the RoCE protocol on one host without RDMA (tests); `verbs` is the real wire. + roce_wire: enum { verbs, shm } = .verbs, + roce_fallback: Fallback = .nccl, + /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). + plan_spin_ns: u64 = 0, + plan_pin: ?PlanPin = null, + + /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. + pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { + var c: Config = .{}; + c.world = try int(u32, get, "TF_TP_WORLD", c.world, 1, 64); + c.rank = try int(u32, get, "TF_TP_RANK", c.rank, 0, c.world - 1); + if (get("TF_TP_MASTER")) |m| c.master = m; + c.port = try int(u16, get, "TF_TP_PORT", c.port, 1, 65535 - 8); + c.link_port = try int(u16, get, "TF_TP_PLAN_PORT", c.port + 7, 1, 65535); + c.connect_timeout_ns = try seconds(get, "TF_TP_PLAN_TIMEOUT", 300, 0.1, 86400); + c.device = try int(u32, get, "TF_TP_DEVICE", c.device, 0, 63); + const backend = "TF_COMM_BACKEND"; + if (get(backend)) |b| c.backend = std.meta.stringToEnum(Backend, std.mem.trim(u8, b, " ")) orelse { + log.warn("{s}={s}: expected nccl, mailbox or roce", .{ backend, b }); + return error.BadSetting; + }; + c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); + if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); + c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); + c.report_ns = try seconds(get, "TF_TP_PEER_MEM_S", 0.5, 0.01, 60); + if (get("TF_NCCL_LIB")) |p| c.nccl_lib = p; + c.roce_hca = get("TF_TP_ROCE_HCA"); + c.roce_hcas = try int(u8, get, "TF_TP_ROCE_HCAS", 2, 1, 2); + c.roce_tc = try int(u8, get, "TF_TP_ROCE_TC", try int(u8, get, "NCCL_IB_TC", 0, 0, 255), 0, 255); + if (try int(u8, get, "TF_TP_ROCE_FAST", 0, 0, 1) == 1) { + const pollers = try int(u32, get, "TF_TP_ROCE_POLLERS", 1, 1, 8); + const stagger = try int(u32, get, "TF_TP_ROCE_STAGGER_NS", 150, 0, 65535); + c.roce_fast = 1 | (pollers << 8) | (stagger << 16); + } + c.roce_blocks = try int(u32, get, "TF_TP_ROCE_BLOCKS", 16, 1, 16); + if (c.roce_blocks & (c.roce_blocks - 1) != 0) { + log.warn("TF_TP_ROCE_BLOCKS={d}: expected a power of two", .{c.roce_blocks}); + return error.BadSetting; + } + if (get("TF_TP_ROCE_CPU")) |_| c.roce_cpu = try int(u32, get, "TF_TP_ROCE_CPU", 0, 0, 4095); + if (get("TF_TP_ROCE_WIRE")) |w| c.roce_wire = std.meta.stringToEnum(@TypeOf(c.roce_wire), w) orelse { + log.warn("TF_TP_ROCE_WIRE={s}: expected verbs or shm", .{w}); + return error.BadSetting; + }; + const fallback = "TF_TP_ROCE_FALLBACK"; + if (get(fallback)) |f| c.roce_fallback = std.meta.stringToEnum(Fallback, f) orelse { + log.warn("{s}={s}: expected nccl or error", .{ fallback, f }); + return error.BadSetting; + }; + if (get("TF_TP_PLAN_LINK")) |l| { + const mode = std.mem.trim(u8, l, " "); + if (std.mem.eql(u8, mode, "rdma")) { + c.plan_spin_ns = 1000 * try int(u64, get, "TF_TP_PLAN_SPIN_US", 5000, 0, 10_000_000); + } else if (!std.mem.eql(u8, mode, "tcp") and !std.mem.eql(u8, mode, "nccl")) { + log.warn("TF_TP_PLAN_LINK={s}: expected nccl, tcp or rdma", .{l}); + return error.BadSetting; + } + } + if (get("TF_TP_PLAN_PIN")) |raw| { + const v = std.mem.trim(u8, raw, " "); + if (std.mem.eql(u8, v, "auto")) { + c.plan_pin = .auto; + } else if (!std.mem.eql(u8, v, "-1") and !std.mem.eql(u8, v, "off")) { + c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; + } + } + return c; + } + + pub fn fromEnv() Error!Config { + return parse(&envGet); + } + + /// The settings both ranks must share (everything but rank, device and timeouts), as bytes to compare. + pub fn digest(c: Config) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(std.mem.asBytes(&c.world)); + h.update(std.mem.asBytes(&c.backend)); + h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + +fn envGet(name: []const u8) ?[]const u8 { + var buf: [64]u8 = undefined; + if (name.len >= buf.len) return null; + @memcpy(buf[0..name.len], name); + buf[name.len] = 0; + const v = std.c.getenv(buf[0..name.len :0]) orelse return null; + const s = std.mem.span(v); + return if (s.len == 0) null else s; +} + +fn int(comptime T: type, get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: T, lo: T, hi: T) Error!T { + const raw = get(name) orelse return default; + const v = std.fmt.parseInt(T, std.mem.trim(u8, raw, " "), 10) catch { + log.warn("{s}={s}: expected an integer", .{ name, raw }); + return error.BadSetting; + }; + if (v < lo or v > hi) { + log.warn("{s}={d}: expected {d}..{d}", .{ name, v, lo, hi }); + return error.BadSetting; + } + return v; +} + +fn seconds(get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: f64, lo: f64, hi: f64) Error!u64 { + const v = if (get(name)) |raw| std.fmt.parseFloat(f64, std.mem.trim(u8, raw, " ")) catch { + log.warn("{s}={s}: expected seconds", .{ name, raw }); + return error.BadSetting; + } else default; + if (!(v >= lo and v <= hi)) { + log.warn("{s}={d}: expected {d}..{d} s", .{ name, v, lo, hi }); + return error.BadSetting; + } + return @intFromFloat(v * std.time.ns_per_s); +} + +const Fake = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "defaults: the plan link 7 ports past the rendezvous, NCCL, fail-fast on" { + Fake.pairs = &.{ .{ "TF_TP_RANK", "1" }, .{ "TF_TP_PORT", "29600" } }; + const c = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u32, 1), c.rank); + try std.testing.expectEqual(@as(u16, 29607), c.link_port); + try std.testing.expectEqual(Backend.nccl, c.backend); + try std.testing.expectEqual(std.time.ns_per_s, c.grace_ns); + try std.testing.expectEqual(std.time.ns_per_s / 2, c.report_ns); +} + +test "bad settings are refused and the digest ignores the rank" { + Fake.pairs = &.{.{ "TF_TP_RANK", "2" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{.{ "TF_COMM_BACKEND", "rdma" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_ROCE_HCA", "rocep1s0f1,roceP2p1s0f1" }, .{ "NCCL_IB_TC", "106" } }; + const r = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, r.backend); + try std.testing.expectEqual(@as(u8, 106), r.roce_tc); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", r.roce_hca.?); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "64" }, .{ "TF_TP_FAILFAST_GRACE_S", "0.25" } }; + const a = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(usize, 64 * 1024), a.mailbox_max_bytes); + try std.testing.expectEqual(std.time.ns_per_s / 4, a.grace_ns); + var b = a; + b.rank = 1; + try std.testing.expectEqualSlices(u8, &a.digest(), &b.digest()); + b.mailbox_max_bytes = 1024; + try std.testing.expect(!std.mem.eql(u8, &a.digest(), &b.digest())); +} + +test "the plan link's spin and pin, the RoCE grid cap" { + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" }, .{ "TF_TP_ROCE_FALLBACK", "error" }, .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_PIN", "auto" } }; + const p = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, p.backend); + try std.testing.expectEqual(@as(usize, 1024 * 1024), p.mailbox_max_bytes); + try std.testing.expectEqual(Fallback.@"error", p.roce_fallback); + try std.testing.expectEqual(@as(u64, 5000 * 1000), p.plan_spin_ns); + try std.testing.expectEqual(PlanPin.auto, p.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "tcp" }, .{ "TF_TP_PLAN_PIN", "7" } }; // tcp blocks at once; an explicit CPU + const o = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 0), o.plan_spin_ns); + try std.testing.expectEqual(PlanPin{ .cpu = 7 }, o.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_SPIN_US", "200" }, .{ "TF_TP_PLAN_PIN", "-1" } }; + const s = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 200 * 1000), s.plan_spin_ns); + try std.testing.expect(s.plan_pin == null); + Fake.pairs = &.{.{ "TF_TP_PLAN_LINK", "udp" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; // the RoCE grid cap: 16 unset, powers of two only + try std.testing.expectEqual(@as(u32, 16), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "4" }}; + try std.testing.expectEqual(@as(u32, 4), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "12" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; +} diff --git a/zig/src/fabric/tp/fate.zig b/zig/src/fabric/tp/fate.zig new file mode 100644 index 0000000000..5084c4648b --- /dev/null +++ b/zig/src/fabric/tp/fate.zig @@ -0,0 +1,261 @@ +//! Fail fast on every rank: a failure or a vanished rank ends all ranks with exit 70 about a second later. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const posix = std.posix; +/// Linux POLLRDHUP: the peer shut its side down (std has no name for it). +const pollrdhup: i16 = 0x2000; + +pub const exit_code: u8 = 70; +const magic: u32 = 0x46463431; +const kind_mem: u32 = 1; +const kind_fail: u32 = 2; +/// rank 0's clean stop: the hang-up that follows is no failure (sent before rank 0 closes its sockets) +const kind_stop: u32 = 3; +const Frame = extern struct { magic: u32, kind: u32, value: i64 }; + +pub const max_ranks = 64; +pub const max_hooks = 8; + +/// A callback with its context (an abort of a collective, rank 0's request failing); `reason` lives until exit. +pub const Hook = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, reason: []const u8) void }; + +pub const Fate = struct { + rank: u32, + world: u32, + fds: [max_ranks]sock.Fd, + grace_ns: u64, + report_ns: u64, + /// This rank's usable memory in bytes for the MEM report (default: MemAvailable). + usable: *const fn () u64 = memAvailable, + /// The process exit; tests replace it (it may return there). + exit_fn: *const fn (u8) void = exitNow, + hooks: [max_hooks]Hook = undefined, + n_hooks: std.atomic.Value(u32) = .init(0), + dead: std.atomic.Value(bool) = .init(false), + closing: std.atomic.Value(bool) = .init(false), + /// Set once the exit was asked for (tests read it; the process is gone in serving). + exited: std.atomic.Value(bool) = .init(false), + reason_buf: [512]u8 = undefined, + reason_len: usize = 0, + send_lock: std.atomic.Value(bool) = .init(false), + /// Rank 0: each rank's last report (bytes) and its monotonic time. + peer_mem: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + peer_mem_at: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + wake: [2]sock.Fd = .{ -1, -1 }, + thread: ?std.Thread = null, + + /// The channel over the bootstrap's sockets (`Bootstrap.release`); `start` runs its thread. + pub fn init(cfg: Config, fds: [max_ranks]sock.Fd) Fate { + return .{ .rank = cfg.rank, .world = cfg.world, .fds = fds, .grace_ns = cfg.grace_ns, .report_ns = cfg.report_ns }; + } + + pub fn start(f: *Fate) !void { + var p: [2]c_int = undefined; + if (std.c.pipe2(&p, .{ .CLOEXEC = true }) != 0) return error.SocketFailed; + f.wake = p; + f.thread = if (f.rank == 0) try std.Thread.spawn(.{}, watch0, .{f}) else try std.Thread.spawn(.{}, watchK, .{f}); + } + + /// Runs `hook` when the ranks go down; collectives register their abort here. + pub fn onFail(f: *Fate, hook: Hook) void { + const i = f.n_hooks.load(.acquire); + std.debug.assert(i < max_hooks); + f.hooks[i] = hook; + f.n_hooks.store(i + 1, .release); + } + + /// This rank failed (`where`: what it was doing): every rank goes down. Returns only in tests. + pub fn fatal(f: *Fate, where: []const u8, err: anyerror) void { + var buf: [256]u8 = undefined; + const msg = std.fmt.bufPrint(&buf, "rank {d} failed in {s}: {t}", .{ f.rank, where, err }) catch "rank failed"; + f.die(msg, true); + } + + /// This rank failed while booting: every rank goes down, and the thread is joined before the caller frees this. + pub fn failBoot(f: *Fate, where: []const u8, err: anyerror) void { + f.fatal(where, err); + f.close(); + } + + /// A normal stop is coming (rank 0 before it tells the others; the others when told): later hang-ups are fine. + pub fn expectStop(f: *Fate) void { + f.closing.store(true, .release); + } + + /// Rank 0's clean stop of every rank: a STOP frame to each ahead of its close, so followers end quietly. + pub fn announceStop(f: *Fate) void { + if (f.closing.swap(true, .acq_rel)) return; + if (f.rank != 0) return; + const h: Frame = .{ .magic = magic, .kind = kind_stop, .value = 0 }; + f.lock(); + defer f.unlock(); + for (f.fds[1..f.world]) |fd| if (fd >= 0) { + _ = sock.trySend(fd, std.mem.asBytes(&h)); + }; + } + + /// A normal stop: the thread ends and the sockets close. + pub fn close(f: *Fate) void { + f.closing.store(true, .release); + if (f.wake[1] >= 0) _ = std.c.write(f.wake[1], "x", 1); + if (f.thread) |t| t.join(); + f.thread = null; + for (&f.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + for (f.wake) |fd| if (fd >= 0) sock.close(fd); + f.wake = .{ -1, -1 }; + } + + /// Rank 0: rank `r`'s last reported usable bytes and the age of the report, or null before its first. + pub fn peerMem(f: *const Fate, r: u32) ?struct { bytes: u64, age_ns: u64 } { + const at = f.peer_mem_at[r].load(.acquire); + if (at == 0) return null; + return .{ .bytes = f.peer_mem[r].load(.acquire), .age_ns = sock.nowNs() -| at }; + } + + pub fn reason(f: *const Fate) []const u8 { + return f.reason_buf[0..f.reason_len]; + } + + fn lock(f: *Fate) void { + while (f.send_lock.swap(true, .acquire)) std.atomic.spinLoopHint(); + } + + fn unlock(f: *Fate) void { + f.send_lock.store(false, .release); + } + + fn die(f: *Fate, why: []const u8, own: bool) void { + if (f.dead.swap(true, .acq_rel)) { + // the first caller exits the process; a second (the fate thread or a round thread) waits for it + while (!f.exited.load(.acquire)) sock.sleepNs(10 * std.time.ns_per_ms); + return; + } + const n = @min(why.len, f.reason_buf.len); + @memcpy(f.reason_buf[0..n], why[0..n]); + f.reason_len = n; + const r = f.reason(); + log.warn("[tensorfold] tp fail-fast: rank {d}: {s}; exiting ({d}) in {d} ms so every rank restarts", .{ f.rank, r, exit_code, f.grace_ns / std.time.ns_per_ms }); + // aborts first: kernels waiting on the dead peer end now; a hook that blocks cannot hold the exit + if (std.Thread.spawn(.{}, runHooks, .{ f, r })) |t| t.detach() else |_| runHooks(f, r); + if (!f.closing.load(.acquire)) { + const head: Frame = .{ .magic = magic, .kind = kind_fail, .value = @intCast(r.len) }; + f.lock(); + for (f.fds[0..f.world]) |fd| if (fd >= 0) { + if (own or f.rank == 0) { + _ = sock.trySend(fd, std.mem.asBytes(&head)); + _ = sock.trySend(fd, r); + } + sock.hangup(fd); + }; + f.unlock(); + } + sock.sleepNs(f.grace_ns); + f.exited.store(true, .release); + f.exit_fn(exit_code); + } + + fn runHooks(f: *Fate, r: []const u8) void { + const n = f.n_hooks.load(.acquire); + for (f.hooks[0..n]) |h| h.call(h.ctx, r); + } + + /// Reads one frame from `fd`: a MEM report, a FAIL (the ranks go down), or the end (the peer is gone). + fn readFrame(f: *Fate, peer: u32, fd: sock.Fd) bool { + var h: Frame = undefined; + const deadline = sock.nowNs() + std.time.ns_per_s; + sock.recvExact(fd, std.mem.asBytes(&h), deadline) catch { + if (f.closing.load(.acquire)) return false; + var buf: [128]u8 = undefined; + f.die(std.fmt.bufPrint(&buf, "rank {d} closed the fate channel (it failed, exited or was killed)", .{peer}) catch "peer gone", false); + return false; + }; + if (h.magic != magic) { + f.die("a frame that is not the fate channel's", false); + return false; + } + if (h.kind == kind_stop) { + // rank 0 stops every rank: its hang-up is next, and no failure + f.closing.store(true, .release); + return false; + } + if (h.kind == kind_mem) { + f.peer_mem[peer].store(@intCast(@max(h.value, 0)), .release); + f.peer_mem_at[peer].store(sock.nowNs(), .release); + return true; + } + var text: [512]u8 = undefined; + const len: usize = @intCast(std.math.clamp(h.value, 0, text.len)); + sock.recvExact(fd, text[0..len], deadline) catch {}; + f.die(if (len > 0) text[0..len] else "a peer failed", false); + return false; + } + + fn watch0(f: *Fate) void { + var p: [max_ranks + 1]std.c.pollfd = undefined; + var who: [max_ranks + 1]u32 = undefined; + var n: usize = 0; + for (f.fds[1..f.world], 1..) |fd, r| if (fd >= 0) { + p[n] = .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }; + who[n] = @intCast(r); + n += 1; + }; + p[n] = .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }; + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + if (std.c.poll(&p, @intCast(n + 1), -1) <= 0) continue; + for (p[0..n], who[0..n]) |*q, r| if (q.revents != 0) { + if (!f.readFrame(r, q.fd)) return; + }; + } + } + + fn watchK(f: *Fate) void { + const fd = f.fds[0]; + var p = [_]std.c.pollfd{ + .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }, + .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }, + }; + const ms: c_int = @intCast(@max(f.report_ns / std.time.ns_per_ms, 1)); + f.report(fd); + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + const ready = std.c.poll(&p, 2, ms); + if (ready == 0) { + f.report(fd); + continue; + } + if (ready > 0 and p[0].revents != 0 and !f.readFrame(0, fd)) return; + } + } + + /// One MEM frame, never blocking (a full buffer drops it: the next comes one period later). + fn report(f: *Fate, fd: sock.Fd) void { + const h: Frame = .{ .magic = magic, .kind = kind_mem, .value = @intCast(f.usable()) }; + f.lock(); + defer f.unlock(); + _ = sock.trySend(fd, std.mem.asBytes(&h)); + } +}; + +fn exitNow(code: u8) void { + std.c._exit(code); +} + +/// MemAvailable from /proc/meminfo in bytes (0 when unreadable). +pub fn memAvailable() u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return 0; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return 0; + const text = buf[0..@intCast(n)]; + const key = "MemAvailable:"; + const at = std.mem.indexOf(u8, text, key) orelse return 0; + var it = std.mem.tokenizeAny(u8, text[at + key.len ..], " \n"); + const kb = std.fmt.parseInt(u64, it.next() orelse return 0, 10) catch return 0; + return kb * 1024; +} diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig new file mode 100644 index 0000000000..6019055f50 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,146 @@ +//! The fate channel between two ranks over real TCP: a failure or a vanished rank ends both, a stop ends neither. +const std = @import("std"); +const sock = @import("sock.zig"); +const fate = @import("fate.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Fate = fate.Fate; + +var exits: std.atomic.Value(u32) = .init(0); +var hooks: std.atomic.Value(u32) = .init(0); + +fn recordExit(code: u8) void { + std.debug.assert(code == fate.exit_code); + _ = exits.fetchAdd(1, .acq_rel); +} + +fn hook(_: *anyopaque, reason: []const u8) void { + std.debug.assert(reason.len > 0); + _ = hooks.fetchAdd(1, .acq_rel); +} + +const grace = 50 * std.time.ns_per_ms; + +/// Two connected ranks' fate channels (rank 0 first), started, with recorded exits and hooks. +fn pair(fates: *[2]Fate) !void { + return pairAt(.{ &fates[0], &fates[1] }); +} + +fn pairAt(fates: [2]*Fate) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var fds0: [fate.max_ranks]sock.Fd = @splat(-1); + var fds1: [fate.max_ranks]sock.Fd = @splat(-1); + fds0[1] = b; + fds1[0] = a; + for (fates, [2][fate.max_ranks]sock.Fd{ fds0, fds1 }, 0..) |f, fds, r| { + f.* = Fate.init(.{ .world = 2, .rank = @intCast(r), .grace_ns = grace, .report_ns = 20 * std.time.ns_per_ms }, fds); + f.exit_fn = recordExit; + f.usable = fakeUsable; + f.onFail(.{ .ctx = f, .call = hook }); + } + for (fates) |f| try f.start(); +} + +fn fakeUsable() u64 { + return 123 << 20; +} + +fn waitExits(want: u32) !u64 { + const t0 = sock.nowNs(); + while (exits.load(.acquire) < want) { + if (sock.nowNs() - t0 > 5 * std.time.ns_per_s) return error.Timeout; + sock.sleepNs(std.time.ns_per_ms); + } + return sock.nowNs() - t0; +} + +fn reset() void { + exits.store(0, .release); + hooks.store(0, .release); +} + +test "rank 1's failure ends both ranks with its reason" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + const mem = fates[0].peerMem(1) orelse return error.NoReport; + try std.testing.expectEqual(@as(u64, 123 << 20), mem.bytes); + fates[1].fatal("prefill", error.OutOfMemory); + const took = try waitExits(2); + try std.testing.expect(took < grace + 500 * std.time.ns_per_ms); + try std.testing.expectEqual(@as(u32, 2), hooks.load(.acquire)); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 failed in prefill: OutOfMemory") != null); +} + +test "rank 0's failure ends rank 1" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + fates[0].fatal("window", error.Unexpected); + _ = try waitExits(2); + try std.testing.expect(std.mem.indexOf(u8, fates[1].reason(), "rank 0 failed in window") != null); +} + +test "a rank that vanishes without a word ends the other" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + // rank 1 killed: its socket closes with no FAIL frame (its own fate thread ends quietly first) + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); +} + +test "a normal stop ends nobody" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + fates[0].expectStop(); + fates[1].expectStop(); + fates[1].close(); + fates[0].close(); + sock.sleepNs(2 * grace); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); +} + +test "a rank failing at boot ends both with its reason, and its channel is down before its owner frees it" { + // a follower's boot failed after its session came up; its Fate is freed while the fate thread could still read it + reset(); + // rank 1's Fate where a Model keeps it: on the heap, its pages unmapped when freed (a stale read faults) + const pa = std.heap.page_allocator; + var f0: Fate = undefined; + const f1 = try pa.create(Fate); + try pairAt(.{ &f0, f1 }); + defer f0.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + f1.failBoot("boot", error.RopeTooShort); + try std.testing.expect(f1.thread == null); + pa.destroy(f1); + _ = try waitExits(2); + // several report periods with rank 1's Fate gone: its thread would have reported from freed memory + sock.sleepNs(5 * 20 * std.time.ns_per_ms); + try std.testing.expect(std.mem.indexOf(u8, f0.reason(), "rank 1 failed in boot: RopeTooShort") != null); +} + +test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order the ranks close in" { + // rank 0 ended its run and closed; the follower, still tearing down, must not see the hang-up as a failure + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + sock.sleepNs(60 * std.time.ns_per_ms); + fates[0].announceStop(); + fates[0].close(); // rank 0 gone before rank 1 even started its own stop + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + fates[1].close(); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} diff --git a/zig/src/fabric/tp/gpu_bench.zig b/zig/src/fabric/tp/gpu_bench.zig new file mode 100644 index 0000000000..99831e6d77 --- /dev/null +++ b/zig/src/fabric/tp/gpu_bench.zig @@ -0,0 +1,82 @@ +//! Latency and bandwidth of each backend's collectives, eager and graphed; RESULT lines are JSON, busbw as nccl-tests. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Ctx = @import("gpu_checks.zig").Ctx; +const Collective = tp.Collective; + +pub const Kind = enum { all_reduce, all_gather }; + +fn issue(c: *Ctx, comm: Collective, kind: Kind, bytes: usize) !void { + const count = bytes / 2; + switch (kind) { + .all_reduce => try comm.allReduce(c.in.ptr, c.out.ptr, count, .bf16, .sum, c.stream.handle), + .all_gather => try comm.allGather(c.in.ptr, c.out.ptr, count, .bf16, c.stream.handle), + } +} + +/// GPU microseconds a call, `iters` calls issued back to back on the stream. +fn eager(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, iters: u32) !f64 { + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + for (0..10) |_| try issue(c, comm, kind, bytes); + try c.stream.synchronize(); + try comm.barrier(); + try e0.record(c.stream); + for (0..iters) |_| try issue(c, comm, kind, bytes); + try e1.record(c.stream); + try e1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(iters)); +} + +/// GPU microseconds a call inside one graph of `per_graph` calls, replayed `replays` times. +fn graphed(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, per_graph: u32, replays: u32) !f64 { + try cuda.graph.beginCapture(c.stream, .thread_local); + const err: ?tp.collective.Error = blk: { + for (0..per_graph) |_| issue(c, comm, kind, bytes) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (err) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + try exec.upload(c.stream); + try exec.launchOn(c.stream); + try c.stream.synchronize(); + try comm.barrier(); + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + try e0.record(c.stream); + for (0..replays) |_| try exec.launchOn(c.stream); + try e1.record(c.stream); + try e1.synchronize(); + try comm.check(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(per_graph * replays)); +} + +/// Every power of two from `lo` to `hi` bytes a rank, both kinds, eager and graphed. +pub fn sweep(c: *Ctx, comm: Collective, label: []const u8, lo: usize, hi: usize) !void { + for ([_]Kind{ .all_reduce, .all_gather }) |kind| { + var bytes = lo; + while (bytes <= hi) : (bytes *= 2) { + const per_graph: u32 = if (bytes >= 1 << 20) 20 else 100; + const e = try eager(c, comm, kind, bytes, if (bytes >= 1 << 20) 50 else 200); + const g = try graphed(c, comm, kind, bytes, per_graph, 5); + const moved: f64 = @floatFromInt(if (kind == .all_gather) 2 * bytes else bytes); + const factor: f64 = if (kind == .all_gather) 0.5 else 1.0; + const busbw = moved / (g * 1e-6) * factor / 1e9; + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"{t}\",\"bytes\":{d},\"eager_us\":{d:.2},\"graph_us\":{d:.2},\"busbw_GBps\":{d:.2}}}\n", .{ label, kind, bytes, e, g, busbw }); + } + } +} + +/// One decode window's exchanges: `n` all-gathers of one row's bf16 partials (`row_bytes` a rank) in one graph. +pub fn decodeWindow(c: *Ctx, comm: Collective, label: []const u8, n: u32, row_bytes: usize) !void { + const per = try graphed(c, comm, .all_gather, row_bytes, n, 20); + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"decode_window\",\"exchanges\":{d},\"bytes\":{d},\"per_exchange_us\":{d:.2},\"window_us\":{d:.1}}}\n", .{ label, n, row_bytes, per, per * @as(f64, @floatFromInt(n)) }); +} diff --git a/zig/src/fabric/tp/gpu_checks.zig b/zig/src/fabric/tp/gpu_checks.zig new file mode 100644 index 0000000000..5dc4ed6d4c --- /dev/null +++ b/zig/src/fabric/tp/gpu_checks.zig @@ -0,0 +1,188 @@ +//! Every collective on real GPUs against the host reference, eager and replayed from a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Collective = tp.Collective; +const DType = tp.DType; +const Op = tp.Op; +const reference = tp.reference; + +pub const Ctx = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + stream: cuda.Stream, + rank: u32, + /// Device buffers sized for the largest case: input, output (world x), a spare. + in: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + /// Host scratch: this rank's input, the peer's, the expected and the downloaded output. + mine: []u8, + peer: []u8, + want: []u8, + got: []u8, + failures: u32 = 0, + passes: u32 = 0, + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, rank: u32, max_bytes: usize) !Ctx { + return .{ + .d = d, + .gpa = gpa, + .stream = try cuda.Stream.init(d, true), + .rank = rank, + .in = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .out = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .mine = try gpa.alloc(u8, max_bytes), + .peer = try gpa.alloc(u8, max_bytes), + .want = try gpa.alloc(u8, 2 * max_bytes), + .got = try gpa.alloc(u8, 2 * max_bytes), + }; + } + + pub fn deinit(c: *Ctx) void { + c.stream.deinit(); + c.in.free(); + c.out.free(); + for ([_][]u8{ c.mine, c.peer, c.want, c.got }) |b| c.gpa.free(b); + } + + fn verdict(c: *Ctx, ok: bool, comptime fmt: []const u8, args: anytype) void { + if (ok) c.passes += 1 else c.failures += 1; + std.debug.print("{s} " ++ fmt ++ "\n", .{if (ok) "PASS" else "FAIL"} ++ args); + } + + /// Both ranks' inputs for case `seed`: this rank's into `mine`, the peer's into `peer` (n bytes each). + fn inputs(c: *Ctx, n: usize, t: DType, seed: u64) void { + reference.fill(c.mine[0..n], t, c.rank, seed); + reference.fill(c.peer[0..n], t, 1 - c.rank, seed); + } + + fn rankInput(c: *Ctx, r: u32, n: usize) []const u8 { + return if (r == c.rank) c.mine[0..n] else c.peer[0..n]; + } + + fn compare(c: *Ctx, comm: Collective, label: []const u8, what: []const u8, t: DType, n: usize, out_off: usize, want: []const u8) !void { + try c.stream.synchronize(); + comm.check() catch |e| { + c.verdict(false, "{s} {s} {t} {d} B: transport error {t}", .{ label, what, t, n, e }); + return; + }; + try c.out.download(out_off, c.got[0..want.len]); + const ok = std.mem.eql(u8, c.got[0..want.len], want); + var first: usize = 0; + if (!ok) first = std.mem.indexOfDiff(u8, c.got[0..want.len], want) orelse 0; + c.verdict(ok, "{s} {s} {t} {d} B a rank{s}", .{ label, what, t, n, if (ok) "" else " (first difference at a byte shown below)" }); + if (!ok) std.debug.print(" byte {d}: got {x} want {x}\n", .{ first, c.got[first], want[first] }); + } + + /// all-reduce (out of place, in place, misaligned by `skew` bytes), all-gather (out of place, in place), exchange. + pub fn collectives(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const reduce_types = [_]DType{ .f32, .bf16, .f16, .i32 }; + for (reduce_types, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + const count = n / t.size(); + c.inputs(n, t, seed + ti); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, t, .sum); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum", t, n, 0, c.want[0..n]); + // in place, at an offset of one element (not 16-byte aligned for 2-byte types) + const skew = t.size(); + try c.out.uploadAsync(skew, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.out.ptr + skew, c.out.ptr + skew, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum in place, skewed", t, n, skew, c.want[0..n]); + } + for ([_]Op{ .max, .min }) |op| { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed + 11); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .f32, op); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, n / 4, .f32, op, c.stream.handle); + try c.compare(comm, label, if (op == .max) "all-reduce max" else "all-reduce min", .f32, n, 0, c.want[0..n]); + } + for ([_]DType{ .u8, .bf16, .f32 }, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + c.inputs(n, t, seed + 20 + ti); + @memcpy(c.want[0..n], c.rankInput(0, n)); + @memcpy(c.want[n..][0..n], c.rankInput(1, n)); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allGather(c.in.ptr, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather", t, n, 0, c.want[0 .. 2 * n]); + try c.out.fill8(0, c.stream.handle); + try c.out.uploadAsync(c.rank * n, c.mine[0..n], c.stream.handle); + try comm.allGather(c.out.ptr + c.rank * n, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather in place", t, n, 0, c.want[0 .. 2 * n]); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.exchange(c.in.ptr, c.out.ptr, n / t.size(), t, 1 - c.rank, c.stream.handle); + try c.compare(comm, label, "exchange", t, n, 0, c.peer[0..n]); + } + } + + /// send then recv one way, then the other way (NCCL behind every backend). + pub fn pointToPoint(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + for (0..2) |dir| { + const sender: u32 = @intCast(dir); + if (c.rank == sender) { + try comm.send(c.in.ptr, n / 4, .f32, 1 - c.rank, c.stream.handle); + try c.stream.synchronize(); + } else { + try comm.recv(c.out.ptr, n / 4, .f32, sender, c.stream.handle); + try c.compare(comm, label, "send/recv", .f32, n, 0, c.peer[0..n]); + } + } + } + + /// A captured chain replayed with fresh inputs, eager calls between replays (the mailbox epoch spans both). + pub fn graphs(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, replays: u32) !void { + const n = n_bytes - n_bytes % 16; + const half = n / 2; // bf16 elements + try cuda.graph.beginCapture(c.stream, .thread_local); + // x1 = sum(x0) [bf16], x2 = gather(x1), x3 = sum(x2) [f32 view], x4 = exchange(x3) + const x0 = c.in.ptr; + const x1 = c.out.ptr; + const x2 = c.in.ptr + n; + const x3 = c.out.ptr + n; + const x4 = c.out.ptr + 3 * n; + const ops: ?tp.collective.Error = blk: { + comm.allReduce(x0, x1, half, .bf16, .sum, c.stream.handle) catch |e| break :blk e; + comm.allGather(x1, x2, half, .bf16, c.stream.handle) catch |e| break :blk e; + comm.allReduce(x2, x3, n / 2, .f32, .sum, c.stream.handle) catch |e| break :blk e; + comm.exchange(x3, x4, n / 2, .f32, 1 - c.rank, c.stream.handle) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (ops) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + const nodes = try g.nodeCount(); + const want1 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want1); + const want3 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want3); + for (0..replays) |k| { + c.inputs(n, .bf16, 1000 + k); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + reference.reduce(want1[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + @memcpy(want1[n..][0..n], want1[0..n]); // the gather of equal sums: twice the same + reference.reduce(want3[0 .. 2 * n], &.{ want1[0 .. 2 * n], want1[0 .. 2 * n] }, .f32, .sum); + try exec.launchOn(c.stream); + try c.compare(comm, label, "graph replay: sum, gather, sum, exchange", .bf16, n, n, want3[0 .. 2 * n]); + try c.compare(comm, label, "graph replay: exchange of the sum", .f32, 2 * n, 3 * n, want3[0 .. 2 * n]); + // an eager call between replays + try c.collectivesQuick(comm, label, k); + } + std.debug.print(" ({s}: {d} graph nodes, {d} replays)\n", .{ label, nodes, replays }); + } + + fn collectivesQuick(c: *Ctx, comm: Collective, label: []const u8, k: usize) !void { + const n = 4096; + c.inputs(n, .bf16, 5000 + k); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + try c.in.uploadAsync(5 * n, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr + 5 * n, c.out.ptr + 5 * n, n / 2, .bf16, .sum, c.stream.handle); + try c.compare(comm, label, "eager between replays", .bf16, n, 5 * n, c.want[0..n]); + } +}; diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..c5668a7b06 --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,92 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host`: one rank of a two-process TP=2 test on real GPUs. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const tp_kernel = @import("tp_kernel"); +const checks = @import("gpu_checks.zig"); +const bench = @import("gpu_bench.zig"); + +const max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host\n", .{}); + return 2; + } + const cfg = try tp.Config.fromEnv(); + if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); + defer ctx.deinit(); + var name_buf: [128]u8 = undefined; + const t0 = tp.sock.nowNs(); + var s: tp.Session = undefined; + try s.open(&driver, cfg, tp_kernel.mailbox); + const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; + std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); + if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); + std.debug.print("unknown command {s}\n", .{args[1]}); + return 2; +} + +fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); + defer c.deinit(); + var backends: [2]struct { comm: tp.Collective, label: []const u8 } = undefined; + var n_backends: usize = 1; + backends[0] = .{ .comm = s.nccl.iface(), .label = "nccl" }; + if (s.comm().kind() != .nccl) { + backends[1] = .{ .comm = s.comm(), .label = if (s.roce != null) "roce" else "mailbox" }; + n_backends = 2; + } + const sizes = [_]usize{ 6, 1000, 4096, 65536, 256 * 1024, 256 * 1024 + 16, 1 << 20, 4 << 20 }; + for (backends[0..n_backends]) |b| { + for (sizes, 0..) |n, i| try c.collectives(b.comm, b.label, n, 100 * i); + try c.pointToPoint(b.comm, b.label, 64 * 1024, 7); + for ([_]usize{ 4096, 65536, 1 << 20 }) |n| try c.graphs(b.comm, b.label, n, 4); + try b.comm.barrier(); + } + for (backends[0..n_backends]) |b| { + try bench.sweep(&c, b.comm, b.label, 4096, 4 << 20); + try bench.decodeWindow(&c, b.comm, b.label, 86, 8192); + } + // the stop: rank 0 marks it expected first, then both meet and close + if (s.fate) |*f| f.expectStop(); + try s.comm().barrier(); + std.debug.print("SUMMARY rank {d}: {d} passed, {d} failed\n", .{ cfg.rank, c.passes, c.failures }); + s.close(); + return if (c.failures == 0) 0 else 1; +} + +fn realNow() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.REALTIME, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; +} + +/// Exchanges until the victim fails; the survivor is left waiting inside one, which only the fate channel ends. +fn failfast(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, mode: []const u8) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, 1 << 20); + defer c.deinit(); + const comm = s.comm(); + var k: u32 = 0; + while (true) : (k += 1) { + if (k == 50 and cfg.rank == victim) { + std.debug.print("INJECT {s} rank {d} at {d:.6}\n", .{ mode, cfg.rank, realNow() }); + if (std.mem.eql(u8, mode, "kill")) _ = std.c.raise(.KILL); + s.fatal("window", error.Injected); + return 3; // not reached: the process exits 70 + } + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + c.stream.synchronize() catch |e| { + std.debug.print("rank {d}: stream error {t} after the peer failed\n", .{ cfg.rank, e }); + }; + comm.check() catch |e| { + std.debug.print("rank {d}: transport reports {t} at exchange {d}; waiting for the fate channel\n", .{ cfg.rank, e, k }); + while (true) tp.sock.sleepNs(std.time.ns_per_s); + }; + } +} diff --git a/zig/src/fabric/tp/host.zig b/zig/src/fabric/tp/host.zig new file mode 100644 index 0000000000..051daea537 --- /dev/null +++ b/zig/src/fabric/tp/host.zig @@ -0,0 +1,293 @@ +//! In-process ranks over host memory: the `Collective` contract without a GPU (desyncs refused, abort wakes waiters). +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub const max_ranks = 8; + +const Call = enum(u8) { none, all_reduce, all_gather, all_gather_v, exchange, barrier }; + +const Post = struct { call: Call = .none, send: DevicePtr = 0, recv: DevicePtr = 0, count: usize = 0, t: DType = .u8, op: Op = .sum, peer: u32 = 0, lens: DevicePtr = 0 }; + +/// One-way message slot from rank `src` to rank `dst` (send / recv). +const Slot = struct { buf: std.atomic.Value(DevicePtr) = .init(0), bytes: usize = 0, taken: std.atomic.Value(bool) = .init(false) }; + +pub const Group = struct { + world: u32, + /// A wait longer than this is `error.Timeout` (a peer that never calls). + timeout_ns: u64 = 10 * std.time.ns_per_s, + posts: [max_ranks]Post = @splat(.{}), + arrived: std.atomic.Value(u32) = .init(0), + generation: std.atomic.Value(u32) = .init(0), + aborted: std.atomic.Value(bool) = .init(false), + /// Ranks (a bit each) whose transport refuses device lengths: their `all_gather_v` moves whole strides (tests). + var_refuse: u32 = 0, + slots: [max_ranks][max_ranks]Slot = @splat(@splat(.{})), + + pub fn init(world: u32) Group { + std.debug.assert(world >= 1 and world <= max_ranks); + return .{ .world = world }; + } + + pub fn rank(g: *Group, r: u32) Rank { + std.debug.assert(r < g.world); + return .{ .group = g, .me = r }; + } + + /// Sense-reversing barrier over every rank; `error.Aborted` once any rank aborts. + fn meet(g: *Group) Error!void { + const gen = g.generation.load(.acquire); + if (g.arrived.fetchAdd(1, .acq_rel) + 1 == g.world) { + g.arrived.store(0, .release); + _ = g.generation.fetchAdd(1, .acq_rel); + return if (g.aborted.load(.acquire)) error.Aborted else {}; + } + const t0 = now(); + while (g.generation.load(.acquire) == gen) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + if (g.aborted.load(.acquire)) return error.Aborted; + } +}; + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn bytesAt(p: DevicePtr, n: usize) []u8 { + if (n == 0) return &.{}; + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +pub const Rank = struct { + group: *Group, + me: u32, + /// `varAgreed`'s answer, once this communicator agreed. + var_agreed: ?bool = null, + + pub fn iface(self: *Rank) collective.Collective { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .host, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + + fn self_(ptr: *anyopaque) *Rank { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).group.world; + } + + /// Posts this rank's call, waits for all, and refuses calls that do not match rank 0's. + fn enter(self: *Rank, p: Post) Error!void { + const g = self.group; + if (g.aborted.load(.acquire)) return error.Aborted; + g.posts[self.me] = p; + try g.meet(); + const r0 = g.posts[0]; + if (r0.call != p.call or r0.count != p.count or r0.t != p.t or r0.op != p.op) { + abort(self); + return error.Invalid; + } + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_reduce, .send = s, .recv = r, .count = count, .t = t, .op = op }); + const n = count * t.size(); + var ins: [max_ranks][]const u8 = undefined; + for (0..g.world) |i| ins[i] = bytesAt(g.posts[i].send, n); + // the output may alias this rank's input: reduce into scratch first, then meet before writing + var scratch: [1 << 16]u8 = undefined; + if (n > scratch.len) return error.Unsupported; + if (n > 0) reference.reduce(scratch[0..n], ins[0..g.world], t, op); + try g.meet(); + @memcpy(bytesAt(r, n), scratch[0..n]); + try g.meet(); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather, .send = s, .recv = r, .count = count, .t = t }); + const n = count * t.size(); + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| @memcpy(out[i * n ..][0..n], bytesAt(g.posts[i].send, n)); + try g.meet(); + } + + fn lensAt(p: DevicePtr, world: u32) []const i32 { + return @as([*]const i32, @ptrFromInt(p))[0..world]; + } + + /// Copies only each rank's lens[r] bytes (`var_refuse` ranks take whole strides); unequal lengths are refused. + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather_v, .send = s, .recv = r, .count = stride, .t = t, .lens = lens }); + const n = stride * t.size(); + const mine = lensAt(lens, g.world); + for (0..g.world) |i| if (!std.mem.eql(i32, mine, lensAt(g.posts[i].lens, g.world))) { + abort(self); + return error.Invalid; + }; + const whole = g.var_refuse & (@as(u32, 1) << @intCast(self.me)) != 0; + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| { + if (mine[i] < 0 or mine[i] > n) { + abort(self); + return error.Invalid; + } + const len: usize = if (whole) n else @intCast(mine[i]); + @memcpy(out[i * n ..][0..len], bytesAt(g.posts[i].send, len)); + } + try g.meet(); + } + + fn varFits(ptr: *anyopaque, bytes: usize) bool { + const self = self_(ptr); + return bytes > 0 and self.group.var_refuse & (@as(u32, 1) << @intCast(self.me)) == 0; + } + + /// One int all-gather of `varFits(16)` the first time, then the kept answer. + fn varAgreed(ptr: *anyopaque) Error!bool { + const self = self_(ptr); + if (self.var_agreed) |v| return v; + const mine: i32 = @intFromBool(varFits(ptr, 16)); + var all: [max_ranks]i32 = undefined; + try allGather(ptr, @intFromPtr(&mine), @intFromPtr(&all), 1, .i32, null); + self.var_agreed = std.mem.min(i32, all[0..self.group.world]) == 1; + return self.var_agreed.?; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + try self.enter(.{ .call = .exchange, .send = s, .recv = r, .count = count, .t = t, .peer = peer }); + if (g.posts[peer].peer != self.me) { + abort(self); + return error.Invalid; + } + const n = count * t.size(); + @memcpy(bytesAt(r, n), bytesAt(g.posts[peer].send, n)); + try g.meet(); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[self.me][peer]; + slot.bytes = count * t.size(); + slot.taken.store(false, .release); + slot.buf.store(if (buf == 0) 1 else buf, .release); + try waitFor(g, &slot.taken, true); + slot.buf.store(0, .release); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[peer][self.me]; + const t0 = now(); + var src = slot.buf.load(.acquire); + while (src == 0) : (src = slot.buf.load(.acquire)) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + const n = count * t.size(); + if (slot.bytes != n) { + abort(self); + return error.Invalid; + } + @memcpy(bytesAt(buf, n), bytesAt(src, n)); + slot.taken.store(true, .release); + } + + fn waitFor(g: *Group, flag: *std.atomic.Value(bool), want: bool) Error!void { + const t0 = now(); + while (flag.load(.acquire) != want) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + } + + fn barrier(ptr: *anyopaque) Error!void { + const self = self_(ptr); + try self.enter(.{ .call = .barrier }); + try self.group.meet(); + } + + fn check(ptr: *anyopaque) Error!void { + if (self_(ptr).group.aborted.load(.acquire)) return error.Aborted; + } + + fn abort(ptr: *anyopaque) void { + self_(ptr).group.aborted.store(true, .release); + } +}; + +/// Runs `body(rank_collective, rank, ctx)` on `world` threads at once and returns the first error any rank hit. +pub fn run(world: u32, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + return runWith(world, .{}, ctx, body); +} + +pub const Options = struct { + /// `Group.var_refuse`: ranks whose transport takes no device lengths. + var_refuse: u32 = 0, +}; + +pub fn runWith(world: u32, opts: Options, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + var group = Group.init(world); + group.var_refuse = opts.var_refuse; + var ranks: [max_ranks]Rank = undefined; + var errs: [max_ranks]?anyerror = @splat(null); + var threads: [max_ranks]std.Thread = undefined; + const Wrap = struct { + fn go(r: *Rank, c: @TypeOf(ctx), e: *?anyerror) void { + body(r.iface(), r.me, c) catch |err| { + e.* = err; + r.group.aborted.store(true, .release); + }; + } + }; + for (0..world) |i| { + ranks[i] = group.rank(@intCast(i)); + threads[i] = try std.Thread.spawn(.{}, Wrap.go, .{ &ranks[i], ctx, &errs[i] }); + } + for (threads[0..world]) |t| t.join(); + for (errs[0..world]) |e| if (e) |err| return err; +} diff --git a/zig/src/fabric/tp/host_test.zig b/zig/src/fabric/tp/host_test.zig new file mode 100644 index 0000000000..7de32008af --- /dev/null +++ b/zig/src/fabric/tp/host_test.zig @@ -0,0 +1,154 @@ +//! The `Collective` contract on the in-process ranks: every call's bits against the reference, desyncs and aborts. +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const host = @import("host.zig"); +const Collective = collective.Collective; +const DType = collective.DType; + +fn addr(b: []u8) collective.DevicePtr { + return @intFromPtr(b.ptr); +} + +const n = 96; + +const Calls = struct { + fn go(c: Collective, rank: u32, _: void) !void { + const world = c.world(); + inline for (.{ DType.f32, DType.bf16, DType.i32, DType.f16 }) |t| { + const bytes = n * t.size(); + var mine: [n * 8]u8 = undefined; + var out: [n * 8 * host.max_ranks]u8 = undefined; + reference.fill(mine[0..bytes], t, rank, 7); + // all-reduce against the reference sum of every rank's input, in rank order + var ins: [host.max_ranks][n * 8]u8 = undefined; + var views: [host.max_ranks][]const u8 = undefined; + for (0..world) |r| { + reference.fill(ins[r][0..bytes], t, @intCast(r), 7); + views[r] = ins[r][0..bytes]; + } + var want: [n * 8]u8 = undefined; + reference.reduce(want[0..bytes], views[0..world], t, .sum); + try c.allReduce(addr(&mine), addr(&out), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], out[0..bytes]); + // in place + try c.allReduce(addr(&mine), addr(&mine), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], mine[0..bytes]); + // all-gather: rank order + reference.fill(mine[0..bytes], t, rank, 7); + try c.allGather(addr(&mine), addr(&out), n, t, null); + for (0..world) |r| try std.testing.expectEqualSlices(u8, views[r], out[r * bytes ..][0..bytes]); + } + try c.barrier(); + } +}; + +test "every call on two and three in-process ranks equals the reference" { + try host.run(2, {}, Calls.go); + try host.run(3, {}, Calls.go); +} + +const Pairwise = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var mine: [64]u8 = @splat(@intCast(rank + 1)); + var got: [64]u8 = undefined; + const peer = 1 - rank; + try c.exchange(addr(&mine), addr(&got), 64, .u8, peer, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(@intCast(peer + 1))), &got); + // send / recv, ordered so neither waits forever + if (rank == 0) { + try c.send(addr(&mine), 16, .f32, 1, null); + } else { + var buf: [64]u8 = undefined; + try c.recv(addr(&buf), 16, .f32, 0, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(1)), &buf); + } + try std.testing.expectError(error.Invalid, c.exchange(addr(&mine), addr(&got), 64, .u8, rank, null)); + } +}; + +test "exchange, send and recv between two ranks" { + try host.run(2, {}, Pairwise.go); +} + +const Desync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + // the ranks disagree on the count: both see it instead of hanging + const r = c.allGather(addr(&a), addr(&b), if (rank == 0) 16 else 32, .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "a desynchronised call is refused on every rank" { + try host.run(2, {}, Desync.go); +} + +const Aborts = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [16]u8 = @splat(0); + var b: [32]u8 = undefined; + if (rank == 1) { + // rank 1 never calls: it aborts after a while, which must free rank 0's waiting call + @import("sock.zig").sleepNs(50 * std.time.ns_per_ms); + c.abort(); + return; + } + try std.testing.expectError(error.Aborted, c.allGather(addr(&a), addr(&b), 16, .u8, null)); + try std.testing.expectError(error.Aborted, c.check()); + try std.testing.expectError(error.Aborted, c.barrier()); + } +}; + +test "abort frees a rank waiting in a collective" { + try host.run(2, {}, Aborts.go); +} + +const VarGather = struct { + fn go(c: Collective, rank: u32, refuse: u32) !void { + const world = c.world(); + const stride = 200; + const whole = refuse & (@as(u32, 1) << @intCast(rank)) != 0; + try std.testing.expectEqual(!whole, c.varFits(stride)); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); // kept: no second collective + var ins: [host.max_ranks][stride]u8 = undefined; + for (0..world) |r| reference.fill(&ins[r], .u8, @intCast(r), 11); + // lengths of 0, odd, 16 and the whole stride, in bytes; .f32 strides count elements + const cases = [_][3]i32{ .{ 0, 17, 200 }, .{ 200, 0, 3 }, .{ 16, 16, 16 }, .{ 199, 200, 1 } }; + inline for (.{ DType.u8, DType.f32 }) |t| for (cases) |lens| { + var out: [stride * host.max_ranks]u8 = @splat(0xEE); + try c.allGatherV(addr(&ins[rank]), addr(&out), stride / t.size(), @intFromPtr(&lens), t, null); + for (0..world) |r| { + const len: usize = @intCast(lens[r]); + try std.testing.expectEqualSlices(u8, ins[r][0..len], out[r * stride ..][0..len]); + // past a length: the transport's choice (untouched here, or the whole stride where lengths are refused) + const past = out[r * stride + len .. (r + 1) * stride]; + if (whole) try std.testing.expectEqualSlices(u8, ins[r][len..], past) else for (past) |x| try std.testing.expectEqual(@as(u8, 0xEE), x); + } + }; + try c.barrier(); + } +}; + +test "a variable-length all-gather moves each rank's length, and the ranks agree on device lengths once" { + for ([_]u32{ 2, 3 }) |w| { + try host.runWith(w, .{}, @as(u32, 0), VarGather.go); + try host.runWith(w, .{ .var_refuse = 2 }, @as(u32, 2), VarGather.go); + } +} + +const VarDesync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + const lens = [2]i32{ 8, if (rank == 0) 8 else 16 }; + const r = c.allGatherV(addr(&a), addr(&b), 64, @intFromPtr(&lens), .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "lengths that differ between ranks are refused on every rank" { + try host.run(2, {}, VarDesync.go); +} diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig new file mode 100644 index 0000000000..2cdeb0e4a4 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,167 @@ +//! Small exchanges through the mailbox (one host) or RoCE (two hosts), the rest through NCCL; ranks always agree. +const std = @import("std"); +const collective = @import("collective.zig"); +const Mailbox = @import("mailbox.zig").Mailbox; +const Roce = @import("roce.zig").Roce; +const Nccl = @import("nccl.zig").Nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +/// The one-kernel path for small shards: both expose the same calls. +pub const Small = union(enum) { + mailbox: *Mailbox, + roce: *Roce, + + fn fits(s: Small, n: usize) bool { + return switch (s) { + inline else => |p| p.fits(n), + }; + } + + fn sum(s: Small, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.sum(src, dst, count, t, stream), + }; + } + + fn gather(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.allGatherBytes(src, dst, n, stream), + }; + } + + fn exchange(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.exchangeBytes(src, dst, n, stream), + }; + } + + fn check(s: Small) Error!void { + return switch (s) { + inline else => |p| p.check(), + }; + } + + fn abort(s: Small) void { + switch (s) { + inline else => |p| p.abort(), + } + } +}; + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. + var_agreed: bool = false, + + pub fn iface(h: *Hybrid) collective.Collective { + return .{ .ptr = h, .vtable = if (h.small == .roce) &roce_vtable else &mailbox_vtable }; + } + + fn vtable(comptime kind: collective.Kind) collective.Collective.VTable { + return .{ + .kind = kind, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + } + + const mailbox_vtable = vtable(.hybrid); + const roce_vtable = vtable(.roce); + + fn self_(ptr: *anyopaque) *Hybrid { + return @ptrCast(@alignCast(ptr)); + } + + fn routed(h: *const Hybrid, n: usize) bool { + return n <= h.max_bytes and h.small.fits(n); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.size; + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (op == .sum and (t == .f32 or t == .bf16) and h.routed(count * t.size())) return h.small.sum(s, r, count, t, stream); + return Nccl.allReduce(h.nccl, s, r, count, t, op, stream); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (h.routed(count * t.size())) return h.small.gather(s, r, count * t.size(), stream); + return Nccl.allGather(h.nccl, s, r, count, t, stream); + } + + /// RoCE moves each rank's device length; the mailbox and NCCL move whole strides (the same bytes where lens says). + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + const n = stride * t.size(); + if (varFits(ptr, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + /// RoCE with its device side, for a stride the small path takes (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .mailbox => false, + }; + } + + fn varAgreed(ptr: *anyopaque) Error!bool { + return self_(ptr).var_agreed; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const h = self_(ptr); + try collective.checkPeer(h.nccl.me, h.nccl.size, peer); + if (h.routed(count * t.size())) return h.small.exchange(s, r, count * t.size(), stream); + return Nccl.exchange(h.nccl, s, r, count, t, peer, stream); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).send(buf, count, t, peer, stream); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).recv(buf, count, t, peer, stream); + } + + fn barrier(ptr: *anyopaque) Error!void { + return Nccl.iface(self_(ptr).nccl).barrier(); + } + + fn check(ptr: *anyopaque) Error!void { + const h = self_(ptr); + try h.small.check(); + try Nccl.iface(h.nccl).check(); + } + + fn abort(ptr: *anyopaque) void { + const h = self_(ptr); + h.small.abort(); + Nccl.iface(h.nccl).abort(); + } +}; diff --git a/zig/src/fabric/tp/log.zig b/zig/src/fabric/tp/log.zig new file mode 100644 index 0000000000..4e3234ab49 --- /dev/null +++ b/zig/src/fabric/tp/log.zig @@ -0,0 +1,15 @@ +//! TP's log lines, silent under test: a test's stderr is the build runner's channel. +const std = @import("std"); +const quiet = @import("builtin").is_test; + +pub fn err(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.err(format, args); +} + +pub fn warn(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.warn(format, args); +} + +pub fn info(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.info(format, args); +} diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu new file mode 100644 index 0000000000..63d53b2484 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,556 @@ +// One-shot two-rank exchange through pinned host memory both GPUs map (the TP mailbox; mailbox.zig is its host side). + +#include +#include + +// Protocol adapted from b12x's RoCEnante (Apache-2.0): a device epoch, slots by seq parity, a seq flag after the data. +namespace tftp { + +constexpr int kMaxBlocks = 16; +constexpr int kFlagStride = 32; // words: one 128-byte line per flag +// ctrl words: [0] abort (any rank, any host thread); per rank r at 16 + 8 r: failed, seq, seen, block, why, waited us +constexpr int kCtrlAbort = 0; +constexpr int kCtrlRank = 16; + +enum Mode : uint32_t { kGather = 0, kSumF32 = 1, kSumBf16 = 2, kExchange = 3 }; + +struct MailArgs { + const uint8_t *in; // this rank's shard (device) + uint8_t *out; // gather [2 n], exchange / sums [n] + uint8_t *peer_slots; // the peer's inbox data, as this GPU maps it + uint32_t *peer_flags; // the peer's inbox flags + const uint8_t *my_slots; // this rank's inbox data + const uint32_t *my_flags; // this rank's inbox flags + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison + uint32_t *ctrl; // shared control words + uint64_t nbytes; // bytes a rank + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; +}; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 ld_relaxed_sys_v4(const void *p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u8(const uint8_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u8 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_relaxed_sys(uint32_t *p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_sys(uint32_t *p, uint32_t v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_gpu(uint32_t *p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_gpu(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_sys(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.sys.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_cas_gpu(uint32_t *p, uint32_t cmp, uint32_t v) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.cas.b32 %0, [%1], %2, %3;" : "=r"(old) : "l"(p), "r"(cmp), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ void fence_sc_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } + +__device__ __forceinline__ uint64_t globaltimer() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// dst[i] <- src[i] for i in [0, n), grid-stride, 16-byte vectors when aligned; SysSrc: system-scope loads. +template +__device__ __forceinline__ void copy_bytes(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 v = SysSrc ? ld_relaxed_sys_v4(src + (i << 4)) : reinterpret_cast(src)[i]; + reinterpret_cast(dst)[i] = v; + } + } else { + for (uint64_t i = index; i < n; i += stride) { + dst[i] = static_cast(SysSrc ? ld_relaxed_sys_u8(src + i) : src[i]); + } + } +} + +// TF_TP_ROCE_FAST: copy_bytes with two grid strides' loads issued before their stores (same bytes, one round trip). +template +__device__ __forceinline__ void copy_bytes2(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) != 0u) { + copy_bytes(dst, src, n, index, stride); + return; + } + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += 2u * stride) { + const uint64_t j = i + stride; + uint4 v0, v1 = make_uint4(0u, 0u, 0u, 0u); + if (SysSrc) { + v0 = ld_relaxed_sys_v4(src + (i << 4)); + if (j < packs) v1 = ld_relaxed_sys_v4(src + (j << 4)); + } else { + v0 = reinterpret_cast(src)[i]; + if (j < packs) v1 = reinterpret_cast(src)[j]; + } + reinterpret_cast(dst)[i] = v0; + if (j < packs) reinterpret_cast(dst)[j] = v1; + } +} + +template +__device__ __forceinline__ void copy_any(bool fast, uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + if (fast) { + copy_bytes2(dst, src, n, index, stride); + } else { + copy_bytes(dst, src, n, index, stride); + } +} + +__device__ __forceinline__ uint32_t add_bf16x2(uint32_t x, uint32_t y) { + const float2 a = __bfloat1622float2(*reinterpret_cast(&x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&y)); + const __nv_bfloat162 r = __floats2bfloat162_rn(a.x + b.x, a.y + b.y); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t add_f32(uint32_t x, uint32_t y) { + return __float_as_uint(__uint_as_float(x) + __uint_as_float(y)); +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u16(const uint8_t *p) { + uint16_t v; + asm volatile("ld.relaxed.sys.global.u16 %0, [%1];" : "=h"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint16_t add_bf16(uint16_t x, uint16_t y) { + const __nv_bfloat16 r = __float2bfloat16_rn(__bfloat162float(__ushort_as_bfloat16(x)) + + __bfloat162float(__ushort_as_bfloat16(y))); + return __bfloat16_as_ushort(r); +} + +// out <- mine + peer's, rank 0's first; IEEE addition commutes, so both ranks get the same bits. +template +__device__ __forceinline__ void sum_bytes(uint8_t *out, const uint8_t *mine, const uint8_t *peer, uint64_t n, + uint32_t rank, uint64_t index, uint64_t stride) { + const uint64_t al = reinterpret_cast(out) | reinterpret_cast(mine) | + reinterpret_cast(peer) | n; + if ((al & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 m = reinterpret_cast(mine)[i]; + const uint4 p = ld_relaxed_sys_v4(peer + (i << 4)); + const uint4 a = rank == 0 ? m : p; + const uint4 b = rank == 0 ? p : m; + uint4 r; + r.x = Bf16 ? add_bf16x2(a.x, b.x) : add_f32(a.x, b.x); + r.y = Bf16 ? add_bf16x2(a.y, b.y) : add_f32(a.y, b.y); + r.z = Bf16 ? add_bf16x2(a.z, b.z) : add_f32(a.z, b.z); + r.w = Bf16 ? add_bf16x2(a.w, b.w) : add_f32(a.w, b.w); + reinterpret_cast(out)[i] = r; + } + } else if (Bf16) { + for (uint64_t i = index; i < (n >> 1); i += stride) { + const uint16_t m = reinterpret_cast(mine)[i]; + const uint16_t p = static_cast(ld_relaxed_sys_u16(peer + (i << 1))); + reinterpret_cast(out)[i] = rank == 0 ? add_bf16(m, p) : add_bf16(p, m); + } + } else { + for (uint64_t i = index; i < (n >> 2); i += stride) { + const uint32_t m = reinterpret_cast(mine)[i]; + const uint32_t p = ld_relaxed_sys(reinterpret_cast(peer) + i); + reinterpret_cast(out)[i] = rank == 0 ? add_f32(m, p) : add_f32(p, m); + } + } +} + +// Thread 0 of a block: spin until the peer block's flag holds seq; false on abort (why 1) or timeout (why 2). +__device__ __noinline__ bool wait_flag(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, uint64_t timeout_ns, + uint32_t *seen, uint32_t *why, uint32_t *waited_us) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *seen = v; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + return *seen == seq; + } + } + } +} + +// TF_TP_ROCE_FAST: wait_flag for one of several pollers of a flag; `done` (shared) ends the others once one saw it. +__device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, + uint64_t timeout_ns, uint32_t *seen, uint32_t *why, uint32_t *waited_us, + volatile uint32_t *done) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *done = 1u; + *seen = v; + return true; + } + if (*done != 0u) { + *seen = seq; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + if (*seen == seq) { + *done = 1u; + return true; + } + return *done != 0u; + } + } + } +} + +} // namespace tftp + +using namespace tftp; + +extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + // a poisoned transport does nothing (the host raises on its next check); poison and epoch in one round trip + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const uint64_t flag_at = (slot * kMaxBlocks + blockIdx.x) * kFlagStride; + + // 1. this block's share of the shard into the peer's inbox, then its flag after it (system-scope release) + copy_bytes(a.peer_slots + slot * a.slot_bytes, a.in, a.nbytes, index, stride); + __syncthreads(); + if (tid == 0) { + fence_sc_sys(); + st_release_sys(a.peer_flags + flag_at, seq); + } + // the local shard while the peer's writes travel + if (a.mode == kGather) { + copy_bytes(a.out + static_cast(a.rank) * a.nbytes, a.in, a.nbytes, index, stride); + } + + // 2. the peer's block of the same index (it wrote exactly the bytes this block reads) + if (tid == 0) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { // one failure record a rank + uint32_t *rec = a.ctrl + kCtrlRank + 8u * a.rank; + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, blockIdx.x); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } + } + } + __syncthreads(); + + // 3. copy out, gated by this block's own wait + if (s_bad == 0u) { + const uint8_t *inbox = a.my_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_bytes(a.out + static_cast(1u - a.rank) * a.nbytes, inbox, a.nbytes, index, + stride); + break; + case kExchange: + copy_bytes(a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the last block publishes the epoch, unless a wait of this launch failed (later launches stay no-ops) + __syncthreads(); + if (tid == 0) { + // the counter is back at 0 for the next launch (stream order), whatever its grid + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + } + } +} + +// The RoCE variant (roce.zig): blocks stage into the send slot, the last rings the doorbell, a host proxy RDMA-writes. +namespace tftp { + +constexpr int kRoceHcas = 2; +// ctrl words shared with roce.zig's proxy: [0] doorbell seq, [4 + slot] padded bytes, [8] completed, [16..21] failure +constexpr int kRoceSeq = 0; +constexpr int kRoceSlotBytes = 4; +constexpr int kRoceCompleted = 8; +constexpr int kRoceFailed = 16; +constexpr int kRoceAbort = 12; // wait_flag reads ctrl[kCtrlAbort]: the RoCE abort word is passed as that ctrl + +struct RoceArgs { + const uint8_t *in; + uint8_t *out; + uint8_t *send_slots; // this rank's send slots (pinned host, read by the NIC) + const uint8_t *recv_slots; // this rank's receive slots (written by the peer's NIC) + const uint32_t *flags; // this rank's flags [2][kRoceHcas], written by the peer's NIC + uint32_t *ctrl; + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison, [3] stage counter + uint64_t nbytes; + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; + uint32_t n_hca; + uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic + // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire + const int32_t *lens; +}; + +// Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. +__device__ __forceinline__ uint64_t roce_shard(const RoceArgs &a, uint32_t r) { + if (a.lens == nullptr) { + return a.nbytes; + } + const int32_t v = __ldg(a.lens + r); + const uint64_t n = v > 0 ? static_cast(v) : 0u; + return n < a.nbytes ? n : a.nbytes; +} + +// The bytes a copy of an n-byte shard moves: n rounded up to 16 when that stays inside a 16-aligned stride, else n. +__device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { + if (a.lens == nullptr || (a.nbytes & 15u) != 0u) { + return n; + } + return (n + 15u) & ~uint64_t(15u); +} + +} // namespace tftp + +extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + __shared__ uint32_t s_done[kRoceHcas]; // fast: some poller of HCA h saw its flag + const bool fast = (a.opts & 1u) != 0u; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + if (tid < static_cast(kRoceHcas)) { + s_done[tid] = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const bool var = a.mode == kGather && a.lens != nullptr; + const uint64_t mine = var ? roce_shard(a, a.rank) : a.nbytes; + const uint64_t theirs = var ? roce_shard(a, 1u - a.rank) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + + // 1. stage this block's share into the send slot; the last block to finish rings the doorbell + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in, var ? roce_span(a, mine) : a.nbytes, index, stride); + __syncthreads(); + if (tid == 0 && fast) { + // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) + const bool last = gdim == 1u || atom_add_acq_rel_sys(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } else if (tid == 0) { + fence_sc_sys(); + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + fence_sc_sys(); // every block's staging (seen through the counter) before the doorbell + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } + if (a.mode == kGather) { + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes, a.in, + var ? roce_span(a, mine) : a.nbytes, index, stride); + } + + // 2. every HCA's flag for this slot; fast: poller k of HCA h is thread 32 k + h, started k x stagger ns late + const uint32_t pollers = (a.opts >> 8) & 0xffu; + const uint32_t stagger_ns = a.opts >> 16; + if (fast) { + const uint32_t lane = tid & 31u; + const uint32_t k = tid >> 5; + if (lane < a.n_hca && k < (pollers == 0u ? 1u : pollers)) { + if (k > 0u && stagger_ns != 0u) { + __nanosleep(k * stagger_ns); + } + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, + a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, lane); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + } + } + } else if (tid == 0) { + for (uint32_t h = 0; h < a.n_hca; h++) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, + &why, &waited)) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, h); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + break; + } + } + } + __syncthreads(); + + // 3. copy out of the receive slot the peer's NIC wrote + if (s_bad == 0u) { + const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes, inbox, + var ? roce_span(a, theirs) : a.nbytes, index, stride); + break; + case kExchange: + copy_any(fast, a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the epoch, as tp_mailbox; the host also sees the last completed seq + __syncthreads(); + if (tid == 0) { + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + st_relaxed_sys(a.ctrl + kRoceCompleted, seq); + } + } +} diff --git a/zig/src/fabric/tp/mailbox.zig b/zig/src/fabric/tp/mailbox.zig new file mode 100644 index 0000000000..902bb79393 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.zig @@ -0,0 +1,215 @@ +//! The mailbox backend: two-rank exchanges in one kernel through pinned host memory both GPUs map (host side). +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const flag_bytes = 2 * max_blocks * 128; +const ctrl_bytes = 4096; +const ctrl_abort = 0; +const ctrl_rank = 16; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `MailArgs`, field for field. +const MailArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + peer_slots: DevicePtr, + peer_flags: DevicePtr, + my_slots: DevicePtr, + my_flags: DevicePtr, + state: DevicePtr, + ctrl: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, +}; + +/// Driver calls the runtime's table does not carry (page-locking memory this process mapped itself). +pub const HostRegister = struct { + lib: std.DynLib, + register: *const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + + pub fn open() !HostRegister { + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + return .{ + .register = lib.lookup(@FieldType(HostRegister, "register"), "cuMemHostRegister_v2") orelse return error.MissingSymbol, + .unregister = lib.lookup(@FieldType(HostRegister, "unregister"), "cuMemHostUnregister") orelse return error.MissingSymbol, + .lib = lib, + }; + } +}; + +/// Grid for `n` bytes a rank: one block per 16 KiB, a power of two, at most `max_blocks` (both ranks pick the same). +pub fn blocksFor(n: usize) u32 { + const want = std.math.divCeil(usize, @max(n, 1), 16 * 1024) catch 1; + return @intCast(@min(std.math.ceilPowerOfTwoAssert(usize, want), max_blocks)); +} + +pub const Mailbox = struct { + d: *const cuda.Driver, + me: u32, + slot_bytes: usize, + timeout_ns: u64, + region: []align(std.heap.page_size_min) u8, + dev: DevicePtr, + hr: HostRegister, + state: cuda.DeviceBuffer, + module: cuda.Module, + func: cuda.Function, + aborted: std.atomic.Value(bool) = .init(false), + + /// Both ranks call it after the bootstrap (two ranks, one host); `image` is mailbox.cu's fatbin. + pub fn init(d: *const cuda.Driver, b: *Bootstrap, image: []const u8, max_bytes: usize, timeout_ns: u64) !Mailbox { + if (b.world != 2) return error.Unsupported; + if (image.len == 0) return error.NoKernel; + const slot = std.mem.alignForward(usize, max_bytes, 4096); + const len = ctrl_bytes + 2 * (flag_bytes + 2 * slot); + // rank 0 creates and sizes the file before anyone learns its name, the others open it, then rank 0 unlinks it + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + const path = name[0..name_len :0]; + if (b.rank != 0) { + fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + try b.agree("mapped", "mailbox file"); + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + try b.agree("open", "mailbox mapping"); + if (b.rank == 0) _ = std.c.unlink(path); + var hr = try HostRegister.open(); + errdefer hr.lib.close(); + try d.check(hr.register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = hr.unregister(region.ptr); + var dev: DevicePtr = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + var state = try cuda.DeviceBuffer.alloc(d, 16); + errdefer state.free(); + try state.fill8(0, null); + var module = try cuda.Module.load(d, image); + errdefer module.unload(); + const func = try module.function("tp_mailbox"); + try b.agree("ready", "mailbox setup"); + return .{ .d = d, .me = b.rank, .slot_bytes = slot, .timeout_ns = timeout_ns, .region = region, .dev = dev, .hr = hr, .state = state, .module = module, .func = func }; + } + + pub fn deinit(m: *Mailbox) void { + m.module.unload(); + m.state.free(); + _ = m.hr.unregister(m.region.ptr); + m.hr.lib.close(); + posix.munmap(m.region); + } + + fn inbox(m: *const Mailbox, r: u32) usize { + return ctrl_bytes + r * (flag_bytes + 2 * m.slot_bytes); + } + + fn ctrlWord(m: *const Mailbox, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(m.region[4 * i ..].ptr)); + } + + /// Whether `n` bytes a rank fit one slot. + pub fn fits(m: *const Mailbox, n: usize) bool { + return n <= m.slot_bytes; + } + + fn launch(m: *Mailbox, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, stream: collective.Stream) Error!void { + if (m.aborted.load(.acquire)) return error.Aborted; + if (n > m.slot_bytes) return error.Unsupported; + if (n == 0) return; + const peer = 1 - m.me; + var a: MailArgs = .{ + .in = in, + .out = out, + .peer_slots = m.dev + m.inbox(peer) + flag_bytes, + .peer_flags = m.dev + m.inbox(peer), + .my_slots = m.dev + m.inbox(m.me) + flag_bytes, + .my_flags = m.dev + m.inbox(m.me), + .state = m.state.ptr, + .ctrl = m.dev, + .nbytes = n, + .slot_bytes = m.slot_bytes, + .timeout_ns = m.timeout_ns, + .rank = m.me, + .mode = @backingInt(mode), + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const r = m.d.api.cuLaunchKernel(m.func.handle, blocksFor(n), 1, 1, threads, 1, 1, 0, stream, ¶ms, null); + m.d.check(r, "cuLaunchKernel tp_mailbox") catch return error.BackendFailed; + } + + pub fn allGatherBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.gather, s, r, n, stream); + } + + pub fn exchangeBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.exchange, s, r, n, stream); + } + + /// Sums only, f32 or bf16 (what the model's partials are); the rest is `error.Unsupported`. + pub fn sum(m: *Mailbox, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return m.launch(mode, s, r, count * t.size(), stream); + } + + /// Ends pending and later exchanges on both ranks (the abort word is shared); this rank's launches stop. + pub fn abort(m: *Mailbox) void { + m.aborted.store(true, .release); + m.ctrlWord(ctrl_abort).store(1, .release); + } + + /// A wait that timed out or was aborted on either rank. + pub fn check(m: *Mailbox) Error!void { + for (0..2) |r| { + const rec = ctrl_rank + 8 * r; + if (m.ctrlWord(rec).load(.acquire) != 0) { + log.warn("tp mailbox: rank {d}'s wait for seq {d} failed ({s} after {d} us; block {d}, flag held {d})", .{ r, m.ctrlWord(rec + 1).load(.acquire), if (m.ctrlWord(rec + 4).load(.acquire) == 1) "aborted" else "timed out", m.ctrlWord(rec + 5).load(.acquire), m.ctrlWord(rec + 3).load(.acquire), m.ctrlWord(rec + 2).load(.acquire) }); + return error.PeerFailed; + } + } + if (m.aborted.load(.acquire) or m.ctrlWord(ctrl_abort).load(.acquire) != 0) return error.Aborted; + } +}; + +test "grids grow with the shard and stay powers of two" { + try std.testing.expectEqual(@as(u32, 1), blocksFor(4096)); + try std.testing.expectEqual(@as(u32, 1), blocksFor(16 * 1024)); + try std.testing.expectEqual(@as(u32, 2), blocksFor(16 * 1024 + 1)); + try std.testing.expectEqual(@as(u32, 4), blocksFor(64 * 1024)); + try std.testing.expectEqual(@as(u32, 16), blocksFor(4 << 20)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(MailArgs)); +} diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig new file mode 100644 index 0000000000..b641839aad --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,163 @@ +//! The NCCL backend: one communicator a rank, every collective one enqueue on the caller's stream. +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const nccl = cuda.nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub fn dataType(t: DType) nccl.DataType { + return switch (t) { + .u8 => .u8, + .i32 => .i32, + .i64 => .i64, + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + .f64 => .f64, + }; +} + +pub fn redOp(op: Op) nccl.RedOp { + return switch (op) { + .sum => .sum, + .max => .max, + .min => .min, + }; +} + +pub const Nccl = struct { + lib: *const nccl.Library, + d: *const cuda.Driver, + comm: nccl.Comm, + me: u32, + size: u32, + aborted: std.atomic.Value(bool) = .init(false), + /// The barrier's one-float buffers and its own stream (host path only). + scratch: cuda.DeviceBuffer, + stream: cuda.Stream, + + /// Joins the communicator `id` names (every rank calls it with the same id; it returns once all have). + pub fn init(lib: *const nccl.Library, d: *const cuda.Driver, id: nccl.UniqueId, rank: u32, world: u32) !Nccl { + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + errdefer _ = lib.api.ncclCommDestroy(comm); + var scratch = try cuda.DeviceBuffer.alloc(d, 4 * (world + 1)); + errdefer scratch.free(); + try scratch.fill8(0, null); + return .{ .lib = lib, .d = d, .comm = comm, .me = rank, .size = world, .scratch = scratch, .stream = try cuda.Stream.init(d, true) }; + } + + pub fn deinit(n: *Nccl) void { + _ = n.stream.synchronize() catch {}; + n.stream.deinit(); + n.scratch.free(); + if (!n.aborted.load(.acquire)) _ = n.lib.api.ncclCommDestroy(n.comm); + } + + pub fn iface(n: *Nccl) collective.Collective { + return .{ .ptr = n, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .nccl, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + }; + + fn self_(ptr: *anyopaque) *Nccl { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).size; + } + + fn call(n: *Nccl, r: nccl.Result, what: []const u8) Error!void { + if (r == 0) return; + if (n.aborted.load(.acquire)) return error.Aborted; + n.lib.check(r, what) catch return error.BackendFailed; + } + + fn live(n: *Nccl) Error!void { + if (n.aborted.load(.acquire)) return error.Aborted; + } + + pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllReduce(s, r, count, dataType(t), redOp(op), n.comm, stream), "ncclAllReduce"); + } + + pub fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(s, r, count, dataType(t), n.comm, stream), "ncclAllGather"); + } + + pub fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclGroupStart(), "ncclGroupStart"); + const a = n.lib.api.ncclSend(s, count, dataType(t), @intCast(peer), n.comm, stream); + const b = n.lib.api.ncclRecv(r, count, dataType(t), @intCast(peer), n.comm, stream); + const e = n.lib.api.ncclGroupEnd(); + try n.call(a, "ncclSend"); + try n.call(b, "ncclRecv"); + try n.call(e, "ncclGroupEnd"); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclSend(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclSend"); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclRecv(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclRecv"); + } + + /// One float all-gathered on the barrier's own stream, then a sync of it (host path, never captured). + fn barrier(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(n.scratch.ptr, n.scratch.ptr + 4, 1, .f32, n.comm, n.stream.handle), "ncclAllGather"); + n.stream.synchronize() catch return if (n.aborted.load(.acquire)) error.Aborted else error.BackendFailed; + } + + fn check(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + var r: nccl.Result = 0; + try n.call(n.lib.api.ncclCommGetAsyncError(n.comm, &r), "ncclCommGetAsyncError"); + if (r != 0 and r != nccl.in_progress) { + n.lib.check(r, "NCCL async error") catch {}; + return error.PeerFailed; + } + } + + fn abort(ptr: *anyopaque) void { + const n = self_(ptr); + if (n.aborted.swap(true, .acq_rel)) return; + _ = n.lib.api.ncclCommAbort(n.comm); + } +}; diff --git a/zig/src/fabric/tp/planlink.zig b/zig/src/fabric/tp/planlink.zig new file mode 100644 index 0000000000..44dbb61a0c --- /dev/null +++ b/zig/src/fabric/tp/planlink.zig @@ -0,0 +1,467 @@ +//! The round plan from rank 0 to the other ranks over its own TCP connection: a 4-byte length, then int64 words. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const builtin = @import("builtin"); +const config = @import("config.zig"); +const linux = std.os.linux; + +pub const max_ranks = 64; + +/// The CPU a plan thread pinned itself to (TF_TP_PLAN_PIN), -1 until one did; an engine's CPU isolation keeps off it. +pub var pinned_cpu: std.atomic.Value(i32) = .init(-1); +const more: u32 = 0x8000; // MSG_MORE: the length and the ints leave as one segment + +pub const Error = error{ PeerClosed, TooLarge, NotRankZero, NotFollower }; + +/// A batch frame's first word (a plan's op is >= 0): [batch, n, len 1, plan 1, ..., len n, plan n]. +pub const batch: i64 = std.math.minInt(i64); +/// Words rank 0 holds at most (more flushes the held plans first), and the most a batch frame carries. +pub const hold_words: usize = 1024; +pub const batch_words: usize = 2 * hold_words + 2; + +/// Largest plan either side accepts: 64M words (512 MiB), so a garbled length cannot become a 4 GiB allocation. +pub const max_frame_words: usize = 64 << 20; + +comptime { + std.debug.assert(builtin.cpu.arch.endian() == .little); // the ints go out as they are in memory +} + +pub const PlanLink = struct { + rank: u32, + world: u32, + /// Rank 0: rank i's socket at index i; the others: rank 0's at index 0 (the bootstrap's `releasePlan`). + fds: [max_ranks]sock.Fd, + sent: u64 = 0, + received: u64 = 0, + /// TF_TP_PLAN_LINK=rdma (config.plan_spin_ns): the follower polls the socket this long before it blocks. + spin_ns: u64 = 0, + /// TF_TP_PLAN_PIN: the sending / receiving thread pins itself at its first plan; `avoid`: the RoCE proxy's CPU. + pin: ?config.PlanPin = null, + avoid: ?u32 = null, + pinned: bool = false, + /// rank 0: plans held for the next send (`hold`), each as [len, plan...]; `held_n` of them + held: [hold_words]i64 = undefined, + held_len: usize = 0, + held_n: usize = 0, + /// the follower: a batch frame's plans not handed out yet, as [len, plan...] from `inbox_at` + inbox: [batch_words]i64 = undefined, + inbox_len: usize = 0, + inbox_at: usize = 0, + batches: u64 = 0, + + pub fn init(rank: u32, world: u32, fds: [max_ranks]sock.Fd) PlanLink { + return .{ .rank = rank, .world = world, .fds = fds }; + } + + /// Rank 0: one plan to every other rank (host only); held plans go first, in the same frame when they fit. + pub fn send(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (ints.len > max_frame_words) return error.TooLarge; + if (p.held_n > 0) { + if (2 + p.held_len + 1 + ints.len <= batch_words) { + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n + 1); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + frame[2 + p.held_len] = @intCast(ints.len); + @memcpy(frame[3 + p.held_len ..][0..ints.len], ints); + const n = 3 + p.held_len + ints.len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + try p.flush(); + } + return p.sendFrame(ints); + } + + /// Rank 0: a plan the followers run before the next sent one, carried in that plan's frame (`send` or `flush`). + pub fn hold(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (1 + ints.len > hold_words) { + try p.flush(); + return p.sendFrame(ints); + } + if (p.held_len + 1 + ints.len > hold_words) try p.flush(); + p.held[p.held_len] = @intCast(ints.len); + @memcpy(p.held[p.held_len + 1 ..][0..ints.len], ints); + p.held_len += 1 + ints.len; + p.held_n += 1; + } + + /// Rank 0: the held plans now, as one batch frame (nothing held: nothing sent). + pub fn flush(p: *PlanLink) Error!void { + if (p.held_n == 0) return; + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + const n = 2 + p.held_len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + + fn sendFrame(p: *PlanLink, ints: []const i64) Error!void { + p.pinOnce(); + const bytes = std.mem.sliceAsBytes(ints); + const len: u32 = @intCast(bytes.len); + for (p.fds[1..p.world]) |fd| { + if (std.c.send(fd, &len, 4, std.posix.MSG.NOSIGNAL | more) != 4) return error.PeerClosed; + sock.sendAll(fd, bytes) catch return error.PeerClosed; + } + p.sent += 1; + } + + /// The other ranks: the next plan into `out`; `error.PeerClosed` once rank 0 is gone, `error.TooLarge` skips one. + pub fn recv(p: *PlanLink, out: []i64) Error![]i64 { + if (p.inboxNext()) |m| { + if (m.len > out.len) return error.TooLarge; + @memcpy(out[0..m.len], m); + return out[0..m.len]; + } + const n = try p.header(); + if (n > out.len) { + try p.drop(n); + return error.TooLarge; + } + const got = try p.body(out[0..n]); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recv(out); + } + return got; + } + + /// The next plan into `buf`, grown to the frame; the buffer is kept, so steady plans allocate nothing. + pub fn recvGrow(p: *PlanLink, gpa: std.mem.Allocator, buf: *std.ArrayList(i64)) (Error || std.mem.Allocator.Error)![]i64 { + if (p.inboxNext()) |m| { + try buf.resize(gpa, m.len); + @memcpy(buf.items, m); + return buf.items; + } + const n = try p.header(); + if (n > max_frame_words) { + try p.drop(n); + return error.TooLarge; + } + try buf.resize(gpa, n); + const got = try p.body(buf.items); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recvGrow(gpa, buf); + } + return got; + } + + /// The follower: a batch frame's plans into the inbox (checked whole first: a garbled batch ends the link). + fn unpack(p: *PlanLink, frame: []const i64) Error!void { + if (frame.len < 2 or frame.len > batch_words or frame[1] < 1) return error.PeerClosed; + var at: usize = 2; + var k: i64 = 0; + while (k < frame[1]) : (k += 1) { + if (at >= frame.len or frame[at] < 0) return error.PeerClosed; + at += 1 + @as(usize, @intCast(frame[at])); + if (at > frame.len) return error.PeerClosed; + } + if (at != frame.len) return error.PeerClosed; + @memcpy(p.inbox[0 .. frame.len - 2], frame[2..]); + p.inbox_len = frame.len - 2; + p.inbox_at = 0; + p.batches += 1; + } + + /// The follower: the next plan of the last batch frame, if any (a slice of the inbox, valid until the next recv). + fn inboxNext(p: *PlanLink) ?[]const i64 { + if (p.inbox_at >= p.inbox_len) return null; + const n: usize = @intCast(p.inbox[p.inbox_at]); + const m = p.inbox[p.inbox_at + 1 ..][0..n]; + p.inbox_at += 1 + n; + return m; + } + + /// The next frame's length in words. + fn header(p: *PlanLink) Error!usize { + if (p.rank == 0) return error.NotFollower; + p.pinOnce(); + if (p.spin_ns > 0) p.spin(); + var len: u32 = undefined; + sock.recvExact(p.fds[0], std.mem.asBytes(&len), null) catch return error.PeerClosed; + if (len % 8 != 0) return error.PeerClosed; // not a frame boundary: the stream is lost + return len / 8; + } + + fn body(p: *PlanLink, out: []i64) Error![]i64 { + sock.recvExact(p.fds[0], std.mem.sliceAsBytes(out), null) catch return error.PeerClosed; + p.received += 1; + return out; + } + + /// Reads and drops `n` words. + fn drop(p: *PlanLink, n: usize) Error!void { + var left: usize = n * 8; + var scratch: [4096]u8 = undefined; + while (left > 0) : (left -= @min(left, scratch.len)) sock.recvExact(p.fds[0], scratch[0..@min(left, scratch.len)], null) catch return error.PeerClosed; + } + + /// Polls rank 0's socket until it is readable (or hung up) or `spin_ns` passed; the read that follows blocks. + fn spin(p: *PlanLink) void { + if (comptime builtin.os.tag != .linux) return; + var fd = [1]linux.pollfd{.{ .fd = p.fds[0], .events = linux.POLL.IN, .revents = 0 }}; + const end = sock.nowNs() + p.spin_ns; + while (true) { + const rc = linux.poll(&fd, 1, 0); + if (linux.errno(rc) == .SUCCESS and rc > 0) return; + if (sock.nowNs() >= end) return; + std.atomic.spinLoopHint(); + } + } + + /// Pins the calling thread once: an explicit CPU, or auto's highest cpu_capacity. + fn pinOnce(p: *PlanLink) void { + if (p.pinned) return; + p.pinned = true; + if (comptime builtin.os.tag != .linux) return; + const want = p.pin orelse return; + const cpu = switch (want) { + .cpu => |c| c, + .auto => bestCpu(p.avoid) orelse { + log.info("[tensorfold] tp: TF_TP_PLAN_PIN=auto: no cpu_capacity differs among the allowed CPUs; the plan thread stays unpinned", .{}); + return; + }, + }; + pinTo(cpu) catch |e| { + log.warn("[tensorfold] tp: TF_TP_PLAN_PIN: CPU {d} refused ({t}); the plan thread stays unpinned", .{ cpu, e }); + return; + }; + pinned_cpu.store(@intCast(cpu), .release); + log.info("[tensorfold] tp: plan link rank {d} pinned to CPU {d}", .{ p.rank, cpu }); + } + + pub fn close(p: *PlanLink) void { + for (&p.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } +}; + +/// The calling thread onto `cpu` alone. +pub fn pinTo(cpu: u32) !void { + if (cpu >= linux.CPU_SETSIZE) return error.BadCpu; + var set: linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + try linux.sched_setaffinity(0, &set); +} + +/// Auto pin: the allowed CPU (not `avoid`) with the highest cpu_capacity, lowest on ties; null when all are equal. +pub fn bestCpu(avoid: ?u32) ?u32 { + var set: linux.cpu_set_t = undefined; + if (linux.errno(linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set)) != .SUCCESS) return null; + var best: ?u32 = null; + var best_cap: u32 = 0; + var low_cap: u32 = std.math.maxInt(u32); + for (set, 0..) |word, wi| { + var bits = word; + while (bits != 0) : (bits &= bits - 1) { + const cpu: u32 = @intCast(wi * @bitSizeOf(usize) + @ctz(bits)); + if (avoid == cpu) continue; + const c = capacity(cpu) orelse continue; + low_cap = @min(low_cap, c); + if (best == null or c > best_cap) { + best = cpu; + best_cap = c; + } + } + } + return if (best != null and best_cap > low_cap) best else null; +} + +fn capacity(cpu: u32) ?u32 { + var pb: [64]u8 = undefined; + const path = std.fmt.bufPrintSentinel(&pb, "/sys/devices/system/cpu/cpu{d}/cpu_capacity", .{cpu}, 0) catch return null; + const rc = linux.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return null; + const fd: i32 = @intCast(rc); + defer _ = linux.close(fd); + var b: [16]u8 = undefined; + const n = linux.read(fd, &b, b.len); + if (linux.errno(n) != .SUCCESS) return null; + return std.fmt.parseInt(u32, std.mem.trim(u8, b[0..n], " \n"), 10) catch null; +} + +test "plans arrive whole and in order; a closed rank 0 ends the follower" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @as(i64, @intCast(i)) * -3 + 0x7FFF_0000_0000; + try lead.send(&.{ 0x504C414E, 2, 1 }); + try lead.send(&big); + try lead.send(&.{}); + var buf: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 2, 1 }, try follow.recv(&buf)); + try std.testing.expectEqualSlices(i64, &big, try follow.recv(&buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&buf)).len); + try std.testing.expectError(error.NotRankZero, follow.send(&.{1})); + var small: [2]i64 = undefined; + try lead.send(&.{ 1, 2, 3 }); + try std.testing.expectError(error.TooLarge, follow.recv(&small)); + lead.close(); + try std.testing.expectError(error.PeerClosed, follow.recv(&buf)); +} + +test "an admission-sized plan (1.2M words) arrives whole through the growable path" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + const gpa = std.testing.allocator; + const prompt = try gpa.alloc(i64, 1_200_000); + defer gpa.free(prompt); + for (prompt, 0..) |*v, i| v.* = @intCast((i *% 2654435761) % 163840); + // the socket buffer holds far less than 9.6 MB: rank 0 sends from its own thread, as it would + const Send = struct { + fn go(link: *PlanLink, ints: []const i64) void { + link.send(&.{ 0x504C414E, 3, 20 }) catch {}; + link.send(ints) catch {}; + link.send(&.{ 0x504C414E, 3, 21 }) catch {}; + } + }; + const t = try std.Thread.spawn(.{}, Send.go, .{ &lead, prompt }); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(gpa); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 20 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, prompt, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 21 }, try follow.recvGrow(gpa, &buf)); + t.join(); + try std.testing.expect(buf.capacity >= prompt.len); // kept for the next admission +} + +test "TF_TP_PLAN_LINK=rdma: a spinning follower gets every plan, late or early; TF_TP_PLAN_PIN pins the thread" { + if (builtin.os.tag != .linux) return error.SkipZigTest; + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + follow.spin_ns = 2 * std.time.ns_per_ms; + // the CPU this thread may run on first: the pin keeps it allowed, so the test runner is not starved + var set: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set); + var first: u32 = 0; + for (set, 0..) |w, i| if (w != 0) { + first = @intCast(i * @bitSizeOf(usize) + @ctz(w)); + break; + }; + const Follower = struct { + fn run(f: *PlanLink, cpu: u32, out: *[3]i64, ok: *bool) void { + f.pin = .{ .cpu = cpu }; + var buf: [8]i64 = undefined; + const x = (f.recv(&buf) catch return)[0]; // arrives within the spin + const y = (f.recv(&buf) catch return)[0]; // arrives after it: the blocking read + var now: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &now); + var count: usize = 0; + for (now) |w| count += @popCount(w); + out.* = .{ x, y, @intCast(count) }; + ok.* = f.pinned and (now[cpu / @bitSizeOf(usize)] >> @intCast(cpu % @bitSizeOf(usize))) & 1 == 1; + } + }; + var out: [3]i64 = .{ 0, 0, 0 }; + var ok = false; + const t = try std.Thread.spawn(.{}, Follower.run, .{ &follow, first, &out, &ok }); + try lead.send(&.{11}); + var ts: linux.timespec = .{ .sec = 0, .nsec = 20 * std.time.ns_per_ms }; + _ = linux.nanosleep(&ts, null); + try lead.send(&.{22}); + t.join(); + try std.testing.expectEqual([3]i64{ 11, 22, 1 }, out); + try std.testing.expect(ok); + if (bestCpu(null)) |c| try std.testing.expect((set[c / @bitSizeOf(usize)] >> @intCast(c % @bitSizeOf(usize))) & 1 == 1); +} + +test "held plans leave with the next plan, in order, as one frame; a large plan after them flushes them first" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("127.0.0.1", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + const lead = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(lead); + lead.* = PlanLink.init(0, 2, f0); + const follow = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(follow); + follow.* = PlanLink.init(1, 2, f1); + defer follow.close(); + // a round: three keeps held, the window sent: one frame + try lead.hold(&.{ 22, 0, 3 }); + try lead.hold(&.{ 22, 1, 0 }); + try lead.hold(&.{ 22, 2, 5 }); + try lead.send(&.{ 21, 4, 9, 9, 9 }); + try std.testing.expectEqual(@as(u64, 1), lead.sent); + // a held plan, then one too large for a batch: two frames, the held one first + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @intCast(i); + try lead.hold(&.{ 22, 3, 1 }); + try lead.send(&big); + try std.testing.expectEqual(@as(u64, 3), lead.sent); + // flush alone, an empty plan held, a plain send after + try lead.hold(&.{}); + try lead.flush(); + try lead.flush(); // nothing held: nothing sent + try lead.send(&.{7}); + try std.testing.expectEqual(@as(u64, 5), lead.sent); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(std.testing.allocator); + const gpa = std.testing.allocator; + try std.testing.expectEqualSlices(i64, &.{ 22, 0, 3 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 1, 0 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 2, 5 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 21, 4, 9, 9, 9 }, try follow.recvGrow(gpa, &buf)); + var out: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 22, 3, 1 }, try follow.recv(&out)); + try std.testing.expectEqualSlices(i64, &big, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&out)).len); + try std.testing.expectEqualSlices(i64, &.{7}, try follow.recv(&out)); + try std.testing.expectEqual(@as(u64, 3), follow.batches); + try std.testing.expectError(error.NotRankZero, follow.hold(&.{1})); + lead.close(); +} diff --git a/zig/src/fabric/tp/reference.zig b/zig/src/fabric/tp/reference.zig new file mode 100644 index 0000000000..b98966fcc8 --- /dev/null +++ b/zig/src/fabric/tp/reference.zig @@ -0,0 +1,121 @@ +//! Host references for the collectives: reductions in rank order, bf16 / f16 added in fp32 and rounded once. +const std = @import("std"); +const collective = @import("collective.zig"); +const DType = collective.DType; +const Op = collective.Op; + +/// bf16 bits to fp32 (exact). +pub fn bf16ToF32(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 to bf16 bits, round to nearest even; NaN stays a quiet NaN. +pub fn f32ToBf16(f: f32) u16 { + const u: u32 = @bitCast(f); + if (std.math.isNan(f)) return @intCast((u >> 16) | 0x40); + const round = 0x7FFF + ((u >> 16) & 1); + return @intCast((u +% round) >> 16); +} + +fn combine(comptime T: type, a: T, b: T, op: Op) T { + return switch (op) { + .sum => if (@typeInfo(T) == .int) a +% b else a + b, + .max => @max(a, b), + .min => @min(a, b), + }; +} + +fn reduceTyped(comptime T: type, acc: []u8, src: []const u8, op: Op) void { + const n = acc.len / @sizeOf(T); + for (0..n) |i| { + const a = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const b = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), src[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const r: T = combine(T, @bitCast(a), @bitCast(b), op); + std.mem.writeInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], @bitCast(r), .little); + } +} + +fn reduceHalf(acc: []u8, src: []const u8, op: Op, bf: bool) void { + const n = acc.len / 2; + for (0..n) |i| { + const a = std.mem.readInt(u16, acc[2 * i ..][0..2], .little); + const b = std.mem.readInt(u16, src[2 * i ..][0..2], .little); + const fa: f32 = if (bf) bf16ToF32(a) else @floatCast(@as(f16, @bitCast(a))); + const fb: f32 = if (bf) bf16ToF32(b) else @floatCast(@as(f16, @bitCast(b))); + const r = combine(f32, fa, fb, op); + const out: u16 = if (bf) f32ToBf16(r) else @bitCast(@as(f16, @floatCast(r))); + std.mem.writeInt(u16, acc[2 * i ..][0..2], out, .little); + } +} + +/// acc <- acc (op) src, element-wise; `acc` holds the lower ranks' result, `src` the next rank's input. +pub fn accumulate(acc: []u8, src: []const u8, t: DType, op: Op) void { + std.debug.assert(acc.len == src.len and acc.len % t.size() == 0); + switch (t) { + .u8 => reduceTyped(u8, acc, src, op), + .i32 => reduceTyped(i32, acc, src, op), + .i64 => reduceTyped(i64, acc, src, op), + .f32 => reduceTyped(f32, acc, src, op), + .f64 => reduceTyped(f64, acc, src, op), + .bf16 => reduceHalf(acc, src, op, true), + .f16 => reduceHalf(acc, src, op, false), + } +} + +/// out <- inputs[0] (op) inputs[1] (op) ..., in rank order. +pub fn reduce(out: []u8, inputs: []const []const u8, t: DType, op: Op) void { + @memcpy(out, inputs[0]); + for (inputs[1..]) |in| accumulate(out, in, t, op); +} + +/// Deterministic test data for `rank`'s input of `bytes`: finite values of type `t` that differ per rank and position. +pub fn fill(out: []u8, t: DType, rank: u32, seed: u64) void { + var rng = std.Random.DefaultPrng.init(seed ^ (@as(u64, rank) *% 0x9E3779B97F4A7C15)); + const r = rng.random(); + const n = out.len / t.size(); + for (0..n) |i| { + const x = r.float(f32) * 8.0 - 4.0; + switch (t) { + .u8 => out[i] = r.int(u8), + .i32 => std.mem.writeInt(i32, out[4 * i ..][0..4], r.int(i32), .little), + .i64 => std.mem.writeInt(i64, out[8 * i ..][0..8], r.int(i64), .little), + .f32 => std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x), .little), + .f64 => std.mem.writeInt(u64, out[8 * i ..][0..8], @bitCast(@as(f64, x) * 1.000000123), .little), + .bf16 => std.mem.writeInt(u16, out[2 * i ..][0..2], f32ToBf16(x), .little), + .f16 => std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(x))), .little), + } + } +} + +test "bf16 rounding is nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.0)); + // 1 + 2^-8 is exactly halfway between bf16 1.0 and 1.0078125: rounds to the even one (1.0) + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.00390625)); + try std.testing.expectEqual(@as(u16, 0x3F82), f32ToBf16(1.01171875)); + try std.testing.expect(std.math.isNan(bf16ToF32(f32ToBf16(std.math.nan(f32))))); +} + +test "rank-order reduction of bf16 adds in fp32 and rounds once" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(u16, a[0..2], f32ToBf16(1.0), .little); + std.mem.writeInt(u16, b[0..2], f32ToBf16(0.00390625), .little); + std.mem.writeInt(u16, a[2..4], f32ToBf16(-2.5), .little); + std.mem.writeInt(u16, b[2..4], f32ToBf16(3.0), .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .bf16, .sum); + try std.testing.expectEqual(f32ToBf16(1.0), std.mem.readInt(u16, out[0..2], .little)); + try std.testing.expectEqual(f32ToBf16(0.5), std.mem.readInt(u16, out[2..4], .little)); + reduce(&out, &.{ &a, &b }, .bf16, .max); + try std.testing.expectEqual(f32ToBf16(3.0), std.mem.readInt(u16, out[2..4], .little)); +} + +test "integer sums wrap" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(i32, &a, std.math.maxInt(i32), .little); + std.mem.writeInt(i32, &b, 1, .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .i32, .sum); + try std.testing.expectEqual(@as(i32, std.math.minInt(i32)), std.mem.readInt(i32, &out, .little)); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig new file mode 100644 index 0000000000..58f88e3441 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,458 @@ +//! The RoCE backend between two hosts: the tp_roce kernel stages and rings a doorbell, a proxy thread moves the slot. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const rv = @import("roce_verbs.zig"); +const HostRegister = @import("mailbox.zig").HostRegister; +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const page = std.heap.page_size_min; +const flag_stride = 128; +// region: ctrl [4 KiB], flags [2 slots][2 HCAs] at 128 B, send [2][slot], recv [2][slot] (mailbox.cu's tp_roce) +const ctrl_off = 0; +const flag_off = 4096; +const send_off = 8192; +// ctrl words (mailbox.cu): doorbell, padded bytes per slot, completed, abort, proxy failure, a failed wait's record +const w_seq = 0; +const w_slot_bytes = 4; +const w_completed = 8; +const w_abort = 12; +const w_proxy_failed = 13; +const w_failed = 16; +/// Model graphs leave sub-millisecond gaps: the proxy stays hot this many idle polls, then naps 20 us at a time. +const idle_spins: u64 = 20_000_000; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `RoceArgs`, field for field. +const RoceArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + send_slots: DevicePtr, + recv_slots: DevicePtr, + flags: DevicePtr, + ctrl: DevicePtr, + state: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, + n_hca: u32, + /// TF_TP_ROCE_FAST (config.roce_fast; 0: the classic path) + opts: u32 = 0, + /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) + lens: DevicePtr = 0, +}; + +pub const WireKind = enum { + /// RDMA writes over the HCAs. + verbs, + /// The proxy copies into the peer's region in shared memory: the same protocol on one host without RDMA (tests). + shm, +}; + +pub const Options = struct { + wire: WireKind = .verbs, + hcas: []const rv.HcaSpec = &.{}, + traffic_class: u8 = 0, + cpu: ?u32 = null, + max_bytes: usize = 256 * 1024, + timeout_ns: u64 = 120 * std.time.ns_per_s, + /// TF_TP_ROCE_FAST: tp_roce's `opts` word (config.roce_fast) + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) + blocks: u32 = 16, +}; + +/// The shm wire: the peer's region is mapped here; a post copies the slot, then stores the flag with release order. +const ShmWire = struct { + peer: []u8, + local: []u8, + slot_bytes: usize, + + fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) void { + const recv = send_off + 2 * w.slot_bytes; + @memcpy(w.peer[recv + slot * w.slot_bytes ..][0..padded], w.local[send_off + slot * w.slot_bytes ..][0..padded]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[flag_off + slot * 2 * flag_stride ..].ptr)); + flag.store(seq, .release); + } +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + shm: ShmWire, + + fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { + switch (w.*) { + .verbs => |*v| try v.post(slot, seq, padded), + .shm => |*s| s.post(slot, seq, padded), + } + } + + fn reap(w: *Wire) !void { + switch (w.*) { + .verbs => |*v| try v.reap(), + .shm => {}, + } + } + + fn hcas(w: *const Wire) u32 { + return switch (w.*) { + .verbs => |v| v.n_hca, + .shm => 1, + }; + } +}; + +pub const Roce = struct { + me: u32, + slot_bytes: usize, + timeout_ns: u64, + /// TF_TP_ROCE_FAST: the kernel's `opts` word + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid + max_grid: u32 = 16, + /// This rank's region (for shm: half of the mapping in `mapping`). + region: []align(page) u8, + mapping: []align(page) u8, + wire: Wire, + cpu: ?u32, + // device side (absent in host-only use: tests and `tf-tp-test roce-host`) + d: ?*const cuda.Driver = null, + hr: ?HostRegister = null, + dev: DevicePtr = 0, + state: ?cuda.DeviceBuffer = null, + module: ?cuda.Module = null, + func: ?cuda.Function = null, + // the proxy thread + thread: ?std.Thread = null, + running: std.atomic.Value(bool) = .init(false), + last_seq: u32 = 0, + err_buf: [160]u8 = undefined, + err_len: usize = 0, + // host emulation of the kernel's side (tests) + host_epoch: u32 = 0, + aborted: std.atomic.Value(bool) = .init(false), + + pub fn regionBytes(slot: usize) usize { + return send_off + 4 * slot; + } + + /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). + pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { + if (b.world != 2) return error.Unsupported; + const slot = std.mem.alignForward(usize, o.max_bytes, 4096); + const len = regionBytes(slot); + const r = try std.heap.page_allocator.create(Roce); + errdefer std.heap.page_allocator.destroy(r); + r.* = .{ .me = b.rank, .slot_bytes = slot, .timeout_ns = o.timeout_ns, .fast = o.fast, .max_grid = o.blocks, .region = undefined, .mapping = undefined, .wire = undefined, .cpu = o.cpu }; + // 1. the regions and the wire, each rank alone; every rank learns whether all did + const ok = r.openWire(b, o, len) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, ok, "RoCE setup")) { + if (ok) r.closeWire(); + return error.RoceUnavailable; + } + errdefer r.closeWire(); + // 2. the verbs connection records, through the bootstrap + const connected = r.connectWire(b) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} connect failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, connected, "RoCE connect")) return error.RoceUnavailable; + // 3. the device side + if (d) |drv| { + const dev_ok = r.openDevice(drv, image) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} device setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, dev_ok, "RoCE device setup")) { + if (dev_ok) r.closeDevice(); + return error.RoceUnavailable; + } + } + r.running.store(true, .release); + r.thread = try std.Thread.spawn(.{}, proxyLoop, .{r}); + try b.agree("proxy", "RoCE proxy start"); + return r; + } + + fn agreeAll(b: *Bootstrap, ok: bool, _: []const u8) !bool { + var blobs: [2][]const u8 = undefined; + var buf: [2]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(ok)}, &buf, &blobs); + return all[0][0] == 1 and all[1][0] == 1; + } + + fn openWire(r: *Roce, b: *Bootstrap, o: Options, len: usize) !bool { + switch (o.wire) { + .verbs => { + r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + r.region = r.mapping; + errdefer posix.munmap(r.mapping); + r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; + }, + .shm => { + r.mapping = try sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .local = r.region, .slot_bytes = r.slot_bytes } }; + }, + } + return true; + } + + fn connectWire(r: *Roce, b: *Bootstrap) !bool { + switch (r.wire) { + .shm => return true, + .verbs => |*v| { + const mine = v.record(); + var blobs: [2][]const u8 = undefined; + var buf: [2 * @sizeOf(rv.Record)]u8 = undefined; + const all = try b.allGather(std.mem.asBytes(&mine), &buf, &blobs); + var peer: rv.Record = undefined; + if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; + @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); + try v.connect(peer); + return true; + }, + } + } + + fn openDevice(r: *Roce, d: *const cuda.Driver, image: []const u8) !bool { + if (image.len == 0) return error.NoKernel; + r.d = d; + r.hr = try HostRegister.open(); + try d.check(r.hr.?.register(r.region.ptr, r.region.len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + try d.check(d.api.cuMemHostGetDevicePointer_v2(&r.dev, r.region.ptr, 0), "cuMemHostGetDevicePointer"); + r.state = try cuda.DeviceBuffer.alloc(d, 16); + try r.state.?.fill8(0, null); + r.module = try cuda.Module.load(d, image); + r.func = try r.module.?.function("tp_roce"); + return true; + } + + fn closeDevice(r: *Roce) void { + if (r.module) |*m| m.unload(); + if (r.state) |*s| s.free(); + if (r.hr) |*h| { + _ = h.unregister(r.region.ptr); + h.lib.close(); + } + r.module = null; + r.state = null; + r.hr = null; + } + + fn closeWire(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| v.close(), + .shm => {}, + } + posix.munmap(r.mapping); + } + + pub fn deinit(r: *Roce) void { + r.running.store(false, .release); + if (r.thread) |t| t.join(); + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); + } + + fn fail(r: *Roce, comptime fmt: []const u8, args: anytype) void { + const s = std.fmt.bufPrint(&r.err_buf, fmt, args) catch "proxy failure"; + r.err_len = s.len; + log.err("[tensorfold] tp roce proxy: {s}", .{s}); + r.word(w_proxy_failed).store(1, .release); + r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + } + + /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. + fn proxyLoop(r: *Roce) void { + if (r.cpu) |cpu| pin(cpu); + var idle: u64 = 0; + while (r.running.load(.monotonic)) { + if (r.word(w_abort).load(.monotonic) != 0) return; + const seq = r.word(w_seq).load(.acquire); + if (seq == r.last_seq) { + idle += 1; + if (idle % 64 == 0) r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + if (idle >= idle_spins) sock.sleepNs(20 * std.time.ns_per_us) else std.atomic.spinLoopHint(); + continue; + } + idle = 0; + var pending = seq -% r.last_seq; + if (pending > 2) return r.fail("doorbell skipped {d} ops (posted {d}, now {d})", .{ pending, r.last_seq, seq }); + var s = r.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const padded = r.word(w_slot_bytes + (s & 1)).load(.acquire); + if (padded == 0 or padded % 16 != 0 or padded > r.slot_bytes) return r.fail("op {d} of {d} bytes", .{ s, padded }); + r.wire.post(s & 1, s, padded) catch |e| return r.fail("post of op {d}: {t}", .{ s, e }); + r.last_seq = s; + } + r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + } + } + + pub fn fits(r: *const Roce, n: usize) bool { + return n <= r.slot_bytes; + } + + fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes) return error.Unsupported; + if (n == 0) return; + const d = r.d orelse return error.Unsupported; + var a: RoceArgs = .{ + .in = in, + .out = out, + .send_slots = r.dev + send_off, + .recv_slots = r.dev + send_off + 2 * r.slot_bytes, + .flags = r.dev + flag_off, + .ctrl = r.dev + ctrl_off, + .state = r.state.?.ptr, + .nbytes = n, + .slot_bytes = r.slot_bytes, + .timeout_ns = r.timeout_ns, + .rank = r.me, + .mode = @backingInt(mode), + .n_hca = r.wire.hcas(), + .opts = r.fast, + .lens = lens, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(n), r.max_grid); + d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; + } + + pub fn allGatherBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.gather, s, o, n, 0, stream); + } + + /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). + pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, stream); + } + + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.exchange, s, o, n, 0, stream); + } + + pub fn sum(r: *Roce, s: DevicePtr, o: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return r.launch(mode, s, o, count * t.size(), 0, stream); + } + + /// Ends this rank's waits and the proxy (the peer learns through the fate channel and aborts its own). + pub fn abort(r: *Roce) void { + r.aborted.store(true, .release); + r.word(w_abort).store(1, .release); + } + + pub fn check(r: *Roce) Error!void { + if (r.word(w_failed).load(.acquire) != 0) { + log.warn("tp roce: rank {d}'s wait for seq {d} failed ({s} after {d} us; HCA {d}, flag held {d}); doorbell {d}, posted {d}, completed {d}, abort {d}, proxy failed {d}, flags {d} {d}", .{ r.me, r.word(w_failed + 1).load(.acquire), if (r.word(w_failed + 4).load(.acquire) == 1) "aborted" else "timed out", r.word(w_failed + 5).load(.acquire), r.word(w_failed + 3).load(.acquire), r.word(w_failed + 2).load(.acquire), r.word(w_seq).load(.acquire), r.last_seq, r.word(w_completed).load(.acquire), r.word(w_abort).load(.acquire), r.word(w_proxy_failed).load(.acquire), r.flagWord(0), r.flagWord(1) }); + return error.PeerFailed; + } + if (r.word(w_proxy_failed).load(.acquire) != 0) return error.BackendFailed; + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return error.Aborted; + } + + fn flagWord(r: *const Roce, slot: usize) u32 { + const p: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + slot * 2 * flag_stride ..].ptr)); + return p.load(.acquire); + } + + pub fn proxyError(r: *const Roce) []const u8 { + return r.err_buf[0..r.err_len]; + } + + /// The kernel's side of one op done by the host (tests, `roce-host`): a gather through the same slots and flags. + pub fn hostGather(r: *Roce, in: []const u8, out: []u8) Error!void { + return r.hostGatherV(in, out, null); + } + + /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. + pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + const n = in.len; + if (n == 0 or n > r.slot_bytes or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const mine: usize = if (lens) |l| l[r.me] else n; + const theirs: usize = if (lens) |l| l[1 - r.me] else n; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[0..mine]); + r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); + r.word(w_seq).store(seq, .release); + @memcpy(out[r.me * n ..][0..mine], in[0..mine]); + const t0 = sock.nowNs(); + for (0..r.wire.hcas()) |h| { + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); + while (flag.load(.acquire) != seq) { + if (r.word(w_abort).load(.monotonic) != 0) return error.Aborted; + if (sock.nowNs() - t0 > r.timeout_ns) return error.Timeout; + std.atomic.spinLoopHint(); + } + } + @memcpy(out[(1 - r.me) * n ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} + +fn pin(cpu: u32) void { + var set: std.os.linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + std.os.linux.sched_setaffinity(0, &set) catch log.warn("tp roce: could not pin the proxy to CPU {d}", .{cpu}); +} diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..afb705d41b --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,58 @@ +//! `tf-tp-test roce-host`: the RoCE wire and proxy between two hosts without a GPU, checked bit for bit, then timed. +const std = @import("std"); +const tp = @import("tp"); +const roce = tp.roce; +const rv = tp.roce_verbs; + +const sizes = [_]usize{ 16, 4096, 8192, 16384, 65536, 262144 }; + +pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { + var b = try tp.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var n: usize = 0; + if (cfg.roce_hca) |text| n = try rv.parseSpecs(text, &specs) else { + var v = try tp.Verbs.open(null); + defer v.close(); + n = rv.autoSpecs(&v, &specs, cfg.roce_hcas); + } + const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); + defer r.deinit(); + const in = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(in); + const peer = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(peer); + const out = try gpa.alloc(u8, 512 * 1024); + defer gpa.free(out); + // 1. every size from 1 byte to the slot, both slots, bit for bit + var checked: u32 = 0; + for (0..1000) |k| { + const len = 1 + (k * 7919) % (256 * 1024); + tp.reference.fill(in[0..len], .u8, cfg.rank, k); + tp.reference.fill(peer[0..len], .u8, 1 - cfg.rank, k); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + const mine_at = cfg.rank * len; + const peer_at = (1 - cfg.rank) * len; + if (!std.mem.eql(u8, out[mine_at..][0..len], in[0..len]) or !std.mem.eql(u8, out[peer_at..][0..len], peer[0..len])) { + std.debug.print("FAIL roce-host gather {d} of {d} bytes\n", .{ k, len }); + return 1; + } + checked += 1; + } + try r.check(); + std.debug.print("PASS roce-host: {d} gathers of 1 B to 256 KiB bit-exact over {s}\n", .{ checked, @tagName(cfg.roce_wire) }); + // 2. one exchange's host-to-host time by size: median and p90 of 2,000 back-to-back ops + var times: [2000]u64 = undefined; + for (sizes) |len| { + for (×) |*t| { + const t0 = tp.sock.nowNs(); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + t.* = tp.sock.nowNs() - t0; + } + std.mem.sort(u64, ×, {}, std.sort.asc(u64)); + if (cfg.rank == 0) std.debug.print("RESULT {{\"backend\":\"roce-host\",\"wire\":\"{s}\",\"bytes\":{d},\"p50_us\":{d:.2},\"p90_us\":{d:.2}}}\n", .{ @tagName(cfg.roce_wire), len, @as(f64, @floatFromInt(times[1000])) / 1e3, @as(f64, @floatFromInt(times[1800])) / 1e3 }); + } + try r.check(); + try b.agree("done", "roce-host end"); + return 0; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..51e9d79133 --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,116 @@ +//! The RoCE protocol end to end on host threads: two ranks, TCP bootstrap, proxy, shm or verbs (TF_TP_ROCE_TEST_HCA). +const std = @import("std"); +const roce = @import("roce.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const reference = @import("reference.zig"); +const Bootstrap = bootstrap.Bootstrap; + +const Run = struct { + port: u16, + wire: roce.WireKind, + devs: [2][]const u8 = .{ "", "" }, + ops: u32, + timeout_ns: u64 = 10 * std.time.ns_per_s, + /// Rank 1 stops after this many ops (rank 0 must then fail its next wait, by timeout or abort). + stop_after: ?u32 = null, + abort_after_ms: ?u64 = null, + errs: [2]?anyerror = .{ null, null }, + done: [2]u32 = .{ 0, 0 }, + max_bytes: usize = 16 * 1024, + /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. + var_lens: bool = false, +}; + +fn rank(run: *Run, me: u32, lfd: sock.Fd) void { + run.errs[me] = body(run, me, lfd); +} + +fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { + var b = (if (me == 0) Bootstrap.acceptOn(lfd, 2, 5 * std.time.ns_per_s) else Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = run.port, .connect_timeout_ns = 5 * std.time.ns_per_s })) catch |e| return e; + defer b.close(); + var spec: [1]rv.HcaSpec = .{.{}}; + if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns }) catch |e| return e; + defer r.deinit(); + var abort_thread: ?std.Thread = null; + defer if (abort_thread) |t| t.join(); + var in: [64 * 1024]u8 = undefined; + var peer: [64 * 1024]u8 = undefined; + var out: [128 * 1024]u8 = undefined; + for (0..run.ops) |k| { + if (run.stop_after) |s| if (k == s) { + // rank 1 skips op s; rank 0 waits in it until its timeout, or until the abort armed now + if (me == 1) { + sock.sleepNs(300 * std.time.ns_per_ms); + return null; + } + if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; + }; + // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds + const n = 1 + (k * 7919) % run.max_bytes; + reference.fill(in[0..n], .u8, me, k); + reference.fill(peer[0..n], .u8, 1 - me, k); + const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; + const l: [2]usize = if (run.var_lens) .{ lens[0], lens[1] } else .{ n, n }; + r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const r0 = if (me == 0) in[0..n] else peer[0..n]; + const r1 = if (me == 1) in[0..n] else peer[0..n]; + if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; + run.done[me] += 1; + } + r.check() catch |e| return e; + // both ranks finish before either tears down (the peer may still read a flag this rank's wire writes) + b.agree("done", "test end") catch |e| return e; + return null; +} + +fn abortLater(r: *roce.Roce, ms: u64) void { + sock.sleepNs(ms * std.time.ns_per_ms); + r.abort(); +} + +fn twoRanks(run: *Run) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + run.port = l.port; + const t1 = try std.Thread.spawn(.{}, rank, .{ run, 1, l.fd }); + rank(run, 0, l.fd); + t1.join(); +} + +test "the RoCE protocol over the shm wire: 2,000 gathers of every size, bit for bit" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the RoCE protocol with variable lengths over the shm wire: each rank posts its own padded count, empty shards included" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; + try twoRanks(&run); + try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); + try std.testing.expectEqual(@as(u32, 20), run.done[0]); + var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + try twoRanks(&run2); + try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); + try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); +} + +test "the RoCE protocol over verbs (TF_TP_ROCE_TEST_HCA=dev0,dev1: soft-RoCE or real HCAs)" { + const raw = std.c.getenv("TF_TP_ROCE_TEST_HCA") orelse return error.SkipZigTest; + var it = std.mem.splitScalar(u8, std.mem.span(raw), ','); + var run: Run = .{ .port = 0, .wire = .verbs, .ops = 2000 }; + run.devs = .{ it.next() orelse return error.SkipZigTest, it.next() orelse return error.SkipZigTest }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig new file mode 100644 index 0000000000..bc99681b5e --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,297 @@ +//! The RoCE wire over verbs: one RC queue pair a HCA to the peer, a send slot striped over the HCAs with seq flags. +const std = @import("std"); +const log = @import("log.zig"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const Verbs = verbs.Verbs; + +pub const max_hcas = 2; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const flag_stride = 128; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 1; // "TFRD", bumped once for this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord }; + +/// One HCA as configured: its device name and RoCE v2 GID index (null: detected). +pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; + +/// What a rank tells its peer to connect (both ranks build it the same way; sent through the bootstrap). +pub const Record = extern struct { + magic: u32 = record_magic, + n_hca: u32 = 0, + region: u64 = 0, + rkey: [max_hcas]u32 = @splat(0), + lid: [max_hcas]u32 = @splat(0), + mtu: [max_hcas]u32 = @splat(0), + qpn: [max_hcas]u32 = @splat(0), + gid: [max_hcas][16]u8 = @splat(@splat(0)), +}; + +const Hca = struct { + ctx: *abi.Context, + pd: *abi.Pd, + mr: *abi.Mr, + cq: *abi.Cq, + qp: *abi.Qp, + gid_index: u8, + gid: [16]u8, + mtu: u32, + lid: u16, + outstanding: u32 = 0, +}; + +/// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). +pub fn parseSpecs(text: []const u8, out: *[max_hcas]HcaSpec) !usize { + var n: usize = 0; + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + if (n == max_hcas) return error.TooManyHcas; + var parts = std.mem.splitScalar(u8, std.mem.trim(u8, item, " "), ':'); + const name = parts.next().?; + if (name.len == 0 or name.len >= 64) return error.BadHcaName; + out[n] = .{ .name_len = @intCast(name.len) }; + @memcpy(out[n].name[0..name.len], name); + if (parts.next()) |g| out[n].gid = try std.fmt.parseInt(u8, g, 10); + n += 1; + } + return n; +} + +/// RoCE v2 by sysfs (the kernel's own word for each GID); null when sysfs cannot say. +fn sysfsRoceV2(name: []const u8, index: u32) ?bool { + var path: [256]u8 = undefined; + const p = std.fmt.bufPrintSentinel(&path, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, index }, 0) catch return null; + const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [32]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return false; + return std.mem.startsWith(u8, buf[0..@intCast(n)], "RoCE v2"); +} + +/// An IPv4-mapped GID (::ffff:a.b.c.d): the address the HCA's netdev carries. +pub fn ipv4Mapped(g: [16]u8) bool { + return std.mem.allEqual(u8, g[0..10], 0) and g[10] == 0xff and g[11] == 0xff; +} + +/// The lowest RoCE v2 GID index that is IPv4-mapped (GID indices move on reboot: never hard-coded). +fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) ?u8 { + var i: u32 = 0; + while (i < @min(@as(u32, @intCast(@max(table, 0))), 255)) : (i += 1) { + const g = v.gid(ctx, 1, i) catch continue; + if (!ipv4Mapped(g.raw)) continue; + const v2 = sysfsRoceV2(name, i) orelse (if (v.gidType(ctx, 1, i)) |t| t == 2 else true); + if (v2) return @intCast(i); + } + return null; +} + +/// Every device whose port 1 is an active Ethernet port with an IPv4 RoCE v2 GID, by name order, at most `limit`. +pub fn autoSpecs(v: *const Verbs, out: *[max_hcas]HcaSpec, limit: usize) usize { + var count: c_int = 0; + const list = v.api.ibv_get_device_list(&count) orelse return 0; + defer v.api.ibv_free_device_list(list); + var n: usize = 0; + var i: usize = 0; + while (list[i]) |dev| : (i += 1) { + if (n == limit) break; + const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); + const ctx = v.api.ibv_open_device(dev) orelse continue; + defer _ = v.api.ibv_close_device(ctx); + const port = v.port(ctx, 1) catch continue; + if (port.state != abi.port_active or port.link_layer != link_ethernet) continue; + const g = detectGid(v, ctx, name, port.gid_tbl_len) orelse continue; + out[n] = .{ .name_len = @intCast(@min(name.len, 63)), .gid = g }; + @memcpy(out[n].name[0..out[n].name_len], name[0..out[n].name_len]); + n += 1; + } + return n; +} + +pub const VerbsWire = struct { + v: Verbs, + hcas: [max_hcas]Hca = undefined, + n_hca: u32 = 0, + traffic_class: u8, + peer_region: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + /// Offsets in both ranks' regions (the same layout): receive slots and flags. + recv_off: usize, + flag_off: usize, + send_off: usize, + slot_bytes: usize, + region: []u8, + posted: u64 = 0, + completed: u64 = 0, + + /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. + pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { + if (specs.len == 0 or specs.len > max_hcas) return error.NoHca; + var w: VerbsWire = .{ .v = try Verbs.open(null), .traffic_class = traffic_class, .recv_off = layout.recv, .flag_off = layout.flags, .send_off = layout.send, .slot_bytes = layout.slot, .region = region }; + errdefer w.close(); + for (specs) |s| { + try w.openHca(s); + } + return w; + } + + fn openHca(w: *VerbsWire, s: HcaSpec) !void { + const api = w.v.api; + const name = s.name[0..s.name_len]; + const ctx = try w.v.openDevice(name); + errdefer _ = api.ibv_close_device(ctx); + const port = try w.v.port(ctx, 1); + if (port.state != abi.port_active) { + log.warn("tp roce: {s} port 1 is not active", .{name}); + return error.NoHca; + } + const gi = s.gid orelse detectGid(&w.v, ctx, name, port.gid_tbl_len) orelse { + log.warn("tp roce: {s} has no IPv4 RoCE v2 GID", .{name}); + return error.NoHca; + }; + const gid = try w.v.gid(ctx, 1, gi); + const pd = api.ibv_alloc_pd(ctx) orelse return error.VerbsFailed; + errdefer _ = api.ibv_dealloc_pd(pd); + const mr = api.ibv_reg_mr(pd, w.region.ptr, w.region.len, abi.access_local_write | abi.access_remote_write) orelse { + log.warn("tp roce: ibv_reg_mr on {s} failed (memlock limit? ulimit -l unlimited)", .{name}); + return error.VerbsFailed; + }; + errdefer _ = api.ibv_dereg_mr(mr); + const cq = api.ibv_create_cq(ctx, send_depth, null, null, 0) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_cq(cq); + var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; + const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_qp(qp); + var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; + if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; + w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; + w.n_hca += 1; + log.info("[tensorfold] tp roce: {s} GID {d} ({d}.{d}.{d}.{d}), MTU enum {d}", .{ name, gi, gid.raw[12], gid.raw[13], gid.raw[14], gid.raw[15], port.active_mtu }); + } + + pub fn close(w: *VerbsWire) void { + const api = w.v.api; + for (w.hcas[0..w.n_hca]) |h| { + _ = api.ibv_destroy_qp(h.qp); + _ = api.ibv_destroy_cq(h.cq); + _ = api.ibv_dereg_mr(h.mr); + _ = api.ibv_dealloc_pd(h.pd); + _ = api.ibv_close_device(h.ctx); + } + w.n_hca = 0; + w.v.close(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr) }; + for (w.hcas[0..w.n_hca], 0..) |h, i| { + r.rkey[i] = h.mr.rkey; + r.lid[i] = h.lid; + r.mtu[i] = h.mtu; + r.qpn[i] = h.qp.qp_num; + r.gid[i] = h.gid; + } + return r; + } + + /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. + pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { + if (peer.magic != record_magic) return error.BadRecord; + if (peer.n_hca != mine.n_hca) return error.BadRecord; + var out: [max_hcas]u32 = undefined; + for (0..mine.n_hca) |i| { + out[i] = @intCast(i); + var hits: u32 = 0; + for (0..peer.n_hca) |j| if (std.mem.eql(u8, mine.gid[i][12..15], peer.gid[j][12..15])) { + out[i] = @intCast(j); + hits += 1; + }; + if (hits > 1) out[i] = @intCast(i); // one subnet for both: keep the order + } + return out; + } + + /// INIT to RTR to RTS against the peer's record. + pub fn connect(w: *VerbsWire, peer: Record) !void { + const pairs = try pairing(w.record(), peer); + const api = w.v.api; + const m = abi.mask; + w.peer_region = peer.region; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + const j = pairs[i]; + w.peer_rkey[i] = peer.rkey[j]; + var rtr: abi.QpAttr = .{ .qp_state = abi.qps_rtr, .path_mtu = @intCast(@min(peer.mtu[j], h.mtu)), .dest_qp_num = peer.qpn[j], .rq_psn = 0, .max_dest_rd_atomic = 1, .min_rnr_timer = 12 }; + rtr.ah_attr = .{ .grh = .{ .dgid = .{ .raw = peer.gid[j] }, .sgid_index = h.gid_index, .hop_limit = 64, .traffic_class = w.traffic_class }, .dlid = @intCast(peer.lid[j]), .is_global = 1, .port_num = 1 }; + if (api.ibv_modify_qp(h.qp, &rtr, m.state | m.av | m.path_mtu | m.dest_qpn | m.rq_psn | m.max_dest_rd_atomic | m.min_rnr_timer) != 0) return error.VerbsFailed; + var rts: abi.QpAttr = .{ .qp_state = abi.qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .sq_psn = 0, .max_rd_atomic = 1 }; + if (api.ibv_modify_qp(h.qp, &rts, m.state | m.timeout | m.retry_cnt | m.rnr_retry | m.sq_psn | m.max_qp_rd_atomic) != 0) return error.VerbsFailed; + } + } + + /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. + pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { + const packs = padded / 16; + var offset: u64 = 0; + var seq_copy = seq; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + while (h.outstanding >= send_depth / 4) try w.reapHca(h); + var stripe = packs / w.n_hca; + if (i < packs % w.n_hca) stripe += 1; + const bytes: u32 = stripe * 16; + var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + w.flag_off + (@as(u64, slot) * max_hcas + i) * flag_stride, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + w.send_off + slot * w.slot_bytes + offset, .length = bytes, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + w.recv_off + slot * w.slot_bytes + offset, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (bytes > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.outstanding += 1; + offset += bytes; + } + w.posted += 1; + } + + fn reapHca(w: *VerbsWire, h: *Hca) !void { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return error.Completion; + for (wc[0..@intCast(n)]) |c| if (c.status != abi.wc_success) { + log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); + return error.Completion; + }; + h.outstanding -= @intCast(n); + w.completed += @intCast(n); + } + + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). + pub fn reap(w: *VerbsWire) !void { + for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); + } +}; + +test "HCA lists parse and subnets pair the functions" { + var specs: [max_hcas]HcaSpec = undefined; + const n = try parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqual(@as(usize, 2), n); + try std.testing.expectEqualStrings("rocep1s0f1", specs[0].name[0..specs[0].name_len]); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); + try std.testing.expectEqual(@as(?u8, null), specs[1].gid); + var a: Record = .{ .n_hca = 2 }; + var b: Record = .{ .n_hca = 2 }; + const sub100 = [16]u8{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 192, 168, 100, 10 }; + var sub101 = sub100; + sub101[14] = 101; + a.gid = .{ sub100, sub101 }; + b.gid = .{ sub101, sub100 }; // the peer lists its functions the other way round + b.gid[0][15] = 11; + b.gid[1][15] = 11; + const p = try VerbsWire.pairing(a, b); + try std.testing.expectEqual(@as(u32, 1), p[0]); + try std.testing.expectEqual(@as(u32, 0), p[1]); + try std.testing.expect(ipv4Mapped(sub100)); + b.n_hca = 1; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..17a4ccbe91 --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# One rank of the two-host TP=2 gate over RoCE: `run_hosts.sh RANK RANK0_ADDR ./tf-tp-test OUTDIR` on both hosts. +set -u +RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} +mkdir -p "$OUT" +[ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; } > "$OUT/host.txt" 2>&1 +PORT=29700 +step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line + local name=$1 backend=$2; shift 2 + PORT=$((PORT + 10)) + ( TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend TF_TP_ROCE_FALLBACK=error \ + timeout 900 "$BIN" "$@"; echo "EXIT rank $RANK rc $? at $(date +%s.%N)" ) > "$OUT/$name.log" 2>&1 + echo "$name: $(grep -h '^SUMMARY\|^PASS roce-host\|^INJECT\|^EXIT' "$OUT/$name.log" | tr '\n' ' ')" +} +step roce-host roce roce-host +step suite-roce roce suite +step suite-nccl nccl suite +for backend in roce nccl; do + for mode in fatal kill; do + for victim in 0 1; do step "ff-$backend-$mode-$victim" "$backend" failfast "$victim" "$mode"; done + done +done diff --git a/zig/src/fabric/tp/run_pair.sh b/zig/src/fabric/tp/run_pair.sh new file mode 100644 index 0000000000..112232816b --- /dev/null +++ b/zig/src/fabric/tp/run_pair.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Two ranks of tf-tp-test on one two-GPU host: `[SUITES=..] [FAILFAST=..] run_pair.sh BIN OUTDIR`. +set -u +BIN=${1:?bin}; OUT=${2:?outdir} +mkdir -p "$OUT" +nvidia-smi > "$OUT/smi.txt" 2>&1 +nvidia-smi topo -m > "$OUT/topo.txt" 2>&1 +ldconfig -p | grep -E "libnccl|libcuda" > "$OUT/libs.txt" 2>&1 +PORT=29500 +pair() { # NAME BACKEND ARGS...: both ranks, logs OUT/NAME.rR.log ending with an EXIT line + local name=$1 backend=$2 r; shift 2 + PORT=$((PORT + 10)) + for r in 0 1; do + ( TF_TP_RANK=$r TF_TP_DEVICE=$r TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend NCCL_DEBUG=${NCCL_DEBUG:-WARN} \ + timeout 600 "$BIN" "$@"; echo "EXIT rank $r rc $? at $(date +%s.%N)" ) > "$OUT/$name.r$r.log" 2>&1 & + done + wait +} +for backend in ${SUITES:-mailbox nccl}; do + pair "suite-$backend" "$backend" suite +done +: > "$OUT/failfast.txt" +for backend in ${FAILFAST:-nccl mailbox}; do + for mode in fatal kill; do + for victim in 0 1; do + name=ff-$backend-$mode-$victim + pair "$name" "$backend" failfast "$victim" "$mode" + python3 - "$OUT" "$name" "$victim" >> "$OUT/failfast.txt" <<'EOF' +import re, sys +out, name, victim = sys.argv[1], sys.argv[2], int(sys.argv[3]) +logs = [open(f"{out}/{name}.r{r}.log").read() for r in (0, 1)] +inj = re.search(r"INJECT \w+ rank \d+ at ([\d.]+)", logs[victim]) +ex = [re.search(r"EXIT rank \d rc (\d+) at ([\d.]+)", l) for l in logs] +if not inj or not all(ex): + print(f"{name}: FAIL (no inject / exit line)"); sys.exit() +t0 = float(inj.group(1)) +rcs = [int(e.group(1)) for e in ex] +dt = [float(e.group(2)) - t0 for e in ex] +want = [137 if (r == victim and "kill" in name) else 70 for r in (0, 1)] +ok = rcs == want and dt[1 - victim] < 3.0 +print(f"{name}: {'PASS' if ok else 'FAIL'} rc {rcs} (want {want}), exit after injection: rank0 {dt[0]:.3f} s, rank1 {dt[1]:.3f} s") +EOF + done + done +done +grep -h "^SUMMARY\|^FAIL" "$OUT"/suite*.log +cat "$OUT/failfast.txt" diff --git a/zig/src/fabric/tp/sock.zig b/zig/src/fabric/tp/sock.zig new file mode 100644 index 0000000000..b382884034 --- /dev/null +++ b/zig/src/fabric/tp/sock.zig @@ -0,0 +1,169 @@ +//! TCP for the bootstrap and the fate channel: listen, accept and connect with deadlines, whole sends and receives. +const std = @import("std"); +const c = std.c; +const posix = std.posix; + +pub const Fd = c.fd_t; + +pub const Error = error{ SocketFailed, AddressUnresolved, AddressInUse, Timeout, Closed, Interrupted }; + +pub fn nowNs() u64 { + var ts: c.timespec = undefined; + _ = c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +pub fn sleepNs(ns: u64) void { + var ts: c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + while (c.nanosleep(&ts, &ts) != 0) {} +} + +/// TCP_NODELAY (frames are tiny and latency matters) and keepalive (a vanished host fails a read in ~90 s, not never). +fn tune(fd: Fd) void { + const one: c_int = 1; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &one, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &one, @sizeOf(c_int)); + const idle: c_int = 30; + const intvl: c_int = 10; + const cnt: c_int = 6; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPIDLE, &idle, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPINTVL, &intvl, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPCNT, &cnt, @sizeOf(c_int)); +} + +/// Rank 0's listening socket on every IPv4 address at `port`. +pub fn listen(port: u16) Error!Fd { + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + errdefer _ = c.close(fd); + const one: c_int = 1; + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, &one, @sizeOf(c_int)); + var addr: posix.sockaddr.in = .{ .port = std.mem.nativeToBig(u16, port), .addr = 0 }; + if (c.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in)) != 0) return error.AddressInUse; + if (c.listen(fd, 64) != 0) return error.SocketFailed; + return fd; +} + +/// One connection before `deadline_ns` (monotonic). +pub fn accept(fd: Fd, deadline_ns: u64) Error!Fd { + while (true) { + try waitReadable(fd, deadline_ns); + const conn = c.accept(fd, null, null); + if (conn >= 0) { + tune(conn); + return conn; + } + } +} + +/// Connects to `host:port`, retrying every 50 ms (the master may not listen yet) until `deadline_ns`. +pub fn connect(host: []const u8, port: u16, deadline_ns: u64) Error!Fd { + var name: [256]u8 = undefined; + if (host.len >= name.len) return error.AddressUnresolved; + @memcpy(name[0..host.len], host); + name[host.len] = 0; + var service: [8]u8 = undefined; + const svc = std.fmt.bufPrintSentinel(&service, "{d}", .{port}, 0) catch unreachable; + const hints: c.addrinfo = .{ .flags = .{}, .family = posix.AF.INET, .socktype = posix.SOCK.STREAM, .protocol = 0, .addrlen = 0, .addr = null, .canonname = null, .next = null }; + while (true) { + var res: ?*c.addrinfo = null; + if (@backingInt(c.getaddrinfo(name[0..host.len :0], svc, &hints, &res)) == 0) { + defer c.freeaddrinfo(res.?); + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + if (c.connect(fd, res.?.addr.?, res.?.addrlen) == 0) { + tune(fd); + return fd; + } + _ = c.close(fd); + } + if (nowNs() >= deadline_ns) return error.Timeout; + sleepNs(50 * std.time.ns_per_ms); + } +} + +fn waitReadable(fd: Fd, deadline_ns: u64) Error!void { + while (true) { + const t = nowNs(); + if (t >= deadline_ns) return error.Timeout; + const ms: c_int = @intCast(@min((deadline_ns - t) / std.time.ns_per_ms + 1, 1000)); + var p = [_]c.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const n = c.poll(&p, 1, ms); + if (n > 0) return; + } +} + +/// Every byte of `bytes`, or `error.Closed`. +pub fn sendAll(fd: Fd, bytes: []const u8) Error!void { + var off: usize = 0; + while (off < bytes.len) { + const n = c.send(fd, bytes[off..].ptr, bytes.len - off, posix.MSG.NOSIGNAL); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// Fills `out` before `deadline_ns` (null: no deadline), or `error.Closed` at the end of the stream. +pub fn recvExact(fd: Fd, out: []u8, deadline_ns: ?u64) Error!void { + var off: usize = 0; + while (off < out.len) { + if (deadline_ns) |d| try waitReadable(fd, d); + const n = c.recv(fd, out[off..].ptr, out.len - off, 0); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// A non-blocking send of the whole of `bytes` if the socket buffer has room (false: dropped, or the peer is gone). +pub fn trySend(fd: Fd, bytes: []const u8) bool { + const n = c.send(fd, bytes.ptr, bytes.len, posix.MSG.NOSIGNAL | posix.MSG.DONTWAIT); + return n == @as(isize, @intCast(bytes.len)); +} + +pub fn hangup(fd: Fd) void { + _ = c.shutdown(fd, posix.SHUT.RDWR); +} + +pub fn close(fd: Fd) void { + _ = c.close(fd); +} + +/// This host's identity (the kernel's boot id, else the host name): ranks with equal ids share memory. +pub fn hostId(out: *[64]u8) []const u8 { + @memset(out, 0); + const f = c.open("/proc/sys/kernel/random/boot_id", .{ .ACCMODE = .RDONLY }); + if (f >= 0) { + defer _ = c.close(f); + const n = c.read(f, out, out.len); + if (n > 0) return std.mem.trim(u8, out[0..@intCast(n)], " \n"); + } + if (c.gethostname(out, out.len) != 0) return ""; + return std.mem.sliceTo(out, 0); +} + +test "a listener and a client trade bytes and see the hang-up" { + const l = try listen(0); + defer close(l); + var addr: posix.sockaddr.in = undefined; + var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); + try std.testing.expectEqual(@as(c_int, 0), c.getsockname(l, @ptrCast(&addr), &len)); + const port = std.mem.bigToNative(u16, addr.port); + const deadline = nowNs() + 5 * std.time.ns_per_s; + const a = try connect("localhost", port, deadline); + defer close(a); + const b = try accept(l, deadline); + defer close(b); + try sendAll(a, "hello"); + var got: [5]u8 = undefined; + try recvExact(b, &got, deadline); + try std.testing.expectEqualStrings("hello", &got); + try std.testing.expectError(error.Timeout, recvExact(b, &got, nowNs() + 20 * std.time.ns_per_ms)); + hangup(a); + try std.testing.expectError(error.Closed, recvExact(b, &got, deadline)); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig new file mode 100644 index 0000000000..9094429c29 --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,212 @@ +//! Tensor parallelism across processes, one GPU each: bootstrap, collectives, fail-fast; `Session.open` starts a rank. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); + +pub const collective = @import("collective.zig"); +pub const reference = @import("reference.zig"); +pub const config = @import("config.zig"); +pub const sock = @import("sock.zig"); +pub const bootstrap = @import("bootstrap.zig"); +pub const fate = @import("fate.zig"); +pub const host = @import("host.zig"); +pub const nccl = @import("nccl.zig"); +pub const mailbox = @import("mailbox.zig"); +pub const hybrid = @import("hybrid.zig"); +pub const roce = @import("roce.zig"); +pub const roce_verbs = @import("roce_verbs.zig"); +pub const planlink = @import("planlink.zig"); +pub const Verbs = @import("verbs").Verbs; + +pub const Collective = collective.Collective; +pub const DType = collective.DType; +pub const Op = collective.Op; +pub const Config = config.Config; +pub const Bootstrap = bootstrap.Bootstrap; +pub const Fate = fate.Fate; + +pub const Session = struct { + cfg: Config, + lib: cuda.nccl.Library, + nccl: nccl.Nccl, + mail: ?mailbox.Mailbox = null, + roce: ?*roce.Roce = null, + route: hybrid.Hybrid = undefined, + /// The round plan's link (rank 0 sends, the others receive), on its own TCP connection. + plan: planlink.PlanLink = undefined, + fate: ?Fate = null, + + /// Brings this rank up in place; `kernel_image` is mailbox.cu's fatbin, empty sends those backends to NCCL. + pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { + var b = try Bootstrap.open(cfg); + defer b.close(); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); + errdefer lib.close(); + const v = try lib.version(); + try b.agree(std.mem.asBytes(&v), "NCCL version"); + var id: cuda.nccl.UniqueId = undefined; + if (cfg.rank == 0) try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try b.broadcast(&id.internal, &id.internal); + s.* = .{ .cfg = cfg, .lib = lib, .nccl = undefined }; + s.nccl = try nccl.Nccl.init(&s.lib, d, id, cfg.rank, cfg.world); + errdefer s.nccl.deinit(); + switch (cfg.backend) { + .nccl => {}, + .mailbox => if (try mailboxPossible(&b, kernel_image)) { + s.mail = try mailbox.Mailbox.init(d, &b, kernel_image, cfg.mailbox_max_bytes, cfg.mailbox_timeout_ns); + s.route = .{ .small = .{ .mailbox = &s.mail.? }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + } else if (cfg.rank == 0) { + log.warn("[tensorfold] tp: TF_COMM_BACKEND=mailbox needs two ranks on one host and the kernel; every rank uses NCCL", .{}); + }, + .roce => try s.openRoce(d, &b, kernel_image), + } + try s.agreeVar(&b); + s.plan = planlink.PlanLink.init(cfg.rank, cfg.world, b.releasePlan()); + s.plan.spin_ns = cfg.plan_spin_ns; + s.plan.pin = cfg.plan_pin; + s.plan.avoid = cfg.roce_cpu; + if (cfg.failfast) { + s.fate = Fate.init(cfg, b.release()); + s.fate.?.onFail(.{ .ctx = s, .call = abortHook }); + try s.fate.?.start(); + } + log.info("[tensorfold] tp: rank {d} of {d} up ({t}, NCCL {d}, fail-fast {s}, roce fast {s}: {d} poller(s), stagger {d} ns)", .{ cfg.rank, cfg.world, s.comm().kind(), v, if (cfg.failfast) "on" else "off", if (cfg.roce_fast & 1 != 0) "on" else "off", (cfg.roce_fast >> 8) & 0xff, cfg.roce_fast >> 16 }); + } + + /// RoCE for the small exchanges, probed against NCCL at load; on failure every rank obeys TF_TP_ROCE_FALLBACK. + fn openRoce(s: *Session, d: *const cuda.Driver, b: *Bootstrap, image: []const u8) !void { + const cfg = s.cfg; + var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; + var n_specs: usize = 0; + if (cfg.roce_wire == .verbs) { + if (cfg.roce_hca) |text| { + n_specs = try roce_verbs.parseSpecs(text, &specs); + } else if (Verbs.open(null)) |vb| { + var v = vb; + defer v.close(); + n_specs = roce_verbs.autoSpecs(&v, &specs, cfg.roce_hcas); + } else |_| {} + } + const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; + const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); + s.roce = r; + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + const good = probe(s, d) catch false; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(good)}, &buf, &blobs); + for (all) |x| if (x[0] != 1) { + r.deinit(); + s.roce = null; + return s.roceFallback(error.ProbeMismatch); + }; + } + + /// `Collective.varAgreed`: whether every rank's transport moves device lengths, agreed once a session. + fn agreeVar(s: *Session, b: *Bootstrap) !void { + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(s.comm().varFits(16))}, &buf, &blobs); + var ok = true; + for (all) |x| ok = ok and x[0] == 1; + if (s.mail != null or s.roce != null) s.route.var_agreed = ok; + } + + fn roceFallback(s: *Session, e: anyerror) !void { + if (s.cfg.roce_fallback == .@"error") return e; + if (s.cfg.rank == 0) log.warn("[tensorfold] tp: RoCE unavailable ({t}); every rank uses NCCL (TF_TP_ROCE_FALLBACK=error refuses instead)", .{e}); + } + + /// One 4 KiB all-gather and one bf16 sum through RoCE and through NCCL: the bytes must be equal. + fn probe(s: *Session, d: *const cuda.Driver) !bool { + const n = 4096; + // input [0, n), gathers [n, 3n) and [3n, 5n), sums [5n, 6n) and [6n, 7n) + var a = try cuda.DeviceBuffer.alloc(d, 7 * n); + defer a.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + var bytes: [n]u8 = undefined; + reference.fill(&bytes, .bf16, s.cfg.rank, 0x7072_6f62); + try a.uploadAsync(0, &bytes, stream.handle); + const fast = s.route.iface(); + const slow = s.nccl.iface(); + try fast.allGather(a.ptr, a.ptr + n, n, .u8, stream.handle); + try slow.allGather(a.ptr, a.ptr + 3 * n, n, .u8, stream.handle); + try fast.allReduce(a.ptr, a.ptr + 5 * n, n / 2, .bf16, .sum, stream.handle); + try slow.allReduce(a.ptr, a.ptr + 6 * n, n / 2, .bf16, .sum, stream.handle); + try stream.synchronize(); + try fast.check(); + var got: [6 * n]u8 = undefined; + try a.download(n, &got); + const gathers = std.mem.eql(u8, got[0 .. 2 * n], got[2 * n .. 4 * n]); + const sums = std.mem.eql(u8, got[4 * n .. 5 * n], got[5 * n .. 6 * n]); + if (!gathers or !sums) log.warn("[tensorfold] tp roce probe: gather {s}, sum {s} against NCCL", .{ if (gathers) "equal" else "DIFFERS", if (sums) "equal" else "DIFFERS" }); + return gathers and sums; + } + + fn mailboxPossible(b: *Bootstrap, image: []const u8) !bool { + if (b.world != 2) return false; + var id_buf: [64]u8 = undefined; + const id = sock.hostId(&id_buf); + var mine: [65]u8 = @splat(0); + mine[0] = @intFromBool(image.len > 0); + @memcpy(mine[1..][0..id.len], id); + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [2 * 65]u8 = undefined; + const all = try b.allGather(&mine, &buf, &blobs); + return all[0][0] == 1 and std.mem.eql(u8, all[0], all[1]); + } + + fn abortHook(ctx: *anyopaque, _: []const u8) void { + const s: *Session = @ptrCast(@alignCast(ctx)); + s.comm().abort(); + } + + /// The communicator the model uses: NCCL, or the hybrid route over the mailbox / RoCE. + pub fn comm(s: *Session) Collective { + if (s.mail != null or s.roce != null) return s.route.iface(); + return s.nccl.iface(); + } + + /// This rank failed inside a round: every rank goes down (exit 70). Without fail-fast it only aborts. + pub fn fatal(s: *Session, where: []const u8, err: anyerror) void { + if (s.fate) |*f| return f.fatal(where, err); + s.comm().abort(); + } + + /// A clean stop announced by rank 0 before it closes: followers end with no fail-fast exit. + pub fn announceStop(s: *Session) void { + if (s.fate) |*f| f.announceStop(); + } + + /// This rank failed after `open`, before serving: the other ranks learn why and go down (fail-fast). + pub fn bootFailed(s: *Session, err: anyerror) void { + if (s.fate) |*f| f.failBoot("boot", err); + s.close(); + } + + /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). + pub fn close(s: *Session) void { + if (s.fate) |*f| f.close(); + s.plan.close(); + if (s.roce) |r| r.deinit(); + if (s.mail) |*m| m.deinit(); + s.nccl.deinit(); + s.lib.close(); + } +}; + +test { + _ = collective; + _ = reference; + _ = config; + _ = sock; + _ = bootstrap; + _ = mailbox; + _ = roce_verbs; + _ = planlink; + _ = @import("roce_test.zig"); + _ = @import("host_test.zig"); + _ = @import("fate_test.zig"); +} diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..f3af36587b 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,7 +1,7 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); const builtin = @import("builtin"); -const abi = @import("verbs_abi.zig"); +pub const abi = @import("verbs_abi.zig"); pub const Error = error{ LibraryUnavailable, MissingSymbol, NoDeviceList, NoSuchDevice, VerbsFailed, BadDeviceList, NoUsableGid }; pub const max_links = 8; From db3bf32f8c97f089bbc3dadbfdff76b2eeefac18 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:56 +0700 Subject: [PATCH 5/7] TP tests: a drain on rank 0 then its clean stop ends rank 1's loop at the stop word with no fail-fast exit; a rank killed mid-drain still fails the other fast. --- zig/src/fabric/tp/fate_test.zig | 95 +++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig index 6019055f50..728f4bfdf3 100644 --- a/zig/src/fabric/tp/fate_test.zig +++ b/zig/src/fabric/tp/fate_test.zig @@ -144,3 +144,98 @@ test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); } + +// -- a server's drain on two ranks: the plan link and the fate channel in the order a family's stop uses them -- + +const planlink = @import("planlink.zig"); + +/// Plan words a family's rank 0 sends: a round's operation, and a stop word that ends a follower's loop. +const op_round: i64 = 1; +const op_stop: i64 = 4; + +fn linkPair(links: *[2]planlink.PlanLink) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [planlink.max_ranks]sock.Fd = @splat(-1); + var f1: [planlink.max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + links[0] = planlink.PlanLink.init(0, 2, f0); + links[1] = planlink.PlanLink.init(1, 2, f1); +} + +/// Rank 1's loop: each round as it comes, until the stop word, or until the link ends. +const Follower = struct { + link: *planlink.PlanLink, + rounds: u32 = 0, + ended: enum { running, stopped, peer_closed, failed } = .running, + + fn run(f: *Follower) void { + var buf: [16]i64 = undefined; + while (true) { + const msg = f.link.recv(&buf) catch |e| { + f.ended = if (e == error.PeerClosed) .peer_closed else .failed; + return; + }; + if (msg.len == 0) continue; + if (msg[0] == op_stop) { + f.ended = .stopped; + return; + } + if (msg[0] == op_round) f.rounds += 1; + } + } +}; + +test "a drain on rank 0, then its clean stop at a round boundary: rank 1 ends its loop there and nobody exits 70" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + var follower: Follower = .{ .link = &links[1] }; + const t = try std.Thread.spawn(.{}, Follower.run, .{&follower}); + // serving, then draining (rank 0 refuses new work, rounds go on; the fate channel is untouched), past a grace + var sent: u32 = 0; + for (0..30) |_| { + try links[0].send(&.{ op_round, sent }); + sent += 1; + sock.sleepNs(5 * std.time.ns_per_ms); + } + try std.testing.expect(!fates[0].closing.load(.acquire)); + // the stop's order: the stop word first, so the follower ends its loop before the fate STOP and the hang-up + try links[0].send(&.{op_stop}); + fates[0].announceStop(); + links[0].close(); + fates[0].close(); + t.join(); + // rank 1's loop returned: its model closes and the process exits 0 + fates[1].close(); + links[1].close(); + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + try std.testing.expectEqual(.stopped, follower.ended); + try std.testing.expectEqual(sent, follower.rounds); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +test "a drain on rank 0 does not disarm fail-fast: rank 1 killed mid-drain still ends rank 0 (exit 70)" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + defer links[0].close(); + for (0..5) |i| try links[0].send(&.{ op_round, @intCast(i) }); + // rank 0 is draining (no stop announced); rank 1 is killed: its sockets close with no word + links[1].close(); + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); + try std.testing.expectEqual(@as(u32, 1), exits.load(.acquire)); +} From 62459b116a951750487551dd9ca3b6e55cd0d778 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:42:56 +0700 Subject: [PATCH 6/7] Session memory for paged-KV families: page pool, prefix trie, RAM residency budget, streamed NVMe tier. zig/src/sessions keeps a chat's KV pages after its request ends: entries are found by a trie over chained page digests, held in RAM within a residency budget (chats and pool bytes), and the least recently used idle chat is parked to an NVMe tier and streamed back on resume. Ranks that split KV run every store call in the same order, so their books stay identical. It complements core/prompt_cache, which snapshots a whole family state; a family that uses it reports the resident bytes as Memory.cache and its budget as Info.prompt_cache_plan. No family uses it yet; zig build test-sessions runs the host tests, the NVMe tier on real files. --- build.zig | 2 + zig/build/sessions.zig | 13 + zig/src/sessions/disk.zig | 587 ++++++++++++++++++++++++++++++++ zig/src/sessions/disk_test.zig | 125 +++++++ zig/src/sessions/format.zig | 197 +++++++++++ zig/src/sessions/freeset.zig | 114 +++++++ zig/src/sessions/jobs.zig | 150 ++++++++ zig/src/sessions/lanes.zig | 336 ++++++++++++++++++ zig/src/sessions/pagestore.zig | 129 +++++++ zig/src/sessions/pool.zig | 475 ++++++++++++++++++++++++++ zig/src/sessions/prefix.zig | 457 +++++++++++++++++++++++++ zig/src/sessions/sessions.zig | 33 ++ zig/src/sessions/store.zig | 570 +++++++++++++++++++++++++++++++ zig/src/sessions/store_test.zig | 216 ++++++++++++ 14 files changed, 3404 insertions(+) create mode 100644 zig/build/sessions.zig create mode 100644 zig/src/sessions/disk.zig create mode 100644 zig/src/sessions/disk_test.zig create mode 100644 zig/src/sessions/format.zig create mode 100644 zig/src/sessions/freeset.zig create mode 100644 zig/src/sessions/jobs.zig create mode 100644 zig/src/sessions/lanes.zig create mode 100644 zig/src/sessions/pagestore.zig create mode 100644 zig/src/sessions/pool.zig create mode 100644 zig/src/sessions/prefix.zig create mode 100644 zig/src/sessions/sessions.zig create mode 100644 zig/src/sessions/store.zig create mode 100644 zig/src/sessions/store_test.zig diff --git a/build.zig b/build.zig index 30659f7b7b..9940d84bca 100644 --- a/build.zig +++ b/build.zig @@ -5,6 +5,7 @@ const dist_build = @import("zig/build/dist.zig"); const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const tp_build = @import("zig/build/tp.zig"); +const sessions_build = @import("zig/build/sessions.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -174,6 +175,7 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); tp_build.add(b, target, optimize, test_step); + sessions_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig new file mode 100644 index 0000000000..a5d8ea6f3c --- /dev/null +++ b/zig/build/sessions.zig @@ -0,0 +1,13 @@ +//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS. + +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/src/sessions/disk.zig b/zig/src/sessions/disk.zig new file mode 100644 index 0000000000..f5ab58f06f --- /dev/null +++ b/zig/src/sessions/disk.zig @@ -0,0 +1,587 @@ +//! The NVMe tier: a file an entry under //rank/, written via .tmp + rename, streamed via the lanes. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const lanes_mod = @import("lanes.zig"); +const Lanes = lanes_mod.Lanes; +const Staging = lanes_mod.Staging; +const Op = lanes_mod.Op; +const up = lanes_mod.up; +const format = @import("format.zig"); +const Layout = format.Layout; +const Head = format.Head; +const Seg = format.Seg; +const PageStore = @import("pagestore.zig").PageStore; +const pool_mod = @import("pool.zig"); +const Family = pool_mod.Family; +const Digest = @import("prefix.zig").Digest; + +pub const none = std.math.maxInt(u32); + +pub const Options = struct { + root: []const u8, + compat: Digest, + rank: u32 = 0, + world: u32 = 1, + /// bytes of files the tier keeps (least recently used files without dependents go first) + budget: u64 = 64 << 30, + /// entries shorter than this are dropped, not parked + min_tokens: u64 = 1024, + /// O_DIRECT where the filesystem takes it (buffered with the page cache dropped otherwise) + direct: bool = true, + /// staging buffers and their size: the host transient of any park or restore + staging: u32 = 16, + chunk: u32 = 8 << 20, + lanes: u32 = 4, + /// delta files a chain may stack before a park writes the entry whole + max_chain: u32 = 8, +}; + +/// A piece of bounded state (rings, carries, taps): small, host bytes. +pub const Blob = struct { + name: []const u8, + bytes: []const u8, +}; + +/// What a park writes: the entry, its tokens from first_page on, and this rank's pages per family, in logical order. +pub const WriteSrc = struct { + key: Digest, + /// the file holding pages [0, first_page) (a delta), or null (a whole entry: first_page 0) + base: ?Digest = null, + tag: u32, + kind: u32 = 0, + page: u32, + pos: u64, + first_page: u32 = 0, + npages: u32, + /// tokens [first_page x page, pos) + ids: []const i32, + store: PageStore, + fence: u64 = 0, + families: []const Family, + /// per family: family-tensor pages of this rank's logical pages [first_page, npages), in logical order + pages: []const []const u32, + blobs: []const Blob = &.{}, +}; + +/// One file of a restore: read logical pages [lo, hi) of it into the slot. +pub const Extent = struct { + key: Digest, + first_page: u32, + lo: u32, + hi: u32, +}; + +/// Where a restore puts rows: per family, the family-tensor page of each of the slot's logical pages. +pub const ReadSink = struct { + store: PageStore, + families: []const Family, + page: u32, + world: u32, + rank: u32, + /// per family: family-tensor pages for logical pages [0, npages) + pages: []const []const u32, +}; + +/// The bounded state a restore returns (owned). +pub const Restored = struct { + blobs: []Blob, + + pub fn deinit(r: *Restored, gpa: std.mem.Allocator) void { + for (r.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(r.blobs); + } +}; + +/// The index record of one file. +pub const Rec = struct { + key: Digest, + base: u32 = none, + tag: u32, + kind: u32, + pos: u64, + first_page: u32, + npages: u32, + size: u64, + /// files stacked on this one + refs: u32 = 0, + /// files under it, this one included + depth: u32 = 1, + older: u32 = none, + newer: u32 = none, + live: bool = true, +}; + +pub const Stats = struct { + writes: u64 = 0, + reads: u64 = 0, + bytes_written: u64 = 0, + bytes_read: u64 = 0, + bad: u64 = 0, + evicted: u64 = 0, +}; + +pub const Disk = struct { + gpa: std.mem.Allocator, + io: Io, + opts: Options, + dir: Io.Dir, + staging: Staging, + lanes: *Lanes, + recs: std.ArrayList(Rec) = .empty, + free_recs: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// least recently used first + oldest: u32 = none, + newest: u32 = none, + used: u64 = 0, + stats: Stats = .{}, + + pub fn open(gpa: std.mem.Allocator, io: Io, opts: Options) !*Disk { + const d = try gpa.create(Disk); + errdefer gpa.destroy(d); + var name_buf: [64]u8 = undefined; + const sub = try std.fmt.bufPrint(&name_buf, "{x}/rank{d}", .{ opts.compat[0..8], opts.rank }); + var root = try Io.Dir.cwd().createDirPathOpen(io, opts.root, .{}); + defer root.close(io); + const dir = try root.createDirPathOpen(io, sub, .{ .open_options = .{ .iterate = true } }); + errdefer dir.close(io); + d.* = .{ .gpa = gpa, .io = io, .opts = opts, .dir = dir, .staging = try Staging.init(gpa, io, @max(2, opts.staging), up(opts.chunk)), .lanes = undefined }; + errdefer d.staging.deinit(); + d.lanes = try Lanes.start(gpa, io, &d.staging, @max(1, opts.lanes)); + return d; + } + + pub fn close(d: *Disk) void { + d.lanes.shutdown(); + d.staging.deinit(); + d.dir.close(d.io); + d.recs.deinit(d.gpa); + d.free_recs.deinit(d.gpa); + d.by_key.deinit(d.gpa); + const gpa = d.gpa; + gpa.destroy(d); + } + + pub fn accepts(d: *const Disk, pos: u64) bool { + return pos >= d.opts.min_tokens; + } + + pub fn get(d: *Disk, key: Digest) ?*Rec { + const i = d.by_key.get(key) orelse return null; + return &d.recs.items[i]; + } + + fn fileName(key: Digest, ext: []const u8, buf: *[48]u8) [:0]const u8 { + const n = (std.fmt.bufPrint(buf[0..47], "{x}.{s}", .{ &key, ext }) catch unreachable).len; + buf[n] = 0; + return buf[0..n :0]; + } + + fn openFile(d: *Disk, name: [:0]const u8, write: bool) !struct { fd: linux.fd_t, direct: bool } { + var flags: linux.O = .{ .ACCMODE = if (write) .WRONLY else .RDONLY, .CLOEXEC = true }; + if (write) { + flags.CREAT = true; + flags.TRUNC = true; + } + if (d.opts.direct) { + var f2 = flags; + f2.DIRECT = true; + const rc = linux.openat(d.dir.handle, name, f2, 0o644); + if (linux.errno(rc) == .SUCCESS) return .{ .fd = @intCast(rc), .direct = true }; + } + const rc = linux.openat(d.dir.handle, name, flags, 0o644); + if (linux.errno(rc) != .SUCCESS) return if (linux.errno(rc) == .NOENT) error.FileNotFound else error.OpenFailed; + return .{ .fd = @intCast(rc), .direct = false }; + } + + /// The layout a WriteSrc gets: segments, chunk sizes (whole pages a chunk for pool segments). + pub fn plan(d: *Disk, src: *const WriteSrc) !Layout { + var segs: std.ArrayList(Seg) = .empty; + defer segs.deinit(d.gpa); + const cap = d.opts.chunk; + try segs.append(d.gpa, .{ .name = format.segName("ids"), .kind = .ids, .family = 0, .nbytes = src.ids.len * 4, .chunk = cap }); + for (src.families, 0..) |f, i| { + const pb: u32 = @intCast(f.pageBytes(src.page)); + const k = @max(1, cap / pb); + if (@as(u64, k) * pb > d.staging.size) return error.PageTooLarge; + try segs.append(d.gpa, .{ .name = format.segName(f.name), .kind = .pool, .family = @intCast(i), .nbytes = @as(u64, src.pages[i].len) * pb, .chunk = k * pb }); + } + for (src.blobs) |b| try segs.append(d.gpa, .{ .name = format.segName(b.name), .kind = .blob, .family = 0, .nbytes = b.bytes.len, .chunk = cap }); + const head: Head = .{ .key = src.key, .base = src.base orelse @splat(0), .compat = d.opts.compat, .tag = src.tag, .kind = src.kind, .page = src.page, .world = d.opts.world, .rank = d.opts.rank, .first_page = src.first_page, .npages = src.npages, .nseg = 0, .ndigest = 0, .pos = src.pos }; + return Layout.build(d.gpa, head, segs.items); + } + + /// Writes an entry's file streamed (the pump copies pages to staging, lanes hash and write); returns its bytes. + pub fn writeFile(d: *Disk, src: *const WriteSrc) !u64 { + var lay = try d.plan(src); + defer lay.deinit(d.gpa); + var tb: [48]u8 = undefined; + var fb: [48]u8 = undefined; + const tmp = fileName(src.key, "tmp", &tb); + const f = try d.openFile(tmp, true); + var ok = false; + defer { + _ = linux.close(f.fd); + if (!ok) _ = linux.unlinkat(d.dir.handle, tmp, 0); + } + try src.store.waitFence(src.fence); + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + for (lay.segs) |*s| { + for (0..s.nchunks) |ci| { + if (op.failed()) |e| return e; + const c = s.chunkAt(@intCast(ci)); + const sb = try d.staging.get(); + errdefer d.staging.put(sb); + const buf = d.staging.buf(sb)[0..c.data]; + const at = @as(u64, ci) * s.chunk; + switch (s.kind) { + .ids => @memcpy(buf, std.mem.sliceAsBytes(src.ids)[at..][0..c.data]), + .blob => @memcpy(buf, blobOf(src, s).bytes[at..][0..c.data]), + .pool => { + const pb: u32 = @intCast(src.families[s.family].pageBytes(src.page)); + const p0: usize = @intCast(at / pb); + try src.store.read(s.family, src.pages[s.family][p0..][0 .. c.data / pb], buf); + }, + _ => unreachable, + } + d.lanes.submit(.{ .kind = .write, .fd = f.fd, .staging = sb, .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + ci], .op = &op }); + } + } + try op.wait(); + const hb = try d.staging.get(); + defer d.staging.put(hb); + if (lay.head_span > d.staging.size) return error.HeaderTooLarge; + lay.encode(d.staging.buf(hb)); + try lanes_mod.pwriteAll(f.fd, d.staging.buf(hb)[0..lay.head_span], 0); + if (linux.errno(linux.fdatasync(f.fd)) != .SUCCESS) return error.SyncFailed; + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + if (linux.errno(linux.renameat(d.dir.handle, tmp, d.dir.handle, fileName(src.key, "tfs", &fb))) != .SUCCESS) return error.RenameFailed; + ok = true; + return lay.head.total; + } + + fn blobOf(src: *const WriteSrc, s: *const Seg) Blob { + for (src.blobs) |b| if (std.mem.eql(u8, b.name, s.nameOf())) return b; + unreachable; + } + + /// A file's header, checked. + pub fn readHead(d: *Disk, fd: linux.fd_t) !Layout { + const hb = try d.staging.get(); + defer d.staging.put(hb); + const buf = d.staging.buf(hb); + try lanes_mod.preadAll(fd, buf[0..4096], 0); + const hn = std.mem.readInt(u32, buf[12..16], .little); + if (up(hn) > buf.len) return error.BadHeader; + if (hn > 4096) try lanes_mod.preadAll(fd, buf[4096..up(hn)], 4096); + return Layout.decode(d.gpa, buf[0..up(hn)]); + } + + /// Reads pages [lo, hi) of one restore file into the sink, each chunk checked first; blobs too when want_blobs. + pub fn readFile(d: *Disk, ext: Extent, sink: *const ReadSink, want_blobs: bool) !?Restored { + var nb: [48]u8 = undefined; + const f = try d.openFile(fileName(ext.key, "tfs", &nb), false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + if (!std.mem.eql(u8, &lay.head.key, &ext.key) or lay.head.first_page != ext.first_page or lay.head.page != sink.page) return error.BadHeader; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS or st.size < lay.head.total) return error.Truncated; + var out: std.ArrayList(Blob) = .empty; + errdefer { + for (out.items) |b| { + d.gpa.free(b.name); + d.gpa.free(b.bytes); + } + out.deinit(d.gpa); + } + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + // (segment, chunk) jobs; the pump keeps at most staging - 1 in flight (one buffer for the header path) + var jobs: std.ArrayList(struct { seg: u32, chunk: u32 }) = .empty; + defer jobs.deinit(d.gpa); + for (lay.segs, 0..) |*s, si| switch (s.kind) { + .pool => { + if (s.family >= sink.families.len) return error.BadHeader; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + if (s.chunk % pb != 0 or !std.mem.eql(u8, s.nameOf(), fam.name)) return error.BadHeader; + const need = ownedIn(fam, sink, ext.first_page, ext.hi) - ownedIn(fam, sink, ext.first_page, ext.lo); + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + if (s.nbytes < (@as(u64, skip) + need) * pb) return error.BadHeader; + const kpp = s.chunk / pb; + const c0 = skip / kpp; + const c1 = if (need == 0) c0 else (skip + need + kpp - 1) / kpp; + for (c0..c1) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + .blob => if (want_blobs) { + const buf = try d.gpa.alloc(u8, s.nbytes); + errdefer d.gpa.free(buf); + const name = try d.gpa.dupe(u8, s.nameOf()); + errdefer d.gpa.free(name); + try out.append(d.gpa, .{ .name = name, .bytes = buf }); + for (0..s.nchunks) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + else => {}, + }; + const window = d.staging.count - 1; + var sent: usize = 0; + var seen: usize = 0; + var bytes: u64 = 0; + while (seen < jobs.items.len) { + while (sent < jobs.items.len and sent - seen < window and op.failed() == null) : (sent += 1) { + const jb = jobs.items[sent]; + const s = &lay.segs[jb.seg]; + const c = s.chunkAt(jb.chunk); + d.lanes.submit(.{ .kind = .read, .fd = f.fd, .staging = try d.staging.get(), .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + jb.chunk], .op = &op, .tag = @intCast(sent) }); + bytes += c.span; + } + const j = op.next() orelse break; + seen += 1; + defer d.staging.put(j.staging); + if (op.failed() != null) continue; + const jb = jobs.items[j.tag]; + const s = &lay.segs[jb.seg]; + const data = d.staging.buf(j.staging)[0..j.data]; + if (s.kind == .blob) { + for (out.items) |b| if (std.mem.eql(u8, b.name, s.nameOf())) @memcpy(@constCast(b.bytes)[@as(u64, jb.chunk) * s.chunk ..][0..data.len], data); + continue; + } + try d.drain(s, jb.chunk, data, ext, sink); + } + try op.wait(); + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + d.stats.bytes_read += bytes; + if (!want_blobs) { + out.deinit(d.gpa); + return null; + } + return .{ .blobs = try out.toOwnedSlice(d.gpa) }; + } + + /// A checked pool chunk's pages that fall in [lo, hi) into the sink's pages. + fn drain(d: *Disk, s: *const Seg, chunk: u32, data: []const u8, ext: Extent, sink: *const ReadSink) !void { + _ = d; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + const kpp = s.chunk / pb; + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + const end = ownedIn(fam, sink, ext.first_page, ext.hi); + const a = @max(chunk * kpp, skip); + const b = @min(chunk * kpp + @as(u32, @intCast(data.len / pb)), end); + if (a >= b) return; + var list: [512]u32 = undefined; + var n: u32 = 0; + var j = a; + while (j < b) : (j += 1) { + list[n] = sink.pages[s.family][logicalOf(fam, sink, ext.first_page, j)]; + n += 1; + if (n == list.len or j + 1 == b) { + const first = j + 1 - n; + try sink.store.write(s.family, list[0..n], data[(first - chunk * kpp) * pb ..][0 .. n * pb]); + n = 0; + } + } + } + + /// Records a written file in the index (round thread, at the plan's point); its base gains a dependent. + pub fn register(d: *Disk, src: *const WriteSrc, size: u64) !u32 { + const id = try d.add(src.key, src.base, src.tag, src.kind, src.pos, src.first_page, src.npages, size); + d.stats.writes += 1; + d.stats.bytes_written += size; + return id; + } + + fn add(d: *Disk, key: Digest, base: ?Digest, tag: u32, kind: u32, pos: u64, first_page: u32, npages: u32, size: u64) !u32 { + var r: Rec = .{ .key = key, .tag = tag, .kind = kind, .pos = pos, .first_page = first_page, .npages = npages, .size = size }; + if (base) |b| { + const bi = d.by_key.get(b) orelse return error.NoBase; + r.base = bi; + r.depth = d.recs.items[bi].depth + 1; + } + const id: u32 = if (d.free_recs.pop()) |i| i else blk: { + try d.recs.append(d.gpa, undefined); + break :blk @intCast(d.recs.items.len - 1); + }; + try d.by_key.put(d.gpa, key, id); + if (r.base != none) d.recs.items[r.base].refs += 1; + d.recs.items[id] = r; + d.link(id); + d.used += size; + return id; + } + + /// A valid file found on restart: its header, size and verified ids (tokens [first_page x page, pos), owned). + pub const Found = struct { + head: Head, + size: u64, + mtime: i64, + ids: []i32, + + fn older(_: void, a: Found, b: Found) bool { + return a.mtime < b.mtime or (a.mtime == b.mtime and a.head.first_page < b.head.first_page); + } + }; + + /// Indexes the directory on restart: drops .tmp and bad files, checks ids, registers bases before their deltas. + pub fn scan(d: *Disk, out: *std.ArrayList(Found)) !void { + var it = d.dir.iterate(); + var nb: [48]u8 = undefined; + while (try it.next(d.io)) |ent| { + if (ent.name.len >= nb.len) continue; + @memcpy(nb[0..ent.name.len], ent.name); + nb[ent.name.len] = 0; + const name = nb[0..ent.name.len :0]; + if (std.mem.endsWith(u8, name, ".tmp")) { + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + } + if (!std.mem.endsWith(u8, name, ".tfs")) continue; + const f = d.found(name) catch { + d.stats.bad += 1; + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + }; + try out.append(d.gpa, f); + } + std.mem.sort(Found, out.items, {}, Found.older); + var i: usize = 0; + while (i < out.items.len) { + const f = out.items[i]; + const base: ?Digest = if (f.head.hasBase()) f.head.base else null; + if (base != null and d.by_key.get(base.?) == null) { + d.stats.bad += 1; + var kb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(f.head.key, "tfs", &kb), 0); + d.gpa.free(f.ids); + _ = out.orderedRemove(i); + continue; + } + _ = try d.add(f.head.key, base, f.head.tag, f.head.kind, f.head.pos, f.head.first_page, f.head.npages, f.size); + i += 1; + } + } + + fn found(d: *Disk, name: [:0]const u8) !Found { + const f = try d.openFile(name, false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + const h = lay.head; + var kb: [48]u8 = undefined; + if (!std.mem.eql(u8, fileName(h.key, "tfs", &kb), name)) return error.BadHeader; + if (!std.mem.eql(u8, &h.compat, &d.opts.compat) or h.rank != d.opts.rank or h.world != d.opts.world) return error.Foreign; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true, .MTIME = true }, &st)) != .SUCCESS or st.size < h.total) return error.Truncated; + const seg = lay.find(.ids, "ids") orelse return error.BadHeader; + if (seg.nbytes % 4 != 0 or seg.nbytes / 4 != h.pos - @as(u64, h.first_page) * h.page) return error.BadHeader; + const ids = try d.gpa.alloc(i32, seg.nbytes / 4); + errdefer d.gpa.free(ids); + const sb = try d.staging.get(); + defer d.staging.put(sb); + for (0..seg.nchunks) |ci| { + const c = seg.chunkAt(@intCast(ci)); + const buf = d.staging.buf(sb); + try lanes_mod.preadAll(f.fd, buf[0..c.span], c.off); + var dg: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(buf[0..c.data], &dg, .{}); + if (!std.mem.eql(u8, &dg, &lay.digests[seg.first_digest + ci])) return error.ChecksumMismatch; + @memcpy(std.mem.sliceAsBytes(ids)[@as(usize, ci) * seg.chunk ..][0..c.data], buf[0..c.data]); + } + return .{ .head = h, .size = st.size, .mtime = st.mtime.sec * std.time.ns_per_s + st.mtime.nsec, .ids = ids }; + } + + fn link(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + r.older = d.newest; + r.newer = none; + if (d.newest != none) d.recs.items[d.newest].newer = id else d.oldest = id; + d.newest = id; + } + + fn unlink(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + if (r.older != none) d.recs.items[r.older].newer = r.newer else d.oldest = r.newer; + if (r.newer != none) d.recs.items[r.newer].older = r.older else d.newest = r.older; + } + + /// Marks a file (and its bases) used now: O(chain). + pub fn touch(d: *Disk, key: Digest) void { + var id = d.by_key.get(key) orelse return; + while (id != none) : (id = d.recs.items[id].base) { + d.unlink(id); + d.link(id); + } + // bases end up newest: walk order keeps the entry itself older than its bases, so trims reach dependents first + } + + /// The files a restore of key reads, base first, with the logical pages each gives. + pub fn extents(d: *Disk, key: Digest, out: *std.ArrayList(Extent)) !void { + out.clearRetainingCapacity(); + var id = d.by_key.get(key) orelse return error.FileNotFound; + var hi = d.recs.items[id].npages; + while (id != none) { + const r = &d.recs.items[id]; + try out.append(d.gpa, .{ .key = r.key, .first_page = r.first_page, .lo = r.first_page, .hi = hi }); + hi = r.first_page; + id = r.base; + } + std.mem.reverse(Extent, out.items); + } + + /// Deletes a file without dependents (error when one stacks on it). + pub fn remove(d: *Disk, key: Digest) !void { + const id = d.by_key.get(key) orelse return; + const r = &d.recs.items[id]; + if (r.refs > 0) return error.HasDependents; + if (r.base != none) d.recs.items[r.base].refs -= 1; + d.unlink(id); + _ = d.by_key.remove(key); + d.used -= r.size; + r.live = false; + var nb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(key, "tfs", &nb), 0); + try d.free_recs.append(d.gpa, id); + } + + /// Deletes least recently used files without dependents while over budget; their keys go to out. + pub fn trim(d: *Disk, out: *std.ArrayList(Digest)) !void { + var id = d.oldest; + while (d.used > d.opts.budget and id != none) { + const next = d.recs.items[id].newer; + if (d.recs.items[id].refs == 0) { + const key = d.recs.items[id].key; + try d.remove(key); + try out.append(d.gpa, key); + d.stats.evicted += 1; + id = d.oldest; // a removed file may free its base + continue; + } + id = next; + } + } +}; + +/// Owned (this rank's) logical pages of family f in [first, k): the file's index of logical page k. +pub fn ownedIn(f: Family, sink: *const ReadSink, first: u32, k: u32) u32 { + if (!(f.split and sink.world > 1)) return k - first; + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, k, sink.world, &by); + return by[sink.rank]; +} + +/// The logical page of a family's j-th stored page of a file that starts at first. +pub fn logicalOf(f: Family, sink: *const ReadSink, first: u32, j: u32) u32 { + if (!(f.split and sink.world > 1)) return first + j; + const w = sink.world; + const r0 = first + (sink.rank + w - first % w) % w; // the first owned logical page at or after first + return r0 + j * w; +} diff --git a/zig/src/sessions/disk_test.zig b/zig/src/sessions/disk_test.zig new file mode 100644 index 0000000000..6ac0525e45 --- /dev/null +++ b/zig/src/sessions/disk_test.zig @@ -0,0 +1,125 @@ +//! The NVMe tier on real files: whole and delta files round-trip bit for bit, split ranks, damage. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk = @import("disk.zig"); + +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Fills a page of every family with bytes that name (family, logical page, byte). +fn fill(h: *HostPool, p: *const Pool, logical: u32, pg: u32, salt: u8) void { + for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pg)) continue; + const v = h.pageOf(@intCast(fi), p.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 131 +% logical *% 7 +% fi *% 29 +% salt); + } +} + +fn familyPages(gpa: std.mem.Allocator, p: *const Pool, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const out = try gpa.alloc([]u32, fams.len); + for (fams, 0..) |f, fi| { + var l: std.ArrayList(u32) = .empty; + for (pages[first..], first..) |pg, k| { + if (f.split and p.world > 1 and owned_only and k % p.world != p.rank) continue; + try l.append(gpa, p.familyPage(f, pg)); + } + out[fi] = try l.toOwnedSlice(gpa); + } + return out; +} + +fn freePages(gpa: std.mem.Allocator, x: [][]u32) void { + for (x) |l| gpa.free(l); + gpa.free(x); +} + +fn roundTrip(world: u32, rank: u32) !void { + const gpa = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(root); + var p = try Pool.init(gpa, .{ .families = &fams, .page = 256 }, 64 * 256, world, rank); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + // small chunks so pool segments span many chunks; few staging buffers so the pump waits on the lanes + const d = try disk.Disk.open(gpa, io, .{ .root = root, .compat = @splat(5), .rank = rank, .world = world, .chunk = 3 * 74_752, .staging = 3, .lanes = 3, .min_tokens = 1 }); + defer d.close(); + const a = try p.newSlot(32 * 256); + const pos: u64 = 9 * 256 + 100; + try a.ensure(pos); + for (a.mapped(), 0..) |pg, k| fill(&h, &p, @intCast(k), pg, 1); + const ids = try gpa.alloc(i32, pos); + defer gpa.free(ids); + for (ids, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + var blob: [5000]u8 = undefined; + for (&blob, 0..) |*x, i| x.* = @truncate(9 - i % 5); + const fp = try familyPages(gpa, &p, a.mapped(), 0, true); + defer freePages(gpa, fp); + const src: disk.WriteSrc = .{ .key = @splat(1), .tag = 4, .page = 256, .pos = pos, .npages = a.len, .ids = ids, .store = h.store(), .families = &fams, .pages = fp, .blobs = &.{.{ .name = "rings", .bytes = &blob }} }; + const size = try d.writeFile(&src); + _ = try d.register(&src, size); + // a newer turn stacks a delta on it: pages [9, 14) only (page 9 was partial in the base) + const pos2: u64 = 13 * 256 + 7; + try a.ensure(pos2); + for (a.mapped()[9..], 9..) |pg, k| fill(&h, &p, @intCast(k), pg, 2); + const ids2 = try gpa.alloc(i32, pos2); + defer gpa.free(ids2); + for (ids2, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + const fp2 = try familyPages(gpa, &p, a.mapped(), 9, true); + defer freePages(gpa, fp2); + const src2: disk.WriteSrc = .{ .key = @splat(2), .base = @splat(1), .tag = 4, .page = 256, .pos = pos2, .first_page = 9, .npages = a.len, .ids = ids2[9 * 256 ..], .store = h.store(), .families = &fams, .pages = fp2, .blobs = &.{.{ .name = "rings", .bytes = blob[0..77] }} }; + const size2 = try d.writeFile(&src2); + _ = try d.register(&src2, size2); + try testing.expect(size2 < size); + try testing.expectError(error.HasDependents, d.remove(@splat(1))); + // restore the delta onto other pages through its chain + const b = try p.newSlot(32 * 256); + try b.ensure(pos2); + var ext: std.ArrayList(disk.Extent) = .empty; + defer ext.deinit(gpa); + try d.extents(@splat(2), &ext); + try testing.expectEqual(@as(usize, 2), ext.items.len); + try testing.expectEqual(@as(u32, 9), ext.items[0].hi); + const tp = try familyPages(gpa, &p, b.mapped(), 0, false); + defer freePages(gpa, tp); + const sink: disk.ReadSink = .{ .store = h.store(), .families = &fams, .page = 256, .world = world, .rank = rank, .pages = tp }; + try testing.expect(try d.readFile(ext.items[0], &sink, false) == null); + var got = (try d.readFile(ext.items[1], &sink, true)).?; + defer got.deinit(gpa); + try testing.expectEqualSlices(u8, blob[0..77], got.blobs[0].bytes); + for (a.mapped(), b.mapped()) |pa, pb| for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pa)) continue; + try testing.expectEqualSlices(u8, h.pageOf(@intCast(fi), p.familyPage(f, pa)), h.pageOf(@intCast(fi), p.familyPage(f, pb))); + }; + // damage one chunk: the read refuses it before or as its rows land + var nb: [48]u8 = undefined; + const key2: [16]u8 = @splat(2); + const name = try std.fmt.bufPrint(&nb, "{x}.tfs", .{&key2}); + const f = try d.dir.openFile(io, name, .{ .mode = .read_write }); + var one = [_]u8{0}; + _ = try f.readPositional(io, &.{&one}, 3 * 4096 + 100); + one[0] ^= 0xff; + try f.writePositionalAll(io, &one, 3 * 4096 + 100); + f.close(io); + try testing.expectError(error.ChecksumMismatch, d.readFile(ext.items[1], &sink, true)); +} + +test "whole and delta files round-trip bit for bit (replicated)" { + try roundTrip(1, 0); +} + +test "split ranks write and restore their own pages (rank 0 and 1 of 2)" { + try roundTrip(2, 0); + try roundTrip(2, 1); +} diff --git a/zig/src/sessions/format.zig b/zig/src/sessions/format.zig new file mode 100644 index 0000000000..2ac391d156 --- /dev/null +++ b/zig/src/sessions/format.zig @@ -0,0 +1,197 @@ +//! A session file (format 2): a header, then segments of 4 KiB-aligned chunks, each with its own SHA-256. + +const std = @import("std"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const up = @import("lanes.zig").up; + +pub const magic = "TFSESS02".*; +pub const version: u32 = 2; + +pub const SegKind = enum(u32) { ids = 1, pool = 2, blob = 3, _ }; + +pub const Head = extern struct { + magic: [8]u8 = magic, + version: u32 = version, + /// bytes of Head + segments + digests (before the pad) + head_bytes: u32 = 0, + /// SHA-256 of those bytes with this field zero + head_sha: [32]u8 = @splat(0), + key: [16]u8, + /// the entry whose files hold logical pages [0, first_page); zero: none + base: [16]u8 = @splat(0), + /// what makes these bytes mean the same thing to a build (family, knobs, layout, split) + compat: [16]u8, + tag: u32, + kind: u32, + page: u32, + world: u32, + rank: u32, + first_page: u32, + npages: u32, + nseg: u32, + ndigest: u32, + _pad: u32 = 0, + pos: u64, + total: u64 = 0, + + pub fn hasBase(h: *const Head) bool { + return !std.mem.allEqual(u8, &h.base, 0); + } +}; + +pub const Seg = extern struct { + name: [40]u8, + kind: SegKind, + /// pool: the family's index in the layout + family: u32, + nbytes: u64, + offset: u64 = 0, + /// bytes a chunk (the last may be shorter) and file bytes a chunk (4 KiB multiple) + chunk: u32, + stride: u32 = 0, + first_digest: u32 = 0, + nchunks: u32 = 0, + + pub fn nameOf(s: *const Seg) []const u8 { + return std.mem.sliceTo(&s.name, 0); + } + + /// Chunk i's data bytes and file span. + pub fn chunkAt(s: *const Seg, i: u32) struct { off: u64, data: u32, span: u32 } { + const data: u32 = @intCast(@min(@as(u64, s.chunk), s.nbytes - @as(u64, i) * s.chunk)); + return .{ .off = s.offset + @as(u64, i) * s.stride, .data = data, .span = @intCast(up(data)) }; + } +}; + +pub const Digest32 = [32]u8; + +pub fn segName(name: []const u8) [40]u8 { + var n: [40]u8 = @splat(0); + @memcpy(n[0..@min(name.len, 39)], name[0..@min(name.len, 39)]); + return n; +} + +/// A file's layout: offsets, chunks and the digest slots, before or after its chunks are written. +pub const Layout = struct { + head: Head, + segs: []Seg, + digests: []Digest32, + /// file bytes of the padded header + head_span: u64, + + /// Places segments (each with name, kind, family, nbytes, chunk) after the header. + pub fn build(gpa: std.mem.Allocator, head: Head, segs_in: []const Seg) !Layout { + const segs = try gpa.dupe(Seg, segs_in); + errdefer gpa.free(segs); + var nd: u32 = 0; + for (segs) |*s| { + if (s.chunk == 0) return error.BadSegment; + s.nchunks = @intCast((s.nbytes + s.chunk - 1) / s.chunk); + s.stride = @intCast(up(s.chunk)); + s.first_digest = nd; + nd += s.nchunks; + } + var h = head; + h.nseg = @intCast(segs.len); + h.ndigest = nd; + h.head_bytes = @intCast(@sizeOf(Head) + segs.len * @sizeOf(Seg) + nd * @sizeOf(Digest32)); + const span = up(h.head_bytes); + var off = span; + for (segs) |*s| { + s.offset = off; + if (s.nchunks > 0) off += @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk); + } + h.total = off; + const d = try gpa.alloc(Digest32, nd); + @memset(d, @splat(0)); + return .{ .head = h, .segs = segs, .digests = d, .head_span = span }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.segs); + gpa.free(l.digests); + l.* = undefined; + } + + /// The header's bytes with its checksum into out (len >= head_span; the pad is zeroed). + pub fn encode(l: *Layout, out: []u8) void { + @memset(out[0..l.head_span], 0); + var h = l.head; + h.head_sha = @splat(0); + const hb = std.mem.asBytes(&h); + const sb = std.mem.sliceAsBytes(l.segs); + const db = std.mem.sliceAsBytes(l.digests); + @memcpy(out[0..hb.len], hb); + @memcpy(out[hb.len..][0..sb.len], sb); + @memcpy(out[hb.len + sb.len ..][0..db.len], db); + var sha: [32]u8 = undefined; + Sha256.hash(out[0..h.head_bytes], &sha, .{}); + @memcpy(out[@offsetOf(Head, "head_sha")..][0..32], &sha); + l.head.head_sha = sha; + } + + /// The header read back from bytes (at least head_bytes of them), every field and offset checked. + pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) !Layout { + if (bytes.len < @sizeOf(Head)) return error.BadHeader; + var h: Head = undefined; + @memcpy(std.mem.asBytes(&h), bytes[0..@sizeOf(Head)]); + if (!std.mem.eql(u8, &h.magic, &magic) or h.version != version) return error.BadHeader; + const want = @sizeOf(Head) + @as(u64, h.nseg) * @sizeOf(Seg) + @as(u64, h.ndigest) * @sizeOf(Digest32); + if (h.head_bytes != want or bytes.len < want) return error.BadHeader; + var check: [32]u8 = undefined; + var sha = Sha256.init(.{}); + sha.update(bytes[0..@offsetOf(Head, "head_sha")]); + sha.update(&@as([32]u8, @splat(0))); + sha.update(bytes[@offsetOf(Head, "head_sha") + 32 .. want]); + sha.final(&check); + if (!std.mem.eql(u8, &check, &h.head_sha)) return error.ChecksumMismatch; + const segs = try gpa.alloc(Seg, h.nseg); + errdefer gpa.free(segs); + @memcpy(std.mem.sliceAsBytes(segs), bytes[@sizeOf(Head)..][0 .. segs.len * @sizeOf(Seg)]); + const d = try gpa.alloc(Digest32, h.ndigest); + errdefer gpa.free(d); + @memcpy(std.mem.sliceAsBytes(d), bytes[@sizeOf(Head) + segs.len * @sizeOf(Seg) ..][0 .. d.len * 32]); + var nd: u64 = 0; + for (segs) |s| { + if (s.chunk == 0 or s.stride != up(s.chunk) or s.first_digest != nd or s.offset % 4096 != 0) return error.BadHeader; + if (s.nchunks != (s.nbytes + s.chunk - 1) / s.chunk) return error.BadHeader; + if (s.nchunks > 0 and s.offset + @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk) > h.total) return error.BadHeader; + nd += s.nchunks; + } + if (nd != h.ndigest) return error.BadHeader; + return .{ .head = h, .segs = segs, .digests = d, .head_span = up(h.head_bytes) }; + } + + pub fn find(l: *const Layout, kind: SegKind, name: []const u8) ?*const Seg { + for (l.segs) |*s| if (s.kind == kind and std.mem.eql(u8, s.nameOf(), name)) return s; + return null; + } +}; + +test "a layout encodes, decodes, and refuses a damaged header" { + const gpa = std.testing.allocator; + const head: Head = .{ .key = @splat(1), .compat = @splat(2), .tag = 3, .kind = 0, .page = 256, .world = 2, .rank = 1, .first_page = 4, .npages = 9, .nseg = 0, .ndigest = 0, .pos = 2200 }; + var l = try Layout.build(gpa, head, &.{ + .{ .name = segName("ids"), .kind = .ids, .family = 0, .nbytes = 4 * (2200 - 1024), .chunk = 1 << 20 }, + .{ .name = segName("comp.20"), .kind = .pool, .family = 0, .nbytes = 3 * 149_504, .chunk = 2 * 149_504 }, + .{ .name = segName("taps"), .kind = .blob, .family = 0, .nbytes = 0, .chunk = 4096 }, + }); + defer l.deinit(gpa); + try std.testing.expectEqual(@as(u32, 1 + 2 + 0), l.head.ndigest); + try std.testing.expectEqual(@as(u64, 4096), l.segs[0].offset); + try std.testing.expectEqual(@as(u64, 4096 + up(4 * 1176)), l.segs[1].offset); + const c1 = l.segs[1].chunkAt(1); + try std.testing.expectEqual(@as(u32, 149_504), c1.data); + try std.testing.expectEqual(l.segs[1].offset + up(2 * 149_504), c1.off); + try std.testing.expectEqual(l.segs[2].offset, l.head.total); + l.digests[1][5] = 9; + var buf: [8192]u8 = undefined; + l.encode(&buf); + var back = try Layout.decode(gpa, &buf); + defer back.deinit(gpa); + try std.testing.expectEqual(l.head.total, back.head.total); + try std.testing.expectEqual(@as(u8, 9), back.digests[1][5]); + try std.testing.expectEqualStrings("comp.20", back.find(.pool, "comp.20").?.nameOf()); + buf[200] ^= 0x40; + try std.testing.expectError(error.ChecksumMismatch, Layout.decode(gpa, &buf)); +} diff --git a/zig/src/sessions/freeset.zig b/zig/src/sessions/freeset.zig new file mode 100644 index 0000000000..b8cce323af --- /dev/null +++ b/zig/src/sessions/freeset.zig @@ -0,0 +1,114 @@ +//! A set of free page indices, lowest first in O(log64 n): a bit a page, then a bit a non-empty word a level up. + +const std = @import("std"); + +const max_levels = 6; // 64^6 pages: more than any pool + +pub const FreeSet = struct { + /// levels[0] holds a bit a page, levels[k] a bit a non-empty word of levels[k - 1] + levels: [max_levels][]u64 = @splat(&.{}), + depth: u8 = 0, + n: u32 = 0, + free: u32 = 0, + + /// All of 0 .. n - 1 free. + pub fn init(gpa: std.mem.Allocator, n: u32) !FreeSet { + var s: FreeSet = .{ .n = n }; + errdefer s.deinit(gpa); + var bits: usize = @max(n, 1); + while (true) { + const words = (bits + 63) / 64; + s.levels[s.depth] = try gpa.alloc(u64, words); + @memset(s.levels[s.depth], 0); + s.depth += 1; + if (words == 1) break; + bits = words; + } + for (0..n) |i| s.setBit(@intCast(i)); + s.free = n; + return s; + } + + pub fn deinit(s: *FreeSet, gpa: std.mem.Allocator) void { + for (s.levels[0..s.depth]) |l| gpa.free(l); + s.* = .{}; + } + + pub fn count(s: *const FreeSet) u32 { + return s.free; + } + + pub fn contains(s: *const FreeSet, i: u32) bool { + return s.levels[0][i / 64] & (@as(u64, 1) << @intCast(i % 64)) != 0; + } + + /// The lowest free index, taken; null when none is free. + pub fn takeLowest(s: *FreeSet) ?u32 { + if (s.free == 0) return null; + var idx: usize = 0; + var lv: usize = s.depth; + while (lv > 0) { + lv -= 1; + const w = s.levels[lv][idx]; + idx = idx * 64 + @ctz(w); + } + const i: u32 = @intCast(idx); + s.clearBit(i); + s.free -= 1; + return i; + } + + /// Takes a given free index (error when it is not free). + pub fn take(s: *FreeSet, i: u32) error{NotFree}!void { + if (i >= s.n or !s.contains(i)) return error.NotFree; + s.clearBit(i); + s.free -= 1; + } + + /// Gives an index back (error when it is already free: a double free). + pub fn give(s: *FreeSet, i: u32) error{DoubleFree}!void { + if (i >= s.n or s.contains(i)) return error.DoubleFree; + s.setBit(i); + s.free += 1; + } + + fn setBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + const was = l[idx / 64]; + l[idx / 64] = was | (@as(u64, 1) << @intCast(idx % 64)); + if (was != 0) return; // the word was non-empty: the levels above already say so + idx /= 64; + } + } + + fn clearBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + l[idx / 64] &= ~(@as(u64, 1) << @intCast(idx % 64)); + if (l[idx / 64] != 0) return; // the word still has a free bit: the levels above stay set + idx /= 64; + } + } +}; + +test "lowest first, give back, double free refused, across levels" { + const gpa = std.testing.allocator; + var s = try FreeSet.init(gpa, 70_000); // three levels + defer s.deinit(gpa); + try std.testing.expectEqual(@as(u8, 3), s.depth); + for (0..5000) |i| try std.testing.expectEqual(@as(?u32, @intCast(i)), s.takeLowest()); + try s.give(4097); + try s.give(17); + try std.testing.expectError(error.DoubleFree, s.give(17)); + try std.testing.expectEqual(@as(?u32, 17), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 4097), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 5000), s.takeLowest()); + try s.take(69_999); + try std.testing.expectError(error.NotFree, s.take(69_999)); + try std.testing.expectEqual(@as(u32, 70_000 - 5001 - 1), s.count()); + while (s.takeLowest()) |_| {} + try std.testing.expectEqual(@as(u32, 0), s.count()); + try s.give(69_999); + try std.testing.expectEqual(@as(?u32, 69_999), s.takeLowest()); +} diff --git a/zig/src/sessions/jobs.zig b/zig/src/sessions/jobs.zig new file mode 100644 index 0000000000..5b445f20e7 --- /dev/null +++ b/zig/src/sessions/jobs.zig @@ -0,0 +1,150 @@ +//! Parks and restores off the round: a tier thread runs each job's file I/O until the plan settles it (Store.settle). + +const std = @import("std"); +const Io = std.Io; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; + +pub const Kind = enum { park, restore }; +pub const State = enum(u8) { queued, running, done, failed }; + +pub const Job = struct { + kind: Kind, + /// the store's entry + entry: u32, + state: std.atomic.Value(State) = .init(.queued), + err: ?anyerror = null, + started_ns: i96 = 0, + elapsed_ns: u64 = 0, + /// park: the file to write (owned buffers below) and its size once written + src: disk_mod.WriteSrc = undefined, + size: u64 = 0, + /// restore: files base first, the sink, the bounded state read + extents: []disk_mod.Extent = &.{}, + sink: disk_mod.ReadSink = undefined, + restored: ?disk_mod.Restored = null, + /// restore: the slot it fills + slot: u32 = 0, + /// pool pages the job holds until settled + held: []u32 = &.{}, + /// owned buffers + ids: []i32 = &.{}, + lists: [][]u32 = &.{}, + blobs: []disk_mod.Blob = &.{}, + + pub fn done(j: *const Job) bool { + const s = j.state.load(.acquire); + return s == .done or s == .failed; + } + + pub fn deinit(j: *Job, gpa: std.mem.Allocator) void { + gpa.free(j.held); + gpa.free(j.ids); + for (j.lists) |l| gpa.free(l); + gpa.free(j.lists); + for (j.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(j.blobs); + gpa.free(j.extents); + if (j.restored) |*r| r.deinit(gpa); + gpa.destroy(j); + } + + fn run(j: *Job, d: *Disk, io: Io) void { + j.state.store(.running, .release); + const t0 = Io.Clock.awake.now(io).toNanoseconds(); + j.err = switch (j.kind) { + .park => if (d.writeFile(&j.src)) |n| blk: { + j.size = n; + break :blk null; + } else |e| e, + .restore => j.restore(d), + }; + j.elapsed_ns = @intCast(Io.Clock.awake.now(io).toNanoseconds() - t0); + j.state.store(if (j.err == null) .done else .failed, .release); + } + + fn restore(j: *Job, d: *Disk) ?anyerror { + for (j.extents, 0..) |e, i| { + const last = i + 1 == j.extents.len; + const r = d.readFile(e, &j.sink, last) catch |err| return err; + if (last) j.restored = r; + } + j.sink.store.publish() catch |err| return err; + return null; + } +}; + +/// The tier thread and its queue (null thread: jobs run inline). +pub const Runner = struct { + gpa: std.mem.Allocator, + io: Io, + disk: *Disk, + thread: ?std.Thread = null, + queue: std.ArrayList(*Job) = .empty, + head: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + + pub fn init(gpa: std.mem.Allocator, io: Io, d: *Disk) Runner { + return .{ .gpa = gpa, .io = io, .disk = d }; + } + + /// Starts the tier thread: jobs then run off the caller's thread. + pub fn startThread(r: *Runner) !void { + if (r.thread != null) return; + r.thread = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, loop, .{r}); + } + + pub fn deinit(r: *Runner) void { + if (r.thread) |t| { + r.mutex.lockUncancelable(r.io); + r.stop = true; + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + t.join(); + } + r.queue.deinit(r.gpa); + } + + pub fn submit(r: *Runner, j: *Job) !void { + if (r.thread == null) return j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (r.head == r.queue.items.len) { + r.queue.clearRetainingCapacity(); + r.head = 0; + } + try r.queue.append(r.gpa, j); + r.cond.signal(r.io); + } + + /// Waits for a job (the plan settles it next). + pub fn wait(r: *Runner, j: *Job) void { + while (!j.done()) { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (!j.done()) r.cond.waitTimeout(r.io, &r.mutex, .{ .duration = .{ .raw = .fromMilliseconds(5), .clock = .awake } }) catch {}; + } + } + + fn loop(r: *Runner) void { + while (true) { + const j = blk: { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + while (r.head == r.queue.items.len and !r.stop) r.cond.waitUncancelable(r.io, &r.mutex); + if (r.head == r.queue.items.len) return; + r.head += 1; + break :blk r.queue.items[r.head - 1]; + }; + j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + } + } +}; diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig new file mode 100644 index 0000000000..38fad8cd89 --- /dev/null +++ b/zig/src/sessions/lanes.zig @@ -0,0 +1,336 @@ +//! The NVMe tier's lane threads, which hash and write (or read and check) 4 KiB chunks, and their staging buffers. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const block = 4096; + +pub fn up(n: u64) u64 { + return (n + block - 1) / block * block; +} + +/// Bounded, aligned staging buffers (allocated once; `release` gives the memory back while idle). +pub const Staging = struct { + io: Io, + size: usize, + count: u32, + bufs: [][]align(block) u8, + free_list: []u32, + free_n: u32, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, io: Io, count: u32, size: usize) !Staging { + if (count == 0 or size % block != 0) return error.BadStaging; + const bufs = try gpa.alloc([]align(block) u8, count); + errdefer gpa.free(bufs); + const fl = try gpa.alloc(u32, count); + errdefer gpa.free(fl); + for (bufs, 0..) |*b, i| { + b.* = &.{}; + fl[i] = @intCast(i); + } + return .{ .io = io, .size = size, .count = count, .bufs = bufs, .free_list = fl, .free_n = count, .gpa = gpa }; + } + + pub fn deinit(s: *Staging) void { + s.release(); + s.gpa.free(s.bufs); + s.gpa.free(s.free_list); + s.* = undefined; + } + + /// A free buffer's index, waiting for one; its memory is mapped on first use. + pub fn get(s: *Staging) !u32 { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + while (s.free_n == 0) s.cond.waitUncancelable(s.io, &s.mutex); + s.free_n -= 1; + const i = s.free_list[s.free_n]; + if (s.bufs[i].len == 0) s.bufs[i] = try s.gpa.alignedAlloc(u8, .fromByteUnits(block), s.size); + return i; + } + + pub fn put(s: *Staging, i: u32) void { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + s.free_list[s.free_n] = i; + s.free_n += 1; + s.cond.signal(s.io); + } + + pub fn buf(s: *Staging, i: u32) []align(block) u8 { + return s.bufs[i]; + } + + /// Unmaps every buffer (all must be free): the memory goes back between operations. + pub fn release(s: *Staging) void { + std.debug.assert(s.free_n == s.count); + for (s.bufs) |*b| if (b.len != 0) { + s.gpa.free(b.*); + b.* = &.{}; + }; + } + + pub fn bytes(s: *const Staging) u64 { + return @as(u64, s.count) * s.size; + } +}; + +pub const Kind = enum { write, read }; + +/// One chunk's I/O: buf[0 .. span] at file offset off (span: the chunk padded to 4 KiB; data: the chunk's own bytes). +pub const Job = struct { + kind: Kind, + fd: linux.fd_t, + staging: u32, + data: u32, + span: u32, + off: u64, + /// write: the chunk's digest goes here; read: the digest it must have + digest: *[32]u8, + op: *Op, + /// the pump's tag for a finished read + tag: u32 = 0, +}; + +/// A park's or restore's chunks in flight: how many are pending, the first error, and finished reads for the pump. +pub const Op = struct { + io: Io, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + pending: u32 = 0, + err: ?anyerror = null, + /// finished reads (their jobs), drained by the pump + done: std.ArrayList(Job) = .empty, + gpa: std.mem.Allocator, + /// write jobs hand their staging back themselves + staging: *Staging, + + pub fn init(gpa: std.mem.Allocator, io: Io, staging: *Staging) !Op { + var o: Op = .{ .io = io, .gpa = gpa, .staging = staging }; + try o.done.ensureTotalCapacity(gpa, staging.count); + return o; + } + + pub fn deinit(o: *Op) void { + o.done.deinit(o.gpa); + } + + fn finish(o: *Op, j: Job, err: ?anyerror) void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + o.pending -= 1; + if (err) |e| { + if (o.err == null) o.err = e; + } + if (j.kind == .read) o.done.appendAssumeCapacity(j) else o.staging.put(j.staging); + o.cond.broadcast(o.io); + } + + /// The next finished read, waiting for one; null when none is pending. + pub fn next(o: *Op) ?Job { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.done.items.len == 0 and o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + return o.done.pop(); + } + + /// Waits for every pending job; returns the first error. + pub fn wait(o: *Op) !void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + if (o.err) |e| return e; + } + + pub fn failed(o: *Op) ?anyerror { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + return o.err; + } +}; + +/// The lane threads and their queue. +pub const Lanes = struct { + io: Io, + gpa: std.mem.Allocator, + staging: *Staging, + threads: []std.Thread, + /// queued jobs: a ring as long as the staging set (every queued job holds a buffer) + ring: []Job, + head: usize = 0, + len: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + /// seconds the lanes spent hashing and in I/O (summed over lanes) + hash_ns: std.atomic.Value(u64) = .init(0), + io_ns: std.atomic.Value(u64) = .init(0), + + pub fn start(gpa: std.mem.Allocator, io: Io, staging: *Staging, n: u32) !*Lanes { + const l = try gpa.create(Lanes); + errdefer gpa.destroy(l); + const ring = try gpa.alloc(Job, staging.count); + errdefer gpa.free(ring); + l.* = .{ .io = io, .gpa = gpa, .staging = staging, .threads = try gpa.alloc(std.Thread, n), .ring = ring }; + errdefer gpa.free(l.threads); + var made: usize = 0; + errdefer { + l.halt(); + for (l.threads[0..made]) |t| t.join(); + } + for (l.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, run, .{l}); + made += 1; + } + return l; + } + + fn halt(l: *Lanes) void { + l.mutex.lockUncancelable(l.io); + l.stop = true; + l.cond.broadcast(l.io); + l.mutex.unlock(l.io); + } + + pub fn shutdown(l: *Lanes) void { + l.halt(); + for (l.threads) |t| t.join(); + l.gpa.free(l.ring); + l.gpa.free(l.threads); + l.gpa.destroy(l); + } + + /// Queues a job (each holds a staging buffer, so the ring never overflows). + pub fn submit(l: *Lanes, j: Job) void { + { + j.op.mutex.lockUncancelable(l.io); + defer j.op.mutex.unlock(l.io); + j.op.pending += 1; + } + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + std.debug.assert(l.len < l.ring.len); + l.ring[(l.head + l.len) % l.ring.len] = j; + l.len += 1; + l.cond.signal(l.io); + } + + fn run(l: *Lanes) void { + while (true) { + const j = blk: { + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + while (l.len == 0 and !l.stop) l.cond.waitUncancelable(l.io, &l.mutex); + if (l.len == 0) return; + const j = l.ring[l.head]; + l.head = (l.head + 1) % l.ring.len; + l.len -= 1; + break :blk j; + }; + j.op.finish(j, l.work(j)); + } + } + + fn work(l: *Lanes, j: Job) ?anyerror { + const b = l.staging.buf(j.staging); + var t0 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + switch (j.kind) { + .write => { + Sha256.hash(b[0..j.data], j.digest, .{}); + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.hash_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + t0 = t1; + @memset(b[j.data..j.span], 0); + pwriteAll(j.fd, b[0..j.span], j.off) catch |e| return e; + _ = l.io_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t0), .monotonic); + }, + .read => { + preadAll(j.fd, b[0..j.span], j.off) catch |e| return e; + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.io_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + var d: [32]u8 = undefined; + Sha256.hash(b[0..j.data], &d, .{}); + _ = l.hash_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t1), .monotonic); + if (!std.mem.eql(u8, &d, j.digest)) return error.ChecksumMismatch; + }, + } + return null; + } +}; + +pub fn pwriteAll(fd: linux.fd_t, data: []const u8, off: u64) !void { + var done: usize = 0; + while (done < data.len) { + const rc = linux.pwrite(fd, data[done..].ptr, data.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + .NOSPC => return error.NoSpaceLeft, + else => return error.WriteFailed, + } + if (rc == 0) return error.WriteFailed; + done += rc; + } +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.ReadFailed, + } + if (rc == 0) return error.ShortRead; + done += rc; + } +} + +test "lanes write, read back and check chunks through bounded staging" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var st = try Staging.init(gpa, io, 3, 2 * block); + defer st.deinit(); + const l = try Lanes.start(gpa, io, &st, 2); + defer l.shutdown(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const f = try tmp.dir.createFile(io, "lanes.bin", .{ .read = true }); + defer f.close(io); + const fd = f.handle; + var digests: [7][32]u8 = undefined; + var op = try Op.init(gpa, io, &st); + defer op.deinit(); + for (0..7) |i| { + const s = try st.get(); + const n: u32 = if (i == 6) 1000 else 2 * block - 7; + for (st.buf(s)[0..n], 0..) |*b, k| b.* = @truncate(k * 31 + i); + l.submit(.{ .kind = .write, .fd = fd, .staging = s, .data = n, .span = @intCast(up(n)), .off = i * 2 * block, .digest = &digests[i], .op = &op }); + } + try op.wait(); + var rop = try Op.init(gpa, io, &st); + defer rop.deinit(); + var seen: u32 = 0; + var sent: u32 = 0; + while (seen < 7) { + while (sent < 7 and sent - seen < 3) : (sent += 1) { + const n: u32 = if (sent == 6) 1000 else 2 * block - 7; + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = n, .span = @intCast(up(n)), .off = sent * 2 * block, .digest = &digests[sent], .op = &rop, .tag = sent }); + } + const j = rop.next().?; + for (st.buf(j.staging)[0..j.data], 0..) |b, k| try std.testing.expectEqual(@as(u8, @truncate(k * 31 + j.tag)), b); + st.put(j.staging); + seen += 1; + } + try rop.wait(); + digests[2][0] ^= 1; // a damaged chunk is refused + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = 2 * block - 7, .span = 2 * block, .off = 2 * 2 * block, .digest = &digests[2], .op = &rop }); + st.put(rop.next().?.staging); + try std.testing.expectError(error.ChecksumMismatch, rop.wait()); +} diff --git a/zig/src/sessions/pagestore.zig b/zig/src/sessions/pagestore.zig new file mode 100644 index 0000000000..0e4792699b --- /dev/null +++ b/zig/src/sessions/pagestore.zig @@ -0,0 +1,129 @@ +//! The bytes behind the pool's pages: a family's tensor of pages, read and written a page run at a time. + +const std = @import("std"); +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; + +pub const PageStore = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// family fam's pages, back to back into dst (pages.len x a page's bytes) + read: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void, + /// src (pages.len x a page's bytes) into family fam's pages + write: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void, + /// one page's rows of family fam into another page (copy on write of a RAM resume's partial page) + copyPage: *const fn (ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void, + /// round thread: a point after every row the compute stream has written so far + fence: *const fn (ptr: *anyopaque) anyerror!u64, + /// before reads: wait until the fence's rows are visible + waitFence: *const fn (ptr: *anyopaque, token: u64) anyerror!void, + /// after writes: make them visible to the compute stream + publish: *const fn (ptr: *anyopaque) anyerror!void, + }; + + pub fn read(s: PageStore, fam: u32, pages: []const u32, dst: []u8) !void { + return s.vtable.read(s.ptr, fam, pages, dst); + } + pub fn write(s: PageStore, fam: u32, pages: []const u32, src: []const u8) !void { + return s.vtable.write(s.ptr, fam, pages, src); + } + pub fn copyPage(s: PageStore, fam: u32, src: u32, dst: u32) !void { + return s.vtable.copyPage(s.ptr, fam, src, dst); + } + pub fn fence(s: PageStore) !u64 { + return s.vtable.fence(s.ptr); + } + pub fn waitFence(s: PageStore, token: u64) !void { + return s.vtable.waitFence(s.ptr, token); + } + pub fn publish(s: PageStore) !void { + return s.vtable.publish(s.ptr); + } +}; + +/// Every family's pages in host memory: [familyTensorPages x page bytes] a family, zeroed (the null page stays zero). +pub const HostPool = struct { + gpa: std.mem.Allocator, + page: u32, + families: []const Family, + tensors: [][]u8, + + pub fn init(gpa: std.mem.Allocator, p: *const Pool) !HostPool { + const t = try gpa.alloc([]u8, p.layout.families.len); + errdefer gpa.free(t); + var made: usize = 0; + errdefer for (t[0..made]) |x| gpa.free(x); + for (p.layout.families, 0..) |f, i| { + t[i] = try gpa.alloc(u8, @as(usize, p.familyTensorPages(f)) * f.pageBytes(p.page)); + @memset(t[i], 0); + made += 1; + } + return .{ .gpa = gpa, .page = p.page, .families = p.layout.families, .tensors = t }; + } + + pub fn deinit(h: *HostPool) void { + for (h.tensors) |x| h.gpa.free(x); + h.gpa.free(h.tensors); + h.* = undefined; + } + + pub fn store(h: *HostPool) PageStore { + return .{ .ptr = h, .vtable = &vtable }; + } + + /// One page of family fam (a view). + pub fn pageOf(h: *HostPool, fam: u32, pg: u32) []u8 { + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + return h.tensors[fam][pg * n ..][0..n]; + } + + const vtable: PageStore.VTable = .{ .read = read, .write = write, .copyPage = copyPage, .fence = fence, .waitFence = waitFence, .publish = publish }; + + fn read(ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (dst.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(dst[i * n ..][0..n], h.pageOf(fam, pg)); + } + + fn write(ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (src.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(h.pageOf(fam, pg), src[i * n ..][0..n]); + } + + fn copyPage(ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + @memcpy(h.pageOf(fam, dst), h.pageOf(fam, src)); + } + + fn fence(_: *anyopaque) anyerror!u64 { + return 0; + } + fn waitFence(_: *anyopaque, _: u64) anyerror!void {} + fn publish(_: *anyopaque) anyerror!void {} +}; + +test "host pool pages round-trip through the interface" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 1, .row_bytes = 8 }, .{ .name = "ik", .ratio = 2, .row_bytes = 4 } }; + var p = try Pool.init(gpa, .{ .families = &fams, .page = 16 }, 64, 1, 0); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + const s = h.store(); + var src: [2 * 128]u8 = undefined; + for (&src, 0..) |*b, i| b.* = @truncate(i * 7 + 3); + try s.write(0, &.{ 3, 1 }, &src); + var dst: [2 * 128]u8 = undefined; + try s.read(0, &.{ 3, 1 }, &dst); + try std.testing.expectEqualSlices(u8, &src, &dst); + try s.copyPage(0, 3, 0); + try std.testing.expectEqualSlices(u8, src[0..128], h.pageOf(0, 0)); + try std.testing.expect(std.mem.allEqual(u8, h.pageOf(0, p.null_page), 0)); + try std.testing.expectError(error.BadLength, s.read(1, &.{0}, &dst)); +} diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig new file mode 100644 index 0000000000..d3e67516af --- /dev/null +++ b/zig/src/sessions/pool.zig @@ -0,0 +1,475 @@ +//! The paged KV pool's host book: lowest-free pages per residue mod W, holder counts, slot page tables, quotas. + +const std = @import("std"); +const FreeSet = @import("freeset.zig").FreeSet; + +pub const Error = error{ PoolExhausted, PastCapacity, SlotNotEmpty, DoubleFree, BadPool, ResidueMismatch } || std.mem.Allocator.Error; + +/// A cache family: rows of `row_bytes` every `ratio` tokens; `split` families shard by logical page over the ranks. +pub const Family = struct { + name: []const u8, + ratio: u32, + row_bytes: u32, + split: bool = false, + + pub fn pageRows(f: Family, page: u32) u32 { + return page / f.ratio; + } + + /// Bytes of one page of this family (a whole page's rows). + pub fn pageBytes(f: Family, page: u32) u64 { + return @as(u64, f.pageRows(page)) * f.row_bytes; + } +}; + +/// The pool's families and page size: what a token costs and how pages are laid out. +pub const Layout = struct { + families: []const Family, + /// tokens a page (a power of two; 256 for every family we serve) + page: u32 = 256, + + /// A rank's bytes a token: the split families count 1 / W. + pub fn bytesPerToken(l: Layout, world: u32) f64 { + var b: f64 = 0; + for (l.families) |f| b += @as(f64, @floatFromInt(f.row_bytes)) / @as(f64, @floatFromInt(f.ratio)) / @as(f64, @floatFromInt(if (f.split and world > 1) world else 1)); + return b; + } + + /// A rank's average bytes a page (split families 1 / W): what the residency budget counts. + pub fn pageBytes(l: Layout, world: u32) u64 { + var b: u64 = 0; + for (l.families) |f| b += f.pageBytes(l.page) / (if (f.split and world > 1) world else 1); + return b; + } +}; + +pub fn pagesFor(tokens: u64, page: u32) u32 { + return @intCast((tokens + page - 1) / page); +} + +/// Logical pages [k0, k1) counted by residue mod `w` (O(w)). +pub fn span(k0: u32, k1: u32, w: u32, out: []u32) void { + for (out[0..w], 0..) |*o, r| { + const rr: u32 = @intCast(r); + const first = k0 + (rr + w - k0 % w) % w; + o.* = if (first >= k1) 0 else (k1 - 1 - first) / w + 1; + } +} + +pub const max_world = 8; + +/// A file whose logical pages [0, pages) match a slot's or entry's rows: a later park writes only past it. +pub const Origin = struct { + key: [16]u8, + pages: u32, +}; + +/// One slot's page table: logical page i -> physical page; the device tables follow through the dirty range. +pub const Slot = struct { + pool: *Pool, + id: u32, + capacity: u64, + max_pages: u32, + /// physical page of each mapped logical page; len is the mapped count + pages: []u32, + len: u32 = 0, + /// pages the admission reserved for this slot's request (null: no admission control) + quota: ?u32 = null, + /// pages taken past the quota, from pages nobody had reserved + over: u32 = 0, + /// table entries changed since the device side last synced them: [lo, hi) + dirty_lo: u32 = std.math.maxInt(u32), + dirty_hi: u32 = 0, + /// the file the leading pages came from (truncation shortens it, release clears it) + origin: ?Origin = null, + + pub fn mapped(s: *const Slot) []const u32 { + return s.pages[0..s.len]; + } + + pub fn tokens(s: *const Slot) u64 { + return @as(u64, s.len) * s.pool.page; + } + + fn mark(s: *Slot, lo: u32, hi: u32) void { + s.dirty_lo = @min(s.dirty_lo, lo); + s.dirty_hi = @max(s.dirty_hi, hi); + } + + /// The table range to upload since the last call (empty when nothing changed): entries past len are the null page. + pub fn takeDirty(s: *Slot) ?struct { lo: u32, hi: u32 } { + if (s.dirty_lo >= s.dirty_hi) return null; + defer { + s.dirty_lo = std.math.maxInt(u32); + s.dirty_hi = 0; + } + return .{ .lo = s.dirty_lo, .hi = s.dirty_hi }; + } + + /// The device table's value at logical page i (the page, or the null page past len). + pub fn tableAt(s: *const Slot, i: u32) u32 { + return if (i < s.len) s.pages[i] else s.pool.null_page; + } + + /// The split families' device table at logical page i: the local page this rank stores, else the discard page. + pub fn localTableAt(s: *const Slot, i: u32) u32 { + return s.pool.localPage(s.tableAt(i)); + } + + fn reservedSpan(s: *const Slot, out: []u32) void { + const q = s.quota orelse 0; + span(s.len, @max(s.len, q), s.pool.world, out); + } + + /// Maps every page holding positions < end (a no-op when they are: the common case). + pub fn ensure(s: *Slot, end: u64) Error!void { + const p = s.pool; + const need = pagesFor(end, p.page); + if (need <= s.len) return; + if (need > s.max_pages) return error.PastCapacity; + const w = p.world; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + if (s.quota) |q| if (need > q) { + // past the reservation: only pages nobody else reserved, residue by residue + var extra: [max_world]u32 = undefined; + span(@max(q, s.len), need, w, &extra); + for (0..w) |r| if (extra[r] > p.free[r].count() -| (p.reserved[r] - before[r])) return error.PoolExhausted; + s.over += need - @max(q, s.len); + }; + var want: [max_world]u32 = undefined; + span(s.len, need, w, &want); + for (0..w) |r| if (want[r] > p.free[r].count()) return error.PoolExhausted; + const lo = s.len; + for (lo..need) |k| { + const pg = p.takeAt(@intCast(k)) catch unreachable; // counted above + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[k] = pg; + } + s.len = need; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(lo, need); + } + + /// Maps shared pages as the first pages of an empty slot (a RAM resume: read only, the slot writes past them). + pub fn adopt(s: *Slot, shared: []const u32) Error!void { + if (s.len != 0) return error.SlotNotEmpty; + if (shared.len > s.max_pages) return error.PastCapacity; + const p = s.pool; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + for (shared, 0..) |pg, k| if (pg >= p.npages or pg % p.world != k % p.world) return error.ResidueMismatch; + for (shared, 0..) |pg, k| { + p.hold(pg, 1, 1); + s.pages[k] = pg; + } + s.len = @intCast(shared.len); + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(0, s.len); + } + + /// Appends one page taken for logical page len (the copy-on-write page of a RAM resume). + pub fn push(s: *Slot, pg: u32) Error!void { + if (s.len >= s.max_pages) return error.PastCapacity; + const p = s.pool; + if (p.world > 1 and pg % p.world != s.len % p.world) return error.ResidueMismatch; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[s.len] = pg; + s.len += 1; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(s.len - 1, s.len); + } + + /// Lets go of the pages past position keep (rows nothing reads again); returns how many. + pub fn truncate(s: *Slot, keep: u64) Error!u32 { + const p = s.pool; + // positions >= keep get written again: a page holding one no longer equals the origin file's + if (s.origin) |*o| o.pages = @min(o.pages, @as(u32, @intCast(keep / p.page))); + const k = pagesFor(keep, p.page); + if (k >= s.len) return 0; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + const n = s.len - k; + for (s.pages[k..s.len]) |pg| try p.release(pg, 1); + const old = s.len; + s.len = k; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(k, old); + return n; + } + + /// Sets the admission's reservation (null: none). + pub fn reserve(s: *Slot, quota: ?u32) void { + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + s.quota = quota; + s.pool.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + s.pool.addReserved(&after); + } + + /// Frees the slot's pages and its reservation. + pub fn releaseAll(s: *Slot) Error!u32 { + s.reserve(null); + s.over = 0; + s.origin = null; + return s.truncate(0); + } +}; + +/// The pool's book: free sets by residue, holders a page, the slots and their reservations. +pub const Pool = struct { + gpa: std.mem.Allocator, + layout: Layout, + npages: u32, + page: u32, + /// log2 of page tokens + shift: u5, + /// ranks the split families are sharded over (1: every rank stores every row) + world: u32, + rank: u32, + /// the page unmapped table entries point at (all zeros: a stray read reads zeros) + null_page: u32, + /// a split family's local page count (npages / W), its local null page and its discard page (other ranks' rows) + local_pages: u32, + local_null: u32, + discard: u32, + free: [max_world]FreeSet = @splat(.{}), + /// pages each residue's slots reserved and have not mapped yet + reserved: [max_world]u32 = @splat(0), + /// holders a page (slots + entries + in-flight jobs) and how many are slots + refs: []u16, + slot_refs: []u16, + /// pages held by something and by no slot (session entries alone): what parking can free + entry_only: u32 = 0, + slots: std.ArrayList(*Slot) = .empty, + + pub fn init(gpa: std.mem.Allocator, layout: Layout, tokens: u64, world: u32, rank: u32) Error!Pool { + if (layout.page == 0 or layout.page & (layout.page - 1) != 0) return error.BadPool; + if (tokens % layout.page != 0 or tokens == 0) return error.BadPool; + if (world == 0 or world > max_world or rank >= world) return error.BadPool; + const npages: u32 = @intCast(tokens / layout.page); + if (npages % world != 0) return error.BadPool; + for (layout.families) |f| if (f.ratio == 0 or layout.page % f.ratio != 0) return error.BadPool; + var p: Pool = .{ + .gpa = gpa, + .layout = layout, + .npages = npages, + .page = layout.page, + .shift = @intCast(std.math.log2_int(u32, layout.page)), + .world = world, + .rank = rank, + .null_page = npages, + .local_pages = npages / world, + .local_null = npages / world, + .discard = npages / world + 1, + .refs = &.{}, + .slot_refs = &.{}, + }; + errdefer p.deinit(); + for (0..world) |r| p.free[r] = try FreeSet.init(gpa, npages / world); + p.refs = try gpa.alloc(u16, npages); + @memset(p.refs, 0); + p.slot_refs = try gpa.alloc(u16, npages); + @memset(p.slot_refs, 0); + return p; + } + + pub fn deinit(p: *Pool) void { + for (p.slots.items) |s| { + p.gpa.free(s.pages); + p.gpa.destroy(s); + } + p.slots.deinit(p.gpa); + for (0..p.world) |r| p.free[r].deinit(p.gpa); + p.gpa.free(p.refs); + p.gpa.free(p.slot_refs); + p.* = undefined; + } + + pub fn newSlot(p: *Pool, capacity: u64) Error!*Slot { + const s = try p.gpa.create(Slot); + errdefer p.gpa.destroy(s); + const mp = pagesFor(capacity, p.page); + s.* = .{ .pool = p, .id = @intCast(p.slots.items.len), .capacity = capacity, .max_pages = mp, .pages = try p.gpa.alloc(u32, mp) }; + errdefer p.gpa.free(s.pages); + try p.slots.append(p.gpa, s); + return s; + } + + /// A free physical page for logical page k (its residue when split), held once by the caller. + pub fn takeAt(p: *Pool, k: u32) Error!u32 { + const r = k % p.world; + const local = p.free[r].takeLowest() orelse return error.PoolExhausted; + const pg = local * p.world + r; + p.hold(pg, 1, 0); // the caller's hold (an entry-like holder until a slot maps it) + return pg; + } + + /// Holders change on page pg: dref total, dslot of them slots; keeps entry_only and gives the page back at zero. + fn hold(p: *Pool, pg: u32, dref: i32, dslot: i32) void { + const r0 = p.refs[pg]; + const s0 = p.slot_refs[pg]; + const r1: u16 = @intCast(@as(i32, r0) + dref); + const s1: u16 = @intCast(@as(i32, s0) + dslot); + p.refs[pg] = r1; + p.slot_refs[pg] = s1; + const was = r0 > 0 and s0 == 0; + const now = r1 > 0 and s1 == 0; + if (was and !now) p.entry_only -= 1; + if (!was and now) p.entry_only += 1; + } + + /// One holder lets go (slot: from a slot's table); the page returns to its free set when nobody holds it. + fn release(p: *Pool, pg: u32, slot: u1) Error!void { + if (p.refs[pg] == 0) return error.DoubleFree; + p.hold(pg, -1, -@as(i32, slot)); + if (p.refs[pg] == 0) p.free[pg % p.world].give(pg / p.world) catch return error.DoubleFree; + } + + /// One more holder (a session entry or an in-flight job) of each page. + pub fn share(p: *Pool, pages: []const u32) void { + for (pages) |pg| p.hold(pg, 1, 0); + } + + /// Those holders let go. + pub fn drop(p: *Pool, pages: []const u32) Error!void { + for (pages) |pg| try p.release(pg, 0); + } + + pub fn holders(p: *const Pool, pg: u32) u16 { + return p.refs[pg]; + } + + fn unreserve(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] -= by[r]; + } + + fn addReserved(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] += by[r]; + } + + /// Free pages no admission reserved (split: W x the scarcer residue, so any logical run fits). + pub fn available(p: *const Pool) u32 { + var low: u32 = std.math.maxInt(u32); + for (0..p.world) |r| low = @min(low, p.free[r].count() -| p.reserved[r]); + return low * p.world; + } + + pub fn freePages(p: *const Pool) u32 { + var n: u32 = 0; + for (0..p.world) |r| n += p.free[r].count(); + return n; + } + + /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). + pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { + return pagesFor(@min(prompt + max_tokens + slack, capacity), p.page); + } + + /// This rank stores the split families' rows of page pg. + pub fn owns(p: *const Pool, pg: u32) bool { + return pg != p.null_page and pg % p.world == p.rank; + } + + /// A split family's local page of pg: its own when owned, the local null for the null page, else the discard page. + pub fn localPage(p: *const Pool, pg: u32) u32 { + if (pg == p.null_page) return p.local_null; + return if (pg % p.world == p.rank) pg / p.world else p.discard; + } + + /// The page a family's tensor indexes for physical page pg (split families: localPage). + pub fn familyPage(p: *const Pool, f: Family, pg: u32) u32 { + return if (f.split and p.world > 1) p.localPage(pg) else pg; + } + + /// Pages a family's device tensor holds: npages + null, or (split) local pages + local null + discard page. + pub fn familyTensorPages(p: *const Pool, f: Family) u32 { + return if (f.split and p.world > 1) p.local_pages + 2 else p.npages + 1; + } + + /// Bytes this rank stores of family f for logical pages [first, first + n) of a slot or entry. + pub fn familyBytes(p: *const Pool, f: Family, first: u32, n: u32) u64 { + if (!(f.split and p.world > 1)) return @as(u64, n) * f.pageBytes(p.page); + var by: [max_world]u32 = undefined; + span(first, first + n, p.world, &by); + return @as(u64, by[p.rank]) * f.pageBytes(p.page); + } + + /// Bytes of every family's device tensor on this rank. + pub fn tensorBytes(p: *const Pool) u64 { + var b: u64 = 0; + for (p.layout.families) |f| b += @as(u64, p.familyTensorPages(f)) * f.pageBytes(p.page); + return b; + } + + /// Pages held by session entries alone and their bytes on this rank (the residency budget's measure): O(1). + pub fn entryOnlyBytes(p: *const Pool) u64 { + return @as(u64, p.entry_only) * p.layout.pageBytes(p.world); + } +}; + +test "residues, holders, quotas and the device tables" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 2, .row_bytes = 584, .split = true }, .{ .name = "index_k", .ratio = 2, .row_bytes = 132 } }; + var p = try Pool.init(gpa, .{ .families = &fams }, 16 * 256, 2, 1); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 8 + 2), p.familyTensorPages(fams[0])); + try std.testing.expectEqual(@as(u32, 16 + 1), p.familyTensorPages(fams[1])); + try std.testing.expectApproxEqAbs(@as(f64, 146 + 66), p.layout.bytesPerToken(2), 1e-9); + const a = try p.newSlot(8 * 256); + const b = try p.newSlot(8 * 256); + try a.ensure(3 * 256 + 1); // 4 pages: residues 0,1,0,1 + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, a.mapped()); + try b.ensure(256); + try std.testing.expectEqualSlices(u32, &.{4}, b.mapped()); + try std.testing.expectEqual(p.discard, a.localTableAt(0)); // page 0 is rank 0's + try std.testing.expectEqual(@as(u32, 0), a.localTableAt(1)); // page 1 -> rank 1's local page 0 + try std.testing.expectEqual(p.local_null, a.localTableAt(7)); + const d = a.takeDirty().?; + try std.testing.expectEqual(@as(u32, 0), d.lo); + try std.testing.expectEqual(@as(u32, 4), d.hi); + try std.testing.expect(a.takeDirty() == null); + // an entry shares a's pages; a lets go; the pages stay held by the entry alone + p.share(a.mapped()[0..3]); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + _ = try a.releaseAll(); + try std.testing.expectEqual(@as(u32, 3), p.entry_only); + try std.testing.expectEqual(@as(u32, 16 - 1 - 3), p.freePages()); + try p.drop(&.{ 0, 1, 2 }); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + try std.testing.expectError(error.DoubleFree, p.drop(&.{0})); + // quotas: b reserves 6 pages; available counts W x the scarcer residue + b.reserve(6); + var by: [2]u32 = undefined; + span(1, 6, 2, &by); + try std.testing.expectEqual(@as(u32, 2), by[0]); + try std.testing.expectEqual(@as(u32, 3), by[1]); + try std.testing.expectEqual(@as(u32, 2 * @min(7 - 2, 8 - 3)), p.available()); + try b.ensure(6 * 256); + try std.testing.expectEqual(@as(u32, 0), p.reserved[0] + p.reserved[1]); + for (b.mapped(), 0..) |pg, k| try std.testing.expectEqual(k % 2, pg % 2); + // RAM resume: adopt full pages of the right residues, refuse the wrong ones + const c = try p.newSlot(8 * 256); + try std.testing.expectError(error.ResidueMismatch, c.adopt(&.{b.pages[1]})); + p.share(b.mapped()[0..2]); + try c.adopt(b.mapped()[0..2]); + try c.push(try p.takeAt(2)); + try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); + try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); +} diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig new file mode 100644 index 0000000000..28c639fbd7 --- /dev/null +++ b/zig/src/sessions/prefix.zig @@ -0,0 +1,457 @@ +//! Session entries by prefix: a trie over chained page digests, one node per distinct run of full pages a tag. + +const std = @import("std"); + +pub const Digest = [16]u8; +pub const none = std.math.maxInt(u32); +const Blake = std.crypto.hash.blake2.Blake2b128; + +/// Chained digests of the full pages of ids under a tag: out[j] determines tag and ids[0 .. page (j + 1)]. +pub fn chain(tag: u32, ids: []const i32, page: u32, out: []Digest) void { + var prev = root(tag); + for (out, 0..) |*o, j| { + o.* = step(prev, ids[j * page ..][0..page]); + prev = o.*; + } +} + +/// The digest of a tag's empty prefix. +pub fn root(tag: u32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-root"); + h.update(std.mem.asBytes(&tag)); + return h.finalResult(); +} + +/// The next page's digest. +pub fn step(prev: Digest, page_ids: []const i32) Digest { + var h = Blake.init(.{}); + h.update(&prev); + h.update(std.mem.sliceAsBytes(page_ids)); + return h.finalResult(); +} + +/// An entry's identity: its node's digest (tag and full pages), its length and tail. +pub fn entryKey(node: Digest, pos: u64, tail: []const i32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-entry"); + h.update(&node); + h.update(std.mem.asBytes(&pos)); + h.update(std.mem.sliceAsBytes(tail)); + return h.finalResult(); +} + +pub const Tier = enum { ram, any }; + +pub fn Index(comptime Extra: type) type { + return struct { + const Self = @This(); + + pub const Node = struct { + digest: Digest, + parent: u32, + depth: u32, + first_child: u32 = none, + next_sib: u32 = none, + prev_sib: u32 = none, + first_entry: u32 = none, + /// entries in the subtree (this node's included): all, and those in RAM + sub_all: u32 = 0, + sub_ram: u32 = 0, + /// RAM entries at this node itself + own_ram: u32 = 0, + /// the tokens of the page leading into this node, while sub_ram > 0 + tokens: ?[]i32 = null, + }; + + pub const Entry = struct { + key: Digest, + tag: u32, + pos: u64, + node: u32, + tail: []i32, + ram: bool, + next_at: u32 = none, + prev_at: u32 = none, + live: bool = true, + extra: Extra, + }; + + /// An entry inserted, or the one already there with the same key. + pub const Placed = struct { id: u32, new: bool }; + + gpa: std.mem.Allocator, + page: u32, + nodes: std.ArrayList(Node) = .empty, + free_nodes: std.ArrayList(u32) = .empty, + by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + entries: std.ArrayList(Entry) = .empty, + free_entries: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// page-sized token blocks, reused + blocks: std.ArrayList([]i32) = .empty, + + pub fn init(gpa: std.mem.Allocator, page: u32) Self { + return .{ .gpa = gpa, .page = page }; + } + + pub fn deinit(x: *Self) void { + for (x.nodes.items) |n| if (n.tokens) |t| x.gpa.free(t); + for (x.blocks.items) |b| x.gpa.free(b); + for (x.entries.items) |e| if (e.live) x.gpa.free(e.tail); + x.nodes.deinit(x.gpa); + x.free_nodes.deinit(x.gpa); + x.by_digest.deinit(x.gpa); + x.entries.deinit(x.gpa); + x.free_entries.deinit(x.gpa); + x.by_key.deinit(x.gpa); + x.blocks.deinit(x.gpa); + x.* = undefined; + } + + pub fn get(x: *Self, e: u32) *Entry { + return &x.entries.items[e]; + } + + pub fn node(x: *Self, n: u32) *Node { + return &x.nodes.items[n]; + } + + pub fn lookup(x: *const Self, key: Digest) ?u32 { + return x.by_key.get(key); + } + + fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { + const id: u32 = if (x.free_nodes.pop()) |i| i else blk: { + try x.nodes.append(x.gpa, undefined); + break :blk @intCast(x.nodes.items.len - 1); + }; + x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; + try x.by_digest.put(x.gpa, digest, id); + if (parent != none) { + const p = &x.nodes.items[parent]; + x.nodes.items[id].next_sib = p.first_child; + if (p.first_child != none) x.nodes.items[p.first_child].prev_sib = id; + p.first_child = id; + } + return id; + } + + fn takeBlock(x: *Self) ![]i32 { + if (x.blocks.pop()) |b| return b; + return x.gpa.alloc(i32, x.page); + } + + /// Adds an entry for ids (chain: its full pages' digests), or returns the one with the same key. + pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { + const full: u32 = @intCast(ids.len / x.page); + std.debug.assert(digests.len >= full); + const r = root(tag); + var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + for (0..full) |j| { + const d = digests[j]; + n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); + } + const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); + if (r2.new and ram) try x.setRam(r2.id, true, ids); + return r2; + } + + /// Adds a disk entry continuing from node anchor (none: the tag's root); suffix: tokens [depth x page, pos). + pub fn insertFrom(x: *Self, tag: u32, anchor: u32, suffix: []const i32, extra: Extra) !Placed { + var n = if (anchor != none) anchor else blk: { + const r = root(tag); + break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); + }; + const d0 = x.nodes.items[n].depth; + const full: u32 = @intCast(suffix.len / x.page); + for (0..full) |j| { + const d = step(x.nodes.items[n].digest, suffix[j * x.page ..][0..x.page]); + n = x.by_digest.get(d) orelse try x.newNode(d, n, d0 + @as(u32, @intCast(j)) + 1); + } + return x.place(tag, n, @as(u64, d0) * x.page + suffix.len, suffix[full * x.page ..], extra); + } + + /// The node at depth on an entry's path (its first `depth` full pages). + pub fn nodeAt(x: *Self, id: u32, depth: u32) ?u32 { + var n = x.entries.items[id].node; + while (n != none and x.nodes.items[n].depth > depth) n = x.nodes.items[n].parent; + return if (n != none and x.nodes.items[n].depth == depth) n else null; + } + + fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { + const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); + if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; + const tail = try x.gpa.dupe(i32, tail_ids); + errdefer x.gpa.free(tail); + const id: u32 = if (x.free_entries.pop()) |i| i else blk: { + try x.entries.append(x.gpa, undefined); + break :blk @intCast(x.entries.items.len - 1); + }; + try x.by_key.put(x.gpa, key, id); + const nd = &x.nodes.items[n]; + x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; + if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; + nd.first_entry = id; + x.addPath(n, 1, 0); + return .{ .id = id, .new = true }; + } + + fn addPath(x: *Self, from: u32, d_all: i32, d_ram: i32) void { + var n = from; + while (n != none) { + const nd = &x.nodes.items[n]; + nd.sub_all = @intCast(@as(i64, nd.sub_all) + d_all); + nd.sub_ram = @intCast(@as(i64, nd.sub_ram) + d_ram); + n = nd.parent; + } + } + + /// Moves an entry in or out of RAM; going in needs its ids (to keep its path's tokens). + pub fn setRam(x: *Self, id: u32, ram: bool, ids: ?[]const i32) !void { + const e = &x.entries.items[id]; + if (e.ram == ram) return; + if (ram) { + const all = ids orelse return error.IdsNeeded; + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.tokens != null) continue; + const b = try x.takeBlock(); + @memcpy(b, all[(nd.depth - 1) * x.page ..][0..x.page]); + nd.tokens = b; + } + } + e.ram = ram; + x.nodes.items[e.node].own_ram = if (ram) x.nodes.items[e.node].own_ram + 1 else x.nodes.items[e.node].own_ram - 1; + x.addPath(e.node, 0, if (ram) 1 else -1); + if (!ram) x.dropTokens(e.node); + } + + fn dropTokens(x: *Self, from: u32) void { + var n = from; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_ram > 0) return; + if (nd.tokens) |t| { + x.blocks.append(x.gpa, t) catch x.gpa.free(t); + nd.tokens = null; + } + } + } + + /// Removes an entry and the nodes no entry needs any more. + pub fn remove(x: *Self, id: u32) void { + x.setRam(id, false, null) catch unreachable; + const e = &x.entries.items[id]; + const n = e.node; + if (e.prev_at != none) x.entries.items[e.prev_at].next_at = e.next_at else x.nodes.items[n].first_entry = e.next_at; + if (e.next_at != none) x.entries.items[e.next_at].prev_at = e.prev_at; + _ = x.by_key.remove(e.key); + x.gpa.free(e.tail); + e.live = false; + x.free_entries.append(x.gpa, id) catch {}; + x.addPath(n, -1, 0); + var m = n; + while (m != none and x.nodes.items[m].sub_all == 0) { + const nd = x.nodes.items[m]; + if (nd.parent != none) { + if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; + if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; + } + _ = x.by_digest.remove(nd.digest); + x.free_nodes.append(x.gpa, m) catch {}; + m = nd.parent; + } + } + + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). + fn tailOf(e: *const Entry, prompt: []const i32) bool { + const at: usize = @intCast(e.pos - e.tail.len); + return e.pos < prompt.len and std.mem.eql(i32, e.tail, prompt[at..][0..e.tail.len]); + } + + /// The longest entry of tag whose ids strictly prefix prompt (RAM first on a tie); digests: the prompt's chain. + pub fn find(x: *Self, tag: u32, prompt: []const i32, digests: []const Digest, tier: Tier) ?u32 { + const full: u32 = @intCast(prompt.len / x.page); + const r = x.by_digest.get(root(tag)) orelse return null; + // the deepest node on the prompt's path: each prefix of an entry's path has a node: binary search + var lo: u32 = 0; + var hi: u32 = full; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (x.by_digest.contains(digests[mid - 1])) lo = mid else hi = mid - 1; + } + var n: u32 = if (lo == 0) r else x.by_digest.get(digests[lo - 1]).?; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_all == 0 or (tier == .ram and nd.own_ram == 0)) continue; + var best: u32 = none; + var e = nd.first_entry; + while (e != none) : (e = x.entries.items[e].next_at) { + const en = &x.entries.items[e]; + if (tier == .ram and !en.ram) continue; + if (!tailOf(en, prompt)) continue; + if (best == none or en.pos > x.entries.items[best].pos or (en.pos == x.entries.items[best].pos and en.ram)) best = e; + } + if (best != none) return best; + } + return null; + } + + /// A newer entry of id's chain holds every position it does (tier .ram: in RAM), so id need not be written. + pub fn covered(x: *Self, id: u32, tier: Tier) bool { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == id or (tier == .ram and !o.ram)) continue; + if (o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) return true; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + const n = if (tier == .ram) ch.sub_ram else ch.sub_all; + if (n == 0) continue; + if (e.tail.len == 0) return true; + if (ch.tokens) |t| if (std.mem.startsWith(i32, t, e.tail)) return true; + } + return false; + } + + /// How many RAM entries extend a RAM entry (longer, same tag, its ids their prefix): branch points have more. + pub fn extensions(x: *Self, id: u32) u32 { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var n: u32 = 0; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f != id and o.ram and o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) n += 1; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + if (ch.sub_ram == 0) continue; + if (std.mem.startsWith(i32, ch.tokens.?, e.tail)) n += ch.sub_ram; + } + return n; + } + + /// A RAM entry no other RAM entry extends: a chat's newest state. + pub fn isLeaf(x: *Self, id: u32) bool { + return x.entries.items[id].ram and !x.covered(id, .ram); + } + + /// A RAM leaf's older turns (the RAM entries it extends) into out; returns the chat's recency (newest `used`). + pub fn members(x: *Self, leaf: u32, out: *std.ArrayList(u32), used: *const fn (*const Extra) u64) !u64 { + const l = &x.entries.items[leaf]; + var recency = used(&l.extra); + var below: u32 = none; // the path's node under n: its tokens are the page an entry at n must prefix + var n = l.node; + while (n != none) : ({ + below = n; + n = x.nodes.items[n].parent; + }) { + const nd = &x.nodes.items[n]; + if (nd.own_ram == 0) continue; + const page_ids: []const i32 = if (below == none) l.tail else x.nodes.items[below].tokens.?; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == leaf or !o.ram or o.pos >= l.pos) continue; + if (!std.mem.startsWith(i32, page_ids, o.tail)) continue; + try out.append(x.gpa, f); + recency = @max(recency, used(&o.extra)); + } + } + return recency; + } + + /// A RAM entry's ids, rebuilt from its path's pages and tail, into out (len = pos). + pub fn idsOf(x: *Self, id: u32, out: []i32) !void { + const e = &x.entries.items[id]; + if (out.len != e.pos) return error.BadLength; + const at = e.pos - e.tail.len; + @memcpy(out[at..], e.tail); + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + const t = nd.tokens orelse return error.NotInRam; + @memcpy(out[(nd.depth - 1) * x.page ..][0..x.page], t); + } + } + }; +} + +const T = Index(u64); + +fn usedOf(e: *const u64) u64 { + return e.*; +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + chain(tag, ids, x.page, d[0 .. ids.len / x.page]); + return (try x.insert(tag, ids, d[0 .. ids.len / x.page], ram, used)).id; +} + +test "longest strict prefix, ties, tags and coverage" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + const chat = try seq(gpa, 200, 0); + defer gpa.free(chat); + const t1 = try add(&x, 7, chat[0..40], true, 1); // a prompt's replay point + const t2 = try add(&x, 7, chat[0..72], true, 2); // a turn's end + const t3 = try add(&x, 7, chat[0..128], false, 3); // on disk only, page aligned + _ = try add(&x, 8, chat[0..100], true, 4); // another tag + var d: [16]Digest = undefined; + chain(7, chat[0..40], 16, d[0..2]); + const again = try x.insert(7, chat[0..40], d[0..2], true, 9); + try std.testing.expect(again.id == t1 and !again.new); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t3), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..150], d[0..9], .ram)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..128], d[0..8], .any)); // strict: 128 is not a prefix of 128 + chain(7, chat[0..72], 16, d[0..4]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, chat[0..72], d[0..4], .any)); + var other = try gpa.dupe(i32, chat[0..150]); + defer gpa.free(other); + other[50] = -1; // diverges inside t2's tail page: t1 still prefixes it + chain(7, other, 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, other, d[0..9], .any)); + // t3 extends t2 on disk only: its page's tokens are gone, so a tail there is undecided (the store marks covered_by) + try std.testing.expect(x.covered(t1, .ram) and !x.covered(t2, .any) and !x.covered(t2, .ram) and !x.covered(t3, .any)); + const t4 = try add(&x, 7, chat[0..64], false, 5); // page aligned: any deeper entry covers it + try std.testing.expect(x.covered(t4, .any) and x.covered(t4, .ram)); + x.remove(t4); + try std.testing.expect(x.isLeaf(t2) and !x.isLeaf(t1)); + try std.testing.expectEqual(@as(u32, 1), x.extensions(t1)); + try std.testing.expectEqual(@as(u32, 0), x.extensions(t2)); + var m: std.ArrayList(u32) = .empty; + defer m.deinit(gpa); + try std.testing.expectEqual(@as(u64, 2), try x.members(t2, &m, usedOf)); + try std.testing.expectEqualSlices(u32, &.{t1}, m.items); + const ids = try gpa.alloc(i32, 72); + defer gpa.free(ids); + try x.idsOf(t2, ids); + try std.testing.expectEqualSlices(i32, chat[0..72], ids); + // going to disk drops the path's tokens once no RAM entry needs them; removal prunes the nodes + try x.setRam(t2, false, null); + try x.setRam(t1, false, null); + try std.testing.expectError(error.NotInRam, x.idsOf(t2, ids)); + x.remove(t3); + x.remove(t2); + x.remove(t1); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig new file mode 100644 index 0000000000..7bdc46b386 --- /dev/null +++ b/zig/src/sessions/sessions.zig @@ -0,0 +1,33 @@ +//! Session memory, family-neutral: the KV pool book, page bytes, prefix entries, the RAM tier and the NVMe tier. + +pub const freeset = @import("freeset.zig"); +pub const pool = @import("pool.zig"); +pub const Pool = pool.Pool; +pub const Slot = pool.Slot; +pub const Family = pool.Family; +pub const Layout = pool.Layout; +pub const pagestore = @import("pagestore.zig"); +pub const PageStore = pagestore.PageStore; +pub const HostPool = pagestore.HostPool; +pub const prefix = @import("prefix.zig"); +pub const lanes = @import("lanes.zig"); +pub const format = @import("format.zig"); +pub const disk = @import("disk.zig"); +pub const Disk = disk.Disk; +pub const jobs = @import("jobs.zig"); +pub const store = @import("store.zig"); +pub const Store = store.Store; + +test { + _ = freeset; + _ = pool; + _ = pagestore; + _ = prefix; + _ = lanes; + _ = format; + _ = disk; + _ = jobs; + _ = store; + _ = @import("disk_test.zig"); + _ = @import("store_test.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig new file mode 100644 index 0000000000..851ce509aa --- /dev/null +++ b/zig/src/sessions/store.zig @@ -0,0 +1,570 @@ +//! The session store: RAM entries holding pool pages and bounded state, a residency budget, and the NVMe tier. + +const std = @import("std"); +const Io = std.Io; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const PageStore = @import("pagestore.zig").PageStore; +const prefix = @import("prefix.zig"); +const Digest = prefix.Digest; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; +const jobs = @import("jobs.zig"); +const Job = jobs.Job; + +pub const none = prefix.none; + +/// A slot's bounded state (rings, carries, lookback, taps) as the family snapshots it: opaque to the store. +pub const Bounded = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + nbytes: *const fn (ptr: *anyopaque) u64, + /// host copies of its pieces for a park (owned by the caller: name and bytes allocated with gpa) + blobs: *const fn (ptr: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob, + release: *const fn (ptr: *anyopaque) void, + }; +}; + +pub const Policy = struct { + /// bounded state the RAM tier keeps (least recently used entries go past it) + ram_bytes: u64 = 512 << 20, + /// idle chats resident in the pool (0: no limit) and pool bytes held by entries alone (0: no limit) + chats: u32 = 0, + resident_bytes: u64 = 0, + /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) + park_prefixes: bool = false, + /// drop, never write, an entry a newer entry of its chain covers + skip_covered: bool = true, + /// write a newer turn as a delta on the file its leading pages came from + delta: bool = true, + /// parks and restores on the tier thread (settled by the plan) + async_io: bool = false, +}; + +pub const Extra = struct { + /// RAM: the pool pages it holds, logical [0, npages) + pages: []u32 = &.{}, + bounded: ?Bounded = null, + /// a file with this key is registered on the NVMe tier + file: bool = false, + used: u64 = 0, + origin: ?pool_mod.Origin = null, + /// a newer entry of its chain was parked from RAM: writing this one again is redundant + covered_by: bool = false, + kind: u32 = 0, + /// an unsettled park of it + job: ?*Job = null, + older: u32 = none, + newer: u32 = none, +}; + +const Index = prefix.Index(Extra); + +pub const Hit = struct { id: u32, ram: bool, pos: u64 }; + +pub const Stats = struct { + saves: u64 = 0, + dups: u64 = 0, + hits_ram: u64 = 0, + hits_disk: u64 = 0, + parks: u64 = 0, + delta_parks: u64 = 0, + drops: u64 = 0, + covered_drops: u64 = 0, + chat_evictions: u64 = 0, + park_ns: u64 = 0, + restore_ns: u64 = 0, +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: Io, + pool: *Pool, + pages: PageStore, + index: Index, + disk: ?*Disk, + /// heap-held: the tier thread keeps its address + runner: ?*jobs.Runner = null, + policy: Policy, + clock: u64 = 0, + bounded_total: u64 = 0, + /// RAM entries, least recently used first + oldest: u32 = none, + newest: u32 = none, + ram_count: u32 = 0, + pending: std.ArrayList(*Job) = .empty, + scratch: std.ArrayList(Digest) = .empty, + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, io: Io, pool: *Pool, pages: PageStore, d: ?*Disk, policy: Policy) !Store { + var s: Store = .{ .gpa = gpa, .io = io, .pool = pool, .pages = pages, .index = Index.init(gpa, pool.page), .disk = d, .policy = policy }; + if (d) |dd| { + const r = try gpa.create(jobs.Runner); + r.* = jobs.Runner.init(gpa, io, dd); + s.runner = r; + if (policy.async_io) r.startThread() catch |e| { + gpa.destroy(r); + return e; + }; + } + return s; + } + + pub fn deinit(s: *Store) void { + if (s.runner) |r| { + for (s.pending.items) |j| r.wait(j); + r.deinit(); + s.gpa.destroy(r); + } + for (s.pending.items) |j| j.deinit(s.gpa); + s.pending.deinit(s.gpa); + for (s.index.entries.items) |*e| if (e.live) { + if (e.extra.bounded) |b| b.vtable.release(b.ptr); + s.pool.drop(e.extra.pages) catch {}; + s.gpa.free(e.extra.pages); + }; + s.index.deinit(); + s.scratch.deinit(s.gpa); + s.* = undefined; + } + + fn digests(s: *Store, tag: u32, ids: []const i32) ![]const Digest { + const n = ids.len / s.pool.page; + try s.scratch.resize(s.gpa, n); + prefix.chain(tag, ids, s.pool.page, s.scratch.items); + return s.scratch.items; + } + + pub fn entry(s: *Store, id: u32) *Index.Entry { + return s.index.get(id); + } + + /// The longest stored entry (RAM or NVMe) whose ids strictly prefix prompt, with the same tag. + pub fn find(s: *Store, tag: u32, prompt: []const i32) !?Hit { + const d = try s.digests(tag, prompt); + const id = s.index.find(tag, prompt, d, .any) orelse return null; + const e = s.index.get(id); + return .{ .id = id, .ram = e.ram, .pos = e.pos }; + } + + fn link(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + x.older = s.newest; + x.newer = none; + if (s.newest != none) s.index.get(s.newest).extra.newer = id else s.oldest = id; + s.newest = id; + } + + fn unlink(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + if (x.older != none) s.index.get(x.older).extra.newer = x.newer else s.oldest = x.newer; + if (x.newer != none) s.index.get(x.newer).extra.older = x.older else s.newest = x.older; + } + + fn touch(s: *Store, id: u32) void { + s.clock += 1; + const e = s.index.get(id); + e.extra.used = s.clock; + if (e.ram) { + s.unlink(id); + s.link(id); + } + } + + /// An entry for a slot's state at len(ids), sharing its pages; null for a duplicate (bounded is released then). + pub fn save(s: *Store, slot: *Slot, bounded: Bounded, ids: []const i32, tag: u32, kind: u32) !?u32 { + if (ids.len == 0) return error.EmptyEntry; + const need = pool_mod.pagesFor(ids.len, s.pool.page); + if (slot.len < need) return error.SlotTooShort; + const d = try s.digests(tag, ids); + const ins = try s.index.insert(tag, ids, d, false, .{ .kind = kind }); + errdefer if (ins.new) s.index.remove(ins.id); + const e = s.index.get(ins.id); + if (e.ram) { + bounded.vtable.release(bounded.ptr); + s.touch(ins.id); + s.stats.dups += 1; + return null; + } + const pages = try s.gpa.dupe(u32, slot.mapped()[0..need]); + errdefer s.gpa.free(pages); + try s.index.setRam(ins.id, true, ids); + s.pool.share(pages); + const x = &s.index.get(ins.id).extra; + x.pages = pages; + x.bounded = bounded; + x.kind = kind; + if (slot.origin) |o| x.origin = .{ .key = o.key, .pages = @min(o.pages, @as(u32, @intCast(ids.len / s.pool.page))) }; + s.bounded_total += bounded.vtable.nbytes(bounded.ptr); + s.ram_count += 1; + s.link(ins.id); + s.touch(ins.id); + s.stats.saves += 1; + try s.trimRam(); + _ = try s.enforce(); + return ins.id; + } + + /// A RAM entry into an empty slot: full pages shared, the partial last page copied; returns its bounded state. + pub fn restoreRam(s: *Store, id: u32, slot: *Slot) !Bounded { + const e = s.index.get(id); + if (!e.ram) return error.NotInRam; + const full: u32 = @intCast(e.pos / s.pool.page); + try slot.adopt(e.extra.pages[0..full]); + errdefer _ = slot.releaseAll() catch {}; + if (e.pos % s.pool.page != 0) { + const src = e.extra.pages[full]; + const dst = try s.pool.takeAt(full); + for (s.pool.layout.families, 0..) |f, fi| { + if (f.split and !s.pool.owns(src)) continue; + try s.pages.copyPage(@intCast(fi), s.pool.familyPage(f, src), s.pool.familyPage(f, dst)); + } + try slot.push(dst); + } + if (e.extra.origin) |o| slot.origin = .{ .key = o.key, .pages = @min(o.pages, full) }; + s.touch(id); + s.stats.hits_ram += 1; + return e.extra.bounded.?; + } + + /// Starts an NVMe resume into an empty slot (pages mapped now, filled by the job); settle it before the slot runs. + pub fn beginRestore(s: *Store, id: u32, slot: *Slot) !*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.extra.file) return error.NotOnDisk; + if (slot.len != 0) return error.SlotNotEmpty; + try slot.ensure(e.pos); + const j = try s.gpa.create(Job); + j.* = .{ .kind = .restore, .entry = id, .slot = slot.id }; + errdefer j.deinit(s.gpa); + var ext: std.ArrayList(disk_mod.Extent) = .empty; + errdefer ext.deinit(s.gpa); + try d.extents(e.key, &ext); + j.extents = try ext.toOwnedSlice(s.gpa); + const np = pool_mod.pagesFor(e.pos, s.pool.page); + j.lists = try s.familyLists(slot.mapped()[0..np], 0, false); + j.sink = .{ .store = s.pages, .families = s.pool.layout.families, .page = s.pool.page, .world = s.pool.world, .rank = s.pool.rank, .pages = j.lists }; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Per family: family-tensor pages of logical [first, len(pages)) (owned_only: this rank's of split families). + fn familyLists(s: *Store, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const fams = s.pool.layout.families; + const out = try s.gpa.alloc([]u32, fams.len); + var made: usize = 0; + errdefer { + for (out[0..made]) |l| s.gpa.free(l); + s.gpa.free(out); + } + for (fams, 0..) |f, fi| { + const w = s.pool.world; + const sharded = f.split and w > 1 and owned_only; + var n: usize = pages.len - first; + if (sharded) { + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, @intCast(pages.len), w, &by); + n = by[s.pool.rank]; + } + const l = try s.gpa.alloc(u32, n); + var i: usize = 0; + var k: usize = first; + if (sharded) k = first + (s.pool.rank + w - first % w) % w; + while (k < pages.len) : (k += if (sharded) w else 1) { + l[i] = s.pool.familyPage(f, pages[k]); + i += 1; + } + out[fi] = l; + made += 1; + } + return out; + } + + /// Starts a park of a RAM entry (null: nothing to write, the entry is already settled out of RAM). + pub fn beginPark(s: *Store, id: u32) !?*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.ram or e.extra.job != null) return error.NotParkable; + if (e.extra.file and d.get(e.key) != null) { + try s.leaveRam(id); + return null; + } + const page = s.pool.page; + const npages = pool_mod.pagesFor(e.pos, page); + var base: ?Digest = null; + var first: u32 = 0; + if (s.policy.delta) if (e.extra.origin) |o| if (o.pages > 0) { + var r = d.get(o.key); + while (r) |rec| : (r = if (rec.base == disk_mod.none) null else &d.recs.items[rec.base]) { + if (rec.first_page < o.pages) break; + } + if (r) |rec| if (rec.depth < d.opts.max_chain) { + base = rec.key; + first = o.pages; + }; + }; + const j = try s.gpa.create(Job); + j.* = .{ .kind = .park, .entry = id }; + errdefer j.deinit(s.gpa); + j.ids = try s.gpa.alloc(i32, e.pos); + try s.index.idsOf(id, j.ids); + j.lists = try s.familyLists(e.extra.pages, first, true); + j.blobs = try e.extra.bounded.?.vtable.blobs(e.extra.bounded.?.ptr, s.gpa); + j.src = .{ .key = e.key, .base = base, .tag = e.tag, .kind = e.extra.kind, .page = page, .pos = e.pos, .first_page = first, .npages = npages, .ids = j.ids[first * page ..], .store = s.pages, .fence = try s.pages.fence(), .families = s.pool.layout.families, .pages = j.lists, .blobs = j.blobs }; + if (base) |b| d.get(b).?.refs += 1; // held while the job runs: a trim cannot take the base + e.extra.job = j; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Settles a finished job at the plan's point: a park registers its file, a restore returns the bounded state. + pub fn settle(s: *Store, j: *Job) !?disk_mod.Restored { + s.runner.?.wait(j); + defer { + for (s.pending.items, 0..) |p, i| if (p == j) { + _ = s.pending.orderedRemove(i); + break; + }; + j.deinit(s.gpa); + } + const d = s.disk.?; + switch (j.kind) { + .park => { + if (j.src.base) |b| d.get(b).?.refs -= 1; + const e = s.index.get(j.entry); + e.extra.job = null; + if (j.err) |err| return err; // the entry stays in RAM; the .tmp is gone + _ = try d.register(&j.src, j.size); + e.extra.file = true; + s.stats.parks += 1; + if (j.src.base != null) s.stats.delta_parks += 1; + s.stats.park_ns += j.elapsed_ns; + if (s.policy.skip_covered) { + var m: std.ArrayList(u32) = .empty; + defer m.deinit(s.gpa); + _ = try s.index.members(j.entry, &m, usedOf); + for (m.items) |o| s.index.get(o).extra.covered_by = true; + } + try s.leaveRam(j.entry); + try s.trimDisk(); + return null; + }, + .restore => { + if (j.err) |err| { + // the slot as before the restore: empty, its reservation kept + _ = try s.pool.slots.items[j.slot].truncate(0); + return err; + } + const e = s.index.get(j.entry); + d.touch(e.key); + s.touch(j.entry); + s.pool.slots.items[j.slot].origin = .{ .key = e.key, .pages = @intCast(e.pos / s.pool.page) }; + s.stats.hits_disk += 1; + s.stats.restore_ns += j.elapsed_ns; + const r = j.restored; + j.restored = null; + return r; + }, + } + } + + /// Settles every pending job in submission order (one rank, or the plan's "settle all"). + pub fn settleAll(s: *Store) !void { + while (s.pending.items.len > 0) { + if (try s.settle(s.pending.items[0])) |r| { + var rr = r; + rr.deinit(s.gpa); + } + } + } + + fn usedOf(x: *const Extra) u64 { + return x.used; + } + + /// A RAM entry lets go of its pages and bounded state; it stays indexed if a file holds it. + fn leaveRam(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (!e.ram) return; + try s.pool.drop(e.extra.pages); + s.gpa.free(e.extra.pages); + e.extra.pages = &.{}; + if (e.extra.bounded) |b| { + s.bounded_total -= b.vtable.nbytes(b.ptr); + b.vtable.release(b.ptr); + } + e.extra.bounded = null; + s.unlink(id); + s.ram_count -= 1; + try s.index.setRam(id, false, null); + if (!e.extra.file) s.index.remove(id); + } + + /// An entry whose file could not be read back: the file deleted, the entry gone unless RAM holds it (every rank). + pub fn forget(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (e.extra.job != null) return error.Busy; + if (e.extra.file) if (s.disk) |d| { + d.stats.bad += 1; + // a delta's base stays while a newer file stacks on it (an opt-in delta mode) + d.remove(e.key) catch |err| return if (err == error.HasDependents) {} else err; + e.extra.file = false; + }; + if (!e.ram) s.index.remove(id); + } + + /// Drops an entry from RAM without writing it. + pub fn drop(s: *Store, id: u32) !void { + try s.leaveRam(id); + s.stats.drops += 1; + } + + /// Lets go of a RAM entry: dropped when covered, parked if the tier takes it, else dropped; returns a started park. + pub fn evict(s: *Store, id: u32) !?*Job { + const e = s.index.get(id); + if (e.extra.job != null) return null; + if (s.policy.skip_covered and (e.extra.covered_by or s.index.covered(id, .any))) { + try s.drop(id); + s.stats.covered_drops += 1; + return null; + } + if (s.disk) |d| if (d.accepts(e.pos)) { + const j = try s.beginPark(id) orelse return null; + if (!s.policy.async_io) _ = try s.settle(j); + return if (s.policy.async_io) j else null; + }; + try s.drop(id); + return null; + } + + /// Evicts least recently used RAM entries past the bounded-state budget. + fn trimRam(s: *Store) !void { + var id = s.oldest; + while (s.bounded_total > s.policy.ram_bytes and s.ram_count > 1 and id != none) { + const next = s.index.get(id).extra.newer; + if (s.index.get(id).extra.job == null) _ = try s.evict(id); + id = next; + } + } + + /// Files the NVMe budget deleted leave the index (a RAM entry with that key stays, unwritten again). + fn trimDisk(s: *Store) !void { + const d = s.disk orelse return; + var gone: std.ArrayList(Digest) = .empty; + defer gone.deinit(s.gpa); + try d.trim(&gone); + for (gone.items) |k| if (s.index.lookup(k)) |id| { + s.index.get(id).extra.file = false; + if (!s.index.get(id).ram) s.index.remove(id); + }; + } + + const Chat = struct { leaf: u32, recency: u64 }; + + fn chatLess(_: void, a: Chat, b: Chat) bool { + return a.recency < b.recency or (a.recency == b.recency and a.leaf < b.leaf); + } + + /// A chat is live while a slot maps its leaf's first page. + fn idleLeaf(s: *Store, id: u32) bool { + const e = s.index.get(id); + return e.extra.job == null and s.index.isLeaf(id) and (e.extra.pages.len == 0 or s.pool.slot_refs[e.extra.pages[0]] == 0); + } + + /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. + pub fn enforce(s: *Store) !u32 { + const pol = s.policy; + if (pol.chats == 0 and pol.resident_bytes == 0) return 0; + var chats: std.ArrayList(Chat) = .empty; + defer chats.deinit(s.gpa); + var members: std.ArrayList(u32) = .empty; + defer members.deinit(s.gpa); + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) { + if (!s.idleLeaf(id)) continue; + members.clearRetainingCapacity(); + try chats.append(s.gpa, .{ .leaf = id, .recency = try s.index.members(id, &members, usedOf) }); + } + std.mem.sort(Chat, chats.items, {}, chatLess); + var idle: u32 = @intCast(chats.items.len); + var gone: u32 = 0; + for (chats.items) |c| { + const over_n = pol.chats > 0 and idle > pol.chats; + const over_b = pol.resident_bytes > 0 and s.pool.entryOnlyBytes() > pol.resident_bytes; + if (!over_n and !over_b) break; + if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; + members.clearRetainingCapacity(); + _ = try s.index.members(c.leaf, &members, usedOf); + _ = try s.evict(c.leaf); + // deepest first: a member more RAM entries extend than this chat's leaf is a branch point, as is all above + const own: u32 = if (s.index.get(c.leaf).ram) 1 else 0; + for (members.items) |m| { + if (s.index.get(m).extra.job != null) continue; + if (s.index.extensions(m) > own) break; + if (pol.park_prefixes) _ = try s.evict(m) else try s.drop(m); + } + idle -= 1; + gone += 1; + s.stats.chat_evictions += 1; + } + return gone; + } + + /// RAM entries to let go of (coldest first, never keep) so need pages are free: [] if they are, null if impossible. + pub fn planSpill(s: *Store, need: u32, keep: []const u32, out: *std.ArrayList(u32)) !bool { + out.clearRetainingCapacity(); + var avail = s.pool.available(); + if (avail >= need) return true; + var left: std.AutoHashMapUnmanaged(u32, u16) = .empty; + defer left.deinit(s.gpa); + var id = s.oldest; + while (id != none and avail < need) : (id = s.index.get(id).extra.newer) { + if (std.mem.indexOfScalar(u32, keep, id) != null or s.index.get(id).extra.job != null) continue; + try out.append(s.gpa, id); + for (s.index.get(id).extra.pages) |pg| { + const g = try left.getOrPut(s.gpa, pg); + if (!g.found_existing) g.value_ptr.* = s.pool.refs[pg]; + g.value_ptr.* -= 1; + if (g.value_ptr.* == 0) avail += 1; // split: a page of the scarcer residue may count for less; the check below decides + } + } + return avail >= need; + } + + /// Rebuilds NVMe entries on restart (ranks then intersect them), deltas on their bases' paths; returns how many. + pub fn reconcile(s: *Store) !u32 { + const d = s.disk orelse return 0; + var found: std.ArrayList(Disk.Found) = .empty; + defer { + for (found.items) |f| s.gpa.free(f.ids); + found.deinit(s.gpa); + } + try d.scan(&found); + for (found.items) |f| { + var anchor: u32 = none; + if (f.head.hasBase()) { + const b = s.index.lookup(f.head.base) orelse return error.NoBase; + anchor = s.index.nodeAt(b, f.head.first_page) orelse return error.BadChain; + } + const r = try s.index.insertFrom(f.head.tag, anchor, f.ids, .{ .file = true, .kind = f.head.kind }); + s.clock += 1; + s.index.get(r.id).extra.used = s.clock; + } + return @intCast(found.items.len); + } + + /// Idle resident chats and the pool bytes held by entries alone. + pub fn resident(s: *Store) struct { chats: u32, bytes: u64 } { + var n: u32 = 0; + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) if (s.idleLeaf(id)) { + n += 1; + }; + return .{ .chats = n, .bytes = s.pool.entryOnlyBytes() }; + } +}; diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig new file mode 100644 index 0000000000..86b58a4f59 --- /dev/null +++ b/zig/src/sessions/store_test.zig @@ -0,0 +1,216 @@ +//! The session store on a host pool and a real NVMe tier: chats, evictions, resumes bit for bit, deltas, two ranks. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk_mod = @import("disk.zig"); +const store_mod = @import("store.zig"); +const Store = store_mod.Store; +const Bounded = store_mod.Bounded; + +const page = 64; +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Bounded state in host memory: one blob naming its chat and length. +const HostBounded = struct { + gpa: std.mem.Allocator, + bytes: []u8, + + fn make(gpa: std.mem.Allocator, chat: u32, pos: u64) !Bounded { + const h = try gpa.create(HostBounded); + h.* = .{ .gpa = gpa, .bytes = try gpa.alloc(u8, 3000) }; + for (h.bytes, 0..) |*b, i| b.* = @truncate(i +% chat *% 17 +% pos); + return .{ .ptr = h, .vtable = &vt }; + } + const vt: Bounded.VTable = .{ .nbytes = nbytes, .blobs = blobs, .release = release }; + fn nbytes(p: *anyopaque) u64 { + const h: *HostBounded = @ptrCast(@alignCast(p)); + return h.bytes.len; + } + fn blobs(p: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob { + const h: *HostBounded = @ptrCast(@alignCast(p)); + const out = try gpa.alloc(disk_mod.Blob, 1); + out[0] = .{ .name = try gpa.dupe(u8, "rings"), .bytes = try gpa.dupe(u8, h.bytes) }; + return out; + } + fn release(p: *anyopaque) void { + const h: *HostBounded = @ptrCast(@alignCast(p)); + h.gpa.free(h.bytes); + h.gpa.destroy(h); + } +}; + +/// One rank: its pool, host pages and store, plus what each chat's slot wrote (the reference rows by logical page). +const Rank = struct { + pool: Pool, + host: HostPool, + disk: *disk_mod.Disk, + store: Store, + + fn init(r: *Rank, gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, policy: store_mod.Policy) !void { + r.pool = try Pool.init(gpa, .{ .families = &fams, .page = page }, 96 * page, world, rank); + r.host = try HostPool.init(gpa, &r.pool); + r.disk = try disk_mod.Disk.open(gpa, testing.io, .{ .root = root, .compat = @splat(7), .rank = rank, .world = world, .chunk = 64 << 10, .staging = 4, .lanes = 2, .min_tokens = 1 }); + r.store = try Store.init(gpa, testing.io, &r.pool, r.host.store(), r.disk, policy); + } + + fn deinit(r: *Rank) void { + r.store.deinit(); + r.disk.close(); + r.host.deinit(); + r.pool.deinit(); + } + + /// The rows a chat's slot writes at a logical page: a function of (chat, logical page, family) alone. + fn write(r: *Rank, s: *const Slot, chat: u32, from_page: u32) void { + for (s.mapped()[from_page..], from_page..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi); + }; + } + + /// Every stored row of the slot's pages below pos equals what write() put there. + fn check(r: *Rank, s: *const Slot, chat: u32, pos: u64) !void { + const n = pool_mod.pagesFor(pos, page); + for (s.mapped()[0..n], 0..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + const rows: usize = if (k + 1 == n and pos % page != 0) (pos % page) / f.ratio * f.row_bytes else v.len; + for (v[0..rows], 0..) |b, i| try testing.expectEqual(@as(u8, @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi)), b); + }; + } +}; + +fn chatIds(gpa: std.mem.Allocator, chat: u32, n: usize) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @intCast(chat * 100_000 + i); + return v; +} + +/// Prefill or resume a chat to len(ids) in slot s, save it, free the slot, enforce residency; returns the resume point. +fn turn(gpa: std.mem.Allocator, r: *Rank, s: *Slot, chat: u32, ids: []const i32) ![]const u8 { + var tier: []const u8 = "fresh"; + var from: u32 = 0; + if (try r.store.find(1, ids)) |hit| { + if (hit.ram) { + const b = try r.store.restoreRam(hit.id, s); + try testing.expectEqual(@as(u64, 3000), b.vtable.nbytes(b.ptr)); + tier = "ram"; + } else { + const j = try r.store.beginRestore(hit.id, s); + var got = (try r.store.settle(j)).?; + defer got.deinit(gpa); + try testing.expectEqual(@as(u8, @truncate(chat *% 17 +% hit.pos)), got.blobs[0].bytes[0]); + tier = "disk"; + } + try r.check(s, chat, hit.pos); + from = @intCast(hit.pos / page); + } + try s.ensure(ids.len); + r.write(s, chat, from); + _ = try r.store.save(s, try HostBounded.make(gpa, chat, ids.len), ids, 1, 0); + _ = try s.releaseAll(); + _ = try r.store.enforce(); + try r.store.settleAll(); // the plan's settle point (async parks finish here) + return tier; +} + +/// The scenario on one rank; returns a trace of its decisions (stats and page ids) to compare ranks. +fn scenario(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, async_io: bool, trace: *std.ArrayList(u64)) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1, .async_io = async_io }); + defer r.deinit(); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + const b = try chatIds(gpa, 2, 20 * page); + defer gpa.free(b); + const st = &r.store.stats; + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 1, a[0 .. 5 * page + 10])); + try testing.expectEqualStrings("ram", try turn(gpa, &r, s, 1, a[0 .. 8 * page + 3])); + try testing.expectEqual(@as(u64, 0), st.parks); + // a second idle chat: chat 1 (least recent) parked whole, its older turn dropped (not written) + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 2, b[0 .. 6 * page])); + try testing.expectEqual(@as(u64, 1), st.parks); + try testing.expectEqual(@as(u64, 1), st.drops); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + const whole = r.disk.used; + // chat 1 resumes from NVMe bit for bit; chat 2 goes + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 12 * page + 40])); + try testing.expectEqual(@as(u64, 2), st.parks); + const before = r.disk.used; + // chat 2 resumes; chat 1 goes again, written as a delta on its first file (pages 8..12 of 13) + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 2, b[0 .. 9 * page])); + try testing.expectEqual(@as(u64, 3), st.parks); + try testing.expectEqual(@as(u64, 1), st.delta_parks); + // the delta holds 5 of the entry's 13 pages: less than chat 1's first whole file of 9 pages + try testing.expect(r.disk.used - before < whole); + // chat 1's newest turn restores through its two files, bit for bit + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 15 * page])); + try testing.expectEqual(@as(u64, 3), st.hits_disk); + try trace.appendSlice(gpa, &.{ st.parks, st.delta_parks, st.drops, st.covered_drops, st.hits_ram, st.hits_disk, r.pool.freePages(), r.pool.entry_only }); + var id = r.store.oldest; + while (id != store_mod.none) : (id = r.store.entry(id).extra.newer) for (r.store.entry(id).extra.pages) |pg| try trace.append(gpa, pg); +} + +/// After the scenario: a new process on the same tier finds every file and resumes chat 1 through its delta chain. +fn restart(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1 }); + defer r.deinit(); + try testing.expectEqual(@as(u32, 4), try r.store.reconcile()); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 16 * page])); + try testing.expectEqual(@as(u64, 1), r.store.stats.hits_disk); +} + +fn tmpRoot(gpa: std.mem.Allocator, tmp: *std.testing.TmpDir, sub: []const u8) ![:0]u8 { + try tmp.dir.createDirPath(testing.io, sub); + return tmp.dir.realPathFileAlloc(testing.io, sub, gpa); +} + +test "chats over turns: residency to NVMe, bit-identical resumes, deltas, older turns dropped (sync and async)" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const r0 = try tmpRoot(gpa, &tmp, "sync"); + defer gpa.free(r0); + try scenario(gpa, r0, 1, 0, false, &t0); + try restart(gpa, r0, 1, 0); + const r1 = try tmpRoot(gpa, &tmp, "async"); + defer gpa.free(r1); + try scenario(gpa, r1, 1, 0, true, &t1); + try testing.expectEqualSlices(u64, t0.items, t1.items); // the tier thread changes when bytes move, not what is decided +} + +test "two split ranks make the same decisions on the same calls and each restores its own rows" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const root = try tmpRoot(gpa, &tmp, "tier"); + defer gpa.free(root); + try scenario(gpa, root, 2, 0, false, &t0); + try scenario(gpa, root, 2, 1, true, &t1); + try restart(gpa, root, 2, 0); + try restart(gpa, root, 2, 1); + try testing.expectEqualSlices(u64, t0.items, t1.items); +} From a5fda7a63b4366e16d46e8357015feceaa8a463a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:48:20 +0700 Subject: [PATCH 7/7] Grammar-constrained decoding on the Zig side: xgrammar 0.2.8's C++ core and a grammar module over it. zig/vendor/xgrammar holds XGrammar 0.2.8's C++ sources (tag v0.2.8, commit 97787376), byte-identical to the release with no local patches; its README records the version, commit and patch list. Zig's C++ toolchain builds them into a static library with a small C ABI (zig/src/core/grammar/xgr_c.cc). The grammar module compiles a request's spec (JSON, JSON schema, regex, choice, EBNF, or a tools spec as a structural tag), keeps one matcher a reply, fills token bitmasks for drafted rows and cuts a draft where the grammar ends, packs a spec for other ranks, and carries a PTX mask kernel for logits rows. zig build test-grammar checks masks, cuts, compiles and refusals against goldens recorded from xgrammar's Python package on a synthetic tokenizer. No family uses it yet. lean_check.py skips exactly the vendored prefix (VENDORED, beside the Flash Next generator outputs), so the upstream copy can stay byte-identical and be updated by replacing it. lean_check_test.py shows the prefix is skipped and that files right beside it, our grammar wrappers included, are still checked. --- NOTICE | 3 + THIRD_PARTY_NOTICES.md | 9 + build.zig | 2 + tools/zig/lean_check.py | 8 +- tools/zig/lean_check_test.py | 38 + zig/build/grammar.zig | 59 + zig/src/core/grammar/constraint.zig | 120 + zig/src/core/grammar/engine.zig | 70 + zig/src/core/grammar/golden_test.zig | 169 + zig/src/core/grammar/grammars.zig | 219 + zig/src/core/grammar/mask_kernel.zig | 13 + zig/src/core/grammar/pack.zig | 52 + zig/src/core/grammar/root.zig | 22 + zig/src/core/grammar/rows.zig | 62 + zig/src/core/grammar/tags.zig | 165 + zig/src/core/grammar/vocab.zig | 74 + zig/src/core/grammar/xgr.zig | 193 + zig/src/core/grammar/xgr_c.cc | 157 + zig/src/core/grammar/xgr_c.h | 58 + .../grammar/fixtures/synthetic.golden.json | 1 + .../grammar/fixtures/synthetic/tokenizer.json | 1 + zig/tests/grammar/gen_goldens.py | 302 + zig/vendor/xgrammar/3rdparty/dlpack/LICENSE | 201 + .../3rdparty/dlpack/include/dlpack/dlpack.h | 332 ++ .../xgrammar/3rdparty/picojson/picojson.h | 1318 +++++ zig/vendor/xgrammar/LICENSE | 201 + zig/vendor/xgrammar/NOTICE | 3 + zig/vendor/xgrammar/README.md | 14 + zig/vendor/xgrammar/cpp/compiled_grammar.cc | 279 + .../xgrammar/cpp/compiled_grammar_impl.h | 148 + zig/vendor/xgrammar/cpp/config.cc | 21 + .../xgrammar/cpp/converter_ext/cohere.cc | 653 +++ .../xgrammar/cpp/converter_ext/deepseek.cc | 117 + zig/vendor/xgrammar/cpp/converter_ext/glm.cc | 14 + .../xgrammar/cpp/converter_ext/kimi_k3.cc | 47 + .../xgrammar/cpp/converter_ext/minimax.cc | 14 + .../xgrammar/cpp/converter_ext/minimax_m3.cc | 315 + zig/vendor/xgrammar/cpp/converter_ext/qwen.cc | 14 + zig/vendor/xgrammar/cpp/earley_parser.cc | 1272 ++++ zig/vendor/xgrammar/cpp/earley_parser.h | 802 +++ zig/vendor/xgrammar/cpp/ebnf_script_creator.h | 188 + zig/vendor/xgrammar/cpp/fsm.cc | 2061 +++++++ zig/vendor/xgrammar/cpp/fsm.h | 1051 ++++ zig/vendor/xgrammar/cpp/fsm_builder.cc | 1814 ++++++ zig/vendor/xgrammar/cpp/fsm_builder.h | 119 + zig/vendor/xgrammar/cpp/grammar.cc | 359 ++ zig/vendor/xgrammar/cpp/grammar_builder.cc | 425 ++ zig/vendor/xgrammar/cpp/grammar_builder.h | 292 + zig/vendor/xgrammar/cpp/grammar_compiler.cc | 1725 ++++++ zig/vendor/xgrammar/cpp/grammar_functor.cc | 4016 +++++++++++++ zig/vendor/xgrammar/cpp/grammar_functor.h | 520 ++ zig/vendor/xgrammar/cpp/grammar_impl.h | 487 ++ zig/vendor/xgrammar/cpp/grammar_matcher.cc | 2856 +++++++++ zig/vendor/xgrammar/cpp/grammar_parser.cc | 1593 +++++ zig/vendor/xgrammar/cpp/grammar_parser.h | 107 + zig/vendor/xgrammar/cpp/grammar_printer.cc | 315 + zig/vendor/xgrammar/cpp/grammar_printer.h | 86 + .../xgrammar/cpp/json_schema_converter.cc | 5126 +++++++++++++++++ .../xgrammar/cpp/json_schema_converter.h | 666 +++ .../xgrammar/cpp/json_schema_converter_ext.h | 322 ++ zig/vendor/xgrammar/cpp/lark_converter.cc | 2714 +++++++++ zig/vendor/xgrammar/cpp/lark_converter.h | 26 + zig/vendor/xgrammar/cpp/regex_converter.cc | 413 ++ zig/vendor/xgrammar/cpp/regex_converter.h | 21 + zig/vendor/xgrammar/cpp/structural_tag.cc | 2546 ++++++++ zig/vendor/xgrammar/cpp/structural_tag.h | 415 ++ zig/vendor/xgrammar/cpp/suffix_automata.cc | 88 + zig/vendor/xgrammar/cpp/suffix_automata.h | 36 + .../xgrammar/cpp/support/compact_2d_array.h | 397 ++ zig/vendor/xgrammar/cpp/support/container.h | 166 + zig/vendor/xgrammar/cpp/support/cpptrace.h | 39 + .../xgrammar/cpp/support/dynamic_bitset.h | 363 ++ zig/vendor/xgrammar/cpp/support/encoding.h | 459 ++ zig/vendor/xgrammar/cpp/support/int_set.h | 132 + zig/vendor/xgrammar/cpp/support/json_parse.h | 118 + .../xgrammar/cpp/support/json_serializer.h | 640 ++ zig/vendor/xgrammar/cpp/support/logging.cc | 24 + zig/vendor/xgrammar/cpp/support/logging.h | 236 + zig/vendor/xgrammar/cpp/support/memory_size.h | 129 + .../xgrammar/cpp/support/recursion_guard.cc | 50 + .../xgrammar/cpp/support/recursion_guard.h | 127 + zig/vendor/xgrammar/cpp/support/reflection.h | 300 + zig/vendor/xgrammar/cpp/support/thread_pool.h | 239 + .../xgrammar/cpp/support/thread_safe_cache.h | 404 ++ .../xgrammar/cpp/support/union_find_set.h | 105 + zig/vendor/xgrammar/cpp/support/utils.h | 451 ++ zig/vendor/xgrammar/cpp/testing.cc | 63 + zig/vendor/xgrammar/cpp/testing.h | 30 + zig/vendor/xgrammar/cpp/tokenizer_info.cc | 562 ++ zig/vendor/xgrammar/cpp/tokenizer_info_impl.h | 143 + zig/vendor/xgrammar/include/module.modulemap | 4 + .../xgrammar/include/xgrammar/compiler.h | 125 + zig/vendor/xgrammar/include/xgrammar/config.h | 35 + .../xgrammar/include/xgrammar/exception.h | 80 + .../xgrammar/include/xgrammar/grammar.h | 228 + .../xgrammar/include/xgrammar/matcher.h | 300 + zig/vendor/xgrammar/include/xgrammar/object.h | 51 + .../include/xgrammar/tokenizer_info.h | 86 + .../xgrammar/include/xgrammar/xgrammar.h | 17 + 99 files changed, 44051 insertions(+), 1 deletion(-) create mode 100644 tools/zig/lean_check_test.py create mode 100644 zig/build/grammar.zig create mode 100644 zig/src/core/grammar/constraint.zig create mode 100644 zig/src/core/grammar/engine.zig create mode 100644 zig/src/core/grammar/golden_test.zig create mode 100644 zig/src/core/grammar/grammars.zig create mode 100644 zig/src/core/grammar/mask_kernel.zig create mode 100644 zig/src/core/grammar/pack.zig create mode 100644 zig/src/core/grammar/root.zig create mode 100644 zig/src/core/grammar/rows.zig create mode 100644 zig/src/core/grammar/tags.zig create mode 100644 zig/src/core/grammar/vocab.zig create mode 100644 zig/src/core/grammar/xgr.zig create mode 100644 zig/src/core/grammar/xgr_c.cc create mode 100644 zig/src/core/grammar/xgr_c.h create mode 100644 zig/tests/grammar/fixtures/synthetic.golden.json create mode 100644 zig/tests/grammar/fixtures/synthetic/tokenizer.json create mode 100644 zig/tests/grammar/gen_goldens.py create mode 100644 zig/vendor/xgrammar/3rdparty/dlpack/LICENSE create mode 100644 zig/vendor/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h create mode 100644 zig/vendor/xgrammar/3rdparty/picojson/picojson.h create mode 100644 zig/vendor/xgrammar/LICENSE create mode 100644 zig/vendor/xgrammar/NOTICE create mode 100644 zig/vendor/xgrammar/README.md create mode 100644 zig/vendor/xgrammar/cpp/compiled_grammar.cc create mode 100644 zig/vendor/xgrammar/cpp/compiled_grammar_impl.h create mode 100644 zig/vendor/xgrammar/cpp/config.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/cohere.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/deepseek.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/glm.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/kimi_k3.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/minimax.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/minimax_m3.cc create mode 100644 zig/vendor/xgrammar/cpp/converter_ext/qwen.cc create mode 100644 zig/vendor/xgrammar/cpp/earley_parser.cc create mode 100644 zig/vendor/xgrammar/cpp/earley_parser.h create mode 100644 zig/vendor/xgrammar/cpp/ebnf_script_creator.h create mode 100644 zig/vendor/xgrammar/cpp/fsm.cc create mode 100644 zig/vendor/xgrammar/cpp/fsm.h create mode 100644 zig/vendor/xgrammar/cpp/fsm_builder.cc create mode 100644 zig/vendor/xgrammar/cpp/fsm_builder.h create mode 100644 zig/vendor/xgrammar/cpp/grammar.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_builder.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_builder.h create mode 100644 zig/vendor/xgrammar/cpp/grammar_compiler.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_functor.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_functor.h create mode 100644 zig/vendor/xgrammar/cpp/grammar_impl.h create mode 100644 zig/vendor/xgrammar/cpp/grammar_matcher.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_parser.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_parser.h create mode 100644 zig/vendor/xgrammar/cpp/grammar_printer.cc create mode 100644 zig/vendor/xgrammar/cpp/grammar_printer.h create mode 100644 zig/vendor/xgrammar/cpp/json_schema_converter.cc create mode 100644 zig/vendor/xgrammar/cpp/json_schema_converter.h create mode 100644 zig/vendor/xgrammar/cpp/json_schema_converter_ext.h create mode 100644 zig/vendor/xgrammar/cpp/lark_converter.cc create mode 100644 zig/vendor/xgrammar/cpp/lark_converter.h create mode 100644 zig/vendor/xgrammar/cpp/regex_converter.cc create mode 100644 zig/vendor/xgrammar/cpp/regex_converter.h create mode 100644 zig/vendor/xgrammar/cpp/structural_tag.cc create mode 100644 zig/vendor/xgrammar/cpp/structural_tag.h create mode 100644 zig/vendor/xgrammar/cpp/suffix_automata.cc create mode 100644 zig/vendor/xgrammar/cpp/suffix_automata.h create mode 100644 zig/vendor/xgrammar/cpp/support/compact_2d_array.h create mode 100644 zig/vendor/xgrammar/cpp/support/container.h create mode 100644 zig/vendor/xgrammar/cpp/support/cpptrace.h create mode 100644 zig/vendor/xgrammar/cpp/support/dynamic_bitset.h create mode 100644 zig/vendor/xgrammar/cpp/support/encoding.h create mode 100644 zig/vendor/xgrammar/cpp/support/int_set.h create mode 100644 zig/vendor/xgrammar/cpp/support/json_parse.h create mode 100644 zig/vendor/xgrammar/cpp/support/json_serializer.h create mode 100644 zig/vendor/xgrammar/cpp/support/logging.cc create mode 100644 zig/vendor/xgrammar/cpp/support/logging.h create mode 100644 zig/vendor/xgrammar/cpp/support/memory_size.h create mode 100644 zig/vendor/xgrammar/cpp/support/recursion_guard.cc create mode 100644 zig/vendor/xgrammar/cpp/support/recursion_guard.h create mode 100644 zig/vendor/xgrammar/cpp/support/reflection.h create mode 100644 zig/vendor/xgrammar/cpp/support/thread_pool.h create mode 100644 zig/vendor/xgrammar/cpp/support/thread_safe_cache.h create mode 100644 zig/vendor/xgrammar/cpp/support/union_find_set.h create mode 100644 zig/vendor/xgrammar/cpp/support/utils.h create mode 100644 zig/vendor/xgrammar/cpp/testing.cc create mode 100644 zig/vendor/xgrammar/cpp/testing.h create mode 100644 zig/vendor/xgrammar/cpp/tokenizer_info.cc create mode 100644 zig/vendor/xgrammar/cpp/tokenizer_info_impl.h create mode 100644 zig/vendor/xgrammar/include/module.modulemap create mode 100644 zig/vendor/xgrammar/include/xgrammar/compiler.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/config.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/exception.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/grammar.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/matcher.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/object.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/tokenizer_info.h create mode 100644 zig/vendor/xgrammar/include/xgrammar/xgrammar.h diff --git a/NOTICE b/NOTICE index 50e45b9ac9..b91c3d4f5e 100644 --- a/NOTICE +++ b/NOTICE @@ -6,3 +6,6 @@ TensorFold is licensed under the Apache License, Version 2.0 (LICENSE) from 0.6. Releases up to 0.5.0 were MIT-licensed, and code written before 0.6.0 keeps its MIT notice in LICENSES/MIT.txt. Third-party code and its licenses are listed in THIRD_PARTY_NOTICES.md. + +This product includes XGrammar (zig/vendor/xgrammar), Copyright (c) 2024 by XGrammar Contributors, +licensed under the Apache License, Version 2.0. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index d02e4cea6d..142f7fefe7 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -93,6 +93,15 @@ Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships The MTP layer TensorFold drafts with comes from that checkpoint's last shard (MIT), converted by `families/deepseek_v4/convert.py`. +`zig/vendor/xgrammar/` is the C++ core of [XGrammar](https://github.com/mlc-ai/xgrammar) 0.2.8 (tag `v0.2.8`, commit +`97787376faee5ed8466cfad57c99855e4ce2f6aa`): the sdist's `cpp/`, `include/` and `3rdparty/` sources, with no local +patches (build files, tests and bindings left out; see its `README.md`). The Zig grammar module +(`zig/src/core/grammar/`) builds it with Zig's C++ toolchain and calls it through its own C ABI (`xgr_c.cc`). +XGrammar is Apache License 2.0, Copyright (c) 2024 by XGrammar Contributors (its `LICENSE` and `NOTICE` sit beside +it). It bundles [DLPack](https://github.com/dmlc/dlpack)'s header (Apache License 2.0, `3rdparty/dlpack/LICENSE`) and +[picojson](https://github.com/kazuho/picojson) (BSD-2-Clause, Copyright 2009-2010 Cybozu Labs, Inc. and 2011-2014 +Kazuho Oku; the text is in `3rdparty/picojson/picojson.h`). + TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card states Apache-2.0. The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use without derivatives. Each checkpoint keeps its own license. diff --git a/build.zig b/build.zig index 9940d84bca..a8e9944ec1 100644 --- a/build.zig +++ b/build.zig @@ -6,6 +6,7 @@ const qwen27_build = @import("zig/build/qwen27.zig"); const hip_build = @import("zig/build/hip.zig"); const tp_build = @import("zig/build/tp.zig"); const sessions_build = @import("zig/build/sessions.zig"); +const grammar_build = @import("zig/build/grammar.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -176,6 +177,7 @@ pub fn build(b: *std.Build) void { b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); tp_build.add(b, target, optimize, test_step); sessions_build.add(b, test_step); + grammar_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index a97aa37701..df6b5f9ab8 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -19,6 +19,7 @@ "zig/src/families/flashnext/roles_gen.zig", "zig/kernels/metal/flashnext/sources_gen.zig", ) +VENDORED = "zig/vendor/xgrammar/" # upstream XGrammar 0.2.8, byte-identical (its README); our wrappers stay checked def generator_output(path: Path) -> bool: @@ -29,9 +30,14 @@ def generator_output(path: Path) -> bool: return rel.startswith("zig/kernels/metal/flashnext/") and HEX_METAL.search(path.name) is not None +def vendored(path: Path) -> bool: + """A file of the vendored XGrammar copy (exactly the VENDORED prefix).""" + return path.relative_to(ROOT).as_posix().startswith(VENDORED) + + def problems(path: Path) -> list[str]: """Every rule break in one file, as 'path:line: reason'.""" - if generator_output(path): + if generator_output(path) or vendored(path): return [] lines = path.read_text(errors="replace").splitlines() rel, found = path.relative_to(ROOT), [] diff --git a/tools/zig/lean_check_test.py b/tools/zig/lean_check_test.py new file mode 100644 index 0000000000..9e2822d833 --- /dev/null +++ b/tools/zig/lean_check_test.py @@ -0,0 +1,38 @@ +"""The lean checker's exemptions: the vendored prefix is skipped, a file beside it is still checked.""" +import sys +import tempfile +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import lean_check # noqa: E402 + +BAD = "// one\n// two\n" # two full-line comments in a row: one finding + + +class Vendored(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.root = Path(self.tmp.name) + self.saved, lean_check.ROOT = lean_check.ROOT, self.root + + def tearDown(self): + lean_check.ROOT = self.saved + self.tmp.cleanup() + + def write(self, rel: str) -> Path: + path = self.root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(BAD) + return path + + def test_vendored_file_is_skipped(self): + self.assertEqual(lean_check.problems(self.write(lean_check.VENDORED + "cpp/fsm.h")), []) + + def test_file_beside_the_vendored_prefix_is_checked(self): + for rel in ("zig/vendor/xgrammar_glue.h", "zig/vendor/other/fsm.h", "zig/src/core/grammar/xgr_c.h"): + self.assertEqual(len(lean_check.problems(self.write(rel))), 1, rel) + + +if __name__ == "__main__": + unittest.main() diff --git a/zig/build/grammar.zig b/zig/build/grammar.zig new file mode 100644 index 0000000000..020ac57f1a --- /dev/null +++ b/zig/build/grammar.zig @@ -0,0 +1,59 @@ +//! xgrammar 0.2.8's C++ core (zig/vendor/xgrammar, Apache-2.0) with the C ABI of xgr_c.cc, and the `grammar` module. + +const std = @import("std"); + +const xgr_sources = [_][]const u8{ + "compiled_grammar.cc", "config.cc", "earley_parser.cc", "fsm.cc", + "fsm_builder.cc", "grammar.cc", "grammar_builder.cc", "grammar_compiler.cc", + "grammar_functor.cc", "grammar_matcher.cc", "grammar_parser.cc", "grammar_printer.cc", + "json_schema_converter.cc", "lark_converter.cc", "regex_converter.cc", "structural_tag.cc", + "suffix_automata.cc", "testing.cc", "tokenizer_info.cc", "support/logging.cc", + "support/recursion_guard.cc", "converter_ext/cohere.cc", "converter_ext/deepseek.cc", "converter_ext/glm.cc", + "converter_ext/kimi_k3.cc", "converter_ext/minimax.cc", "converter_ext/minimax_m3.cc", "converter_ext/qwen.cc", +}; + +/// xgrammar's static library with the C ABI, always ReleaseFast (a debug Earley parser is too slow for the tests). +pub fn library(b: *std.Build, target: std.Build.ResolvedTarget) *std.Build.Step.Compile { + const root = "zig/vendor/xgrammar/"; + const m = b.createModule(.{ .target = target, .optimize = .ReleaseFast, .link_libc = true, .link_libcpp = true }); + m.addIncludePath(b.path(root ++ "include")); + m.addIncludePath(b.path(root ++ "cpp")); + m.addSystemIncludePath(b.path(root ++ "3rdparty/picojson")); + m.addSystemIncludePath(b.path(root ++ "3rdparty/dlpack/include")); + m.addIncludePath(b.path("zig/src/core/grammar")); + const flags = [_][]const u8{ "-std=c++17", "-DXGRAMMAR_ENABLE_CPPTRACE=0", "-fexceptions", "-Wno-everything" }; + m.addCSourceFiles(.{ .root = b.path(root ++ "cpp"), .files = &xgr_sources, .flags = &flags }); + m.addCSourceFile(.{ .file = b.path("zig/src/core/grammar/xgr_c.cc"), .flags = &flags }); + return b.addLibrary(.{ .name = "tf_xgrammar", .root_module = m, .linkage = .static }); +} + +/// The `grammar` module (zig/src/core/grammar/root.zig) linked with xgrammar. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/core/grammar/root.zig"), .target = target, .optimize = optimize, .link_libc = true, .link_libcpp = true }); + m.addIncludePath(b.path("zig/src/core/grammar")); + m.linkLibrary(library(b, target)); + return m; +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const host = b.graph.host; + const t = b.addTest(.{ .root_module = module(b, host, .debug) }); + const run = b.addRunArtifact(t); + run.setCwd(b.path(".")); + for ([_][]const u8{ "synthetic.golden.json", "synthetic/tokenizer.json" }) |f| run.addFileInput(b.path(b.fmt("zig/tests/grammar/fixtures/{s}", .{f}))); + b.step("test-grammar", "Grammar masks / cuts / compiles vs xgrammar 0.2.8's Python package on recorded goldens (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} + +/// The mask kernel as sm_80 PTX (JIT-compiled by the driver), for `@embedFile` as the import `grammar_mask_ptx`. +pub fn maskPtx(b: *std.Build) std.Build.LazyPath { + const nvptx = b.resolveTargetQuery(.{ .cpu_arch = .nvptx64, .os_tag = .cuda, .cpu_model = .{ .explicit = &std.Target.nvptx.cpu.sm_80 } }); + const obj = b.addObject(.{ .name = "tf_grammar_mask", .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/core/grammar/mask_kernel.zig"), .target = nvptx, .optimize = .ReleaseFast }) }); + return obj.getEmittedAsm(); +} + +/// Adds `grammar` (and the mask kernel's PTX as `grammar_mask_ptx`) to a module that serves grammars. +pub fn addTo(b: *std.Build, m: *std.Build.Module, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) void { + m.addImport("grammar", module(b, target, optimize)); + m.addAnonymousImport("grammar_mask_ptx", .{ .root_source_file = maskPtx(b) }); +} diff --git a/zig/src/core/grammar/constraint.zig b/zig/src/core/grammar/constraint.zig new file mode 100644 index 0000000000..ea32481445 --- /dev/null +++ b/zig/src/core/grammar/constraint.zig @@ -0,0 +1,120 @@ +//! One reply's grammar: `cut` drops rejected drafts, `fill` writes allowed bits, `advance` follows chosen tokens. + +const std = @import("std"); +const engine = @import("engine.zig"); + +pub const Error = engine.Error || error{ Rejected, Empty, RowsChanged }; + +/// A chain window as the grammar keeps it: rows [0, keep) of the window run; rows [first, keep) are constrained. +pub const Cut = struct { + keep: u32, + first: u32, + stop: bool = false, // row `keep` is the grammar's stop token: it and every later row are cut + + pub fn rows(c: Cut) u32 { + return c.keep - c.first; + } +}; + +pub const Constraint = struct { + m: engine.Matcher, + /// tokens the grammar never takes (special tokens other than the stop tokens): a draft of one is cut unasked + never: *const std.DynamicBitSetUnmanaged, + think_end: u32, + active: bool, // false: unconstrained until `think_end` is chosen + words: u32, + + /// A fresh reply over `compiled`; `active` false: the grammar starts after `think_end` is chosen. + pub fn init(compiled: engine.Compiled, never: *const std.DynamicBitSetUnmanaged, think_end: u32, active: bool, words: u32) engine.Error!Constraint { + return .{ .m = try compiled.matcher(), .never = never, .think_end = think_end, .active = active, .words = words }; + } + + pub fn deinit(c: *Constraint) void { + c.m.deinit(); + } + + /// The grammar has taken its stop token: the reply is complete (nothing after it is constrained). + pub fn finished(c: *const Constraint) bool { + return c.active and c.m.terminated(); + } + + fn never_(c: *const Constraint, t: u32) bool { + return t < c.never.bit_length and c.never.isSet(t); + } + + /// Window rows (row 0: the pending token, already followed) an accepted path can use; the matcher ends unchanged. + pub fn cut(c: *Constraint, tokens: []const u32) Error!Cut { + const n: u32 = @intCast(tokens.len); + if (c.finished()) return .{ .keep = n, .first = n }; + var active = c.active; + var first: u32 = if (active) 0 else n; + var keep: u32 = 1; + var accepted: u32 = 0; + var stop = false; + defer if (accepted > 0) c.m.rollback(accepted) catch {}; + for (tokens[1..], 1..) |t, r| { + if (active) { + if (c.never_(t) or !try c.m.accept(t)) break; // the parent row's masked draw could never choose it + accepted += 1; + if (c.m.terminated()) { // its stop token: nothing is verified after it + stop = true; + break; + } + } else { + active = t == c.think_end; // the grammar starts at the row after + if (active) first = @intCast(r); + } + keep += 1; + } + if (first > keep) first = keep; + return .{ .keep = keep, .first = first, .stop = stop }; + } + + /// Each constrained row r's allowed bits at `bits[(r - k.first) * words ..]`; error.Empty when a row allows none. + pub fn fill(c: *Constraint, tokens: []const u32, k: Cut, bits: []u32) Error!void { + if (k.rows() == 0) return; + if (bits.len < @as(usize, k.rows()) * c.words) return error.RowsChanged; + var active = c.active; + var accepted: u32 = 0; + var j: u32 = 0; + defer if (accepted > 0) c.m.rollback(accepted) catch {}; + for (tokens[0..k.keep], 0..) |t, r| { + if (r > 0) { + if (active) { + if (!try c.m.accept(t)) return error.Rejected; + accepted += 1; + } else active = t == c.think_end; + } + if (!active) continue; + if (r != k.first + j) return error.RowsChanged; + const row = bits[@as(usize, j) * c.words ..][0..c.words]; + try c.m.fill(row); + if (std.mem.allEqual(u32, row, 0)) return error.Empty; + j += 1; + } + if (j != k.rows()) return error.RowsChanged; + } + + /// Follow chosen tokens (each chosen under this grammar's mask); after the stop token nothing follows. + pub fn advance(c: *Constraint, tokens: []const u32) Error!void { + for (tokens) |t| { + if (!c.active) { + c.active = t == c.think_end; + continue; + } + if (c.m.terminated()) return; + if (!try c.m.accept(t)) return error.Rejected; + } + } +}; + +/// False when `prompt` ends inside an open think block (its last `open` after its last `end`). +pub fn thinkActive(prompt: []const u32, open: ?u32, end: ?u32) bool { + var i = prompt.len; + while (i > 0) { + i -= 1; + if (end != null and prompt[i] == end.?) return true; + if (open != null and prompt[i] == open.?) return false; + } + return true; +} diff --git a/zig/src/core/grammar/engine.zig b/zig/src/core/grammar/engine.zig new file mode 100644 index 0000000000..1aeb9f6ca7 --- /dev/null +++ b/zig/src/core/grammar/engine.zig @@ -0,0 +1,70 @@ +//! Grammar engine interfaces: `Compiler` -> `Compiled` -> `Matcher`s that walk tokens and fill allowed bits (xgr.zig). + +const std = @import("std"); + +pub const Error = error{ Xgrammar, OutOfMemory }; + +/// What xgrammar compiles: a JSON schema (with ``max_whitespace_cnt``), a regex, EBNF (root ``root``) or a tag JSON. +pub const Source = enum { json_schema, regex, ebnf, structural_tag }; + +pub const Compiler = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + compile: *const fn (ptr: *anyopaque, kind: Source, text: []const u8, max_whitespace: ?u32) Error!Compiled, + }; + + pub fn compile(x: Compiler, kind: Source, text: []const u8, max_whitespace: ?u32) Error!Compiled { + return x.vtable.compile(x.ptr, kind, text, max_whitespace); + } +}; + +pub const Compiled = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + matcher: *const fn (ptr: *anyopaque) Error!Matcher, + deinit: *const fn (ptr: *anyopaque) void, + }; + + /// A fresh matcher at the grammar's start. + pub fn matcher(g: Compiled) Error!Matcher { + return g.vtable.matcher(g.ptr); + } + pub fn deinit(g: Compiled) void { + g.vtable.deinit(g.ptr); + } +}; + +/// One reply's walk through a grammar. +pub const Matcher = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + accept: *const fn (ptr: *anyopaque, token: u32) Error!bool, + /// The next token's allowed bits (token t: bit t % 32 of word t / 32), one bitmask row. + fill: *const fn (ptr: *anyopaque, words: []u32) Error!void, + rollback: *const fn (ptr: *anyopaque, n: u32) Error!void, + terminated: *const fn (ptr: *anyopaque) bool, + deinit: *const fn (ptr: *anyopaque) void, + }; + + pub fn accept(m: Matcher, token: u32) Error!bool { + return m.vtable.accept(m.ptr, token); + } + pub fn fill(m: Matcher, words: []u32) Error!void { + return m.vtable.fill(m.ptr, words); + } + pub fn rollback(m: Matcher, n: u32) Error!void { + return m.vtable.rollback(m.ptr, n); + } + pub fn terminated(m: Matcher) bool { + return m.vtable.terminated(m.ptr); + } + pub fn deinit(m: Matcher) void { + m.vtable.deinit(m.ptr); + } +}; diff --git a/zig/src/core/grammar/golden_test.zig b/zig/src/core/grammar/golden_test.zig new file mode 100644 index 0000000000..50eff101d9 --- /dev/null +++ b/zig/src/core/grammar/golden_test.zig @@ -0,0 +1,169 @@ +//! `test-grammar`: gen_goldens.py's cuts, mask digests, structural tags and refusals replayed through `Grammars`. + +const std = @import("std"); +const testing = std.testing; +const grammars_mod = @import("grammars.zig"); +const vocab_mod = @import("vocab.zig"); +const tags = @import("tags.zig"); +const xgr = @import("xgr.zig"); +const Constraint = @import("constraint.zig").Constraint; + +const fixtures = "zig/tests/grammar/fixtures/"; + +const Step = struct { + window: ?[]const u32 = null, + keep: u32 = 0, + first: u32 = 0, + stop: bool = false, + rows: []const [2]std.json.Value = &.{}, + advance: []const u32, +}; +const Case = struct { name: []const u8, kind: []const u8, text: []const u8, tag: ?[]const u8, active: bool, steps: []const Step, finished: bool }; +const Bad = struct { name: []const u8, kind: []const u8, text: []const u8, field: []const u8, message: ?[]const u8 }; +const View = struct { vocab_type: i32, add_prefix_space: bool, vocab: []const u8, never_digest: []const u8, never_count: u32 }; +const Golden = struct { + vocab_size: u32, + eos: u32, + think_open: ?u32, + think_end: ?u32, + words: u32, + views: struct { text: View, tools: View }, + cases: []const Case, + bad: []const Bad, +}; + +fn hex16(bytes: []const u8) [16]u8 { + var d: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(bytes, &d, .{}); + var out: [16]u8 = undefined; + _ = std.fmt.bufPrint(&out, "{x}", .{d[0..8]}) catch unreachable; + return out; +} + +fn vocabDigest(tokens: []const []const u8) [16]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (tokens) |s| { + var n: [4]u8 = undefined; + std.mem.writeInt(u32, &n, @intCast(s.len), .little); + h.update(&n); + h.update(s); + } + var d: [32]u8 = undefined; + h.final(&d); + var out: [16]u8 = undefined; + _ = std.fmt.bufPrint(&out, "{x}", .{d[0..8]}) catch unreachable; + return out; +} + +fn read(a: std.mem.Allocator, path: []const u8) ![]u8 { + return std.Io.Dir.cwd().readFileAlloc(testing.io, path, a, .limited(1 << 30)); +} + +fn run(golden_path: []const u8, tokenizer_path: []const u8) !void { + const gpa = testing.allocator; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const g = try std.json.parseFromSliceLeaky(Golden, a, try read(a, golden_path), .{ .ignore_unknown_fields = true }); + const tok_json = try read(a, tokenizer_path); + + // the views: the vocabulary each gives xgrammar, its metadata, the tokens it never takes + const meta = try xgr.detectMetadata(gpa, tok_json); + const dsml = [_][]const u8{"|DSML|"}; + for ([_]View{ g.views.text, g.views.tools }, [_][]const []const u8{ &.{}, &dsml }) |want, keep| { + try testing.expectEqual(want.vocab_type, meta.vocab_type); + try testing.expectEqual(want.add_prefix_space, meta.add_prefix_space); + var v = try vocab_mod.build(gpa, tok_json, g.vocab_size, keep); + defer v.deinit(); + try testing.expectEqualStrings(want.vocab, &vocabDigest(v.tokens)); + try testing.expectEqual(g.think_open, v.think_open); + try testing.expectEqual(g.think_end, v.think_end); + } + const gs = try grammars_mod.Grammars.init(gpa, tok_json, .{ .vocab_size = g.vocab_size, .stops = &.{g.eos}, .tool_tokens = &dsml, .tag_model = .deepseek_v4_1 }); + defer gs.deinit(); + try testing.expectEqual(g.words, gs.words); + for ([_]View{ g.views.text, g.views.tools }, [_]grammars_mod.Kind{ .json, .tools }) |want, kind| { + const never = &gs.views[if (kind == .tools) 1 else 0].never; + try testing.expectEqual(want.never_count, @as(u32, @intCast(never.count()))); + var list: std.Io.Writer.Allocating = .init(gpa); + defer list.deinit(); + try list.writer.writeByte('['); + var it = never.iterator(.{}); + var i: usize = 0; + while (it.next()) |t| : (i += 1) try list.writer.print("{s}{d}", .{ if (i > 0) ", " else "", t }); + try list.writer.writeByte(']'); + try testing.expectEqualStrings(want.never_digest, &hex16(list.written())); + } + + const bits = try gpa.alloc(u32, 16 * gs.words); + defer gpa.free(bits); + var rows_checked: usize = 0; + for (g.cases) |c| { + const kind = std.meta.stringToEnum(grammars_mod.Kind, c.kind).?; + if (c.tag) |want| { + var problem: tags.Problem = .{}; + const got = try tags.build(gpa, .deepseek_v4_1, c.text, &problem); + defer gpa.free(got); + testing.expectEqualStrings(want, got) catch |e| { + std.debug.print("case {s}: the structural tag differs\n", .{c.name}); + return e; + }; + } + var why: grammars_mod.Grammars.Why = .{}; + const compiled = gs.compile(.{ .kind = kind, .text = c.text }, &why) catch |e| { + std.debug.print("case {s}: {s}\n", .{ c.name, why.text() }); + return e; + }; + defer compiled.deinit(); + var con = try gs.bound(kind, compiled, c.active); + defer con.deinit(); + for (c.steps, 0..) |s, si| { + if (s.window) |w| { + const k = try con.cut(w); + const bad = k.keep != s.keep or (k.rows() > 0 and k.first != s.first) or k.stop != s.stop; + if (bad) { + std.debug.print("case {s} step {d}: cut keep {d} first {d} stop {} vs reference {d} {d} {}\n", .{ c.name, si, k.keep, k.first, k.stop, s.keep, s.first, s.stop }); + return error.TestUnexpectedResult; + } + try testing.expectEqual(s.rows.len, k.rows()); + try con.fill(w, k, bits); + for (s.rows, 0..) |want, j| { + const row = bits[j * gs.words ..][0..gs.words]; + var allowed: usize = 0; + for (row, 0..) |word, wi| { + const valid: u32 = @min(32, g.vocab_size - @as(u32, @intCast(32 * wi))); + const m: u32 = if (valid == 32) 0xffff_ffff else (@as(u32, 1) << @intCast(valid)) - 1; + allowed += @popCount(word & m); + } + const digest = hex16(std.mem.sliceAsBytes(row)); + if (!std.mem.eql(u8, &digest, want[0].string) or allowed != @as(usize, @intCast(want[1].integer))) { + std.debug.print("case {s} step {d} row {d}: mask {s} ({d} allowed) vs reference {s} ({d})\n", .{ c.name, si, j, &digest, allowed, want[0].string, want[1].integer }); + return error.TestUnexpectedResult; + } + rows_checked += 1; + } + } + try con.advance(s.advance); + } + try testing.expectEqual(c.finished, con.finished()); + } + for (g.bad) |b| { + const kind = std.meta.stringToEnum(grammars_mod.Kind, b.kind).?; + var why: grammars_mod.Grammars.Why = .{}; + if (gs.compile(.{ .kind = kind, .text = b.text }, &why)) |compiled| { + compiled.deinit(); + try testing.expect(b.message == null); + } else |e| { + try testing.expectEqual(error.Invalid, e); + const got = try std.fmt.allocPrint(gpa, "{s}: the grammar cannot be enforced: {s}", .{ b.field, why.text() }); + defer gpa.free(got); + try testing.expectEqualStrings(b.message.?, got); + } + } + // every case has masked rows (a walk that never masked would prove nothing) + try testing.expect(rows_checked >= 10 * g.cases.len); +} + +test "grammar goldens: synthetic tokenizer" { + try run(fixtures ++ "synthetic.golden.json", fixtures ++ "synthetic/tokenizer.json"); +} diff --git a/zig/src/core/grammar/grammars.zig b/zig/src/core/grammar/grammars.zig new file mode 100644 index 0000000000..de80b61973 --- /dev/null +++ b/zig/src/core/grammar/grammars.zig @@ -0,0 +1,219 @@ +//! A checkpoint's grammar compilers over two tokenizer views (`text`: no added tokens; `tools`: call markup kept). + +const std = @import("std"); +const engine = @import("engine.zig"); +const xgr = @import("xgr.zig"); +const vocab_mod = @import("vocab.zig"); +const tags = @import("tags.zig"); +const constraint = @import("constraint.zig"); + +/// The request kinds (their order is the packed form's). +pub const Kind = enum(u8) { json, json_schema, regex, choice, grammar, tools }; + +/// A request's structured output: its kind and canonical text (json: empty; choice: a JSON array; tools: tags.zig). +pub const Spec = struct { kind: Kind, text: []const u8 = "" }; + +/// json_object: any JSON object (OpenAI's contract), not an array. +pub const object_schema = "{\"type\": \"object\"}"; +/// The most blank characters between JSON tokens: pretty-printing, not endless. +pub const blanks = 32; +/// Compiled grammars each view's compiler keeps, by their text. +pub const cache_bytes: i64 = 128 << 20; + +pub const Options = struct { + /// the logits' columns (config.json's vocab_size) + vocab_size: u32, + /// the engine's eos ids: allowed where the grammar may end + stops: []const u32, + /// the tools view's kept added tokens (the model's call markup) and its structural tag + tool_tokens: []const []const u8, + tag_model: tags.Model, +}; + +pub const Grammars = struct { + gpa: std.mem.Allocator, + vocab_size: u32, + words: u32, + think_open: ?u32, + think_end: ?u32, + views: [2]View, + + const View = struct { + tok: xgr.Tokenizer, + compiler: xgr.Compiler, + /// tokens no grammar of this view takes (special, not a stop token) + never: std.DynamicBitSetUnmanaged, + + fn deinit(v: *View, gpa: std.mem.Allocator) void { + v.compiler.deinit(); + v.tok.deinit(); + v.never.deinit(gpa); + } + }; + + /// The compilers for a tokenizer.json text. + pub fn init(gpa: std.mem.Allocator, tokenizer_json: []const u8, o: Options) !*Grammars { + const g = try gpa.create(Grammars); + errdefer gpa.destroy(g); + const meta = try xgr.detectMetadata(gpa, tokenizer_json); + const stops = try gpa.alloc(i32, o.stops.len); + defer gpa.free(stops); + for (stops, o.stops) |*s, t| s.* = @intCast(t); + g.* = .{ .gpa = gpa, .vocab_size = o.vocab_size, .words = xgr.bitmaskWords(o.vocab_size), .think_open = null, .think_end = null, .views = undefined }; + var made: usize = 0; + errdefer for (g.views[0..made]) |*v| v.deinit(gpa); + for (&g.views, [_][]const []const u8{ &.{}, o.tool_tokens }) |*v, keep| { + var voc = try vocab_mod.build(gpa, tokenizer_json, o.vocab_size, keep); + defer voc.deinit(); + g.think_open = voc.think_open; + g.think_end = voc.think_end; + const tok = try xgr.Tokenizer.init(gpa, voc.tokens, meta, stops); + errdefer tok.deinit(); + var never = try std.DynamicBitSetUnmanaged.initEmpty(gpa, o.vocab_size); + errdefer never.deinit(gpa); + const special = try tok.special(gpa); + defer gpa.free(special); + for (special) |t| if (t >= 0 and t < o.vocab_size and std.mem.indexOfScalar(u32, o.stops, @intCast(t)) == null) never.set(@intCast(t)); + v.* = .{ .tok = tok, .compiler = try xgr.Compiler.init(tok, o.vocab_size, cache_bytes), .never = never }; + made += 1; + } + return g; + } + + pub fn deinit(g: *Grammars) void { + for (&g.views) |*v| v.deinit(g.gpa); + g.gpa.destroy(g); + } + + fn view(g: *Grammars, kind: Kind) *View { + return &g.views[if (kind == .tools) 1 else 0]; + } + + /// Why a spec does not compile: xgrammar's trimmed message or the tools spec's problem. + pub const Why = struct { + buf: [1024]u8 = undefined, + len: usize = 0, + + pub fn text(w: *const Why) []const u8 { + return w.buf[0..w.len]; + } + fn set(w: *Why, s: []const u8) void { + w.len = @min(s.len, w.buf.len); + @memcpy(w.buf[0..w.len], s[0..w.len]); + } + }; + + /// The compiled grammar of `spec` (cached by text); error.Invalid with `why` when it cannot be enforced. + pub fn compile(g: *Grammars, spec: Spec, why: *Why) !engine.Compiled { + const c = g.view(spec.kind).compiler.compiler(); + var problem: tags.Problem = .{}; + return switch (spec.kind) { + .json => compiled(c.compile(.json_schema, object_schema, blanks), why), + .json_schema => compiled(c.compile(.json_schema, spec.text, blanks), why), + .regex => compiled(c.compile(.regex, spec.text, null), why), + .grammar => compiled(c.compile(.ebnf, spec.text, null), why), + .choice => { + const ebnf = choiceGrammar(g.gpa, spec.text) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => { + why.set("choice must be a JSON array of strings"); + return error.Invalid; + }, + }; + defer g.gpa.free(ebnf); + return compiled(c.compile(.ebnf, ebnf, null), why); + }, + .tools => { + const tag = tags.build(g.gpa, .deepseek_v4_1, spec.text, &problem) catch |e| switch (e) { + error.InvalidTools => { + why.set(problem.message); + return error.Invalid; + }, + else => |x| return x, + }; + defer g.gpa.free(tag); + return compiled(c.compile(.structural_tag, tag, null), why); + }, + }; + } + + fn compiled(r: engine.Error!engine.Compiled, why: *Why) !engine.Compiled { + return r catch |e| switch (e) { + error.OutOfMemory => error.OutOfMemory, + error.Xgrammar => { + why.set(message(xgr.lastError())); + return error.Invalid; + }, + }; + } + + /// A fresh reply's grammar state for `compiled`, bound to its prompt's think state. + pub fn constraint_(g: *Grammars, kind: Kind, compiled_: engine.Compiled, prompt: []const u32) !constraint.Constraint { + const active = constraint.thinkActive(prompt, g.think_open, g.think_end); + return g.bound(kind, compiled_, active); + } + + /// A fresh reply's state at a given think state (a follower's, from the leader's packed grammar). + pub fn bound(g: *Grammars, kind: Kind, compiled_: engine.Compiled, active: bool) !constraint.Constraint { + return constraint.Constraint.init(compiled_, &g.view(kind).never, g.think_end orelse std.math.maxInt(u32), active, g.words); + } + + pub fn tokenizer(g: *Grammars, kind: Kind) xgr.Tokenizer { + return g.view(kind).tok; + } +}; + +/// xgrammar's message without its timestamp and source location: the first line, ``^\[[^\]]*\]\s*\S+:\d+:\s*`` removed. +pub fn message(full: []const u8) []const u8 { + const trimmed = std.mem.trim(u8, full, " \t\r\n"); + var s = trimmed[0 .. std.mem.indexOfScalar(u8, trimmed, '\n') orelse trimmed.len]; + if (s.len == 0) return "XGrammarError"; + if (s[0] != '[') return s; + const close = std.mem.indexOfScalar(u8, s, ']') orelse return s; + var i = close + 1; + while (i < s.len and std.ascii.isWhitespace(s[i])) i += 1; + const loc = i; + while (i < s.len and !std.ascii.isWhitespace(s[i])) i += 1; + // \S+:\d+: — the location's last '::' ends it (\S+ is greedy, backtracking to the last such match) + const word = s[loc..i]; + var end: ?usize = null; + var k: usize = word.len; + while (k > 0) : (k -= 1) { + if (word[k - 1] != ':') continue; + var d = k - 1; + while (d > 0 and std.ascii.isDigit(word[d - 1])) d -= 1; + if (d < k - 1 and d > 1 and word[d - 1] == ':') { + end = k; + break; + } + } + if (end == null) { + // the location may be followed by the message in the same word (no space after the colon) + return s; + } + var j = loc + end.?; + while (j < s.len and std.ascii.isWhitespace(s[j])) j += 1; + s = s[j..]; + return s; +} + +/// ``root ::= "a" | "b"``: each option a JSON string, non-ASCII written raw. +pub fn choiceGrammar(gpa: std.mem.Allocator, choices_json: []const u8) ![]u8 { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, choices_json, .{}); + defer parsed.deinit(); + if (parsed.value != .array) return error.InvalidChoices; + var out: std.Io.Writer.Allocating = .init(gpa); + errdefer out.deinit(); + try out.writer.writeAll("root ::= "); + for (parsed.value.array.items, 0..) |v, i| { + if (v != .string) return error.InvalidChoices; + if (i > 0) try out.writer.writeAll(" | "); + try pyString(&out.writer, v.string); + } + return out.toOwnedSlice(); +} + +/// A JSON string with only `"`, `\` and control characters escaped (non-ASCII raw). +pub fn pyString(w: *std.Io.Writer, s: []const u8) !void { + try tags.string(w, s); +} diff --git a/zig/src/core/grammar/mask_kernel.zig b/zig/src/core/grammar/mask_kernel.zig new file mode 100644 index 0000000000..e58f03373b --- /dev/null +++ b/zig/src/core/grammar/mask_kernel.zig @@ -0,0 +1,13 @@ +//! The grammar mask kernel (PTX from Zig's NVPTX backend): grid (ceil(span / 128), rows), a thread per 32-token word. + +const rows = @import("rows.zig"); + +/// Logits rows idx[r] (stride ld) under bits row r: this rank's `n` words from `w0`, `width` columns. +export fn tf_grammar_mask(logits: [*]addrspace(.global) f32, ld: u64, idx: [*]addrspace(.global) const u32, bits: [*]addrspace(.global) const u32, words: u32, w0: u32, n: u32, width: u32) callconv(.kernel) void { + const r = @workGroupId(1); + const k = @workGroupId(0) * @workGroupSize(0) + @workItemId(0); + if (k >= n) return; + const word = bits[@as(u64, r) * words + w0 + k]; + const base: [*]f32 = @addrSpaceCast(logits + @as(u64, idx[r]) * ld); + rows.maskWord(base, word, 32 * k, width); +} diff --git a/zig/src/core/grammar/pack.zig b/zig/src/core/grammar/pack.zig new file mode 100644 index 0000000000..3832ebd603 --- /dev/null +++ b/zig/src/core/grammar/pack.zig @@ -0,0 +1,52 @@ +//! A request's grammar as non-negative int64 words for another rank, so both ranks compile and fill the same masks. + +const std = @import("std"); +const grammars = @import("grammars.zig"); + +/// Words `encode` writes for a spec of `text_len` bytes. +pub fn len(text_len: usize) usize { + return 3 + (text_len + 6) / 7; +} + +/// [kind, active, n, bytes 7 a word...] into `out` (len(spec.text.len) words). +pub fn encode(spec: grammars.Spec, active: bool, out: []i64) void { + out[0] = @backingInt(spec.kind); + out[1] = @intFromBool(active); + out[2] = @intCast(spec.text.len); + var i: usize = 0; + while (i < spec.text.len) : (i += 7) { + var v: i64 = 0; + const chunk = spec.text[i..@min(i + 7, spec.text.len)]; + for (chunk, 0..) |b, k| v |= @as(i64, b) << @intCast(8 * k); + out[3 + i / 7] = v; + } +} + +pub const Decoded = struct { spec: grammars.Spec, active: bool, used: usize }; + +/// The spec `encode` wrote at the front of `words`; its text in `buf` (resized). +pub fn decode(gpa: std.mem.Allocator, words: []const i64, buf: *std.ArrayList(u8)) !Decoded { + if (words.len < 3) return error.BadPlan; + const kind = std.enums.fromInt(grammars.Kind, words[0]) orelse return error.BadPlan; + const n: usize = @intCast(words[2]); + const used = len(n); + if (words.len < used) return error.BadPlan; + try buf.resize(gpa, n); + for (buf.items, 0..) |*b, i| b.* = @truncate(@as(u64, @bitCast(words[3 + i / 7])) >> @intCast(8 * (i % 7))); + return .{ .spec = .{ .kind = kind, .text = buf.items }, .active = words[1] != 0, .used = used }; +} + +test "pack round trip" { + const gpa = std.testing.allocator; + const text = "{\"type\": \"object\", \"é\": \"|DSML|\"}"; + var words: [len(text.len)]i64 = undefined; + encode(.{ .kind = .json_schema, .text = text }, true, &words); + for (words) |w| try std.testing.expect(w >= 0); + var buf: std.ArrayList(u8) = .empty; + defer buf.deinit(gpa); + const d = try decode(gpa, &words, &buf); + try std.testing.expectEqual(grammars.Kind.json_schema, d.spec.kind); + try std.testing.expect(d.active); + try std.testing.expectEqualStrings(text, d.spec.text); + try std.testing.expectEqual(words.len, d.used); +} diff --git a/zig/src/core/grammar/root.zig b/zig/src/core/grammar/root.zig new file mode 100644 index 0000000000..6d569cd65d --- /dev/null +++ b/zig/src/core/grammar/root.zig @@ -0,0 +1,22 @@ +//! Grammar-constrained decoding over xgrammar's C++ core: `Grammars`, `Constraint`, `engine`/`xgr` and `rows`. + +pub const engine = @import("engine.zig"); +pub const xgr = @import("xgr.zig"); +pub const vocab = @import("vocab.zig"); +pub const tags = @import("tags.zig"); +pub const constraint = @import("constraint.zig"); +pub const grammars = @import("grammars.zig"); +pub const pack = @import("pack.zig"); +pub const rows = @import("rows.zig"); + +pub const Grammars = grammars.Grammars; +pub const Spec = grammars.Spec; +pub const Kind = grammars.Kind; +pub const Constraint = constraint.Constraint; +pub const Cut = constraint.Cut; + +test { + _ = @import("golden_test.zig"); + _ = pack; + _ = rows; +} diff --git a/zig/src/core/grammar/rows.zig b/zig/src/core/grammar/rows.zig new file mode 100644 index 0000000000..e24d3b02d9 --- /dev/null +++ b/zig/src/core/grammar/rows.zig @@ -0,0 +1,62 @@ +//! A grammar bitmask on a rank's logits columns (disallowed tokens -inf before top-k), device and host forms. + +const std = @import("std"); + +/// Row `row`'s columns [col0, col0 + 32) under `word` (bit b: token col0 + b allowed); columns >= width are skipped. +pub inline fn maskWord(row: [*]f32, word: u32, col0: u32, width: u32) void { + if (word == 0xffff_ffff) return; + var b: u32 = 0; + while (b < 32) : (b += 1) { + const col = col0 + b; + if (col < width and (word >> @intCast(b)) & 1 == 0) row[col] = -std.math.inf(f32); + } +} + +/// Words of this rank's columns: [lo / 32, (lo + width + 31) / 32). `lo` must sit on a word (others are refused). +pub fn span(lo: u32, width: u32) struct { w0: u32, n: u32 } { + return .{ .w0 = lo / 32, .n = (width + 31) / 32 }; +} + +/// Host form: logits rows `idx[r]` (stride `ld`) under bitmask rows `bits[r * words ..]`, columns from token `lo`. +pub fn maskRows(logits: []f32, ld: usize, idx: []const u32, bits: []const u32, words: u32, lo: u32, width: u32) void { + const s = span(lo, width); + for (idx, 0..) |row, r| { + const base = logits[row * ld ..].ptr; + const words_r = bits[r * words + s.w0 ..][0..s.n]; + for (words_r, 0..) |w, k| maskWord(base, w, @intCast(32 * k), width); + } +} + +test "maskRows equals the per-token rule on a rank's half" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(3); + const rnd = prng.random(); + const vocab: u32 = 1000; + const words: u32 = (vocab + 31) / 32; + const world: u32 = 2; + const width: u32 = 512; // the halves of a 1024-column head (the last 24 columns past the vocabulary) + const n = 5; + const bits = try gpa.alloc(u32, n * words); + defer gpa.free(bits); + for (bits) |*b| b.* = switch (rnd.uintLessThan(u32, 3)) { + 0 => 0xffff_ffff, + 1 => 0, + else => rnd.int(u32), + }; + for (0..world) |rank| { + const lo: u32 = @intCast(rank * width); + const logits = try gpa.alloc(f32, 8 * width); + defer gpa.free(logits); + for (logits) |*x| x.* = rnd.float(f32); + const want = try gpa.dupe(f32, logits); + defer gpa.free(want); + const idx = [_]u32{ 0, 2, 3, 6, 7 }; + for (idx, 0..) |row, r| for (0..width) |c| { + const t = lo + c; + const ok = t < words * 32 and (bits[r * words + t / 32] >> @intCast(t % 32)) & 1 == 1; + if (!ok) want[row * width + c] = -std.math.inf(f32); + }; + maskRows(logits, width, &idx, bits, words, lo, width); + try std.testing.expectEqualSlices(f32, want, logits); + } +} diff --git a/zig/src/core/grammar/tags.zig b/zig/src/core/grammar/tags.zig new file mode 100644 index 0000000000..2a69963048 --- /dev/null +++ b/zig/src/core/grammar/tags.zig @@ -0,0 +1,165 @@ +//! Tool-call structural tags: xgrammar 0.2.8's ``get_model_structural_tag`` JSON for a tools spec, byte for byte. + +const std = @import("std"); + +pub const Model = enum { deepseek_v4_1 }; + +pub const Error = error{ InvalidTools, OutOfMemory, WriteFailed }; + +/// Why a tools spec cannot be built (xgrammar's error text where it has one), set with error.InvalidTools. +pub const Problem = struct { + message: []const u8 = "", + buf: [512]u8 = undefined, +}; + +const blanks = 32; + +/// The structural tag JSON for `spec_text` (gpa-owned). +pub fn build(gpa: std.mem.Allocator, model: Model, spec_text: []const u8, problem: *Problem) Error![]u8 { + _ = model; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const spec = std.json.parseFromSliceLeaky(std.json.Value, a, spec_text, .{ .parse_numbers = false }) catch return fail(problem, "the tools spec is not valid JSON"); + if (spec != .object) return fail(problem, "the tools spec is not an object"); + const tools_v = spec.object.get("tools") orelse return fail(problem, "the tools spec has no tools"); + if (tools_v != .array) return fail(problem, "The 'tools' argument must be a list."); + const parallel = if (spec.object.get("parallel_tool_calls")) |p| p == .bool and p.bool else true; + + var fns: std.ArrayList(std.json.ObjectMap) = .empty; + for (tools_v.array.items) |t| { + const f = (if (t == .object) t.object.get("function") else null) orelse continue; + if (f != .object) continue; + try fns.append(a, f.object); + } + const Choice = enum { auto, required, forced }; + var choice: Choice = .auto; + const cv = spec.object.get("tool_choice") orelse std.json.Value.null; + if (cv == .string) { + if (std.mem.eql(u8, cv.string, "required")) choice = .required else if (std.mem.eql(u8, cv.string, "none")) fns.clearRetainingCapacity(); + } else if (cv == .object) { + const name = nameOf(cv.object) orelse return fail(problem, "tool_choice names no function"); + var kept: std.ArrayList(std.json.ObjectMap) = .empty; + for (fns.items) |f| if (std.mem.eql(u8, nameOf(f) orelse "", name)) try kept.append(a, f); + if (kept.items.len == 0) return fail(problem, std.fmt.bufPrint(&problem.buf, "The tool with name '{s}' is not found in the tools list.", .{name}) catch "The tool is not found in the tools list."); + if (kept.items.len != 1) return fail(problem, "Forced tool choice must resolve to exactly one tool."); + fns = kept; + choice = .forced; + } + if (choice == .required and fns.items.len == 0) return fail(problem, "The 'tools' list is empty, which is not allowed when 'tool_choice' is 'required'."); + + var out: std.Io.Writer.Allocating = .init(gpa); + errdefer out.deinit(); + const w = &out.writer; + try w.writeAll("{\"type\":\"structural_tag\",\"format\":"); + switch (choice) { + .auto => if (fns.items.len == 0) { + try w.writeAll("{\"type\":\"any_text\",\"excludes\":[\"\",\"\",\"<|DSML| calls>\"]}"); + } else { + try w.writeAll("{\"type\":\"triggered_tags\",\"triggers\":[\"<|DSML| calls>\"],\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| calls>\\n\",\"content\":"); + try invokes(w, fns.items, parallel); + try w.print(",\"end\":\"\"}}],\"at_least_one\":false,\"stop_after_first\":{},\"excludes\":[\"\",\"\"]}}", .{!parallel}); + }, + .required => { + try w.writeAll("{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},"); + try invokes(w, fns.items, parallel); + try w.writeAll(",{\"type\":\"const_string\",\"value\":\"\"}]}"); + }, + .forced => { + try w.writeAll("{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},"); + try invoke(w, fns.items[0]); + try w.writeAll(",{\"type\":\"const_string\",\"value\":\"\"}]}"); + }, + } + try w.writeByte('}'); + return out.toOwnedSlice(); +} + +fn fail(problem: *Problem, message: []const u8) Error { + problem.message = message; + return error.InvalidTools; +} + +fn nameOf(o: std.json.ObjectMap) ?[]const u8 { + if (o.get("name")) |n| if (n == .string) return n.string; + if (o.get("function")) |f| if (f == .object) if (f.object.get("name")) |n| if (n == .string) return n.string; + return null; +} + +fn invokes(w: *std.Io.Writer, fns: []const std.json.ObjectMap, parallel: bool) !void { + try w.writeAll("{\"type\":\"tags_with_separator\",\"tags\":["); + for (fns, 0..) |f, i| { + if (i > 0) try w.writeByte(','); + try invoke(w, f); + } + try w.print("],\"separator\":\"\",\"at_least_one\":true,\"stop_after_first\":{}}}", .{!parallel}); +} + +/// One tool's invoke tag, its arguments held to its parameters (``strict: false`` or none: any JSON). +fn invoke(w: *std.Io.Writer, f: std.json.ObjectMap) !void { + try w.writeAll("{\"type\":\"tag\",\"begin\":"); + var begin: [1024]u8 = undefined; + const name = nameOf(f) orelse ""; + try string(w, std.fmt.bufPrint(&begin, "<|DSML| invoke name=\"{s}\">\n", .{name}) catch return error.WriteFailed); + try w.writeAll(",\"content\":{\"type\":\"json_schema\",\"json_schema\":"); + const strict_false = if (f.get("strict")) |s| s == .bool and !s.bool else false; + const params = f.get("parameters"); + if (strict_false or params == null or params.? == .null) try w.writeAll("true") else try value(w, params.?); + try w.print(",\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":{d},\"excludes\":[]}},\"end\":\"\\n\"}}", .{blanks}); +} + +/// A JSON value as pydantic writes it: compact, keys in order, numbers as given. +fn value(w: *std.Io.Writer, v: std.json.Value) !void { + switch (v) { + .null => try w.writeAll("null"), + .bool => |b| try w.writeAll(if (b) "true" else "false"), + .integer => |i| try w.print("{d}", .{i}), + .float => |x| try w.print("{d}", .{x}), + .number_string => |s| try w.writeAll(s), + .string => |s| try string(w, s), + .array => |arr| { + try w.writeByte('['); + for (arr.items, 0..) |x, i| { + if (i > 0) try w.writeByte(','); + try value(w, x); + } + try w.writeByte(']'); + }, + .object => |o| { + try w.writeByte('{'); + var it = o.iterator(); + var i: usize = 0; + while (it.next()) |e| : (i += 1) { + if (i > 0) try w.writeByte(','); + try string(w, e.key_ptr.*); + try w.writeByte(':'); + try value(w, e.value_ptr.*); + } + try w.writeByte('}'); + }, + } +} + +/// A JSON string as serde_json writes it: `"` `\` and control characters escaped, everything else raw. +pub fn string(w: *std.Io.Writer, s: []const u8) !void { + try w.writeByte('"'); + var from: usize = 0; + for (s, 0..) |ch, i| { + const esc: ?[]const u8 = switch (ch) { + '"' => "\\\"", + '\\' => "\\\\", + '\n' => "\\n", + '\r' => "\\r", + '\t' => "\\t", + 0x08 => "\\b", + 0x0c => "\\f", + else => null, + }; + if (esc == null and ch >= 0x20) continue; + try w.writeAll(s[from..i]); + if (esc) |e| try w.writeAll(e) else try w.print("\\u{x:0>4}", .{ch}); + from = i + 1; + } + try w.writeAll(s[from..]); + try w.writeByte('"'); +} diff --git a/zig/src/core/grammar/vocab.zig b/zig/src/core/grammar/vocab.zig new file mode 100644 index 0000000000..e7ab7b8c18 --- /dev/null +++ b/zig/src/core/grammar/vocab.zig @@ -0,0 +1,74 @@ +//! The grammar's view of a vocabulary: tokenizer.json's strings by id, added tokens and padding empty unless kept. + +const std = @import("std"); + +pub const Vocab = struct { + arena: std.heap.ArenaAllocator, + /// `size` entries, id order + tokens: [][]const u8, + think_open: ?u32, + think_end: ?u32, + + pub fn deinit(v: *Vocab) void { + v.arena.deinit(); + } +}; + +pub const Error = error{ InvalidTokenizer, VocabTooSmall, OutOfMemory }; + +/// The vocabulary of `json_text` (a tokenizer.json) at `size` columns, added tokens blanked but `keep`. +pub fn build(gpa: std.mem.Allocator, json_text: []const u8, size: u32, keep: []const []const u8) Error!Vocab { + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + const a = arena.allocator(); + var scratch = std.heap.ArenaAllocator.init(gpa); + defer scratch.deinit(); + const doc = std.json.parseFromSliceLeaky(std.json.Value, scratch.allocator(), json_text, .{}) catch return error.InvalidTokenizer; + if (doc != .object) return error.InvalidTokenizer; + const model = doc.object.get("model") orelse return error.InvalidTokenizer; + if (model != .object) return error.InvalidTokenizer; + const vocab = model.object.get("vocab") orelse return error.InvalidTokenizer; + if (vocab != .object) return error.InvalidTokenizer; + const added: []const std.json.Value = if (doc.object.get("added_tokens")) |x| (if (x == .array) x.array.items else &.{}) else &.{}; + + // get_vocab(with_added_tokens=True): the model's map, then each added token's content -> its id + var map: std.StringArrayHashMapUnmanaged(u32) = .empty; + defer map.deinit(gpa); + try map.ensureTotalCapacity(gpa, vocab.object.count() + added.len); + var it = vocab.object.iterator(); + while (it.next()) |e| { + if (e.value_ptr.* != .integer) return error.InvalidTokenizer; + map.putAssumeCapacity(e.key_ptr.*, @intCast(e.value_ptr.integer)); + } + for (added) |t| { + const id = idOf(t) orelse return error.InvalidTokenizer; + map.putAssumeCapacity(contentOf(t) orelse return error.InvalidTokenizer, id); + } + const tokens = try a.alloc([]const u8, size); + @memset(tokens, ""); + for (map.keys(), map.values()) |s, id| { + if (id >= size) return error.VocabTooSmall; + tokens[id] = try a.dupe(u8, s); + } + for (added) |t| { + const s = contentOf(t).?; + const id = idOf(t).?; + const kept = for (keep) |k| { + if (std.mem.eql(u8, k, s)) break true; + } else false; + if (!kept and id < size) tokens[id] = ""; + } + return .{ .arena = arena, .tokens = tokens, .think_open = map.get(""), .think_end = map.get("") }; +} + +fn idOf(t: std.json.Value) ?u32 { + if (t != .object) return null; + const id = t.object.get("id") orelse return null; + return if (id == .integer and id.integer >= 0) @intCast(id.integer) else null; +} + +fn contentOf(t: std.json.Value) ?[]const u8 { + if (t != .object) return null; + const s = t.object.get("content") orelse return null; + return if (s == .string) s.string else null; +} diff --git a/zig/src/core/grammar/xgr.zig b/zig/src/core/grammar/xgr.zig new file mode 100644 index 0000000000..f1c6509ebd --- /dev/null +++ b/zig/src/core/grammar/xgr.zig @@ -0,0 +1,193 @@ +//! xgrammar's C++ core through xgr_c.h, implementing engine.zig's interfaces: masks are xgrammar's own, bit for bit. + +const std = @import("std"); +const engine = @import("engine.zig"); + +/// xgr_c.h's declarations (xgr_c.cc). +const c = struct { + const xgr_tok = opaque {}; + const xgr_compiler = opaque {}; + const xgr_compiled = opaque {}; + const xgr_matcher = opaque {}; + const XGR_JSON_SCHEMA: i32 = 0; + const XGR_REGEX: i32 = 1; + const XGR_EBNF: i32 = 2; + const XGR_STRUCTURAL_TAG: i32 = 3; + extern fn xgr_last_error() [*:0]const u8; + extern fn xgr_detect_metadata(backend: [*]const u8, len: usize, out: [*]u8, cap: usize) i64; + extern fn xgr_tok_new(vocab: [*]const [*]const u8, lens: [*]const usize, n: i32, vocab_type: i32, vocab_size: i32, stops: [*]const i32, n_stops: i32, add_prefix_space: i32) ?*xgr_tok; + extern fn xgr_tok_free(t: *xgr_tok) void; + extern fn xgr_tok_special(t: *const xgr_tok, out: ?[*]i32, cap: i32) i32; + extern fn xgr_bitmask_words(vocab_size: i32) i32; + extern fn xgr_compiler_new(t: *const xgr_tok, max_threads: i32, cache_bytes: i64) ?*xgr_compiler; + extern fn xgr_compiler_free(x: *xgr_compiler) void; + extern fn xgr_compile(x: *xgr_compiler, kind: i32, text: [*]const u8, len: usize, max_whitespace: i32) ?*xgr_compiled; + extern fn xgr_compiled_free(g: *xgr_compiled) void; + extern fn xgr_matcher_new(g: *const xgr_compiled) ?*xgr_matcher; + extern fn xgr_matcher_free(m: *xgr_matcher) void; + extern fn xgr_matcher_accept(m: *xgr_matcher, token: i32) i32; + extern fn xgr_matcher_fill(m: *xgr_matcher, words: [*]i32, n_words: i32) i32; + extern fn xgr_matcher_rollback(m: *xgr_matcher, n: i32) i32; + extern fn xgr_matcher_terminated(m: *const xgr_matcher) i32; +}; + +pub const Error = error{ Xgrammar, OutOfMemory }; + +/// xgrammar's message for the calling thread's last failure. +pub fn lastError() []const u8 { + return std.mem.span(c.xgr_last_error()); +} + +pub fn bitmaskWords(vocab: u32) u32 { + return @intCast(c.xgr_bitmask_words(@intCast(vocab))); +} + +pub const Metadata = struct { vocab_type: i32, add_prefix_space: bool }; + +/// TokenizerInfo._detect_metadata_from_hf over a tokenizer.json text. +pub fn detectMetadata(gpa: std.mem.Allocator, tokenizer_json: []const u8) !Metadata { + var buf: [256]u8 = undefined; + const n = c.xgr_detect_metadata(tokenizer_json.ptr, tokenizer_json.len, &buf, buf.len); + if (n < 0) return error.Xgrammar; + const M = struct { vocab_type: i32, add_prefix_space: bool }; + const p = try std.json.parseFromSlice(M, gpa, buf[0..@intCast(n)], .{}); + defer p.deinit(); + return .{ .vocab_type = p.value.vocab_type, .add_prefix_space = p.value.add_prefix_space }; +} + +pub const Tokenizer = struct { + h: *c.xgr_tok, + + /// TokenizerInfo(encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space). + pub fn init(gpa: std.mem.Allocator, vocab: []const []const u8, meta: Metadata, stops: []const i32) !Tokenizer { + const ptrs = try gpa.alloc([*]const u8, vocab.len); + defer gpa.free(ptrs); + const lens = try gpa.alloc(usize, vocab.len); + defer gpa.free(lens); + for (vocab, ptrs, lens) |v, *p, *l| { + p.* = v.ptr; + l.* = v.len; + } + const h = c.xgr_tok_new(ptrs.ptr, lens.ptr, @intCast(vocab.len), meta.vocab_type, @intCast(vocab.len), stops.ptr, @intCast(stops.len), @intFromBool(meta.add_prefix_space)) orelse return error.Xgrammar; + return .{ .h = h }; + } + + pub fn deinit(t: Tokenizer) void { + c.xgr_tok_free(t.h); + } + + /// TokenizerInfo.special_token_ids (gpa-owned). + pub fn special(t: Tokenizer, gpa: std.mem.Allocator) ![]i32 { + const n = c.xgr_tok_special(t.h, null, 0); + const out = try gpa.alloc(i32, @intCast(n)); + _ = c.xgr_tok_special(t.h, out.ptr, n); + return out; + } +}; + +/// One tokenizer view's GrammarCompiler (max_threads 8, its cache capped at `cache_bytes`), as `engine.Compiler`. +pub const Compiler = struct { + h: *c.xgr_compiler, + vocab: u32, + /// compiles come from HTTP threads (rare, not on the round loop): a spin lock is enough + mutex: std.atomic.Mutex = .unlocked, + + pub fn init(t: Tokenizer, vocab: u32, cache_bytes: i64) !Compiler { + return .{ .h = c.xgr_compiler_new(t.h, 8, cache_bytes) orelse return error.Xgrammar, .vocab = vocab }; + } + + pub fn deinit(x: *Compiler) void { + c.xgr_compiler_free(x.h); + } + + pub fn compiler(x: *Compiler) engine.Compiler { + return .{ .ptr = x, .vtable = &.{ .compile = compileFn } }; + } + + /// `kind`'s grammar from `text`; on failure null with xgrammar's message copied into `why` (gpa-owned). + pub fn compile(x: *Compiler, kind: engine.Source, text: []const u8, max_whitespace: ?u32) Error!*Compiled { + while (!x.mutex.tryLock()) std.atomic.spinLoopHint(); + defer x.mutex.unlock(); + const k: i32 = switch (kind) { + .json_schema => c.XGR_JSON_SCHEMA, + .regex => c.XGR_REGEX, + .ebnf => c.XGR_EBNF, + .structural_tag => c.XGR_STRUCTURAL_TAG, + }; + const ws: i32 = if (max_whitespace) |w| @intCast(w) else -1; + const g = c.xgr_compile(x.h, k, text.ptr, text.len, ws) orelse return error.Xgrammar; + const out = std.heap.c_allocator.create(Compiled) catch { + c.xgr_compiled_free(g); + return error.OutOfMemory; + }; + out.* = .{ .h = g, .vocab = x.vocab }; + return out; + } + + fn compileFn(ptr: *anyopaque, kind: engine.Source, text: []const u8, max_whitespace: ?u32) Error!engine.Compiled { + const x: *Compiler = @ptrCast(@alignCast(ptr)); + return (try x.compile(kind, text, max_whitespace)).compiled(); + } +}; + +pub const Compiled = struct { + h: *c.xgr_compiled, + vocab: u32, + + pub fn compiled(g: *Compiled) engine.Compiled { + return .{ .ptr = g, .vtable = &.{ .matcher = matcherFn, .deinit = deinitFn } }; + } + + fn matcherFn(ptr: *anyopaque) Error!engine.Matcher { + const g: *Compiled = @ptrCast(@alignCast(ptr)); + const m = std.heap.c_allocator.create(Matcher) catch return error.OutOfMemory; + m.* = .{ .h = c.xgr_matcher_new(g.h) orelse { + std.heap.c_allocator.destroy(m); + return error.Xgrammar; + }, .words = bitmaskWords(g.vocab) }; + return m.matcher(); + } + + fn deinitFn(ptr: *anyopaque) void { + const g: *Compiled = @ptrCast(@alignCast(ptr)); + c.xgr_compiled_free(g.h); + std.heap.c_allocator.destroy(g); + } +}; + +/// GrammarMatcher(compiled) with xgrammar's defaults (the compiled grammar's stop tokens, unlimited rollback). +pub const Matcher = struct { + h: *c.xgr_matcher, + words: u32, + + pub fn matcher(m: *Matcher) engine.Matcher { + return .{ .ptr = m, .vtable = &vtable }; + } + + const vtable: engine.Matcher.VTable = .{ .accept = accept, .fill = fill, .rollback = rollback, .terminated = terminated, .deinit = deinit }; + + fn self(p: *anyopaque) *Matcher { + return @ptrCast(@alignCast(p)); + } + fn accept(p: *anyopaque, token: u32) Error!bool { + const r = c.xgr_matcher_accept(self(p).h, @intCast(token)); + if (r < 0) return error.Xgrammar; + return r == 1; + } + fn fill(p: *anyopaque, words: []u32) Error!void { + const m = self(p); + if (words.len != m.words) return error.Xgrammar; + if (c.xgr_matcher_fill(m.h, @ptrCast(words.ptr), @intCast(words.len)) < 0) return error.Xgrammar; + } + fn rollback(p: *anyopaque, n: u32) Error!void { + if (c.xgr_matcher_rollback(self(p).h, @intCast(n)) < 0) return error.Xgrammar; + } + fn terminated(p: *anyopaque) bool { + return c.xgr_matcher_terminated(self(p).h) == 1; + } + fn deinit(p: *anyopaque) void { + const m = self(p); + c.xgr_matcher_free(m.h); + std.heap.c_allocator.destroy(m); + } +}; diff --git a/zig/src/core/grammar/xgr_c.cc b/zig/src/core/grammar/xgr_c.cc new file mode 100644 index 0000000000..8d004b135c --- /dev/null +++ b/zig/src/core/grammar/xgr_c.cc @@ -0,0 +1,157 @@ +// The C ABI of xgr_c.h over xgrammar's C++ API, with the defaults of xgrammar 0.2.8's Python bindings. +#include "xgr_c.h" + +#include + +#include +#include +#include +#include + +namespace { +thread_local std::string last_error; + +template +auto guard(F&& f, decltype(f()) failed) -> decltype(f()) { + try { + last_error.clear(); + return f(); + } catch (const std::exception& e) { + last_error = e.what(); + } catch (...) { + last_error = "unknown xgrammar error"; + } + return failed; +} +} // namespace + +struct xgr_tok { + xgrammar::TokenizerInfo info; +}; +struct xgr_compiler { + xgrammar::GrammarCompiler c; +}; +struct xgr_compiled { + xgrammar::CompiledGrammar g; +}; +struct xgr_matcher { + xgrammar::GrammarMatcher m; +}; + +extern "C" { + +const char* xgr_last_error(void) { return last_error.c_str(); } + +int64_t xgr_detect_metadata(const char* backend, size_t len, char* out, size_t cap) { + return guard( + [&]() -> int64_t { + std::string s = xgrammar::TokenizerInfo::DetectMetadataFromHF(std::string(backend, len)); + if (s.size() + 1 > cap) throw std::runtime_error("metadata buffer too small"); + std::memcpy(out, s.data(), s.size()); + out[s.size()] = 0; + return static_cast(s.size()); + }, + -1 + ); +} + +xgr_tok* xgr_tok_new(const char* const* vocab, const size_t* lens, int32_t n, int32_t vocab_type, int32_t vocab_size, + const int32_t* stops, int32_t n_stops, int32_t add_prefix_space) { + return guard( + [&]() -> xgr_tok* { + std::vector enc; + enc.reserve(n); + for (int32_t i = 0; i < n; ++i) enc.emplace_back(vocab[i], lens[i]); + std::vector st(stops, stops + n_stops); + return new xgr_tok{xgrammar::TokenizerInfo(enc, static_cast(vocab_type), vocab_size, st, + add_prefix_space != 0)}; + }, + nullptr + ); +} + +void xgr_tok_free(xgr_tok* t) { delete t; } + +int32_t xgr_tok_special(const xgr_tok* t, int32_t* out, int32_t cap) { + const auto& ids = t->info.GetSpecialTokenIds(); + int32_t n = static_cast(ids.size()); + for (int32_t i = 0; i < n && i < cap; ++i) out[i] = ids[i]; + return n; +} + +int32_t xgr_tok_vocab_type(const xgr_tok* t) { return static_cast(t->info.GetVocabType()); } +int32_t xgr_tok_add_prefix_space(const xgr_tok* t) { return t->info.GetAddPrefixSpace() ? 1 : 0; } +int32_t xgr_bitmask_words(int32_t vocab_size) { return xgrammar::GetBitmaskSize(vocab_size); } + +xgr_compiler* xgr_compiler_new(const xgr_tok* t, int32_t max_threads, int64_t cache_bytes) { + return guard([&]() -> xgr_compiler* { return new xgr_compiler{xgrammar::GrammarCompiler(t->info, max_threads, true, cache_bytes)}; }, + nullptr); +} + +void xgr_compiler_free(xgr_compiler* c) { delete c; } + +xgr_compiled* xgr_compile(xgr_compiler* c, int32_t kind, const char* text, size_t len, int32_t max_whitespace) { + return guard( + [&]() -> xgr_compiled* { + std::string s(text, len); + switch (kind) { + case XGR_JSON_SCHEMA: { + std::optional ws = max_whitespace < 0 ? std::nullopt : std::optional(max_whitespace); + return new xgr_compiled{c->c.CompileJSONSchema(s, true, std::nullopt, std::nullopt, true, ws, false)}; + } + case XGR_REGEX: + return new xgr_compiled{c->c.CompileRegex(s)}; + case XGR_EBNF: + return new xgr_compiled{c->c.CompileGrammar(s, "root")}; + case XGR_STRUCTURAL_TAG: + return new xgr_compiled{c->c.CompileStructuralTag(s)}; + } + throw std::runtime_error("unknown grammar kind"); + }, + nullptr + ); +} + +void xgr_compiled_free(xgr_compiled* g) { delete g; } + +xgr_matcher* xgr_matcher_new(const xgr_compiled* g) { + return guard([&]() -> xgr_matcher* { return new xgr_matcher{xgrammar::GrammarMatcher(g->g)}; }, nullptr); +} + +void xgr_matcher_free(xgr_matcher* m) { delete m; } + +int32_t xgr_matcher_accept(xgr_matcher* m, int32_t token) { + return guard([&]() -> int32_t { return m->m.AcceptToken(token) ? 1 : 0; }, -1); +} + +int32_t xgr_matcher_fill(xgr_matcher* m, int32_t* words, int32_t n_words) { + return guard( + [&]() -> int32_t { + int64_t shape[2] = {1, n_words}; + DLTensor t; + t.data = words; + t.device = DLDevice{kDLCPU, 0}; + t.ndim = 2; + t.dtype = xgrammar::GetBitmaskDLType(); + t.shape = shape; + t.strides = nullptr; + t.byte_offset = 0; + return m->m.FillNextTokenBitmask(&t, 0) ? 1 : 0; + }, + -1 + ); +} + +int32_t xgr_matcher_rollback(xgr_matcher* m, int32_t n) { + return guard( + [&]() -> int32_t { + m->m.Rollback(n); + return 0; + }, + -1 + ); +} + +int32_t xgr_matcher_terminated(const xgr_matcher* m) { return m->m.IsTerminated() ? 1 : 0; } + +} // extern "C" diff --git a/zig/src/core/grammar/xgr_c.h b/zig/src/core/grammar/xgr_c.h new file mode 100644 index 0000000000..8391d0480d --- /dev/null +++ b/zig/src/core/grammar/xgr_c.h @@ -0,0 +1,58 @@ +// A C ABI over xgrammar 0.2.8's C++ core; failing calls return NULL or -1 with the message in xgr_last_error(). +#ifndef TF_XGR_C_H_ +#define TF_XGR_C_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct xgr_tok xgr_tok; +typedef struct xgr_compiler xgr_compiler; +typedef struct xgr_compiled xgr_compiled; +typedef struct xgr_matcher xgr_matcher; + +enum xgr_kind { + XGR_JSON_SCHEMA = 0, // any_whitespace, strict_mode, max_whitespace_cnt (< 0: none) + XGR_REGEX = 1, + XGR_EBNF = 2, // root rule "root" + XGR_STRUCTURAL_TAG = 3 // the structural tag's JSON +}; + +// The calling thread's last error message (empty when none); valid until the thread's next failing call. +const char *xgr_last_error(void); + +// {"vocab_type": int, "add_prefix_space": bool} for a tokenizer.json into out (NUL-terminated); its length or -1. +int64_t xgr_detect_metadata(const char *backend, size_t len, char *out, size_t cap); + +// TokenizerInfo(encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space). +xgr_tok *xgr_tok_new(const char *const *vocab, const size_t *lens, int32_t n, int32_t vocab_type, int32_t vocab_size, + const int32_t *stops, int32_t n_stops, int32_t add_prefix_space); +void xgr_tok_free(xgr_tok *t); +// special_token_ids into out (at most cap); the count. +int32_t xgr_tok_special(const xgr_tok *t, int32_t *out, int32_t cap); +int32_t xgr_tok_vocab_type(const xgr_tok *t); +int32_t xgr_tok_add_prefix_space(const xgr_tok *t); +// The bitmask's int32 words for a vocabulary of vocab_size tokens. +int32_t xgr_bitmask_words(int32_t vocab_size); + +xgr_compiler *xgr_compiler_new(const xgr_tok *t, int32_t max_threads, int64_t cache_bytes); +void xgr_compiler_free(xgr_compiler *c); +xgr_compiled *xgr_compile(xgr_compiler *c, int32_t kind, const char *text, size_t len, int32_t max_whitespace); +void xgr_compiled_free(xgr_compiled *g); + +xgr_matcher *xgr_matcher_new(const xgr_compiled *g); +void xgr_matcher_free(xgr_matcher *m); +// 1 accepted, 0 rejected, -1 failed. +int32_t xgr_matcher_accept(xgr_matcher *m, int32_t token); +// The next token's allowed bits (token t: bit t % 32 of word t / 32); 1 some token masked, 0 none, -1 failed. +int32_t xgr_matcher_fill(xgr_matcher *m, int32_t *words, int32_t n_words); +int32_t xgr_matcher_rollback(xgr_matcher *m, int32_t n); +int32_t xgr_matcher_terminated(const xgr_matcher *m); + +#ifdef __cplusplus +} +#endif +#endif diff --git a/zig/tests/grammar/fixtures/synthetic.golden.json b/zig/tests/grammar/fixtures/synthetic.golden.json new file mode 100644 index 0000000000..11664f0740 --- /dev/null +++ b/zig/tests/grammar/fixtures/synthetic.golden.json @@ -0,0 +1 @@ +{"xgrammar":"0.2.8","reference":"python-serving","vocab_size":288,"eos":282,"think_open":283,"think_end":284,"words":9,"views":{"text":{"vocab_type":2,"add_prefix_space":false,"vocab":"230dd81d9035c1ae","never_digest":"efd94cf53593dd43","never_count":6},"tools":{"vocab_type":2,"add_prefix_space":false,"vocab":"da271f92bf814402","never_digest":"980723c3fb61bf6c","never_count":5}},"cases":[{"name":"json_object","kind":"json","text":"","tag":null,"active":true,"steps":[{"window":[165,90],"keep":2,"first":0,"stop":false,"rows":[["1d5ede0d5993de19",2],["40d231f08f2ac3e5",11]],"advance":[90,275]},{"window":[275],"keep":1,"first":0,"stop":false,"rows":[["40d231f08f2ac3e5",11]],"advance":[1]},{"window":[1],"keep":1,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169]],"advance":[139]},{"window":[139],"keep":1,"first":0,"stop":false,"rows":[["7721af923e2625d6",64]],"advance":[105]},{"window":[105,31,258,133,113],"keep":5,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["7721af923e2625d6",64],["b7c8fcfe406ef69d",169]],"advance":[174]},{"window":[174,109],"keep":2,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["7721af923e2625d6",64]],"advance":[109,117]},{"window":[117],"keep":1,"first":0,"stop":false,"rows":[["7721af923e2625d6",64]],"advance":[239]},{"window":[239],"keep":1,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169]],"advance":[176]},{"window":[176],"keep":1,"first":0,"stop":false,"rows":[["a0c6e552712bf01e",16]],"advance":[237]},{"window":[237,251,230,62,153],"keep":5,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["7721af923e2625d6",64],["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["7721af923e2625d6",64]],"advance":[251,230,62,153,255]},{"window":[255,18,214],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169]],"advance":[87]},{"window":[87,19],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169]],"advance":[19,87]},{"window":[87,148,47],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["7721af923e2625d6",64]],"advance":[148,252]},{"window":[252,79,146,236,236],"keep":4,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["7721af923e2625d6",64],["b7c8fcfe406ef69d",169]],"advance":[79,146,115]},{"window":[115,30],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169]],"advance":[73]},{"window":[73,159],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["7721af923e2625d6",64]],"advance":[42]},{"window":[42,35,71,129],"keep":4,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169],["7721af923e2625d6",64]],"advance":[35,38]},{"window":[38],"keep":1,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169]],"advance":[38]},{"window":[38,259],"keep":2,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169],["b7c8fcfe406ef69d",169]],"advance":[67]},{"window":[67,214],"keep":1,"first":0,"stop":false,"rows":[["b7c8fcfe406ef69d",169]],"advance":[265]},{"window":[265,275,268,275,275,201],"keep":6,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[201]},{"window":[201,197,220,0],"keep":3,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[197,220,197]},{"window":[197,186],"keep":1,"first":0,"stop":false,"rows":[["01b401fc4219b614",7]],"advance":[25]},{"window":[25,90],"keep":2,"first":0,"stop":false,"rows":[["e12e8bd769f11711",32],["40d231f08f2ac3e5",11]],"advance":[90,275]}],"finished":false},{"name":"json_schema","kind":"json_schema","text":"{\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}","tag":null,"active":true,"steps":[{"window":[186,265,66,175],"keep":3,"first":0,"stop":false,"rows":[["1d5ede0d5993de19",2],["940300c8f1b2b534",2],["0195320b34a1e3b8",1]],"advance":[265,66,72]},{"window":[72,83,88,1],"keep":4,"first":0,"stop":false,"rows":[["d1d61829e98de7e7",1],["449ca9e633bd281f",1],["3405cabc04f6f848",2],["01b401fc4219b614",7]],"advance":[83,88,1]},{"window":[1,197,220,197,275,275,197,25,268],"keep":9,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["6854501f0132b612",10],["6854501f0132b612",10]],"advance":[197,220,197,275,197]},{"window":[197,275,198,52],"keep":3,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[201]},{"window":[201,220],"keep":2,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[201]},{"window":[201,275,275,197,25,266,21,32,238],"keep":8,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["6854501f0132b612",10],["8cf2d085632efb8e",197],["8cf2d085632efb8e",197],["8cf2d085632efb8e",197]],"advance":[275,275,220]},{"window":[220,198,25,52],"keep":3,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["6854501f0132b612",10]],"advance":[198,25,198]},{"window":[198,201,108],"keep":2,"first":0,"stop":false,"rows":[["6854501f0132b612",10],["6854501f0132b612",10]],"advance":[201,275]},{"window":[275,275],"keep":2,"first":0,"stop":false,"rows":[["6854501f0132b612",10],["6854501f0132b612",10]],"advance":[275,198]},{"window":[198],"keep":1,"first":0,"stop":false,"rows":[["6854501f0132b612",10]],"advance":[268]},{"window":[268,266,33,4,200,76],"keep":6,"first":0,"stop":false,"rows":[["6854501f0132b612",10],["8cf2d085632efb8e",197],["8cf2d085632efb8e",197],["8cf2d085632efb8e",197],["3f509174ac09f493",196],["a9c31053618c6136",195]],"advance":[266,36]},{"window":[36,15,175,101],"keep":4,"first":0,"stop":false,"rows":[["8cf2d085632efb8e",197],["8cf2d085632efb8e",197],["7721af923e2625d6",64],["7721af923e2625d6",64]],"advance":[15,48]},{"window":[48,50],"keep":2,"first":0,"stop":false,"rows":[["3f509174ac09f493",196],["a9c31053618c6136",195]],"advance":[50,176]},{"window":[176],"keep":1,"first":0,"stop":false,"rows":[["a0c6e552712bf01e",16]],"advance":[236]},{"window":[236,119,253],"keep":3,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["7721af923e2625d6",64],["a9c31053618c6136",195]],"advance":[119,253,64]},{"window":[64,67,51,31,217,217,1,267,201],"keep":9,"first":0,"stop":false,"rows":[["759979dd3deb5878",193],["41bcc0fdcf789137",191],["94ae616c304096a4",186],["f61d3673d08954f5",184],["0e122487dc6467b3",178],["3bab55816f9c1c63",1],["26b878181b44e3ed",8],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[67,51,31,24]},{"window":[24,216,1,268],"keep":4,"first":0,"stop":false,"rows":[["0e122487dc6467b3",178],["3bab55816f9c1c63",1],["26b878181b44e3ed",8],["26b878181b44e3ed",8]],"advance":[216,1]},{"window":[1,187],"keep":1,"first":0,"stop":false,"rows":[["26b878181b44e3ed",8]],"advance":[267]},{"window":[267,201,220,275,1,84],"keep":6,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["796ed0ebcbbf2a29",1],["b2ff1bbe17acf6b2",1]],"advance":[201,220,275]},{"window":[275,268],"keep":2,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[268,201]},{"window":[201,201],"keep":2,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[201,1]},{"window":[1,93],"keep":1,"first":0,"stop":false,"rows":[["796ed0ebcbbf2a29",1]],"advance":[84]},{"window":[84,77,113],"keep":2,"first":0,"stop":false,"rows":[["b2ff1bbe17acf6b2",1],["0195320b34a1e3b8",1]],"advance":[77]},{"window":[77],"keep":1,"first":0,"stop":false,"rows":[["0195320b34a1e3b8",1]],"advance":[72]}],"finished":false},{"name":"array","kind":"json_schema","text":"{\"type\": \"array\", \"items\": {\"type\": \"integer\"}, \"minItems\": 1, \"maxItems\": 3}","tag":null,"active":true,"steps":[{"window":[173,58,122],"keep":2,"first":0,"stop":false,"rows":[["2b3be16762ba092d",1],["58af9c4fafe828e4",18]],"advance":[58]},{"window":[58,103],"keep":1,"first":0,"stop":false,"rows":[["58af9c4fafe828e4",18]],"advance":[23]},{"window":[23,267,12,16,20,201],"keep":6,"first":0,"stop":false,"rows":[["69f86cbb1989355e",20],["58af9c4fafe828e4",18],["f814d9278b80255a",10],["69f86cbb1989355e",20],["69f86cbb1989355e",20],["31c274f071c46a73",9]],"advance":[198]},{"window":[198,267,20,18,15,22,197,201,11],"keep":9,"first":0,"stop":false,"rows":[["31c274f071c46a73",9],["58af9c4fafe828e4",18],["69f86cbb1989355e",20],["69f86cbb1989355e",20],["69f86cbb1989355e",20],["69f86cbb1989355e",20],["31c274f071c46a73",9],["31c274f071c46a73",9],["58af9c4fafe828e4",18]],"advance":[267,20,18,15,22,197,201,197]},{"window":[197,220],"keep":2,"first":0,"stop":false,"rows":[["31c274f071c46a73",9],["31c274f071c46a73",9]],"advance":[220,275]},{"window":[275,201,220],"keep":3,"first":0,"stop":false,"rows":[["31c274f071c46a73",9],["31c274f071c46a73",9],["31c274f071c46a73",9]],"advance":[201,60]},{"window":[60,282,45],"keep":1,"first":0,"stop":true,"rows":[["d584ce97eece608f",1]],"advance":[282]}],"finished":true},{"name":"regex","kind":"regex","text":"\\d{3}-[a-z]{2,5}( ok)?","tag":null,"active":true,"steps":[{"window":[201],"keep":1,"first":0,"stop":false,"rows":[["ec9ed08f9e8bac06",11]],"advance":[19]},{"window":[19,23,24,12,73,262],"keep":6,"first":0,"stop":false,"rows":[["ec9ed08f9e8bac06",11],["e8ba6adca6ac8d56",10],["26f871a1d2e6b370",1],["83acbedcaa775442",34],["a318cb4064d21351",33],["2cd3160f7888fa43",28]],"advance":[23,24,12,73,262,86]},{"window":[86,220,78,21],"keep":3,"first":0,"stop":false,"rows":[["9421e19538973902",2],["fde77043ec5b1050",1],["478d904b5aab5b8e",1]],"advance":[220]},{"window":[220,78,74,282],"keep":3,"first":0,"stop":true,"rows":[["fde77043ec5b1050",1],["478d904b5aab5b8e",1],["d584ce97eece608f",1]],"advance":[78,74,282]}],"finished":true},{"name":"choice","kind":"choice","text":"[\"yes\", \"no\", \"maybe \\\"so\\\"\", \"été\"]","tag":null,"active":true,"steps":[{"window":[129],"keep":1,"first":0,"stop":false,"rows":[["6f5ff761f5291e4a",6]],"advance":[77]},{"window":[77,78,282,282,282,282,75],"keep":2,"first":0,"stop":true,"rows":[["fde77043ec5b1050",1],["d584ce97eece608f",1]],"advance":[78,282]}],"finished":true},{"name":"ebnf","kind":"grammar","text":"root ::= expr\nexpr ::= term (\"+\" term)*\nterm ::= [0-9]+ | \"(\" expr \")\"","tag":null,"active":true,"steps":[{"window":[68],"keep":1,"first":0,"stop":false,"rows":[["bb75d1c061c4a28c",12]],"advance":[7]},{"window":[7,269,24,20],"keep":4,"first":0,"stop":false,"rows":[["bb75d1c061c4a28c",12],["045b3e39731e59c5",13],["045b3e39731e59c5",13],["045b3e39731e59c5",13]],"advance":[22]},{"window":[22,20],"keep":2,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13],["045b3e39731e59c5",13]],"advance":[17]},{"window":[17,259],"keep":1,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13]],"advance":[17]},{"window":[17,16,46],"keep":2,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13],["045b3e39731e59c5",13]],"advance":[24]},{"window":[24,22,10,21,20,24],"keep":6,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13],["045b3e39731e59c5",13],["bb75d1c061c4a28c",12],["045b3e39731e59c5",13],["045b3e39731e59c5",13],["045b3e39731e59c5",13]],"advance":[22,19]},{"window":[19,18,173],"keep":2,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13],["045b3e39731e59c5",13]],"advance":[18,10]},{"window":[10,18],"keep":2,"first":0,"stop":false,"rows":[["bb75d1c061c4a28c",12],["045b3e39731e59c5",13]],"advance":[18,18]},{"window":[18],"keep":1,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13]],"advance":[19]},{"window":[19,219],"keep":1,"first":0,"stop":false,"rows":[["045b3e39731e59c5",13]],"advance":[8]},{"window":[8,146],"keep":1,"first":0,"stop":false,"rows":[["c694e0b5a2bfb880",2]],"advance":[282]}],"finished":true},{"name":"tools_required","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}, {\"type\": \"function\", \"function\": {\"name\": \"loose\", \"parameters\": {\"type\": \"object\"}, \"strict\": false}}], \"tool_choice\": \"required\", \"parallel_tool_calls\": true}","tag":"{\"type\":\"structural_tag\",\"format\":{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},{\"type\":\"tags_with_separator\",\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"get_weather\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\",\"maxLength\":12},\"unit\":{\"enum\":[\"c\",\"f\"]},\"days\":{\"type\":\"integer\",\"minimum\":1,\"maximum\":7}},\"required\":[\"city\",\"unit\"],\"additionalProperties\":false},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"note\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"text\":{\"type\":\"string\"},\"tags\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"maxItems\":2}},\"required\":[\"text\"]},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"loose\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":true,\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"}],\"separator\":\"\",\"at_least_one\":true,\"stop_after_first\":false},{\"type\":\"const_string\",\"value\":\"\"}]}}","active":true,"steps":[{"window":[223,275,27,171,121,250],"keep":6,"first":0,"stop":false,"rows":[["e969cc599117b642",2],["18dc8579859e6bb3",1],["01134894a381097e",2],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1]],"advance":[275,27,171,121]},{"window":[121,250,35,50,44,43],"keep":6,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1],["78edc49140ab8efb",1],["0c25f068e5a271aa",1],["e1dfcc033d5b0d2a",1]],"advance":[250,35,50,44,43]},{"window":[43,171,121,175],"keep":3,"first":0,"stop":false,"rows":[["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1]],"advance":[171,121,250]},{"window":[250,220,66],"keep":3,"first":0,"stop":false,"rows":[["1ac3cfee443070d4",2],["f03868427cf71418",1],["7d507dccdd7bf2e2",1]],"advance":[220,66]},{"window":[66,64,75,75,127],"keep":4,"first":0,"stop":false,"rows":[["7d507dccdd7bf2e2",1],["15fc28c5cf74bc3c",1],["15fc28c5cf74bc3c",1],["5e1fa56f07739472",1]],"advance":[64,75,75,82]},{"window":[82,29,198,27,285,220,72,77,20],"keep":8,"first":0,"stop":false,"rows":[["0f3ff8d4a6ba23d2",1],["f44d7b1682b3347f",1],["18dc8579859e6bb3",1],["01134894a381097e",2],["16ac141c6832090d",1],["39759c9b02c87122",2],["b2ff1bbe17acf6b2",1],["901c50779bff8a9b",1]],"advance":[29]},{"window":[29,198],"keep":2,"first":0,"stop":false,"rows":[["f44d7b1682b3347f",1],["18dc8579859e6bb3",1]],"advance":[198]},{"window":[198,27,153],"keep":2,"first":0,"stop":false,"rows":[["18dc8579859e6bb3",1],["01134894a381097e",2]],"advance":[27]},{"window":[27,171,121],"keep":3,"first":0,"stop":false,"rows":[["01134894a381097e",2],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1]],"advance":[171,121]},{"window":[121,250,35,99],"keep":3,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1]],"advance":[250,35]},{"window":[35],"keep":1,"first":0,"stop":false,"rows":[["cd5ed2a33e593c70",1]],"advance":[50]},{"window":[50,44,173],"keep":2,"first":0,"stop":false,"rows":[["78edc49140ab8efb",1],["0c25f068e5a271aa",1]],"advance":[44,43]},{"window":[43,171,121,250,220,253],"keep":5,"first":0,"stop":false,"rows":[["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["16ac141c6832090d",1],["39759c9b02c87122",2]],"advance":[171,121,250,220,72]},{"window":[72,77],"keep":2,"first":0,"stop":false,"rows":[["b2ff1bbe17acf6b2",1],["901c50779bff8a9b",1]],"advance":[77]},{"window":[77],"keep":1,"first":0,"stop":false,"rows":[["901c50779bff8a9b",1]],"advance":[85]},{"window":[85,78,74,68,220,193],"keep":5,"first":0,"stop":false,"rows":[["fde77043ec5b1050",1],["478d904b5aab5b8e",1],["71b0cde5e953ddaf",1],["16ac141c6832090d",1],["b84ef774785a90c9",2]],"advance":[78,74,68]},{"window":[68,220,215],"keep":2,"first":0,"stop":false,"rows":[["16ac141c6832090d",1],["b84ef774785a90c9",2]],"advance":[220]},{"window":[220,259,242],"keep":2,"first":0,"stop":false,"rows":[["b84ef774785a90c9",2],["2ac4811b6f080943",2]],"advance":[259]},{"window":[259,276,263,83,17],"keep":4,"first":0,"stop":false,"rows":[["2ac4811b6f080943",2],["0acdc0b465020e3c",5],["d1d61829e98de7e7",1],["71b0cde5e953ddaf",1]],"advance":[276,263,83,68]},{"window":[68],"keep":1,"first":0,"stop":false,"rows":[["cd0bdf7e5dadc655",2]],"advance":[1]},{"window":[1,173],"keep":1,"first":0,"stop":false,"rows":[["0f3ff8d4a6ba23d2",1]],"advance":[29]},{"window":[29,275,198],"keep":3,"first":0,"stop":false,"rows":[["e969cc599117b642",2],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[275,27]},{"window":[27,171,119],"keep":2,"first":0,"stop":false,"rows":[["01134894a381097e",2],["ac6acd14b749f97f",1]],"advance":[285]},{"window":[285,220,274,220,77,64,76,68,28],"keep":9,"first":0,"stop":false,"rows":[["16ac141c6832090d",1],["9cac882ea6761c02",2],["16ac141c6832090d",1],["b84ef774785a90c9",2],["7d507dccdd7bf2e2",1],["f86c655263c87439",1],["71b0cde5e953ddaf",1],["2ac4811b6f080943",2],["3bab55816f9c1c63",1]],"advance":[220,274,220,77,64,76,68,276]}],"finished":false},{"name":"tools_required_single","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}, {\"type\": \"function\", \"function\": {\"name\": \"loose\", \"parameters\": {\"type\": \"object\"}, \"strict\": false}}], \"tool_choice\": \"required\", \"parallel_tool_calls\": false}","tag":"{\"type\":\"structural_tag\",\"format\":{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},{\"type\":\"tags_with_separator\",\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"get_weather\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\",\"maxLength\":12},\"unit\":{\"enum\":[\"c\",\"f\"]},\"days\":{\"type\":\"integer\",\"minimum\":1,\"maximum\":7}},\"required\":[\"city\",\"unit\"],\"additionalProperties\":false},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"note\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"text\":{\"type\":\"string\"},\"tags\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"maxItems\":2}},\"required\":[\"text\"]},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"loose\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":true,\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"}],\"separator\":\"\",\"at_least_one\":true,\"stop_after_first\":true},{\"type\":\"const_string\",\"value\":\"\"}]}}","active":true,"steps":[{"window":[82,275,27,285],"keep":4,"first":0,"stop":false,"rows":[["e969cc599117b642",2],["18dc8579859e6bb3",1],["01134894a381097e",2],["1ac3cfee443070d4",2]],"advance":[275,27]},{"window":[27,171],"keep":2,"first":0,"stop":false,"rows":[["01134894a381097e",2],["ac6acd14b749f97f",1]],"advance":[285]},{"window":[285,273,29],"keep":3,"first":0,"stop":false,"rows":[["1ac3cfee443070d4",2],["0f3ff8d4a6ba23d2",1],["f44d7b1682b3347f",1]],"advance":[220]},{"window":[220,66],"keep":2,"first":0,"stop":false,"rows":[["f03868427cf71418",1],["7d507dccdd7bf2e2",1]],"advance":[66,64]},{"window":[64,75,75,82],"keep":4,"first":0,"stop":false,"rows":[["15fc28c5cf74bc3c",1],["15fc28c5cf74bc3c",1],["5e1fa56f07739472",1],["0f3ff8d4a6ba23d2",1]],"advance":[75,75]},{"window":[75,150],"keep":1,"first":0,"stop":false,"rows":[["5e1fa56f07739472",1]],"advance":[82]},{"window":[82],"keep":1,"first":0,"stop":false,"rows":[["0f3ff8d4a6ba23d2",1]],"advance":[29]},{"window":[29],"keep":1,"first":0,"stop":false,"rows":[["f44d7b1682b3347f",1]],"advance":[198]},{"window":[198,27,179],"keep":2,"first":0,"stop":false,"rows":[["18dc8579859e6bb3",1],["01134894a381097e",2]],"advance":[27]},{"window":[27,285,104],"keep":2,"first":0,"stop":false,"rows":[["01134894a381097e",2],["16ac141c6832090d",1]],"advance":[285]},{"window":[285,220,16],"keep":2,"first":0,"stop":false,"rows":[["16ac141c6832090d",1],["39759c9b02c87122",2]],"advance":[220,272]},{"window":[272],"keep":1,"first":0,"stop":false,"rows":[["16ac141c6832090d",1]],"advance":[220]},{"window":[220,77,64,76],"keep":4,"first":0,"stop":false,"rows":[["b84ef774785a90c9",2],["7d507dccdd7bf2e2",1],["f86c655263c87439",1],["71b0cde5e953ddaf",1]],"advance":[77,64,76,68]},{"window":[68,276,77],"keep":3,"first":0,"stop":false,"rows":[["2ac4811b6f080943",2],["0acdc0b465020e3c",5],["fde77043ec5b1050",1]],"advance":[276,77,78]},{"window":[78,83,68,277],"keep":4,"first":0,"stop":false,"rows":[["d1d61829e98de7e7",1],["71b0cde5e953ddaf",1],["cd0bdf7e5dadc655",2],["e969cc599117b642",2]],"advance":[83,68,277,198]},{"window":[198,282],"keep":1,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7]],"advance":[268]},{"window":[268,189],"keep":1,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7]],"advance":[197]},{"window":[197,27,251],"keep":2,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["01134894a381097e",2]],"advance":[198]},{"window":[198,201],"keep":2,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[201,27]},{"window":[27,171,121,250,35,50],"keep":6,"first":0,"stop":false,"rows":[["01134894a381097e",2],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1],["78edc49140ab8efb",1]],"advance":[171]},{"window":[171,121],"keep":2,"first":0,"stop":false,"rows":[["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1]],"advance":[121]},{"window":[121,250,35,50,44,43,171,121,250],"keep":9,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1],["78edc49140ab8efb",1],["0c25f068e5a271aa",1],["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["16ac141c6832090d",1]],"advance":[250,35]},{"window":[35,50],"keep":2,"first":0,"stop":false,"rows":[["cd5ed2a33e593c70",1],["78edc49140ab8efb",1]],"advance":[50]},{"window":[50,44,43,171,33],"keep":4,"first":0,"stop":false,"rows":[["78edc49140ab8efb",1],["0c25f068e5a271aa",1],["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1]],"advance":[44]}],"finished":false},{"name":"tools_named","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}, {\"type\": \"function\", \"function\": {\"name\": \"loose\", \"parameters\": {\"type\": \"object\"}, \"strict\": false}}], \"tool_choice\": {\"type\": \"function\", \"function\": {\"name\": \"note\"}}, \"parallel_tool_calls\": true}","tag":"{\"type\":\"structural_tag\",\"format\":{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"note\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"text\":{\"type\":\"string\"},\"tags\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"maxItems\":2}},\"required\":[\"text\"]},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"const_string\",\"value\":\"\"}]}}","active":true,"steps":[{"window":[67,198,198,179],"keep":3,"first":0,"stop":false,"rows":[["e969cc599117b642",2],["f44d7b1682b3347f",1],["18dc8579859e6bb3",1]],"advance":[198,198,27]},{"window":[27,285],"keep":2,"first":0,"stop":false,"rows":[["01134894a381097e",2],["1ac3cfee443070d4",2]],"advance":[285]},{"window":[285,273,29,93],"keep":3,"first":0,"stop":false,"rows":[["1ac3cfee443070d4",2],["0f3ff8d4a6ba23d2",1],["f44d7b1682b3347f",1]],"advance":[273,29]},{"window":[29,198,27,285,220,272,220,259,28],"keep":9,"first":0,"stop":false,"rows":[["f44d7b1682b3347f",1],["18dc8579859e6bb3",1],["01134894a381097e",2],["16ac141c6832090d",1],["39759c9b02c87122",2],["16ac141c6832090d",1],["b84ef774785a90c9",2],["2ac4811b6f080943",2],["3bab55816f9c1c63",1]],"advance":[198,27,171]},{"window":[171,121,250],"keep":3,"first":0,"stop":false,"rows":[["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1]],"advance":[121]},{"window":[121,250,35,50,181],"keep":4,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1],["78edc49140ab8efb",1]],"advance":[250,35,50]},{"window":[50,44,43,171,121,250],"keep":6,"first":0,"stop":false,"rows":[["78edc49140ab8efb",1],["0c25f068e5a271aa",1],["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["16ac141c6832090d",1]],"advance":[44,43,171,121]},{"window":[121],"keep":1,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1]],"advance":[250]},{"window":[250,220,72,12],"keep":3,"first":0,"stop":false,"rows":[["16ac141c6832090d",1],["39759c9b02c87122",2],["b2ff1bbe17acf6b2",1]],"advance":[220,72]},{"window":[72,77],"keep":2,"first":0,"stop":false,"rows":[["b2ff1bbe17acf6b2",1],["901c50779bff8a9b",1]],"advance":[77]},{"window":[77,85],"keep":2,"first":0,"stop":false,"rows":[["901c50779bff8a9b",1],["fde77043ec5b1050",1]],"advance":[85]},{"window":[85],"keep":1,"first":0,"stop":false,"rows":[["fde77043ec5b1050",1]],"advance":[78]},{"window":[78,74,68,231],"keep":3,"first":0,"stop":false,"rows":[["478d904b5aab5b8e",1],["71b0cde5e953ddaf",1],["16ac141c6832090d",1]],"advance":[74]},{"window":[74],"keep":1,"first":0,"stop":false,"rows":[["71b0cde5e953ddaf",1]],"advance":[68]},{"window":[68,171],"keep":1,"first":0,"stop":false,"rows":[["16ac141c6832090d",1]],"advance":[220]},{"window":[220,259,28],"keep":3,"first":0,"stop":false,"rows":[["b84ef774785a90c9",2],["2ac4811b6f080943",2],["3bab55816f9c1c63",1]],"advance":[77]},{"window":[77,64],"keep":2,"first":0,"stop":false,"rows":[["7d507dccdd7bf2e2",1],["f86c655263c87439",1]],"advance":[64]},{"window":[64,76],"keep":2,"first":0,"stop":false,"rows":[["f86c655263c87439",1],["71b0cde5e953ddaf",1]],"advance":[76,68]},{"window":[68,28,1,77],"keep":4,"first":0,"stop":false,"rows":[["2ac4811b6f080943",2],["3bab55816f9c1c63",1],["a290755ab1a135cc",2],["fde77043ec5b1050",1]],"advance":[28,1]},{"window":[1],"keep":1,"first":0,"stop":false,"rows":[["a290755ab1a135cc",2]],"advance":[77]},{"window":[77],"keep":1,"first":0,"stop":false,"rows":[["fde77043ec5b1050",1]],"advance":[78]},{"window":[78,83],"keep":2,"first":0,"stop":false,"rows":[["d1d61829e98de7e7",1],["71b0cde5e953ddaf",1]],"advance":[83]},{"window":[83,186],"keep":1,"first":0,"stop":false,"rows":[["71b0cde5e953ddaf",1]],"advance":[68]},{"window":[68,1,29,275,275,198,198,268,275],"keep":9,"first":0,"stop":false,"rows":[["cd0bdf7e5dadc655",2],["0f3ff8d4a6ba23d2",1],["e969cc599117b642",2],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[1,29,275,275,198,198,268,275,201]}],"finished":false},{"name":"tools_strict_auto","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}, \"strict\": true}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}], \"tool_choice\": \"auto\", \"parallel_tool_calls\": true}","tag":"{\"type\":\"structural_tag\",\"format\":{\"type\":\"triggered_tags\",\"triggers\":[\"<|DSML| calls>\"],\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| calls>\\n\",\"content\":{\"type\":\"tags_with_separator\",\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"get_weather\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\",\"maxLength\":12},\"unit\":{\"enum\":[\"c\",\"f\"]},\"days\":{\"type\":\"integer\",\"minimum\":1,\"maximum\":7}},\"required\":[\"city\",\"unit\"],\"additionalProperties\":false},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"note\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"text\":{\"type\":\"string\"},\"tags\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"maxItems\":2}},\"required\":[\"text\"]},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"}],\"separator\":\"\",\"at_least_one\":true,\"stop_after_first\":false},\"end\":\"\"}],\"at_least_one\":false,\"stop_after_first\":false,\"excludes\":[\"\",\"\"]}}","active":true,"steps":[{"window":[147,265,275,147,282,282],"keep":4,"first":0,"stop":true,"rows":[["1ef3b8a733c25f1c",283],["1ef3b8a733c25f1c",283],["1ef3b8a733c25f1c",283],["1ef3b8a733c25f1c",283]],"advance":[282]}],"finished":true},{"name":"json_schema_thinking","kind":"json_schema","text":"{\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}","tag":null,"active":false,"steps":[{"window":[251],"keep":1,"first":1,"stop":false,"rows":[],"advance":[94]},{"window":[94,263,81,284],"keep":4,"first":3,"stop":false,"rows":[["1d5ede0d5993de19",2]],"advance":[263,255]},{"window":[255,41],"keep":2,"first":2,"stop":false,"rows":[],"advance":[41,287]},{"window":[287],"keep":1,"first":1,"stop":false,"rows":[],"advance":[284]},{"window":[284,90,198,220],"keep":4,"first":0,"stop":false,"rows":[["1d5ede0d5993de19",2],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[90,201]},{"window":[201,220,220,198,201,197,198,201,220],"keep":9,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[220,220,198,198]},{"window":[198,220,197],"keep":3,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[220,197,198]},{"window":[198,198,268,268,196],"keep":4,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[198,197]},{"window":[197,197,54],"keep":2,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[197,197]},{"window":[197,275,220,1],"keep":4,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["940300c8f1b2b534",2]],"advance":[275,198]},{"window":[198,197,220,275,275,197],"keep":6,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7]],"advance":[197,220,275,275,197,197]},{"window":[197,198,201,268,201,275,275,1,66],"keep":9,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["3bab55816f9c1c63",1],["940300c8f1b2b534",2],["0195320b34a1e3b8",1]],"advance":[198]},{"window":[198,197,220,201,275,268],"keep":6,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2e9ffa3ec890148",7],["a2dba1bfbc7cfa8c",5]],"advance":[197,220,201,197]},{"window":[197,1,66,72,214],"keep":4,"first":0,"stop":false,"rows":[["a2e9ffa3ec890148",7],["940300c8f1b2b534",2],["0195320b34a1e3b8",1],["d1d61829e98de7e7",1]],"advance":[1,66,72,83]},{"window":[83,88,1,201,275,98],"keep":5,"first":0,"stop":false,"rows":[["449ca9e633bd281f",1],["3405cabc04f6f848",2],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[88,1,201,268]},{"window":[268,197,275,198,275,275],"keep":6,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[197,198]},{"window":[198,201,275,197,201,268,198,275,197],"keep":9,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7],["01b401fc4219b614",7]],"advance":[201,275,197,201,268,198,275,197,275]},{"window":[275,25,220],"keep":3,"first":0,"stop":false,"rows":[["01b401fc4219b614",7],["6854501f0132b612",10],["6854501f0132b612",10]],"advance":[25,277]},{"window":[277],"keep":1,"first":0,"stop":false,"rows":[["8cf2d085632efb8e",197]],"advance":[60]},{"window":[60,39],"keep":2,"first":0,"stop":false,"rows":[["8cf2d085632efb8e",197],["8cf2d085632efb8e",197]],"advance":[136]},{"window":[136,105,208,257,145,109],"keep":6,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["8cf2d085632efb8e",197],["3f509174ac09f493",196],["94ae616c304096a4",186],["7721af923e2625d6",64],["f61d3673d08954f5",184]],"advance":[105,208,140]},{"window":[140,101],"keep":2,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["a9c31053618c6136",195]],"advance":[101,152]},{"window":[152,115],"keep":2,"first":0,"stop":false,"rows":[["7721af923e2625d6",64],["a9c31053618c6136",195]],"advance":[114]},{"window":[114,216,142],"keep":3,"first":0,"stop":false,"rows":[["a9c31053618c6136",195],["759979dd3deb5878",193],["7721af923e2625d6",64]],"advance":[216,219]}],"finished":false},{"name":"tools_required_thinking","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}, {\"type\": \"function\", \"function\": {\"name\": \"loose\", \"parameters\": {\"type\": \"object\"}, \"strict\": false}}], \"tool_choice\": \"required\", \"parallel_tool_calls\": true}","tag":"{\"type\":\"structural_tag\",\"format\":{\"type\":\"sequence\",\"elements\":[{\"type\":\"const_string\",\"value\":\"\\n\\n<|DSML| calls>\\n\"},{\"type\":\"tags_with_separator\",\"tags\":[{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"get_weather\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"city\":{\"type\":\"string\",\"maxLength\":12},\"unit\":{\"enum\":[\"c\",\"f\"]},\"days\":{\"type\":\"integer\",\"minimum\":1,\"maximum\":7}},\"required\":[\"city\",\"unit\"],\"additionalProperties\":false},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"note\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":{\"type\":\"object\",\"properties\":{\"text\":{\"type\":\"string\"},\"tags\":{\"type\":\"array\",\"items\":{\"type\":\"string\"},\"maxItems\":2}},\"required\":[\"text\"]},\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"},{\"type\":\"tag\",\"begin\":\"<|DSML| invoke name=\\\"loose\\\">\\n\",\"content\":{\"type\":\"json_schema\",\"json_schema\":true,\"style\":\"deepseek_v4_1_xml\",\"any_order\":false,\"max_whitespace_cnt\":32,\"excludes\":[]},\"end\":\"\\n\"}],\"separator\":\"\",\"at_least_one\":true,\"stop_after_first\":false},{\"type\":\"const_string\",\"value\":\"\"}]}}","active":false,"steps":[{"window":[38,14],"keep":2,"first":2,"stop":false,"rows":[],"advance":[23]},{"window":[23,169,48,284,275,27,171,121,250],"keep":9,"first":3,"stop":false,"rows":[["e969cc599117b642",2],["18dc8579859e6bb3",1],["01134894a381097e",2],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1]],"advance":[169,48,284,275,27,285]},{"window":[285,273,29,198],"keep":4,"first":0,"stop":false,"rows":[["1ac3cfee443070d4",2],["0f3ff8d4a6ba23d2",1],["f44d7b1682b3347f",1],["18dc8579859e6bb3",1]],"advance":[273,29]},{"window":[29],"keep":1,"first":0,"stop":false,"rows":[["f44d7b1682b3347f",1]],"advance":[198]},{"window":[198,27,97],"keep":2,"first":0,"stop":false,"rows":[["18dc8579859e6bb3",1],["01134894a381097e",2]],"advance":[27,171]},{"window":[171,121,250,35,50,44],"keep":6,"first":0,"stop":false,"rows":[["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1],["a1c1987d9c6afc01",1],["cd5ed2a33e593c70",1],["78edc49140ab8efb",1],["0c25f068e5a271aa",1]],"advance":[121,250,35,50,44,43]},{"window":[43,171,121],"keep":3,"first":0,"stop":false,"rows":[["e1dfcc033d5b0d2a",1],["ac6acd14b749f97f",1],["d63af49dc26bcaa7",1]],"advance":[171]},{"window":[171],"keep":1,"first":0,"stop":false,"rows":[["ac6acd14b749f97f",1]],"advance":[121]},{"window":[121,250,220,272],"keep":4,"first":0,"stop":false,"rows":[["d63af49dc26bcaa7",1],["16ac141c6832090d",1],["39759c9b02c87122",2],["16ac141c6832090d",1]],"advance":[250,220,272,220]},{"window":[220],"keep":1,"first":0,"stop":false,"rows":[["b84ef774785a90c9",2]],"advance":[77]},{"window":[77,64,76,68],"keep":4,"first":0,"stop":false,"rows":[["7d507dccdd7bf2e2",1],["f86c655263c87439",1],["71b0cde5e953ddaf",1],["2ac4811b6f080943",2]],"advance":[64]},{"window":[64,76],"keep":2,"first":0,"stop":false,"rows":[["f86c655263c87439",1],["71b0cde5e953ddaf",1]],"advance":[76,68]},{"window":[68,28,1,279],"keep":4,"first":0,"stop":false,"rows":[["2ac4811b6f080943",2],["3bab55816f9c1c63",1],["0acdc0b465020e3c",5],["f6839908d75a2d3f",2]],"advance":[28,1,279,280]},{"window":[280,1,29,198,220,201,220,197,198],"keep":9,"first":0,"stop":false,"rows":[["cd0bdf7e5dadc655",2],["0f3ff8d4a6ba23d2",1],["e969cc599117b642",2],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[1,29,275]},{"window":[275,198,198,268],"keep":4,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[198,198,275]},{"window":[275,197,201,201,197,27],"keep":6,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["01134894a381097e",2]],"advance":[197,201,198]},{"window":[198,275,198],"keep":3,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[275]},{"window":[275,275,201,220,275,197],"keep":6,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[275,201,220]},{"window":[220],"keep":1,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7]],"advance":[201]},{"window":[201,275,197],"keep":3,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7]],"advance":[275,220]},{"window":[220],"keep":1,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7]],"advance":[275]},{"window":[275],"keep":1,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7]],"advance":[201]},{"window":[201,275,220,27,285,220],"keep":6,"first":0,"stop":false,"rows":[["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["5fd16ef2e271af66",7],["01134894a381097e",2],["16ac141c6832090d",1],["9cac882ea6761c02",2]],"advance":[275,27]},{"window":[27,285,220,79],"keep":4,"first":0,"stop":false,"rows":[["01134894a381097e",2],["16ac141c6832090d",1],["9cac882ea6761c02",2],["7d507dccdd7bf2e2",1]],"advance":[171]}],"finished":false}],"bad":[{"name":"bad_schema","kind":"json_schema","text":"{\"type\": \"object\", \"properties\": {\"a\": {\"type\": \"nosuch\"}}}","field":"guided_json","message":"guided_json: the grammar cannot be enforced: Unsupported type \"nosuch\""},{"name":"bad_regex","kind":"regex","text":"(ab","field":"guided_regex","message":"guided_regex: the grammar cannot be enforced: Regex parsing error at position 4: The parenthesis is not closed."},{"name":"bad_ebnf","kind":"grammar","text":"root ::= undefined_rule","field":"guided_grammar","message":"guided_grammar: the grammar cannot be enforced: EBNF parser error at line 1, column 10: Rule \"undefined_rule\" is not defined"},{"name":"named_missing","kind":"tools","text":"{\"tools\": [{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"maxLength\": 12}, \"unit\": {\"enum\": [\"c\", \"f\"]}, \"days\": {\"type\": \"integer\", \"minimum\": 1, \"maximum\": 7}}, \"required\": [\"city\", \"unit\"], \"additionalProperties\": false}}}, {\"type\": \"function\", \"function\": {\"name\": \"note\", \"parameters\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}, \"tags\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}, \"maxItems\": 2}}, \"required\": [\"text\"]}}}, {\"type\": \"function\", \"function\": {\"name\": \"loose\", \"parameters\": {\"type\": \"object\"}, \"strict\": false}}], \"tool_choice\": {\"type\": \"function\", \"function\": {\"name\": \"nope\"}}, \"parallel_tool_calls\": true}","field":"tools","message":"tools: the grammar cannot be enforced: The tool with name 'nope' is not found in the tools list."}],"tokenizer":"synthetic"} diff --git a/zig/tests/grammar/fixtures/synthetic/tokenizer.json b/zig/tests/grammar/fixtures/synthetic/tokenizer.json new file mode 100644 index 0000000000..ceda72104a --- /dev/null +++ b/zig/tests/grammar/fixtures/synthetic/tokenizer.json @@ -0,0 +1 @@ +{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [{"id": 281, "content": "<|begin▁of▁sentence|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 282, "content": "<|end▁of▁sentence|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 283, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 284, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 285, "content": "|DSML|", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 286, "content": "<|User|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}], "normalizer": null, "pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": true}, "post_processor": null, "decoder": {"type": "ByteLevel", "add_prefix_space": true, "trim_offsets": true, "use_regex": true}, "model": {"type": "BPE", "dropout": null, "unk_token": null, "continuing_subword_prefix": null, "end_of_word_suffix": null, "fuse_unk": false, "byte_fallback": false, "ignore_merges": false, "vocab": {"!": 0, "\"": 1, "#": 2, "$": 3, "%": 4, "&": 5, "'": 6, "(": 7, ")": 8, "*": 9, "+": 10, ",": 11, "-": 12, ".": 13, "/": 14, "0": 15, "1": 16, "2": 17, "3": 18, "4": 19, "5": 20, "6": 21, "7": 22, "8": 23, "9": 24, ":": 25, ";": 26, "<": 27, "=": 28, ">": 29, "?": 30, "@": 31, "A": 32, "B": 33, "C": 34, "D": 35, "E": 36, "F": 37, "G": 38, "H": 39, "I": 40, "J": 41, "K": 42, "L": 43, "M": 44, "N": 45, "O": 46, "P": 47, "Q": 48, "R": 49, "S": 50, "T": 51, "U": 52, "V": 53, "W": 54, "X": 55, "Y": 56, "Z": 57, "[": 58, "\\": 59, "]": 60, "^": 61, "_": 62, "`": 63, "a": 64, "b": 65, "c": 66, "d": 67, "e": 68, "f": 69, "g": 70, "h": 71, "i": 72, "j": 73, "k": 74, "l": 75, "m": 76, "n": 77, "o": 78, "p": 79, "q": 80, "r": 81, "s": 82, "t": 83, "u": 84, "v": 85, "w": 86, "x": 87, "y": 88, "z": 89, "{": 90, "|": 91, "}": 92, "~": 93, "¡": 94, "¢": 95, "£": 96, "¤": 97, "¥": 98, "¦": 99, "§": 100, "¨": 101, "©": 102, "ª": 103, "«": 104, "¬": 105, "®": 106, "¯": 107, "°": 108, "±": 109, "²": 110, "³": 111, "´": 112, "µ": 113, "¶": 114, "·": 115, "¸": 116, "¹": 117, "º": 118, "»": 119, "¼": 120, "½": 121, "¾": 122, "¿": 123, "À": 124, "Á": 125, "Â": 126, "Ã": 127, "Ä": 128, "Å": 129, "Æ": 130, "Ç": 131, "È": 132, "É": 133, "Ê": 134, "Ë": 135, "Ì": 136, "Í": 137, "Î": 138, "Ï": 139, "Ð": 140, "Ñ": 141, "Ò": 142, "Ó": 143, "Ô": 144, "Õ": 145, "Ö": 146, "×": 147, "Ø": 148, "Ù": 149, "Ú": 150, "Û": 151, "Ü": 152, "Ý": 153, "Þ": 154, "ß": 155, "à": 156, "á": 157, "â": 158, "ã": 159, "ä": 160, "å": 161, "æ": 162, "ç": 163, "è": 164, "é": 165, "ê": 166, "ë": 167, "ì": 168, "í": 169, "î": 170, "ï": 171, "ð": 172, "ñ": 173, "ò": 174, "ó": 175, "ô": 176, "õ": 177, "ö": 178, "÷": 179, "ø": 180, "ù": 181, "ú": 182, "û": 183, "ü": 184, "ý": 185, "þ": 186, "ÿ": 187, "Ā": 188, "ā": 189, "Ă": 190, "ă": 191, "Ą": 192, "ą": 193, "Ć": 194, "ć": 195, "Ĉ": 196, "ĉ": 197, "Ċ": 198, "ċ": 199, "Č": 200, "č": 201, "Ď": 202, "ď": 203, "Đ": 204, "đ": 205, "Ē": 206, "ē": 207, "Ĕ": 208, "ĕ": 209, "Ė": 210, "ė": 211, "Ę": 212, "ę": 213, "Ě": 214, "ě": 215, "Ĝ": 216, "ĝ": 217, "Ğ": 218, "ğ": 219, "Ġ": 220, "ġ": 221, "Ģ": 222, "ģ": 223, "Ĥ": 224, "ĥ": 225, "Ħ": 226, "ħ": 227, "Ĩ": 228, "ĩ": 229, "Ī": 230, "ī": 231, "Ĭ": 232, "ĭ": 233, "Į": 234, "į": 235, "İ": 236, "ı": 237, "IJ": 238, "ij": 239, "Ĵ": 240, "ĵ": 241, "Ķ": 242, "ķ": 243, "ĸ": 244, "Ĺ": 245, "ĺ": 246, "Ļ": 247, "ļ": 248, "Ľ": 249, "ľ": 250, "Ŀ": 251, "ŀ": 252, "Ł": 253, "ł": 254, "Ń": 255, "true": 256, "false": 257, "null": 258, "name": 259, "city": 260, "Paris": 261, "yes": 262, "no": 263, "\":": 264, "{\"": 265, "\"}": 266, ", ": 267, " ": 268, "12": 269, "Ġthe": 270, "Ġ{": 271, "invoke": 272, "Ġcalls": 273, "parameter": 274, "ĊĊ": 275, "=\"": 276, "\">": 277, " str: + return subprocess.run(["git", "-C", repo, "show", f"{commit}:{path}"], check=True, capture_output=True, + text=True).stdout + + (spark / "grammar.py").write_text(show("src/tensorfold/families/glm5_next/spark/grammar.py")) + (spark / "server.py").write_text("def tool_fixes(raw=None):\n return frozenset()\n") + (ds / "structured.py").write_text(show("src/tensorfold/families/deepseek_v41/cuda/structured.py")) + sys.path.insert(0, str(root)) + from tensorfold.families.deepseek_v41.cuda import structured + from tensorfold.families.glm5_next.spark import grammar as gm + return gm, structured + + +def digest(words: np.ndarray) -> str: + return hashlib.sha256(np.ascontiguousarray(words, dtype=" str: + h = hashlib.sha256() + for s in enc: + b = s.encode() + h.update(len(b).to_bytes(4, "little")) + h.update(b) + return h.hexdigest()[:16] + + +def allowed(bits_row: np.ndarray, vocab: int) -> np.ndarray: + u = bits_row.view(np.uint32) + ids = np.nonzero(((u[:, None] >> np.arange(32, dtype=np.uint32)) & 1).reshape(-1)[:vocab])[0] + return ids + + +def synthetic_tokenizer(path: Path) -> None: + """A small byte-level BPE tokenizer.json: 256 byte symbols, some words, eos/think/tool-markup added tokens.""" + + from tokenizers.pre_tokenizers import ByteLevel + + alphabet = ByteLevel.alphabet() + byte_map = sorted(alphabet) + words = ["true", "false", "null", "name", "city", "Paris", "yes", "no", "\":", "{\"", "\"}", ", ", " ", "12", + "Ġthe", "Ġ{", "invoke", "Ġcalls", "parameter", "ĊĊ", "Ċ", "=\"", "\">", "", "<|end▁of▁sentence|>", "", "", "|DSML|", "<|User|>"]: + added.append({"id": len(vocab) + len(added), "content": content, "single_word": False, "lstrip": False, + "rstrip": False, "normalized": False, "special": True}) + doc = {"version": "1.0", "truncation": None, "padding": None, "added_tokens": added, "normalizer": None, + "pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": False, "trim_offsets": True, "use_regex": True}, + "post_processor": None, + "decoder": {"type": "ByteLevel", "add_prefix_space": True, "trim_offsets": True, "use_regex": True}, + "model": {"type": "BPE", "dropout": None, "unk_token": None, "continuing_subword_prefix": None, + "end_of_word_suffix": None, "fuse_unk": False, "byte_fallback": False, "ignore_merges": False, + "vocab": vocab, "merges": []}} + path.write_text(json.dumps(doc, ensure_ascii=False)) + + +WEATHER = {"type": "object", "properties": {"city": {"type": "string", "maxLength": 12}, + "unit": {"enum": ["c", "f"]}, + "days": {"type": "integer", "minimum": 1, "maximum": 7}}, + "required": ["city", "unit"], "additionalProperties": False} +TOOLS = [{"type": "function", "function": {"name": "get_weather", "description": "w", "parameters": WEATHER}}, + {"type": "function", "function": {"name": "note", "parameters": {"type": "object", "properties": { + "text": {"type": "string"}, "tags": {"type": "array", "items": {"type": "string"}, "maxItems": 2}}, + "required": ["text"]}}}, + {"type": "function", "function": {"name": "loose", "strict": False, "parameters": {"type": "object"}}}] + + +def cases(gm): + out = [] + + def body(**kw): + return kw + + out.append(("json_object", body(response_format={"type": "json_object"}), False)) + out.append(("json_schema", body(response_format={"type": "json_schema", "json_schema": {"name": "w", "schema": WEATHER}}), False)) + out.append(("array", body(guided_json={"type": "array", "items": {"type": "integer"}, "minItems": 1, "maxItems": 3}), False)) + out.append(("regex", body(guided_regex=r"\d{3}-[a-z]{2,5}( ok)?"), False)) + out.append(("choice", body(guided_choice=["yes", "no", "maybe \"so\"", "été"]), False)) + out.append(("ebnf", body(guided_grammar='root ::= expr\nexpr ::= term ("+" term)*\nterm ::= [0-9]+ | "(" expr ")"'), False)) + out.append(("tools_required", body(tools=TOOLS, tool_choice="required"), False)) + out.append(("tools_required_single", body(tools=TOOLS, tool_choice="required", parallel_tool_calls=False), False)) + out.append(("tools_named", body(tools=TOOLS, tool_choice={"type": "function", "function": {"name": "note"}}), False)) + strict = [dict(TOOLS[0], function=dict(TOOLS[0]["function"], strict=True)), TOOLS[1]] + out.append(("tools_strict_auto", body(tools=strict), False)) + out.append(("json_schema_thinking", body(response_format={"type": "json_schema", "json_schema": {"schema": WEATHER}}), True)) + out.append(("tools_required_thinking", body(tools=TOOLS, tool_choice="required"), True)) + return out + + +BAD = [("bad_schema", {"guided_json": {"type": "object", "properties": {"a": {"type": "nosuch"}}}}), + ("bad_regex", {"guided_regex": "(ab"}), + ("bad_ebnf", {"guided_grammar": "root ::= undefined_rule"}), + ("named_missing", {"tools": TOOLS, "tool_choice": {"type": "function", "function": {"name": "nope"}}})] + + +def walk(gm, grammars, spec, active, rng: random.Random, steps: int, vocab: int, eos: int, think_end: int): + """Drafted reply rounds under Constraint: scratch-matcher draft chains (some wrong), cut, fill, then advance.""" + + bound = gm.Bound(spec, grammars.compile(spec), think_end, active) + con = grammars.constraint(bound) + hist: list[int] = [] + recs = [] + pending = rng.randrange(vocab) + for step in range(steps): + if con.finished: + break + scratch = clone(grammars, bound, hist) + n = rng.choice([1, 2, 3, 4, 6, 9]) + drafts = [] + for i in range(n - 1): + if rng.random() < 0.12: + drafts.append(rng.randrange(vocab)) # a wrong draft (cut unless the grammar takes it) + break + t = pick(rng, scratch, con.words, vocab, eos, think_end, step + i) + drafts.append(t) + try: + scratch.advance([t]) + except gm.GrammarError: + break + tokens = [pending] + drafts + win = con.cut(tokens) + keep = len(win.tokens) + stop = stop_at(con, tokens, keep, think_end) + con.fill(win) + rows = [] + for j in range(len(win.rows)): + rows.append([digest(win.bits[j]), int(len(allowed(win.bits[j], vocab)))]) + acc = rng.randrange(keep) + if win.rows and acc >= win.rows[0]: + ok = allowed(win.bits[acc - win.rows[0]], vocab) + bonus = eos if eos in ok and rng.random() < 0.15 else int(ok[rng.randrange(min(len(ok), 400))]) + else: + bonus = think_end if step >= 2 and rng.random() < 0.4 else rng.randrange(vocab) + adv = [int(t) for t in tokens[1:acc + 1]] + [int(bonus)] + con.advance(adv) + hist.extend(adv) + recs.append({"window": [int(t) for t in tokens], "keep": keep, "first": win.rows[0] if win.rows else keep, + "stop": stop, "rows": rows, "advance": adv}) + pending = bonus + return recs, con + + +def clone(grammars, bound, hist): + """A constraint at the committed path `hist` (replayed).""" + + c = grammars.constraint(bound) + c.advance(hist) + return c + + +def stop_at(con, tokens, keep, think_end) -> bool: + """Whether the cut ended at a draft that completes the grammar (it and every later row are cut).""" + + if keep >= len(tokens): + return False + m = con.m + active = con.active + acc = 0 + try: + for t in tokens[1:keep]: + if active: + m.accept_token(t) + acc += 1 + else: + active = t == think_end + t = tokens[keep] + if not active or t in con.never: + return False + if not m.accept_token(t): + return False + acc += 1 + return bool(m.is_terminated()) + finally: + if acc: + m.rollback(acc) + + +def pick(rng, scratch, words, vocab, eos, think_end, unthink) -> int: + """A token the scratch matcher takes next (sometimes the stop token); before , any token or .""" + + if not scratch.active: + return think_end if unthink >= 3 and rng.random() < 0.5 else rng.randrange(vocab) + if scratch.m.is_terminated(): + return eos + bits = np.full((1, words), -1, dtype=np.int32) + scratch.m.fill_next_token_bitmask(bits, 0) + ok = allowed(bits[0], vocab) + if eos in ok and rng.random() < 0.15: + return eos + if len(ok) == 0: + return eos + # prefer short ASCII tokens so walks reach the grammar's end + return int(ok[rng.randrange(min(len(ok), 400))]) if rng.random() < 0.7 else int(ok[rng.randrange(len(ok))]) + + +def main() -> None: + ap = argparse.ArgumentParser(description=( + "Write test-grammar goldens: the tokenizer views (vocab type, prefix space, vocabulary and special-id digests, " + "think ids); per case the spec, the tools spec's structural tag JSON, the think state and each step's window, " + "cut (keep, first constrained row, stop), per-row bitmask digest (first 16 hex of sha256 over the int32 words) " + "and allowed count, and advanced tokens; and the refusal message of each bad spec. Needs xgrammar==0.2.8, " + "tokenizers and numpy.")) + ap.add_argument("--reference-repo", required=True, + help="checkout of a Python serving implementation (structured.py and grammar.py)") + ap.add_argument("--commit", required=True, help="reference commit to read structured.py and grammar.py from") + ap.add_argument("--tokenizer", help="dir with tokenizer.json") + ap.add_argument("--synthetic", action="store_true", + help="use a small synthetic tokenizer (written as synthetic/tokenizer.json beside --out)") + ap.add_argument("--vocab-size", type=int, default=129280, help="logits columns (ignored with --synthetic)") + ap.add_argument("--eos", type=int, default=1, help="stop token id (ignored with --synthetic)") + ap.add_argument("--steps", type=int, default=24, help="most rounds walked per case") + ap.add_argument("--seed", type=int, default=7, help="walk RNG seed") + ap.add_argument("--out", required=True, help="golden JSON to write") + a = ap.parse_args() + gm, structured = reference_modules(a.reference_repo, a.commit) + import xgrammar as xgr + + out = Path(a.out) + if a.synthetic: + tok_dir = out.parent / "synthetic" + tok_dir.mkdir(parents=True, exist_ok=True) + synthetic_tokenizer(tok_dir / "tokenizer.json") + vocab = 288 + eos = 1 + json.loads((tok_dir / "tokenizer.json").read_text())["added_tokens"][0]["id"] + else: + tok_dir = Path(a.tokenizer) + vocab, eos = a.vocab_size, a.eos + grammars = structured.Dsv41Grammars.from_model(tok_dir, vocab, (eos,)) + views = {} + path = tok_dir / "tokenizer.json" + for name, keep in (("text", ()), ("tools", structured.DSML_TOKENS)): + enc, backend, t_open, t_end = gm._vocab(path, vocab, keep) + meta = xgr.TokenizerInfo._detect_metadata_from_hf(backend) + views[name] = {"vocab_type": int(meta["vocab_type"].value), "add_prefix_space": bool(meta["add_prefix_space"]), + "vocab": vocab_digest(enc), "never_digest": hashlib.sha256(json.dumps(sorted(int(t) for t in grammars.special[name])).encode()).hexdigest()[:16], + "never_count": len(grammars.special[name])} + rng = random.Random(a.seed) + recs = [] + for name, body, thinking in cases(gm): + spec = structured.Host.spec(structured.Host.__new__(structured.Host), body) + tag = None + if spec.kind == "tools": + d = json.loads(spec.text) + tag = xgr.get_model_structural_tag(structured.MODEL_TAG, tools=d["tools"], tool_choice=d["tool_choice"], + reasoning="disabled", max_whitespace_cnt=gm.BLANKS, + parallel_tool_calls=bool(d["parallel_tool_calls"])).model_dump_json(indent=None) + active = not thinking + steps, con = walk(gm, grammars, spec, active, rng, a.steps, vocab, eos, grammars.think_end) + recs.append({"name": name, "kind": spec.kind, "text": spec.text, "tag": tag, "active": active, + "steps": steps, "finished": bool(con.finished)}) + print(f"{name}: {len(steps)} steps, {sum(len(s.get('rows', [])) for s in steps)} rows, finished " + f"{con.finished}", file=sys.stderr) + bad = [] + for name, body in BAD: + try: + spec = structured.Host.spec(structured.Host.__new__(structured.Host), body) + grammars.compile(spec) + msg = None + except ValueError as exc: + msg = str(exc) + bad.append({"name": name, "kind": spec.kind, "text": spec.text, "field": spec.field, "message": msg}) + doc = {"xgrammar": "0.2.8", "reference": a.commit, "vocab_size": vocab, "eos": eos, "think_open": grammars.think_open, + "think_end": grammars.think_end, "words": (vocab + 31) // 32, "views": views, "cases": recs, "bad": bad, + "tokenizer": "synthetic" if a.synthetic else "model"} + out.write_text(json.dumps(doc, ensure_ascii=False, indent=None, separators=(",", ":")) + "\n") + print(f"wrote {out} ({out.stat().st_size} bytes)", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/zig/vendor/xgrammar/3rdparty/dlpack/LICENSE b/zig/vendor/xgrammar/3rdparty/dlpack/LICENSE new file mode 100644 index 0000000000..20a9c8a7b4 --- /dev/null +++ b/zig/vendor/xgrammar/3rdparty/dlpack/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017 by Contributors + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/zig/vendor/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h b/zig/vendor/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h new file mode 100644 index 0000000000..bcb77949a8 --- /dev/null +++ b/zig/vendor/xgrammar/3rdparty/dlpack/include/dlpack/dlpack.h @@ -0,0 +1,332 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +/** + * \brief Compatibility with C++ + */ +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current major version of dlpack */ +#define DLPACK_MAJOR_VERSION 1 + +/*! \brief The current minor version of dlpack */ +#define DLPACK_MINOR_VERSION 0 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! + * \brief The DLPack version. + * + * A change in major version indicates that we have changed the + * data layout of the ABI - DLManagedTensorVersioned. + * + * A change in minor version indicates that we have added new + * code, such as a new device type, but the ABI is kept the same. + * + * If an obtained DLPack tensor has a major version that disagrees + * with the version number specified in this header file + * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter + * (and it is safe to do so). It is not safe to access any other fields + * as the memory layout will have changed. + * + * In the case of a minor version mismatch, the tensor can be safely used as + * long as the consumer knows how to interpret all fields. Minor version + * updates indicate the addition of enumeration values. + */ +typedef struct { + /*! \brief DLPack major version. */ + uint32_t major; + /*! \brief DLPack minor version. */ + uint32_t minor; +} DLPackVersion; + +/*! + * \brief The device type in DLDevice. + */ +#ifdef __cplusplus +typedef enum : int32_t { +#else +typedef enum { +#endif + /*! \brief CPU device */ + kDLCPU = 1, + /*! \brief CUDA GPU device */ + kDLCUDA = 2, + /*! + * \brief Pinned CUDA CPU memory by cudaMallocHost + */ + kDLCUDAHost = 3, + /*! \brief OpenCL devices. */ + kDLOpenCL = 4, + /*! \brief Vulkan buffer for next generation graphics. */ + kDLVulkan = 7, + /*! \brief Metal for Apple GPU. */ + kDLMetal = 8, + /*! \brief Verilog simulator buffer */ + kDLVPI = 9, + /*! \brief ROCm GPUs for AMD GPUs */ + kDLROCM = 10, + /*! + * \brief Pinned ROCm CPU memory allocated by hipMallocHost + */ + kDLROCMHost = 11, + /*! + * \brief Reserved extension device type, + * used for quickly test extension device + * The semantics can differ depending on the implementation. + */ + kDLExtDev = 12, + /*! + * \brief CUDA managed/unified memory allocated by cudaMallocManaged + */ + kDLCUDAManaged = 13, + /*! + * \brief Unified shared memory allocated on a oneAPI non-partititioned + * device. Call to oneAPI runtime is required to determine the device + * type, the USM allocation type and the sycl context it is bound to. + * + */ + kDLOneAPI = 14, + /*! \brief GPU support for next generation WebGPU standard. */ + kDLWebGPU = 15, + /*! \brief Qualcomm Hexagon DSP */ + kDLHexagon = 16, + /*! \brief Microsoft MAIA devices */ + kDLMAIA = 17, +} DLDeviceType; + +/*! + * \brief A Device for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! + * \brief The device index. + * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. + */ + int32_t device_id; +} DLDevice; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + /*! \brief signed integer */ + kDLInt = 0U, + /*! \brief unsigned integer */ + kDLUInt = 1U, + /*! \brief IEEE floating point */ + kDLFloat = 2U, + /*! + * \brief Opaque handle type, reserved for testing purposes. + * Frameworks need to agree on the handle data type for the exchange to be well-defined. + */ + kDLOpaqueHandle = 3U, + /*! \brief bfloat16 */ + kDLBfloat = 4U, + /*! + * \brief complex number + * (C/C++/Python layout: compact struct per complex number) + */ + kDLComplex = 5U, + /*! \brief boolean */ + kDLBool = 6U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. The data type is assumed to follow the + * native endian-ness. An explicit error message should be raised when attempting to + * export an array with non-native endianness + * + * Examples + * - float: type_code = 2, bits = 32, lanes = 1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 + * - int8: type_code = 0, bits = 8, lanes = 1 + * - std::complex: type_code = 5, bits = 64, lanes = 1 + * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits) + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The data pointer points to the allocated data. This will be CUDA + * device pointer or cl_mem handle in OpenCL. It may be opaque on some device + * types. This pointer is always aligned to 256 bytes as in CUDA. The + * `byte_offset` field should be used to point to the beginning of the data. + * + * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow, + * TVM, perhaps others) do not adhere to this 256 byte aligment requirement + * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed + * (after which this note will be updated); at the moment it is recommended + * to not rely on the data pointer being correctly aligned. + * + * For given DLTensor, the size of memory required to store the contents of + * data is calculated as follows: + * + * \code{.c} + * static inline size_t GetDataSize(const DLTensor* t) { + * size_t size = 1; + * for (tvm_index_t i = 0; i < t->ndim; ++i) { + * size *= t->shape[i]; + * } + * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; + * return size; + * } + * \endcode + * + * Note that if the tensor is of size zero, then the data pointer should be + * set to `NULL`. + */ + void* data; + /*! \brief The device of the tensor */ + DLDevice device; + /*! \brief Number of dimensions */ + int32_t ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor (in number of elements, not bytes) + * can be NULL, indicating tensor is compact and row-majored. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + * + * \note This data structure is used as Legacy DLManagedTensor + * in DLPack exchange and is deprecated after DLPack v0.8 + * Use DLManagedTensorVersioned instead. + * This data structure may get renamed or deleted in future versions. + * + * \sa DLManagedTensorVersioned + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void * manager_ctx; + /*! + * \brief Destructor - this should be called + * to destruct the manager_ctx which backs the DLManagedTensor. It can be + * NULL if there is no way for the caller to provide a reasonable destructor. + * The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensor * self); +} DLManagedTensor; + +// bit masks used in in the DLManagedTensorVersioned + +/*! \brief bit mask to indicate that the tensor is read only. */ +#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) + +/*! + * \brief bit mask to indicate that the tensor is a copy made by the producer. + * + * If set, the tensor is considered solely owned throughout its lifetime by the + * consumer, until the producer-provided deleter is invoked. + */ +#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) + +/*! + * \brief A versioned and managed C Tensor object, manage memory of DLTensor. + * + * This data structure is intended to facilitate the borrowing of DLTensor by + * another framework. It is not meant to transfer the tensor. When the borrowing + * framework doesn't need the tensor, it should call the deleter to notify the + * host that the resource is no longer needed. + * + * \note This is the current standard DLPack exchange data structure. + */ +struct DLManagedTensorVersioned { + /*! + * \brief The API and ABI version of the current managed Tensor + */ + DLPackVersion version; + /*! + * \brief the context of the original host framework. + * + * Stores DLManagedTensorVersioned is used in the + * framework. It can also be NULL. + */ + void *manager_ctx; + /*! + * \brief Destructor. + * + * This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned. + * It can be NULL if there is no way for the caller to provide a reasonable + * destructor. The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensorVersioned *self); + /*! + * \brief Additional bitmask flags information about the tensor. + * + * By default the flags should be set to 0. + * + * \note Future ABI changes should keep everything until this field + * stable, to ensure that deleter can be correctly called. + * + * \sa DLPACK_FLAG_BITMASK_READ_ONLY + * \sa DLPACK_FLAG_BITMASK_IS_COPIED + */ + uint64_t flags; + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; +}; + +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/zig/vendor/xgrammar/3rdparty/picojson/picojson.h b/zig/vendor/xgrammar/3rdparty/picojson/picojson.h new file mode 100644 index 0000000000..f9e8e117d0 --- /dev/null +++ b/zig/vendor/xgrammar/3rdparty/picojson/picojson.h @@ -0,0 +1,1318 @@ +/* + * Copyright 2009-2010 Cybozu Labs, Inc. + * Copyright 2011-2014 Kazuho Oku + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + */ +#pragma once + +#ifndef PICOJSON_USE_INT64 +#define PICOJSON_USE_INT64 +#ifndef __STDC_FORMAT_MACROS +#define __STDC_FORMAT_MACROS 1 +#endif +#endif + +// If PICOJSON_USE_ORDERED_OBJECT is set, picojson uses object_with_ordered_keys, which maintains +// the insertion order of keys, i.e. the order of keys in the json string. +// This macro is set by default. +#ifndef PICOJSON_USE_ORDERED_OBJECT +#define PICOJSON_USE_ORDERED_OBJECT 1 +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// for isnan/isinf +#if __cplusplus >= 201103L +#include +#else +extern "C" { +#ifdef _MSC_VER +#include +#elif defined(__INTEL_COMPILER) +#include +#else +#include +#endif +} +#endif + +#ifndef PICOJSON_USE_RVALUE_REFERENCE +#if (defined(__cpp_rvalue_references) && __cpp_rvalue_references >= 200610) || \ + (defined(_MSC_VER) && _MSC_VER >= 1600) +#define PICOJSON_USE_RVALUE_REFERENCE 1 +#else +#define PICOJSON_USE_RVALUE_REFERENCE 0 +#endif +#endif // PICOJSON_USE_RVALUE_REFERENCE + +#ifndef PICOJSON_NOEXCEPT +#if PICOJSON_USE_RVALUE_REFERENCE +#define PICOJSON_NOEXCEPT noexcept +#else +#define PICOJSON_NOEXCEPT throw() +#endif +#endif + +// experimental support for int64_t (see README.mkdn for detail) +#ifdef PICOJSON_USE_INT64 +#include +#include +#endif + +// to disable the use of localeconv(3), set PICOJSON_USE_LOCALE to 0 +#ifndef PICOJSON_USE_LOCALE +#define PICOJSON_USE_LOCALE 1 +#endif +#if PICOJSON_USE_LOCALE +extern "C" { +#include +} +#endif + +#ifndef PICOJSON_ASSERT +#ifndef PICOJSON_DISABLE_EXCEPTION +#define PICOJSON_ASSERT(e) \ + do { \ + if (!(e)) throw std::runtime_error(#e); \ + } while (0) +#else +#define PICOJSON_ASSERT(e) \ + do { \ + if (!(e)) std::abort(); \ + } while (0) +#endif // PICOJSON_DISABLE_EXCEPTION +#endif + +#ifdef _MSC_VER +#define SNPRINTF _snprintf_s +#pragma warning(push) +#pragma warning(disable : 4244) // conversion from int to char +#pragma warning(disable : 4127) // conditional expression is constant +#pragma warning(disable : 4702) // unreachable code +#else +#define SNPRINTF snprintf +#endif + +namespace picojson { + +enum { + null_type, + boolean_type, + number_type, + string_type, + array_type, + object_type +#ifdef PICOJSON_USE_INT64 + , + int64_type +#endif +}; + +enum { INDENT_WIDTH = 2 }; + +struct null {}; + +class object_with_ordered_keys; + +class value { + public: + typedef std::vector array; +#ifdef PICOJSON_USE_ORDERED_OBJECT + typedef object_with_ordered_keys object; +#else + typedef std::unordered_map object; +#endif + + union _storage { + bool boolean_; + double number_; +#ifdef PICOJSON_USE_INT64 + int64_t int64_; +#endif + std::string* string_; + array* array_; + object* object_; + }; + + protected: + int type_; + _storage u_; + + public: + value(); + value(int type, bool); + explicit value(bool b); +#ifdef PICOJSON_USE_INT64 + explicit value(int64_t i); +#endif + explicit value(double n); + explicit value(const std::string& s); + explicit value(const array& a); + explicit value(const object& o); +#if PICOJSON_USE_RVALUE_REFERENCE + explicit value(std::string&& s); + explicit value(array&& a); + explicit value(object&& o); +#endif + explicit value(const char* s); + value(const char* s, size_t len); + ~value(); + value(const value& x); + value& operator=(const value& x); +#if PICOJSON_USE_RVALUE_REFERENCE + value(value&& x) PICOJSON_NOEXCEPT; + value& operator=(value&& x) PICOJSON_NOEXCEPT; +#endif + void swap(value& x) PICOJSON_NOEXCEPT; + template + bool is() const; + template + const T& get() const; + template + T& get(); + template + void set(const T&); +#if PICOJSON_USE_RVALUE_REFERENCE + template + void set(T&&); +#endif + bool evaluate_as_boolean() const; + const value& get(const size_t idx) const; + const value& get(const std::string& key) const; + value& get(const size_t idx); + value& get(const std::string& key); + + bool contains(const size_t idx) const; + bool contains(const std::string& key) const; + std::string to_str() const; + template + void serialize(Iter os, bool prettify = false) const; + std::string serialize(bool prettify = false) const; + + private: + template + // NOLINTNEXTLINE(runtime/explicit) + value(const T*); // intentionally defined to block implicit conversion of + // pointer to bool + template + static void _indent(Iter os, int indent); + template + void _serialize(Iter os, int indent) const; + std::string _serialize(int indent) const; + void clear(); +}; + +// The ordered version of hashmap. It has the same interface as std::unordered_map, but provides +// ordered_keys() to return the keys in the order they were inserted. +class object_with_ordered_keys : private std::unordered_map { + public: + using typename std::unordered_map::value_type; + using typename std::unordered_map::iterator; + using typename std::unordered_map::const_iterator; + + object_with_ordered_keys() = default; + object_with_ordered_keys(const object_with_ordered_keys&) = default; + object_with_ordered_keys(object_with_ordered_keys&&) = default; + object_with_ordered_keys(std::initializer_list init) + : std::unordered_map(init) { + for (const auto& pair : init) { + ordered_keys_.push_back(pair.first); + } + } + object_with_ordered_keys& operator=(const object_with_ordered_keys&) = default; + object_with_ordered_keys& operator=(object_with_ordered_keys&&) = default; + + using std::unordered_map::begin; + using std::unordered_map::end; + using std::unordered_map::cbegin; + using std::unordered_map::cend; + using std::unordered_map::empty; + using std::unordered_map::size; + using std::unordered_map::at; + using std::unordered_map::count; + using std::unordered_map::find; + using std::unordered_map::reserve; + + value& operator[](const std::string& key) { + if (count(key) == 0) { + ordered_keys_.push_back(key); + } + return std::unordered_map::operator[](key); + } + + const value& operator[](const std::string& key) const { + return std::unordered_map::at(key); + } + + void clear() { + std::unordered_map::clear(); + ordered_keys_.clear(); + } + + std::pair insert(const value_type& kv) { + if (!count(kv.first)) { + ordered_keys_.push_back(kv.first); + } + return std::unordered_map::insert(kv); + } + + template + std::pair emplace(Args&&... args) { + return insert(value_type(std::forward(args)...)); + } + + iterator erase(const_iterator it) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), it->first)); + return std::unordered_map::erase(it); + } + + iterator erase(iterator it) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), it->first)); + return std::unordered_map::erase(it); + } + + size_t erase(const std::string& key) { + if (std::unordered_map::erase(key)) { + ordered_keys_.erase(std::find(ordered_keys_.begin(), ordered_keys_.end(), key)); + return 1; + } else { + return 0; + } + } + + const std::vector& ordered_keys() const { return ordered_keys_; } + + friend bool operator==(const object_with_ordered_keys& lhs, const object_with_ordered_keys& rhs); + + private: + std::vector ordered_keys_; +}; + +inline bool operator==(const object_with_ordered_keys& lhs, const object_with_ordered_keys& rhs) { + return static_cast&>(lhs) == + static_cast&>(rhs); +} + +typedef value::array array; +typedef value::object object; + +inline value::value() : type_(null_type), u_() {} + +inline value::value(int type, bool) : type_(type), u_() { + switch (type) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(boolean_, false); + INIT(number_, 0.0); +#ifdef PICOJSON_USE_INT64 + INIT(int64_, 0); +#endif + INIT(string_, new std::string()); + INIT(array_, new array()); + INIT(object_, new object()); +#undef INIT + default: + break; + } +} + +inline value::value(bool b) : type_(boolean_type), u_() { u_.boolean_ = b; } + +#ifdef PICOJSON_USE_INT64 +inline value::value(int64_t i) : type_(int64_type), u_() { u_.int64_ = i; } +#endif + +inline value::value(double n) : type_(number_type), u_() { + if ( +#ifdef _MSC_VER + !_finite(n) +#elif __cplusplus >= 201103L + std::isnan(n) || std::isinf(n) +#else + isnan(n) || isinf(n) +#endif + ) { +#ifndef PICOJSON_DISABLE_EXCEPTION + throw std::overflow_error(""); +#else + std::abort(); +#endif + } + u_.number_ = n; +} + +inline value::value(const std::string& s) : type_(string_type), u_() { + u_.string_ = new std::string(s); +} + +inline value::value(const array& a) : type_(array_type), u_() { u_.array_ = new array(a); } + +inline value::value(const object& o) : type_(object_type), u_() { u_.object_ = new object(o); } + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(std::string&& s) : type_(string_type), u_() { + u_.string_ = new std::string(std::move(s)); +} + +inline value::value(array&& a) : type_(array_type), u_() { u_.array_ = new array(std::move(a)); } + +inline value::value(object&& o) : type_(object_type), u_() { + u_.object_ = new object(std::move(o)); +} +#endif + +inline value::value(const char* s) : type_(string_type), u_() { u_.string_ = new std::string(s); } + +inline value::value(const char* s, size_t len) : type_(string_type), u_() { + u_.string_ = new std::string(s, len); +} + +inline void value::clear() { + switch (type_) { +#define DEINIT(p) \ + case p##type: \ + delete u_.p; \ + break + DEINIT(string_); + DEINIT(array_); + DEINIT(object_); +#undef DEINIT + default: + break; + } +} + +inline value::~value() { clear(); } + +inline value::value(const value& x) : type_(x.type_), u_() { + switch (type_) { +#define INIT(p, v) \ + case p##type: \ + u_.p = v; \ + break + INIT(string_, new std::string(*x.u_.string_)); + INIT(array_, new array(*x.u_.array_)); + INIT(object_, new object(*x.u_.object_)); +#undef INIT + default: + u_ = x.u_; + break; + } +} + +inline value& value::operator=(const value& x) { + if (this != &x) { + value t(x); + swap(t); + } + return *this; +} + +#if PICOJSON_USE_RVALUE_REFERENCE +inline value::value(value&& x) PICOJSON_NOEXCEPT : type_(null_type), u_() { swap(x); } +inline value& value::operator=(value&& x) PICOJSON_NOEXCEPT { + swap(x); + return *this; +} +#endif +inline void value::swap(value& x) PICOJSON_NOEXCEPT { + std::swap(type_, x.type_); + std::swap(u_, x.u_); +} + +#define IS(ctype, jtype) \ + template <> \ + inline bool value::is() const { \ + return type_ == jtype##_type; \ + } +IS(null, null) +IS(bool, boolean) +#ifdef PICOJSON_USE_INT64 +IS(int64_t, int64) +#endif +IS(std::string, string) +IS(array, array) +IS(object, object) +#undef IS +template <> +inline bool value::is() const { + return type_ == number_type +#ifdef PICOJSON_USE_INT64 + || type_ == int64_type +#endif + // NOLINTNEXTLINE(whitespace/semicolon) + ; +} + +#define GET(ctype, var) \ + template <> \ + inline const ctype& value::get() const { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } \ + template <> \ + inline ctype& value::get() { \ + PICOJSON_ASSERT("type mismatch! call is() before get()" && is()); \ + return var; \ + } +GET(bool, u_.boolean_) +GET(std::string, *u_.string_) +GET(array, *u_.array_) +GET(object, *u_.object_) +#ifdef PICOJSON_USE_INT64 +GET(double, + (type_ == int64_type && (const_cast(this)->type_ = number_type, + (const_cast(this)->u_.number_ = u_.int64_)), + u_.number_)) +GET(int64_t, u_.int64_) +#else +GET(double, u_.number_) +#endif +#undef GET + +#define SET(ctype, jtype, setter) \ + template <> \ + inline void value::set(const ctype& _val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +SET(bool, boolean, u_.boolean_ = _val;) +SET(std::string, string, u_.string_ = new std::string(_val);) +SET(array, array, u_.array_ = new array(_val);) +SET(object, object, u_.object_ = new object(_val);) +SET(double, number, u_.number_ = _val;) +#ifdef PICOJSON_USE_INT64 +SET(int64_t, int64, u_.int64_ = _val;) +#endif +#undef SET + +#if PICOJSON_USE_RVALUE_REFERENCE +#define MOVESET(ctype, jtype, setter) \ + template <> \ + inline void value::set(ctype && _val) { \ + clear(); \ + type_ = jtype##_type; \ + setter \ + } +MOVESET(std::string, string, u_.string_ = new std::string(std::move(_val));) +MOVESET(array, array, u_.array_ = new array(std::move(_val));) +MOVESET(object, object, u_.object_ = new object(std::move(_val));) +#undef MOVESET +#endif + +inline bool value::evaluate_as_boolean() const { + switch (type_) { + case null_type: + return false; + case boolean_type: + return u_.boolean_; + case number_type: + return u_.number_ != 0; +#ifdef PICOJSON_USE_INT64 + case int64_type: + return u_.int64_ != 0; +#endif + case string_type: + return !u_.string_->empty(); + default: + return true; + } +} + +inline const value& value::get(const size_t idx) const { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline value& value::get(const size_t idx) { + static value s_null; + PICOJSON_ASSERT(is()); + return idx < u_.array_->size() ? (*u_.array_)[idx] : s_null; +} + +inline const value& value::get(const std::string& key) const { + static value s_null; + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline value& value::get(const std::string& key) { + static value s_null; + PICOJSON_ASSERT(is()); + object::iterator i = u_.object_->find(key); + return i != u_.object_->end() ? i->second : s_null; +} + +inline bool value::contains(const size_t idx) const { + PICOJSON_ASSERT(is()); + return idx < u_.array_->size(); +} + +inline bool value::contains(const std::string& key) const { + PICOJSON_ASSERT(is()); + object::const_iterator i = u_.object_->find(key); + return i != u_.object_->end(); +} + +inline std::string value::to_str() const { + switch (type_) { + case null_type: + return "null"; + case boolean_type: + return u_.boolean_ ? "true" : "false"; +#ifdef PICOJSON_USE_INT64 + case int64_type: { + char buf[sizeof("-9223372036854775808")]; + SNPRINTF(buf, sizeof(buf), "%" PRId64, u_.int64_); + return buf; + } +#endif + case number_type: { + char buf[256]; + double tmp; + SNPRINTF( + buf, + sizeof(buf), + fabs(u_.number_) < (1ULL << 53) && modf(u_.number_, &tmp) == 0 ? "%.f" : "%.17g", + u_.number_ + ); +#if PICOJSON_USE_LOCALE + char* decimal_point = localeconv()->decimal_point; + if (strcmp(decimal_point, ".") != 0) { + size_t decimal_point_len = strlen(decimal_point); + for (char* p = buf; *p != '\0'; ++p) { + if (strncmp(p, decimal_point, decimal_point_len) == 0) { + return std::string(buf, p) + "." + (p + decimal_point_len); + } + } + } +#endif + return buf; + } + case string_type: + return *u_.string_; + case array_type: + return "array"; + case object_type: + return "object"; + default: + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + } + return std::string(); +} + +template +void copy(const std::string& s, Iter oi) { + std::copy(s.begin(), s.end(), oi); +} + +template +struct serialize_str_char { + Iter oi; + void operator()(char c) { + switch (c) { +#define MAP(val, sym) \ + case val: \ + copy(sym, oi); \ + break + MAP('"', "\\\""); + MAP('\\', "\\\\"); + MAP('\b', "\\b"); + MAP('\f', "\\f"); + MAP('\n', "\\n"); + MAP('\r', "\\r"); + MAP('\t', "\\t"); +#undef MAP + default: + if (static_cast(c) < 0x20 || c == 0x7f) { + char buf[7]; + SNPRINTF(buf, sizeof(buf), "\\u%04x", c & 0xff); + copy(buf, buf + 6, oi); + } else { + *oi++ = c; + } + break; + } + } +}; + +template +void serialize_str(const std::string& s, Iter oi) { + *oi++ = '"'; + serialize_str_char process_char = {oi}; + std::for_each(s.begin(), s.end(), process_char); + *oi++ = '"'; +} + +template +void value::serialize(Iter oi, bool prettify) const { + return _serialize(oi, prettify ? 0 : -1); +} + +inline std::string value::serialize(bool prettify) const { return _serialize(prettify ? 0 : -1); } + +template +void value::_indent(Iter oi, int indent) { + *oi++ = '\n'; + for (int i = 0; i < indent * INDENT_WIDTH; ++i) { + *oi++ = ' '; + } +} + +template +void value::_serialize(Iter oi, int indent) const { + switch (type_) { + case string_type: + serialize_str(*u_.string_, oi); + break; + case array_type: { + *oi++ = '['; + if (indent != -1) { + ++indent; + } + for (array::const_iterator i = u_.array_->begin(); i != u_.array_->end(); ++i) { + if (i != u_.array_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + i->_serialize(oi, indent); + } + if (indent != -1) { + --indent; + if (!u_.array_->empty()) { + _indent(oi, indent); + } + } + *oi++ = ']'; + break; + } + case object_type: { + *oi++ = '{'; + if (indent != -1) { + ++indent; + } + +#if PICOJSON_USE_ORDERED_OBJECT + for (auto i = u_.object_->ordered_keys().begin(); i != u_.object_->ordered_keys().end(); + ++i) { + if (i != u_.object_->ordered_keys().begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + serialize_str(*i, oi); + *oi++ = ':'; + if (indent != -1) { + *oi++ = ' '; + } + u_.object_->at(*i)._serialize(oi, indent); + } +#else + for (object::const_iterator i = u_.object_->begin(); i != u_.object_->end(); ++i) { + if (i != u_.object_->begin()) { + *oi++ = ','; + } + if (indent != -1) { + _indent(oi, indent); + } + serialize_str(i->first, oi); + *oi++ = ':'; + if (indent != -1) { + *oi++ = ' '; + } + i->second._serialize(oi, indent); + } +#endif + if (indent != -1) { + --indent; + if (!u_.object_->empty()) { + _indent(oi, indent); + } + } + *oi++ = '}'; + break; + } + default: + copy(to_str(), oi); + break; + } + if (indent == 0) { + *oi++ = '\n'; + } +} + +inline std::string value::_serialize(int indent) const { + std::string s; + _serialize(std::back_inserter(s), indent); + return s; +} + +template +class input { + protected: + Iter cur_, end_; + bool consumed_; + int line_; + + public: + input(const Iter& first, const Iter& last) + : cur_(first), end_(last), consumed_(false), line_(1) {} + int getc() { + if (consumed_) { + if (*cur_ == '\n') { + ++line_; + } + ++cur_; + } + if (cur_ == end_) { + consumed_ = false; + return -1; + } + consumed_ = true; + return *cur_ & 0xff; + } + void ungetc() { consumed_ = false; } + Iter cur() const { + if (consumed_) { + input* self = const_cast*>(this); + self->consumed_ = false; + ++self->cur_; + } + return cur_; + } + int line() const { return line_; } + void skip_ws() { + while (1) { + int ch = getc(); + if (!(ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r')) { + ungetc(); + break; + } + } + } + bool expect(const int expected) { + skip_ws(); + if (getc() != expected) { + ungetc(); + return false; + } + return true; + } + bool match(const std::string& pattern) { + for (std::string::const_iterator pi(pattern.begin()); pi != pattern.end(); ++pi) { + if (getc() != *pi) { + ungetc(); + return false; + } + } + return true; + } +}; + +template +// NOLINTNEXTLINE(runtime/references) +inline int _parse_quadhex(input& in) { + int uni_ch = 0, hex; + for (int i = 0; i < 4; i++) { + if ((hex = in.getc()) == -1) { + return -1; + } + if ('0' <= hex && hex <= '9') { + hex -= '0'; + } else if ('A' <= hex && hex <= 'F') { + hex -= 'A' - 0xa; + } else if ('a' <= hex && hex <= 'f') { + hex -= 'a' - 0xa; + } else { + in.ungetc(); + return -1; + } + uni_ch = uni_ch * 16 + hex; + } + return uni_ch; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_codepoint(String& out, input& in) { + int uni_ch; + if ((uni_ch = _parse_quadhex(in)) == -1) { + return false; + } + if (0xd800 <= uni_ch && uni_ch <= 0xdfff) { + if (0xdc00 <= uni_ch) { + // a second 16-bit of a surrogate pair appeared + return false; + } + // first 16-bit of surrogate pair, get the next one + if (in.getc() != '\\' || in.getc() != 'u') { + in.ungetc(); + return false; + } + int second = _parse_quadhex(in); + if (!(0xdc00 <= second && second <= 0xdfff)) { + return false; + } + uni_ch = ((uni_ch - 0xd800) << 10) | ((second - 0xdc00) & 0x3ff); + uni_ch += 0x10000; + } + if (uni_ch < 0x80) { + out.push_back(static_cast(uni_ch)); + } else { + if (uni_ch < 0x800) { + out.push_back(static_cast(0xc0 | (uni_ch >> 6))); + } else { + if (uni_ch < 0x10000) { + out.push_back(static_cast(0xe0 | (uni_ch >> 12))); + } else { + out.push_back(static_cast(0xf0 | (uni_ch >> 18))); + out.push_back(static_cast(0x80 | ((uni_ch >> 12) & 0x3f))); + } + out.push_back(static_cast(0x80 | ((uni_ch >> 6) & 0x3f))); + } + out.push_back(static_cast(0x80 | (uni_ch & 0x3f))); + } + return true; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_string(String& out, input& in) { + while (1) { + int ch = in.getc(); + if (ch < ' ') { + in.ungetc(); + return false; + } else if (ch == '"') { + return true; + } else if (ch == '\\') { + if ((ch = in.getc()) == -1) { + return false; + } + switch (ch) { +#define MAP(sym, val) \ + case sym: \ + out.push_back(val); \ + break + MAP('"', '\"'); + MAP('\\', '\\'); + MAP('/', '/'); + MAP('b', '\b'); + MAP('f', '\f'); + MAP('n', '\n'); + MAP('r', '\r'); + MAP('t', '\t'); +#undef MAP + case 'u': + if (!_parse_codepoint(out, in)) { + return false; + } + break; + default: + return false; + } + } else { + out.push_back(static_cast(ch)); + } + } + return false; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_array(Context& ctx, input& in) { + if (!ctx.parse_array_start()) { + return false; + } + size_t idx = 0; + if (in.expect(']')) { + return ctx.parse_array_stop(idx); + } + do { + if (!ctx.parse_array_item(in, idx)) { + return false; + } + idx++; + } while (in.expect(',')); + return in.expect(']') && ctx.parse_array_stop(idx); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse_object(Context& ctx, input& in) { + if (!ctx.parse_object_start()) { + return false; + } + if (in.expect('}')) { + return true; + } + do { + std::string key; + if (!in.expect('"') || !_parse_string(key, in) || !in.expect(':')) { + return false; + } + if (!ctx.parse_object_item(in, key)) { + return false; + } + } while (in.expect(',')); + return in.expect('}'); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline std::string _parse_number(input& in) { + std::string num_str; + while (1) { + int ch = in.getc(); + if (('0' <= ch && ch <= '9') || ch == '+' || ch == '-' || ch == 'e' || ch == 'E') { + num_str.push_back(static_cast(ch)); + } else if (ch == '.') { +#if PICOJSON_USE_LOCALE + num_str += localeconv()->decimal_point; +#else + num_str.push_back('.'); +#endif + } else { + in.ungetc(); + break; + } + } + return num_str; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline bool _parse(Context& ctx, input& in) { + in.skip_ws(); + int ch = in.getc(); + switch (ch) { +#define IS(ch, text, op) \ + case ch: \ + if (in.match(text) && op) { \ + return true; \ + } else { \ + return false; \ + } + IS('n', "ull", ctx.set_null()); + IS('f', "alse", ctx.set_bool(false)); + IS('t', "rue", ctx.set_bool(true)); +#undef IS + case '"': + return ctx.parse_string(in); + case '[': + return _parse_array(ctx, in); + case '{': + return _parse_object(ctx, in); + default: + if (('0' <= ch && ch <= '9') || ch == '-') { + double f; + char* endp; + in.ungetc(); + std::string num_str(_parse_number(in)); + if (num_str.empty()) { + return false; + } +#ifdef PICOJSON_USE_INT64 + { + errno = 0; + intmax_t ival = strtoimax(num_str.c_str(), &endp, 10); + if (errno == 0 && std::numeric_limits::min() <= ival && + ival <= std::numeric_limits::max() && + endp == num_str.c_str() + num_str.size()) { + ctx.set_int64(ival); + return true; + } + } +#endif + f = strtod(num_str.c_str(), &endp); + if (endp == num_str.c_str() + num_str.size()) { + ctx.set_number(f); + return true; + } + return false; + } + break; + } + in.ungetc(); + return false; +} + +class deny_parse_context { + public: + bool set_null() { return false; } + bool set_bool(bool) { return false; } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { return false; } +#endif + bool set_number(double) { return false; } + template + bool parse_string(input&) { + return false; + } + bool parse_array_start() { return false; } + template + bool parse_array_item(input&, size_t) { + return false; + } + bool parse_array_stop(size_t) { return false; } + bool parse_object_start() { return false; } + template + bool parse_object_item(input&, const std::string&) { + return false; + } +}; + +class default_parse_context { + protected: + value* out_; + + public: + // NOLINTNEXTLINE(runtime/explicit) + default_parse_context(value* out) : out_(out) {} + bool set_null() { + *out_ = value(); + return true; + } + bool set_bool(bool b) { + *out_ = value(b); + return true; + } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t i) { + *out_ = value(i); + return true; + } +#endif + bool set_number(double f) { + *out_ = value(f); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_string(input& in) { + *out_ = value(string_type, false); + return _parse_string(out_->get(), in); + } + bool parse_array_start() { + *out_ = value(array_type, false); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_array_item(input& in, size_t) { + array& a = out_->get(); + a.push_back(value()); + default_parse_context ctx(&a.back()); + return _parse(ctx, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { + *out_ = value(object_type, false); + return true; + } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_object_item(input& in, const std::string& key) { + object& o = out_->get(); + default_parse_context ctx(&o[key]); + return _parse(ctx, in); + } + + private: + default_parse_context(const default_parse_context&); + default_parse_context& operator=(const default_parse_context&); +}; + +class null_parse_context { + public: + struct dummy_str { + void push_back(int) {} + }; + + public: + null_parse_context() {} + bool set_null() { return true; } + bool set_bool(bool) { return true; } +#ifdef PICOJSON_USE_INT64 + bool set_int64(int64_t) { return true; } +#endif + bool set_number(double) { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_string(input& in) { + dummy_str s; + return _parse_string(s, in); + } + bool parse_array_start() { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_array_item(input& in, size_t) { + return _parse(*this, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { return true; } + template + // NOLINTNEXTLINE(runtime/references) + bool parse_object_item(input& in, const std::string&) { + return _parse(*this, in); + } + + private: + null_parse_context(const null_parse_context&); + null_parse_context& operator=(const null_parse_context&); +}; + +// obsolete, use the version below +template +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, Iter& pos, const Iter& last) { + std::string err; + pos = parse(out, pos, last, &err); + return err; +} + +template +// NOLINTNEXTLINE(runtime/references) +inline Iter _parse(Context& ctx, const Iter& first, const Iter& last, std::string* err) { + input in(first, last); + if (!_parse(ctx, in) && err != NULL) { + char buf[64]; + SNPRINTF(buf, sizeof(buf), "syntax error at line %d near: ", in.line()); + *err = buf; + while (1) { + int ch = in.getc(); + if (ch == -1 || ch == '\n') { + break; + } else if (ch >= ' ') { + err->push_back(static_cast(ch)); + } + } + } + return in.cur(); +} + +template +// NOLINTNEXTLINE(runtime/references) +inline Iter parse(value& out, const Iter& first, const Iter& last, std::string* err) { + default_parse_context ctx(&out); + return _parse(ctx, first, last, err); +} + +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, const std::string& s) { + std::string err; + parse(out, s.begin(), s.end(), &err); + return err; +} + +// NOLINTNEXTLINE(runtime/references) +inline std::string parse(value& out, std::istream& is) { + std::string err; + parse(out, std::istreambuf_iterator(is.rdbuf()), std::istreambuf_iterator(), &err); + return err; +} + +template +struct last_error_t { + static std::string s; +}; +template +// NOLINTNEXTLINE(runtime/string) +std::string last_error_t::s; + +inline void set_last_error(const std::string& s) { last_error_t::s = s; } + +inline const std::string& get_last_error() { return last_error_t::s; } + +inline bool operator==(const value& x, const value& y) { + if (x.is()) return y.is(); +#define PICOJSON_CMP(type) \ + if (x.is()) return y.is() && x.get() == y.get() + PICOJSON_CMP(bool); + PICOJSON_CMP(double); + PICOJSON_CMP(std::string); + PICOJSON_CMP(array); + PICOJSON_CMP(object); +#undef PICOJSON_CMP + PICOJSON_ASSERT(0); +#ifdef _MSC_VER + __assume(0); +#endif + return false; +} + +inline bool operator!=(const value& x, const value& y) { return !(x == y); } +} // namespace picojson + +#if !PICOJSON_USE_RVALUE_REFERENCE +namespace std { +template <> +inline void swap(picojson::value& x, picojson::value& y) { + x.swap(y); +} +} // namespace std +#endif + +inline std::istream& operator>>(std::istream& is, picojson::value& x) { + picojson::set_last_error(std::string()); + const std::string err(picojson::parse(x, is)); + if (!err.empty()) { + picojson::set_last_error(err); + is.setstate(std::ios::failbit); + } + return is; +} + +inline std::ostream& operator<<(std::ostream& os, const picojson::value& x) { + x.serialize(std::ostream_iterator(os)); + return os; +} +#ifdef _MSC_VER +#pragma warning(pop) +#endif diff --git a/zig/vendor/xgrammar/LICENSE b/zig/vendor/xgrammar/LICENSE new file mode 100644 index 0000000000..261eeb9e9f --- /dev/null +++ b/zig/vendor/xgrammar/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/zig/vendor/xgrammar/NOTICE b/zig/vendor/xgrammar/NOTICE new file mode 100644 index 0000000000..2d56df4754 --- /dev/null +++ b/zig/vendor/xgrammar/NOTICE @@ -0,0 +1,3 @@ +XGrammar + +Copyright (c) 2024 by XGrammar Contributors diff --git a/zig/vendor/xgrammar/README.md b/zig/vendor/xgrammar/README.md new file mode 100644 index 0000000000..045891dd0c --- /dev/null +++ b/zig/vendor/xgrammar/README.md @@ -0,0 +1,14 @@ +# Vendored: XGrammar 0.2.8 + +Upstream: https://github.com/mlc-ai/xgrammar, tag `v0.2.8`, commit `97787376faee5ed8466cfad57c99855e4ce2f6aa`. + +This folder holds the C++ core from the 0.2.8 source distribution: `cpp/`, `include/` and `3rdparty/` (DLPack's header and +picojson). The build files, tests and Python bindings are left out. **Local patches: none.** Every file is byte-identical +to the upstream release, so an update means replacing the folder with the next release's same subset. + +It is upstream code we do not maintain, so `tools/zig/lean_check.py` exempts exactly this prefix (`VENDORED`). Our own +code around it (`zig/src/core/grammar/`, `zig/build/grammar.zig`) is checked like any other source. + +Licence: Apache License 2.0, Copyright (c) 2024 by XGrammar Contributors (`LICENSE` and `NOTICE` here). DLPack: Apache +License 2.0 (`3rdparty/dlpack/LICENSE`). picojson: BSD-2-Clause (text in `3rdparty/picojson/picojson.h`). See also the +repository's `THIRD_PARTY_NOTICES.md` and `NOTICE`. diff --git a/zig/vendor/xgrammar/cpp/compiled_grammar.cc b/zig/vendor/xgrammar/cpp/compiled_grammar.cc new file mode 100644 index 0000000000..466a0235a9 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/compiled_grammar.cc @@ -0,0 +1,279 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/compiled_grammar.cc + */ + +#include + +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "support/json_parse.h" +#include "support/json_serializer.h" +#include "testing.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* AdaptiveTokenMask *******************/ + +AdaptiveTokenMask::AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& rejected_indices, + const std::vector& uncertain_indices +) { + auto size_acc = accepted_indices.size(); + auto size_rej = rejected_indices.size(); + + store_type = size_acc >= USE_BITSET_THRESHOLD && size_rej >= USE_BITSET_THRESHOLD + ? StoreType::kAcceptedBitset + : size_acc < size_rej ? StoreType::kAccepted + : StoreType::kRejected; + + if (store_type == StoreType::kAcceptedBitset) { + accepted_bitset = DynamicBitset(vocab_size); + for (auto idx : accepted_indices) { + accepted_bitset.Set(sorted_decoded_vocab[idx].first, true); + } + } else if (store_type == StoreType::kAccepted) { + this->accepted_indices = accepted_indices; + } else { + this->rejected_indices = rejected_indices; + } + + this->uncertain_indices = uncertain_indices; +} + +AdaptiveTokenMask::AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& uncertain_indices +) { + auto size_acc = accepted_indices.size(); + + store_type = size_acc >= USE_BITSET_THRESHOLD ? StoreType::kAcceptedBitset : StoreType::kAccepted; + + if (store_type == StoreType::kAcceptedBitset) { + accepted_bitset = DynamicBitset(vocab_size); + for (auto idx : accepted_indices) { + accepted_bitset.Set(sorted_decoded_vocab[idx].first, true); + } + } else { + XGRAMMAR_DCHECK(store_type == StoreType::kAccepted); + this->accepted_indices = accepted_indices; + } + this->uncertain_indices = uncertain_indices; +} + +std::string AdaptiveTokenMask::Print(const TokenizerInfo& tokenizer_info) const { + constexpr int kMaxPrintTokens = 100; + std::stringstream ss; + const auto& sorted_decoded_vocab = tokenizer_info.GetSortedDecodedVocab(); + std::vector accepted_indices; + std::vector rejected_indices; + std::unordered_set uncertain_indices_set( + uncertain_indices.begin(), uncertain_indices.end() + ); + + accepted_indices.reserve(sorted_decoded_vocab.size()); + rejected_indices.reserve(sorted_decoded_vocab.size()); + + if (store_type == StoreType::kAcceptedBitset) { + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + if (uncertain_indices_set.count(i)) { + continue; + } + if (accepted_bitset[sorted_decoded_vocab[i].first]) { + accepted_indices.push_back(i); + } else { + rejected_indices.push_back(i); + } + } + } else if (store_type == StoreType::kAccepted) { + accepted_indices = this->accepted_indices; + // Reject indices = [0, sorted_decoded_vocab.size()) \ accepted_indices \ uncertain_indices + int acc_ptr = 0; + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + while (acc_ptr < static_cast(accepted_indices.size()) && accepted_indices[acc_ptr] < i) { + ++acc_ptr; + } + if (acc_ptr < static_cast(accepted_indices.size()) && accepted_indices[acc_ptr] == i) { + continue; + } + if (uncertain_indices_set.count(i)) { + continue; + } + rejected_indices.push_back(i); + } + } else { + XGRAMMAR_DCHECK(store_type == StoreType::kRejected); + rejected_indices = this->rejected_indices; + // Accepted indices = [0, sorted_decoded_vocab.size()) \ rejected_indices \ uncertain_indices + int rej_ptr = 0; + for (int i = 0; i < static_cast(sorted_decoded_vocab.size()); ++i) { + while (rej_ptr < static_cast(rejected_indices.size()) && rejected_indices[rej_ptr] < i) { + ++rej_ptr; + } + if (rej_ptr < static_cast(rejected_indices.size()) && rejected_indices[rej_ptr] == i) { + continue; + } + if (uncertain_indices_set.count(i)) { + continue; + } + accepted_indices.push_back(i); + } + } + + std::string storage_type_str = store_type == StoreType::kAcceptedBitset ? "AcceptedBitset" + : store_type == StoreType::kAccepted ? "Accepted" + : "Rejected"; + + ss << "AdaptiveTokenMask(num_tokens=" << sorted_decoded_vocab.size() + << ", accepted_num=" << accepted_indices.size() << ", rejected_num=" << rejected_indices.size() + << ", uncertain_num=" << uncertain_indices.size() << ", storage_type=" << storage_type_str + << ",\n"; + + // Convert indices to token ids for printing + std::vector accepted_token_ids; + std::vector rejected_token_ids; + std::vector uncertain_token_ids; + accepted_token_ids.reserve(accepted_indices.size()); + rejected_token_ids.reserve(rejected_indices.size()); + uncertain_token_ids.reserve(uncertain_indices.size()); + + for (auto idx : accepted_indices) { + accepted_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(accepted_token_ids.begin(), accepted_token_ids.end()); + for (auto idx : rejected_indices) { + rejected_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(rejected_token_ids.begin(), rejected_token_ids.end()); + for (auto idx : uncertain_indices) { + uncertain_token_ids.push_back(sorted_decoded_vocab[idx].first); + } + std::sort(uncertain_token_ids.begin(), uncertain_token_ids.end()); + + ss << "accepted=" << PrintTokenByIds(accepted_token_ids, tokenizer_info, kMaxPrintTokens) + << ",\nrejected=" << PrintTokenByIds(rejected_token_ids, tokenizer_info, kMaxPrintTokens) + << ",\nuncertain=" << PrintTokenByIds(uncertain_token_ids, tokenizer_info, kMaxPrintTokens) + << "\n)"; + return ss.str(); +} + +/************** CompiledGrammar::Impl **************/ + +picojson::value SerializeJSONValue(const CompiledGrammar::Impl& impl) { + auto result = picojson::object{}; + result["grammar"] = AutoSerializeJSONValue(impl.grammar); + result["tokenizer_metadata"] = impl.tokenizer_info->DumpMetadataValue(); + result["adaptive_token_mask_cache"] = AutoSerializeJSONValue(impl.adaptive_token_mask_cache); + return picojson::value(result); +} + +std::optional DeserializeJSONValue( + CompiledGrammar::Impl* impl, + const picojson::value& json_value, + const TokenizerInfo& tokenizer_info +) { + const auto& type_name = "CompiledGrammar"; + if (!json_value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& object = json_value.get(); + if (object.find("grammar") == object.end()) { + return ConstructDeserializeError("Expect a 'grammar' field", type_name); + } + if (auto error = AutoDeserializeJSONValue(&(impl->grammar), object["grammar"], type_name)) { + return error; + } + if (impl->grammar.IsNull()) { + return ConstructDeserializeError("Expect a non-null grammar", type_name); + } + if (object.find("tokenizer_metadata") == object.end()) { + return ConstructDeserializeError("Expect a 'tokenizer_metadata' field", type_name); + } + const auto& tokenizer_metadata = object["tokenizer_metadata"]; + if (auto error = tokenizer_info->CheckMetadataMatch(tokenizer_metadata)) { + return ConstructDeserializeError( + std::string("Tokenizer metadata mismatch: ") + error->what(), type_name + ); + } + impl->tokenizer_info = tokenizer_info; + if (object.find("adaptive_token_mask_cache") == object.end()) { + return ConstructDeserializeError("Expect a 'adaptive_token_mask_cache' field", type_name); + } + if (auto error = AutoDeserializeJSONValue( + &(impl->adaptive_token_mask_cache), object["adaptive_token_mask_cache"], type_name + )) { + return error; + } + // The masks index sorted_decoded_vocab and are OR-ed into a vocab_size-bit bitset, so their + // contents must match the tokenizer they are deserialized with. + const int64_t num_sorted_tokens = tokenizer_info.GetSortedDecodedVocab().size(); + auto indices_ok = [&](const std::vector& indices) { + return std::all_of(indices.begin(), indices.end(), [&](int32_t index) { + return index >= 0 && index < num_sorted_tokens; + }); + }; + for (const auto& [state, mask] : impl->adaptive_token_mask_cache) { + using StoreType = AdaptiveTokenMask::StoreType; + const bool store_type_ok = mask.store_type == StoreType::kAccepted || + mask.store_type == StoreType::kRejected || + mask.store_type == StoreType::kAcceptedBitset; + const bool bitset_ok = mask.store_type != StoreType::kAcceptedBitset || + mask.accepted_bitset.Size() == tokenizer_info.GetVocabSize(); + if (!store_type_ok || !bitset_ok || !indices_ok(mask.accepted_indices) || + !indices_ok(mask.rejected_indices) || !indices_ok(mask.uncertain_indices)) { + return ConstructDeserializeError( + "adaptive_token_mask_cache contains a mask that does not match the tokenizer", type_name + ); + } + } + return std::nullopt; +} + +/************** CompiledGrammar **************/ + +std::size_t MemorySize(const CompiledGrammar::Impl& impl) { + return MemorySize(impl.grammar) + MemorySize(impl.adaptive_token_mask_cache); +} + +std::size_t CompiledGrammar::MemorySizeBytes() const { return MemorySize(*pimpl_); } + +Grammar CompiledGrammar::GetGrammar() const { return pimpl_->GetGrammar(); } + +TokenizerInfo CompiledGrammar::GetTokenizerInfo() const { return pimpl_->GetTokenizerInfo(); } + +/*! \brief Return the serialized JSON string of the compiled grammar. */ +std::string CompiledGrammar::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +/*! \brief Deserialize a compiled grammar from a JSON string and tokenizer info. */ +std::variant CompiledGrammar::DeserializeJSON( + const std::string& json_string, const TokenizerInfo& tokenizer_info +) { + picojson::value json_value; + if (auto error = ParseJSON(json_value, json_string); !error.empty()) { + return InvalidJSONError("Failed to parse JSON: " + error); + } + if (!json_value.is()) { + return DeserializeFormatError("Expect an object"); + } + const auto& object = json_value.get(); + if (auto error = SerializeVersion::Check(object)) { + return error.value(); + } + auto impl = std::make_shared(); + if (auto error = DeserializeJSONValue(impl.get(), json_value, tokenizer_info)) { + return error.value(); + } + return CompiledGrammar(std::move(impl)); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/compiled_grammar_impl.h b/zig/vendor/xgrammar/cpp/compiled_grammar_impl.h new file mode 100644 index 0000000000..413c946b0d --- /dev/null +++ b/zig/vendor/xgrammar/cpp/compiled_grammar_impl.h @@ -0,0 +1,148 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiled_grammar_impl.h + * \brief The header for the data structures of the compiled grammar. + */ +#ifndef XGRAMMAR_COMPILED_GRAMMAR_IMPL_H_ +#define XGRAMMAR_COMPILED_GRAMMAR_IMPL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "earley_parser.h" +#include "support/dynamic_bitset.h" +#include "support/reflection.h" +#include "xgrammar/compiler.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* CompiledGrammar Datastructures *******************/ + +/*! + * \brief Preprocessed information, for a given specific ParserState, divides the token set + * into three categories: accepted, rejected, and uncertain. + * Accepted: tokens that can be determined by the current ParserState to be acceptable + * Rejected: tokens that can be determined by the current ParserState to be unacceptable + * Uncertain: tokens that need the state of the parent ParserStates to determine if acceptable + * + * \note uncertain indices are stored directly. Accepted / rejected indices have three ways to + * store to reduce memory and computation usage. See StoreType. + * \note These indices are the indices of sorted_decoded_vocab in the CompiledGrammar + * object, instead of the token ids. That helps the matching process. + */ +struct AdaptiveTokenMask { + enum class StoreType { + // Only store all accepted token indices. Then rejected indices = all_indices - accepted_indices + // - uncertain_indices. This is useful when |accepted_indices| < |rejected_indices|. + kAccepted = 0, + // Only store all rejected token indices. Then accepted indices = all_indices - rejected_indices + // - uncertain_indices. This is useful when |accepted_indices| > |rejected_indices|. + kRejected = 1, + // Store all accepted token indices in a bitset. This is useful when both |accepted_indices| and + // |rejected_indices| are large. + kAcceptedBitset = 2 + }; + StoreType store_type; + + static constexpr int USE_BITSET_THRESHOLD = 1000; + + std::vector accepted_indices; + std::vector rejected_indices; + DynamicBitset accepted_bitset; + + std::vector uncertain_indices; + + /*! \brief Default constructor. Only for deserialization. */ + AdaptiveTokenMask() = default; + + AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& rejected_indices, + const std::vector& uncertain_indices + ); + + AdaptiveTokenMask( + size_t vocab_size, + const std::vector>& sorted_decoded_vocab, + const std::vector& accepted_indices, + const std::vector& uncertain_indices + ); + + std::string Print(const TokenizerInfo& tokenizer_info) const; + + friend std::size_t MemorySize(const AdaptiveTokenMask& mask) { + return MemorySize(mask.uncertain_indices) + MemorySize(mask.accepted_indices) + + MemorySize(mask.rejected_indices) + MemorySize(mask.accepted_bitset); + } +}; + +XGRAMMAR_MEMBER_TABLE( + AdaptiveTokenMask, + "store_type", + &AdaptiveTokenMask::store_type, + "accepted_indices", + &AdaptiveTokenMask::accepted_indices, + "rejected_indices", + &AdaptiveTokenMask::rejected_indices, + "accepted_bitset", + &AdaptiveTokenMask::accepted_bitset, + "uncertain_indices", + &AdaptiveTokenMask::uncertain_indices +); + +/*! + * \brief All information that we need to match tokens in the tokenizer to the specified grammar. + * It is the result of preprocessing. + * \sa xgrammar::GrammarMatcher + */ +class CompiledGrammar::Impl { + public: + /*! \brief The grammar for the GrammarMatcher. */ + Grammar grammar{NullObj{}}; + + /*! \brief The tokenizer information. */ + TokenizerInfo tokenizer_info{NullObj{}}; + + /*! \brief Default constructor. */ + Impl() = default; + + /*! \brief Mapping from the parser state to the adaptive token mask. */ + std::unordered_map + adaptive_token_mask_cache; + + Grammar GetGrammar() const { return grammar; } + + TokenizerInfo GetTokenizerInfo() const { return tokenizer_info; } + + friend struct member_trait; + friend picojson::value SerializeJSONValue(const Impl& impl); + friend std::optional DeserializeJSONValue( + CompiledGrammar::Impl* impl, + const picojson::value& json_value, + const TokenizerInfo& tokenizer_info + ); + friend std::size_t MemorySize(const Impl& impl); +}; + +XGRAMMAR_MEMBER_TABLE( + CompiledGrammar::Impl, + "grammar", + &CompiledGrammar::Impl::grammar, + "tokenizer_info", + &CompiledGrammar::Impl::tokenizer_info, + "adaptive_token_mask_cache", + &CompiledGrammar::Impl::adaptive_token_mask_cache +); + +} // namespace xgrammar + +#endif // XGRAMMAR_COMPILED_GRAMMAR_IMPL_H_ diff --git a/zig/vendor/xgrammar/cpp/config.cc b/zig/vendor/xgrammar/cpp/config.cc new file mode 100644 index 0000000000..b8778d2572 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/config.cc @@ -0,0 +1,21 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/config.cc + */ + +#include + +#include "support/json_serializer.h" +#include "support/recursion_guard.h" + +namespace xgrammar { + +void SetMaxRecursionDepth(int max_recursion_depth) { + RecursionGuard::SetMaxRecursionDepth(max_recursion_depth); +} + +int GetMaxRecursionDepth() { return RecursionGuard::GetMaxRecursionDepth(); } + +std::string GetSerializationVersion() { return std::string(SerializeVersion::GetVersion()); } + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/cohere.cc b/zig/vendor/xgrammar/cpp/converter_ext/cohere.cc new file mode 100644 index 0000000000..798f827e7d --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/cohere.cc @@ -0,0 +1,653 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/cohere.cc + * \brief Implementation of the Cohere XML Tool Calling converter. + */ +#include + +#include +#include +#include +#include +#include +#include + +#include "../json_schema_converter_ext.h" +#include "../support/encoding.h" +#include "../support/json_parse.h" +#include "../support/logging.h" + +namespace xgrammar { + +namespace converter_ext { + +XMLWrapper GetCohereXMLWrapper() { return {"", "", ""}; } + +} // namespace converter_ext + +namespace { + +constexpr std::array, 4> kCohereKeyEntities = { + std::pair{'&', "&"}, + std::pair{'<', "<"}, + std::pair{'>', ">"}, + std::pair{'"', """}, +}; + +std::string SerializeCohereKeyCodepoint(TCodepoint codepoint) { + for (const auto& [entity_codepoint, entity] : kCohereKeyEntities) { + if (codepoint == entity_codepoint) { + return entity; + } + } + return CharToUTF8(codepoint); +} + +std::vector ParseCohereKeyCodepoints(const std::string& key) { + XGRAMMAR_CHECK(key.find('\0') == std::string::npos) << "Cohere property names cannot contain NUL"; + auto codepoints = ParseUTF8(key.c_str()); + XGRAMMAR_CHECK(codepoints.size() != 1 || codepoints[0] != CharHandlingError::kInvalidUTF8) + << "Cohere property names must be valid UTF-8"; + return codepoints; +} + +std::string SerializeCohereKey(const std::string& key) { + std::string serialized; + for (TCodepoint codepoint : ParseCohereKeyCodepoints(key)) { + serialized += SerializeCohereKeyCodepoint(codepoint); + } + return serialized; +} + +template +std::vector CohereOrdinaryKeyRangesExcluding( + const Children& children +) { + constexpr TCodepoint kMaxUnicodeCodepoint = 0x10FFFF; + + // Ordinary key characters must not consume NUL, XML-sensitive characters (which are + // represented by entity alternatives), or a codepoint handled by a child trie branch. + std::vector excluded; + excluded.reserve(1 + kCohereKeyEntities.size() + children.size()); + excluded.push_back('\0'); + for (const auto& entry : kCohereKeyEntities) { + excluded.push_back(entry.first); + } + for (const auto& entry : children) { + excluded.push_back(entry.first); + } + + // Sorting makes duplicate exclusions adjacent so that unique + erase can remove them. + std::sort(excluded.begin(), excluded.end()); + excluded.erase(std::unique(excluded.begin(), excluded.end()), excluded.end()); + + // Build the positive character class as the gaps between excluded codepoints. Positive + // ranges preserve full Unicode support when the grammar is lowered to an FSM. + std::vector ranges; + TCodepoint range_start = 0; + for (TCodepoint codepoint : excluded) { + if (codepoint < range_start) { + continue; + } + if (range_start < codepoint) { + ranges.push_back({range_start, codepoint - 1}); + } + range_start = codepoint + 1; + } + if (range_start <= kMaxUnicodeCodepoint) { + ranges.push_back({range_start, kMaxUnicodeCodepoint}); + } + return ranges; +} + +} // namespace + +const std::string CohereXMLToolCallingConverter::kCohereKey = "cohere_key"; +const std::string CohereXMLToolCallingConverter::kCohereAnyScalar = "cohere_any_scalar"; +const std::string CohereXMLToolCallingConverter::kCohereAnyList = "cohere_any_list"; + +CohereXMLToolCallingConverter::CohereXMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order +) + : XMLToolCallingConverter( + indent, + separators, + any_whitespace, + max_whitespace_cnt, + ref_resolver, + JSONFormat::kCohereXML, + any_order + ) {} + +void CohereXMLToolCallingConverter::AddBasicRules() { + // Cohere's dynamic key and recursive Any rules must have stable targets before kXMLObject is + // built, because its additional-property formatting reaches them through virtual dispatch. + builder_.AddEmptyRule(kCohereKey); + builder_.AddEmptyRule(kCohereAnyScalar); + builder_.AddEmptyRule(kCohereAnyList); + + XMLToolCallingConverter::AddBasicRules(); + + builder_.UpdateRuleBody(kCohereKey, RegexExpression(R"(([^\x00"&<>]|&|<|>|")+)")); + + builder_.UpdateRuleBody( + kCohereAnyScalar, Choice({RuleRef(kBasicNumber), RuleRef(kBasicBoolean), RuleRef(kBasicNull)}) + ); + // kXMLObject already provides named_any_value* through its dynamic-property formatting. + // Lists need the corresponding recursive sequence of unnamed Any wrappers explicitly. + int32_t unnamed_any_value = FormatAnyCohereParam(std::nullopt, std::nullopt); + builder_.UpdateRuleBody( + kCohereAnyList, Repeat("cohere_any_list_items", unnamed_any_value, 0, -1) + ); +} + +bool CohereXMLToolCallingConverter::AtCohereRoot() const { + return nested_object_level_ == 0 && object_stack_.empty() && cohere_array_level_ == 0; +} + +bool CohereXMLToolCallingConverter::InCohereValueContext() const { + return nested_object_level_ <= 1 || !object_stack_.empty() || cohere_array_level_ > 0; +} + +std::string CohereXMLToolCallingConverter::RefCacheKey(const std::string& uri) const { + // Recursive Cohere dict/list contents remain tagged values, regardless of nesting depth. + int context = AtCohereRoot() ? 0 : (InCohereValueContext() ? 1 : 2); + return std::to_string(context) + ":" + uri; +} + +int32_t CohereXMLToolCallingConverter::FormatCohereValue(int32_t value_rule_id) { + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + return RuleRef(value_rule_id); + } + return Sequence({WhitespaceExpression(), RuleRef(value_rule_id), WhitespaceExpression()}); +} + +std::string CohereXMLToolCallingConverter::CohereTypeForJSONLiteral(const std::string& json_value) { + picojson::value value; + std::string error = ParseJSON(value, json_value); + // Const/enum object and array literals are emitted as JSON text today, not recursive Cohere + // dict/list bodies, so only JSON strings get the raw Cohere type. + return error.empty() && value.is() ? "raw" : "json"; +} + +std::optional CohereXMLToolCallingConverter::CommonCohereTypeForJSONLiterals( + const std::vector& json_values +) { + std::optional common_type; + for (const auto& json_value : json_values) { + auto type = CohereTypeForJSONLiteral(json_value); + if (!common_type.has_value()) { + common_type = type; + } else if (*common_type != type) { + return std::nullopt; + } + } + return common_type; +} + +int32_t CohereXMLToolCallingConverter::GetCohereTypePattern(const SchemaSpecPtr& schema) { + return std::visit( + [this](const auto& spec) -> int32_t { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return ByteString("raw"); + } else if constexpr (std::is_same_v) { + return ByteString("dict"); + } else if constexpr (std::is_same_v) { + return ByteString("list"); + } else if constexpr (std::is_same_v) { + return ByteString(CohereTypeForJSONLiteral(spec.json_value)); + } else if constexpr (std::is_same_v) { + auto common_type = CommonCohereTypeForJSONLiterals(spec.json_values); + // Mixed enums are branch-correlated by FormatCohereParam. This fallback is only used if + // a mixed enum somehow reaches the single-wrapper path, where there is no one true type. + return ByteString(common_type.has_value() ? *common_type : "json"); + } else { + return ByteString("json"); + } + }, + schema->spec + ); +} + +std::optional> CohereXMLToolCallingConverter::GetCohereCompositeOptions( + const SchemaSpecPtr& schema +) const { + if (schema == nullptr) { + return std::nullopt; + } + return std::visit( + [](const auto& spec) -> std::optional> { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return spec.options; + } else if constexpr (std::is_same_v) { + return spec.options; + } else if constexpr (std::is_same_v) { + if (spec.schemas.size() == 1) { + return spec.schemas; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + return spec.type_schemas; + } else if constexpr (std::is_same_v) { + if (spec.json_values.empty() || + CommonCohereTypeForJSONLiterals(spec.json_values).has_value()) { + return std::nullopt; + } + std::vector options; + options.reserve(spec.json_values.size()); + for (size_t index = 0; index < spec.json_values.size(); ++index) { + const auto& json_value = spec.json_values[index]; + ConstSpec const_spec; + const_spec.json_value = json_value; + options.push_back( + SchemaSpec::Make(std::move(const_spec), "", "enum_case_" + std::to_string(index)) + ); + } + return options; + } else { + return std::nullopt; + } + }, + schema->spec + ); +} + +int32_t CohereXMLToolCallingConverter::CohereParamPrefix( + const std::optional& name, const std::optional& key_pattern_expr +) { + std::vector elements = {ByteString(xml_wrapper_.key_wrapper_prefix)}; + if (name.has_value()) { + elements.push_back(ByteString(" name=\"" + SerializeCohereKey(*name) + "\"")); + } else if (key_pattern_expr.has_value()) { + elements.push_back(ByteString(" name=\"")); + elements.push_back(*key_pattern_expr); + elements.push_back(ByteString("\"")); + } + return Sequence(elements); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereSuffixWithType( + int32_t type_expression, int32_t value_rule_id +) { + std::vector elements = { + ByteString(" type=\""), type_expression, ByteString("\"" + xml_wrapper_.key_wrapper_suffix) + }; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + elements.push_back(FormatCohereValue(value_rule_id)); + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); +} + +int32_t CohereXMLToolCallingConverter::FormatAnyCohereSuffix() { + // kXMLAny is the aggregate body-only union. Wrapping it under every type would create a + // type/body cross product, so each wrapper deliberately references its matching component. + return Choice( + {FormatCohereSuffixWithType(ByteString("raw"), builder_.GetRuleId(kXMLString)), + FormatCohereSuffixWithType(ByteString("json"), builder_.GetRuleId(kCohereAnyScalar)), + FormatCohereSuffixWithType(ByteString("dict"), builder_.GetRuleId(kXMLObject)), + FormatCohereSuffixWithType(ByteString("list"), builder_.GetRuleId(kCohereAnyList))} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatAnyCohereParam( + const std::optional& name, const std::optional& key_pattern_expr +) { + return Sequence({CohereParamPrefix(name, key_pattern_expr), FormatAnyCohereSuffix()}); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereParam( + const std::optional& name, + const std::optional& key_pattern_expr, + const SchemaSpecPtr& schema, + int32_t value_rule_id +) { + return Sequence( + {CohereParamPrefix(name, key_pattern_expr), FormatCohereParamSuffix(schema, value_rule_id)} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatCohereParamSuffix( + const SchemaSpecPtr& schema, int32_t value_rule_id +) { + // Copy the name before generating: GenerateFromSpec may add rules and reallocate the + // builder's rule storage, invalidating references into it. + std::string value_rule_name = builder_.GetRule(value_rule_id).name; + if (const auto* ref = std::get_if(&schema->spec); ref != nullptr) { + auto cached = cohere_param_ref_rules_.find(ref->uri); + if (cached != cohere_param_ref_rules_.end()) { + return RuleRef(cached->second); + } + // Register the rule before resolving the reference. A schema that leads back to this URI, + // directly or through nested dict/list items, then reuses the rule instead of expanding + // again without bound, and shared acyclic references are built only once. + int32_t param_rule_id = builder_.AddEmptyRuleWithHint(value_rule_name + "_cohere_param"); + cohere_param_ref_rules_.emplace(ref->uri, param_rule_id); + SchemaSpecPtr resolved = ResolveRefSchema(*ref, value_rule_name); + builder_.UpdateRuleBody(param_rule_id, FormatCohereParamSuffix(resolved, value_rule_id)); + return RuleRef(param_rule_id); + } + + // CreateRule may return any aggregate rule (cached or freshly generated), but the schema is + // retained along this call path so Any can select correlated wrappers here. + if (std::holds_alternative(schema->spec)) { + return FormatAnyCohereSuffix(); + } + if (const auto* all_of = std::get_if(&schema->spec); + all_of != nullptr && all_of->schemas.size() != 1) { + // The base converter intentionally falls back to Any while multi-branch allOf support is + // incomplete. Keep that fallback canonical instead of wrapping its aggregate body as json. + return FormatAnyCohereSuffix(); + } + + auto options = GetCohereCompositeOptions(schema); + if (!options.has_value()) { + return FormatCohereSuffixWithType(GetCohereTypePattern(schema), value_rule_id); + } + + std::vector choices; + choices.reserve(options->size()); + for (size_t index = 0; index < options->size(); ++index) { + const SchemaSpecPtr& option = (*options)[index]; + int32_t option_rule_id = + CreateRule(option, value_rule_name + "_cohere_case_" + std::to_string(index)); + choices.push_back(FormatCohereParamSuffix(option, option_rule_id)); + } + return choices.size() == 1 ? choices[0] : Choice(choices); +} + +int32_t CohereXMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateString(spec, rule_name); + } + if (!spec.pattern.has_value() && !spec.format.has_value() && spec.min_length == 0 && + spec.max_length == -1) { + return RuleRef(kXMLString); + } + if (spec.format.has_value()) { + const std::string& format = *spec.format; + auto regex_pattern = JSONFormatToRegexPattern(format); + if (regex_pattern.has_value()) { + return RegexExpression(regex_pattern.value(), false, true); + } + } + if (spec.pattern.has_value()) { + return RegexExpression(*spec.pattern, false, /*force_cfg_expansion=*/true); + } + if (spec.min_length != 0 || spec.max_length != -1) { + return Repeat( + rule_name + "_characters", + builder_.AddCharacterClass({{0, 0x10ffff}}), + spec.min_length, + spec.max_length + ); + } + return JSONSchemaConverter::GenerateString(spec, rule_name); +} + +int32_t CohereXMLToolCallingConverter::GenerateAny( + const AnySpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateAny(spec, rule_name); + } + if (AtCohereRoot()) { + return RuleRef(kXMLObject); + } + return Choice( + {RuleRef(kXMLString), RuleRef(kCohereAnyScalar), RuleRef(kXMLObject), RuleRef(kCohereAnyList)} + ); +} + +int32_t CohereXMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateConst(spec, rule_name); + } + return ByteString(XMLValue(spec.json_value)); +} + +int32_t CohereXMLToolCallingConverter::GenerateEnum( + const EnumSpec& spec, const std::string& rule_name +) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + if (!InCohereValueContext()) { + return JSONSchemaConverter::GenerateEnum(spec, rule_name); + } + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + values.push_back(ByteString(XMLValue(value))); + } + return Choice(values); +} + +int32_t CohereXMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + nested_object_level_++; + bool use_cohere_object = InCohereValueContext(); + + int32_t result; + if (use_cohere_object) { + SchemaSpecPtr additional_property; + if (spec.allow_additional_properties && spec.additional_properties_schema) { + additional_property = spec.additional_properties_schema; + } else if (spec.allow_unevaluated_properties && spec.unevaluated_properties_schema) { + additional_property = spec.unevaluated_properties_schema; + } else if (spec.allow_additional_properties || spec.allow_unevaluated_properties) { + additional_property = SchemaSpec::Make(AnySpec{}, "", "any"); + } + + object_stack_.push_back(&spec); + additional_property_stack_.push_back(additional_property); + result = JSONSchemaConverter::GenerateObject(spec, rule_name, false); + additional_property_stack_.pop_back(); + object_stack_.pop_back(); + } else { + result = JSONSchemaConverter::GenerateObject(spec, rule_name, nested_object_level_ > 1); + } + + nested_object_level_--; + return result; +} + +int32_t CohereXMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + if (!InCohereValueContext()) { + nested_object_level_++; + auto result = JSONSchemaConverter::GenerateArray(spec, rule_name); + nested_object_level_--; + return result; + } + + cohere_array_level_++; + std::vector item_patterns; + for (size_t i = 0; i < spec.prefix_items.size(); ++i) { + int32_t item_rule_id = + CreateRule(spec.prefix_items[i], rule_name + "_item_" + std::to_string(i)); + item_patterns.push_back( + FormatCohereParam(std::nullopt, std::nullopt, spec.prefix_items[i], item_rule_id) + ); + } + + std::optional additional_item_pattern; + if (spec.allow_additional_items && spec.additional_items) { + int32_t additional_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + additional_item_pattern = + FormatCohereParam(std::nullopt, std::nullopt, spec.additional_items, additional_rule_id); + } + cohere_array_level_--; + + if (item_patterns.empty()) { + if (!additional_item_pattern.has_value() || spec.max_items == 0) { + return Empty(); + } + return Repeat( + rule_name + "_items", + *additional_item_pattern, + static_cast(spec.min_items), + spec.max_items == -1 ? -1 : static_cast(spec.max_items) + ); + } + + int32_t prefix_part = Sequence(item_patterns); + if (!additional_item_pattern.has_value()) { + return prefix_part; + } + + int64_t min_additional = std::max( + static_cast(0), spec.min_items - static_cast(item_patterns.size()) + ); + int64_t max_additional = + spec.max_items == -1 ? -1 : spec.max_items - static_cast(item_patterns.size()); + return Sequence( + {prefix_part, + Repeat( + rule_name + "_additional_items", + *additional_item_pattern, + static_cast(min_additional), + max_additional == -1 ? -1 : static_cast(max_additional) + )} + ); +} + +int32_t CohereXMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + if (!object_stack_.empty() && idx >= 0 && + idx < static_cast(object_stack_.back()->properties.size())) { + const auto& prop = object_stack_.back()->properties[idx]; + return FormatCohereParam(prop.name, std::nullopt, prop.schema, value_rule_id); + } + return XMLToolCallingConverter::FormatProperty(key, value_rule_id, rule_name, idx, schema); +} + +int32_t CohereXMLToolCallingConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + SchemaSpecPtr value_schema = schema; + if (!value_schema && !additional_property_stack_.empty()) { + value_schema = additional_property_stack_.back(); + } + if (!value_schema && InCohereValueContext()) { + value_schema = SchemaSpec::Make(AnySpec{}, "", "any"); + value_rule_id = CreateRule(value_schema, rule_name + "_" + rule_name_suffix + "_cohere_any"); + } + if (value_schema) { + return FormatCohereParam(std::nullopt, key_pattern_expr, value_schema, value_rule_id); + } + return XMLToolCallingConverter::FormatOtherProperty( + key_pattern_expr, value_rule_id, rule_name, rule_name_suffix, schema + ); +} + +std::string CohereXMLToolCallingConverter::GetKeyPattern() const { + if (InCohereValueContext()) { + return kCohereKey; + } + return JSONSchemaConverter::GetKeyPattern(); +} + +int32_t CohereXMLToolCallingConverter::BuildCohereKeyExcludingBody( + const CohereKeyTrieNode& node, int depth +) { + std::vector choices; + if (depth > 0 && !node.is_terminal) { + choices.push_back(Empty()); + } + + int32_t optional_key_suffix = Choice({Empty(), RuleRef(kCohereKey)}); + int32_t ordinary_key_unit = + builder_.AddCharacterClass(CohereOrdinaryKeyRangesExcluding(node.children)); + choices.push_back(Sequence({ordinary_key_unit, optional_key_suffix})); + for (const auto& [codepoint, entity] : kCohereKeyEntities) { + if (!node.children.count(codepoint)) { + int32_t entity_key_unit = ByteString(entity); + choices.push_back(Sequence({entity_key_unit, optional_key_suffix})); + } + } + + for (const auto& [codepoint, child] : node.children) { + choices.push_back(Sequence( + {ByteString(SerializeCohereKeyCodepoint(codepoint)), + BuildCohereKeyExcludingBody(child, depth + 1)} + )); + } + + return Choice(choices); +} + +int32_t CohereXMLToolCallingConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (InCohereValueContext()) { + if (properties.empty()) { + return RuleRef(GetKeyPattern()); + } + CohereKeyTrieNode root; + for (const auto& prop : properties) { + CohereKeyTrieNode* cur = &root; + auto codepoints = ParseCohereKeyCodepoints(prop.name); + for (TCodepoint codepoint : codepoints) { + cur = &cur->children[codepoint]; + } + if (!codepoints.empty()) { + cur->is_terminal = true; + } + } + int32_t key_rule_id = builder_.AddEmptyRuleWithHint(rule_name + "_cohere_addl_key"); + builder_.UpdateRuleBody(key_rule_id, BuildCohereKeyExcludingBody(root, 0)); + return RuleRef(key_rule_id); + } + return JSONSchemaConverter::GetKeyPatternExcluding(properties, rule_name); +} + +std::string CohereXMLToolCallingConverter::NextSeparator(bool is_end) { + if (InCohereValueContext()) { + return GetWhitespacePattern(); + } + return JSONSchemaConverter::NextSeparator(is_end); +} + +void CohereXMLToolCallingConverter::AddCache(const std::string& key, int32_t rule_id) { + if (key.empty()) { + return; + } + rule_cache_manager_.AddCache(key, nested_object_level_ > 1 && !InCohereValueContext(), rule_id); +} + +std::optional CohereXMLToolCallingConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + // "true" and {} are equivalent schemas. At the tool-arguments root both are unrestricted + // dictionaries, while nested {} keeps using the aggregate Any body rule. + if (AtCohereRoot() && (key == "{}" || key == "true")) { + return builder_.GetRuleId(kXMLObject); + } + return rule_cache_manager_.GetCache(key, nested_object_level_ > 1 && !InCohereValueContext()); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/deepseek.cc b/zig/vendor/xgrammar/cpp/converter_ext/deepseek.cc new file mode 100644 index 0000000000..a7ee874dd9 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/deepseek.cc @@ -0,0 +1,117 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/deepseek.cc + * \brief DeepSeek XML parameter formats and value rendering. + */ +#include + +#include "../json_schema_converter_ext.h" + +namespace xgrammar { +namespace converter_ext { + +XMLWrapper GetDeepSeekXMLWrapper() { + return {"<|DSML|parameter name=\"", "", "", ""}; +} + +XMLWrapper GetDeepSeekV41XMLWrapper() { + return {"<|DSML| parameter name=\"", "", "", ""}; +} + +} // namespace converter_ext + +int32_t XMLToolCallingConverter::FormatDeepSeekParamSuffix( + const SchemaSpecPtr& schema, int32_t value_rule_id, const std::string& rule_name_hint +) { + // Copy the name: creating alternative rules can reallocate the builder's rule storage. + std::string value_rule_name = + value_rule_id >= 0 ? builder_.GetRule(value_rule_id).name : rule_name_hint; + if (schema != nullptr) { + if (const auto* ref = std::get_if(&schema->spec); ref != nullptr) { + auto cached = deepseek_param_ref_rules_.find(ref->uri); + if (cached != deepseek_param_ref_rules_.end()) { + return RuleRef(cached->second); + } + // Cache the rule before descending through references or alternatives. A recursive + // branch then refers back to this rule, and shared acyclic subgraphs are built only once. + int32_t param_rule_id = builder_.AddEmptyRuleWithHint(value_rule_name + "_dsml_param"); + deepseek_param_ref_rules_.emplace(ref->uri, param_rule_id); + auto resolved = ResolveRefSchema(*ref, value_rule_name); + builder_.UpdateRuleBody( + param_rule_id, FormatDeepSeekParamSuffix(resolved, value_rule_id, value_rule_name) + ); + return RuleRef(param_rule_id); + } + } + + // string="true" wraps a raw string whose whitespace is part of the value; string="false" + // wraps a JSON value that may be padded with whitespace. + auto wrap = [&](int32_t value_expr, bool is_string) { + std::vector elements = { + ByteString(is_string ? "\" string=\"true\">" : "\" string=\"false\">") + }; + if (!is_string) elements.push_back(WhitespaceExpression()); + elements.push_back(value_expr); + if (!is_string) elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + }; + + // A schema rendered with a single type keeps the value rule built by the caller. + std::optional pinned_type = GetRenderedJSONType(schema); + if (pinned_type.has_value()) { + if (value_rule_id < 0) { + value_rule_id = CreateRule(schema, value_rule_name); + } + return wrap(RuleRef(value_rule_id), *pinned_type == "string"); + } + + // Unions and mixed enums get one alternative per option so each carries its own attribute. + std::vector options; + if (schema != nullptr) { + std::visit( + [&](const auto& spec) { + using T = std::decay_t; + if constexpr (std::is_same_v || std::is_same_v) { + options = spec.options; + } else if constexpr (std::is_same_v) { + options = spec.type_schemas; + } else if constexpr (std::is_same_v) { + if (spec.schemas.size() == 1) options = spec.schemas; + } else if constexpr (std::is_same_v) { + for (const auto& value : spec.json_values) { + options.push_back(SchemaSpec::Make(ConstSpec{value})); + } + } + }, + schema->spec + ); + } + if (options.empty()) { + // No schema, {} and allOf with several schemas all render any value. + return Choice( + {wrap(RuleRef(kXMLString), true), + wrap( + Choice( + {RuleRef(kBasicNumber), + RuleRef(kBasicBoolean), + RuleRef(kBasicNull), + RuleRef(kBasicArray), + RuleRef(kBasicObject)} + ), + false + )} + ); + } + std::vector choices; + for (size_t index = 0; index < options.size(); ++index) { + // Only leaves need value rules. Rebuilding a nested union's whole value subtree at each + // level would create quadratically many rules before assembling its typed alternatives. + choices.push_back(FormatDeepSeekParamSuffix( + options[index], -1, value_rule_name + "_dsml_case_" + std::to_string(index) + )); + } + return Choice(choices); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/glm.cc b/zig/vendor/xgrammar/cpp/converter_ext/glm.cc new file mode 100644 index 0000000000..2242543b5c --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/glm.cc @@ -0,0 +1,14 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/glm.cc + * \brief GLM XML parameter format. + */ +#include "../json_schema_converter_ext.h" + +namespace xgrammar { +namespace converter_ext { + +XMLWrapper GetGLMXMLWrapper() { return {"", "", "", ""}; } + +} // namespace converter_ext +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/kimi_k3.cc b/zig/vendor/xgrammar/cpp/converter_ext/kimi_k3.cc new file mode 100644 index 0000000000..b244c94826 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/kimi_k3.cc @@ -0,0 +1,47 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/kimi_k3.cc + * \brief Kimi-K3 XML type attributes and key escaping. + */ +#include "../json_schema_converter_ext.h" + +namespace xgrammar { + +namespace converter_ext { + +XMLWrapper GetKimiK3XMLWrapper() { + // The key suffix (type attribute and <|sep|>) is generated in XMLKeySuffix. + return {"<|open|>argument key=\"", "", "", "<|close|>argument<|sep|>"}; +} + +const XMLKeySuffix& GetKimiK3XMLKeySuffix() { + static const XMLKeySuffix suffix = { + "\" type=\"", + {"string", "number", "integer", "boolean", "object", "array", "null"}, + "\"<|sep|>" + }; + return suffix; +} + +} // namespace converter_ext + +std::string XMLToolCallingConverter::EscapeAttrValue(const std::string& value) const { + if (json_format_ != JSONFormat::kKimiK3XML) { + return value; + } + // Kimi-K3's renderer escapes attribute values with & -> & and " -> ". + std::string escaped; + escaped.reserve(value.size()); + for (char c : value) { + if (c == '&') { + escaped += "&"; + } else if (c == '"') { + escaped += """; + } else { + escaped += c; + } + } + return escaped; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/minimax.cc b/zig/vendor/xgrammar/cpp/converter_ext/minimax.cc new file mode 100644 index 0000000000..d2d710707b --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/minimax.cc @@ -0,0 +1,14 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/minimax.cc + * \brief MiniMax XML parameter format. + */ +#include "../json_schema_converter_ext.h" + +namespace xgrammar { +namespace converter_ext { + +XMLWrapper GetMiniMaxXMLWrapper() { return {"", "", ""}; } + +} // namespace converter_ext +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/minimax_m3.cc b/zig/vendor/xgrammar/cpp/converter_ext/minimax_m3.cc new file mode 100644 index 0000000000..241fdc528f --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/minimax_m3.cc @@ -0,0 +1,315 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/converter_ext/minimax_m3.cc + * \brief Implementation of the MiniMax M3 XML Tool Calling converter. + */ +#include + +#include +#include +#include +#include +#include + +#include "../json_schema_converter_ext.h" +#include "../support/encoding.h" +#include "../support/logging.h" + +namespace xgrammar { + +namespace { + +constexpr const char* kStringCacheKey = "{\"type\":\"string\"}"; +constexpr const char* kMiniMaxM3Namespace = "]<]minimax[>["; +constexpr const char* kMiniMaxM3ArrayItemName = "item"; + +bool IsASCIIWhitespace(uint8_t byte) { + return byte == ' ' || byte == '\t' || byte == '\n' || byte == '\r' || byte == '\f' || + byte == '\v'; +} + +bool IsCanonicalUTF8(const std::string& text) { + for (size_t offset = 0; offset < text.size();) { + auto [codepoint, num_bytes] = ParseNextUTF8(text.data() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8 || num_bytes <= 0 || + offset + num_bytes > text.size() || (codepoint >= 0xd800 && codepoint <= 0xdfff) || + codepoint > 0x10ffff || text.compare(offset, num_bytes, CharToUTF8(codepoint)) != 0) { + return false; + } + offset += num_bytes; + } + return true; +} + +} // namespace + +MiniMaxM3XMLToolCallingConverter::MiniMaxM3XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order +) + : JSONSchemaConverter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ) {} + +void MiniMaxM3XMLToolCallingConverter::AddBasicRules() { + for (const auto& name : {kBasicInteger, kBasicNumber, kBasicString, kBasicBoolean, kBasicNull}) { + builder_.AddEmptyRule(name); + } + + builder_.UpdateRuleBody( + kBasicInteger, JSONSchemaConverter::GenerateInteger(IntegerSpec{}, kBasicInteger) + ); + AddCache("{\"type\":\"integer\"}", builder_.GetRuleId(kBasicInteger)); + + builder_.UpdateRuleBody( + kBasicNumber, JSONSchemaConverter::GenerateNumber(NumberSpec{}, kBasicNumber) + ); + AddCache("{\"type\":\"number\"}", builder_.GetRuleId(kBasicNumber)); + + builder_.UpdateRuleBody(kBasicString, TagDispatch(false, {kMiniMaxM3Namespace})); + AddCache(kStringCacheKey, builder_.GetRuleId(kBasicString)); + + builder_.UpdateRuleBody( + kBasicBoolean, JSONSchemaConverter::GenerateBoolean(BooleanSpec{}, kBasicBoolean) + ); + AddCache("{\"type\":\"boolean\"}", builder_.GetRuleId(kBasicBoolean)); + + builder_.UpdateRuleBody(kBasicNull, JSONSchemaConverter::GenerateNull(NullSpec{}, kBasicNull)); + AddCache("{\"type\":\"null\"}", builder_.GetRuleId(kBasicNull)); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + const bool has_known_format = + spec.format.has_value() && JSONFormatToRegexPattern(*spec.format).has_value(); + XGRAMMAR_CHECK( + !spec.pattern.has_value() && !has_known_format && spec.min_length == 0 && + spec.max_length == -1 + ) << "String pattern, recognized format, and length constraints are not supported by " + "minimax_m3_xml because they cannot be combined with the namespace-marker exclusion"; + return RuleRef(kBasicString); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + constexpr int64_t kMaxRepeatCount = std::numeric_limits::max(); + XGRAMMAR_CHECK( + spec.min_items <= kMaxRepeatCount && + (spec.max_items == -1 || spec.max_items <= kMaxRepeatCount) && + spec.prefix_items.size() <= static_cast(kMaxRepeatCount) + ) << "minimax_m3_xml array bounds exceed the supported range"; + XGRAMMAR_CHECK(!spec.allow_additional_items || spec.additional_items != nullptr) + << "minimax_m3_xml requires a fixed schema for array items"; + + std::vector prefix_items; + prefix_items.reserve(spec.prefix_items.size()); + for (size_t index = 0; index < spec.prefix_items.size(); ++index) { + int32_t item_rule_id = + CreateRule(spec.prefix_items[index], rule_name + "_item_" + std::to_string(index)); + prefix_items.push_back(FormatElement(kMiniMaxM3ArrayItemName, item_rule_id)); + } + + std::optional additional_item; + if (spec.allow_additional_items) { + int32_t item_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + additional_item = FormatElement(kMiniMaxM3ArrayItemName, item_rule_id); + } + + int32_t empty = Empty(); + int32_t whitespace = WhitespaceExpression(); + if (prefix_items.empty()) { + if (!additional_item.has_value() || spec.max_items == 0) { + return empty; + } + int32_t min_items = static_cast(spec.min_items); + int32_t max_items = spec.max_items == -1 ? -1 : static_cast(spec.max_items); + int32_t nonempty = Sequence( + {whitespace, + *additional_item, + Repeat( + rule_name + "_items", + Sequence({whitespace, *additional_item}), + std::max(0, min_items - 1), + max_items == -1 ? -1 : std::max(0, max_items - 1) + ), + whitespace} + ); + return min_items == 0 ? Choice({nonempty, empty}) : nonempty; + } + + int32_t prefix_count = static_cast(prefix_items.size()); + int32_t tail = empty; + if (additional_item.has_value()) { + int32_t min_additional = std::max(0, static_cast(spec.min_items) - prefix_count); + int32_t max_additional = spec.max_items == -1 + ? -1 + : std::max(0, static_cast(spec.max_items) - prefix_count); + tail = Repeat( + rule_name + "_additional_items", + Sequence({whitespace, *additional_item}), + min_additional, + max_additional + ); + } + + // A prefixItems entry constrains its position but does not make that position mandatory. Build + // a linear chain whose suffix can stop once minItems is satisfied. + for (int32_t index = prefix_count - 2; index >= 0; --index) { + int32_t emitted_count = index + 1; + bool can_stop = emitted_count >= spec.min_items; + bool can_continue = spec.max_items == -1 || emitted_count < spec.max_items; + int32_t body = empty; + if (can_continue) { + int32_t continuation = Sequence({whitespace, prefix_items[index + 1], tail}); + body = can_stop ? Choice({continuation, empty}) : continuation; + } + int32_t tail_rule_id = + builder_.AddRuleWithHint(rule_name + "_prefix_tail_" + std::to_string(index), body); + tail = RuleRef(tail_rule_id); + } + + if (spec.max_items == 0) { + return empty; + } + int32_t nonempty = Sequence({whitespace, prefix_items[0], tail, whitespace}); + return spec.min_items == 0 ? Choice({nonempty, empty}) : nonempty; +} + +void MiniMaxM3XMLToolCallingConverter::ValidateObject(const ObjectSpec& spec) const { + XGRAMMAR_CHECK( + !spec.allow_additional_properties && spec.additional_properties_schema == nullptr && + !spec.allow_unevaluated_properties && spec.unevaluated_properties_schema == nullptr && + spec.pattern_properties.empty() && spec.property_names == nullptr + ) << "minimax_m3_xml requires fixed object property names; additionalProperties, " + "unevaluatedProperties, patternProperties, and propertyNames are not supported"; + + std::unordered_set property_names; + for (const auto& property : spec.properties) { + XGRAMMAR_CHECK(property.schema != nullptr) + << "minimax_m3_xml property must have a fixed schema: " << property.name; + ValidateElementName(property.name); + property_names.insert(property.name); + } + for (const auto& required : spec.required) { + XGRAMMAR_CHECK(property_names.count(required) != 0) + << "minimax_m3_xml required property has no fixed schema: " << required; + } +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + ValidateObject(spec); + bool saved_any_whitespace = any_whitespace_; + any_whitespace_ = false; + int32_t result = JSONSchemaConverter::GenerateObject(spec, rule_name, /*need_braces=*/false); + any_whitespace_ = saved_any_whitespace; + return result; +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateAny( + const AnySpec& spec, const std::string& rule_name +) { + XGRAMMAR_LOG(FATAL) << "minimax_m3_xml does not support unconstrained schemas"; + XGRAMMAR_UNREACHABLE(); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateLiteral(const picojson::value& value) { + if (value.is()) { + const std::string& text = value.get(); + XGRAMMAR_CHECK(text.find(kMiniMaxM3Namespace) == std::string::npos) + << "A minimax_m3_xml string literal cannot contain the namespace marker"; + return ByteString(text); + } + if (value.is()) { + const auto& object = value.get(); + std::vector properties; + properties.reserve(object.size()); + for (const auto& key : object.ordered_keys()) { + int32_t value_expr = GenerateLiteral(object.at(key)); + int32_t value_rule_id = builder_.AddRuleWithHint("literal_" + key, value_expr); + properties.push_back(FormatElement(key, value_rule_id)); + } + return Sequence(properties); + } + if (value.is()) { + const auto& array = value.get(); + std::vector items; + items.reserve(array.size()); + for (size_t index = 0; index < array.size(); ++index) { + int32_t value_expr = GenerateLiteral(array[index]); + int32_t value_rule_id = + builder_.AddRuleWithHint("literal_item_" + std::to_string(index), value_expr); + items.push_back(FormatElement(kMiniMaxM3ArrayItemName, value_rule_id)); + } + return Sequence(items); + } + return ByteString(value.serialize()); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + picojson::value value; + std::string error = picojson::parse(value, spec.json_value); + XGRAMMAR_CHECK(error.empty()) << "Invalid const JSON value: " << error; + return GenerateLiteral(value); +} + +int32_t MiniMaxM3XMLToolCallingConverter::GenerateEnum( + const EnumSpec& spec, const std::string& rule_name +) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& json_value : spec.json_values) { + picojson::value value; + std::string error = picojson::parse(value, json_value); + XGRAMMAR_CHECK(error.empty()) << "Invalid enum JSON value: " << error; + values.push_back(GenerateLiteral(value)); + } + return Choice(values); +} + +void MiniMaxM3XMLToolCallingConverter::ValidateElementName(const std::string& name) { + XGRAMMAR_CHECK(!name.empty() && name.front() != '/' && name.find('>') == std::string::npos) + << "Invalid minimax_m3_xml element name: " << name; + XGRAMMAR_CHECK(IsCanonicalUTF8(name)) << "minimax_m3_xml element names must be valid UTF-8"; + XGRAMMAR_CHECK(std::any_of(name.begin(), name.end(), [](unsigned char byte) { + return !IsASCIIWhitespace(byte); + })) << "minimax_m3_xml element names cannot be blank"; +} + +int32_t MiniMaxM3XMLToolCallingConverter::FormatElement( + const std::string& name, int32_t value_rule_id +) { + ValidateElementName(name); + return Sequence( + {ByteString(std::string(kMiniMaxM3Namespace) + "<" + name + ">"), + RuleRef(value_rule_id), + ByteString(std::string(kMiniMaxM3Namespace) + "")} + ); +} + +int32_t MiniMaxM3XMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + return FormatElement(key, value_rule_id); +} + +std::string MiniMaxM3XMLToolCallingConverter::NextSeparator(bool is_end) { + return GetWhitespacePattern(); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/converter_ext/qwen.cc b/zig/vendor/xgrammar/cpp/converter_ext/qwen.cc new file mode 100644 index 0000000000..6c649228e7 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/converter_ext/qwen.cc @@ -0,0 +1,14 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/converter_ext/qwen.cc + * \brief Qwen XML parameter format. + */ +#include "../json_schema_converter_ext.h" + +namespace xgrammar { +namespace converter_ext { + +XMLWrapper GetQwenXMLWrapper() { return {"", "", ""}; } + +} // namespace converter_ext +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/earley_parser.cc b/zig/vendor/xgrammar/cpp/earley_parser.cc new file mode 100644 index 0000000000..675ca7bddb --- /dev/null +++ b/zig/vendor/xgrammar/cpp/earley_parser.cc @@ -0,0 +1,1272 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/earley_parser.cc + */ + +#include "earley_parser.h" + +#include +#include +#include +#include +#include + +#include "fsm.h" +#include "grammar_impl.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +using GrammarExprType = Grammar::Impl::GrammarExprType; + +using GrammarExpr = Grammar::Impl::GrammarExpr; + +bool EarleyParser::IsCompleted() const { return is_completed_.back(); } + +bool EarleyParser::CompletionConsumedMarker(const ParserState& state) const { + const auto& body = grammar_->GetGrammarExpr(grammar_->GetRule(state.rule_id).body_expr_id); + if (body.type != GrammarExprType::kTagDispatch) { + return true; + } + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id]->GetFsm().GetFsm(); + return fsm.GetEdges(state.element_id).size() == 0; +} + +std::vector EarleyParser::CollectStopCaptureTargets(const ParserState& state +) const { + // Follow only the parent links of this concrete rule occurrence. A byte-overlap test at + // materialization time cannot distinguish an actual captured ancestor from an unrelated + // Earley branch that happens to cover the same input. + std::vector targets; + std::vector pending{{state.rule_id, state.rule_start_pos}}; + std::unordered_set visited; + while (!pending.empty()) { + CaptureOccurrence occurrence = pending.back(); + pending.pop_back(); + int64_t occurrence_key = (static_cast(occurrence.rule_id) << 32) | + static_cast(occurrence.start_pos); + if (!visited.insert(occurrence_key).second) { + continue; + } + if (RuleHasCapture(occurrence.rule_id)) { + targets.push_back(occurrence); + } + if (occurrence.start_pos == ParserState::kNoPrevInputPos) { + continue; + } + const auto& parent_states = rule_id_to_completable_states_[occurrence.start_pos]; + for (const auto& [ref_rule_id, parent_state] : parent_states) { + if (ref_rule_id != occurrence.rule_id || parent_state.rule_id < 0) { + continue; + } + pending.push_back({parent_state.rule_id, parent_state.rule_start_pos}); + } + } + return targets; +} + +void EarleyParser::RecordCaptureEvent(const ParserState& state, bool marker_present) { + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(state.rule_id); + bool marker_consumed = + marker_present && suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0) && + CompletionConsumedMarker(state); + const int32_t hidden_suffix_bytes = marker_consumed ? suffix_stop_info->hidden_suffix_bytes : 0; + const int32_t hidden_stop_bytes = marker_consumed ? suffix_stop_info->hidden_stop_bytes : 0; + + int32_t event_start_pos = state.rule_start_pos; + if (marker_consumed && suffix_stop_info->body_rule_id == state.rule_id) { + // A self-referencing body helper marks the zero-width event inserted immediately after a + // dynamic string trigger. Its capture span is the fixed-length marker that precedes it. + XGRAMMAR_DCHECK(event_start_pos != ParserState::kNoPrevInputPos); + event_start_pos -= std::max(hidden_suffix_bytes, hidden_stop_bytes); + XGRAMMAR_DCHECK(event_start_pos >= 0); + } + + std::vector stop_capture_targets = + hidden_stop_bytes > 0 ? CollectStopCaptureTargets(state) : std::vector{}; + + capture_event_history_.PushBackInLatestRow( + {state.rule_id, + event_start_pos, + state.rule_start_pos, + hidden_suffix_bytes, + hidden_stop_bytes, + std::move(stop_capture_targets)} + ); +} + +int32_t EarleyParser::ResolveActiveTemperatureRule(int32_t rule_id, int32_t inherited_rule_id) + const { + return grammar_->GetRule(rule_id).temperature.has_value() ? rule_id : inherited_rule_id; +} + +void EarleyParser::PopLastStates(int32_t cnt) { + stop_token_is_accepted_ = false; + if (cnt >= static_cast(rule_id_to_completable_states_.size())) { + XGRAMMAR_LOG(FATAL) << "The number of states to be popped is larger than the size of states."; + } + rule_id_to_completable_states_.PopBack(cnt); + is_completed_.erase(is_completed_.end() - cnt, is_completed_.end()); + scanable_state_history_.PopBack(cnt); + if (capture_tracking_) { + capture_event_history_.PopBack(cnt); + } + if (has_char_budget_rules_) { + char_count_history_.erase(char_count_history_.end() - cnt, char_count_history_.end()); + char_budget_entry_history_.erase( + char_budget_entry_history_.end() - cnt, char_budget_entry_history_.end() + ); + } +} + +void EarleyParser::Complete(const ParserState& state, bool debug_print, bool marker_present) { + // Record capture and hidden-span events. This is only enabled during definitive advances; + // speculative completions (mask computation, lookahead) never record events. + if (capture_recording_ && RuleNeedsCaptureEvent(state.rule_id)) { + RecordCaptureEvent(state, marker_present); + } + if (state.rule_id != -1 && grammar_->GetRule(state.rule_id).is_lazy) { + tmp_completed_lazy_occurrences_.emplace_back(state.rule_id, state.rule_start_pos); + } + // Check if a rule is completed. + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + // assert: if a root rule can achieve here, then it must be completed. + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The root rule is completed."; + } + tmp_accept_stop_token_ = true; + return; + } + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name + << " is completed, trying to complete its parent states."; + } + + // Check all the possible parent states. + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + for (const auto& [ref_id, parent_state] : parent_states_map) { + if (ref_id != state.rule_id) { + continue; + } + XGRAMMAR_DCHECK( + parent_state.rule_id == -1 || grammar_->per_rule_fsms[parent_state.rule_id].has_value() + ); + if (parent_state.rule_id == -1) { + const auto& parent_expr = grammar_->GetGrammarExpr(parent_state.sequence_id); + const auto& element_expr = grammar_->GetGrammarExpr(parent_expr[parent_state.element_id]); + // The new rule is not referenced by a fsm. + XGRAMMAR_DCHECK( + element_expr.type == GrammarExprType::kRuleRef || + element_expr.type == GrammarExprType::kRepeat + ); + if (element_expr.type == GrammarExprType::kRuleRef) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id + 1, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + continue; + } + XGRAMMAR_DCHECK(element_expr.type == GrammarExprType::kRepeat); + // The parent state is a repeat, we need to increase the repeat count. + auto new_state = parent_state; + const int32_t& min_repeat_count = element_expr[1]; + const int32_t& max_repeat_count = element_expr[2]; + new_state.repeat_count++; + // The repeat rule can be completed, and we advance the state. Don't forget to + // reset the repeat count. + if (new_state.repeat_count >= min_repeat_count) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id + 1, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + // If the repeat count is less than the max repeat count, we can continue to + // visit the repeat state for another round. + if (new_state.repeat_count < max_repeat_count) { + Enqueue(new_state); + } + continue; + } + // If the rule is referenced by a fsm, we need to advance the fsm. + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[parent_state.rule_id].has_value()); + + // Check if the parent_state sits on a kRepeatRef edge + bool handled_as_repeat = false; + const auto& parent_fsm = grammar_->per_rule_fsms[parent_state.rule_id].value(); + for (const auto& edge : parent_fsm.GetFsm().GetFsm().GetEdges(parent_state.element_id)) { + // Because of invariance, a state with a kRepeatRef edge has exactly one outgoing edge. + if (!edge.IsRepeatRef()) continue; + auto info = grammar_->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + if (info.RuleId() != ref_id) continue; + handled_as_repeat = true; + int32_t new_count = parent_state.repeat_count + 1; + if (new_count >= info.Lower()) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + edge.target, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + 0, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + if (new_count < info.Upper()) { + Enqueue(ParserState{ + parent_state.rule_id, + parent_state.sequence_id, + parent_state.element_id, + parent_state.rule_start_pos, + parent_state.budget_deadline, + 0, + new_count, + 0, + parent_state.active_temperature_rule_id, + parent_state.char_budget_deadline + }); + } + break; + } + if (!handled_as_repeat) { + Enqueue(parent_state); + } + } +} + +std::pair EarleyParser::Predict( + const ParserState& state, bool debug_print +) { + // Check if the rule has a corresponding FSM. + if (state.rule_id != -1) { + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const uint8_t flags = GetFsmStateFlags(state.rule_id, state.element_id); + if (flags & kFsmStateNonTerminal) { + ExpandNextRuleRefElementOnFSM(state, debug_print); + } + return std::make_pair(flags & kFsmStateScanable, flags & kFsmStateEnd); + } + const GrammarExpr& grammar_expr = grammar_->GetGrammarExpr(state.sequence_id); + XGRAMMAR_DCHECK( + grammar_expr.type == GrammarExprType::kSequence || + grammar_expr.type == GrammarExprType::kEmptyStr + ); + if (state.element_id == grammar_expr.size()) { + // The rule is completed. + return std::make_pair(false, true); + } + const auto& element_expr = grammar_->GetGrammarExpr(grammar_expr[state.element_id]); + switch (element_expr.type) { + case GrammarExprType::kRuleRef: { + ExpandNextRuleRefElement(state, grammar_expr, &element_expr, debug_print); + return std::make_pair(false, false); + } + case GrammarExprType::kCharacterClassStar: { + if (state.sub_element_id == 0) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + return std::make_pair(true, false); + } + case GrammarExprType::kRepeat: { + const int32_t& min_repeat_count = element_expr[1]; + const int32_t& max_repeat_count = element_expr[2]; + // If the current repeat count is less than the max repeat count, + // we can expand the next rule reference element. + XGRAMMAR_DCHECK(state.repeat_count <= max_repeat_count); + ExpandNextRuleRefElement(state, grammar_expr, &element_expr, debug_print); + if (state.repeat_count >= min_repeat_count) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + return std::make_pair(false, false); + } + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: { + return std::make_pair(true, false); // The element is scanable, but not completable. + } + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: { + return std::make_pair(false, false); + } + default: { + XGRAMMAR_LOG(FATAL) << "The element type is not supported! The type is: " + << int(element_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } +} + +void EarleyParser::Scan(const ParserState& state, const uint8_t ch) { + XGRAMMAR_DCHECK(state.rule_id == -1 || grammar_->per_rule_fsms[state.rule_id].has_value()); + if (state.rule_id == -1) { + const auto& cur_rule = grammar_->GetGrammarExpr(state.sequence_id); + const auto& element_expr = grammar_->GetGrammarExpr(cur_rule[state.element_id]); + // The element is a rule reference, we do not need to scan it. + switch (element_expr.type) { + case (GrammarExprType::kByteString): { + AdvanceByteString(state, ch, element_expr); + break; + } + case (GrammarExprType::kCharacterClass): { + AdvanceCharacterClass(state, ch, element_expr); + break; + } + case (GrammarExprType::kCharacterClassStar): { + AdvanceCharacterClassStar(state, ch, element_expr); + break; + } + default: { + XGRAMMAR_LOG(FATAL) << "The element type is not supported! The type is: " + << int(element_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + } else { + AdvanceFsm(state, ch); + } +} + +/*! + \note The workflow of Advance is as follows: + 1. Scan all the states in the latest states. Add all the possible states + to the next states. + 2. If the next states are empty, then the character is not accepted. + 3. If the next states are not empty, then the character is accepted. Moreover, + we need to complete and predict the next states. + + \note Thus, when initializing the Earley parser, we need to add the initial state + to the history_states[0], and perform prediction and completion on the initial state. +*/ +bool EarleyParser::Advance(const uint8_t ch, bool debug_print) { + // Initialize the containers. + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()) + << "The tmp_process_state_queue_ should be empty before the scan."; + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_accept_stop_token_ = false; + tmp_completed_lazy_occurrences_.clear(); + if (has_char_budget_rules_) { + tmp_char_budget_entered_ = char_budget_entry_history_.back(); + char_count_history_.push_back(GetCurrentCharIndex() + StartsUTF8Codepoint(ch)); + } + const auto& latest_states = scanable_state_history_[scanable_state_history_.size() - 1]; + // Scan all the scanable states. + for (const auto& state : latest_states) { + if (skip_expired_states_ && IsExpiredState(state)) { + continue; + } + Scan(state, ch); + } + + // Check if the character is accepted. + if (tmp_process_state_queue_.empty() && tmp_states_to_be_added_.empty()) { + if (has_char_budget_rules_) { + char_count_history_.pop_back(); + } + return false; + } + + // execute Predict and Complete for all states in the queue until empty. + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + + // Check if the grammar is completed, and add the scannable states to the history. + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } + return true; +} + +void EarleyParser::RemoveCommittedLazyStates() { + auto is_committed = [&](const ParserState& state) { + for (const auto& [rule_id, rule_start_pos] : tmp_completed_lazy_occurrences_) { + if (state.rule_id == rule_id && state.rule_start_pos == rule_start_pos) { + return true; + } + } + return false; + }; + tmp_states_to_be_added_.erase( + std::remove_if(tmp_states_to_be_added_.begin(), tmp_states_to_be_added_.end(), is_committed), + tmp_states_to_be_added_.end() + ); +} + +EarleyParser::EarleyParser(const Grammar& grammar, std::optional initial_state) + : grammar_(grammar), + fsm_state_flags_cache_(grammar->NumRules()), + rule_is_nullable_(grammar->NumRules(), 0) { + if (!grammar->optimized) { + XGRAMMAR_LOG(FATAL) << "The grammar is not optimized. Please optimize the grammar before using " + "the Earley parser."; + } + for (int32_t i = 0; i < grammar_->NumRules(); ++i) { + has_budget_rules_ = has_budget_rules_ || grammar_->GetRule(i).max_tokens >= 0; + has_char_budget_rules_ = has_char_budget_rules_ || grammar_->GetRule(i).max_chars >= 0; + if (has_budget_rules_ && has_char_budget_rules_) { + break; + } + } + for (int32_t i = 0; i < grammar_->NumRules(); ++i) { + const auto& rule = grammar_->GetRule(i); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(i); + capture_tracking_ = + capture_tracking_ || !rule.capture_name.empty() || + (suffix_stop_info != nullptr && !suffix_stop_info->stop_capture_name.empty()); + has_hidden_capture_rules_ = + has_hidden_capture_rules_ || + (suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0)); + if (capture_tracking_ && has_hidden_capture_rules_) { + break; + } + } + for (int32_t rule_id : grammar_->allow_empty_rule_ids) { + rule_is_nullable_[rule_id] = true; + } + PushStateAndExpand(initial_state.has_value() ? *initial_state : RootInitialState()); +} + +uint8_t EarleyParser::InitializeFsmStateFlags(int32_t rule_id, int32_t state_id) { + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[rule_id]->GetFsm(); + auto& flags_cache = fsm_state_flags_cache_[rule_id]; + if (flags_cache.empty()) { + flags_cache.resize(fsm.NumStates()); + } + XGRAMMAR_DCHECK(state_id >= 0 && state_id < static_cast(flags_cache.size())); + uint8_t& flags = flags_cache[state_id]; + if (flags & kFsmStateInitialized) { + return flags; + } + + flags = kFsmStateInitialized; + if (fsm.IsEndState(state_id)) { + flags |= kFsmStateEnd; + } + const auto& edges = fsm.GetFsm().GetEdges(state_id); + if (edges.size() != 0) { + flags |= kFsmStateHasEdges; + } + for (const auto& edge : edges) { + if (edge.IsCharRange() || edge.IsToken() || edge.IsExcludeToken()) { + flags |= kFsmStateScanable; + } else if (edge.IsRuleRef() || edge.IsEpsilon() || edge.IsRepeatRef()) { + flags |= kFsmStateNonTerminal; + } + } + return flags; +} + +ParserState EarleyParser::RootInitialState() const { + const auto root_rule_id = grammar_->GetRootRuleId(); + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[root_rule_id].has_value()); + return ParserState( + root_rule_id, + grammar_->GetRule(root_rule_id).body_expr_id, + grammar_->per_rule_fsms[root_rule_id]->GetFsm().GetStart(), + ParserState::kNoPrevInputPos, + DeadlineForRule(root_rule_id, -1), + 0, + 0, + 0, + ResolveActiveTemperatureRule(root_rule_id, -1), + CharDeadlineForRule(root_rule_id, -1) + ); +} + +void EarleyParser::PushStateAndExpand(const ParserState& state) { + tmp_states_visited_in_queue_.Clear(); + tmp_accept_stop_token_ = false; + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + Enqueue(state); + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = tmp_process_state_queue_.front(); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state); + if (completable) { + Complete(state); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_count_history_.push_back(GetCurrentCharIndex()); + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } +} + +void EarleyParser::Reset() { + rule_id_to_completable_states_.PopBack(rule_id_to_completable_states_.size()); + scanable_state_history_.PopBack(scanable_state_history_.size()); + is_completed_.clear(); + stop_token_is_accepted_ = false; + if (capture_tracking_) { + capture_event_history_.PopBack(capture_event_history_.size()); + } + char_count_history_.clear(); + char_budget_entry_history_.clear(); + tmp_char_budget_entered_ = false; + capture_recording_ = false; + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + PushStateAndExpand(RootInitialState()); +} + +void EarleyParser::ExpandNextRuleRefElement( + const ParserState& state, + const GrammarExpr& grammar_expr, + const GrammarExpr* sub_grammar_expr, + bool debug_print +) { + // Path A. The rule has a corresponding FSM. + XGRAMMAR_DCHECK(!(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value())); + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSequence); + XGRAMMAR_DCHECK( + sub_grammar_expr->type == GrammarExprType::kRuleRef || + sub_grammar_expr->type == GrammarExprType::kRepeat + ); + auto ref_rule_id = (*sub_grammar_expr)[0]; + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name << " predict the new rule " + << ref_rule_id << ": " << grammar_->GetRule(ref_rule_id).name << "."; + } + + bool right_recursion_to_root = false; + // The right-recursion optimization elides the completion of the parent rule (and, in the + // to-root case, corrupts the start position of the child rule), so it must be disabled when + // either rule produces capture-history events. + if (state.element_id != grammar_expr.size() - 1 || + sub_grammar_expr->type == GrammarExprType::kRepeat || + (state.rule_start_pos == rule_id_to_completable_states_.size() - 1) || + RuleNeedsCaptureEvent(state.rule_id) || RuleNeedsCaptureEvent(ref_rule_id)) { + // It's not the right recursion, or it's the root rule. + rule_id_to_completable_states_.PushBackInLatestRow(std::make_pair(ref_rule_id, state)); + } else { + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + right_recursion_to_root = true; + } else { + // If it's the right recursion, we need to add the ancestors of the parent state. + const auto in_vec = [&](const ParserState& state_) { + return std::find_if( + rule_id_to_completable_states_.Back().begin(), + rule_id_to_completable_states_.Back().end(), + [&](const auto& s) { + return StateEqualForParsing()(s.second, state_) && s.first == ref_rule_id; + } + ) != rule_id_to_completable_states_.Back().end(); + }; + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + std::vector> to_added_states; + for (const auto& parent_state_iter : parent_states_map) { + if (parent_state_iter.first != state.rule_id) continue; + const auto& parent_state = parent_state_iter.second; + if (!in_vec(parent_state)) { + to_added_states.push_back({ref_rule_id, parent_state}); + } + } + for (const auto& to_add_state : to_added_states) { + rule_id_to_completable_states_.PushBackInLatestRow(to_add_state); + } + } + } + + if (IsRuleNullable(ref_rule_id)) { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSequence); + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + state.element_id + 1, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + + // If the reference rule is not visited, we need to add it to the queue. + const auto& ref_rule = grammar_->GetRule(ref_rule_id); + if (ref_rule.max_chars >= 0) { + tmp_char_budget_entered_ = true; + } + const auto& ref_grammar_expr_id = ref_rule.body_expr_id; + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[ref_rule_id].has_value()); + const auto& ref_fsm = grammar_->per_rule_fsms[ref_rule_id].value(); + Enqueue(ParserState{ + ref_rule_id, + ref_grammar_expr_id, + ref_fsm.GetFsm().GetStart(), + right_recursion_to_root ? ParserState::kNoPrevInputPos + : int32_t(rule_id_to_completable_states_.size() - 1), + DeadlineForRule(ref_rule_id, state.budget_deadline), + 0, + 0, + 0, + ResolveActiveTemperatureRule(ref_rule_id, state.active_temperature_rule_id), + CharDeadlineForRule(ref_rule_id, state.char_budget_deadline) + }); +} + +void EarleyParser::ExpandNextRuleRefElementOnFSM(const ParserState& state, bool debug_print) { + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id].value(); + + // Add the rule reference pairs, and enqueue the epsilon edges. + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + if (edge.IsEpsilon()) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + edge.target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + continue; + } + + int target; + int ref_rule_id; + bool is_repeat = false; + RepeatEdgeRef repeat_info{nullptr}; + + if (edge.IsRuleRef()) { + target = edge.target; + ref_rule_id = edge.GetRefRuleId(); + } else if (edge.IsRepeatRef()) { + is_repeat = true; + repeat_info = grammar_->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + target = edge.target; + ref_rule_id = repeat_info.RuleId(); + + if (state.repeat_count >= repeat_info.Lower()) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + if (state.repeat_count >= repeat_info.Upper()) { + continue; + } + } else { + continue; + } + bool right_recursion_to_root = false; + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The rule " << state.rule_id << ": " + << grammar_->GetRule(state.rule_id).name << " predict the new rule " + << ref_rule_id << ": " << grammar_->GetRule(ref_rule_id).name << "."; + } + const uint8_t target_flags = GetFsmStateFlags(state.rule_id, target); + if (!is_repeat && !(target_flags & kFsmStateHasEdges) && (target_flags & kFsmStateEnd) && + state.rule_start_pos != static_cast(rule_id_to_completable_states_.size() - 1) && + !RuleNeedsCaptureEvent(state.rule_id) && !RuleNeedsCaptureEvent(ref_rule_id)) { + // It's a right recursion. We can optimize it. The optimization elides the completion of + // the parent rule, so it is disabled when either rule produces capture-history events. + // If it's the right recursion, we need to add the ancestors of the parent state. + if (state.rule_start_pos == ParserState::kNoPrevInputPos) { + // In this case, we can mark the new state as the root state to speed up. + right_recursion_to_root = true; + } else { + const auto in_vec = [&](const ParserState& state_) { + return std::find_if( + rule_id_to_completable_states_.Back().begin(), + rule_id_to_completable_states_.Back().end(), + [&](const auto& s) { + return StateEqualForParsing()(s.second, state_) && s.first == ref_rule_id; + } + ) != rule_id_to_completable_states_.Back().end(); + }; + const auto& parent_states_map = rule_id_to_completable_states_[state.rule_start_pos]; + std::vector> to_added_states; + for (const auto& parent_state_iter : parent_states_map) { + if (parent_state_iter.first != state.rule_id) continue; + const auto& parent_state = parent_state_iter.second; + if (!in_vec(parent_state)) { + to_added_states.push_back({ref_rule_id, parent_state}); + } + } + for (const auto& to_add_state : to_added_states) { + rule_id_to_completable_states_.PushBackInLatestRow(to_add_state); + } + } + } else { + if (is_repeat) { + // For kRepeatRef: store element_id = source state, preserve repeat_count + rule_id_to_completable_states_.PushBackInLatestRow( + {ref_rule_id, + ParserState{ + state.rule_id, + state.sequence_id, + state.element_id, + state.rule_start_pos, + state.budget_deadline, + 0, + state.repeat_count, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }} + ); + } else { + // For kRuleRef: store element_id = target (post-transition state) + rule_id_to_completable_states_.PushBackInLatestRow( + {ref_rule_id, + ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }} + ); + } + } + + // Check if the reference rule can be empty. + if (!is_repeat && IsRuleNullable(ref_rule_id)) { + Enqueue(ParserState{ + state.rule_id, + state.sequence_id, + target, + state.rule_start_pos, + state.budget_deadline, + 0, + 0, + 0, + state.active_temperature_rule_id, + state.char_budget_deadline + }); + } + + // If the reference rule is not visited, we need to add it to the queue. + const auto& ref_rule = grammar_->GetRule(ref_rule_id); + if (ref_rule.max_chars >= 0) { + tmp_char_budget_entered_ = true; + } + const auto& ref_grammar_expr_id = ref_rule.body_expr_id; + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[ref_rule_id].has_value()); + const auto& ref_fsm = grammar_->per_rule_fsms[ref_rule_id].value(); + Enqueue(ParserState{ + ref_rule_id, + ref_grammar_expr_id, + ref_fsm.GetFsm().GetStart(), + right_recursion_to_root ? ParserState::kNoPrevInputPos + : int32_t(rule_id_to_completable_states_.size() - 1), + DeadlineForRule(ref_rule_id, state.budget_deadline), + 0, + 0, + 0, + ResolveActiveTemperatureRule(ref_rule_id, state.active_temperature_rule_id), + CharDeadlineForRule(ref_rule_id, state.char_budget_deadline) + }); + } +} + +void EarleyParser::AdvanceByteString( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_rule +) { + XGRAMMAR_DCHECK(sub_rule.type == GrammarExprType::kByteString); + XGRAMMAR_DCHECK(sub_rule.size() > state.sub_element_id); + if (static_cast(sub_rule[state.sub_element_id]) == ch) { + auto new_state = state; + new_state.sub_element_id++; + if (new_state.sub_element_id == sub_rule.size()) { + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + // Assert: In a sequence, the bytestring can't be skipped. So the state can't be repeated. + } else { + tmp_states_to_be_added_.push_back(new_state); + } + } + return; +} + +void EarleyParser::AdvanceCharacterClass( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence +) { + XGRAMMAR_DCHECK(sub_sequence.type == GrammarExprType::kCharacterClass) + << "The element type is not supported!"; + + bool is_negative = static_cast(sub_sequence[0]); + + // The state is matching a UTF8 character (continuation bytes). + if (state.sub_element_id > 0) { + if ((ch & 0xC0) == 0x80) { + auto new_state = state; + new_state.sub_element_id--; + // Accumulate the codepoint from continuation byte + new_state.partial_codepoint = (new_state.partial_codepoint << 6) | (ch & 0x3F); + + // Check if the UTF8 character is completed. + if (new_state.sub_element_id == 0) { + if (is_negative) { + // For negative classes, accept if codepoint is NOT in any range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (!matches_range) { + new_state.element_id++; + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } else { + // For positive classes, accept if codepoint IS in a range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (matches_range) { + new_state.element_id++; + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } + } else { + // Check if partial codepoint could still potentially match any range + int32_t remaining_bytes = new_state.sub_element_id; + int32_t min_codepoint = new_state.partial_codepoint << (6 * remaining_bytes); + int32_t max_codepoint = min_codepoint | ((1 << (6 * remaining_bytes)) - 1); + + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: always continue (will verify on final byte) + // For positive classes: only continue if some range could match + bool should_continue = is_negative ? true : could_match; + if (should_continue) { + tmp_states_to_be_added_.push_back(new_state); + } + } + } + return; + } + + // Handle non-ASCII first bytes + if (!isascii(ch)) { + auto [accepted, num_bytes, partial] = HandleUTF8FirstByte(ch); + if (!accepted) { + return; + } + + XGRAMMAR_DCHECK(num_bytes > 1); + + // Compute possible codepoint range for this first byte + int32_t min_codepoint = partial << (6 * (num_bytes - 1)); + int32_t max_codepoint = min_codepoint | ((1 << (6 * (num_bytes - 1))) - 1); + + // Check if any stored range could potentially match + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + // Check for overlap between [min_codepoint, max_codepoint] and [lower, upper] + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: accept if no range could match (will verify on final byte) + // For positive classes: accept if some range could match (will verify on final byte) + bool should_continue = is_negative ? true : could_match; + + if (should_continue) { + auto new_state = state; + new_state.sub_element_id = num_bytes - 1; + new_state.partial_codepoint = partial; + tmp_states_to_be_added_.push_back(new_state); + } + return; + } + + // ASCII handling (unchanged) + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (static_cast(sub_sequence[i]) <= ch && + ch <= static_cast(sub_sequence[i + 1])) { + if (!is_negative) { + auto new_state = state; + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + } + return; + } + } + if (is_negative) { + auto new_state = state; + new_state.element_id++; + new_state.sub_element_id = 0; + Enqueue(new_state); + } +} + +void EarleyParser::AdvanceCharacterClassStar( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence +) { + XGRAMMAR_DCHECK(sub_sequence.type == GrammarExprType::kCharacterClassStar) + << "The element type is not supported!"; + + bool is_negative = static_cast(sub_sequence[0]); + + // The state is matching a UTF8 character (continuation bytes). + if (state.sub_element_id > 0) { + if ((ch & 0xC0) == 0x80) { + auto new_state = state; + new_state.sub_element_id--; + // Accumulate the codepoint from continuation byte + new_state.partial_codepoint = (new_state.partial_codepoint << 6) | (ch & 0x3F); + + // Check if the UTF8 character is completed. + if (new_state.sub_element_id == 0) { + if (is_negative) { + // For negative classes, accept if codepoint is NOT in any range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (!matches_range) { + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } else { + // For positive classes, accept if codepoint IS in a range + bool matches_range = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (new_state.partial_codepoint >= sub_sequence[i] && + new_state.partial_codepoint <= sub_sequence[i + 1]) { + matches_range = true; + break; + } + } + if (matches_range) { + new_state.partial_codepoint = 0; + Enqueue(new_state); + } + } + } else { + // Check if partial codepoint could still potentially match any range + int32_t remaining_bytes = new_state.sub_element_id; + int32_t min_codepoint = new_state.partial_codepoint << (6 * remaining_bytes); + int32_t max_codepoint = min_codepoint | ((1 << (6 * remaining_bytes)) - 1); + + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: always continue (will verify on final byte) + // For positive classes: only continue if some range could match + bool should_continue = is_negative ? true : could_match; + if (should_continue) { + tmp_states_to_be_added_.push_back(new_state); + } + } + } + return; + } + + // Handle non-ASCII first bytes + if (!isascii(ch)) { + auto [accepted, num_bytes, partial] = HandleUTF8FirstByte(ch); + if (!accepted) { + return; + } + + XGRAMMAR_DCHECK(num_bytes > 1); + + // Compute possible codepoint range for this first byte + int32_t min_codepoint = partial << (6 * (num_bytes - 1)); + int32_t max_codepoint = min_codepoint | ((1 << (6 * (num_bytes - 1))) - 1); + + // Check if any stored range could potentially match + bool could_match = false; + for (int i = 1; i < sub_sequence.size(); i += 2) { + int32_t lower = sub_sequence[i]; + int32_t upper = sub_sequence[i + 1]; + // Check for overlap between [min_codepoint, max_codepoint] and [lower, upper] + if (max_codepoint >= lower && min_codepoint <= upper) { + could_match = true; + break; + } + } + + // For negative classes: accept if no range could match (will verify on final byte) + // For positive classes: accept if some range could match (will verify on final byte) + bool should_continue = is_negative ? true : could_match; + + if (should_continue) { + auto new_state = state; + new_state.sub_element_id = num_bytes - 1; + new_state.partial_codepoint = partial; + tmp_states_to_be_added_.push_back(new_state); + } + return; + } + + // ASCII handling (unchanged) + for (int i = 1; i < sub_sequence.size(); i += 2) { + if (static_cast(sub_sequence[i]) <= ch && + ch <= static_cast(sub_sequence[i + 1])) { + if (!is_negative) { + Enqueue(state); + } + return; + } + } + if (is_negative) { + Enqueue(state); + } +} + +void EarleyParser::AdvanceFsm(const ParserState& state, const uint8_t ch) { + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& current_fsm = grammar_->per_rule_fsms[state.rule_id].value(); + for (const auto& edge : current_fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + if ((!edge.IsCharRange()) || ch < edge.min || ch > edge.max) { + continue; + } + auto new_state = state; + new_state.element_id = edge.target; + const uint8_t flags = GetFsmStateFlags(state.rule_id, edge.target); + if (!(flags & kFsmStateNonTerminal) && !(flags & kFsmStateEnd) && (flags & kFsmStateScanable)) { + EnqueueWithoutProcessing(std::move(new_state)); + } else { + Enqueue(std::move(new_state)); + } + } +} + +void EarleyParser::ScanAtomicToken(const ParserState& state, int32_t token_id) { + if (state.rule_id == -1) return; + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& current_fsm = grammar_->per_rule_fsms[state.rule_id].value(); + for (const auto& edge : current_fsm.GetFsm().GetFsm().GetEdges(state.element_id)) { + bool matched = false; + if (edge.IsToken()) { + auto info = current_fsm.GetFsm().GetFsm().GetTokenEdgeInfo(edge.GetAuxIndex()); + matched = info.Contains(token_id); + } else if (edge.IsExcludeToken()) { + auto info = current_fsm.GetFsm().GetFsm().GetExcludeTokenEdgeInfo(edge.GetAuxIndex()); + matched = info.Accepts(token_id); + } + if (!matched) continue; + auto new_state = state; + new_state.element_id = edge.target; + const uint8_t flags = GetFsmStateFlags(state.rule_id, edge.target); + if (!(flags & kFsmStateNonTerminal) && !(flags & kFsmStateEnd) && (flags & kFsmStateScanable)) { + EnqueueWithoutProcessing(std::move(new_state)); + } else { + Enqueue(std::move(new_state)); + } + } +} + +bool EarleyParser::AdvanceAtomicToken( + int32_t token_id, bool debug_print, int32_t token_char_count +) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()) + << "The tmp_process_state_queue_ should be empty before AdvanceAtomicToken."; + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_accept_stop_token_ = false; + tmp_completed_lazy_occurrences_.clear(); + if (has_char_budget_rules_) { + tmp_char_budget_entered_ = char_budget_entry_history_.back(); + char_count_history_.push_back(GetCurrentCharIndex() + token_char_count); + } + const auto& latest_states = scanable_state_history_[scanable_state_history_.size() - 1]; + for (const auto& state : latest_states) { + if (skip_expired_states_ && IsExpiredState(state)) { + continue; + } + ScanAtomicToken(state, token_id); + } + if (tmp_process_state_queue_.empty() && tmp_states_to_be_added_.empty()) { + if (has_char_budget_rules_) { + char_count_history_.pop_back(); + } + return false; + } + rule_id_to_completable_states_.PushBack(std::vector>()); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + is_completed_.push_back(tmp_accept_stop_token_); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + if (has_char_budget_rules_) { + char_budget_entry_history_.push_back(tmp_char_budget_entered_); + } + return true; +} + +bool RepeatDetector::IsVisited(const ParserState& state) const { + // If the size is larger than the threshold, then we use the set to check. + if (size_ > transition_threshold_) { + return visited_set_.find(state) != visited_set_.end(); + } + return std::find_if( + visited_vector_.begin(), + visited_vector_.begin() + size_, + [&state](const ParserState& s) { return StateEqualForParsing()(state, s); } + ) != visited_vector_.begin() + size_; +} + +void RepeatDetector::Insert(const ParserState& state) { + if (size_ == transition_threshold_) { + for (const auto& s : visited_vector_) { + visited_set_.insert(s); + } + } + size_++; + if (size_ > transition_threshold_) { + visited_set_.insert(state); + } else { + visited_vector_[size_ - 1] = state; + } +} + +void RepeatDetector::Clear() { + if (size_ > transition_threshold_) { + visited_set_.clear(); + } + size_ = 0; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/earley_parser.h b/zig/vendor/xgrammar/cpp/earley_parser.h new file mode 100644 index 0000000000..a3ef3e71ba --- /dev/null +++ b/zig/vendor/xgrammar/cpp/earley_parser.h @@ -0,0 +1,802 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/earley_parser.h + * \brief The header for the definition of the Earley parser. + */ + +#ifndef XGRAMMAR_EARLEY_PARSER_H_ +#define XGRAMMAR_EARLEY_PARSER_H_ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "support/compact_2d_array.h" +#include "support/utils.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief The state of the Earley parser. + * In the implementation, a rule can only be a kchoices or a ktagdispatch. + * A kchoices rule must be composed of some ksequence rules, or a kemptyrule. + * In the ksequence, every element in the sequence must be a kbytestring, a + * kcharacterclass, a kcharacterclassstar, or a rule reference. + * + * - rule_id: The id of the rule. + * - sequence_id: The id of the sequence in the rule. + * - element_id: The id of the element in the sequence, or the id of the node in + * the tag dispatch fsm. + * - rule_start_pos: The id of the parent node in the Earley parser. i.e. the rule + * is predicted from the k-th character. + * - sub_element_id: The id of the sub element in the current element, i.e.: + * - kbytestring: the id of the byte in the string. + * - kcharacterclass: How many bytes are left to be read in the utf8 character. + * - kcharacterclassstar: How many bytes are left to be read in the utf8 character. + */ +struct ParserState { + constexpr ParserState() = default; + + constexpr ParserState( + const int32_t& rule_id, + const int32_t& sequence_id, + const int32_t& element_id, + const int32_t& rule_start_pos, + const int32_t& budget_deadline = -1, + const int32_t& sub_element_id = 0, + const int32_t& repeat_count = 0, + const int32_t& partial_codepoint = 0, + const int32_t& active_temperature_rule_id = -1, + const int32_t& char_budget_deadline = -1 + ) + : rule_id(rule_id), + sequence_id(sequence_id), + element_id(element_id), + rule_start_pos(rule_start_pos), + budget_deadline(budget_deadline), + sub_element_id(sub_element_id), + repeat_count(repeat_count), + partial_codepoint(partial_codepoint), + active_temperature_rule_id(active_temperature_rule_id), + char_budget_deadline(char_budget_deadline) {} + + /*! + * \brief A rule_start_pos value of kNoPrevInputPos means this ParserState is the root of the + * parsing stack. + */ + static constexpr int32_t kNoPrevInputPos = -1; + + /*! \brief The rule's id. */ + int32_t rule_id = -1; + + /*! \brief Which choice in this rule is selected. */ + int32_t sequence_id = -1; + + /*! + * \brief Which element of the choice sequence is to be visited. When the current sequence is + * a tag dispatch rule, this element id is the current node. + */ + int32_t element_id = -1; + + /*! \brief The position of the state, i.e. from which position, the rule starts. */ + int32_t rule_start_pos = -1; + + /*! \brief The last token index this state's derivation may consume, from the token budget + * (Rule::max_tokens) of the rule it is inside; -1 means unlimited. Set when a budgeted rule + * is predicted and inherited by the states inside it. */ + int32_t budget_deadline = -1; + + /*! \brief The id of the sub element in the current element of the sequence. */ + int32_t sub_element_id = 0; + + /*! \brief The number of times the element is repeated. It will be used in kRepeat.*/ + int32_t repeat_count = 0; + + /*! \brief Partial codepoint accumulated during UTF-8 decoding for positive character classes. */ + int32_t partial_codepoint = 0; + + /*! \brief The innermost active rule that specifies a sampling temperature. */ + int32_t active_temperature_rule_id = -1; + + /*! \brief The number of Unicode codepoints this derivation may consume before its active + * character budget expires; -1 means unlimited. Stored as an absolute input position. */ + int32_t char_budget_deadline = -1; + + /*! + * \brief Lexicographic order over all fields. It is only used to sort the states for + * deterministic serialization, and is not needed during parsing. + */ + bool operator<(const ParserState& other) const { + if (rule_id != other.rule_id) return rule_id < other.rule_id; + if (sequence_id != other.sequence_id) return sequence_id < other.sequence_id; + if (element_id != other.element_id) return element_id < other.element_id; + if (rule_start_pos != other.rule_start_pos) return rule_start_pos < other.rule_start_pos; + if (budget_deadline != other.budget_deadline) return budget_deadline < other.budget_deadline; + if (sub_element_id != other.sub_element_id) return sub_element_id < other.sub_element_id; + if (repeat_count != other.repeat_count) return repeat_count < other.repeat_count; + if (partial_codepoint != other.partial_codepoint) { + return partial_codepoint < other.partial_codepoint; + } + if (active_temperature_rule_id != other.active_temperature_rule_id) { + return active_temperature_rule_id < other.active_temperature_rule_id; + } + return char_budget_deadline < other.char_budget_deadline; + } + + friend std::ostream& operator<<(std::ostream& os, const ParserState& state) { + os << state.ToString(); + return os; + } + + std::string ToString() const { + std::string result = "ParserState(rule_id=" + std::to_string(rule_id) + + ", sequence_id=" + std::to_string(sequence_id) + + ", element_id=" + std::to_string(element_id) + + ", rule_start_pos=" + std::to_string(rule_start_pos) + + ", sub_element_id=" + std::to_string(sub_element_id); + if (repeat_count != 0) { + result += ", repeat_count=" + std::to_string(repeat_count); + } + if (partial_codepoint != 0) { + result += ", partial_codepoint=" + std::to_string(partial_codepoint); + } + if (budget_deadline != -1) { + result += ", budget_deadline=" + std::to_string(budget_deadline); + } + if (active_temperature_rule_id != -1) { + result += ", active_temperature_rule_id=" + std::to_string(active_temperature_rule_id); + } + if (char_budget_deadline != -1) { + result += ", char_budget_deadline=" + std::to_string(char_budget_deadline); + } + result += ")"; + return result; + } +}; + +XGRAMMAR_MEMBER_ARRAY( + ParserState, + &ParserState::rule_id, + &ParserState::sequence_id, + &ParserState::element_id, + &ParserState::rule_start_pos, + &ParserState::budget_deadline, + &ParserState::sub_element_id, + &ParserState::repeat_count, + &ParserState::partial_codepoint, + &ParserState::active_temperature_rule_id, + &ParserState::char_budget_deadline +); + +/*! + * \brief Hash of a state used as the key of the adaptive token mask cache. The token mask of a + * state does not depend on rule_start_pos, repeat_count or partial_codepoint, so they are + * ignored. Pairs with StateEqualForCache. + */ +class StateHashForCache { + public: + size_t operator()(const ParserState& state) const { + return HashCombine(state.rule_id, state.sequence_id, state.element_id, state.sub_element_id); + } +}; + +/*! + * \brief Equality of states used as the key of the adaptive token mask cache. Compares the same + * fields as StateHashForCache hashes. + */ +class StateEqualForCache { + public: + bool operator()(const ParserState& lhs, const ParserState& rhs) const { + return lhs.rule_id == rhs.rule_id && lhs.sequence_id == rhs.sequence_id && + lhs.element_id == rhs.element_id && lhs.sub_element_id == rhs.sub_element_id; + } +}; + +/*! + * \brief When matching the state, we need to consider the rule_start_pos, since if two states + * don't have the same rule_start_pos, they are not the same state. + */ +class StateEqualForParsing { + public: + bool operator()(const ParserState& lhs, const ParserState& rhs) const { + return lhs.rule_id == rhs.rule_id && lhs.sequence_id == rhs.sequence_id && + lhs.element_id == rhs.element_id && lhs.rule_start_pos == rhs.rule_start_pos && + lhs.sub_element_id == rhs.sub_element_id && lhs.repeat_count == rhs.repeat_count && + lhs.partial_codepoint == rhs.partial_codepoint && + lhs.budget_deadline == rhs.budget_deadline && + lhs.active_temperature_rule_id == rhs.active_temperature_rule_id && + lhs.char_budget_deadline == rhs.char_budget_deadline; + } +}; + +/*! + * \brief This class is used to hash the ParserState for parsing. + * If two ParserStates don't have the same rule_start_pos, they are not the same state. + */ +class StateHashForParsing { + public: + size_t operator()(const ParserState& state) const { + return HashCombine( + state.rule_id, + state.sequence_id, + state.element_id, + state.rule_start_pos, + state.sub_element_id, + state.repeat_count, + state.partial_codepoint, + state.budget_deadline, + state.active_temperature_rule_id, + state.char_budget_deadline + ); + } +}; + +/*! \brief This class is used to detect the repeated states. */ +class RepeatDetector { + private: + const int transition_threshold_; + + std::vector visited_vector_; + + std::unordered_set visited_set_; + + int size_ = 0; + + public: + RepeatDetector(const int transition_threshold = 50) + : transition_threshold_(transition_threshold), size_(0) { + visited_vector_.resize(transition_threshold_); + } + + /*! + * \brief Check if the element is visited. + * \return True if visited, false otherwise. + */ + bool IsVisited(const ParserState& state) const; + + /*! + * \brief Add the state into the visited states. + * \param state The state to be added. + */ + void Insert(const ParserState& state); + + /*! \brief Reset the detector. */ + void Clear(); +}; + +/*! \brief A concrete occurrence of a captured rule in an Earley parent chain. */ +struct CaptureOccurrence { + /*! \brief The id of the captured rule. */ + int32_t rule_id; + /*! \brief The position where the rule occurrence started. */ + int32_t start_pos; + + bool operator==(const CaptureOccurrence& other) const { + return rule_id == other.rule_id && start_pos == other.start_pos; + } +}; + +/*! + * \brief A completion event of a captured rule, recorded when the rule is completed during + * parsing. The matched span is [start_pos, r) in input positions, where r is the position (i.e. + * the history row) at which the event is recorded. + */ +struct CaptureEvent { + /*! \brief The id of the completed rule. */ + int32_t rule_id; + /*! \brief The position where the rule started matching. kNoPrevInputPos means position 0 (the + * rule acts as the root). */ + int32_t start_pos; + /*! \brief The unadjusted start position of this rule occurrence. This differs from start_pos + * for the zero-width event inserted after a dynamic-dispatch marker. */ + int32_t occurrence_start_pos; + /*! \brief Number of trailing bytes hidden only from this rule's own capture for this + * completion. */ + int32_t hidden_suffix_bytes = 0; + /*! \brief Number of trailing bytes hidden from every containing capture for this completion. */ + int32_t hidden_stop_bytes = 0; + /*! \brief The captured rule occurrences whose concrete Earley parent chains contain this stop + * completion. Includes this rule's occurrence when the rule itself is captured. */ + std::vector stop_capture_targets; +}; + +class EarleyParser { + /*! + * \brief Here is an article about Earley Parser. + * https://en.wikipedia.org/wiki/Earley_parser#Pseudocode + * We divide the parser states into three categories: + * - Scanable (which will be stored in scanable_state_history_). + * - Predictable(If it predict a new rule successfully, then it will be stored in + * rule_id_to_completable_states). + * - completable(which can perform a completion operation). + * A state will be stored in rule_id_to_completable_states_ if it can be completed, + * and it will be stored in scanable_state_history_ if it can be scanned. Otherwise, + * it will be discarded. + */ + protected: + using GrammarExpr = Grammar::Impl::GrammarExpr; + + /*! \brief The grammar to be parsed. */ + Grammar grammar_; + + /*! \brief In this round of advancing, check if the stop token can be accepted. */ + bool tmp_accept_stop_token_ = false; + + /*! \brief store when accepting i characters, if the stop token can be accepted. */ + std::vector is_completed_; + + /*! + * \brief rule_id_to_completable_states[i][j] is the i pos j rule_id states. Earley + * parser needs it to complete. + */ + Compact2DArray> rule_id_to_completable_states_; + + /*! + * \brief The states history. state_stack[i] is a vector storing the states after accepting the + * input[i-1]. + */ + Compact2DArray scanable_state_history_; + + /*! + * \brief A temporary vector only used in Advance, used to add states in the + * scanable_state_history. + */ + std::vector tmp_states_to_be_added_; + + /*! \brief It's the processing queue of the earley parser. */ + std::queue tmp_process_state_queue_; + + /*! \brief The class is used to check if a state has been added into the queue. */ + RepeatDetector tmp_states_visited_in_queue_; + + /*! \brief Check if the stop token is accepted. */ + bool stop_token_is_accepted_ = false; + + enum FsmStateFlag : uint8_t { + kFsmStateInitialized = 1 << 0, + kFsmStateScanable = 1 << 1, + kFsmStateNonTerminal = 1 << 2, + kFsmStateEnd = 1 << 3, + kFsmStateHasEdges = 1 << 4, + }; + + /*! \brief Lazily-computed FSM state properties, indexed by rule id and state id. */ + std::vector> fsm_state_flags_cache_; + + /*! \brief Whether each rule can match the empty string. */ + std::vector rule_is_nullable_; + + /*! \brief Compute and cache properties for a state in a per-rule FSM. */ + uint8_t InitializeFsmStateFlags(int32_t rule_id, int32_t state_id); + + /*! \brief Return cached properties for a state in a per-rule FSM. */ + uint8_t GetFsmStateFlags(int32_t rule_id, int32_t state_id) { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(fsm_state_flags_cache_.size())); + auto& flags_cache = fsm_state_flags_cache_[rule_id]; + if (!flags_cache.empty()) { + XGRAMMAR_DCHECK(state_id >= 0 && state_id < static_cast(flags_cache.size())); + if (flags_cache[state_id] != 0) { + return flags_cache[state_id]; + } + } + return InitializeFsmStateFlags(rule_id, state_id); + } + + bool IsRuleNullable(int32_t rule_id) const { return rule_is_nullable_[rule_id] != 0; } + + /*! \brief The index of the LLM token currently being accepted, set by the matcher; -1 + * before any token. budget_deadline values are compared against it. */ + int32_t current_token_index_ = -1; + + /*! \brief Whether states past their budget deadline are skipped when scanning. Enabled by + * the matcher for accepts that follow an enforcing mask computation. */ + bool skip_expired_states_ = false; + + /*! \brief Whether any rule of the grammar has a token budget. */ + bool has_budget_rules_ = false; + + /*! \brief The number of Unicode codepoints accepted at every parser history row. */ + std::vector char_count_history_; + + /*! \brief Whether any rule of the grammar has a character budget. */ + bool has_char_budget_rules_ = false; + + /*! \brief Whether a character-budgeted occurrence was entered since the initial parser row. */ + std::vector char_budget_entry_history_; + + /*! \brief Entry-history value for the row currently being expanded. */ + bool tmp_char_budget_entered_ = false; + + /*! \brief Whether the state's derivation may not consume the next token. */ + bool IsExpiredState(const ParserState& state) const { + return state.budget_deadline >= 0 && current_token_index_ > state.budget_deadline; + } + + /*! \brief The deadline for a newly predicted occurrence of the rule: its own budget counted + * from the current token, capped by the parent's deadline for nested budgets. */ + int32_t DeadlineForRule(int32_t rule_id, int32_t parent_deadline) const { + int32_t own = grammar_->GetRule(rule_id).max_tokens; + if (own < 0) { + return parent_deadline; + } + int64_t uncapped_deadline = static_cast(current_token_index_) + own; + int32_t deadline = static_cast( + std::min(uncapped_deadline, std::numeric_limits::max()) + ); + return parent_deadline >= 0 ? std::min(deadline, parent_deadline) : deadline; + } + + /*! \brief The character deadline for a newly predicted rule occurrence. */ + int32_t CharDeadlineForRule(int32_t rule_id, int32_t parent_deadline) const { + int32_t own = grammar_->GetRule(rule_id).max_chars; + if (own < 0) { + return parent_deadline; + } + int32_t current_char_index = GetCurrentCharIndex(); + int32_t deadline = own > std::numeric_limits::max() - current_char_index + ? std::numeric_limits::max() + : current_char_index + own; + return parent_deadline >= 0 ? std::min(deadline, parent_deadline) : deadline; + } + + /*! \brief Whether the state's derivation may not consume another Unicode codepoint. */ + bool IsCharExpiredState(const ParserState& state) const { + return state.char_budget_deadline >= 0 && GetCurrentCharIndex() >= state.char_budget_deadline; + } + + static bool StartsUTF8Codepoint(uint8_t byte) { return (byte & 0xC0) != 0x80; } + + /*! \brief Whether any rule of the grammar has a capture or stop_capture name. Fixed at + * construction. When false, the capture machinery is fully disabled and has no overhead. */ + bool capture_tracking_ = false; + + /*! \brief Whether the grammar contains suffix/stop spans that may affect captures. */ + bool has_hidden_capture_rules_ = false; + + /*! + * \brief Whether capture events are currently recorded in Complete(). Only enabled during + * definitive advances (accepting a token or string), not during speculative exploration + * (mask computation, jump-forward search, lookahead checks), so that speculative completions + * never produce capture events. + */ + bool capture_recording_ = false; + + /*! + * \brief The history of capture events. capture_event_history_[i] stores the events recorded + * when input position i was created. Kept aligned with scanable_state_history_ row-by-row + * whenever capture_tracking_ is true, so PopLastStates rolls back events automatically. + */ + Compact2DArray capture_event_history_; + + /*! \brief Returns true if the rule exists and has a capture name. */ + bool RuleHasCapture(int32_t rule_id) const { + return capture_tracking_ && rule_id >= 0 && !grammar_->GetRule(rule_id).capture_name.empty(); + } + + /*! \brief Returns true if completing this rule can hide bytes from a capture. */ + bool RuleHasHiddenBytes(int32_t rule_id) const { + if (!capture_tracking_ || !has_hidden_capture_rules_ || rule_id < 0) { + return false; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(rule_id); + return suffix_stop_info != nullptr && + (suffix_stop_info->hidden_suffix_bytes > 0 || suffix_stop_info->hidden_stop_bytes > 0); + } + + /*! \brief Returns true if completing this rule must produce a capture-history event. */ + bool RuleNeedsCaptureEvent(int32_t rule_id) const { + return RuleHasCapture(rule_id) || RuleHasHiddenBytes(rule_id); + } + + /*! \brief Record a capture or hidden-span event for a completed rule in the current row. */ + void RecordCaptureEvent(const ParserState& state, bool marker_present); + + /*! + * \brief Whether this completion of a suffix/stop rule actually consumed the trailing marker. + * A non-looping TagDispatch can also complete before its trigger is encountered, which is what + * lets a free-text tail end normally. Only the terminal post-dispatch state has no outgoing + * edges; completions in the trigger-scanning states did not consume a suffix/stop marker. + */ + bool CompletionConsumedMarker(const ParserState& state) const; + + /*! \brief Collect the captured rule occurrences whose concrete Earley parent chains contain + * the given stop completion. Includes the completed rule's own occurrence when captured. */ + std::vector CollectStopCaptureTargets(const ParserState& state) const; + + /*! + * \brief The lazy rule occurrences (rule_id, rule_start_pos) completed while building the + * current row. Committed-shortest matching: their remaining states are removed when the row is + * finalized, so the occurrence cannot be extended further. + */ + std::vector> tmp_completed_lazy_occurrences_; + + /*! \brief Remove the states of the lazy occurrences completed in the current row. */ + void RemoveCommittedLazyStates(); + + /*! + * \brief Check if the state has been added into the queue. + * \param state The state to check. + * \return True if in the vector, false otherwise. + */ + bool IsStateVisitedInQueue(const ParserState& state) const { + return tmp_states_visited_in_queue_.IsVisited(state); + } + + /*! + * \brief The scanning operation of the Earley parser. Put the new states in the queue. + */ + void Scan(const ParserState& state, const uint8_t ch); + + /*! + * \brief The completion operation of the Earley parser. + * \param state The state to be completed. + * \param debug_print Whether to print the debug information. + * \details The reason is that if the state can't be scanned, then + * add it into the next states is useless. Moreover, the end + * of the grammar is used to check if the grammar is completed, + * so it should be added into the next states. + */ + void Complete(const ParserState& state, bool debug_print = false, bool marker_present = true); + + /*! + * \brief The prediction operation of the Earley parser. + * \param state The state to be predicted. + * \param debug_print Whether to print the debug information. + * \return First: If the state scanable, or the state is the end of the grammar, + * then return true, otherwise return false. + * \return Second: If the state is completable, then return true, otherwise return false. + */ + std::pair Predict(const ParserState& state, bool debug_print = false); + + /*! \brief The initial state expanded from the root rule of the grammar. */ + ParserState RootInitialState() const; + + /*! \brief Resolve the active temperature rule when entering a rule. */ + int32_t ResolveActiveTemperatureRule(int32_t rule_id, int32_t inherited_rule_id) const; + + /*! + * \brief Expand the rule, used for RuleRef and kTagDispatch. + * \param state The state to be expanded, which is the parent state. + * The type of the state is kTagDispatch or kSequence. Moreover, the + * element of the sequence should be a rule reference; the node in + * the kTagDispatch should be an end node. + * \param grammar_expr The grammar expression to be expanded. + * \param sub_grammar_expr The sub grammar expression to be expanded, especially + * when the rule is a kSequence, and the sub rule is a kRuleRef. + * \param debug_print Whether to print the debug information. + */ + void ExpandNextRuleRefElement( + const ParserState& state, + const GrammarExpr& grammar_expr, + const GrammarExpr* sub_grammar_expr, + bool debug_print = false + ); + + /*! + * \brief Expand the rule, used for RuleRef and kTagDispatch. + * \param state The state to be expanded, and it's should be on the FSM. + * \param debug_print Whether to print the debug information. + */ + void ExpandNextRuleRefElementOnFSM(const ParserState& state, bool debug_print = false); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kCharacterClass. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceCharacterClass( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kByteString. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceByteString( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sub sequence is kCharacterClassStar. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \param sub_sequence The sub sequence to be checked. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceCharacterClassStar( + const ParserState& state, const uint8_t ch, const GrammarExpr& sub_sequence + ); + + /*! + * \brief Advance the parser to the next state, with the sequence is kTagDispatch. + * \param state The state to be advanced. + * \param ch The character to be advanced. + * \note The advanced states are enqueued; nothing is enqueued if the character is not accepted. + */ + void AdvanceFsm(const ParserState& state, const uint8_t ch); + + /*! + * \brief Scan a token edge: check if token_id matches any kToken or kExcludeToken edge from + * state. + */ + void ScanAtomicToken(const ParserState& state, int32_t token_id); + + /*! + * \brief Advance the parser by accepting a whole token via kToken/kExcludeToken edges. + * \param token_id The token ID to accept. + * \param debug_print Whether to print debug info. + * \return True if any state advanced, false otherwise. + */ + bool AdvanceAtomicToken(int32_t token_id, bool debug_print = false, int32_t token_char_count = 0); + + /*! + * \brief Enqueue the state into the queue. + * \param state The state to be enqueued. + * \details The state is enqueued if it is not visited in the queue. + */ + void Enqueue(const ParserState& state) { + if (!IsStateVisitedInQueue(state)) { + tmp_process_state_queue_.push(state); + tmp_states_visited_in_queue_.Insert(state); + } + } + + /*! + * \brief Enqueue the state into the queue, without prediction and completion. + * \param state The state to be enqueued. + */ + void EnqueueWithoutProcessing(const ParserState& state) { + if (!IsStateVisitedInQueue(state)) { + tmp_states_visited_in_queue_.Insert(state); + tmp_states_to_be_added_.push_back(state); + } + } + + public: + /*! + * \brief Constructor of the Earley parser. + * \param grammar The grammar to be parsed. It must be optimized. + * \param initial_state The state to start parsing from. If not provided, parsing starts + * from the root rule of the grammar. + */ + explicit EarleyParser( + const Grammar& grammar, std::optional initial_state = std::nullopt + ); + + /*! + * \brief From the current states, advance to the next state. + * \param ch The character to be advanced. + * \param debug_print Whether to print the debug information. + * \return True if the character is accepted, false otherwise. + * \note If the character isn't accepted, then the states won't be changed. + */ + bool Advance(const uint8_t ch, bool debug_print = false); + + /*! + * \brief Remove the newly added states. + * \param count The number of states to be removed. + */ + void PopLastStates(int32_t count = 1); + + /*! + * \brief Check whether any of the multiple states stored in the parser has already completed. + * \note Since the parser contains multiple parallel states, some may have already completed, + * while others might still be able to accept more characters. + * \return True if the root rule is completed, false otherwise. + */ + bool IsCompleted() const; + + /*! + * \brief Push the initial state into the Earley parser. + * \param state The initial state to be pushed. + */ + void PushStateAndExpand(const ParserState& state); + + /*! + * \brief Reset the parser. + * \note This function is used to reset the parser, and initialize the + * parser with the root rule. + */ + void Reset(); + + /*! + * \brief Get the current scanable states. + * \return The scanable states. + */ + std::vector GetLatestScanableStates() const { + std::vector latest_states; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + latest_states.push_back(state); + } + return latest_states; + } + + /*! + * \brief Push one state to check if it can accept the token. + * \param state The state to be pushed. + */ + void PushOneStateToCheck(const ParserState& state) { + PushStatesToCheck(std::vector{state}, is_completed_.back()); + } + + /*! \brief Push a temporary parser row for token-mask checking. */ + void PushStatesToCheck(const std::vector& states, bool completed) { + rule_id_to_completable_states_.PushBack(std::vector>()); + is_completed_.push_back(completed); + scanable_state_history_.PushBack(states); + if (capture_tracking_) { + capture_event_history_.PushBack(std::vector()); + } + if (has_char_budget_rules_) { + char_count_history_.push_back(GetCurrentCharIndex()); + char_budget_entry_history_.push_back(char_budget_entry_history_.back()); + } + } + + /*! \brief Push a character-count row for a parser row created by the matcher. */ + void PushCharCountRow(int32_t char_count, bool char_budget_entered) { + if (!has_char_budget_rules_) { + return; + } + char_count_history_.push_back(char_count); + char_budget_entry_history_.push_back(char_budget_entered); + } + + int32_t GetCurrentCharIndex() const { + return char_count_history_.empty() ? 0 : char_count_history_.back(); + } + + bool HasEnteredCharBudget() const { + return has_char_budget_rules_ && char_budget_entry_history_.back(); + } + + /*! \brief Whether the grammar has any captured rule. */ + bool IsCaptureTrackingEnabled() const { return capture_tracking_; } + + /*! \brief Copy the capture events of the latest input position. */ + std::vector CopyLastCaptureRow() const { + if (!capture_tracking_) { + return {}; + } + auto row = capture_event_history_[capture_event_history_.size() - 1]; + return std::vector(row.begin(), row.end()); + } + + /*! + * \brief Push a new row of capture events. Used when a new input position is created outside + * of Advance / AdvanceAtomicToken (e.g. when merging parallel advance results), to keep the + * capture history aligned with the state history. + */ + void PushCaptureRow(const std::vector& events) { + if (capture_tracking_) { + capture_event_history_.PushBack(events); + } + } + + std::string PrintStates() const { + std::string result; + result += "There are " + std::to_string(scanable_state_history_.size()) + + " steps in history. Last step: [\n"; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + result += state.ToString() + ", \n"; + } + result += "]"; + return result; + } +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_EARLEY_PARSER_H_ diff --git a/zig/vendor/xgrammar/cpp/ebnf_script_creator.h b/zig/vendor/xgrammar/cpp/ebnf_script_creator.h new file mode 100644 index 0000000000..5ffd06c047 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/ebnf_script_creator.h @@ -0,0 +1,188 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/ebnf_script_creator.h + * \brief The header for the creating EBNF script. + */ + +#ifndef XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ +#define XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ + +#include + +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/*! + * \brief A class for creating EBNF grammar scripts. + * + * This class helps build EBNF (Extended Backus-Naur Form) grammar scripts + * by managing rules and their content. + */ +class EBNFScriptCreator { + public: + /*! \brief Constructor */ + EBNFScriptCreator() = default; + + /*! + * \brief Adds a new rule to the grammar with a suggested name + * \param rule_name_hint Suggested name for the rule + * \param rule_body The EBNF content/definition of the rule + * \return The actual name assigned to the rule + */ + std::string AddRule(const std::string& rule_name_hint, const std::string& rule_body) { + return AddRuleWithAllocatedName(AllocateRuleName(rule_name_hint), rule_body); + } + + /*! + * \brief Generates a new rule name based on a suggested name + * \param rule_name_hint Suggested name for the rule + * \return The actual name assigned to the rule + */ + std::string AllocateRuleName(const std::string& rule_name_hint) { + if (rule_names_.find(rule_name_hint) == rule_names_.end()) { + rule_names_.insert(rule_name_hint); + return rule_name_hint; + } + for (int i = 0; i < NAME_SUFFIX_MAXIMUM; ++i) { + std::string rule_name = rule_name_hint + "_" + std::to_string(i); + if (rule_names_.find(rule_name) == rule_names_.end()) { + rule_names_.insert(rule_name); + return rule_name; + } + } + XGRAMMAR_LOG(FATAL) << "Cannot find a unique rule name for " << rule_name_hint; + XGRAMMAR_UNREACHABLE(); + } + + /*! + * \brief Adds a new rule to the grammar with a allocated name. Used with AllocateRuleName() + * \param rule_name The name of the rule to add + * \param rule_body The EBNF content/definition of the rule + * \return The actual name assigned to the rule + */ + std::string AddRuleWithAllocatedName(const std::string& rule_name, const std::string& rule_body) { + XGRAMMAR_CHECK(rule_names_.find(rule_name) != rule_names_.end()) + << "Rule name " << rule_name << " is not allocated"; + rules_.emplace_back(rule_name, rule_body); + return rule_name; + } + + /*! + * \brief Concatenates a list of strings with a space separator + * \param items The list of strings to concatenate + * \return The concatenated string + */ + static std::string Concat(const std::vector& items) { + std::stringstream ss; + ss << "("; + for (int i = 0; i < static_cast(items.size()); ++i) { + if (i > 0) { + ss << " "; + } + ss << items[i]; + } + ss << ")"; + return ss.str(); + } + + /*! + * \brief Joins a list of strings with an OR operator + * \param items The list of strings to join + * \return The joined string + */ + static std::string Or(const std::vector& items) { + std::stringstream ss; + ss << "("; + for (int i = 0; i < static_cast(items.size()); ++i) { + if (i > 0) { + ss << " | "; + } + ss << items[i]; + } + ss << ")"; + return ss.str(); + } + + /*! + * \brief Escape and quote a string + * \param str The string to escape and quote + * \return The escaped and quoted string + */ + static std::string Str(const std::string& str) { + std::stringstream ss; + ss << "\"" << EscapeString(str) << "\""; + return ss.str(); + } + + /*! + * \brief Repeats an item a given number of times + * \param item The item to repeat + * \param min The minimum number of times to repeat the item + * \param max The maximum number of times to repeat the item + * \return The repeated string + */ + static std::string Repeat(const std::string& item, int min, int max) { + std::stringstream ss; + ss << item; + if (min == 0 && max == 1) { + ss << "?"; + } else if (min == 0 && max == -1) { + ss << "*"; + } else if (min == 1 && max == -1) { + ss << "+"; + } else if (min == 0 && max == 0) { + return ""; + } else if (min == max) { + ss << "{" << min << "}"; + } else if (max == -1) { + ss << "{" << min << ",}"; + } else { + ss << "{" << min << "," << max << "}"; + } + return ss.str(); + } + + /*! + * \brief Gets the complete EBNF grammar script + * \return The full EBNF grammar script as a string + */ + std::string GetScript() { + std::string script = ""; + for (const auto& rule : rules_) { + script += rule.first + " ::= " + rule.second + "\n"; + } + return script; + } + + /*! + * \brief Retrieves the content/definition of a specific rule + * \param rule_name The name of the rule to look up + * \return The EBNF content/definition of the specified rule + */ + std::string GetRuleContent(const std::string& rule_name) { + auto it = std::find_if(rules_.begin(), rules_.end(), [rule_name](const auto& rule) { + return rule.first == rule_name; + }); + if (it != rules_.end()) { + return it->second; + } + return ""; + } + + private: + std::vector> rules_; + std::unordered_set rule_names_; + const int NAME_SUFFIX_MAXIMUM = 10000; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_EBNF_SCRIPT_CREATOR_H_ diff --git a/zig/vendor/xgrammar/cpp/fsm.cc b/zig/vendor/xgrammar/cpp/fsm.cc new file mode 100644 index 0000000000..e6584969e9 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/fsm.cc @@ -0,0 +1,2061 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm.cc + */ +#include "fsm.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "support/union_find_set.h" +#include "support/utils.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/****************** FSMImplBase ******************/ + +template +class FSMImplBase { + static_assert( + std::is_same_v>> || + std::is_same_v>, + "ContainerType must be std::vector> or Compact2DArray" + ); + + public: + /*! \brief Default constructor. */ + FSMImplBase() = default; + + FSMImplBase(const ContainerType& edges, std::vector edge_aux_data = {}) + : edges_(edges), edge_aux_data_(std::move(edge_aux_data)) {} + + FSMImplBase(ContainerType&& edges, std::vector edge_aux_data = {}) + : edges_(std::move(edges)), edge_aux_data_(std::move(edge_aux_data)) {} + + int NumStates() const { return edges_.size(); } + + std::string EdgesToString(std::optional> states = std::nullopt) const; + + const ContainerType& GetEdges() const { return edges_; } + + // For std::vector>, return const std::vector& to avoid copying. + // For Compact2DArray, return Compact2DArray::Row since it is just a simple + // pointer. + decltype(auto) GetEdges(int state) const { return edges_[state]; } + + void GetEpsilonClosure(std::unordered_set* state_set) const; + + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + const std::vector& GetEdgeAuxData() const { return edge_aux_data_; } + + void SetEdgeAuxData(std::vector data) { edge_aux_data_ = std::move(data); } + + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const { return {edge_aux_data_.data() + idx}; } + + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const { return {edge_aux_data_.data() + idx}; } + + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const { + return {edge_aux_data_.data() + idx}; + } + + protected: + ContainerType edges_; + std::vector edge_aux_data_; + friend struct member_trait; +}; + +template +std::string FSMImplBase::EdgesToString(std::optional> states +) const { + std::string result = "[\n"; + auto f_print_one = [&, this](int i) { + result += std::to_string(i) + ": ["; + const auto& edges = edges_[i]; + for (int j = 0; j < static_cast(edges.size()); ++j) { + const auto& edge = edges[j]; + if (edge.min >= 0 && edge.min != edge.max) { + std::string char_min_str = EscapeString(static_cast(edge.min)); + std::string char_max_str = EscapeString(static_cast(edge.max)); + result += "[" + char_min_str + "-" + char_max_str + "]->" + std::to_string(edge.target); + } else if (edge.min >= 0 && edge.min == edge.max) { + std::string char_str = EscapeString(static_cast(edge.min)); + result += "'" + char_str + "'->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kRuleRef) { + result += "Rule(" + std::to_string(edge.max) + ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kEpsilon) { + result += "Eps->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kEOS) { + result += "EOS->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kRepeatRef) { + auto info = GetRepeatEdgeInfo(edge.max); + result += "Repeat(rule=" + std::to_string(info.RuleId()) + + ", min=" + std::to_string(info.Lower()) + + ", max=" + std::to_string(info.Upper()) + ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kToken) { + auto info = GetTokenEdgeInfo(edge.max); + result += "Token("; + for (int32_t k = 0; k < info.Count(); ++k) { + if (k > 0) result += ", "; + result += std::to_string(info.TokenIds()[k]); + } + result += ")->" + std::to_string(edge.target); + } else if (edge.min == FSMEdge::EdgeType::kExcludeToken) { + auto info = GetExcludeTokenEdgeInfo(edge.max); + result += "ExcludeToken("; + for (int32_t k = 0; k < info.Count(); ++k) { + if (k > 0) result += ", "; + result += std::to_string(info.TokenIds()[k]); + } + result += ")->" + std::to_string(edge.target); + } + if (j < static_cast(edges.size()) - 1) { + result += ", "; + } + } + result += "]\n"; + }; + if (states.has_value()) { + for (int i : states.value()) { + f_print_one(i); + } + } else { + for (int i = 0; i < int(NumStates()); ++i) { + f_print_one(i); + } + } + result += "]"; + return result; +} + +template +void FSMImplBase::GetEpsilonClosure(std::unordered_set* state_set) const { + std::queue queue; + for (const auto& state : *state_set) { + queue.push(state); + } + while (!queue.empty()) { + int current = queue.front(); + queue.pop(); + for (const auto& edge : edges_[current]) { + if (!edge.IsEpsilon()) { + continue; + } + if (state_set->find(edge.target) != state_set->end()) { + continue; + } + state_set->insert(edge.target); + queue.push(edge.target); + } + } +} + +template +void FSMImplBase::GetPossibleRules(int state, std::unordered_set* rules) const { + rules->clear(); + for (const auto& edge : edges_[state]) { + if (edge.IsRuleRef()) { + rules->insert(edge.GetRefRuleId()); + } + } +} + +template +void FSMImplBase::GetReachableStates( + const std::vector& from, std::unordered_set* result +) const { + result->clear(); + std::queue queue; + for (const auto& state : from) { + queue.push(state); + result->insert(state); + } + while (!queue.empty()) { + int current = queue.front(); + queue.pop(); + for (const auto& edge : edges_[current]) { + if (result->find(edge.target) != result->end()) { + continue; + } + result->insert(edge.target); + queue.push(edge.target); + } + } +} + +/****************** FSM::Impl ******************/ + +class FSM::Impl : public FSMImplBase>> { + using EdgeType = FSMEdge::EdgeType; + + public: + Impl() = default; + + Impl(int num_states = 0) { edges_.resize(num_states); } + + using FSMImplBase>>::FSMImplBase; + + int GetNextState(int from, int value, EdgeType edge_type) const; + + using FSMImplBase>>::GetEdges; + + std::vector>& GetEdges() { return edges_; } + + std::vector& GetEdges(int state) { return edges_[state]; } + + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + EdgeType edge_type, + bool from_is_closure + ) const; + + int AddState() { + edges_.emplace_back(); + return edges_.size() - 1; + } + + void AddEdge(int from, int to, int32_t min, int32_t max) { + XGRAMMAR_DCHECK(from < static_cast(edges_.size())); + edges_[from].push_back({min, max, to}); + } + + void AddRuleEdge(int from, int to, int32_t rule_id) { + AddEdge(from, to, FSMEdge::EdgeType::kRuleRef, rule_id); + } + + void AddEpsilonEdge(int from, int to) { AddEdge(from, to, FSMEdge::EdgeType::kEpsilon, 0); } + + void AddEOSEdge(int from, int to) { AddEdge(from, to, FSMEdge::EdgeType::kEOS, 0); } + + void AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper) { + XGRAMMAR_DCHECK(edges_[from].empty()) + << "A state with a kRepeatRef edge must have no other outgoing edges."; + XGRAMMAR_DCHECK(edge_aux_data_.size() <= INT32_MAX); + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.reserve(edge_aux_data_.size() + 3); + edge_aux_data_.emplace_back(rule_id); + edge_aux_data_.emplace_back(lower); + edge_aux_data_.emplace_back(upper); + AddEdge(from, to, FSMEdge::EdgeType::kRepeatRef, aux_index); + } + + void AddTokenEdge(int from, int to, const std::vector& token_ids) { + XGRAMMAR_DCHECK(!token_ids.empty()) << "Token set must not be empty"; + XGRAMMAR_CHECK(edge_aux_data_.size() <= INT32_MAX) + << "edge_aux_data_ overflow: too many auxiliary data entries"; + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.push_back(static_cast(token_ids.size())); + for (int32_t id : token_ids) { + edge_aux_data_.push_back(id); + } + edges_[from].push_back(FSMEdge(FSMEdge::EdgeType::kToken, aux_index, to)); + } + + void AddExcludeTokenEdge(int from, int to, const std::vector& token_ids) { + XGRAMMAR_DCHECK(!token_ids.empty()) << "Token exclude set must not be empty"; + XGRAMMAR_CHECK(edge_aux_data_.size() <= INT32_MAX) + << "edge_aux_data_ overflow: too many auxiliary data entries"; + int32_t aux_index = static_cast(edge_aux_data_.size()); + edge_aux_data_.push_back(static_cast(token_ids.size())); + for (int32_t id : token_ids) { + edge_aux_data_.push_back(id); + } + edges_[from].push_back(FSMEdge(FSMEdge::EdgeType::kExcludeToken, aux_index, to)); + } + + void AddFSM(const FSM& fsm, std::vector* state_mapping); + + FSM RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const; + + void SortEdges(); + + CompactFSM ToCompact(); + + friend class FSMWithStartEnd; +}; + +int FSM::Impl::GetNextState(int from, int value, EdgeType edge_type) const { + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call GetNextState with edge type kEpsilon."; + if (edge_type == EdgeType::kCharRange) { + for (const auto& edge : edges_[from]) { + if (edge.min >= EdgeType::kCharRange && edge.min <= value && edge.max >= value) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& edge : edges_[from]) { + if (edge.min == EdgeType::kRuleRef && edge.GetRefRuleId() == value) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kEOS) { + for (const auto& edge : edges_[from]) { + if (edge.min == EdgeType::kEOS) { + return edge.target; + } + } + return FSM::kNoNextState; + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + XGRAMMAR_DCHECK(edges_[from].size() == 1 && edges_[from][0].IsRepeatRef()); + return edges_[from][0].target; + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + XGRAMMAR_UNREACHABLE(); +} + +void FSM::Impl::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + EdgeType edge_type, + bool from_is_closure +) const { + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call Advance with edge type kEpsilon."; + + const std::unordered_set* start_closure; + std::unordered_set start_closure_tmp; + + if (from_is_closure) { + start_closure = &from; + } else { + start_closure_tmp.insert(from.begin(), from.end()); + GetEpsilonClosure(&start_closure_tmp); + start_closure = &start_closure_tmp; + } + + result->clear(); + + if (edge_type == EdgeType::kCharRange) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsCharRange() && edge.min <= value && edge.max >= value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsRuleRef() && edge.GetRefRuleId() == value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsEOS()) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& state : *start_closure) { + if (!edges_[state].empty() && edges_[state][0].IsRepeatRef()) { + result->insert(edges_[state][0].target); + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + + // Get the epsilon closure of the result. + GetEpsilonClosure(result); +} + +void FSM::Impl::AddFSM(const FSM& fsm, std::vector* state_mapping) { + int old_num_states = NumStates(); + int32_t aux_offset = static_cast(edge_aux_data_.size()); + + const auto& other_aux = fsm.GetEdgeAuxData(); + edge_aux_data_.insert(edge_aux_data_.end(), other_aux.begin(), other_aux.end()); + + if (state_mapping != nullptr) { + state_mapping->clear(); + state_mapping->reserve(fsm.NumStates()); + for (int i = 0; i < fsm.NumStates(); ++i) { + state_mapping->push_back(i + old_num_states); + } + } + + edges_.resize(edges_.size() + fsm.NumStates()); + + for (int i = 0; i < fsm.NumStates(); ++i) { + for (const auto& edge : fsm.GetEdges()[i]) { + int32_t max_val = edge.max; + if (edge.IsAuxEdge() && aux_offset > 0) { + max_val = static_cast(edge.max + aux_offset); + } + AddEdge(i + old_num_states, edge.target + old_num_states, edge.min, max_val); + } + } +} + +FSM FSM::Impl::RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const { + std::vector> new_edges(new_num_states); + for (int i = 0; i < static_cast(edges_.size()); ++i) { + for (const auto& edge : edges_[i]) { + if (edge.IsEpsilon() && state_mapping[i] == state_mapping[edge.target]) { + continue; // Skip self-loops for epsilon edges. + } + new_edges[state_mapping[i]].emplace_back(edge.min, edge.max, state_mapping[edge.target]); + } + } + // aux_indices remain stable since only state ids are remapped + for (int i = 0; i < new_num_states; ++i) { + std::sort(new_edges[i].begin(), new_edges[i].end()); + const auto& end_iter = std::unique(new_edges[i].begin(), new_edges[i].end()); + new_edges[i].erase(end_iter, new_edges[i].end()); + } + return FSM(std::move(new_edges), std::vector(edge_aux_data_)); +} + +void FSM::Impl::SortEdges() { + for (int i = 0; i < static_cast(edges_.size()); ++i) { + std::sort(edges_[i].begin(), edges_[i].end()); + } +} + +CompactFSM FSM::Impl::ToCompact() { + SortEdges(); + Compact2DArray edges; + for (int i = 0; i < static_cast(edges_.size()); ++i) { + edges.PushBack(edges_[i]); + } + return CompactFSM(std::move(edges), std::move(edge_aux_data_)); +} + +/****************** FSM ******************/ + +FSM::FSM(int num_states) : pimpl_(std::make_shared(num_states)) {} + +FSM::FSM(const std::vector>& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(edges, std::move(edge_aux_data))) {} + +FSM::FSM(std::vector>&& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(std::move(edges), std::move(edge_aux_data))) {} + +int FSM::NumStates() const { return pimpl_->NumStates(); } + +int FSM::AddState() { return pimpl_->AddState(); } + +void FSM::AddEdge(int from, int to, int32_t min, int32_t max) { + pimpl_->AddEdge(from, to, min, max); +} + +void FSM::AddEdge(int from, int to, FSMEdge::EdgeType type, int32_t value) { + pimpl_->AddEdge(from, to, type, value); +} + +void FSM::AddEpsilonEdge(int from, int to) { pimpl_->AddEpsilonEdge(from, to); } + +void FSM::AddRuleEdge(int from, int to, int32_t rule_id) { pimpl_->AddRuleEdge(from, to, rule_id); } + +void FSM::AddEOSEdge(int from, int to) { pimpl_->AddEOSEdge(from, to); } + +void FSM::AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper) { + pimpl_->AddRepeatEdge(from, to, rule_id, lower, upper); +} + +void FSM::AddTokenEdge(int from, int to, const std::vector& token_ids) { + pimpl_->AddTokenEdge(from, to, token_ids); +} + +void FSM::AddExcludeTokenEdge(int from, int to, const std::vector& token_ids) { + pimpl_->AddExcludeTokenEdge(from, to, token_ids); +} + +const std::vector& FSM::GetEdgeAuxData() const { return pimpl_->GetEdgeAuxData(); } + +void FSM::SetEdgeAuxData(std::vector data) { pimpl_->SetEdgeAuxData(std::move(data)); } + +RepeatEdgeRef FSM::GetRepeatEdgeInfo(int32_t idx) const { return pimpl_->GetRepeatEdgeInfo(idx); } + +TokenEdgeRef FSM::GetTokenEdgeInfo(int32_t idx) const { return pimpl_->GetTokenEdgeInfo(idx); } + +ExcludeTokenEdgeRef FSM::GetExcludeTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetExcludeTokenEdgeInfo(idx); +} + +void FSM::AddFSM(const FSM& fsm, std::vector* state_mapping) { + pimpl_->AddFSM(fsm, state_mapping); +} + +std::string FSM::EdgesToString(std::optional> states) const { + return pimpl_->EdgesToString(states); +} + +const std::vector& FSM::GetEdges(int state) const { return pimpl_->GetEdges(state); } + +std::vector>& FSM::GetEdges() { return pimpl_->GetEdges(); } + +const std::vector>& FSM::GetEdges() const { return pimpl_->GetEdges(); } + +std::vector& FSM::GetEdges(int state) { return pimpl_->GetEdges(state); } + +FSM FSM::Copy() const { return FSM(std::make_shared(*pimpl_)); } + +int FSM::GetNextState(int from, int value, FSMEdge::EdgeType edge_type) const { + return pimpl_->GetNextState(from, value, edge_type); +} + +void FSM::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + pimpl_->Advance(from, value, result, edge_type, from_is_closure); +} + +void FSM::GetPossibleRules(int state, std::unordered_set* rules) const { + pimpl_->GetPossibleRules(state, rules); +} + +void FSM::GetEpsilonClosure(std::unordered_set* state_set) const { + pimpl_->GetEpsilonClosure(state_set); +} + +void FSM::GetReachableStates(const std::vector& from, std::unordered_set* result) const { + pimpl_->GetReachableStates(from, result); +} + +FSM FSM::RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const { + return pimpl_->RebuildWithMapping(state_mapping, new_num_states); +} + +void FSM::SortEdges() { pimpl_->SortEdges(); } + +CompactFSM FSM::ToCompact() { return pimpl_->ToCompact(); } + +/****************** CompactFSM::Impl ******************/ + +class CompactFSM::Impl : public FSMImplBase> { + using EdgeType = FSMEdge::EdgeType; + + public: + Impl() = default; + + Impl(const Compact2DArray& edges, std::vector edge_aux_data = {}) + : FSMImplBase>(edges, std::move(edge_aux_data)), + edge_num_(ComputeEdgeNum(edges_)) {} + + Impl(Compact2DArray&& edges, std::vector edge_aux_data = {}) + : FSMImplBase>(std::move(edges), std::move(edge_aux_data)), + edge_num_(ComputeEdgeNum(edges_)) {} + + void GetNextStates(int from, int value, EdgeType edge_type, std::vector* target) const; + + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure + ) const; + + FSM ToFSM() const; + + size_t GetNumEdges() const { return edge_num_; } + + /*! + * \brief Check that every edge target and every auxiliary data reference is in range. Used after + * deserialization, where the fields are restored verbatim. + * \return An error message if the FSM is malformed. + */ + std::optional Validate() const { + const int64_t aux_size = edge_aux_data_.size(); + for (int state = 0; state < NumStates(); ++state) { + for (const auto& edge : edges_[state]) { + if (edge.target < 0 || edge.target >= NumStates()) { + return "Edge target " + std::to_string(edge.target) + " is out of range"; + } + if (!edge.IsAuxEdge()) { + continue; + } + // A repeat edge owns 3 aux elements; a token edge owns a count followed by count ids. + const int64_t idx = edge.max; + bool in_range = idx >= 0 && idx < aux_size; + if (in_range && edge.IsRepeatRef()) { + in_range = idx + 3 <= aux_size; + } else if (in_range) { + in_range = edge_aux_data_[idx] >= 0 && idx + 1 + edge_aux_data_[idx] <= aux_size; + } + if (!in_range) { + return "Edge aux index " + std::to_string(idx) + " is out of range"; + } + } + } + if (edge_num_ != ComputeEdgeNum(edges_)) { + return "edge_num does not match the number of edges"; + } + return std::nullopt; + } + + size_t edge_num_ = 0; + + friend std::size_t MemorySize(const Impl& impl) { + return MemorySize(impl.edges_) + MemorySize(impl.edge_aux_data_) + sizeof(impl.edge_num_); + } + + private: + static size_t ComputeEdgeNum(const Compact2DArray& edges) { + size_t edge_num = 0; + for (int i = 0; i < edges.size(); ++i) { + edge_num += edges[i].size(); + } + return edge_num; + } +}; + +XGRAMMAR_MEMBER_TABLE( + CompactFSM::Impl, + "edges", + &CompactFSM::Impl::edges_, + "edge_aux_data", + &CompactFSM::Impl::edge_aux_data_, + "edge_num", + &CompactFSM::Impl::edge_num_ +); + +void CompactFSM::Impl::GetNextStates( + int from, int value, EdgeType edge_type, std::vector* targets +) const { + targets->clear(); + XGRAMMAR_DCHECK(edge_type != EdgeType::kEpsilon) + << "Should not call GetNextState with edge type kEpsilon."; + if (edge_type == EdgeType::kCharRange) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kCharRange) { + continue; + } else if (edge.min > value) { + break; + } else if (edge.max >= value) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kRuleRef) { + continue; + } else if (edge.min > EdgeType::kRuleRef) { + break; + } else if (edge.GetRefRuleId() == value) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& edge : edges_[from]) { + if (edge.min < EdgeType::kEOS) { + continue; + } else if (edge.min > EdgeType::kEOS) { + break; + } else if (edge.max >= EdgeType::kEOS) { + targets->push_back(edge.target); + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& edge : edges_[from]) { + if (edge.IsRepeatRef()) { + targets->push_back(edge.target); + break; + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } +} + +void CompactFSM::Impl::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + const std::unordered_set* start_closure; + std::unordered_set start_closure_tmp; + + if (from_is_closure) { + start_closure = &from; + } else { + start_closure_tmp.insert(from.begin(), from.end()); + GetEpsilonClosure(&start_closure_tmp); + start_closure = &start_closure_tmp; + } + + result->clear(); + + if (edge_type == EdgeType::kCharRange) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kCharRange) { + continue; + } else if (edge.min > value) { + break; + } else if (edge.max >= value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRuleRef) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kRuleRef) { + continue; + } else if (edge.min > EdgeType::kRuleRef) { + break; + } else if (edge.GetRefRuleId() == value) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kEOS) { + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.min < EdgeType::kEOS) { + continue; + } else if (edge.min > EdgeType::kEOS) { + break; + } else if (edge.max >= EdgeType::kEOS) { + result->insert(edge.target); + } + } + } + } else if (edge_type == EdgeType::kRepeatRef) { + // By invariant, a state with kRepeatRef has exactly one outgoing edge. + for (const auto& state : *start_closure) { + for (const auto& edge : edges_[state]) { + if (edge.IsRepeatRef()) { + result->insert(edge.target); + break; + } + } + } + } else { + XGRAMMAR_DCHECK(false) << "Invalid edge type: " << static_cast(edge_type); + } + + // Get the epsilon closure of the result. + GetEpsilonClosure(result); +} + +FSM CompactFSM::Impl::ToFSM() const { + std::vector> edges(NumStates()); + for (int i = 0; i < edges_.size(); i++) { + const auto& row = edges_[i]; + edges[i].insert(edges[i].end(), row.begin(), row.end()); + } + return FSM(std::move(edges), std::vector(edge_aux_data_)); +} + +/****************** CompactFSM ******************/ + +CompactFSM::CompactFSM(const Compact2DArray& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(edges, std::move(edge_aux_data))) {} + +CompactFSM::CompactFSM(Compact2DArray&& edges, std::vector edge_aux_data) + : pimpl_(std::make_shared(std::move(edges), std::move(edge_aux_data))) {} + +int CompactFSM::NumStates() const { return pimpl_->NumStates(); } + +const Compact2DArray& CompactFSM::GetEdges() const { return pimpl_->GetEdges(); } + +Compact2DArray::Row CompactFSM::GetEdges(int state) const { + return pimpl_->GetEdges(state); +} + +std::string CompactFSM::EdgesToString(std::optional> states) const { + return pimpl_->EdgesToString(states); +} + +void CompactFSM::GetNextStates( + int from, int value, FSMEdge::EdgeType edge_type, std::vector* targets +) const { + return pimpl_->GetNextStates(from, value, edge_type, targets); +} + +void CompactFSM::Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type, + bool from_is_closure +) const { + pimpl_->Advance(from, value, result, edge_type, from_is_closure); +} + +void CompactFSM::GetPossibleRules(int state_num, std::unordered_set* rules) const { + pimpl_->GetPossibleRules(state_num, rules); +} + +void CompactFSM::GetEpsilonClosure(std::unordered_set* state_set) const { + pimpl_->GetEpsilonClosure(state_set); +} + +void CompactFSM::GetReachableStates(const std::vector& from, std::unordered_set* result) + const { + pimpl_->GetReachableStates(from, result); +} + +size_t CompactFSM::GetNumEdges() const { return pimpl_->GetNumEdges(); } + +FSM CompactFSM::ToFSM() const { return pimpl_->ToFSM(); } + +const std::vector& CompactFSM::GetEdgeAuxData() const { return pimpl_->GetEdgeAuxData(); } + +void CompactFSM::SetEdgeAuxData(std::vector data) { + pimpl_->SetEdgeAuxData(std::move(data)); +} + +RepeatEdgeRef CompactFSM::GetRepeatEdgeInfo(int32_t idx) const { + return pimpl_->GetRepeatEdgeInfo(idx); +} + +TokenEdgeRef CompactFSM::GetTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetTokenEdgeInfo(idx); +} + +ExcludeTokenEdgeRef CompactFSM::GetExcludeTokenEdgeInfo(int32_t idx) const { + return pimpl_->GetExcludeTokenEdgeInfo(idx); +} + +picojson::value SerializeJSONValue(const CompactFSM& value) { + return detail::json_serializer::AutoSerializeJSONValuePImpl(value); +} + +std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name +) { + return detail::json_serializer::AutoDeserializeJSONValuePImpl(result, value, type_name); +} + +struct CompactFSMWithStartEndSerializeHelper { + CompactFSM fsm; + int start; + bool is_dfa; + std::vector end_index; + size_t edge_num; + + CompactFSMWithStartEndSerializeHelper(const CompactFSMWithStartEnd& compact_fsm_with_se) + : fsm(compact_fsm_with_se.fsm_), + start(compact_fsm_with_se.start_), + is_dfa(compact_fsm_with_se.is_dfa_), + end_index(compact_fsm_with_se.ends_), + edge_num(compact_fsm_with_se.edge_num_) {} + + CompactFSMWithStartEndSerializeHelper() = default; + + std::optional Validate() const { + if (fsm.IsNull()) { + return "Expect a non-null fsm"; + } + auto in_range = [&](int32_t state) { return state >= 0 && state < fsm.NumStates(); }; + if (!in_range(start) || !std::all_of(end_index.begin(), end_index.end(), in_range)) { + return "The start or end state is out of range"; + } + return std::nullopt; + } +}; + +XGRAMMAR_MEMBER_ARRAY( + CompactFSMWithStartEndSerializeHelper, + &CompactFSMWithStartEndSerializeHelper::fsm, + &CompactFSMWithStartEndSerializeHelper::start, + &CompactFSMWithStartEndSerializeHelper::end_index, + &CompactFSMWithStartEndSerializeHelper::is_dfa, + &CompactFSMWithStartEndSerializeHelper::edge_num +); + +picojson::value SerializeJSONValue(const CompactFSMWithStartEnd& value) { + return AutoSerializeJSONValue(CompactFSMWithStartEndSerializeHelper(value)); +} +std::optional DeserializeJSONValue( + CompactFSMWithStartEnd* result, const picojson::value& value, const std::string& type_name +) { + CompactFSMWithStartEndSerializeHelper tmp; + auto err = AutoDeserializeJSONValue(&tmp, value, type_name); + if (err.has_value()) { + return err; + } + result->fsm_ = std::move(tmp.fsm); + result->start_ = tmp.start; + result->is_dfa_ = tmp.is_dfa; + result->edge_num_ = tmp.edge_num; + result->SetEndStates(std::move(tmp.end_index)); + return std::nullopt; +} + +struct CompactFSMWithStartEndWithSizeSerializeHelper { + CompactFSMWithStartEnd fsm; + size_t edge_num; + size_t node_num; + + CompactFSMWithStartEndWithSizeSerializeHelper( + const CompactFSMWithStartEndWithSize& compact_fsm_with_size + ) + : fsm(compact_fsm_with_size.fsm_), + edge_num(compact_fsm_with_size.edge_num_), + node_num(compact_fsm_with_size.node_num_) {} + + CompactFSMWithStartEndWithSizeSerializeHelper() = default; +}; + +XGRAMMAR_MEMBER_ARRAY( + CompactFSMWithStartEndWithSizeSerializeHelper, + &CompactFSMWithStartEndWithSizeSerializeHelper::fsm, + &CompactFSMWithStartEndWithSizeSerializeHelper::edge_num, + &CompactFSMWithStartEndWithSizeSerializeHelper::node_num +); + +picojson::value SerializeJSONValue(const CompactFSMWithStartEndWithSize& value) { + return AutoSerializeJSONValue(CompactFSMWithStartEndWithSizeSerializeHelper(value)); +} + +std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name +) { + CompactFSMWithStartEndWithSizeSerializeHelper tmp; + auto err = AutoDeserializeJSONValue(&tmp, value, type_name); + if (err.has_value()) { + return err; + } + result->fsm_ = std::move(tmp.fsm); + result->edge_num_ = tmp.edge_num; + result->node_num_ = tmp.node_num; + return std::nullopt; +} + +/****************** FSMWithStartEnd ******************/ + +std::string FSMWithStartEnd::ToString() const { + std::string result; + result += "FSM(num_states=" + std::to_string(NumStates()) + ", start=" + std::to_string(start_) + + ", end=["; + + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + std::vector reachable_states_vec(reachable_states.begin(), reachable_states.end()); + std::sort(reachable_states_vec.begin(), reachable_states_vec.end()); + + bool first = true; + for (auto end : ends_) { + if (!first) { + result += ", "; + } + first = false; + result += std::to_string(end); + } + + result += "], edges=" + fsm_.EdgesToString(reachable_states_vec) + ")"; + return result; +} + +std::ostream& operator<<(std::ostream& os, const FSMWithStartEnd& fsm) { + os << fsm.ToString(); + return os; +} + +FSMWithStartEnd FSMWithStartEnd::Copy() const { + return FSMWithStartEnd(fsm_.Copy(), start_, ends_, is_dfa_); +} + +FSMWithStartEnd FSMWithStartEnd::RebuildWithMapping( + const std::vector& state_mapping, int new_num_states +) const { + FSM new_fsm = fsm_.RebuildWithMapping(state_mapping, new_num_states); + auto new_start = state_mapping[start_]; + std::vector new_ends; + new_ends.reserve(ends_.size()); + for (auto end : ends_) { + new_ends.push_back(state_mapping[end]); + } + return FSMWithStartEnd(new_fsm, new_start, std::move(new_ends)); +} + +CompactFSMWithStartEnd FSMWithStartEnd::ToCompact() { + return CompactFSMWithStartEnd(fsm_.ToCompact(), start_, ends_, is_dfa_); +} + +FSMWithStartEndWithSize FSMWithStartEnd::AddToCompleteFSM( + FSM* complete_fsm, std::vector* state_mapping +) { + XGRAMMAR_DCHECK(state_mapping != nullptr) << "state_mapping cannot be nullptr"; + complete_fsm->AddFSM(fsm_, state_mapping); + int new_start = (*state_mapping)[start_]; + // Map the end states to the states in the complete FSM. The mapping is monotonic, so the + // sorted invariant is preserved. The sparse representation is important here: this method is + // called once per rule, and the ends of each returned view must not scale with the size of the + // complete FSM, otherwise the total cost is O(num_rules * num_total_states). + std::vector new_ends; + new_ends.reserve(ends_.size()); + for (auto end : ends_) { + new_ends.push_back((*state_mapping)[end]); + } + + int num_edges = 0; + for (int i = 0; i < fsm_.NumStates(); ++i) { + num_edges += fsm_.GetEdges(i).size(); + } + + int num_nodes = fsm_.NumStates(); + + auto fsm_with_se = FSMWithStartEnd(*complete_fsm, new_start, std::move(new_ends), is_dfa_); + + return FSMWithStartEndWithSize(fsm_with_se, num_edges, num_nodes); +} + +FSMWithStartEnd FSMWithStartEnd::Star() const { + FSM fsm = fsm_.Copy(); + auto new_start = fsm.AddState(); + for (auto end : ends_) { + fsm.AddEpsilonEdge(end, new_start); + } + fsm.AddEpsilonEdge(new_start, start_); + return FSMWithStartEnd(fsm, new_start, {new_start}); +} + +FSMWithStartEnd FSMWithStartEnd::Plus() const { + FSM fsm = fsm_.Copy(); + for (auto end : ends_) { + fsm.AddEpsilonEdge(end, start_); + } + return FSMWithStartEnd(fsm, start_, ends_); +} + +FSMWithStartEnd FSMWithStartEnd::Optional() const { + FSM fsm = fsm_.Copy(); + if (!ends_.empty()) { + fsm.AddEpsilonEdge(start_, ends_.front()); + } + return FSMWithStartEnd(fsm, start_, ends_); +} + +Result FSMWithStartEnd::Not(int max_result_num_states) const { + // Check if the FSM contains any rule references. + if (!IsLeaf()) { + XGRAMMAR_LOG(FATAL) << "Not operation is not supported for FSM with rule references."; + } + FSMWithStartEnd result; + if (is_dfa_) { + result = Copy(); + } else { + Result dfa_result = ToDFA(max_result_num_states); + if (dfa_result.IsErr()) { + return dfa_result; + } + result = std::move(dfa_result).Unwrap(); + } + // Reverse all the final states. + std::vector new_final_states; + for (int i = 0; i < result.NumStates(); ++i) { + if (!result.IsEndState(i)) { + new_final_states.push_back(i); // Mark all states as final except the original final states. + } + } + + // Add a new final state that accepts all characters. + int accept_all_new_state = result.AddState(); + new_final_states.push_back(accept_all_new_state); + + std::bitset<256> char_set; + for (int i = 0; i < result.NumStates(); i++) { + char_set.reset(); + // Collect all characters that are not accepted by the original FSM. + for (const auto& edge : result.GetFsm().GetEdges(i)) { + if (edge.IsCharRange()) { + for (int j = edge.min; j <= edge.max; ++j) { + char_set.set(j); + } + } + } + // Add edges for characters that are not accepted. + for (int left_bound = 0; left_bound < 256; ++left_bound) { + if (char_set[left_bound]) { + continue; // Skip characters that are accepted. + } + int right_bound = left_bound + 1; + while (right_bound < 256 && !char_set[right_bound]) { + ++right_bound; + } + result.GetFsm().AddEdge(i, accept_all_new_state, left_bound, right_bound - 1); + left_bound = right_bound; + } + } + + result.SetEndStates(new_final_states); + return ResultOk(result); +} + +FSMWithStartEnd FSMWithStartEnd::Union(const std::vector& fsms) { + // Put all the FSMs in parallel. + // Allocate a new start state. Start state will be linked to the start states of all the FSMs. + // The end states of the new FSM will be the union of the end states of all the FSMs. + if (fsms.size() == 1) { + return fsms[0]; + } + XGRAMMAR_DCHECK(fsms.size() > 1) << "Union of 0 FSMs is not allowed."; + + FSM fsm(1); + int start = 0; + std::vector ends; + + std::vector state_mapping; + + for (const auto& fsm_with_se : fsms) { + fsm.AddFSM(fsm_with_se.GetFsm(), &state_mapping); + fsm.AddEpsilonEdge(start, state_mapping[fsm_with_se.GetStart()]); + for (auto end : fsm_with_se.GetEnds()) { + ends.push_back(state_mapping[end]); + } + } + + return FSMWithStartEnd(fsm, start, std::move(ends)); +} + +FSMWithStartEnd FSMWithStartEnd::Concat(const std::vector& fsms) { + // For each FSM, link the end states to the start state of the next FSM. + // Set the start state of the first FSM as the start state of the result. + // Set the end states of the last FSM as the end states of the result. + if (fsms.size() == 1) { + return fsms[0]; + } + XGRAMMAR_DCHECK(fsms.size() > 1) << "Concatenation of 0 FSMs is not allowed."; + + FSM fsm; + int start = 0; + std::vector ends; + + std::vector state_mapping; + std::vector previous_ends; + + for (int i = 0; i < static_cast(fsms.size()); ++i) { + fsm.AddFSM(fsms[i].GetFsm(), &state_mapping); + if (i == 0) { + start = state_mapping[fsms[i].GetStart()]; + } else { + auto this_start = state_mapping[fsms[i].GetStart()]; + for (const auto& end : previous_ends) { + fsm.AddEpsilonEdge(end, this_start); + } + } + if (i == static_cast(fsms.size()) - 1) { + for (auto end : fsms[i].GetEnds()) { + ends.push_back(state_mapping[end]); + } + } else { + previous_ends.clear(); + previous_ends.reserve(fsms[i].GetEnds().size()); + for (auto end : fsms[i].GetEnds()) { + previous_ends.push_back(state_mapping[end]); + } + } + } + + return FSMWithStartEnd(fsm, start, std::move(ends)); +} + +Result FSMWithStartEnd::Intersect( + const FSMWithStartEnd& lhs, const FSMWithStartEnd& rhs, int max_result_num_states +) { + if (!lhs.IsLeaf() || !rhs.IsLeaf()) { + return ResultErr("Intersect only support leaf fsm!"); + } + auto lhs_dfa_raw = lhs.ToDFA(max_result_num_states); + auto rhs_dfa_raw = rhs.ToDFA(max_result_num_states); + + if (lhs_dfa_raw.IsErr()) { + return lhs_dfa_raw; + } + if (rhs_dfa_raw.IsErr()) { + return rhs_dfa_raw; + } + + auto lhs_dfa = std::move(lhs_dfa_raw).Unwrap(); + auto rhs_dfa = std::move(rhs_dfa_raw).Unwrap(); + // Initialize the result FSM. + FSM result_fsm(0); + FSMWithStartEnd result(result_fsm, 0, std::vector(), true); + std::unordered_map, int> state_map; + std::unordered_set> visited; + std::queue> queue; + queue.push({lhs_dfa.GetStart(), rhs_dfa.GetStart()}); + result.AddState(); + state_map[{lhs_dfa.GetStart(), rhs_dfa.GetStart()}] = 0; + while (!queue.empty()) { + auto [lhs_state, rhs_state] = std::move(queue.front()); + if (lhs_dfa.IsEndState(lhs_state) && rhs_dfa.IsEndState(rhs_state)) { + result.AddEndState(state_map[{lhs_state, rhs_state}]); + } + queue.pop(); + for (const auto& lhs_edge : lhs_dfa.GetFsm().GetEdges(lhs_state)) { + for (const auto& rhs_edge : rhs_dfa.GetFsm().GetEdges(rhs_state)) { + XGRAMMAR_DCHECK(lhs_edge.IsCharRange() && rhs_edge.IsCharRange()); + // Check if the edges intersect. + if (lhs_edge.min > rhs_edge.max || rhs_edge.min > lhs_edge.max) { + continue; // No intersection. + } + int min_value = std::max(lhs_edge.min, rhs_edge.min); + int max_value = std::min(lhs_edge.max, rhs_edge.max); + if (state_map.find(std::make_pair(lhs_edge.target, rhs_edge.target)) == state_map.end()) { + if (result.NumStates() >= max_result_num_states) { + return ResultErr("The number of states in the intersection exceeds the limit."); + } + state_map[{lhs_edge.target, rhs_edge.target}] = result.AddState(); + queue.push({lhs_edge.target, rhs_edge.target}); + } + int target_state = state_map[{lhs_edge.target, rhs_edge.target}]; + result.GetFsm().AddEdge( + state_map[{lhs_state, rhs_state}], target_state, min_value, max_value + ); + } + } + } + return ResultOk(std::move(result)); +} + +bool FSMWithStartEnd::IsDFA() { + if (is_dfa_) { + return true; + } + std::bitset<256> character_transitions; + std::unordered_set rule_transitions; + for (const auto& edges : fsm_->GetEdges()) { + character_transitions.reset(); + rule_transitions.clear(); + for (const auto& edge : edges) { + if (edge.IsEpsilon()) { + return false; // Epsilon transitions are not allowed in DFA. + } + if (edge.IsCharRange()) { + for (int i = edge.min; i <= edge.max; ++i) { + if (character_transitions[i]) { + return false; // Duplicate character transition. + } + character_transitions.set(i); + } + continue; + } + if (edge.IsRuleRef()) { + if (rule_transitions.find(edge.GetRefRuleId()) != rule_transitions.end()) { + return false; // Duplicate rule transition. + } + rule_transitions.insert(edge.GetRefRuleId()); + } + // kRepeatRef: by invariant, a state with kRepeatRef has exactly one edge, always + // deterministic. + } + } + is_dfa_ = true; + return true; +} + +FSMWithStartEnd FSMWithStartEnd::SimplifyEpsilon(int max_num_states) const { + if (is_dfa_) { + return *this; + } + if (NumStates() > max_num_states) { + return *this; + } + + UnionFindSet union_find_set; + std::vector in_degree(NumStates(), 0); + std::vector> epsilon_edges; + + std::vector has_exclude_token(NumStates(), false); + for (int i = 0; i < NumStates(); i++) { + for (const auto& edge : fsm_->GetEdges(i)) { + if (edge.IsExcludeToken()) { + has_exclude_token[i] = true; + break; + } + } + } + + for (int i = 0; i < NumStates(); i++) { + const auto& edges = fsm_->GetEdges(i); + for (const auto& edge : edges) { + in_degree[edge.target]++; + if (edge.IsEpsilon()) { + // a -- epsilon --> b, and a doesn't have other outward edges. Do not merge an + // accepting a into a non-accepting b: the merged state would be accepting, so other + // paths reaching b would be wrongly accepted. (If b is accepting, a effectively + // accepts already via the epsilon edge, so merging is safe.) + if (edges.size() == 1 && !has_exclude_token[i] && !has_exclude_token[edge.target] && + (!IsEndState(i) || IsEndState(edge.target))) { + union_find_set.Add(i); + union_find_set.Add(edge.target); + union_find_set.Union(i, edge.target); + in_degree[edge.target]--; // Remove the inward edge since a and b are merged. + } else { + // Otherwise, we store it to check for the second merge rule. + epsilon_edges.emplace_back(i, edge.target); + } + } + } + } + + // Build the equivalent graph. + std::vector equiv_node(NumStates()); + for (int i = 0; i < NumStates(); i++) { + if (union_find_set.Count(i)) { + equiv_node[i] = union_find_set.Find(i); + if (equiv_node[i] == i) { + continue; + } + in_degree[equiv_node[i]] += in_degree[i]; + } else { + equiv_node[i] = i; + } + } + + // a --> epsilon --> b, and b doesn't have other inward edges. + for (const auto& [from_raw, to_raw] : epsilon_edges) { + const int& from = equiv_node[from_raw]; + const int& to = equiv_node[to_raw]; + if (in_degree[to] == 1 && equiv_node[GetStart()] != to && !has_exclude_token[from_raw] && + !has_exclude_token[to_raw]) { + union_find_set.Add(from); + union_find_set.Add(to); + union_find_set.Union(from, to); + } + } + + // Merge the states. + auto eq_classes = union_find_set.GetAllSets(); + if (eq_classes.empty()) { + return *this; + } + + std::vector new_to_old(NumStates(), -1); + for (size_t i = 0; i < eq_classes.size(); i++) { + for (const auto& state : eq_classes[i]) { + new_to_old[state] = i; + } + } + + int cnt = eq_classes.size(); + for (int i = 0; i < NumStates(); i++) { + if (new_to_old[i] == -1) { + new_to_old[i] = cnt; + cnt++; + } + } + return RebuildWithMapping(new_to_old, cnt); +} + +FSMWithStartEnd FSMWithStartEnd::MergeEquivalentStates(int max_result_num_states) const { + if (max_result_num_states < NumStates()) { + return *this; + } + // No merge is possible with fewer than 4 states (need >=2 sources sharing a target, + // or >=2 sinks sharing a source). + if (NumStates() < 4) { + return Copy(); + } + bool changed = true; + FSMWithStartEnd result = Copy(); + result.GetFsm()->SortEdges(); + UnionFindSet union_find_set; + + // A compact edge view used for incoming/outgoing CSR rows. `peer` means source state in + // incoming_edges and target state in outgoing_edges. + struct EndpointEdge { + int peer; // source in incoming_edges, target in outgoing_edges + int32_t min; + int32_t max; + + bool operator<(const EndpointEdge& other) const { + return std::tie(peer, min, max) < std::tie(other.peer, other.min, other.max); + } + }; + + // Scratch buffers reused across iterations to avoid repeated vector allocation. + // Number of incoming edges for each state, used to size the incoming CSR rows. + std::vector incoming_row_sizes; + // Number of outgoing edges for each state, used to size the outgoing CSR rows. + std::vector outgoing_row_sizes; + // Write positions while filling incoming_edges rows. + std::vector incoming_write_positions; + // Write positions while filling outgoing_edges rows. + std::vector outgoing_write_positions; + // Incoming edges grouped by target state. + Compact2DArray incoming_edges; + // Outgoing edges grouped by source state. + Compact2DArray outgoing_edges; + // Number of distinct predecessor states for each state. + std::vector incoming_distinct_count; + // Number of distinct successor states for each state. + std::vector outgoing_distinct_count; + // The only predecessor state when incoming_distinct_count[state] == 1. + std::vector single_incoming_source; + // The only successor state when outgoing_distinct_count[state] == 1. + std::vector single_outgoing_target; + // Terminal end states collected for leaf-state merging. + std::vector no_successor_end_states; + // Terminal non-end states collected for leaf-state merging. + std::vector no_successor_non_end_states; + + while (changed) { + int n = result.NumStates(); + union_find_set.Clear(); + + // First pass: count row sizes for the incoming/outgoing CSR arrays. + incoming_row_sizes.assign(n, 0); + outgoing_row_sizes.assign(n, 0); + for (int source = 0; source < n; ++source) { + const auto& edges = result.GetFsm().GetEdges(source); + outgoing_row_sizes[source] = static_cast(edges.size()); + for (const auto& edge : edges) { + ++incoming_row_sizes[edge.target]; + } + } + + // Allocate CSR rows. The underlying storage is reset and reused across iterations. + incoming_edges.ResetWithRowSizes(incoming_row_sizes); + outgoing_edges.ResetWithRowSizes(outgoing_row_sizes); + incoming_write_positions.assign(n, 0); + outgoing_write_positions.assign(n, 0); + + // Second pass: fill incoming and outgoing rows. Incoming rows are naturally grouped by + // source because we scan source states in order; outgoing rows are sorted by target below. + for (int source = 0; source < n; ++source) { + const auto& edges = result.GetFsm().GetEdges(source); + auto outgoing_row = outgoing_edges.MutableRowAt(source); + for (const auto& edge : edges) { + incoming_edges.MutableRowAt(edge.target + )[incoming_write_positions[edge.target]++] = {source, edge.min, edge.max}; + outgoing_row[outgoing_write_positions[source]++] = {edge.target, edge.min, edge.max}; + } + std::sort(outgoing_row.begin(), outgoing_row.end()); + } + + // Identify states with exactly one distinct predecessor/successor. These states are the + // candidates for the two local merge rules below. + incoming_distinct_count.assign(n, 0); + outgoing_distinct_count.assign(n, 0); + single_incoming_source.assign(n, -1); + single_outgoing_target.assign(n, -1); + for (int state = 0; state < n; ++state) { + auto incoming_row = incoming_edges[state]; + if (incoming_row.size() > 0) { + incoming_distinct_count[state] = 1; + single_incoming_source[state] = incoming_row[0].peer; + for (int32_t i = 1; i < incoming_row.size(); ++i) { + if (incoming_row[i].peer != incoming_row[i - 1].peer) { + ++incoming_distinct_count[state]; + single_incoming_source[state] = -1; + } + } + } + auto outgoing_row = outgoing_edges[state]; + if (outgoing_row.size() > 0) { + outgoing_distinct_count[state] = 1; + single_outgoing_target[state] = outgoing_row[0].peer; + for (int32_t i = 1; i < outgoing_row.size(); ++i) { + if (outgoing_row[i].peer != outgoing_row[i - 1].peer) { + ++outgoing_distinct_count[state]; + single_outgoing_target[state] = -1; + } + } + } + } + + // Case 1: Like ab | ac | ad, then they can be merged into a(b | c | d). + bool is_equiv_successor = false; + for (int i = 0; i < n; i++) { + // The start state is also entered without taking any edge, so it is never equivalent to + // a state that is only reached through the start state's own self-loop label. + if (incoming_distinct_count[i] != 1 || union_find_set.Count(i) || i == result.GetStart()) { + continue; + } + int previous_state = single_incoming_source[i]; + auto edges_to_i = incoming_edges[i]; + auto siblings = outgoing_edges[previous_state]; + int32_t group_begin = 0; + while (group_begin < siblings.size()) { + int sibling = siblings[group_begin].peer; + int32_t group_end = group_begin + 1; + while (group_end < siblings.size() && siblings[group_end].peer == sibling) { + ++group_end; + } + auto edges_to_sibling = siblings.Slice(group_begin, group_end); + group_begin = group_end; + if (sibling <= i || incoming_distinct_count[sibling] != 1 || sibling == result.GetStart() || + result.IsEndState(sibling) != result.IsEndState(i)) { + continue; + } + // Check if the edges from previous_state to i and sibling are the same. + if (edges_to_i.size() != edges_to_sibling.size()) { + continue; // Different edges, not equivalent. + } + bool is_equiv = true; + for (int32_t j = 0; j < edges_to_i.size(); ++j) { + if (edges_to_i[j].min != edges_to_sibling[j].min || + edges_to_i[j].max != edges_to_sibling[j].max) { + is_equiv = false; + break; // Different edge ranges, not equivalent. + } + } + // Merge the equivalent successor states. + if (is_equiv) { + union_find_set.Add(i); + union_find_set.Add(sibling); + union_find_set.Union(i, sibling); + is_equiv_successor = true; + } + } + } + + // Case 2: Like ba | ca | da, then they can be merged into (b | c | d)a. + bool is_equiv_precursor = false; + no_successor_end_states.clear(); + no_successor_non_end_states.clear(); + + for (int i = 0; i < n; i++) { + int outgoing_count = outgoing_distinct_count[i]; + if (outgoing_count == 0) { + if (result.IsEndState(i)) { + no_successor_end_states.push_back(i); + } else { + no_successor_non_end_states.push_back(i); + } + continue; // Skip states with no successors. + } + if (outgoing_count != 1 || union_find_set.Count(i)) { + continue; // Skip states with multiple successors. + } + int next_state = single_outgoing_target[i]; + auto node_edges = outgoing_edges[i]; + auto siblings = incoming_edges[next_state]; + int32_t group_begin = 0; + while (group_begin < siblings.size()) { + int sibling = siblings[group_begin].peer; + while (group_begin < siblings.size() && siblings[group_begin].peer == sibling) { + ++group_begin; + } + // Avoid chaining a Case 2 merge onto a state already merged earlier in this iteration + // (typically by Case 1), which can over-merge via transitive closure. + if (sibling <= i || union_find_set.Count(sibling) || + outgoing_distinct_count[sibling] != 1 || + result.IsEndState(i) != result.IsEndState(sibling)) { + continue; + } + auto sibling_node_edges = outgoing_edges[sibling]; + if (sibling_node_edges.size() != node_edges.size()) { + continue; // Different number of edges, not equivalent. + } + // Check if the sibling state has the same outgoing edges as i. + bool is_equiv = true; + for (int32_t j = 0; j < node_edges.size(); ++j) { + if (sibling_node_edges[j].min != node_edges[j].min || + sibling_node_edges[j].max != node_edges[j].max) { + is_equiv = false; + break; + } + } + // Merge the equivalent precursor states. + if (is_equiv) { + union_find_set.Add(i); + union_find_set.Add(sibling); + union_find_set.Union(i, sibling); + is_equiv_precursor = true; + } + } + } + + if (no_successor_end_states.size() > 1) { + // Merge all end states with no successors. + for (size_t i = 1; i < no_successor_end_states.size(); ++i) { + union_find_set.Add(no_successor_end_states[0]); + union_find_set.Add(no_successor_end_states[i]); + union_find_set.Union(no_successor_end_states[0], no_successor_end_states[i]); + is_equiv_precursor = true; + } + } + + if (no_successor_non_end_states.size() > 1) { + // Merge all non-end states with no successors. + for (size_t i = 1; i < no_successor_non_end_states.size(); ++i) { + union_find_set.Add(no_successor_non_end_states[0]); + union_find_set.Add(no_successor_non_end_states[i]); + union_find_set.Union(no_successor_non_end_states[0], no_successor_non_end_states[i]); + is_equiv_precursor = true; + } + } + + changed = is_equiv_successor || is_equiv_precursor; + if (changed) { + // Rebuild the FSM with the equivalent states merged, then repeat until no local merge + // rule applies. + auto eq_classes = union_find_set.GetAllSets(); + std::vector old_to_new(result.NumStates(), -1); + for (size_t i = 0; i < eq_classes.size(); i++) { + for (const auto& state : eq_classes[i]) { + old_to_new[state] = i; + } + } + int cnt = eq_classes.size(); + for (int i = 0; i < result.NumStates(); i++) { + if (old_to_new[i] == -1) { + old_to_new[i] = cnt; + cnt++; + } + } + result = result.RebuildWithMapping(old_to_new, cnt); + result.GetFsm()->SortEdges(); + } + } + return result; +} + +Result FSMWithStartEnd::MinimizeDFA(int max_num_states) const { + FSMWithStartEnd now_fsm(FSM(0), 0, std::vector(), true); + if (NumStates() > max_num_states) { + return ResultErr("The number of states exceeds the limit."); + } + // To perform the algorithm, we must make sure the FSM is + // a DFA. + if (!is_dfa_) { + Result dfa_raw = ToDFA(max_num_states); + if (dfa_raw.IsErr()) { + return dfa_raw; + } + now_fsm = std::move(dfa_raw).Unwrap(); + } else { + now_fsm = Copy(); + } + + // Initialize the precursors of nodes. + std::vector, int>>> precursors; + precursors.resize(now_fsm.NumStates()); + for (int i = 0; i < now_fsm.NumStates(); ++i) { + const auto& edges = now_fsm.GetFsm().GetEdges(i); + for (const auto& edge : edges) { + XGRAMMAR_DCHECK(!edge.IsEpsilon()); + precursors[edge.target].push_back(std::make_pair(std::make_pair(edge.min, edge.max), i)); + } + } + + // Initialize the partitions and working set. + std::vector> partitions; + std::vector> working_set; + std::unordered_set final_states; + std::unordered_set non_final_states; + for (int i = 0; i < now_fsm.NumStates(); ++i) { + if (now_fsm.IsEndState(i)) { + final_states.insert(i); + } else { + non_final_states.insert(i); + } + } + partitions.push_back(final_states); + partitions.push_back(non_final_states); + working_set.push_back(std::move(final_states)); + working_set.push_back(std::move(non_final_states)); + + while (!working_set.empty()) { + std::map, std::unordered_set> possible_transitions; + auto current_partition = std::move(working_set.back()); + working_set.pop_back(); + + // Get the possible transitions from the current partition. + for (const auto& state : current_partition) { + const auto& precursor_map = precursors[state]; + for (const auto& precursor : precursor_map) { + if (possible_transitions.find(precursor.first) == possible_transitions.end()) { + possible_transitions[precursor.first] = std::unordered_set(); + } + possible_transitions[precursor.first].insert(precursor.second); + } + } + + // Check each possible transition. + std::vector intersection; + std::vector difference; + for (const auto& [transition, precursors] : possible_transitions) { + for (size_t i = 0; i < partitions.size(); i++) { + const auto& partition = partitions[i]; + intersection.clear(); // partition \cap precursors + difference.clear(); // partition - precursors + for (const auto& partition_state : partition) { + if (precursors.find(partition_state) != precursors.end()) { + intersection.push_back(partition_state); + } else { + difference.push_back(partition_state); + } + } + + // the states in the partition is not equivalent. We need to + // update the working set and the partitions. + if ((!intersection.empty()) && (!difference.empty())) { + bool in_working_set = false; + for (size_t i = 0; i < working_set.size(); i++) { + if (partition == working_set[i]) { + in_working_set = true; + working_set[i].clear(); + for (const auto& state : intersection) { + working_set[i].insert(state); + } + working_set.emplace_back(); + for (const auto& state : difference) { + working_set.back().insert(state); + } + break; + } + } + if (!in_working_set) { + const auto& smaller_set = + difference.size() < intersection.size() ? difference : intersection; + working_set.emplace_back(); + for (const auto& state : smaller_set) { + working_set.back().insert(state); + } + } + partitions[i].clear(); + for (const auto& state : intersection) { + partitions[i].insert(state); + } + partitions.emplace_back(); + for (const auto& state : difference) { + partitions.back().insert(state); + } + } + } + } + } + std::vector state_mapping(now_fsm.NumStates(), -1); + for (size_t i = 0; i < partitions.size(); ++i) { + for (const auto& state : partitions[i]) { + state_mapping[state] = i; + } + } + int new_num_states = partitions.size(); + return ResultOk(now_fsm.RebuildWithMapping(state_mapping, new_num_states)); +} + +Result FSMWithStartEnd::ToDFA(int max_num_states) const { + if (NumStates() > max_num_states) { + return ResultErr("The number of states exceeds the limit."); + } + FSMWithStartEnd dfa(FSM(0), 0, std::vector(), true); + std::vector> closures; + std::unordered_set rules; + std::unordered_set repeat_aux_indices; + int now_process = 0; + std::unordered_set closure; + closure.insert(start_); + fsm_.GetEpsilonClosure(&closure); + closures.push_back(closure); + while (now_process < static_cast(closures.size())) { + rules.clear(); + repeat_aux_indices.clear(); + std::unordered_set token_aux_indices; + std::unordered_set exclude_token_aux_indices; + std::set interval_ends; + std::bitset<256> allowed_characters; + dfa.AddState(); + // Check if the closure is a final state. + for (const auto& state : closures[now_process]) { + if (IsEndState(state)) { + dfa.AddEndState(now_process); + } + const auto& edges = fsm_->GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + interval_ends.insert(edge.min); + interval_ends.insert(edge.max + 1); + for (int i = edge.min; i <= edge.max; ++i) { + allowed_characters.set(i); + } + continue; + } else if (edge.IsRuleRef()) { + rules.insert(edge.GetRefRuleId()); + } else if (edge.IsRepeatRef()) { + repeat_aux_indices.insert(edge.GetAuxIndex()); + } else if (edge.IsToken()) { + token_aux_indices.insert(edge.GetAuxIndex()); + } else if (edge.IsExcludeToken()) { + exclude_token_aux_indices.insert(edge.GetAuxIndex()); + } + } + } + // This part is to get the all possible intervals. + // Which can help reduce the transitions. + using Interval = std::pair; + std::vector intervals; + intervals.reserve(interval_ends.size()); + int last = -1; + for (const auto& end : interval_ends) { + if (last == -1) { + last = end; + continue; + } + bool allowed = true; + for (int i = last; i < end; ++i) { + if (!allowed_characters[i]) { + allowed = false; + break; + } + } + if (allowed) { + intervals.emplace_back(last, end - 1); + } + last = end; + } + for (const auto& interval : intervals) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_->GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + if (interval.first >= edge.min && interval.second <= edge.max) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, interval.first, interval.second); + flag = true; + break; + } + } + if (!flag) { + if (static_cast(closures.size()) >= max_num_states) { + return ResultErr("The number of DFA states exceeds the limit."); + } + dfa.GetFsm().AddEdge(now_process, closures.size(), interval.first, interval.second); + closures.push_back(next_closure); + } + } + for (auto rule : rules) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsRuleRef()) { + if (rule == edge.GetRefRuleId()) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddRuleEdge(now_process, j, rule); + flag = true; + break; + } + } + if (!flag) { + if (static_cast(closures.size()) >= max_num_states) { + return ResultErr("The number of DFA states exceeds the limit."); + } + dfa.GetFsm().AddRuleEdge(now_process, closures.size(), rule); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : repeat_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsRepeatRef() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kRepeatRef, aux_idx); + flag = true; + break; + } + } + if (!flag) { + if (static_cast(closures.size()) >= max_num_states) { + return ResultErr("The number of DFA states exceeds the limit."); + } + dfa.GetFsm().AddEdge(now_process, closures.size(), FSMEdge::EdgeType::kRepeatRef, aux_idx); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : token_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsToken() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kToken, aux_idx); + flag = true; + break; + } + } + if (!flag) { + if (static_cast(closures.size()) >= max_num_states) { + return ResultErr("The number of DFA states exceeds the limit."); + } + dfa.GetFsm().AddEdge(now_process, closures.size(), FSMEdge::EdgeType::kToken, aux_idx); + closures.push_back(next_closure); + } + } + + for (auto aux_idx : exclude_token_aux_indices) { + std::unordered_set next_closure; + for (const auto& state : closures[now_process]) { + const auto& edges = fsm_.GetEdges(state); + for (const auto& edge : edges) { + if (edge.IsExcludeToken() && edge.GetAuxIndex() == aux_idx) { + if (next_closure.find(edge.target) == next_closure.end()) { + std::unordered_set epsilon_closure; + epsilon_closure.insert(edge.target); + fsm_.GetEpsilonClosure(&epsilon_closure); + next_closure.insert(epsilon_closure.begin(), epsilon_closure.end()); + } + } + } + } + bool flag = false; + for (int j = 0; j < static_cast(closures.size()); j++) { + if (closures[j] == next_closure) { + dfa.GetFsm().AddEdge(now_process, j, FSMEdge::EdgeType::kExcludeToken, aux_idx); + flag = true; + break; + } + } + if (!flag) { + if (static_cast(closures.size()) >= max_num_states) { + return ResultErr("The number of DFA states exceeds the limit."); + } + dfa.GetFsm().AddEdge( + now_process, closures.size(), FSMEdge::EdgeType::kExcludeToken, aux_idx + ); + closures.push_back(next_closure); + } + } + + now_process++; + } + dfa.GetFsm().SetEdgeAuxData(std::vector(fsm_.GetEdgeAuxData())); + dfa.is_dfa_ = true; + return ResultOk(dfa); +} + +/****************** CompactFSMWithStartEnd ******************/ + +std::string CompactFSMWithStartEnd::ToString() const { + std::string result; + result += "CompactFSM(num_states=" + std::to_string(NumStates()) + + ", start=" + std::to_string(start_) + ", end=["; + + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + std::vector reachable_states_vec(reachable_states.begin(), reachable_states.end()); + std::sort(reachable_states_vec.begin(), reachable_states_vec.end()); + bool first = true; + for (auto end : ends_) { + if (reachable_states.count(end)) { + if (!first) { + result += ", "; + } + first = false; + result += std::to_string(end); + } + } + + result += "], edges=" + fsm_.EdgesToString(reachable_states_vec) + ")"; + return result; +} + +std::ostream& operator<<(std::ostream& os, const CompactFSMWithStartEnd& fsm) { + os << fsm.ToString(); + return os; +} + +std::size_t MemorySize(const CompactFSM& self) { return MemorySize(*self.ImplPtr()); } + +std::size_t MemorySize(const CompactFSMWithStartEnd& self) { + // The underlying CompactFSM is not counted here: CompactFSMWithStartEnd is a view, and many + // views usually share one CompactFSM (e.g. the per-rule FSMs of a grammar all point to the + // grammar's complete FSM, which is counted by the grammar itself). Counting it per view would + // multiply the shared FSM's size by the number of views. + return MemorySize(self.ends_); +} + +FSMWithStartEnd CompactFSMWithStartEnd::ToFSM() const { + return FSMWithStartEnd(fsm_.ToFSM(), start_, ends_); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/fsm.h b/zig/vendor/xgrammar/cpp/fsm.h new file mode 100644 index 0000000000..7647962938 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/fsm.h @@ -0,0 +1,1051 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm.h + * \note For functions accepting a pointer to a container as result, the container will be cleared + * before the result is stored. + */ +#ifndef XGRAMMAR_FSM_H_ +#define XGRAMMAR_FSM_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/compact_2d_array.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "support/utils.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/*! + * \brief The edge of a FSM. + */ +struct FSMEdge { + /*! + * \brief Edge type is encoded in the `min` field. When min >= 0, the edge is a character range + * [min, max]. When min < 0, it is a special edge type identified by the enum values below. + * + * For each type, `max` has a type-specific meaning (see comments on each enumerator). + */ + enum EdgeType : int32_t { + //! Character range [min, max]. min >= 0. + kCharRange = 0, + //! Epsilon transition. max is unused. + kEpsilon = -1, + //! Rule reference. max = rule_id. + kRuleRef = -2, + //! Accepts the EOS token. max is unused. + kEOS = -3, + //! Repeated rule reference. max = aux index into edge_aux_data + //! (layout: [rule_id, lower, upper]). + //! Invariant: a state with a kRepeatRef edge has exactly one outgoing edge. + kRepeatRef = -4, + //! Accepts a set of token IDs. max = aux index into edge_aux_data + //! (layout: [count, token_id_0, token_id_1, ...]). + kToken = -5, + //! Accepts any token NOT in the given set. max = aux index into edge_aux_data + //! (layout: [count, token_id_0, token_id_1, ...]). + kExcludeToken = -6, + }; + + inline static constexpr int kMaxChar = 255; + + int32_t min, max; + + /*! + * \brief The target state id of the edge. + */ + int32_t target; + + // for serialization only + FSMEdge() = default; + + FSMEdge(int32_t min, int32_t max, int32_t target) : min(min), max(max), target(target) { + XGRAMMAR_DCHECK(!IsCharRange() || min <= max) + << "Invalid FSMEdge: min > max. min=" << min << ", max=" << max; + } + + /*! + * \brief Compare the edges. Used to sort the edges in the FSM. + */ + // TODO(yixin): consider combining the fields to a single int64_t for better efficiency + friend bool operator==(const FSMEdge& lhs, const FSMEdge& rhs) { + return std::make_tuple(lhs.min, lhs.max, lhs.target) == + std::make_tuple(rhs.min, rhs.max, rhs.target); + } + + /*! + * \brief Compare the edges. Used to sort the edges in the FSM. + */ + friend bool operator<(const FSMEdge& lhs, const FSMEdge& rhs) { + return std::make_tuple(lhs.min, lhs.max, lhs.target) < + std::make_tuple(rhs.min, rhs.max, rhs.target); + } + + /*! + * \brief Check if the edge is a character range. + */ + bool IsCharRange() const { return min >= 0; } + + /*! + * \brief Check if the edge is an epsilon transition. + */ + bool IsEpsilon() const { return min == EdgeType::kEpsilon; } + + /*! + * \brief Check if the edge is a rule reference. + */ + bool IsRuleRef() const { return min == EdgeType::kRuleRef; } + + /*! + * \brief Check if the edge is an EOS transition. + */ + bool IsEOS() const { return min == EdgeType::kEOS; } + + /*! + * \brief Check if the edge is a repeat reference. + */ + bool IsRepeatRef() const { return min == EdgeType::kRepeatRef; } + + bool IsToken() const { return min == EdgeType::kToken; } + + bool IsExcludeToken() const { return min == EdgeType::kExcludeToken; } + + /*! + * \brief Get the rule id of the edge. + * \return The rule id of the edge. -1 if the edge is not a rule reference. + */ + int32_t GetRefRuleId() const { return IsRuleRef() ? max : -1; } + + /*! + * \brief Get the auxiliary data index for repeat reference edges. + * \return The index into the owning FSM's edge_aux_data. -1 if not a repeat reference. + */ + int32_t GetAuxIndex() const { + return (IsRepeatRef() || IsToken() || IsExcludeToken()) ? max : -1; + } + + /*! \brief Check if the edge uses auxiliary data. */ + bool IsAuxEdge() const { return IsRepeatRef() || IsToken() || IsExcludeToken(); } + + friend struct member_trait; +}; + +/*! \brief View into edge_aux_data for a repeat edge (layout: [rule_id, lower, upper]). */ +struct RepeatEdgeRef { + const int32_t* data; + int32_t RuleId() const { return data[0]; } + int32_t Lower() const { return data[1]; } + int32_t Upper() const { return data[2]; } +}; + +/*! \brief View into edge_aux_data for a token edge (layout: [count, token_id_0, ...]). */ +struct TokenEdgeRef { + const int32_t* data; + int32_t Count() const { return data[0]; } + const int32_t* TokenIds() const { return data + 1; } + bool Contains(int32_t token_id) const { + for (int32_t i = 0; i < Count(); ++i) { + if (TokenIds()[i] == token_id) return true; + } + return false; + } +}; + +/*! \brief View into edge_aux_data for an exclude-token edge (layout: [count, token_id_0, ...]). */ +struct ExcludeTokenEdgeRef { + const int32_t* data; + int32_t Count() const { return data[0]; } + const int32_t* TokenIds() const { return data + 1; } + bool Contains(int32_t token_id) const { + for (int32_t i = 0; i < Count(); ++i) { + if (TokenIds()[i] == token_id) return true; + } + return false; + } + bool Accepts(int32_t token_id) const { return !Contains(token_id); } +}; + +/*! + * \brief Comparator for FSMEdge. Only compare the min and max. + */ +struct FSMEdgeRangeComparator { + bool operator()(const FSMEdge& lhs, const FSMEdge& rhs) const { + return std::make_tuple(lhs.min, lhs.max) < std::make_tuple(rhs.min, rhs.max); + } +}; + +XGRAMMAR_MEMBER_ARRAY(FSMEdge, &FSMEdge::min, &FSMEdge::max, &FSMEdge::target); + +} // namespace xgrammar + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::FSMEdge, &xgrammar::FSMEdge::min, &xgrammar::FSMEdge::max, &xgrammar::FSMEdge::target +); + +namespace xgrammar { + +class CompactFSM; + +/*! + * \brief FSM is a class that represents a finite state machine, could be a DFA or an NFA. + * \details It's mutable, which means you can add edges and states to it. + */ +class FSM { + public: + /*! + * \brief Construct an FSM with a given number of states. + * \param num_states The number of states in the FSM. + */ + FSM(int num_states = 0); + + /*! + * \brief Construct an FSM with a given set of edges. + */ + FSM(const std::vector>& edges, std::vector edge_aux_data = {}); + + /*! + * \brief Construct an FSM with a given set of edges. + */ + FSM(std::vector>&& edges, std::vector edge_aux_data = {}); + + /****************** FSM Visitors ******************/ + + /*! + * \brief Get the number of states in the FSM. + * \return The number of states in the FSM. + */ + int NumStates() const; + + /*! + * \brief Get the edges of the FSM. + * \return The edges of the FSM. + */ + const std::vector>& GetEdges() const; + + /*! + * \brief Get the edges of the FSM. + * \return The edges of the FSM. + */ + std::vector>& GetEdges(); + + /*! + * \brief Get the edges of the FSM. + * \param state The state to get the edges from. + * \return The edges of the FSM. + */ + std::vector& GetEdges(int state); + + /*! + * \brief Get the edges of the FSM. + * \param state The state to get the edges from. + * \return The edges of the FSM. + */ + const std::vector& GetEdges(int state) const; + + /*! + * \brief Convert the edges of the FSM to a string. Used in printing the FSM. + * \return The string representation of the edges of the FSM. + */ + std::string EdgesToString(std::optional> states = std::nullopt) const; + + /****************** FSM Traversal Visitors ******************/ + + inline static constexpr int kNoNextState = -1; + + /*! + * \brief Advance the FSM from a given state based on an input character. If there are multiple + * transitions, the first one will be returned. + * \param from The source state to transition from. + * \param character The input character. + * \return The target state if a valid transition exists, kNoNextState otherwise. + */ + int GetNextState(int from, int value, FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange) + const; + + /*! + * \brief Advance the FSM to the next state. + * \param from The current states. + * \param value The input value. + * \param result The possible next states. The result is cleared at the beginning. + * \param value_is_rule Whether the input value is a rule id. + * \param from_is_closure Whether from is an epsilon closure. + */ + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + bool from_is_closure = false + ) const; + + /*! + * \brief Get all the possible rule numbers for a given state. + * \param state_num The state number. + * \param rules The set of possible rule numbers. The result is cleared at the beginning. + */ + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + /*! + * \brief Get the epsilon closure of a set of states, i.e. those can be reached by epsilon + * transitions. + * \param state_set The states in the epsilon closure. The result is not cleared. + */ + void GetEpsilonClosure(std::unordered_set* state_set) const; + + /*! + * \brief Get the reachable states from a set of states. + * \param from The current states. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + /****************** FSM Mutators ******************/ + + /*! + * \brief Adds a new state to the FSM. + * \return The index of the newly added state. + */ + int AddState(); + + /*! + * \brief Adds a transition edge between states with given min and max values. For character + * transitions, it accepts any character in range [min, max]. + * \param from The source state. + * \param to The target state. + * \param min The min value of the range. + * \param max The max value of the range. + */ + void AddEdge(int from, int to, int32_t min, int32_t max); + + /*! \brief Add a raw edge with explicit type and value. */ + void AddEdge(int from, int to, FSMEdge::EdgeType type, int32_t value); + + /*! + * \brief Add an epsilon transition between two states. + * \param from The source state. + * \param to The target state. + */ + void AddEpsilonEdge(int from, int to); + + /*! + * \brief Add a rule reference edge between states. + * \param from The source state. + * \param to The target state. + * \param rule_id The rule id to reference. + */ + void AddRuleEdge(int from, int to, int32_t rule_id); + + /*! + * \brief Add an EOS transition between two states. + * \param from The source state. + * \param to The target state. + */ + void AddEOSEdge(int from, int to); + + /*! + * \brief Add a repeat reference edge between states, allocating auxiliary data. + * \param from The source state. + * \param to The target state. + * \param rule_id The rule to repeat. + * \param lower Minimum repeat count. + * \param upper Maximum repeat count (-1 for unlimited). + */ + void AddRepeatEdge(int from, int to, int32_t rule_id, int32_t lower, int32_t upper); + + void AddTokenEdge(int from, int to, const std::vector& token_ids); + + void AddExcludeTokenEdge(int from, int to, const std::vector& token_ids); + + /*! \brief Get the edge auxiliary data. */ + const std::vector& GetEdgeAuxData() const; + + /*! \brief Set the edge auxiliary data (used during FSM construction). */ + void SetEdgeAuxData(std::vector data); + + /*! \brief Get repeat edge info by aux index. */ + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const; + + /*! \brief Get token edge info by aux index. */ + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const; + + /*! \brief Get exclude-token edge info by aux index. */ + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const; + + /*! + * \brief Add a whole FSM to the current FSM. + * \param fsm The FSM to be added. + * \param state_mapping The mapping from the state ids of the added FSM to the new ids in the + * current FSM. The result is cleared at the beginning. If the fsm's state id starts from 0, use + * it for efficiency. + */ + void AddFSM(const FSM& fsm, std::vector* state_mapping = nullptr); + + /****************** FSM Construction Methods ******************/ + + /*! + \brief Return a copy of the FSM. + */ + FSM Copy() const; + + /*! + * \brief Rebuild the FSM with the new state ids. + * \param state_mapping The mapping from the old state ids to the new state ids. + * \param new_num_states The new number of states. + * \return The rebuilt FSM. + */ + FSM RebuildWithMapping(const std::vector& state_mapping, int new_num_states) const; + + /*! + * \brief Sort the edges of the FSM by their min, max and target. + */ + void SortEdges(); + + /*! + * \brief Transform a FSM to a compact FSM. This method will first sort the edges of the FSM, + * then put all the edges into a compact array. + * \return The compact FSM. + */ + CompactFSM ToCompact(); + + XGRAMMAR_DEFINE_PIMPL_METHODS(FSM); +}; + +/*! + * \brief CompactFSM is the compact from of FSM. + * \details It uses Compact2DArray to store the edges, ensuring memory contiguity. It sorts all + * outgoing edges from a node according to their min and max values, so traversal can be faster. + * + * CompactFSM is immutable. If you need to modify a CompactFSM, you need to convert it to a FSM + * first, and convert it back after modification. + * + * It share the same set of visitor methods with FSM. + */ + +class CompactFSM { + public: + // for serialization only + CompactFSM() = default; + + explicit CompactFSM( + const Compact2DArray& edges, std::vector edge_aux_data = {} + ); + + explicit CompactFSM(Compact2DArray&& edges, std::vector edge_aux_data = {}); + + /****************** CompactFSM Visitors ******************/ + + /*! + * \brief Get the number of states in the FSM. + * \return The number of states in the FSM. + */ + int NumStates() const; + + /*! + * \brief Get the edges of the CompactFSM. + * \return The edges of the CompactFSM. + */ + const Compact2DArray& GetEdges() const; + + /*! + * \brief Get the edges of the CompactFSM. + * \param state The state to get the edges from. + * \return The edges of the CompactFSM. + */ + Compact2DArray::Row GetEdges(int state) const; + + /*! + * \brief Convert the edges of the CompactFSM to a string. Used in printing the CompactFSM. + * \return The string representation of the edges of the CompactFSM. + */ + std::string EdgesToString(std::optional> states = std::nullopt) const; + + /*! + * \brief Get the memory size of the CompactFSM. + * \param self The CompactFSM. + * \return The memory size of the CompactFSM. + */ + friend std::size_t MemorySize(const CompactFSM& self); + + /****************** CompactFSM Traversal Visitors ******************/ + + inline static constexpr int kNoNextState = -1; + + /*! + * \brief Advance the FSM from a given state based on an input character. If there are multiple + * transitions, the first one will be returned. + * \param from The source state to transition from. + * \param character The input character. + * \param targets The target states to be filled with the possible next states. + * \return The target state if a valid transition exists, kNoNextState otherwise. + */ + void GetNextStates( + int from, + int value, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + std::vector* targets = nullptr + ) const; + + /*! + * \brief Advance the FSM to the next state. + * \param from The current states. + * \param value The input value. + * \param result The possible next states. The result is cleared at the beginning. + * \param value_is_rule Whether the input value is a rule id. + * \param from_is_closure Whether from is an epsilon closure. + */ + void Advance( + const std::unordered_set& from, + int value, + std::unordered_set* result, + FSMEdge::EdgeType edge_type = FSMEdge::EdgeType::kCharRange, + bool from_is_closure = false + ) const; + + /*! + * \brief Get all the possible rule numbers for a given state. + * \param state_num The state number. + * \param rules The set of possible rule numbers. The result is cleared at the beginning. + */ + void GetPossibleRules(int state_num, std::unordered_set* rules) const; + + /*! + * \brief Get the epsilon closure of a set of states, i.e. those can be reached by epsilon + * transitions. + * \param state_set The states in the epsilon closure. The result is not cleared. + */ + void GetEpsilonClosure(std::unordered_set* state_set) const; + + /*! + * \brief Get the reachable states from a set of states. + * \param from The current states. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(const std::vector& from, std::unordered_set* result) const; + + /*! + * \brief Get the number of edges in the compact FSM. + * \return The number of edges. + */ + size_t GetNumEdges() const; + + /****************** CompactFSM Auxiliary Data ******************/ + + /*! \brief Get the edge auxiliary data. */ + const std::vector& GetEdgeAuxData() const; + + /*! \brief Set the edge auxiliary data (used during FSM construction). */ + void SetEdgeAuxData(std::vector data); + + /*! \brief Get repeat edge info by aux index. */ + RepeatEdgeRef GetRepeatEdgeInfo(int32_t idx) const; + + /*! \brief Get token edge info by aux index. */ + TokenEdgeRef GetTokenEdgeInfo(int32_t idx) const; + + /*! \brief Get exclude-token edge info by aux index. */ + ExcludeTokenEdgeRef GetExcludeTokenEdgeInfo(int32_t idx) const; + + /****************** CompactFSM Construction Methods ******************/ + + /*! + * \brief Transform the compact FSM to a FSM. + * \return The FSM. + */ + FSM ToFSM() const; + + friend picojson::value SerializeJSONValue(const CompactFSM& value); + friend std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(CompactFSM); +}; + +std::optional DeserializeJSONValue( + CompactFSM* result, const picojson::value& value, const std::string& type_name = "" +); + +class FSMWithStartEnd; +class FSMWithStartEndWithSize; +class CompactFSMWithStartEnd; +class CompactFSMWithStartEndWithSize; +struct CompactFSMWithStartEndWithSizeSerializeHelper; + +/*! + * \brief The base class for FSMWithStartEnd and CompactFSMWithStartEnd. It defines the + * common constructor and visitor methods. + */ +template +class FSMWithStartEndBase { + static_assert( + std::is_same_v || std::is_same_v, + "FSMType must be FSM or CompactFSM" + ); + + public: + // For serialization only + FSMWithStartEndBase() = default; + + FSMWithStartEndBase(const FSMType& fsm, int start, std::vector ends, bool is_dfa = false) + : fsm_(fsm), start_(start), ends_(std::move(ends)), is_dfa_(is_dfa) { + NormalizeEnds(); + } + + /****************** Member Accessors and Mutators ******************/ + + /*! \brief Returns the underlying FSM. */ + const FSMType& GetFsm() const { return fsm_; } + + /*! \brief Returns the start state of the FSM. */ + int GetStart() const { return start_; } + + /*! + * \brief Returns the end states of the FSM as a sorted, deduplicated list of state ids. + * \note End states are stored sparsely because an FSM view over a large shared FSM (e.g. the + * per-rule FSMs over the complete FSM of a grammar) usually has only a few end states, while + * the state id space can be very large. + */ + const std::vector& GetEnds() const { return ends_; } + + /*! + * \brief Checks if a given state is an end/accepting state. + * \param state The state to check. + * \return True if the state is an end state, false otherwise. + */ + bool IsEndState(int state) const { return std::binary_search(ends_.begin(), ends_.end(), state); } + + /*! \brief Check if a state is scanable. + * \param state The state to check. + * \return True if the state is scanable, false otherwise. + */ + bool IsScanableState(int state) const { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsCharRange() || edge.IsToken() || edge.IsExcludeToken()) { + return true; + } + } + return false; + } + + /*! + * \brief Check if a state is not terminal. + * \param state The state to check. + * \return True if the state is scanable, false otherwise. + */ + bool IsNonTerminalState(int state) const { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsRuleRef() || edge.IsEpsilon() || edge.IsRepeatRef()) { + return true; + } + } + return false; + } + + /*! + * \brief Sets the start state of the FSM. + * \param state The state to set as the start state. + */ + void SetStartState(int state) { + XGRAMMAR_DCHECK(state < NumStates()); + start_ = state; + } + + /*! + * \brief Adds an end/accepting state to the FSM. + * \param state The state to add as an end state. + */ + void AddEndState(int state) { + XGRAMMAR_DCHECK(state < NumStates()); + auto it = std::lower_bound(ends_.begin(), ends_.end(), state); + if (it == ends_.end() || *it != state) { + ends_.insert(it, state); + } + } + + /*! + * \brief Adds a new state to the FSM and marks it as non-end. + * \return The index of the newly added state. + */ + int AddState() { return fsm_.AddState(); } + + /*! + * \brief Sets the end states of the FSM. + * \param ends The new end states, as a list of state ids. Need not be sorted or deduplicated. + */ + void SetEndStates(std::vector ends) { + ends_ = std::move(ends); + NormalizeEnds(); + } + + /*! \brief Returns the total number of states in the FSM. */ + int NumStates() const { return fsm_.NumStates(); } + + /*! + * \brief Access the methods of the underlying FSM. + */ + FSMType& GetFsm() { return fsm_; } + + /****************** FSM Traversal Algorithms ******************/ + + /*! + * \brief Check if the FSM accepts the string. + * \param str The input string. + * \return True if the FSM accepts the string, false otherwise. + */ + bool AcceptString(const std::string& str) const; + + /*! + * \brief Get the reachable states from the start state. + * \param result The reachable states. The result is cleared at the beginning. + */ + void GetReachableStates(std::unordered_set* result) const; + + /*! + * \brief Check if the FSM is a leaf FSM. + * \return True if the FSM is a leaf FSM, false otherwise. + */ + bool IsLeaf() const; + + protected: + /*! \brief Sort and deduplicate the end states to maintain the sorted invariant. */ + void NormalizeEnds() { + std::sort(ends_.begin(), ends_.end()); + ends_.erase(std::unique(ends_.begin(), ends_.end()), ends_.end()); + } + + /*! \brief The underlying finite state machine. */ + FSMType fsm_; + /*! \brief The start state of the FSM. */ + int start_; + + /*! + * \brief The set of accepting/end states, stored as a sorted, deduplicated list of state ids. + * \note Stored sparsely (rather than as a bitvector over all states) because FSM views over a + * large shared FSM usually have only a few end states, while the state id space can be huge. + */ + std::vector ends_; + + protected: + /*! \brief Whether this FSM is a deterministic finite automaton. */ + bool is_dfa_ = false; +}; + +/*! + * \brief FSMWithStartEnd represents a FSM with start and end states. + * \details It stores a pointer to a FSM, a start state, and a set of end states. Multiple + * FSMWithStartEnd can share the same FSM. It also provides a set of methods to construct FSMs. + */ +class FSMWithStartEnd : public FSMWithStartEndBase { + public: + using FSMWithStartEndBase::FSMWithStartEndBase; + + /*! + * \brief Convert the FSMWithStartEnd to a string. Only considers the nodes approachable from the + * start state. + * \return The string representation of the FSMWithStartEnd. + */ + std::string ToString() const; + + friend std::ostream& operator<<(std::ostream& os, const FSMWithStartEnd& fsm); + + /****************** FSM Construction Methods ******************/ + + /*! + * \brief Return a copy of the FSMWithStartEnd. + */ + FSMWithStartEnd Copy() const; + + /*! + * \brief Rebuild the FSM with the new state ids. + * \param state_mapping The mapping from old state ids to new state ids. + * \param new_num_states The new number of states. + */ + FSMWithStartEnd RebuildWithMapping(const std::vector& state_mapping, int new_num_states) + const; + + /*! + * \brief Add the underlying FSM to another complete FSM that could contain multiple FSMs. + * Return a new FSMWithStartEnd that points to the complete FSM and whose start and ends are + * mapped to the states in the complete FSM. + * \param complete_fsm The complete FSM. + * \param state_mapping The mapping from the old state ids to the new state ids. The result is + * cleared at the beginning. Should not be nullptr. + * \return The FSMWithStartEnd that points to the complete FSM. + */ + FSMWithStartEndWithSize AddToCompleteFSM(FSM* complete_fsm, std::vector* state_mapping); + + /*! + * \brief Transform the FSMWithStartEnd to a CompactFSMWithStartEnd. + * \return The CompactFSMWithStartEnd. + */ + CompactFSMWithStartEnd ToCompact(); + + /****************** FSM Algorithms ******************/ + + /*! + * \brief Return a new FSM representing FSM* + * \return The FSM that accepts FSM*. + */ + FSMWithStartEnd Star() const; + + /*! + * \brief Return a new FSM representing rule1+. + * \return The FSM that accepts rule1+. + */ + FSMWithStartEnd Plus() const; + + /*! + * \brief Return a new FSM representing rule1?. + * \return The FSM that accepts rule1?. + */ + FSMWithStartEnd Optional() const; + + /*! + * \brief Return a new FSM representing the complement of the language. + * \return The complement FSM. + */ + Result Not(int max_result_num_states = 1e6) const; + + /*! + * \brief Intersect the FSMs. + * \param lhs The left FSM. + * \param rhs The right FSM. + * \return The intersection of the FSMs. + */ + static Result Intersect( + const FSMWithStartEnd& lhs, const FSMWithStartEnd& rhs, int max_result_num_states = 1e6 + ); + + /*! + * \brief Union the FSMs. + * \param fsms The FSMs to be unioned. + * \return The union of the FSMs. + */ + static FSMWithStartEnd Union(const std::vector& fsms); + + /*! + * \brief Concatenate the FSMs. + * \param fsms The FSMs to be concatenated, which should be in order. + * \return The concatenation of the FSMs. + */ + static FSMWithStartEnd Concat(const std::vector& fsms); + + /*! + * \brief Check if the FSM is a DFA. + * \return True if the FSM is a DFA, false otherwise. + */ + bool IsDFA(); + + /*! + * \brief Merge some states by removing some epsilon transitions. + * \details If a --\epsilon--> b, and either 1) b doesn't have any other inward edges, or + * 2) a doesn't have any other outward edges, we can merge a and b. + */ + FSMWithStartEnd SimplifyEpsilon(int max_num_states = 1e8) const; + + /*! + * \brief Merge equivalent states in the FSM. + * \details If two states are 1) pointed to by edges with the same label from the same state, and + * 2) they are not pointed to by other edges, then we can merge them. + * \example n0 --(c)--> n1, n0 --(c)--> n2, then we can merge n1 and n2. + */ + FSMWithStartEnd MergeEquivalentStates(int max_num_states = 1e5) const; + + /*! + * \brief Transform the FSM to a DFA. + * \param max_result_num_states The maximum number of states in the DFA. + * \return The DFA. + */ + Result ToDFA(int max_num_states = 1e3) const; + + /*! + * \brief Minimize the DFA. + * \param max_result_num_states The maximum number of states in the DFA. + * \return The minimized DFA. + */ + Result MinimizeDFA(int max_num_states = 1e3) const; +}; + +/*! + * \brief Wrapper that bundles an FSMWithStartEnd with explicit size metadata. It is + * used when we want to store the number of edges and nodes in the part of the FSM, instead + * of the completed FSMWithStartEnd. + */ +class FSMWithStartEndWithSize { + public: + // For serialization only + FSMWithStartEndWithSize() = default; + + explicit FSMWithStartEndWithSize(FSMWithStartEnd fsm, int edge_num, int node_num) + : fsm_(std::move(fsm)), edge_num_(edge_num), node_num_(node_num) {} + + const FSMWithStartEnd& GetFsm() const { return fsm_; } + int GetEdgeNum() const { return edge_num_; } + int GetNodeNum() const { return node_num_; } + + private: + FSMWithStartEnd fsm_; + int edge_num_ = 0; + int node_num_ = 0; +}; + +/*! + * \brief A class that represents a compact-form FSM with a start state and a set of end states. + * \details CompactFSMWithStartEnd stores a pointer to a CompactFSM, a start state, and a set of end + * states. Multiple CompactFSMWithStartEnd can share the same CompactFSM. It share the same set of + * visitor methods with FSMWithStartEnd. + */ +class CompactFSMWithStartEnd : public FSMWithStartEndBase { + public: + // For serialization only + CompactFSMWithStartEnd() = default; + + explicit CompactFSMWithStartEnd(const CompactFSM& fsm, int start, std::vector ends) + : FSMWithStartEndBase(fsm, start, std::move(ends)), + edge_num_(fsm.GetNumEdges()) {} + + using FSMWithStartEndBase::FSMWithStartEndBase; + + /*! + * \brief Convert the FSMWithStartEnd to a string. Only considers the nodes approachable from the + * start state. + * \return The string representation of the FSMWithStartEnd. + */ + std::string ToString() const; + + /*! + * \brief Transform the CompactFSMWithStartEnd to a FSMWithStartEnd. + * \return The FSMWithStartEnd. + */ + FSMWithStartEnd ToFSM() const; + + private: + size_t edge_num_ = 0; + + /*! + * \brief Print the CompactFSMWithStartEnd. + * \param os The output stream. + * \param fsm The CompactFSMWithStartEnd. + * \return The output stream. + */ + friend std::ostream& operator<<(std::ostream& os, const CompactFSMWithStartEnd& fsm); + + /*! + * \brief Get the memory size of the CompactFSMWithStartEnd. + * \param self The CompactFSMWithStartEnd. + * \return The memory size of the CompactFSMWithStartEnd. + */ + friend std::size_t MemorySize(const CompactFSMWithStartEnd& self); + + friend struct member_trait; + + friend struct CompactFSMWithStartEndSerializeHelper; + + friend picojson::value SerializeJSONValue(const CompactFSMWithStartEnd& value); + friend std::optional DeserializeJSONValue( + CompactFSMWithStartEnd* result, const picojson::value& value, const std::string& type_name + ); +}; + +/*! + * \brief Wrapper that bundles a CompactFSMWithStartEnd with explicit size metadata. It is + * used when we want to store the number of edges and nodes in the part of the + * CompactFSMWithStartEnd, instead of the completed CompactFSMWithStartEnd. + */ +class CompactFSMWithStartEndWithSize { + public: + // For serialization only + CompactFSMWithStartEndWithSize() = default; + + explicit CompactFSMWithStartEndWithSize(CompactFSMWithStartEnd fsm, int edge_num, int node_num) + : fsm_(std::move(fsm)), edge_num_(edge_num), node_num_(node_num) {} + + const CompactFSMWithStartEnd& GetFsm() const { return fsm_; } + int GetEdgeNum() const { return edge_num_; } + int GetNodeNum() const { return node_num_; } + + friend picojson::value SerializeJSONValue(const CompactFSMWithStartEndWithSize& value); + friend std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name + ); + + private: + CompactFSMWithStartEnd fsm_; + int edge_num_ = 0; + int node_num_ = 0; + + friend std::size_t MemorySize(const CompactFSMWithStartEndWithSize& self) { + return MemorySize(self.fsm_) + sizeof(self.edge_num_) + sizeof(self.node_num_); + } + + friend struct CompactFSMWithStartEndWithSizeSerializeHelper; +}; + +std::optional DeserializeJSONValue( + CompactFSMWithStartEndWithSize* result, + const picojson::value& value, + const std::string& type_name = "" +); + +/****************** FSMWithStartEndBase Template Implementation ******************/ + +template +inline bool FSMWithStartEndBase::AcceptString(const std::string& str) const { + std::unordered_set start_states{start_}; + fsm_.GetEpsilonClosure(&start_states); + std::unordered_set result_states; + for (const auto& character : str) { + result_states.clear(); + fsm_.Advance( + start_states, + static_cast(static_cast(character)), + &result_states, + FSMEdge::EdgeType::kCharRange, + false + ); + if (result_states.empty()) { + return false; + } + start_states = result_states; + } + return std::any_of(start_states.begin(), start_states.end(), [&](int state) { + return IsEndState(state); + }); +} + +template +inline void FSMWithStartEndBase::GetReachableStates(std::unordered_set* result +) const { + return fsm_.GetReachableStates({start_}, result); +} + +template +inline bool FSMWithStartEndBase::IsLeaf() const { + std::unordered_set reachable_states; + GetReachableStates(&reachable_states); + for (const auto& state : reachable_states) { + for (const auto& edge : fsm_.GetEdges(state)) { + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + return false; + } + } + } + return true; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_FSM_H_ diff --git a/zig/vendor/xgrammar/cpp/fsm_builder.cc b/zig/vendor/xgrammar/cpp/fsm_builder.cc new file mode 100644 index 0000000000..c71f9c8c50 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/fsm_builder.cc @@ -0,0 +1,1814 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm_builder.cc + */ +#include "fsm_builder.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm.h" +#include "grammar_builder.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/******************** Packed UTF-8 range helpers ********************/ + +uint32_t CodepointToPackedUTF8(uint32_t codepoint) { + if (codepoint <= 0x7F) { + // 1-byte sequence (ASCII) + return codepoint; + } else if (codepoint <= 0x7FF) { + // 2-byte sequence: byte0 = 110xxxxx, byte1 = 10xxxxxx + uint8_t byte0 = 0xC0 | ((codepoint >> 6) & 0x1F); + uint8_t byte1 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 8) | byte1; + } else if (codepoint <= 0xFFFF) { + // 3-byte sequence: byte0 = 1110xxxx, byte1 = 10xxxxxx, byte2 = 10xxxxxx + uint8_t byte0 = 0xE0 | ((codepoint >> 12) & 0x0F); + uint8_t byte1 = 0x80 | ((codepoint >> 6) & 0x3F); + uint8_t byte2 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 16) | (static_cast(byte1) << 8) | byte2; + } else { + // 4-byte sequence: byte0 = 11110xxx, byte1-3 = 10xxxxxx + uint8_t byte0 = 0xF0 | ((codepoint >> 18) & 0x07); + uint8_t byte1 = 0x80 | ((codepoint >> 12) & 0x3F); + uint8_t byte2 = 0x80 | ((codepoint >> 6) & 0x3F); + uint8_t byte3 = 0x80 | (codepoint & 0x3F); + return (static_cast(byte0) << 24) | (static_cast(byte1) << 16) | + (static_cast(byte2) << 8) | byte3; + } +} + +// This function will add a range [min, max] of characters to the FSM, and the length +// of the characters are the same. +static void AddSameLengthCharacterRange(FSM& fsm, int from, int to, uint32_t min, uint32_t max) { + uint8_t byte_min[4] = { + static_cast(min & 0xFF), + static_cast(min >> 8), + static_cast(min >> 16), + static_cast(min >> 24) + }; + uint8_t byte_max[4] = { + static_cast(max & 0xFF), + static_cast(max >> 8), + static_cast(max >> 16), + static_cast(max >> 24) + }; + + // ASCII. + if (byte_max[1] == 0) { + fsm.AddEdge(from, to, byte_min[0], byte_max[0]); + return; + } + + if (byte_max[3] != 0) { + // 4-byte unicode. + if (byte_max[3] == byte_min[3]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[3], byte_max[3]); + min = (min & 0x00FFFFFF); + max = (max & 0x00FFFFFF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FFFFFF) != 0x808080) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[3], byte_min[3]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFFFF), 0x00BFBFBF); + } else { + byte_min[3]--; + } + if ((max & 0x00FFFFFF) != 0xBFBFBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[3], byte_max[3]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x00808080, (max & 0x00FFFFFF)); + } else { + byte_max[3]++; + } + if (byte_max[3] - byte_min[3] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[3] + 1, byte_max[3] - 1); + int tmp_state_mid2 = fsm.AddState(); + // Second byte. + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + int tmp_state_mid3 = fsm.AddState(); + // Third byte. + fsm.AddEdge(tmp_state_mid2, tmp_state_mid3, 0x80, 0xBF); + // Last byte. + fsm.AddEdge(tmp_state_mid3, to, 0x80, 0xBF); + } + return; + } + if (byte_max[2] != 0) { + // 3 byte unicode. + if (byte_max[2] == byte_min[2]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[2], byte_max[2]); + min = (min & 0x00FFFF); + max = (max & 0x00FFFF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FFFF) != 0x8080) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[2], byte_min[2]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FFFF), 0x00BFBF); + } else { + byte_min[2]--; + } + if ((max & 0x00FFFF) != 0xBFBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[2], byte_max[2]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x008080, (max & 0x00FFFF)); + } else { + byte_max[2]++; + } + if (byte_max[2] - byte_min[2] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[2] + 1, byte_max[2] - 1); + int tmp_state_mid2 = fsm.AddState(); + // Second byte. + fsm.AddEdge(tmp_state_mid, tmp_state_mid2, 0x80, 0xBF); + // Last byte. + fsm.AddEdge(tmp_state_mid2, to, 0x80, 0xBF); + } + return; + } + + // 2 byte unicode. + if (byte_max[1] == byte_min[1]) { + int tmp_state = fsm.AddState(); + fsm.AddEdge(from, tmp_state, byte_min[1], byte_max[1]); + min = (min & 0x00FF); + max = (max & 0x00FF); + AddSameLengthCharacterRange(fsm, tmp_state, to, min, max); + return; + } + if ((min & 0x00FF) != 0x80) { + int tmp_state_min = fsm.AddState(); + fsm.AddEdge(from, tmp_state_min, byte_min[1], byte_min[1]); + AddSameLengthCharacterRange(fsm, tmp_state_min, to, (min & 0x00FF), 0x00BF); + } else { + byte_min[1]--; + } + if ((max & 0x00FF) != 0xBF) { + int tmp_state_max = fsm.AddState(); + fsm.AddEdge(from, tmp_state_max, byte_max[1], byte_max[1]); + AddSameLengthCharacterRange(fsm, tmp_state_max, to, 0x0080, (max & 0x00FF)); + } else { + byte_max[1]++; + } + if (byte_max[1] - byte_min[1] > 1) { + int tmp_state_mid = fsm.AddState(); + // First byte. + fsm.AddEdge(from, tmp_state_mid, byte_min[1] + 1, byte_max[1] - 1); + fsm.AddEdge(tmp_state_mid, to, 0x80, 0xBF); + } + return; +} + +void AddPackedUTF8RangeEdges(FSM& fsm, int from, int to, uint32_t min, uint32_t max) { + XGRAMMAR_CHECK(min <= max) << "Invalid character range: min (" << min << ") > max (" << max + << ")"; + // Ensure max and min are valid unicode value. + if (max > kMax4BytesUnicode) { + max = kMax4BytesUnicode; + } else if (max > kMax3BytesUnicode) { + if (max < kMin4BytesUnicode) { + max = kMax3BytesUnicode; + } + } else if (max > kMax2BytesUnicode) { + if (max < kMin3BytesUnicode) { + max = kMax2BytesUnicode; + } + } else if (max < kMin2BytesUnicode && (max > kMax1ByteUnicode)) { + max = kMax1ByteUnicode; + } + + if (min > kMax4BytesUnicode) { + min = kMax4BytesUnicode; + } else if (min > kMax3BytesUnicode) { + if (min < kMin4BytesUnicode) { + min = kMin4BytesUnicode; + } + } else if (min > kMax2BytesUnicode) { + if (min < kMin3BytesUnicode) { + min = kMin3BytesUnicode; + } + } else if (min < kMin2BytesUnicode && (min > kMax1ByteUnicode)) { + min = kMin2BytesUnicode; + } + + // Step2. Divide the range into several ranges, which contain characters with different lengths. + if (max <= kMax1ByteUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + return; + } + if (max <= kMax2BytesUnicode) { + if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, max); + } + return; + } + if (max <= kMax3BytesUnicode) { + if (min >= kMin3BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, max); + } + return; + } + XGRAMMAR_CHECK(max <= kMax4BytesUnicode); + if (min >= kMin4BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, max); + } else if (min >= kMin3BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } else if (min >= kMin2BytesUnicode) { + AddSameLengthCharacterRange(fsm, from, to, min, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } else { + AddSameLengthCharacterRange(fsm, from, to, min, kMax1ByteUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin2BytesUnicode, kMax2BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin3BytesUnicode, kMax3BytesUnicode); + AddSameLengthCharacterRange(fsm, from, to, kMin4BytesUnicode, max); + } + return; +} + +std::string RewriteRegexDots(const std::string& pattern, bool dot_matches_newline) { + if (dot_matches_newline) { + return pattern; + } + std::string result; + result.reserve(pattern.size()); + bool escaped = false; + bool in_character_class = false; + for (char c : pattern) { + if (escaped) { + result.push_back(c); + escaped = false; + continue; + } + if (c == '\\') { + result.push_back(c); + escaped = true; + } else if (c == '[') { + result.push_back(c); + in_character_class = true; + } else if (c == ']' && in_character_class) { + result.push_back(c); + in_character_class = false; + } else if (c == '.' && !in_character_class) { + result += "[^\\n]"; + } else { + result.push_back(c); + } + } + return result; +} + +/******************** Codepoint range utilities ********************/ + +namespace { + +constexpr uint32_t kMaxCodepoint = 0x10FFFF; + +/*! \brief Bounded repetitions above this threshold are compiled into a repeat FSM edge (when a + * GrammarBuilder is available) instead of being physically unrolled. Matches the unroll threshold + * of the grammar-level RepetitionRangeExpander. */ +constexpr int kLargeRepeatThreshold = 128; + +/*! \brief Hard cap of the estimated state count when a bounded repetition has to be unrolled + * because no GrammarBuilder is available. */ +constexpr int64_t kMaxUnrolledRepeatStates = 100000; + +using CodepointRange = std::pair; + +/*! \brief Sort the ranges and merge overlapping or adjacent ones. */ +void NormalizeRanges(std::vector* ranges) { + std::sort(ranges->begin(), ranges->end()); + std::vector result; + for (const auto& range : *ranges) { + if (!result.empty() && range.first <= result.back().second + 1 && + range.first >= result.back().first) { + result.back().second = std::max(result.back().second, range.second); + } else { + result.push_back(range); + } + } + *ranges = std::move(result); +} + +/*! \brief Complement normalized ranges over the codepoint domain [0, kMaxCodepoint]. */ +std::vector ComplementRanges(const std::vector& ranges) { + std::vector result; + uint32_t next = 0; + for (const auto& range : ranges) { + if (range.first > next) { + result.push_back({next, range.first - 1}); + } + if (range.second >= kMaxCodepoint) { + return result; + } + next = std::max(next, range.second + 1); + } + result.push_back({next, kMaxCodepoint}); + return result; +} + +/*! \brief Append the ASCII case-folded counterparts of every letter contained in the ranges. */ +void FoldAsciiCaseRanges(std::vector* ranges) { + size_t original_size = ranges->size(); + for (size_t i = 0; i < original_size; ++i) { + uint32_t low = (*ranges)[i].first; + uint32_t high = (*ranges)[i].second; + uint32_t fold_low = std::max(low, 'a'); + uint32_t fold_high = std::min(high, 'z'); + if (fold_low <= fold_high) { + ranges->push_back({fold_low - ('a' - 'A'), fold_high - ('a' - 'A')}); + } + fold_low = std::max(low, 'A'); + fold_high = std::min(high, 'Z'); + if (fold_low <= fold_high) { + ranges->push_back({fold_low + ('a' - 'A'), fold_high + ('a' - 'A')}); + } + } +} + +/*! \brief Add edges from `from` to `to` accepting the UTF-8 encoding of every codepoint in the + * normalized ranges. Multi-byte characters get intermediate states. */ +void AddCodepointRangesToFSM( + FSM* fsm, int from, int to, const std::vector& ranges +) { + for (const auto& [low, high] : ranges) { + if (low <= kMax1ByteUnicode) { + fsm->AddEdge(from, to, low, std::min(high, kMax1ByteUnicode)); + } + if (high > kMax1ByteUnicode) { + uint32_t multi_byte_low = std::max(low, kMax1ByteUnicode + 1); + AddPackedUTF8RangeEdges( + *fsm, from, to, CodepointToPackedUTF8(multi_byte_low), CodepointToPackedUTF8(high) + ); + } + } +} + +/*! \brief One parsed regex escape (or literal): either a single codepoint, or a (possibly + * negated) set of codepoint ranges for class escapes like \d, \D, \w, \W, \s, \S. */ +struct RegexEscapeItem { + std::vector ranges; + bool negated = false; + bool is_single = false; + uint32_t codepoint = 0; +}; + +/*! + * \brief Parse the escape sequence starting at regex[*pos] == '\\'. On success, *pos is advanced + * past the escape sequence. + * \param in_class Whether the escape appears inside a character class ([...]). Inside a class, + * \b means the backspace character instead of a word boundary assertion. + */ +Result ParseRegexEscape(const std::string& regex, size_t* pos, bool in_class) { + XGRAMMAR_DCHECK(regex[*pos] == '\\'); + if (*pos + 1 >= regex.size()) { + return ResultErr("Regex ends with a trailing backslash"); + } + char escaped = regex[*pos + 1]; + *pos += 2; + RegexEscapeItem item; + auto single = [&](uint32_t codepoint) { + item.is_single = true; + item.codepoint = codepoint; + return ResultOk(std::move(item)); + }; + switch (escaped) { + case 'n': + return single('\n'); + case 't': + return single('\t'); + case 'r': + return single('\r'); + case 'f': + return single('\f'); + case 'v': + return single('\v'); + case 'a': + return single('\a'); + case 'e': + return single(0x1B); + case '0': + return single(0); + case 'b': + if (in_class) { + return single(0x08); + } + return ResultErr("Word boundary assertion \\b is not supported in regex"); + case 'B': + return ResultErr("Word boundary assertion \\B is not supported in regex"); + case 'p': + case 'P': + return ResultErr("Unicode property escape \\p / \\P is not supported in regex"); + case 'k': + return ResultErr("Backreference \\k is not supported in regex"); + case '1': + case '2': + case '3': + case '4': + case '5': + case '6': + case '7': + case '8': + case '9': + return ResultErr("Backreference \\" + std::string(1, escaped) + " is not supported in regex"); + case 'd': + item.ranges = {{'0', '9'}}; + return ResultOk(std::move(item)); + case 'D': + item.ranges = {{'0', '9'}}; + item.negated = true; + return ResultOk(std::move(item)); + case 'w': + item.ranges = {{'0', '9'}, {'A', 'Z'}, {'_', '_'}, {'a', 'z'}}; + return ResultOk(std::move(item)); + case 'W': + item.ranges = {{'0', '9'}, {'A', 'Z'}, {'_', '_'}, {'a', 'z'}}; + item.negated = true; + return ResultOk(std::move(item)); + case 's': + item.ranges = {{0x09, 0x0D}, {' ', ' '}}; + return ResultOk(std::move(item)); + case 'S': + item.ranges = {{0x09, 0x0D}, {' ', ' '}}; + item.negated = true; + return ResultOk(std::move(item)); + case 'x': { + if (*pos + 1 >= regex.size()) { + return ResultErr("\\x must be followed by two hexadecimal digits in regex"); + } + int high_digit = HexCharToInt(regex[*pos]); + int low_digit = HexCharToInt(regex[*pos + 1]); + if (high_digit < 0 || low_digit < 0) { + return ResultErr("\\x must be followed by two hexadecimal digits in regex"); + } + *pos += 2; + return single(high_digit * 16 + low_digit); + } + case 'u': { + if (*pos < regex.size() && regex[*pos] == '{') { + size_t close = regex.find('}', *pos + 1); + if (close == std::string::npos || close == *pos + 1 || close > *pos + 7) { + return ResultErr("\\u{...} must contain one to six hexadecimal digits in regex"); + } + uint32_t codepoint = 0; + for (size_t i = *pos + 1; i < close; ++i) { + int digit = HexCharToInt(regex[i]); + if (digit < 0) { + return ResultErr("\\u{...} must contain one to six hexadecimal digits in regex"); + } + codepoint = codepoint * 16 + digit; + } + if (codepoint > kMaxCodepoint) { + return ResultErr("\\u{...} escape is beyond the Unicode range in regex"); + } + *pos = close + 1; + return single(codepoint); + } + if (*pos + 3 >= regex.size()) { + return ResultErr("\\u must be followed by four hexadecimal digits in regex"); + } + uint32_t codepoint = 0; + for (size_t i = *pos; i < *pos + 4; ++i) { + int digit = HexCharToInt(regex[i]); + if (digit < 0) { + return ResultErr("\\u must be followed by four hexadecimal digits in regex"); + } + codepoint = codepoint * 16 + digit; + } + *pos += 4; + return single(codepoint); + } + case 'c': { + if (*pos >= regex.size() || !std::isalpha(static_cast(regex[*pos]))) { + return ResultErr("\\c must be followed by a letter in regex"); + } + uint32_t codepoint = static_cast(regex[*pos]) & 0x1F; + *pos += 1; + return single(codepoint); + } + default: { + if (static_cast(escaped) >= 0x80) { + // Multi-byte UTF-8 character after the backslash: match it literally. + *pos -= 1; + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + *pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || *pos + num_bytes > regex.size()) { + return ResultErr("Invalid UTF-8 in regex escape sequence"); + } + *pos += num_bytes; + return single(static_cast(codepoint)); + } + if (std::isalnum(static_cast(escaped))) { + XGRAMMAR_LOG(WARNING) << "Escape sequence \\" << escaped + << " is not recognized in regex; matching the character literally"; + } + return single(static_cast(escaped)); + } + } +} + +/*! + * \brief Parse a character class leaf "[...]" into the final set of accepted codepoint ranges. + * ASCII case folding (if requested) is applied before negation. + */ +Result> ParseCharacterClassLeaf( + const std::string& regex, bool case_insensitive +) { + XGRAMMAR_DCHECK(regex.size() >= 2 && regex.front() == '[' && regex.back() == ']'); + size_t pos = 1; + size_t content_end = regex.size() - 1; + bool negated = false; + if (pos < content_end && regex[pos] == '^') { + negated = true; + ++pos; + } + if (pos == content_end) { + return ResultErr("Empty character class " + regex + " is not allowed in regex"); + } + + // Parse one class unit: an escape sequence or a literal (possibly multi-byte) character. + auto parse_unit = [&]() -> Result { + if (regex[pos] == '\\') { + return ParseRegexEscape(regex, &pos, /*in_class=*/true); + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || pos + num_bytes > content_end) { + return ResultErr("Invalid UTF-8 in regex character class " + regex); + } + pos += num_bytes; + RegexEscapeItem item; + item.is_single = true; + item.codepoint = static_cast(codepoint); + return ResultOk(std::move(item)); + }; + + std::vector ranges; + while (pos < content_end) { + auto unit_result = parse_unit(); + if (unit_result.IsErr()) { + return ResultErr(std::move(unit_result).UnwrapErr()); + } + auto unit = std::move(unit_result).Unwrap(); + if (!unit.is_single) { + // Class escapes like \d cannot be a range endpoint; a following '-' is literal. + if (unit.negated) { + NormalizeRanges(&unit.ranges); + auto complement = ComplementRanges(unit.ranges); + ranges.insert(ranges.end(), complement.begin(), complement.end()); + } else { + ranges.insert(ranges.end(), unit.ranges.begin(), unit.ranges.end()); + } + continue; + } + if (pos < content_end && regex[pos] == '-' && pos + 1 < content_end) { + ++pos; + auto high_result = parse_unit(); + if (high_result.IsErr()) { + return ResultErr(std::move(high_result).UnwrapErr()); + } + auto high_unit = std::move(high_result).Unwrap(); + if (!high_unit.is_single) { + return ResultErr("Invalid character range endpoint in regex character class " + regex); + } + if (high_unit.codepoint < unit.codepoint) { + return ResultErr( + "Invalid character range (lower bound exceeds upper bound) in regex character class " + + regex + ); + } + ranges.push_back({unit.codepoint, high_unit.codepoint}); + } else { + ranges.push_back({unit.codepoint, unit.codepoint}); + } + } + + if (case_insensitive) { + FoldAsciiCaseRanges(&ranges); + } + NormalizeRanges(&ranges); + if (negated) { + ranges = ComplementRanges(ranges); + } + return ResultOk(std::move(ranges)); +} + +} // namespace + +/******************** RegexIR ********************/ + +class RegexIR { + public: + struct Leaf; + + struct Symbol; + + struct Union; + + struct Bracket; + + struct Repeat; + + struct RuleRefNode; + + struct RepeatSubrule; + + static constexpr int kRepeatNoUpperBound = -1; + + using State = std::variant; + + // This struct is used to store one atom of the regex: the empty string (regex == ""), a + // character class (regex == "[...]"), or a short sequence of literal characters / escapes. + struct Leaf { + std::string regex; + }; + + // This struct is used to store the symbol in regex, i.e. + // +, *, ? + enum class RegexSymbol { + star, + plus, + optional, + }; + + struct Bracket { + std::vector states; + }; + + struct Symbol { + RegexSymbol symbol; + std::vector state; + }; + + // This struct is used to represent a union symbol. + struct Union { + std::vector states; + }; + + struct Repeat { + std::vector states; + int lower_bound = 0; + int upper_bound = 0; + }; + + // A reference to a grammar rule, compiled into a kRuleRef FSM edge. + struct RuleRefNode { + int32_t rule_id; + }; + + // A bounded repetition of a grammar rule, compiled into a kRepeatRef FSM edge. The referenced + // rule holds the repeated sub-pattern; the Earley parser executes the repetition with a + // counter at runtime, so no FSM unrolling happens. + struct RepeatSubrule { + int32_t rule_id; + int lower_bound = 0; + int upper_bound = 0; + }; + + // The top-level sequence of the regex. + std::vector states; + + // Whether matching is ASCII case-insensitive (enabled by a leading "(?i)"). + bool case_insensitive = false; + + /*! + \brief Constructs a NFA from the regex IR. + */ + Result Build() const; + + /*! + \brief the visit function for the variant. + */ + Result visit(const Leaf& state) const; + + Result visit(const Symbol& state) const; + + Result visit(const Union& state) const; + + Result visit(const Bracket& state) const; + + Result visit(const Repeat& state) const; + + Result visit(const RuleRefNode& state) const; + + Result visit(const RepeatSubrule& state) const; + + /*! \brief Whether the IR node can match the empty string. Purely syntactic; no FSM is built. */ + static bool IsNullable(const State& state); + + /*! \brief Whether a sequence of IR nodes can match the empty string. */ + static bool IsNullableSequence(const std::vector& states); + + /*! + * \brief Check repeat in regex. i.e {...} and {...,...} + * \param regex The regex string. + * \param start The start position of the repeat. i.e. regex[start] == '{'. + * After the function, start will be the position of '}'. + * \return The repeat range. + */ + static Result> CheckRepeat(const std::string& regex, int& start); + + private: + /*! + * \brief Construct a FSM from a regex leaf. + * \details The leaf is the empty string, a character class like [a-c0-9], or a sequence of + * literal characters / escapes like "ab\n". Any symbols like "a|b" or "a*b" are not supported. + * \param regex The regex string. + * \return The FSM with start and end states. + */ + Result BuildLeafFSMFromRegex(const std::string& regex) const; + + /*! + * \brief Add the transition(s) accepting a single codepoint (with ASCII case folding when + * case_insensitive is set) from `current` to a new state, and return the new state. + */ + int AddSingleCodepoint(FSMWithStartEnd& result, int current, uint32_t codepoint) const; +}; + +Result> RegexIR::CheckRepeat(const std::string& regex, int& start) { + // 10^9 fits in an int; longer counts would overflow. + constexpr size_t kMaxRepeatDigits = 9; + if (regex[start] != '{') { + return ResultErr("Invalid repetition: expected '{'"); + } + int lower_bound = 0; + int upper_bound = RegexIR::kRepeatNoUpperBound; + std::string num_str; + XGRAMMAR_DCHECK(regex[start] == '{'); + start++; + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + while (static_cast(start) < regex.size() && std::isdigit(regex[start])) { + num_str += regex[start]; + start++; + } + if (num_str.empty()) { + return ResultErr("Invalid repetition count: expected a number after '{'"); + } + if (num_str.size() > kMaxRepeatDigits) { + return ResultErr("Invalid repetition count: the count " + num_str + " is too large"); + } + lower_bound = std::stoi(num_str); + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + // The format is {n} + if (regex[start] == '}') { + upper_bound = lower_bound; + return ResultOk(std::make_pair(lower_bound, upper_bound)); + } + if (regex[start] != ',') { + return ResultErr("Invalid repetition count: expected ',' or '}' after the lower bound"); + } + XGRAMMAR_DCHECK(regex[start] == ','); + start++; + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + // The format is {n,} + if (regex[start] == '}') { + return ResultOk(std::make_pair(lower_bound, upper_bound)); + } + num_str.clear(); + while (static_cast(start) < regex.size() && std::isdigit(regex[start])) { + num_str += regex[start]; + start++; + } + if (num_str.empty()) { + return ResultErr("Invalid repetition count: expected a number or '}' after ','"); + } + if (num_str.size() > kMaxRepeatDigits) { + return ResultErr("Invalid repetition count: the count " + num_str + " is too large"); + } + upper_bound = std::stoi(num_str); + if (upper_bound < lower_bound) { + return ResultErr( + "Invalid repetition count: the lower bound " + std::to_string(lower_bound) + + " is larger than the upper bound " + std::to_string(upper_bound) + ); + } + while (static_cast(start) < regex.size() && regex[start] == ' ') { + start++; + } + if (regex[start] != '}') { + return ResultErr("Invalid repetition count: expected '}' after the upper bound"); + } + XGRAMMAR_DCHECK(regex[start] == '}'); + return ResultOk(std::make_pair(lower_bound, upper_bound)); +} + +bool RegexIR::IsNullableSequence(const std::vector& states) { + return std::all_of(states.begin(), states.end(), [](const State& state) { + return IsNullable(state); + }); +} + +bool RegexIR::IsNullable(const State& state) { + return std::visit( + [](const auto& node) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return node.regex.empty(); + } else if constexpr (std::is_same_v) { + return node.symbol != RegexSymbol::plus || IsNullableSequence(node.state); + } else if constexpr (std::is_same_v) { + return std::any_of(node.states.begin(), node.states.end(), [](const State& child) { + return IsNullable(child); + }); + } else if constexpr (std::is_same_v) { + return IsNullableSequence(node.states); + } else if constexpr (std::is_same_v) { + return node.lower_bound == 0 || IsNullableSequence(node.states); + } else if constexpr (std::is_same_v) { + return false; + } else { + static_assert(std::is_same_v); + return node.lower_bound == 0; + } + }, + state + ); +} + +Result RegexIR::Build() const { + if (states.empty()) { + FSM empty_fsm(1); + FSMWithStartEnd result(empty_fsm, 0, {0}, false); + return ResultOk(std::move(result)); + } + std::vector fsm_list; + for (const auto& state : states) { + auto visited = std::visit([&](auto&& arg) { return visit(arg); }, state); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + if (fsm_list.size() > 1) { + return ResultOk(FSMWithStartEnd::Concat(fsm_list)); + } else { + // If there is only one FSM, return it directly. + return ResultOk(std::move(fsm_list[0])); + } +} + +Result RegexIR::visit(const RegexIR::Leaf& state) const { + return BuildLeafFSMFromRegex(state.regex); +} + +Result RegexIR::visit(const RegexIR::Union& state) const { + std::vector fsm_list; + for (const auto& child : state.states) { + auto visited = std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, child); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + if (fsm_list.size() <= 1) { + return ResultErr("Internal error: a union node in the regex IR has fewer than two branches"); + } + return ResultOk(FSMWithStartEnd::Union(fsm_list)); +} + +Result RegexIR::visit(const RegexIR::Symbol& state) const { + if (state.state.size() != 1) { + return ResultErr("Internal error: a quantifier node in the regex IR must hold exactly one child" + ); + } + Result child_result = + std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, state.state[0]); + if (child_result.IsErr()) { + return child_result; + } + auto child = std::move(child_result).Unwrap(); + + switch (state.symbol) { + case RegexIR::RegexSymbol::plus: { + return ResultOk(child.Plus()); + } + case RegexIR::RegexSymbol::star: { + return ResultOk(child.Star()); + } + case RegexIR::RegexSymbol::optional: { + return ResultOk(child.Optional()); + } + default: { + XGRAMMAR_LOG(FATAL) << "Unknown regex symbol: " << static_cast(state.symbol); + XGRAMMAR_UNREACHABLE(); + } + } +} + +Result RegexIR::visit(const RegexIR::Bracket& state) const { + if (state.states.empty()) { + // An empty group or an empty union branch matches the empty string. + FSM empty_fsm(1); + return ResultOk(FSMWithStartEnd(empty_fsm, 0, {0}, false)); + } + std::vector fsm_list; + for (const auto& child : state.states) { + auto visited = std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, child); + if (visited.IsErr()) { + return visited; + } + fsm_list.push_back(std::move(visited).Unwrap()); + } + return ResultOk(FSMWithStartEnd::Concat(fsm_list)); +} + +Result RegexIR::visit(const RegexIR::RuleRefNode& state) const { + FSM fsm(2); + fsm.AddRuleEdge(0, 1, state.rule_id); + return ResultOk(FSMWithStartEnd(fsm, 0, {1}, false)); +} + +Result RegexIR::visit(const RegexIR::RepeatSubrule& state) const { + // The state holding the repeat edge is padded with epsilon transitions on both sides, so that + // FSM compositions (Concat / Star / Optional / ...) never add another outgoing edge to it: a + // state with a kRepeatRef edge must have no other outgoing edges. + FSM fsm(4); + fsm.AddEpsilonEdge(0, 1); + fsm.AddRepeatEdge(1, 2, state.rule_id, state.lower_bound, state.upper_bound); + fsm.AddEpsilonEdge(2, 3); + return ResultOk(FSMWithStartEnd(fsm, 0, {3}, false)); +} + +Result RegexIR::visit(const RegexIR::Repeat& state) const { + if (state.states.size() != 1) { + return ResultErr("Internal error: a repetition node in the regex IR must hold exactly one child" + ); + } + bool has_upper_bound = state.upper_bound != RegexIR::kRepeatNoUpperBound; + if (has_upper_bound && state.upper_bound == 0) { + // {0} / {0,0}: matches exactly the empty string. The general path below cannot express + // this: it starts from one copy of the child whose end states stay accepting. + FSM empty_fsm(1); + return ResultOk(FSMWithStartEnd(empty_fsm, 0, {0}, false)); + } + Result child_result = + std::visit([&](auto&& arg) { return RegexIR::visit(arg); }, state.states[0]); + if (child_result.IsErr()) { + return child_result; + } + FSMWithStartEnd child = std::move(child_result).Unwrap(); + + // Guard against FSM state explosion when the repetition has to be physically unrolled. When + // a GrammarBuilder is available, large repetitions are compiled into repeat edges instead and + // never reach this point. + int64_t num_copies = has_upper_bound ? state.upper_bound : std::max(state.lower_bound, 1); + if (static_cast(child.NumStates()) * num_copies > kMaxUnrolledRepeatStates) { + return ResultErr( + "The bounded repetition {" + std::to_string(state.lower_bound) + "," + + (has_upper_bound ? std::to_string(state.upper_bound) : "") + + "} is too large to compile into a FSM" + ); + } + + FSMWithStartEnd result = child.Copy(); + std::unordered_set new_ends; + + if (state.lower_bound <= 1 && (!has_upper_bound || state.upper_bound >= 1)) { + // A single copy is accepting when the lower bound is at most 1. + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + new_ends.insert(end); + } + } + } + + // Add a fresh accepting start state so that zero repetitions match. A fresh state is + // required: making the original start accepting would also accept strings that merely + // loop back to the start inside the first copy. + auto allow_zero_repetitions = [](FSMWithStartEnd* fsm) { + int new_start = fsm->AddState(); + fsm->GetFsm().AddEpsilonEdge(new_start, fsm->GetStart()); + fsm->SetStartState(new_start); + fsm->AddEndState(new_start); + }; + + // Handling {n,} + if (!has_upper_bound) { + for (int i = 2; i < state.lower_bound; i++) { + result = FSMWithStartEnd::Concat(std::vector{result, child}); + } + int end_state_of_lower_bound_fsm = -1; + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + end_state_of_lower_bound_fsm = end; + break; + } + } + XGRAMMAR_DCHECK(end_state_of_lower_bound_fsm != -1) + << "No end state found in the lower bound FSM."; + result = FSMWithStartEnd::Concat(std::vector{result, child}); + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + result.GetFsm().AddEpsilonEdge(end, end_state_of_lower_bound_fsm); + } + } + for (const auto& end : new_ends) { + result.AddEndState(end); + } + if (state.lower_bound == 0) { + allow_zero_repetitions(&result); + } + return ResultOk(std::move(result)); + } + // Handling {n, m} or {n} + for (int i = 2; i <= state.upper_bound; i++) { + result = FSMWithStartEnd::Concat(std::vector{result, child}); + if (i >= state.lower_bound) { + for (int end = 0; end < result.NumStates(); ++end) { + if (result.IsEndState(end)) { + new_ends.insert(end); + } + } + } + } + for (const auto& end : new_ends) { + result.AddEndState(end); + } + if (state.lower_bound == 0) { + allow_zero_repetitions(&result); + } + return ResultOk(std::move(result)); +} + +int RegexIR::AddSingleCodepoint(FSMWithStartEnd& result, int current, uint32_t codepoint) const { + int next = result.AddState(); + if (codepoint <= kMax1ByteUnicode) { + result.GetFsm().AddEdge(current, next, codepoint, codepoint); + if (case_insensitive) { + if (codepoint >= 'a' && codepoint <= 'z') { + uint32_t upper = codepoint - ('a' - 'A'); + result.GetFsm().AddEdge(current, next, upper, upper); + } else if (codepoint >= 'A' && codepoint <= 'Z') { + uint32_t lower = codepoint + ('a' - 'A'); + result.GetFsm().AddEdge(current, next, lower, lower); + } + } + return next; + } + std::string utf8_bytes = CharToUTF8(static_cast(codepoint)); + int state = current; + for (size_t i = 0; i < utf8_bytes.size(); ++i) { + int target = (i + 1 == utf8_bytes.size()) ? next : result.AddState(); + uint8_t byte = static_cast(utf8_bytes[i]); + result.GetFsm().AddEdge(state, target, byte, byte); + state = target; + } + return next; +} + +Result RegexIR::BuildLeafFSMFromRegex(const std::string& regex) const { + FSM initial_fsm(1); + FSMWithStartEnd result(initial_fsm, 0, {}, false); + if (regex.empty()) { + // The empty leaf matches the empty string. + result.AddEndState(0); + return ResultOk(std::move(result)); + } + if (regex[0] == '[') { + // Character class. + auto ranges_result = ParseCharacterClassLeaf(regex, case_insensitive); + if (ranges_result.IsErr()) { + return ResultErr(std::move(ranges_result).UnwrapErr()); + } + auto ranges = std::move(ranges_result).Unwrap(); + int end_state = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), 0, end_state, ranges); + result.AddEndState(end_state); + return ResultOk(std::move(result)); + } + // Sequence of literal characters, '.' and escapes. + int current = 0; + size_t pos = 0; + while (pos < regex.size()) { + if (regex[pos] == '.') { + ++pos; + int next = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), current, next, {{0, kMaxCodepoint}}); + current = next; + continue; + } + if (regex[pos] == '\\') { + auto item_result = ParseRegexEscape(regex, &pos, /*in_class=*/false); + if (item_result.IsErr()) { + return ResultErr(std::move(item_result).UnwrapErr()); + } + auto item = std::move(item_result).Unwrap(); + if (item.is_single) { + current = AddSingleCodepoint(result, current, item.codepoint); + } else { + auto ranges = std::move(item.ranges); + if (case_insensitive) { + FoldAsciiCaseRanges(&ranges); + } + NormalizeRanges(&ranges); + if (item.negated) { + ranges = ComplementRanges(ranges); + } + int next = result.AddState(); + AddCodepointRangesToFSM(&result.GetFsm(), current, next, ranges); + current = next; + } + continue; + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + pos); + if (codepoint == CharHandlingError::kInvalidUTF8 || pos + num_bytes > regex.size()) { + // Be permissive with non-UTF-8 patterns: match the raw byte. + int next = result.AddState(); + uint8_t byte = static_cast(regex[pos]); + result.GetFsm().AddEdge(current, next, byte, byte); + current = next; + ++pos; + continue; + } + current = AddSingleCodepoint(result, current, static_cast(codepoint)); + pos += num_bytes; + } + result.AddEndState(current); + return ResultOk(std::move(result)); +} + +/******************** RegexFSMBuilder ********************/ + +namespace { + +/*! \brief One entry of the regex parsing stack: an IR node or a marker character ('(' or '|'), + * together with the source span [span_begin, span_end) of the atom in the regex string. */ +struct RegexStackEntry { + std::variant item; + int span_begin = 0; + int span_end = 0; +}; + +/*! \brief Skip a character class starting at regex[pos] == '['. Returns the position right after + * the closing ']', or std::string::npos if the class is not closed. */ +size_t SkipCharacterClass(const std::string& regex, size_t pos) { + XGRAMMAR_DCHECK(regex[pos] == '['); + ++pos; + if (pos < regex.size() && regex[pos] == '^') { + ++pos; + } + while (pos < regex.size()) { + if (regex[pos] == '\\') { + pos += 2; + continue; + } + if (regex[pos] == ']') { + return pos + 1; + } + ++pos; + } + return std::string::npos; +} + +/*! + * \brief Parse a regex string into a RegexIR. + * \param regex_with_flags The regex. A leading "(?i)" enables ASCII case-insensitive matching. + * \param builder If not null, large bounded repetitions are compiled into subrules added through + * this builder; see RegexFSMBuilder::Build. + * \param rule_hint Name hint for the created subrules. + */ +Result ParseRegexToIR( + const std::string& regex_with_flags, GrammarBuilder* builder, const std::string& rule_hint +) { + RegexIR ir; + std::string regex = regex_with_flags; + int flag_prefix_length = 0; + if (regex.size() >= 4 && regex.compare(0, 4, "(?i)") == 0) { + ir.case_insensitive = true; + regex = regex.substr(4); + flag_prefix_length = 4; + } + + // Mirror the error format of the RegexConverter path: a 1-based position in the pattern + // (including the "(?i)" prefix when present), followed by the description. + auto error_at = [&](int pos, const std::string& message) { + return "Regex parsing error at position " + std::to_string(pos + flag_prefix_length + 1) + + ": " + message; + }; + + std::vector stack; + int size = static_cast(regex.size()); + for (int i = 0; i < size; i++) { + char current_char = regex[i]; + // Handle anchors. + if (current_char == '^' || current_char == '$') { + if (!((current_char == '^' && i == 0) || (current_char == '$' && i == size - 1))) { + XGRAMMAR_LOG(WARNING) << "Anchor '" << current_char + << "' in the middle of regex is ignored: " << regex; + } + continue; + } + // Handle the character class. + if (current_char == '[') { + size_t class_end = SkipCharacterClass(regex, i); + if (class_end == std::string::npos) { + return ResultErr(error_at(i, "Unclosed '['")); + } + size_t content_begin = i + 1; + if (content_begin < regex.size() && regex[content_begin] == '^') { + ++content_begin; + } + if (content_begin + 1 == class_end) { + return ResultErr(error_at(i, "Empty character class is not allowed in regex")); + } + RegexIR::Leaf leaf; + leaf.regex = regex.substr(i, class_end - i); + stack.push_back({leaf, i, static_cast(class_end)}); + i = static_cast(class_end) - 1; + continue; + } + if (current_char == ']') { + return ResultErr(error_at(i, "Unmatched ']'")); + } + // Handle quantifiers. + if (current_char == '+' || current_char == '*' || current_char == '?') { + if (stack.empty() || std::holds_alternative(stack.back().item)) { + return ResultErr( + error_at(i, std::string("There is nothing to repeat before '") + current_char + "'") + ); + } + RegexStackEntry atom = std::move(stack.back()); + stack.pop_back(); + RegexIR::Symbol symbol; + symbol.state.push_back(std::move(std::get(atom.item))); + switch (current_char) { + case '+': { + symbol.symbol = RegexIR::RegexSymbol::plus; + break; + } + case '*': { + symbol.symbol = RegexIR::RegexSymbol::star; + break; + } + case '?': { + symbol.symbol = RegexIR::RegexSymbol::optional; + break; + } + } + // Skip the non-greedy modifier: greedy and non-greedy quantifiers accept the same + // language, and the Earley parser explores all derivations anyway. + if (i + 1 < size && regex[i + 1] == '?') { + i++; + } + stack.push_back({std::move(symbol), atom.span_begin, i + 1}); + continue; + } + // Handle groups and alternation. + if (current_char == '(') { + if (i + 1 < size && regex[i + 1] == '?') { + if (i + 2 >= size) { + return ResultErr(error_at(i, "Group modifier is not finished")); + } + char modifier = regex[i + 2]; + if (modifier == ':') { + stack.push_back({'(', i, i + 1}); + i += 2; + continue; + } + if (modifier == '=' || modifier == '!') { + // Skip the whole lookahead group and treat it as the empty string. + XGRAMMAR_LOG(WARNING) << "Lookahead assertion is not supported and is ignored in regex: " + << regex; + int depth = 1; + size_t j = i + 3; + while (j < regex.size() && depth > 0) { + if (regex[j] == '\\') { + j += 2; + continue; + } + if (regex[j] == '[') { + size_t class_begin = j; + j = SkipCharacterClass(regex, j); + if (j == std::string::npos) { + return ResultErr(error_at(static_cast(class_begin), "Unclosed '['")); + } + continue; + } + if (regex[j] == '(') { + depth++; + } else if (regex[j] == ')') { + depth--; + } + j++; + } + if (depth != 0) { + return ResultErr(error_at(i, "The parenthesis is not closed")); + } + stack.push_back({RegexIR::Leaf{""}, i, static_cast(j)}); + i = static_cast(j) - 1; + continue; + } + if (modifier == '<' || (modifier == 'P' && i + 3 < size && regex[i + 3] == '<')) { + size_t name_begin = (modifier == '<') ? i + 3 : i + 4; + if (name_begin < regex.size() && (regex[name_begin] == '=' || regex[name_begin] == '!')) { + return ResultErr(error_at(i, "Lookbehind assertion is not supported in regex")); + } + size_t j = name_begin; + while (j < regex.size() && + (std::isalnum(static_cast(regex[j])) || regex[j] == '_')) { + j++; + } + if (j == name_begin || j >= regex.size() || regex[j] != '>') { + return ResultErr(error_at(i, "Invalid named capturing group")); + } + // Ignore the group name and compile the content as a normal group. + stack.push_back({'(', i, i + 1}); + i = static_cast(j); + continue; + } + return ResultErr( + error_at(i, "Unsupported group modifier '(?" + std::string(1, modifier) + "'") + ); + } + stack.push_back({'(', i, i + 1}); + continue; + } + if (current_char == '|') { + stack.push_back({'|', i, i + 1}); + continue; + } + if (current_char == ')') { + std::vector popped; + bool paired = false; + bool unioned = false; + int group_begin = 0; + while (!stack.empty()) { + RegexStackEntry entry = std::move(stack.back()); + stack.pop_back(); + if (std::holds_alternative(entry.item)) { + char marker = std::get(entry.item); + if (marker == '(') { + paired = true; + group_begin = entry.span_begin; + break; + } + XGRAMMAR_DCHECK(marker == '|'); + unioned = true; + } + popped.push_back(std::move(entry)); + } + if (!paired) { + return ResultErr(error_at(i, "Unmatched ')'")); + } + // `popped` stores the group content from right to left. + if (!unioned) { + RegexIR::Bracket bracket; + for (auto it = popped.rbegin(); it != popped.rend(); ++it) { + bracket.states.push_back(std::move(std::get(it->item))); + } + // An empty bracket (e.g. "()") matches the empty string. + stack.push_back({std::move(bracket), group_begin, i + 1}); + } else { + RegexIR::Union union_state; + RegexIR::Bracket bracket; + for (auto it = popped.rbegin(); it != popped.rend(); ++it) { + if (std::holds_alternative(it->item)) { + XGRAMMAR_DCHECK(std::get(it->item) == '|'); + // An empty bracket represents an empty alternative, e.g. "(a|)". + union_state.states.push_back(std::move(bracket)); + bracket = RegexIR::Bracket(); + continue; + } + bracket.states.push_back(std::move(std::get(it->item))); + } + union_state.states.push_back(std::move(bracket)); + stack.push_back({std::move(union_state), group_begin, i + 1}); + } + continue; + } + // Handle repetitions. + if (current_char == '{') { + if (stack.empty() || std::holds_alternative(stack.back().item)) { + return ResultErr(error_at(i, "There is nothing to repeat before the repetition")); + } + RegexStackEntry atom = std::move(stack.back()); + stack.pop_back(); + int repeat_begin = i; + auto bounds_result = RegexIR::CheckRepeat(regex, i); + if (bounds_result.IsErr()) { + return ResultErr(error_at(repeat_begin, std::move(bounds_result).UnwrapErr().what())); + } + auto [lower_bound, upper_bound] = std::move(bounds_result).Unwrap(); + // Skip the non-greedy modifier. + if (i + 1 < size && regex[i + 1] == '?') { + i++; + } + bool is_large_repeat = upper_bound == RegexIR::kRepeatNoUpperBound + ? lower_bound > kLargeRepeatThreshold + : upper_bound > kLargeRepeatThreshold; + if (is_large_repeat && builder != nullptr) { + // Compile the repeated sub-pattern into a new rule (with a kRegex body), and represent + // the repetition as a kRepeatRef FSM edge. The Earley parser executes it with a counter + // at runtime, so the FSM is not unrolled. + const auto& atom_state = std::get(atom.item); + std::string inner_regex = regex.substr(atom.span_begin, atom.span_end - atom.span_begin); + if (ir.case_insensitive) { + inner_regex = "(?i)" + inner_regex; + } + if (RegexIR::IsNullable(atom_state)) { + // Mirror RepetitionNormalizer: when the repeated element can match the empty string, + // the repetition count cannot be enforced, so the lower bound is relaxed to 0. + lower_bound = 0; + } + std::string name_hint = (rule_hint.empty() ? "regex" : rule_hint) + "_repeat"; + int32_t inner_rule_id = builder->AddRuleWithHint(name_hint, builder->AddRegex(inner_regex)); + builder->UpdateLookaheadExact(inner_rule_id, true); + if (upper_bound == RegexIR::kRepeatNoUpperBound) { + // {n,} == {n}{0,}: a repeat edge for the mandatory part, then a starred rule + // reference. + RegexIR::Symbol star_symbol; + star_symbol.symbol = RegexIR::RegexSymbol::star; + star_symbol.state.push_back(RegexIR::RuleRefNode{inner_rule_id}); + if (lower_bound > 0) { + RegexIR::Bracket bracket; + bracket.states.push_back(RegexIR::RepeatSubrule{inner_rule_id, lower_bound, lower_bound} + ); + bracket.states.push_back(std::move(star_symbol)); + stack.push_back({std::move(bracket), atom.span_begin, i + 1}); + } else { + stack.push_back({std::move(star_symbol), atom.span_begin, i + 1}); + } + } else { + stack.push_back( + {RegexIR::RepeatSubrule{inner_rule_id, lower_bound, upper_bound}, + atom.span_begin, + i + 1} + ); + } + } else { + RegexIR::Repeat repeat; + repeat.lower_bound = lower_bound; + repeat.upper_bound = upper_bound; + repeat.states.push_back(std::move(std::get(atom.item))); + stack.push_back({std::move(repeat), atom.span_begin, i + 1}); + } + continue; + } + // Handle literal characters and escapes. Each leaf holds exactly one codepoint or escape + // sequence, so that a following quantifier applies to the whole character. + RegexIR::Leaf leaf; + if (current_char == '\\') { + size_t escape_end = i; + auto escape_result = ParseRegexEscape(regex, &escape_end, /*in_class=*/false); + if (escape_result.IsErr()) { + return ResultErr(error_at(i, std::move(escape_result).UnwrapErr().what())); + } + leaf.regex = regex.substr(i, escape_end - i); + stack.push_back({std::move(leaf), i, static_cast(escape_end)}); + i = static_cast(escape_end) - 1; + continue; + } + auto [codepoint, num_bytes] = ParseNextUTF8(regex.c_str() + i); + if (codepoint == CharHandlingError::kInvalidUTF8 || i + num_bytes > size) { + num_bytes = 1; + } + leaf.regex = regex.substr(i, num_bytes); + stack.push_back({std::move(leaf), i, i + num_bytes}); + i += num_bytes - 1; + continue; + } + + // Assemble the top-level sequence / union. `stack` stores the content from left to right. + std::vector segment; + std::vector> union_segments; + bool unioned = false; + for (auto& entry : stack) { + if (std::holds_alternative(entry.item)) { + char marker = std::get(entry.item); + if (marker == '|') { + union_segments.push_back(std::move(segment)); + segment.clear(); + unioned = true; + continue; + } + return ResultErr(error_at(entry.span_begin, "The parenthesis is not closed")); + } + segment.push_back(std::move(std::get(entry.item))); + } + if (!unioned) { + ir.states = std::move(segment); + } else { + union_segments.push_back(std::move(segment)); + RegexIR::Union union_state; + for (auto& branch : union_segments) { + RegexIR::Bracket bracket; + bracket.states = std::move(branch); + union_state.states.push_back(std::move(bracket)); + } + ir.states.push_back(std::move(union_state)); + } + return ResultOk(std::move(ir)); +} + +} // namespace + +Result RegexFSMBuilder::Build( + const std::string& regex, GrammarBuilder* builder, const std::string& rule_hint +) { + auto ir_result = ParseRegexToIR(regex, builder, rule_hint); + if (ir_result.IsErr()) { + return ResultErr(std::move(ir_result).UnwrapErr()); + } + return std::move(ir_result).Unwrap().Build(); +} + +Result RegexFSMBuilder::MatchesEmpty(const std::string& regex) { + auto ir_result = ParseRegexToIR(regex, nullptr, ""); + if (ir_result.IsErr()) { + return ResultErr(std::move(ir_result).UnwrapErr()); + } + auto ir = std::move(ir_result).Unwrap(); + return ResultOk(RegexIR::IsNullableSequence(ir.states)); +} + +Result RegexFSMBuilder::BuildWithForbiddenChars( + const std::string& regex, + const std::bitset<256>& forbidden_chars, + GrammarBuilder* builder, + const std::string& rule_hint +) { + auto build_result = Build(regex, builder, rule_hint); + if (build_result.IsErr() || forbidden_chars.none()) { + return build_result; + } + auto fsm_wse = std::move(build_result).Unwrap(); + const auto& fsm = fsm_wse.GetFsm(); + FSM new_fsm(fsm_wse.NumStates()); + for (int state = 0; state < fsm_wse.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (!edge.IsCharRange()) { + new_fsm.AddEdge(state, edge.target, edge.min, edge.max); + continue; + } + // Split the character range into the maximal sub-ranges of allowed characters. + int range_start = -1; + for (int c = edge.min; c <= edge.max + 1; ++c) { + if (c <= edge.max && !forbidden_chars[c]) { + if (range_start == -1) { + range_start = c; + } + } else if (range_start != -1) { + new_fsm.AddEdge(state, edge.target, range_start, c - 1); + range_start = -1; + } + } + } + } + new_fsm.SetEdgeAuxData(std::vector(fsm.GetEdgeAuxData())); + return ResultOk(FSMWithStartEnd(new_fsm, fsm_wse.GetStart(), fsm_wse.GetEnds())); +} + +class TrieFSMBuilderImpl { + public: + TrieFSMBuilderImpl() = default; + std::optional Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges + ); + void AddBackEdges( + FSM* fsm, + int start, + const std::unordered_set& ends, + std::unordered_set* dead_states + ); +}; + +std::optional TrieFSMBuilderImpl::Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges +) { + FSM fsm(1); + int start = 0; + std::unordered_set ends; + + if (end_states) { + end_states->clear(); + } + + for (const auto& pattern : patterns) { + // Check for empty patterns + if (!allow_overlap && pattern.empty()) { + return std::nullopt; + } + + int current_state = 0; + for (const auto& ch : pattern) { + int32_t ch_int32 = static_cast(static_cast(ch)); + int next_state = fsm.GetNextState(current_state, ch_int32); + if (next_state == FSM::kNoNextState) { + next_state = fsm.AddState(); + fsm.AddEdge(current_state, next_state, ch_int32, ch_int32); + } + current_state = next_state; + if (!allow_overlap && ends.count(current_state) > 0) { + return std::nullopt; + } + } + if (!allow_overlap && fsm.GetEdges(current_state).size() > 0) { + return std::nullopt; + } + ends.insert(current_state); + if (end_states) { + end_states->push_back(current_state); + } + } + + std::unordered_set dead_state_set; + + if (add_back_edges) { + // Build trie for excluded patterns. + for (const auto& excluded_pattern : excluded_patterns) { + if (!allow_overlap && excluded_pattern.empty()) { + return std::nullopt; + } + + int current_state = 0; + for (const auto& ch : excluded_pattern) { + int32_t ch_int32 = static_cast(static_cast(ch)); + int next_state = fsm.GetNextState(current_state, ch_int32); + if (next_state == FSM::kNoNextState) { + next_state = fsm.AddState(); + fsm.AddEdge(current_state, next_state, ch_int32, ch_int32); + } + current_state = next_state; + if (!allow_overlap && ends.count(current_state) > 0) { + return std::nullopt; + } + } + if (!allow_overlap && fsm.GetEdges(current_state).size() > 0) { + return std::nullopt; + } + + ends.insert(current_state); + dead_state_set.insert(current_state); + } + + // Add back edges. + AddBackEdges(&fsm, start, ends, &dead_state_set); + + // Remove the edges to excluded end states. + if (dead_state_set.size() != 0) { + for (int state = 0; state < fsm.NumStates(); state++) { + std::vector& edges = fsm.GetEdges(state); + std::vector new_edges; + for (const auto& edge : edges) { + if (dead_state_set.count(edge.target) == 0) { + new_edges.push_back(edge); + } + } + edges = std::move(new_edges); + } + } + } else if (excluded_patterns.size() > 0) { + XGRAMMAR_LOG(WARNING) << "Excluded patterns are ignored when back edges are not added."; + } + + return FSMWithStartEnd(fsm, start, std::vector(ends.begin(), ends.end())); +} + +void TrieFSMBuilderImpl::AddBackEdges( + FSM* fsm, + int start, + const std::unordered_set& ends, + std::unordered_set* dead_states +) { + // Build an Aho-Corasick automaton by adding back edges. + // When matching on the trie fails at state u on byte b, the matcher must resume from + // the longest proper suffix of u's prefix that is still a path in the trie (the + // failure state), and retry b from there. Falling back only to the start state (or to + // the start state's direct children) loses matches whose start lies inside an + // already-followed branch of another pattern. Example: patterns {"bcd", "abce"} on + // input "abcd" -- after following "abc" of the "abce" branch, 'd' must transition to + // the "bcd" end state via the failure state "bc", not back to the start state. + + int num_states = fsm->NumStates(); + + // Step 1. Record the BFS order of the trie (a tree at this point), so that shallower + // states are always processed first. + std::vector bfs_order; + bfs_order.reserve(num_states); + bfs_order.push_back(start); + for (size_t head = 0; head < bfs_order.size(); head++) { + for (const auto& edge : fsm->GetEdges(bfs_order[head])) { + XGRAMMAR_DCHECK(edge.min == edge.max && edge.min >= 0 && edge.min <= 255); + bfs_order.push_back(edge.target); + } + } + XGRAMMAR_DCHECK(static_cast(bfs_order.size()) == num_states); + + // Step 2. Compute the failure link and the fully resolved transition table with the + // textbook O(num_states * 256) dynamic program: delta[u][b] is the trie child when it + // exists, and delta[fail[u]][b] otherwise -- fail[u] is strictly shallower than u, so + // its row is already final when u is processed in BFS order. + std::vector fail(num_states, start); + std::vector> delta(num_states); + for (auto& row : delta) { + row.fill(FSM::kNoNextState); + } + for (int u = 0; u < num_states; u++) { + for (const auto& edge : fsm->GetEdges(u)) { + delta[u][edge.min] = edge.target; + } + } + for (int u : bfs_order) { + // An explicit trie child can complete an excluded suffix without entering its + // terminal state: with excludes {"abc", "xabcq"}, the state for "xabc" has + // failure state "abc". It must be dead too, even if more trie edges remain. + // BFS ensures the failure state's excluded status has already been propagated. + if (dead_states->count(fail[u]) > 0) { + dead_states->insert(u); + } + for (int byte = 0; byte < 256; byte++) { + // Entries of deeper states are untouched so far, so a non-empty entry here is + // exactly a trie child of u. + int child = delta[u][byte]; + int fallback = (u == start) ? start : delta[fail[u]][byte]; + if (child == FSM::kNoNextState) { + delta[u][byte] = fallback; + } else { + fail[child] = fallback; + } + } + } + + // Step 3. Overwrite the edges of every non-end state with its resolved row, + // compressing consecutive bytes with the same target into range edges. + for (int u = 0; u < num_states; u++) { + if (u != start && ends.count(u) > 0) { + continue; + } + const auto& row = delta[u]; + std::vector new_edges; + for (int byte = 0; byte < 256;) { + int target = row[byte]; + int range_end = byte; + while (range_end + 1 < 256 && row[range_end + 1] == target) { + range_end++; + } + new_edges.push_back(FSMEdge(byte, range_end, target)); + byte = range_end + 1; + } + fsm->GetEdges(u) = std::move(new_edges); + } +} + +std::optional TrieFSMBuilder::Build( + const std::vector& patterns, + const std::vector& exclude_patterns, + std::vector* end_states, + bool allow_overlap, + bool add_back_edges +) { + return TrieFSMBuilderImpl().Build( + patterns, exclude_patterns, end_states, allow_overlap, add_back_edges + ); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/fsm_builder.h b/zig/vendor/xgrammar/cpp/fsm_builder.h new file mode 100644 index 0000000000..1858668f8c --- /dev/null +++ b/zig/vendor/xgrammar/cpp/fsm_builder.h @@ -0,0 +1,119 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/fsm_builder.h + */ +#ifndef XGRAMMAR_FSM_BUILDER_H_ +#define XGRAMMAR_FSM_BUILDER_H_ + +#include +#include +#include +#include + +#include "fsm.h" +#include "support/utils.h" + +namespace xgrammar { + +class GrammarBuilder; + +/*! \brief Bounds of the packed UTF-8 representation for each encoded length. The packed format + * stores the UTF-8 bytes of one codepoint as (byte0 << 24) | (byte1 << 16) | (byte2 << 8) | byte3, + * left-aligned to the actual length (e.g. a 2-byte character is (byte0 << 8) | byte1). */ +constexpr uint32_t kMax1ByteUnicode = 0x7F; +constexpr uint32_t kMin2BytesUnicode = 0xC080; +constexpr uint32_t kMax2BytesUnicode = 0xDFBF; +constexpr uint32_t kMin3BytesUnicode = 0xE08080; +constexpr uint32_t kMax3BytesUnicode = 0xEFBFBF; +constexpr uint32_t kMin4BytesUnicode = 0xF0808080; +constexpr uint32_t kMax4BytesUnicode = 0xF7BFBFBF; + +/*! \brief Convert a Unicode codepoint to the packed UTF-8 format described above. */ +uint32_t CodepointToPackedUTF8(uint32_t codepoint); + +/*! + * \brief Add FSM edges (with intermediate states for multi-byte characters) from `from` to `to` + * accepting every UTF-8 encoded character in the packed range [min, max]. + */ +void AddPackedUTF8RangeEdges(FSM& fsm, int from, int to, uint32_t min, uint32_t max); + +/*! + * \brief Rewrite every unescaped '.' outside character classes to "[^\n]" unless + * `dot_matches_newline` is true. Used to implement the standard regex dot semantics (and the + * dot-all 's' flag) on top of the regex engine, whose '.' matches every codepoint. + */ +std::string RewriteRegexDots(const std::string& pattern, bool dot_matches_newline); + +/*! + * \brief A builder that converts a regex string to a FSM. + */ +class RegexFSMBuilder { + public: + /*! + * \brief Converts a regex string to a FSM. + * \param regex The regex string. A leading "(?i)" makes the match ASCII case-insensitive. + * \param builder If not null, bounded repetitions whose upper bound exceeds the unroll + * threshold are compiled into a kRepeatRef FSM edge referencing a new rule (with a kRegex + * body holding the repeated sub-pattern) added through this builder, instead of being + * physically unrolled. + * \param rule_hint Name hint for rules created through `builder`. + * \return The FSM with start and end states. + */ + static Result Build( + const std::string& regex, GrammarBuilder* builder = nullptr, const std::string& rule_hint = "" + ); + + /*! + * \brief Converts a regex string to a FSM, then removes the forbidden characters from every + * character transition. The result accepts the intersection of the regex language and the + * set of strings that contain no forbidden character. The result language may be empty. + * \param regex The regex string. + * \param forbidden_chars The forbidden characters. + * \param builder See Build(). + * \param rule_hint See Build(). + * \return The FSM with start and end states. + */ + static Result BuildWithForbiddenChars( + const std::string& regex, + const std::bitset<256>& forbidden_chars, + GrammarBuilder* builder = nullptr, + const std::string& rule_hint = "" + ); + + /*! + * \brief Check whether the regex matches the empty string. Only parses the regex; no FSM is + * built, so this is cheap even for regexes with huge bounded repetitions. + */ + static Result MatchesEmpty(const std::string& regex); +}; + +/*! + * \brief A builder that converts a list of patterns to a trie-based FSM. + */ +class TrieFSMBuilder { + public: + /*! + * \brief Build a trie-based FSM from a list of patterns. + * \param patterns The patterns to be built. + * \param excluded_patterns The patterns to be excluded. + * \param end_states The end states of the FSM. This is the terminal state of each pattern and + * the order follows the order of patterns. + * \param allow_overlap Whether to allow overlap between patterns (one being a prefix of the + * other). It does not allow empty patterns either. If false and there is overlap, will return + * std::nullopt. + * \param add_back_edges Whether to add back edges to the FSM. This complements the trie to an + * Aho-Corasick automaton. + * \return If success, the FSM with start and end states. Otherwise, std::nullopt. + */ + static std::optional Build( + const std::vector& patterns, + const std::vector& excluded_patterns, + std::vector* end_states = nullptr, + bool allow_overlap = true, + bool add_back_edges = false + ); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_FSM_BUILDER_H_ diff --git a/zig/vendor/xgrammar/cpp/grammar.cc b/zig/vendor/xgrammar/cpp/grammar.cc new file mode 100644 index 0000000000..59f4de0963 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar.cc @@ -0,0 +1,359 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.cc + */ + +#include + +#include +#include +#include +#include + +#include "grammar_functor.h" +#include "grammar_parser.h" +#include "grammar_printer.h" +#include "json_schema_converter.h" +#include "lark_converter.h" +#include "regex_converter.h" +#include "structural_tag.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/******************* Grammar::Impl *******************/ + +std::size_t MemorySize(const Grammar::Impl& impl) { + /// TODO: Now, we evaluate the memory size of each rule as sizeof(Rule), which counts its + /// string members as sizeof(std::string), with an assumption that the strings are small. + /// This should be improved in the future. + return impl.rules_.size() * sizeof(Grammar::Impl::Rule) + + impl.suffix_stop_infos_.size() * sizeof(Grammar::Impl::SuffixStopInfo) + + MemorySize(impl.grammar_expr_data_) + MemorySize(impl.grammar_expr_indptr_) + + MemorySize(impl.complete_fsm) + MemorySize(impl.per_rule_fsms) + + MemorySize(impl.allow_empty_rule_ids); +} + +/******************* Grammar *******************/ + +std::string Grammar::ToString() const { return GrammarPrinter(*this).ToString(); } + +Grammar Grammar::FromEBNF(const std::string& ebnf_string, const std::string& root_rule_name) { + auto grammar = ParseEBNF(ebnf_string, root_rule_name); + grammar = GrammarNormalizer().Apply(grammar); + return grammar; +} + +Grammar Grammar::FromJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool print_converted_ebnf, + bool any_order +) { + auto grammar = GrammarNormalizer::Apply(JSONSchemaToGrammar( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + )); + if (print_converted_ebnf) { + XGRAMMAR_LOG(INFO) << "Converted EBNF: " << grammar.ToString() << std::endl; + } + return grammar; +} + +Grammar Grammar::FromRegex(const std::string& regex, bool print_converted_ebnf) { + auto ebnf_string = RegexToEBNF(regex); + if (print_converted_ebnf) { + XGRAMMAR_LOG(INFO) << "Converted EBNF: " << ebnf_string << std::endl; + } + return FromEBNF(ebnf_string); +} + +Grammar Grammar::FromLark( + const std::string& lark_string, + const std::optional& tokenizer_info, + const std::vector& named_grammars +) { + return LarkToGrammar(lark_string, tokenizer_info, named_grammars); +} + +std::variant Grammar::FromStructuralTag( + const std::string& structural_tag_json, const std::optional& tokenizer_info +) { + return StructuralTagToGrammar(structural_tag_json, tokenizer_info).ToVariant(); +} + +// Optimized json grammar for the speed of the grammar matcher +const std::string kJSONGrammarString = R"( +root ::= ( + "{" [ \n\r\t]* members_and_embrace | + "[" [ \n\r\t]* elements_or_embrace +) +value_non_str ::= ( + "{" [ \n\r\t]* members_and_embrace | + "[" [ \n\r\t]* elements_or_embrace | + "0" fraction exponent | + [1-9] [0-9]* fraction exponent | + "-" [0-9] fraction exponent | + "-" [1-9] [0-9]* fraction exponent | + "true" | + "false" | + "null" +) (= [ \n\r\t]* member_suffix_suffix) +members_and_embrace ::= ("\"" characters_and_colon [ \n\r\t]* members_suffix | "}") (= [ \n\r\t,}\]]) +members_suffix ::= ( + value_non_str [ \n\r\t]* member_suffix_suffix | + "\"" characters_and_embrace | + "\"" characters_and_comma [ \n\r\t]* "\"" characters_and_colon [ \n\r\t]* members_suffix +) (= [ \n\r\t,}\]]) +member_suffix_suffix ::= ( + "}" | + "," [ \n\r\t]* "\"" characters_and_colon [ \n\r\t]* members_suffix +) (= [ \n\r\t,}\]]) +elements_or_embrace ::= ( + "{" [ \n\r\t]* members_and_embrace elements_rest [ \n\r\t]* "]" | + "[" [ \n\r\t]* elements_or_embrace elements_rest [ \n\r\t]* "]" | + "\"" characters_item elements_rest [ \n\r\t]* "]" | + "0" fraction exponent elements_rest [ \n\r\t]* "]" | + [1-9] [0-9]* fraction exponent elements_rest [ \n\r\t]* "]" | + "-" "0" fraction exponent elements_rest [ \n\r\t]* "]" | + "-" [1-9] [0-9]* fraction exponent elements_rest [ \n\r\t]* "]" | + "true" elements_rest [ \n\r\t]* "]" | + "false" elements_rest [ \n\r\t]* "]" | + "null" elements_rest [ \n\r\t]* "]" | + "]" +) +elements ::= ( + "{" [ \n\r\t]* members_and_embrace elements_rest | + "[" [ \n\r\t]* elements_or_embrace elements_rest | + "\"" characters_item elements_rest | + "0" fraction exponent elements_rest | + [1-9] [0-9]* fraction exponent elements_rest | + "-" [0-9] fraction exponent elements_rest | + "-" [1-9] [0-9]* fraction exponent elements_rest | + "true" elements_rest | + "false" elements_rest | + "null" elements_rest +) +elements_rest ::= ( + "" | + [ \n\r\t]* "," [ \n\r\t]* elements +) +characters_and_colon ::= ( + "\"" [ \n\r\t]* ":" | + [^"\\\x00-\x1F] characters_and_colon | + "\\" escape characters_and_colon +) (=[ \n\r\t]* [\"{[0-9tfn-]) +characters_and_comma ::= ( + "\"" [ \n\r\t]* "," | + [^"\\\x00-\x1F] characters_and_comma | + "\\" escape characters_and_comma +) (=[ \n\r\t]* "\"") +characters_and_embrace ::= ( + "\"" [ \n\r\t]* "}" | + [^"\\\x00-\x1F] characters_and_embrace | + "\\" escape characters_and_embrace +) (=[ \n\r\t]* [},]) +characters_item ::= ( + "\"" | + [^"\\\x00-\x1F] characters_item | + "\\" escape characters_item +) (= [ \n\r\t]* [,\]]) +escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] +fraction ::= "" | "." [0-9] [0-9]* +exponent ::= "" | "e" sign [0-9] [0-9]* | "E" sign [0-9] [0-9]* +sign ::= "" | "+" | "-" +)"; + +Grammar Grammar::BuiltinJSONGrammar() { + static const Grammar grammar = FromEBNF(kJSONGrammarString); + return grammar; +} + +Grammar Grammar::Union(const std::vector& grammars) { + return GrammarUnionFunctor::Apply(grammars); +} + +Grammar Grammar::Concat(const std::vector& grammars) { + return GrammarConcatFunctor::Apply(grammars); +} + +std::ostream& operator<<(std::ostream& os, const Grammar& grammar) { + os << grammar.ToString(); + return os; +} + +std::optional Grammar::Impl::Validate() const { + const int64_t num_rules = rules_.size(); + const int64_t num_exprs = grammar_expr_indptr_.size(); + const int64_t data_size = grammar_expr_data_.size(); + auto rule_ok = [&](int64_t id) { return id >= 0 && id < num_rules; }; + auto expr_ok = [&](int64_t id) { return id >= 0 && id < num_exprs; }; + + // Pass 1: every expr must fit in grammar_expr_data_ before any expr can be read. + for (int64_t expr_id = 0; expr_id < num_exprs; ++expr_id) { + const int64_t start = grammar_expr_indptr_[expr_id]; + if (start < 0 || start + 2 > data_size) { + return "grammar_expr_indptr[" + std::to_string(expr_id) + "] is out of range"; + } + const int64_t type = grammar_expr_data_[start]; + const int64_t len = grammar_expr_data_[start + 1]; + if (len < 0 || start + 2 + len > data_size) { + return "The length of grammar expr " + std::to_string(expr_id) + " is out of range"; + } + if (type < 0 || type > static_cast(GrammarExprType::kSubstring)) { + return "Unknown type of grammar expr " + std::to_string(expr_id); + } + } + + // Pass 2: the ids stored inside each expr must refer to existing rules and exprs. + for (int64_t expr_id = 0; expr_id < num_exprs; ++expr_id) { + const auto expr = GetGrammarExpr(expr_id); + const int64_t size = expr.size(); + bool ok = true; + switch (expr.type) { + case GrammarExprType::kByteString: + case GrammarExprType::kEmptyStr: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + break; + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + // [is_negative, lower0, upper0, ...] + ok = size >= 1 && size % 2 == 1; + break; + case GrammarExprType::kRuleRef: + ok = size == 1 && rule_ok(expr[0]); + break; + case GrammarExprType::kSequence: + case GrammarExprType::kChoices: + ok = std::all_of(expr.begin(), expr.end(), expr_ok); + break; + case GrammarExprType::kTagDispatch: { + // [tag_expr0, rule_id0, ..., loop_after_dispatch, excluded_str_expr_id] + const int64_t extra = TagDispatch::kTagDispatchExtraParameter; + ok = size >= extra && (size - extra) % 2 == 0; + for (int64_t i = 0; ok && i < size - extra; i += 2) { + ok = expr_ok(expr[i]) && rule_ok(expr[i + 1]); + } + // The excluded strings are read as a kChoices expr of byte string exprs. + ok = ok && expr_ok(expr[size - 1]) && + GetGrammarExpr(expr[size - 1]).type == GrammarExprType::kChoices; + break; + } + case GrammarExprType::kRepeat: + ok = size == 3 && rule_ok(expr[0]); + break; + case GrammarExprType::kTokenTagDispatch: { + // [trigger_cnt, (token_id, rule_id) x N, loop_after_dispatch, exclude_cnt, token_id x M] + const int64_t trigger_cnt = size >= 1 ? expr[0] : -1; + ok = trigger_cnt >= 0 && 1 + 2 * trigger_cnt + 2 <= size; + for (int64_t i = 0; ok && i < trigger_cnt; ++i) { + ok = rule_ok(expr[2 + 2 * i]); + } + if (ok) { + const int64_t exclude_cnt = expr[1 + 2 * trigger_cnt + 1]; + ok = exclude_cnt >= 0 && 1 + 2 * trigger_cnt + 2 + exclude_cnt == size; + } + break; + } + case GrammarExprType::kRegex: + ok = size >= 1; + break; + case GrammarExprType::kSubstring: + // [chunk0_len, byte0_0, ..., chunk1_len, ...] + for (int64_t i = 0; ok && i < size;) { + const int64_t chunk_len = expr[i++]; + ok = chunk_len >= 0 && i + chunk_len <= size; + i += chunk_len; + } + break; + } + if (!ok) { + return "Grammar expr " + std::to_string(expr_id) + " is malformed"; + } + } + + for (int64_t rule_id = 0; rule_id < num_rules; ++rule_id) { + const auto& rule = rules_[rule_id]; + if (!expr_ok(rule.body_expr_id) || + (rule.lookahead_assertion_id != -1 && !expr_ok(rule.lookahead_assertion_id))) { + return "Rule " + std::to_string(rule_id) + " refers to a grammar expr out of range"; + } + } + if (!rule_ok(root_rule_id_)) { + return "root_rule_id " + std::to_string(root_rule_id_) + " is out of range"; + } + for (const auto& info : suffix_stop_infos_) { + if (!rule_ok(info.rule_id) || (info.body_rule_id != -1 && !rule_ok(info.body_rule_id)) || + (info.marker_rule_id != -1 && !rule_ok(info.marker_rule_id))) { + return "suffix_stop_infos refers to a rule out of range"; + } + } + if (!std::all_of(allow_empty_rule_ids.begin(), allow_empty_rule_ids.end(), rule_ok)) { + return "allow_empty_rule_ids refers to a rule out of range"; + } + + // The FSMs are internally consistent (see CompactFSM::Impl::Validate); check the rule ids they + // refer to. The parser reads the repeat info of a per-rule FSM edge from complete_fsm, so the + // aux indices of per-rule repeat edges must also be valid for complete_fsm. + auto fsm_rules_ok = [&](const CompactFSM& fsm) { + const auto& aux = fsm.GetEdgeAuxData(); + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if ((edge.IsRuleRef() && !rule_ok(edge.max)) || + (edge.IsRepeatRef() && !rule_ok(aux[edge.max]))) { + return false; + } + } + } + return true; + }; + if (!complete_fsm.IsNull() && !fsm_rules_ok(complete_fsm)) { + return "complete_fsm refers to a rule out of range"; + } + if (optimized && + (complete_fsm.IsNull() || static_cast(per_rule_fsms.size()) != num_rules)) { + return "An optimized grammar must have complete_fsm and one FSM per rule"; + } + for (int64_t rule_id = 0; rule_id < static_cast(per_rule_fsms.size()); ++rule_id) { + if (!per_rule_fsms[rule_id].has_value()) { + if (optimized) { + return "Rule " + std::to_string(rule_id) + " has no FSM in an optimized grammar"; + } + continue; + } + const CompactFSM& fsm = per_rule_fsms[rule_id]->GetFsm().GetFsm(); + if (!fsm_rules_ok(fsm)) { + return "The FSM of rule " + std::to_string(rule_id) + " refers to a rule out of range"; + } + const int64_t complete_aux_size = + complete_fsm.IsNull() ? 0 : complete_fsm.GetEdgeAuxData().size(); + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (edge.IsRepeatRef() && (static_cast(edge.max) + 3 > complete_aux_size || + !rule_ok(complete_fsm.GetEdgeAuxData()[edge.max]))) { + return "The FSM of rule " + std::to_string(rule_id) + + " has a repeat edge that is out of range of complete_fsm"; + } + } + } + } + return std::nullopt; +} + +std::string Grammar::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +std::variant Grammar::DeserializeJSON(const std::string& json_string) { + Grammar result{NullObj()}; + if (auto err = AutoDeserializeJSON(&result, json_string, true, "Grammar")) { + return err.value(); + } + return result; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_builder.cc b/zig/vendor/xgrammar/cpp/grammar_builder.cc new file mode 100644 index 0000000000..e66b98dd32 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_builder.cc @@ -0,0 +1,425 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/grammar_builder.cc + */ + +#include "grammar_builder.h" + +#include +#include +#include +#include +#include +#include + +#include "support/logging.h" + +namespace xgrammar { + +/****************** GrammarBuilder ******************/ + +GrammarBuilder::GrammarBuilder() : grammar_(std::make_shared()) {} + +GrammarBuilder::GrammarBuilder(const Grammar& grammar) + : grammar_(std::make_shared(*grammar.operator->())), + rule_name_map_built_(false) {} + +GrammarBuilder GrammarBuilder::FromMutableGrammar(Grammar* grammar) { + GrammarBuilder builder; + builder.grammar_ = grammar->pimpl_; + builder.rule_name_map_built_ = false; + return builder; +} + +void GrammarBuilder::EnsureRuleNameMap() const { + if (rule_name_map_built_) { + return; + } + for (int32_t i = 0; i < static_cast(grammar_->rules_.size()); ++i) { + rule_name_to_id_[grammar_->rules_[i].name] = i; + } + rule_name_map_built_ = true; +} + +Grammar GrammarBuilder::Get(const std::string& root_rule_name) { + int32_t root_rule_id = GetRuleId(root_rule_name); + XGRAMMAR_CHECK(root_rule_id != -1) + << "The root rule with name \"" << root_rule_name << "\" is not found."; + return Get(root_rule_id); +} + +Grammar GrammarBuilder::Get(int32_t root_rule_id) { + XGRAMMAR_CHECK(root_rule_id >= 0 && root_rule_id < static_cast(grammar_->rules_.size())) + << "The root rule id " << root_rule_id << " is out of bound."; + grammar_->root_rule_id_ = root_rule_id; + return Grammar(grammar_); +} + +int32_t GrammarBuilder::AddGrammarExpr(const GrammarExpr& grammar_expr) { + // Offsets into grammar_expr_data_ are stored as int32. + XGRAMMAR_CHECK( + grammar_->grammar_expr_data_.size() + 2 + grammar_expr.data_len <= + static_cast(std::numeric_limits::max()) + ) << "The grammar is too large: the grammar expr data exceeds 2^31 elements"; + grammar_->grammar_expr_indptr_.push_back(grammar_->grammar_expr_data_.size()); + grammar_->grammar_expr_data_.push_back(static_cast(grammar_expr.type)); + grammar_->grammar_expr_data_.push_back(grammar_expr.data_len); + grammar_->grammar_expr_data_.insert( + grammar_->grammar_expr_data_.end(), + grammar_expr.data, + grammar_expr.data + grammar_expr.data_len + ); + return static_cast(grammar_->grammar_expr_indptr_.size()) - 1; +} + +int32_t GrammarBuilder::AddByteString(const std::vector& bytes) { + return AddGrammarExpr( + {GrammarExprType::kByteString, bytes.data(), static_cast(bytes.size())} + ); +} + +int32_t GrammarBuilder::AddByteString(const std::string& str) { + std::vector bytes; + bytes.reserve(str.size()); + for (char c : str) { + bytes.push_back(static_cast(static_cast(c))); + } + return AddGrammarExpr( + {GrammarExprType::kByteString, bytes.data(), static_cast(bytes.size())} + ); +} + +int32_t GrammarBuilder::AddRegex(const std::string& regex_str, bool json_string) { + std::vector data; + data.reserve(regex_str.size() + 1); + data.push_back(static_cast(json_string)); + for (char c : regex_str) { + data.push_back(static_cast(static_cast(c))); + } + return AddGrammarExpr({GrammarExprType::kRegex, data.data(), static_cast(data.size())}); +} + +int32_t GrammarBuilder::AddSubstring(const std::vector& chunks) { + std::vector data; + for (const std::string& chunk : chunks) { + data.push_back(static_cast(chunk.size())); + for (char c : chunk) { + data.push_back(static_cast(static_cast(c))); + } + } + return AddGrammarExpr( + {GrammarExprType::kSubstring, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddCharacterClass( + const std::vector& elements, bool is_negative +) { + std::vector data; + data.reserve(1 + elements.size() * 2); + data.push_back(static_cast(is_negative)); + for (const auto& range : elements) { + data.push_back(range.lower); + data.push_back(range.upper); + } + return AddGrammarExpr( + {GrammarExprType::kCharacterClass, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddCharacterClassStar( + const std::vector& elements, bool is_negative +) { + std::vector data; + data.reserve(1 + elements.size() * 2); + data.push_back(static_cast(is_negative)); + for (const auto& range : elements) { + data.push_back(range.lower); + data.push_back(range.upper); + } + return AddGrammarExpr( + {GrammarExprType::kCharacterClassStar, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddEmptyStr() { + return AddGrammarExpr({GrammarExprType::kEmptyStr, nullptr, 0}); +} + +int32_t GrammarBuilder::AddTokenSet(const std::vector& token_ids) { + return AddGrammarExpr( + {GrammarExprType::kToken, token_ids.data(), static_cast(token_ids.size())} + ); +} + +int32_t GrammarBuilder::AddExcludeTokenSet(const std::vector& token_ids) { + return AddGrammarExpr( + {GrammarExprType::kExcludeToken, token_ids.data(), static_cast(token_ids.size())} + ); +} + +int32_t GrammarBuilder::AddRuleRef(int32_t rule_id) { + std::vector data; + data.push_back(rule_id); + return AddGrammarExpr({GrammarExprType::kRuleRef, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddSequence(const std::vector& elements) { + return AddGrammarExpr( + {GrammarExprType::kSequence, elements.data(), static_cast(elements.size())} + ); +} + +int32_t GrammarBuilder::AddChoices(const std::vector& choices) { + return AddGrammarExpr( + {GrammarExprType::kChoices, choices.data(), static_cast(choices.size())} + ); +} + +int32_t GrammarBuilder::AddTagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch) { + std::vector data; + data.reserve(tag_dispatch.tag_rule_pairs.size() * 2 + 2); + for (const auto& [tag, rule_id] : tag_dispatch.tag_rule_pairs) { + data.push_back(AddByteString(tag)); + data.push_back(rule_id); + } + data.push_back(static_cast(tag_dispatch.loop_after_dispatch)); + std::vector exclude_str_expr_ids; + for (const auto& exclude_str : tag_dispatch.excludes) { + exclude_str_expr_ids.push_back(AddByteString(exclude_str)); + } + data.push_back(AddChoices(exclude_str_expr_ids)); + return AddGrammarExpr( + {GrammarExprType::kTagDispatch, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch +) { + std::vector data; + data.push_back(static_cast(token_tag_dispatch.trigger_rule_pairs.size())); + for (const auto& [token_id, rule_id] : token_tag_dispatch.trigger_rule_pairs) { + data.push_back(token_id); + data.push_back(rule_id); + } + data.push_back(static_cast(token_tag_dispatch.loop_after_dispatch)); + data.push_back(static_cast(token_tag_dispatch.excludes.size())); + for (auto token_id : token_tag_dispatch.excludes) { + data.push_back(token_id); + } + return AddGrammarExpr( + {GrammarExprType::kTokenTagDispatch, data.data(), static_cast(data.size())} + ); +} + +int32_t GrammarBuilder::AddRepeat( + int32_t ref_rule_id, int32_t min_repeat_count, int32_t max_repeat_count +) { + std::vector data({ref_rule_id, min_repeat_count, max_repeat_count}); + return AddGrammarExpr({GrammarExprType::kRepeat, data.data(), static_cast(data.size())}); +} + +int32_t GrammarBuilder::AddRepeatFromExpr( + const std::string& cur_rule_name, + int32_t grammar_expr_id, + int32_t min_repeat_count, + int32_t max_repeat_count +) { + const auto& expr = GetGrammarExpr(grammar_expr_id); + int32_t ref_rule_id; + if (expr.type == GrammarExprType::kRuleRef) { + ref_rule_id = expr[0]; + } else { + ref_rule_id = AddRule(GetNewRuleName(cur_rule_name), grammar_expr_id); + } + return AddRepeat(ref_rule_id, min_repeat_count, max_repeat_count); +} + +int32_t GrammarBuilder::NumGrammarExprs() const { return grammar_->NumGrammarExprs(); } + +GrammarBuilder::GrammarExpr GrammarBuilder::GetGrammarExpr(int32_t grammar_expr_id) { + return grammar_->GetGrammarExpr(grammar_expr_id); +} + +int32_t GrammarBuilder::AddRule(const Rule& rule) { + EnsureRuleNameMap(); + int32_t id = static_cast(grammar_->rules_.size()); + grammar_->rules_.push_back(rule); + XGRAMMAR_CHECK(rule_name_to_id_.count(rule.name) == 0); + rule_name_to_id_[rule.name] = id; + return id; +} + +int32_t GrammarBuilder::AddRule(const std::string& name, int32_t body_expr_id) { + return AddRule({name, body_expr_id}); +} + +int32_t GrammarBuilder::AddRuleWithHint(const std::string& name_hint, int32_t body_expr_id) { + return AddRule({GetNewRuleName(name_hint), body_expr_id}); +} + +int32_t GrammarBuilder::NumRules() const { return grammar_->NumRules(); } + +const GrammarBuilder::Rule& GrammarBuilder::GetRule(int32_t rule_id) const { + return grammar_->rules_[rule_id]; +} + +int32_t GrammarBuilder::AddEmptyRule(const std::string& name) { return AddRule({name, -1}); } + +int32_t GrammarBuilder::AddEmptyRuleWithHint(const std::string& name_hint) { + return AddRule({GetNewRuleName(name_hint), -1}); +} + +void GrammarBuilder::UpdateRuleBody(int32_t rule_id, int32_t body_expr_id) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].body_expr_id = body_expr_id; +} + +void GrammarBuilder::UpdateRuleBody(std::string rule_name, int32_t body_expr_id) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateRuleBody(rule_id, body_expr_id); +} + +void GrammarBuilder::UpdateLookaheadAssertion(int32_t rule_id, int32_t lookahead_assertion_id) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].lookahead_assertion_id = lookahead_assertion_id; +} + +void GrammarBuilder::UpdateLookaheadExact(int32_t rule_id, bool is_exact) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].is_exact_lookahead = is_exact; +} + +void GrammarBuilder::UpdateRuleTemperature(int32_t rule_id, std::optional temperature) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].temperature = temperature; +} + +void GrammarBuilder::UpdateLookaheadAssertion( + std::string rule_name, int32_t lookahead_assertion_id +) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateLookaheadAssertion(rule_id, lookahead_assertion_id); +} + +void GrammarBuilder::UpdateMaxTokens(int32_t rule_id, int32_t max_tokens) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].max_tokens = max_tokens; +} + +void GrammarBuilder::UpdateMaxTokens(std::string rule_name, int32_t max_tokens) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateMaxTokens(rule_id, max_tokens); +} + +void GrammarBuilder::UpdateMaxChars(int32_t rule_id, int32_t max_chars) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].max_chars = max_chars; +} + +void GrammarBuilder::UpdateMaxChars(std::string rule_name, int32_t max_chars) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateMaxChars(rule_id, max_chars); +} + +void GrammarBuilder::UpdateCaptureName(int32_t rule_id, const std::string& capture_name) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].capture_name = capture_name; +} + +void GrammarBuilder::UpdateCaptureName(std::string rule_name, const std::string& capture_name) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateCaptureName(rule_id, capture_name); +} + +void GrammarBuilder::UpdateSuffixStopInfo(int32_t rule_id, const SuffixStopInfo& info) { + XGRAMMAR_CHECK(rule_id >= 0 && rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + XGRAMMAR_CHECK(info.hidden_suffix_bytes >= 0 && info.hidden_stop_bytes >= 0) + << "The number of hidden suffix/stop bytes must be non-negative."; + int32_t num_rules = static_cast(grammar_->rules_.size()); + XGRAMMAR_CHECK( + (info.body_rule_id == -1 && info.marker_rule_id == -1) || + (info.body_rule_id >= 0 && info.body_rule_id < num_rules && info.marker_rule_id >= 0 && + info.marker_rule_id < num_rules) + ) << "Capture-hidden helper rule ids must both be -1 or valid rule ids."; + + auto it = std::lower_bound( + grammar_->suffix_stop_infos_.begin(), + grammar_->suffix_stop_infos_.end(), + rule_id, + [](const SuffixStopInfo& existing, int32_t id) { return existing.rule_id < id; } + ); + if (info.IsEmpty()) { + if (it != grammar_->suffix_stop_infos_.end() && it->rule_id == rule_id) { + grammar_->suffix_stop_infos_.erase(it); + } + return; + } + SuffixStopInfo updated = info; + updated.rule_id = rule_id; + if (it != grammar_->suffix_stop_infos_.end() && it->rule_id == rule_id) { + *it = std::move(updated); + } else { + grammar_->suffix_stop_infos_.insert(it, std::move(updated)); + } +} + +void GrammarBuilder::UpdateSuffixStopInfo(std::string rule_name, const SuffixStopInfo& info) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateSuffixStopInfo(rule_id, info); +} + +void GrammarBuilder::UpdateLazy(int32_t rule_id, bool is_lazy) { + XGRAMMAR_CHECK(rule_id < static_cast(grammar_->rules_.size())) + << "Rule id " << rule_id << " is out of range."; + grammar_->rules_[rule_id].is_lazy = is_lazy; +} + +void GrammarBuilder::UpdateLazy(std::string rule_name, bool is_lazy) { + int32_t rule_id = GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not found."; + UpdateLazy(rule_id, is_lazy); +} + +std::string GrammarBuilder::GetNewRuleName(const std::string& name_hint) { + EnsureRuleNameMap(); + if (rule_name_to_id_.count(name_hint) == 0) { + return name_hint; + } + int* cnt = &next_cnt_per_hint_[name_hint]; + if (*cnt == 0) { + *cnt = 1; + } + while (rule_name_to_id_.count(name_hint + "_" + std::to_string(*cnt)) != 0) { + ++(*cnt); + } + return name_hint + "_" + std::to_string(*cnt); +} + +int32_t GrammarBuilder::GetRuleId(const std::string& name) const { + EnsureRuleNameMap(); + auto it = rule_name_to_id_.find(name); + if (it == rule_name_to_id_.end()) { + return -1; + } else { + return it->second; + } +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_builder.h b/zig/vendor/xgrammar/cpp/grammar_builder.h new file mode 100644 index 0000000000..7209a567ba --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_builder.h @@ -0,0 +1,292 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_builder.h + * \brief The header for the building the BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_BUILDER_H_ +#define XGRAMMAR_GRAMMAR_BUILDER_H_ + +#include + +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief Helper class to build a BNF grammar. + */ +class GrammarBuilder { + public: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using GrammarExpr = Grammar::Impl::GrammarExpr; + + /*! \brief One element of a character class, containing a lower and a upper bound. Both bounds are + * inclusive. + */ + struct CharacterClassElement { + int32_t lower; + int32_t upper; + }; + + /*! \brief Default constructor. Creates a new grammar object. */ + GrammarBuilder(); + + /*! \brief Constructor. Creates a new grammar object from an existing grammar. */ + GrammarBuilder(const Grammar& grammar); + + /*! + * \brief Create a builder bound to an existing grammar without copying it. Unlike the copy + * constructor above, appended exprs and rule updates are written directly into *grammar. The + * grammar must outlive the builder. + */ + static GrammarBuilder FromMutableGrammar(Grammar* grammar); + + /*! + * \brief Get the result grammar. This function will also set the root rule to the rule with the + * specified name. The rule should be already added to the grammar. + * \param root_rule_name The name of the root rule. Default is "root". + */ + Grammar Get(const std::string& root_rule_name = "root"); + + /*! + * \brief Get the result grammar. This function will also set the root rule to the rule with + * the specified id. The rule should be already added to the grammar. + * \param root_rule_id The id of the root rule. + */ + Grammar Get(int32_t root_rule_id); + + /****************** GrammarExpr handling ******************/ + + /*! \brief Add a grammar_expr and return the grammar_expr id. */ + int32_t AddGrammarExpr(const GrammarExpr& grammar_expr); + + /*! + * \brief Add a GrammarExpr for string stored in bytes. + * \param bytes A vector of int32_t, each representing a byte (0~255) in the string. + * The string is stored in int32 vector to match the storage format of the grammar. + */ + int32_t AddByteString(const std::vector& bytes); + + /*! + * \brief Add a GrammarExpr for string stored in bytes. + * \param str The string to be added. + */ + int32_t AddByteString(const std::string& str); + + /*! + * \brief Add a GrammarExpr for a regex. The pattern is stored as-is and compiled into an + * automaton by GrammarFSMBuilder. + * \param regex_str The regex pattern string. + * \param json_string Whether the regex matches the body of a JSON string literal. If true, + * the characters that must be escaped in a JSON string (the control characters, '"' and + * '\\') are excluded from every character match of the compiled automaton. + */ + int32_t AddRegex(const std::string& regex_str, bool json_string = false); + + /*! + * \brief Add a GrammarExpr for a substring expression, which matches every contiguous + * subsequence of the chunk list (including the empty one). The chunks are stored as-is and + * compiled into an automaton by GrammarFSMBuilder. + * \param chunks The list of byte string chunks. Chunks may be empty or repeated. + */ + int32_t AddSubstring(const std::vector& chunks); + + /*! + * \brief Add a GrammarExpr for a character class. + * \param elements A vector of CharacterClassElement, each containing a lower and a upper bound. + * \param is_negative Whether the character class is negated. + */ + int32_t AddCharacterClass( + const std::vector& elements, bool is_negative = false + ); + + /*! + * \brief Add a GrammarExpr for a star quantifier of a character class. + * \param elements A vector of CharacterClassElement, each containing a lower and a upper bound. + * \param is_negative Whether the character class is negated. + */ + int32_t AddCharacterClassStar( + const std::vector& elements, bool is_negative = false + ); + + /*! \brief Add a GrammarExpr for empty string.*/ + int32_t AddEmptyStr(); + + /*! \brief Add a GrammarExpr for kToken (token-level matching). */ + int32_t AddTokenSet(const std::vector& token_ids); + + /*! \brief Add a GrammarExpr for kExcludeToken (excluded token-level matching). */ + int32_t AddExcludeTokenSet(const std::vector& token_ids); + + /*! \brief Add a GrammarExpr for rule reference.*/ + int32_t AddRuleRef(int32_t rule_id); + + /*! \brief Add a GrammarExpr for GrammarExpr sequence.*/ + int32_t AddSequence(const std::vector& elements); + + /*! \brief Add a GrammarExpr for GrammarExpr choices.*/ + int32_t AddChoices(const std::vector& choices); + + /*! + * \brief Add a GrammarExpr for tag dispatch. + * \param tag_dispatch_list A list of pairs of tag_expr_id and rule_id. + */ + int32_t AddTagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + + /*! \brief Encode a TokenTagDispatch struct into a kTokenTagDispatch expr. */ + int32_t AddTokenTagDispatch(const Grammar::Impl::TokenTagDispatch& token_tag_dispatch); + + int32_t AddRepeat(int32_t ref_rule_id, int32_t min_repeat_count, int32_t max_repeat_count); + + /*! + * \brief Add a repeat GrammarExpr from an arbitrary grammar expression. If the expression is + * not a rule reference, a new rule is created to wrap it. + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param min_repeat_count Minimum repeat count (inclusive). + * \param max_repeat_count Maximum repeat count (inclusive), or -1 for unbounded. + */ + int32_t AddRepeatFromExpr( + const std::string& cur_rule_name, + int32_t grammar_expr_id, + int32_t min_repeat_count, + int32_t max_repeat_count + ); + + /*! \brief Get the number of grammar_exprs. */ + int32_t NumGrammarExprs() const; + + /*! \brief Get the grammar_expr with the given id. */ + GrammarExpr GetGrammarExpr(int32_t grammar_expr_id); + + /****************** Rule handling ******************/ + + /*! \brief Add a rule and return the rule id. */ + int32_t AddRule(const Rule& rule); + + int32_t AddRule(const std::string& name, int32_t body_expr_id); + + int32_t AddRuleWithHint(const std::string& name_hint, int32_t body_expr_id); + + int32_t NumRules() const; + + /*! \brief Get the rule with the given id. */ + const Rule& GetRule(int32_t rule_id) const; + + /*! + * \brief Add an rule without body, and return the rule id. The rule body should be set later + * with GrammarBuilder::UpdateRuleBody. This method is useful for cases where the rule id is + * required to build the rule body. + * \sa GrammarBuilder::UpdateRuleBody + */ + int32_t AddEmptyRule(const std::string& name); + + int32_t AddEmptyRuleWithHint(const std::string& name_hint); + + /*! + * \brief Update the rule body of the given rule, specified by rule id. Can be used to set the + * rule body of a rule inserted by GrammarBuilder::AddEmptyRule. + */ + void UpdateRuleBody(int32_t rule_id, int32_t body_expr_id); + + /*! + * \brief Update the rule body of the given rule, specified by rule name. Can be used to set the + * rule body of a rule inserted by GrammarBuilder::AddEmptyRule. + */ + void UpdateRuleBody(std::string rule_name, int32_t body_expr_id); + + /*! + * \brief Add a lookahead assertion to a rule referred by the given rule_id. The lookahead + * assertion should be a sequence GrammarExpr id. An id of -1 means no lookahead assertion. + */ + void UpdateLookaheadAssertion(int32_t rule_id, int32_t lookahead_assertion_id); + + void UpdateLookaheadExact(int32_t rule_id, bool is_exact = true); + + /*! \brief Set the sampling temperature associated with a rule. */ + void UpdateRuleTemperature(int32_t rule_id, std::optional temperature); + + /*! + * \brief Add a lookahead assertion to a rule referred by the given name. The lookahead + * assertion should be a sequence GrammarExpr id. An id of -1 means no lookahead assertion. + */ + void UpdateLookaheadAssertion(std::string rule_name, int32_t lookahead_assertion_id); + + /*! \brief Update the token budget of the rule referred by the given rule_id. -1 means none. */ + void UpdateMaxTokens(int32_t rule_id, int32_t max_tokens); + + /*! \brief Update the token budget of the rule referred by the given name. -1 means none. */ + void UpdateMaxTokens(std::string rule_name, int32_t max_tokens); + + /*! \brief Update the character budget of the rule referred by the given rule_id. -1 means none. + */ + void UpdateMaxChars(int32_t rule_id, int32_t max_chars); + + /*! \brief Update the character budget of the rule referred by the given name. -1 means none. */ + void UpdateMaxChars(std::string rule_name, int32_t max_chars); + + /*! + * \brief Update the capture group name of the rule referred by the given rule_id. An empty + * string means no capture. + */ + void UpdateCaptureName(int32_t rule_id, const std::string& capture_name); + + /*! + * \brief Update the capture group name of the rule referred by the given name. An empty string + * means no capture. + */ + void UpdateCaptureName(std::string rule_name, const std::string& capture_name); + + /*! \brief Set or clear the sparse suffix/stop metadata for a rule. */ + void UpdateSuffixStopInfo(int32_t rule_id, const SuffixStopInfo& info); + + /*! \brief Set or clear sparse suffix/stop metadata on the rule referred by name. */ + void UpdateSuffixStopInfo(std::string rule_name, const SuffixStopInfo& info); + + /*! \brief Set whether the rule referred by the given rule_id is lazy (committed-shortest). */ + void UpdateLazy(int32_t rule_id, bool is_lazy); + + /*! \brief Set whether the rule referred by the given name is lazy (committed-shortest). */ + void UpdateLazy(std::string rule_name, bool is_lazy); + + /*! + * \brief Find a name for a new rule starting with the given name hint. Some integer suffix (_1, + * _2, ...) may be added to avoid name conflict. + */ + std::string GetNewRuleName(const std::string& name_hint); + + /*! + * \brief Get the rule id of the rule with the given name. Return -1 if not found. + */ + int32_t GetRuleId(const std::string& name) const; + + private: + /*! + * \brief Build rule_name_to_id_ from the existing rules if it has not been built yet. + * Builders bound to an existing grammar (copy constructor, FromMutableGrammar) defer building + * the map to the first name-based operation, so id-only rewriting pays no name hashing cost. + */ + void EnsureRuleNameMap() const; + + // Mutable pointer to the grammar object. + std::shared_ptr grammar_; + // Map from rule name to rule id. Built lazily; see EnsureRuleNameMap. + mutable std::unordered_map rule_name_to_id_; + mutable bool rule_name_map_built_ = true; + // Cache of next suffix index per name_hint for GetNewRuleName. + std::unordered_map next_cnt_per_hint_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_BUILDER_H_ diff --git a/zig/vendor/xgrammar/cpp/grammar_compiler.cc b/zig/vendor/xgrammar/cpp/grammar_compiler.cc new file mode 100644 index 0000000000..dcd6264de1 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_compiler.cc @@ -0,0 +1,1725 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiler.cc + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "earley_parser.h" +#include "fsm.h" +#include "grammar_functor.h" +#include "grammar_impl.h" +#include "support/dynamic_bitset.h" +#include "support/int_set.h" +#include "support/logging.h" +#include "support/thread_pool.h" +#include "support/thread_safe_cache.h" +#include "support/utils.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/grammar.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +/************** AdaptiveTokenMaskCache Generator **************/ + +/*! \brief The concrete implementation of GrammarMatcherNode. */ +class GrammarMatcherForTokenMaskCache : public EarleyParser { + public: + GrammarMatcherForTokenMaskCache( + const Grammar& grammar, + const ParserState& init_state, + const std::unordered_map& + tag_dispatch_rule_id_to_second_slicing_bitset, + const TokenizerInfo& tokenizer_info, + std::optional& rule_level_cache + ) + : EarleyParser(grammar, init_state), + init_rule_id_(init_state.rule_id), + initial_state_(init_state), + tag_dispatch_rule_id_to_second_slicing_bitset_(tag_dispatch_rule_id_to_second_slicing_bitset + ), + tokenizer_info_(tokenizer_info), + rule_level_cache_(rule_level_cache) {} + /*! + * \brief Get the adaptive token mask for the given ParserState. + * \param is_root_rule Whether to consider the parent rule. If false, there will be + * no uncertain tokens. Useful for the root rule. + */ + AdaptiveTokenMask GetAdaptiveTokenMask(bool is_root_rule); + + /*! + * \brief Get the token mask for the given ParserState. + * \param first_char_mask The first character mask. + * \param is_root_rule Whether to consider the parent rule. If false, there will be + * no uncertain tokens. Useful for the root rule. + * \returns True if the rejected indices are filled as usual, False otherwise. + * It's used to determine which construction function will be used. + */ + bool GetTokenMaskWithFirstCharacterCheck( + const std::bitset<256>& first_char_mask, + bool is_root_rule, + const std::vector& token_edge_accepted + ); + + /*! + * \brief Adapt the cache with lookahead assertion. + * \param cache The adaptive token mask to be adapted. + * \param is_root_rule Whether to consider the parent rule. + */ + void AdaptCacheWithLookahead(AdaptiveTokenMask* cache, bool is_root_rule); + + private: + /*! \brief Check if a token can pass the lookahead assertion. */ + std::pair IsTokenPassLookaheadAssertion( + const std::string& token, const std::vector& can_reach_end_stack + ); + + /*! + * \brief Check if speculative calculation will be applied. + * \return first: whether speculative calculation is applicable. + * \return second: part of the first character mask, + * which can be used in speculative calculation. + */ + std::pair> GetSpeculativeCalculation(); + + /*! + * \brief Get the first character mask. + * \param first_character_mask the bitset to store the first character mask. + */ + void GetFirstCharacterMask(std::bitset<256>& first_character_mask); + + /*! + * \brief Compute sorted vocab indices accepted by token edges at the current FSM state. + * Token(ids) edges accept listed token IDs. + * ExcludeToken(ids) edges accept all tokens except listed IDs. + * \return Sorted, deduplicated vector of accepted sorted vocab indices. + */ + const std::vector& GetTokenEdgeAcceptedIndices(); + + // The id of the initial rule. + int32_t init_rule_id_; + + // The initial state of the parser. + ParserState initial_state_; + + /*! + * \brief This is a mapping from TagDispatch rule id to the bitset used for second slicing. + * \note If a rule is a TagDispatch rule, then there will be an AC automaton for its triggers. + * Which means that it can accept a lot of tokens. However, it will be slow to check a lot of + * tokens. The DynamicBitset here is used to do a second slicing: if a token's substr(1, n - 1) + * can be accepted by the start state of the AC automaton, then it will be True in the bitset. + * When we check a token, we first check if its first character can transit to the start state. + * If yes, then we check if it is in the bitset. If yes, then we accept it directly. + */ + const std::unordered_map& tag_dispatch_rule_id_to_second_slicing_bitset_; + + const TokenizerInfo& tokenizer_info_; + + std::optional rule_level_cache_; + + // Temporary data for GetAdaptiveTokenMask. + std::vector tmp_accepted_indices_; + std::vector tmp_rejected_indices_; + std::vector tmp_uncertain_indices_; + std::vector tmp_rejected_by_lookahead_indices_; + std::vector tmp_accepted_by_lookahead_indices_; + std::vector tmp_can_reach_end_stack_; + std::vector tmp_can_reach_end_prefix_or_stack_; + // Temporary data for GetTokenEdgeAcceptedIndices. + std::vector tmp_token_edge_accepted_; + std::vector tmp_token_edge_excluded_; +}; + +void GrammarMatcherForTokenMaskCache::AdaptCacheWithLookahead( + AdaptiveTokenMask* cache_ptr, bool is_root_rule +) { + AdaptiveTokenMask& cache = *cache_ptr; + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_nodes_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + const std::string* prev_token = nullptr; + bool is_exact_lookahead = grammar_->GetRule(init_rule_id_).is_exact_lookahead; + int prev_matched_size = 0; + int last_rejected_range = 0; + int last_uncertain_range = 0; + if (is_root_rule) { + tmp_rejected_indices_ = cache.uncertain_indices; + } else { + const auto& lookahead_id = grammar_->GetRule(init_rule_id_).lookahead_assertion_id; + if (lookahead_id == -1) { + return; + } + for (const auto& uncertain_index : cache.uncertain_indices) { + const auto& token = sorted_decoded_vocab[uncertain_index].second; + // Many tokens may contain the same prefix, so we will avoid unnecessary matching + // by finding the longest common prefix with the previous token. + bool accepted = true; + if (uncertain_index < last_rejected_range) { + tmp_rejected_indices_.push_back(uncertain_index); + continue; + } + if (uncertain_index < last_uncertain_range) { + // This token is already marked as uncertain. + continue; + } + if (prev_token != nullptr) { + int lcp_len = + std::mismatch(token.begin(), token.end(), prev_token->begin(), prev_token->end()) + .first - + token.begin(); + if (lcp_len > prev_matched_size) { + // Case 1. The common prefix is rejected by the matcher in the last token. Reject + // directly. + accepted = false; + } else if (lcp_len < prev_matched_size) { + // Case 2. The common prefix is shorter than the previous matched size. Rollback + // the non-common part. + PopLastStates(prev_matched_size - lcp_len); + tmp_can_reach_end_stack_.erase( + tmp_can_reach_end_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_stack_.end() + ); + tmp_can_reach_end_prefix_or_stack_.erase( + tmp_can_reach_end_prefix_or_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_prefix_or_stack_.end() + ); + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + prev_token = &token; + + if (accepted) { + // Accept the rest chars one by one. + for (int j = prev_matched_size; j < static_cast(token.size()); ++j) { + if (!Advance(token[j])) { + accepted = false; + break; + } + tmp_can_reach_end_stack_.push_back(IsCompleted()); + tmp_can_reach_end_prefix_or_stack_.push_back( + tmp_can_reach_end_stack_.back() || tmp_can_reach_end_prefix_or_stack_.back() + ); + prev_matched_size = j + 1; + } + } + + XGRAMMAR_DCHECK(!tmp_can_reach_end_prefix_or_stack_.empty()); + bool can_reach_end = tmp_can_reach_end_prefix_or_stack_.back(); + + XGRAMMAR_DCHECK(!accepted) << "All the tokens are at least uncertain!"; + if (can_reach_end && prev_matched_size > 0) { + auto [lookahead_accepted, lookahead_completed] = + IsTokenPassLookaheadAssertion(token, tmp_can_reach_end_stack_); + if ((!is_root_rule) && lookahead_accepted) { + if (lookahead_completed || !is_exact_lookahead) { + tmp_uncertain_indices_.push_back(uncertain_index); + } else { + tmp_accepted_indices_.push_back(uncertain_index); + } + } else { + tmp_rejected_indices_.push_back(uncertain_index); + last_rejected_range = subtree_nodes_range[uncertain_index]; + } + } else { + tmp_rejected_indices_.push_back(uncertain_index); + last_rejected_range = subtree_nodes_range[uncertain_index]; + } + } + } + + // This strategy ensures the consistency of the cache storage type in most cases. + // However, in this case, the storage type is inconsistent: + // 1. The original cache is accepted_indices, and rejected_indices is also small. + // After adapting with lookahead, |accepted_indices| + |accepted_by_lookahead_indices| > + // |rejected_indices| + |rejected_by_lookahead_indices|, and |rejected_indices| + + // |rejected_by_lookahead_indices| < AdaptiveTokenMask::USE_BITSET_THRESHOLD. In this case, it + // should be kRejected, but ignored. + // 2. The original cache is rejected_indices, and accepted_indices is also small. + // After adapting with lookahead, |accepted_indices| + |accepted_by_lookahead_indices| < + // |rejected_indices| + |rejected_by_lookahead_indices|, and |accepted_indices| + + // |accepted_by_lookahead_indices| < AdaptiveTokenMask::USE_BITSET_THRESHOLD. In this case, it + // should be kAccepted, but ignored. These two cases are very rare in practice, and the impact is + // very limited, so we ignore them for simplicity. + cache.uncertain_indices = tmp_uncertain_indices_; + switch (cache.store_type) { + case AdaptiveTokenMask::StoreType::kAccepted: { + if (cache.accepted_indices.size() + tmp_accepted_indices_.size() < + AdaptiveTokenMask::USE_BITSET_THRESHOLD) { + IntsetUnion(&cache.accepted_indices, tmp_accepted_indices_); + break; + } + // Transform to bitset. + cache.store_type = AdaptiveTokenMask::StoreType::kAcceptedBitset; + cache.accepted_bitset = DynamicBitset(tokenizer_info_.GetVocabSize()); + for (const auto& accepted_index : cache.accepted_indices) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + for (const auto& accepted_index : tmp_accepted_indices_) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + cache.accepted_indices.clear(); + break; + } + case AdaptiveTokenMask::StoreType::kRejected: { + if (cache.rejected_indices.size() + tmp_rejected_indices_.size() < + AdaptiveTokenMask::USE_BITSET_THRESHOLD) { + IntsetUnion(&cache.rejected_indices, tmp_rejected_indices_); + break; + } + // Transform to bitset. + cache.store_type = AdaptiveTokenMask::StoreType::kAcceptedBitset; + cache.accepted_bitset = DynamicBitset(tokenizer_info_.GetVocabSize()); + cache.accepted_bitset.Set(); + for (const auto& special_index : tokenizer_info_.GetSpecialTokenIds()) { + cache.accepted_bitset.Reset(special_index); + } + for (const auto& uncertain_index : cache.uncertain_indices) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[uncertain_index].first); + } + for (const auto& rejected_index : cache.rejected_indices) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[rejected_index].first); + } + for (const auto& rejected_index : tmp_rejected_indices_) { + cache.accepted_bitset.Reset(sorted_decoded_vocab[rejected_index].first); + } + cache.rejected_indices.clear(); + break; + } + case AdaptiveTokenMask::StoreType::kAcceptedBitset: { + for (const auto& accepted_index : tmp_accepted_indices_) { + cache.accepted_bitset.Set(sorted_decoded_vocab[accepted_index].first); + } + break; + } + } +} + +std::pair GrammarMatcherForTokenMaskCache::IsTokenPassLookaheadAssertion( + const std::string& token, const std::vector& can_reach_end_stack +) { + bool accepted = true; + bool can_reach_end = true; + auto lookahead_assertion_id = grammar_->GetRule(init_rule_id_).lookahead_assertion_id; + if (lookahead_assertion_id == -1) { + return {accepted, can_reach_end}; + } + auto lookahead_state = + ParserState(/*rule_id*/ -1, lookahead_assertion_id, 0, ParserState::kNoPrevInputPos, 0); + PushStateAndExpand(lookahead_state); + int token_len = token.size(); + if (IsCompleted()) { + // If the lookahead assertion is already completed, we can accept the token. + PopLastStates(1); + return {accepted, can_reach_end}; + } + + // Find all positions that can come to and end. Then check if the suffix from that position + // can be accepted by the lookahead assertion. + for (int i = static_cast(can_reach_end_stack.size()) - 1; i >= 0; --i) { + if (!can_reach_end_stack[i]) { + continue; + } + int last_accept_pos = i - 1; + for (int pos = i; pos < token_len; ++pos) { + if (!Advance(token[pos])) { + break; + } + last_accept_pos = pos; + // Case 1. The whole rule is finished. + if (IsCompleted()) { + // accepted chars: pos - i + 1 + // we need to rollback the pushed initial state as well + PopLastStates(pos - i + 2); + return {accepted, can_reach_end}; + } + } + // Case 2. The whole token is accepted + if (last_accept_pos == token_len - 1) { + PopLastStates(last_accept_pos - i + 2); + can_reach_end = false; + return {accepted, can_reach_end}; + } + // Case 3. The token is not accepted. Check the next position. + PopLastStates(last_accept_pos - i + 1); + } + + PopLastStates(1); + can_reach_end = false; + accepted = false; + return {accepted, can_reach_end}; +} + +// Comparator for std::pair based on the string value. +class IntStringPairComparator { + public: + bool operator()( + const std::pair& lhs, const std::pair& rhs + ) const { + return lhs.second < rhs.second; + } +}; + +int GetPossibleTokenIntervals( + const std::vector>& sorted_decoded_vocab, + const std::bitset<256>& first_char_mask, + std::vector>& possible_intervals +) { + int possible_token_num = 0; + int matched_size = 0; + int last_interval_end = -1; + for (int32_t i = 0; i < 256; i++) { + if (first_char_mask[i]) { + if (last_interval_end == -1) { + last_interval_end = i; + } + } else { + if (last_interval_end != -1) { + int32_t interval_left_end = + std::lower_bound( + sorted_decoded_vocab.begin() + matched_size, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(last_interval_end))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + int32_t interval_right_end = std::lower_bound( + sorted_decoded_vocab.begin() + interval_left_end, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(i))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + possible_intervals.emplace_back(interval_left_end, interval_right_end); + possible_token_num += interval_right_end - interval_left_end; + last_interval_end = -1; + matched_size = interval_right_end; + } + } + } + + if (last_interval_end != -1) { + // If the last interval is not closed, we need to close it. + int32_t interval_left_end = + std::lower_bound( + sorted_decoded_vocab.begin() + matched_size, + sorted_decoded_vocab.end(), + std::make_pair(0, std::string(1, static_cast(last_interval_end))), + IntStringPairComparator() + ) - + sorted_decoded_vocab.begin(); + possible_intervals.emplace_back(interval_left_end, sorted_decoded_vocab.size()); + possible_token_num += sorted_decoded_vocab.size() - interval_left_end; + } + return possible_token_num; +} + +std::pair> GrammarMatcherForTokenMaskCache::GetSpeculativeCalculation() { + using GrammarExprType = Grammar::Impl::GrammarExprType; + // If the initial rule is a tag dispatch, we will check if it can achieve its initial state. + const auto& rule = grammar_->GetRule(init_rule_id_); + if (rule.is_lazy) { + // The fast path assumes greedy self-loop extension is always legal, which does not hold for + // committed-shortest rules; they must go through the full per-token simulation. + return {false, std::bitset<256>()}; + } + const auto& rule_body = grammar_->GetGrammarExpr(rule.body_expr_id); + if (rule_body.type == GrammarExprType::kTagDispatch) { + std::bitset<256> speculative_mask; + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id)) { + if (edge.target != fsm.GetFsm().GetStart()) { + continue; + } + if (!edge.IsCharRange()) { + continue; + } + for (int32_t ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + } + return {true, speculative_mask}; + } + + // Check if the initial state is self-recursive-like via FSM. + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + bool can_be_applied = false; + std::bitset<256> speculative_mask; + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + XGRAMMAR_DCHECK(initial_state_.element_id < fsm.GetFsm().NumStates()) + << "Initial State's element id cannot exceed the whole FSM's number of states."; + for (const auto& edge : fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id)) { + if (edge.IsCharRange()) { + // Case A: The edge is towards itself. + if (edge.target == initial_state_.element_id) { + can_be_applied = true; + for (int ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + continue; + } + + // Case B: The state is the start state, and there's an edge to another state, + // which calls the fsm itself. + if (fsm.GetFsm().GetStart() == initial_state_.element_id) { + for (const auto& next_edge : fsm.GetFsm().GetFsm().GetEdges(edge.target)) { + if ((next_edge.IsRuleRef() && next_edge.GetRefRuleId() == init_rule_id_) || + (next_edge.IsRepeatRef() && + fsm.GetFsm().GetFsm().GetRepeatEdgeInfo(next_edge.GetAuxIndex()).RuleId() == + init_rule_id_)) { + can_be_applied = true; + for (int ch = edge.min; ch <= edge.max; ++ch) { + speculative_mask.set(ch); + } + break; + } + } + } + } + } + return {can_be_applied, speculative_mask}; +} + +bool GrammarMatcherForTokenMaskCache::GetTokenMaskWithFirstCharacterCheck( + const std::bitset<256>& first_char_mask, + bool is_root_rule, + const std::vector& token_edge_accepted +) { + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_nodes_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + // the pair (a, b) means [a, b). Intialize the possible intervals. + std::vector> possible_intervals; + int possible_token_num = + GetPossibleTokenIntervals(sorted_decoded_vocab, first_char_mask, possible_intervals); + + // Check if the type of the mask can be rejected. + tmp_accepted_indices_.reserve(possible_token_num); + bool fill_reject_indices = + (sorted_decoded_vocab.size() - possible_token_num) < AdaptiveTokenMask::USE_BITSET_THRESHOLD; + + XGRAMMAR_DCHECK(possible_intervals.size() > 0) + << "There should be at least one possible interval for the first character mask."; + + if (possible_intervals[0].first != 0 && fill_reject_indices) { + for (int i = 0; i < possible_intervals[0].first; ++i) { + tmp_rejected_indices_.push_back(i); + } + } + + XGRAMMAR_DCHECK(init_rule_id_ != -1 && grammar_->per_rule_fsms[init_rule_id_].has_value()); + auto [speculative_calculation, speculative_mask] = GetSpeculativeCalculation(); + if (has_char_budget_rules_) { + speculative_calculation = false; + } + + int prev_matched_size = 0; + int last_rejected_range = 0; + const bool& is_exact_lookahead = grammar_->GetRule(init_rule_id_).is_exact_lookahead; + std::optional definite_accepted_bitset = std::nullopt; + const bool is_tag_dispatch_rule = + grammar_->GetGrammarExpr(grammar_->GetRule(init_rule_id_).body_expr_id).type == + Grammar::Impl::GrammarExprType::kTagDispatch; + if (is_tag_dispatch_rule) { + XGRAMMAR_DCHECK(tag_dispatch_rule_id_to_second_slicing_bitset_.count(init_rule_id_) > 0); + definite_accepted_bitset = &tag_dispatch_rule_id_to_second_slicing_bitset_.at(init_rule_id_); + } + + const std::string* prev_token = nullptr; + int32_t skip_ptr = 0; + const int32_t skip_size = static_cast(token_edge_accepted.size()); + for (size_t interval_idx = 0; interval_idx < possible_intervals.size(); ++interval_idx) { + const auto& interval = possible_intervals[interval_idx]; + for (int i = interval.first; i < interval.second; ++i) { + // Skip tokens already accepted by token edges (avoid expensive Earley simulation). + while (skip_ptr < skip_size && token_edge_accepted[skip_ptr] < i) ++skip_ptr; + if (skip_ptr < skip_size && token_edge_accepted[skip_ptr] == i) continue; + + // Check if the current token is in the rejected range. i.e. check if the current token + // is on the subtree of the rejected token. + if (i < last_rejected_range) { + if (fill_reject_indices) { + tmp_rejected_indices_.push_back(i); + fill_reject_indices = + tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } else { + i = last_rejected_range - 1; + } + continue; + } + const auto& token = sorted_decoded_vocab[i].second; + // This optimization is useful for simple self-recursive rules, like string content. + if (speculative_calculation) { + // Optimization for tag dispatch rules. + if (definite_accepted_bitset.has_value()) { + // If the token is empty, it must be accepted. + if (token.empty()) { + tmp_accepted_indices_.push_back(i); + continue; + } + // If the token doesn't contain tags or stop strings since the second character, and it + // will transit to the start state after consuming the first character, it must be + // accepted. + if (speculative_mask[static_cast(token[0])] && + (*definite_accepted_bitset.value())[i]) { + tmp_accepted_indices_.push_back(i); + continue; + } + } else { + bool all_accepted = true; + for (char ch : token) { + // If the first character is not the ascii character or can't be accepted by the + // first character mask, we need to check them in the parser. + if (isascii(ch) == 0 || !speculative_mask[static_cast(ch)]) { + all_accepted = false; + break; + } + } + if (all_accepted) { + tmp_accepted_indices_.push_back(i); + continue; + } + } + } + // Many tokens may contain the same prefix, so we will avoid unnecessary matching + // by finding the longest common prefix with the previous token. + bool accepted = true; + if (prev_token != nullptr) { + int lcp_len = + std::mismatch(token.begin(), token.end(), prev_token->begin(), prev_token->end()) + .first - + token.begin(); + if (lcp_len > prev_matched_size) { + // Case 1. The common prefix is rejected by the matcher in the last token. Reject + // directly. + accepted = false; + } else if (lcp_len < prev_matched_size) { + // Case 2. The common prefix is shorter than the previous matched size. Rollback + // the non-common part. + PopLastStates(prev_matched_size - lcp_len); + tmp_can_reach_end_stack_.erase( + tmp_can_reach_end_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_stack_.end() + ); + tmp_can_reach_end_prefix_or_stack_.erase( + tmp_can_reach_end_prefix_or_stack_.end() - (prev_matched_size - lcp_len), + tmp_can_reach_end_prefix_or_stack_.end() + ); + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + prev_token = &token; + + if (accepted) { + // Accept the rest chars one by one. + for (int j = prev_matched_size; j < static_cast(token.size()); ++j) { + if (!Advance(token[j])) { + accepted = false; + break; + } + tmp_can_reach_end_stack_.push_back(IsCompleted()); + tmp_can_reach_end_prefix_or_stack_.push_back( + tmp_can_reach_end_stack_.back() || tmp_can_reach_end_prefix_or_stack_.back() + ); + prev_matched_size = j + 1; + } + } + + bool can_reach_end = tmp_can_reach_end_prefix_or_stack_.back(); + + if (accepted) { + if (HasEnteredCharBudget()) { + tmp_uncertain_indices_.push_back(i); + } else { + tmp_accepted_indices_.push_back(i); + } + } else if (can_reach_end && prev_matched_size > 0) { + auto [lookahead_accepted, lookahead_completed] = + IsTokenPassLookaheadAssertion(token, tmp_can_reach_end_stack_); + if ((!is_root_rule) && lookahead_accepted) { + if (lookahead_completed || !is_exact_lookahead) { + tmp_uncertain_indices_.push_back(i); + } else if (HasEnteredCharBudget()) { + tmp_uncertain_indices_.push_back(i); + } else { + tmp_accepted_indices_.push_back(i); + tmp_accepted_by_lookahead_indices_.push_back(i); + } + } else { + for (int j = i; j < subtree_nodes_range[i]; j++) { + tmp_rejected_indices_.push_back(j); + tmp_rejected_by_lookahead_indices_.push_back(j); + } + i = subtree_nodes_range[i] - 1; // Skip the subtree nodes. + } + } else { + tmp_rejected_indices_.push_back(i); + last_rejected_range = subtree_nodes_range[i]; + fill_reject_indices = + tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } + } + if (interval_idx != possible_intervals.size() - 1 && fill_reject_indices) { + const auto& next_interval = possible_intervals[interval_idx + 1]; + for (int i = interval.second; i < next_interval.first; ++i) { + tmp_rejected_indices_.push_back(i); + } + fill_reject_indices = tmp_rejected_indices_.size() >= AdaptiveTokenMask::USE_BITSET_THRESHOLD + ? false + : fill_reject_indices; + } + } + + // Rollback the last matched part. + PopLastStates(prev_matched_size); + + if (possible_intervals.back().second != static_cast(sorted_decoded_vocab.size()) && + fill_reject_indices) { + // If the last interval is not closed, we need to reject the rest tokens. + for (int i = possible_intervals.back().second; + i < static_cast(sorted_decoded_vocab.size()); + ++i) { + tmp_rejected_indices_.push_back(i); + } + } + + return fill_reject_indices; +} + +void GrammarMatcherForTokenMaskCache::GetFirstCharacterMask(std::bitset<256>& first_character_mask +) { + first_character_mask.reset(); + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + const auto& edges = fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id); + for (const auto& edge : edges) { + if (edge.IsCharRange()) { + for (int c = edge.min; c <= edge.max; ++c) { + first_character_mask[c] = true; + } + } + } +} + +const std::vector& GrammarMatcherForTokenMaskCache::GetTokenEdgeAcceptedIndices() { + // Compute sorted vocab indices accepted by Token(ids) and ExcludeToken(ids) edges. + // Result is stored in tmp_token_edge_accepted_. + + tmp_token_edge_accepted_.clear(); + tmp_token_edge_excluded_.clear(); + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[init_rule_id_].has_value()); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + const auto& edges = fsm.GetFsm().GetFsm().GetEdges(initial_state_.element_id); + + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + int32_t sorted_size = static_cast(sorted_decoded_vocab.size()); + const auto& tid_to_sorted = tokenizer_info_.ImplPtr()->GetTokenIdToSortedVocabIndex(); + // Out-of-vocabulary ids are rejected before compilation starts (CheckTokenIdsInVocab); guard + // here as well because this runs on worker threads, where an out-of-bounds read cannot be + // reported as an error. + auto sorted_index = [&](int32_t tid) { + return tid >= 0 && tid < static_cast(tid_to_sorted.size()) ? tid_to_sorted[tid] : -1; + }; + + bool has_exclude_token = false; + + for (const auto& edge : edges) { + if (edge.IsToken()) { + auto info = fsm.GetFsm().GetFsm().GetTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + if (int32_t index = sorted_index(info.TokenIds()[i]); index >= 0) { + tmp_token_edge_accepted_.push_back(index); + } + } + } else if (edge.IsExcludeToken()) { + has_exclude_token = true; + auto info = fsm.GetFsm().GetFsm().GetExcludeTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + if (int32_t index = sorted_index(info.TokenIds()[i]); index >= 0) { + tmp_token_edge_excluded_.push_back(index); + } + } + } + } + + // Token-only: result = token_accepted + if (!has_exclude_token) { + if (!tmp_token_edge_accepted_.empty()) { + std::sort(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()); + tmp_token_edge_accepted_.erase( + std::unique(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()), + tmp_token_edge_accepted_.end() + ); + } + return tmp_token_edge_accepted_; + } + + // ExcludeToken: result = [0, sorted_size) - (excluded - token_accepted) + // Token(ids) overrides ExcludeToken(ids) when both present. + if (!tmp_token_edge_accepted_.empty()) { + std::sort(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()); + tmp_token_edge_accepted_.erase( + std::unique(tmp_token_edge_accepted_.begin(), tmp_token_edge_accepted_.end()), + tmp_token_edge_accepted_.end() + ); + } + std::sort(tmp_token_edge_excluded_.begin(), tmp_token_edge_excluded_.end()); + tmp_token_edge_excluded_.erase( + std::unique(tmp_token_edge_excluded_.begin(), tmp_token_edge_excluded_.end()), + tmp_token_edge_excluded_.end() + ); + IntsetDifference(&tmp_token_edge_excluded_, tmp_token_edge_accepted_); + IntsetComplement(&tmp_token_edge_accepted_, sorted_size, tmp_token_edge_excluded_); + return tmp_token_edge_accepted_; +} + +AdaptiveTokenMask GrammarMatcherForTokenMaskCache::GetAdaptiveTokenMask(bool is_root_rule) { + tmp_accepted_indices_.clear(); + tmp_rejected_indices_.clear(); + tmp_uncertain_indices_.clear(); + tmp_rejected_by_lookahead_indices_.clear(); + tmp_accepted_by_lookahead_indices_.clear(); + tmp_can_reach_end_prefix_or_stack_.clear(); + tmp_can_reach_end_stack_.clear(); + // For every character in the current token, stores whether it is possible to reach the end of + // the rule when matching until this character. Store it in a stack for later rollback. + tmp_can_reach_end_stack_.push_back(false); + tmp_can_reach_end_prefix_or_stack_.push_back(false); + + // Try to get the crossing cache. + bool rule_level_cache_is_available = !has_char_budget_rules_ && rule_level_cache_.has_value() && + grammar_->per_rule_fsm_hashes[init_rule_id_].has_value(); + std::optional fsm_hash = std::nullopt; + int32_t new_state_id = -1; + std::optional crossing_cache = std::nullopt; + int lookahead_id = grammar_->GetRule(initial_state_.rule_id).lookahead_assertion_id; + bool is_exact_lookahead = grammar_->GetRule(initial_state_.rule_id).is_exact_lookahead; + std::optional lookahead_hash = std::nullopt; + if (rule_level_cache_is_available) { + lookahead_hash = GrammarFSMHasher::HashSequence(grammar_, lookahead_id); + const auto& original_to_new_id = grammar_->per_rule_fsm_new_state_ids[init_rule_id_]; + fsm_hash = grammar_->per_rule_fsm_hashes[init_rule_id_].value(); + for (const auto& original_new_pair : original_to_new_id) { + if (original_new_pair.first == initial_state_.element_id) { + new_state_id = original_new_pair.second; + break; + } + } + XGRAMMAR_DCHECK(new_state_id != -1); + const auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + if (lookahead_hash.has_value()) { + crossing_cache = rule_level_cache_->GetCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum() + ); + if (crossing_cache.has_value()) { + // A perfect match. + return crossing_cache.value(); + } + } + crossing_cache = rule_level_cache_->GetCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum() + ); + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + if (crossing_cache.has_value()) { + AdaptCacheWithLookahead(&crossing_cache.value(), is_root_rule); + return std::move(crossing_cache.value()); + } + } + + std::bitset<256> first_character_mask; + GetFirstCharacterMask(first_character_mask); + + // Token edge accepted indices (for byte path skip + merge). + const auto& token_edge_accepted = GetTokenEdgeAcceptedIndices(); + + // Byte path: skip tokens already accepted by token edges. + bool rejected_filled; + if (first_character_mask.none()) { + rejected_filled = false; + } else { + rejected_filled = GetTokenMaskWithFirstCharacterCheck( + first_character_mask, is_root_rule, token_edge_accepted + ); + } + + // Token edges are rechecked at runtime when a character budget is present because accepting + // one can enter a budgeted rule within the same token. + if (!token_edge_accepted.empty()) { + if (has_char_budget_rules_) { + IntsetUnion(&tmp_uncertain_indices_, token_edge_accepted); + } else { + IntsetUnion(&tmp_accepted_indices_, token_edge_accepted); + IntsetDifference(&tmp_uncertain_indices_, token_edge_accepted); + } + IntsetDifference(&tmp_rejected_indices_, token_edge_accepted); + } + if (rejected_filled) { + auto return_value = AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + tmp_accepted_indices_, + tmp_rejected_indices_, + tmp_uncertain_indices_ + ); + if (rule_level_cache_is_available) { + if (lookahead_id == -1 && !is_root_rule) { + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum(), return_value + ); + return return_value; + } + + // We can add a cache for basic fsm, and a better one for lookahead. + // All the tokens rejected by lookahead should be uncertain. + IntsetUnion(&tmp_uncertain_indices_, tmp_rejected_by_lookahead_indices_); + IntsetUnion(&tmp_uncertain_indices_, tmp_accepted_by_lookahead_indices_); + std::vector rejected_indices_without_lookahead; + std::vector accepted_indices_without_lookahead; + rejected_indices_without_lookahead.reserve( + tmp_rejected_indices_.size() - tmp_rejected_by_lookahead_indices_.size() + ); + accepted_indices_without_lookahead.reserve( + tmp_accepted_indices_.size() - tmp_accepted_by_lookahead_indices_.size() + ); + std::set_difference( + tmp_rejected_indices_.begin(), + tmp_rejected_indices_.end(), + tmp_rejected_by_lookahead_indices_.begin(), + tmp_rejected_by_lookahead_indices_.end(), + std::back_inserter(rejected_indices_without_lookahead) + ); + std::set_difference( + tmp_accepted_indices_.begin(), + tmp_accepted_indices_.end(), + tmp_accepted_by_lookahead_indices_.begin(), + tmp_accepted_by_lookahead_indices_.end(), + std::back_inserter(accepted_indices_without_lookahead) + ); + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + fsm_hash.value(), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + accepted_indices_without_lookahead, + rejected_indices_without_lookahead, + tmp_uncertain_indices_ + ) + ); + if (lookahead_hash.has_value()) { + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + rule_level_cache_->AddCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + return_value + ); + } + } + return return_value; + } else { + auto return_value = AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + tmp_accepted_indices_, + tmp_uncertain_indices_ + ); + + if (rule_level_cache_is_available) { + // Prepare for cache. + auto& fsm = grammar_->per_rule_fsms[init_rule_id_].value(); + if (lookahead_id == -1 && !is_root_rule) { + // If the rule doesn't have a lookahead, then it is exactly the same fsm. + rule_level_cache_->AddCache( + fsm_hash.value(), new_state_id, fsm.GetNodeNum(), fsm.GetEdgeNum(), return_value + ); + return return_value; + } + + // Add 2 caches. + IntsetUnion(&tmp_uncertain_indices_, tmp_rejected_by_lookahead_indices_); + IntsetUnion(&tmp_uncertain_indices_, tmp_accepted_by_lookahead_indices_); + std::vector accepted_indices_without_lookahead; + accepted_indices_without_lookahead.reserve( + tmp_accepted_indices_.size() - tmp_accepted_by_lookahead_indices_.size() + ); + std::set_difference( + tmp_accepted_indices_.begin(), + tmp_accepted_indices_.end(), + tmp_accepted_by_lookahead_indices_.begin(), + tmp_accepted_by_lookahead_indices_.end(), + std::back_inserter(accepted_indices_without_lookahead) + ); + rule_level_cache_->AddCache( + fsm_hash.value(), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + AdaptiveTokenMask( + tokenizer_info_.GetVocabSize(), + tokenizer_info_.GetSortedDecodedVocab(), + accepted_indices_without_lookahead, + tmp_uncertain_indices_ + ) + ); + + if (lookahead_hash.has_value()) { + rule_level_cache_->AddCache( + HashCombine(fsm_hash.value(), lookahead_hash.value(), is_exact_lookahead), + new_state_id, + fsm.GetNodeNum(), + fsm.GetEdgeNum(), + return_value + ); + } + } + return return_value; + } +} + +/******************* GrammarCompilerNoCache *******************/ + +/*! + * \brief The base class for the grammar compiler. Handles the compilation logic without cache. + */ +class GrammarCompilerSub { + public: + GrammarCompilerSub( + const TokenizerInfo& tokenizer_info, + int max_threads, + std::optional rule_level_cache + ) + : tokenizer_info_(tokenizer_info), + max_threads_(max_threads), + rule_level_cache_(rule_level_cache) {} + + CompiledGrammar CompileBuiltinJSONGrammar(); + + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order + ); + + CompiledGrammar CompileRegex(const std::string& regex); + + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars + ); + + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + CompiledGrammar CompileGrammar(const Grammar& grammar); + + CompiledGrammar CompileGrammar(const std::string& ebnf_str, std::string root_rule_name); + + private: + /*! \brief The main logic. Compile the grammar with multi-threading. */ + CompiledGrammar MultiThreadCompileGrammar(Grammar grammar); + /*! \brief Optimization for TagDispatch. + * \param compiled_grammar_impl the compiled_grammar to be optimized. + * \param tag_dispatch_rule_id_to_second_slicing_bitset Return value. Mapping from the rule_id to + * the definite accepted token mask. + */ + void TagDispatchOptimization( + std::shared_ptr compiled_grammar_impl, + std::unordered_map* tag_dispatch_rule_id_to_second_slicing_bitset + ); + + /*! \brief The vocabulary associated with this storage class. */ + const TokenizerInfo tokenizer_info_; + /*! \brief The maximum number of threads to use. */ + const int max_threads_; + + /*! \brief The manager of the rule level cache.*/ + std::optional rule_level_cache_; +}; + +/*! + * \brief Check that every token id written in Token(...), ExcludeToken(...) and TokenTagDispatch + * exists in the vocabulary. They are used as indices into per-token arrays during compilation, + * which runs on worker threads where an error cannot be reported. + */ +static void CheckTokenIdsInVocab(const Grammar& grammar, int vocab_size) { + const CompactFSM& fsm = grammar->complete_fsm; + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (!edge.IsToken() && !edge.IsExcludeToken()) { + continue; + } + // Token and exclude-token edges share the [count, token_id_0, ...] aux layout. + const auto info = fsm.GetTokenEdgeInfo(edge.GetAuxIndex()); + for (int32_t i = 0; i < info.Count(); ++i) { + XGRAMMAR_CHECK(info.TokenIds()[i] >= 0 && info.TokenIds()[i] < vocab_size) + << "Token id " << info.TokenIds()[i] + << " in the grammar is out of the vocabulary range [0, " << vocab_size << ")"; + } + } + } +} + +CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_unoptimized) { + auto compiled_grammar_impl = std::make_shared(); + compiled_grammar_impl->grammar = GrammarOptimizer::Apply(grammar_unoptimized); + compiled_grammar_impl->tokenizer_info = tokenizer_info_; + CheckTokenIdsInVocab(compiled_grammar_impl->grammar, tokenizer_info_.GetVocabSize()); + if (tokenizer_info_.GetVocabSize() == 0) { + return CompiledGrammar(compiled_grammar_impl); + } + std::unordered_map tag_dispatch_rule_id_to_second_slicing_bitset; + TagDispatchOptimization(compiled_grammar_impl, &tag_dispatch_rule_id_to_second_slicing_bitset); + + // If the compiler is cache-enabled, then we hash the grammars for crossing-grammar caching. + if (rule_level_cache_.has_value()) { + GrammarFSMHasher().Apply(&compiled_grammar_impl->grammar); + } + // Step 3. Compute the adaptive token mask cache + // The token mask cache is computed for these positions in the grammar: + // 1. All character class or character class star (with last_utf8_bytes=0, 1, 2, 3) + // 2. All byte strings (with element_in_string=0, 1, 2, ...) + // since other positions will be expanded to the above positions + + // TODO(Charlie): Figure out how to support ThreadPool and std::mutex in WebAssembly. + // Only declare ThreadPool and mutex if max_threads > 1, so when max_threads = 1, we do + // not need ThreadPool or std::mutex, which throws error in runtime in WebAssembly. + std::optional thread_pool; + std::optional adaptive_token_mask_cache_mutex; + if (max_threads_ > 1) { + thread_pool.emplace(max_threads_); + adaptive_token_mask_cache_mutex.emplace(); + } + + auto add_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) { + auto grammar_matcher = GrammarMatcherForTokenMaskCache( + compiled_grammar_impl->grammar, + state, + tag_dispatch_rule_id_to_second_slicing_bitset, + tokenizer_info_, + rule_level_cache_ + ); + auto cur_adaptive_token_mask_cache = grammar_matcher.GetAdaptiveTokenMask(is_root_rule); + if (max_threads_ > 1) { + std::lock_guard lock(adaptive_token_mask_cache_mutex.value()); + compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache; + } else { + compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache; + } + }; + + auto add_task_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) { + // Execute depending on whether we use thread_pool + if (max_threads_ > 1) { + thread_pool->Execute([add_adaptive_token_mask, state, is_root_rule]() { + add_adaptive_token_mask(state, is_root_rule); + }); + } else { + add_adaptive_token_mask(state, is_root_rule); + } + }; + + auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId(); + + for (int32_t rule_id = 0; rule_id < static_cast(compiled_grammar_impl->grammar->NumRules()); + ++rule_id) { + auto rule = compiled_grammar_impl->grammar->GetRule(rule_id); + const auto& rule_fsm = compiled_grammar_impl->grammar->per_rule_fsms[rule_id]; + XGRAMMAR_DCHECK(rule_fsm.has_value()); + auto cur_stack_element = + ParserState(rule_id, rule.body_expr_id, 0, ParserState::kNoPrevInputPos, 0); + std::unordered_set reachable_states; + rule_fsm->GetFsm().GetReachableStates(&reachable_states); + for (int i : reachable_states) { + cur_stack_element.element_id = i; + if (!rule_fsm->GetFsm().IsScanableState(i)) { + continue; + } + add_task_adaptive_token_mask(cur_stack_element, rule_id == root_rule_id); + } + } + + if (max_threads_ > 1) { + thread_pool->Join(); + } + + return CompiledGrammar(compiled_grammar_impl); +} + +CompiledGrammar GrammarCompilerSub::CompileBuiltinJSONGrammar() { + return MultiThreadCompileGrammar(Grammar::BuiltinJSONGrammar()); +} + +CompiledGrammar GrammarCompilerSub::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + return MultiThreadCompileGrammar(Grammar::FromJSONSchema( + schema, + any_whitespace, + indent, + separators, + strict_mode, + max_whitespace_cnt, + /*print_converted_ebnf=*/false, + any_order + )); +} + +CompiledGrammar GrammarCompilerSub::CompileStructuralTag(const std::string& structural_tag_json) { + auto result = Grammar::FromStructuralTag(structural_tag_json, tokenizer_info_); + XGRAMMAR_CHECK(std::holds_alternative(result)) + << GetMessageFromVariantError(std::get<1>(result)); + return MultiThreadCompileGrammar(std::get<0>(result)); +} + +CompiledGrammar GrammarCompilerSub::CompileRegex(const std::string& regex) { + return MultiThreadCompileGrammar(Grammar::FromRegex(regex)); +} + +CompiledGrammar GrammarCompilerSub::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + return MultiThreadCompileGrammar(Grammar::FromLark(lark_string, tokenizer_info_, named_grammars)); +} + +CompiledGrammar GrammarCompilerSub::CompileGrammar(const Grammar& grammar) { + return MultiThreadCompileGrammar(grammar); +} + +CompiledGrammar GrammarCompilerSub::CompileGrammar( + const std::string& ebnf_str, std::string root_rule_name +) { + return MultiThreadCompileGrammar(Grammar::FromEBNF(ebnf_str, root_rule_name)); +} + +void GrammarCompilerSub::TagDispatchOptimization( + std::shared_ptr compiled_grammar_impl, + std::unordered_map* tag_dispatch_rule_id_to_second_slicing_bitset +) { + using GrammarExprType = Grammar::Impl::GrammarExprType; + tag_dispatch_rule_id_to_second_slicing_bitset->clear(); + + // Optimization for TagDispatch: Precompute the definitely accepted tokens. + for (int i = 0; i < compiled_grammar_impl->grammar->NumRules(); i++) { + const auto& rule = compiled_grammar_impl->grammar->GetRule(i); + const auto& rule_body = compiled_grammar_impl->grammar->GetGrammarExpr(rule.body_expr_id); + if (rule_body.type != GrammarExprType::kTagDispatch) { + continue; + } + XGRAMMAR_DCHECK(rule_body.type == GrammarExprType::kTagDispatch); + Grammar::Impl::TagDispatch tag_dispatch = + compiled_grammar_impl->GetGrammar()->GetTagDispatch(rule.body_expr_id); + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + DynamicBitset definite_accepted_tokens_since_second_char(sorted_decoded_vocab.size()); + for (int j = 0; j < static_cast(sorted_decoded_vocab.size()); j++) { + bool definite_accept_since_second_char = true; + const auto& token = sorted_decoded_vocab[j].second; + if (token.empty()) { + definite_accepted_tokens_since_second_char.Set(j); + continue; + } + + // Check if the token contains any string trigger or exclude string after first char. + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + if (token.find(trigger, 1) != std::string::npos) { + definite_accept_since_second_char = false; + break; + } + } + if (definite_accept_since_second_char) { + for (const auto& excl : tag_dispatch.excludes) { + if (token.find(excl, 1) != std::string::npos) { + definite_accept_since_second_char = false; + break; + } + } + } + + if (definite_accept_since_second_char) { + definite_accepted_tokens_since_second_char.Set(j); + } + } + (*tag_dispatch_rule_id_to_second_slicing_bitset)[i] = + definite_accepted_tokens_since_second_char; + } +} + +/******************* GrammarCompiler::Impl *******************/ + +/*! + * \brief The keys for the cache. This is defined here instead of inside the GrammarCompiler::Impl + * class due C++ template specialization and hash specialization rules. + */ +class GrammarCompilerCacheKeys { + public: + struct SchemaKey { + std::string schema; + bool any_whitespace; + std::optional indent; + std::optional> separators; + bool strict_mode; + std::optional max_whitespace_cnt; + bool any_order; + + XGRAMMAR_EQUAL_BY_MEMBERS( + SchemaKey, + &SchemaKey::schema, + &SchemaKey::any_whitespace, + &SchemaKey::indent, + &SchemaKey::separators, + &SchemaKey::strict_mode, + &SchemaKey::max_whitespace_cnt, + &SchemaKey::any_order + ); + }; + + struct StructuralTagKey { + std::string structural_tag_json; + + XGRAMMAR_EQUAL_BY_MEMBERS(StructuralTagKey, &StructuralTagKey::structural_tag_json); + }; + + struct GrammarKey { + std::string ebnf_str; + std::string root_rule_name; + + XGRAMMAR_EQUAL_BY_MEMBERS(GrammarKey, &GrammarKey::ebnf_str, &GrammarKey::root_rule_name); + }; + + struct RegexKey { + std::string regex; + + XGRAMMAR_EQUAL_BY_MEMBERS(RegexKey, &RegexKey::regex); + }; + + struct LarkNamedGrammarKey { + std::string name; + bool is_lark_source; + std::string source_or_ebnf; + std::string root_rule_name; + std::variant value; + + bool operator==(const LarkNamedGrammarKey& other) const { + return std::tie(name, is_lark_source, source_or_ebnf, root_rule_name) == + std::tie(other.name, other.is_lark_source, other.source_or_ebnf, other.root_rule_name); + } + }; + + struct LarkKey { + std::string lark_string; + std::vector named_grammars; + + XGRAMMAR_EQUAL_BY_MEMBERS(LarkKey, &LarkKey::lark_string, &LarkKey::named_grammars); + }; + + struct BuiltinJSONGrammarKey { + XGRAMMAR_EQUAL_BY_MEMBERS_EMPTY(BuiltinJSONGrammarKey); + }; + + using UnionKey = std:: + variant; +}; + +} // namespace xgrammar + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::SchemaKey, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::schema, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::any_whitespace, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::indent, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::separators, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::strict_mode, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::max_whitespace_cnt, + &xgrammar::GrammarCompilerCacheKeys::SchemaKey::any_order +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::StructuralTagKey, + &xgrammar::GrammarCompilerCacheKeys::StructuralTagKey::structural_tag_json +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::GrammarKey, + &xgrammar::GrammarCompilerCacheKeys::GrammarKey::ebnf_str, + &xgrammar::GrammarCompilerCacheKeys::GrammarKey::root_rule_name +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::RegexKey, + &xgrammar::GrammarCompilerCacheKeys::RegexKey::regex +); + +XGRAMMAR_HASH_BY_MEMBERS( + xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::name, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::is_lark_source, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::source_or_ebnf, + &xgrammar::GrammarCompilerCacheKeys::LarkNamedGrammarKey::root_rule_name +); + +namespace std { +template <> +struct hash { + std::size_t operator()(const xgrammar::GrammarCompilerCacheKeys::LarkKey& key) const noexcept { + uint64_t seed = std::hash{}(key.lark_string); + for (const auto& named_grammar : key.named_grammars) { + xgrammar::HashCombineBinary( + seed, std::hash{}(named_grammar) + ); + } + return seed; + } +}; +} // namespace std + +XGRAMMAR_HASH_BY_MEMBERS_EMPTY(xgrammar::GrammarCompilerCacheKeys::BuiltinJSONGrammarKey); + +namespace xgrammar { + +/*! + * \brief The implementation of the grammar compiler with cache. It calls the no cache compiler + * to compile the grammar, and implements the cache logic upon it. + */ +class GrammarCompiler::Impl { + public: + Impl( + const TokenizerInfo& tokenizer_info, + int max_threads, + bool cache_enabled, + int64_t max_memory_bytes + ) + : cache_enabled_(cache_enabled), + rule_level_cache_( + cache_enabled + ? std::optional( + max_memory_bytes == -1 + ? static_cast(-1) + : static_cast(max_memory_bytes - max_memory_bytes / 3 * 2) + ) + : std::nullopt + ), + no_cache_compiler_(tokenizer_info, max_threads, rule_level_cache_), + grammar_level_cache_( + max_memory_bytes == -1 ? static_cast(-1) + : static_cast(max_memory_bytes / 3 * 2), + Computer(*this) + ) { + if (max_memory_bytes < -1) { + XGRAMMAR_LOG(FATAL) << "Invalid max_memory_bytes: " << max_memory_bytes << ". " + << "It should be -1 (unlimited) or a non-negative integer."; + } + } + + CompiledGrammar CompileBuiltinJSONGrammar(); + + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order + ); + + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + CompiledGrammar CompileRegex(const std::string& regex); + + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars + ); + + CompiledGrammar CompileGrammar(const Grammar& grammar); + + CompiledGrammar CompileGrammar(const std::string& ebnf_str, std::string root_rule_name); + + void ClearCache(); + + int64_t GetCacheSizeBytes() const; + + int64_t CacheLimitBytes() const; + + private: + using SchemaKey = GrammarCompilerCacheKeys::SchemaKey; + using StructuralTagKey = GrammarCompilerCacheKeys::StructuralTagKey; + using GrammarKey = GrammarCompilerCacheKeys::GrammarKey; + using RegexKey = GrammarCompilerCacheKeys::RegexKey; + using LarkNamedGrammarKey = GrammarCompilerCacheKeys::LarkNamedGrammarKey; + using LarkKey = GrammarCompilerCacheKeys::LarkKey; + using BuiltinJSONGrammarKey = GrammarCompilerCacheKeys::BuiltinJSONGrammarKey; + using UnionKey = GrammarCompilerCacheKeys::UnionKey; + + CompiledGrammar Compute(const UnionKey& key); + + struct Computer { + Computer(Impl& compiler) : compiler(compiler) {} + // Forward the key to GrammarCompiler::Impl::Compute(key) + CompiledGrammar operator()(const UnionKey& key) const { return compiler.Compute(key); } + GrammarCompiler::Impl& compiler; + }; + + struct SizeEstimator { + std::size_t operator()(const CompiledGrammar& value) const { return value.MemorySizeBytes(); } + }; + + /*! \brief Whether the cache is enabled. */ + const bool cache_enabled_; + + /*! \brief The crossing cache manager for compiled grammars. */ + std::optional rule_level_cache_ = std::nullopt; + + /*! \brief The no cache compiler. */ + GrammarCompilerSub no_cache_compiler_; + + /*! \brief The cache for compiled grammars. */ + ThreadSafeLRUCache grammar_level_cache_; +}; + +CompiledGrammar GrammarCompiler::Impl::Compute(const UnionKey& key) { + return std::visit( + [this](const auto& key) -> CompiledGrammar { + using KeyType = std::decay_t; + if constexpr (std::is_same_v) { + const auto& [ebnf_str, root_rule_name] = key; + return this->no_cache_compiler_.CompileGrammar(ebnf_str, root_rule_name); + } else if constexpr (std::is_same_v) { + const auto& [schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order] = + key; + return this->no_cache_compiler_.CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); + } else if constexpr (std::is_same_v) { + const auto& [structural_tag_json] = key; + return this->no_cache_compiler_.CompileStructuralTag(structural_tag_json); + } else if constexpr (std::is_same_v) { + const auto& [regex] = key; + return this->no_cache_compiler_.CompileRegex(regex); + } else if constexpr (std::is_same_v) { + std::vector named_grammars; + named_grammars.reserve(key.named_grammars.size()); + for (const auto& named_grammar : key.named_grammars) { + named_grammars.push_back({named_grammar.name, named_grammar.value}); + } + return this->no_cache_compiler_.CompileLark(key.lark_string, named_grammars); + } else if constexpr (std::is_same_v) { + return this->no_cache_compiler_.CompileBuiltinJSONGrammar(); + } else { + XGRAMMAR_UNREACHABLE(); + } + }, + key + ); +} + +CompiledGrammar GrammarCompiler::Impl::CompileBuiltinJSONGrammar() { + if (!cache_enabled_) { + return no_cache_compiler_.CompileBuiltinJSONGrammar(); + } + return grammar_level_cache_.Get(BuiltinJSONGrammarKey{}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); + } + return grammar_level_cache_.Get(SchemaKey{ + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + }); +} + +CompiledGrammar GrammarCompiler::Impl::CompileStructuralTag(const std::string& structural_tag_json +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileStructuralTag(structural_tag_json); + } + return grammar_level_cache_.Get(StructuralTagKey{structural_tag_json}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileRegex(const std::string& regex) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileRegex(regex); + } + return grammar_level_cache_.Get(RegexKey{regex}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileLark(lark_string, named_grammars); + } + + std::vector named_grammar_keys; + named_grammar_keys.reserve(named_grammars.size()); + for (const auto& named_grammar : named_grammars) { + if (std::holds_alternative(named_grammar.grammar)) { + const auto& source = std::get(named_grammar.grammar); + named_grammar_keys.push_back( + {named_grammar.name, /*is_lark_source=*/true, source, "", named_grammar.grammar} + ); + } else { + const auto& grammar = std::get(named_grammar.grammar); + named_grammar_keys.push_back( + {named_grammar.name, + /*is_lark_source=*/false, + grammar.ToString(), + grammar->GetRootRule().name, + named_grammar.grammar} + ); + } + } + std::sort( + named_grammar_keys.begin(), + named_grammar_keys.end(), + [](const LarkNamedGrammarKey& lhs, const LarkNamedGrammarKey& rhs) { + return std::tie(lhs.name, lhs.is_lark_source, lhs.source_or_ebnf, lhs.root_rule_name) < + std::tie(rhs.name, rhs.is_lark_source, rhs.source_or_ebnf, rhs.root_rule_name); + } + ); + return grammar_level_cache_.Get(LarkKey{lark_string, std::move(named_grammar_keys)}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileGrammar(const Grammar& grammar) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileGrammar(grammar); + } + return grammar_level_cache_.Get(GrammarKey{grammar.ToString(), grammar->GetRootRule().name}); +} + +CompiledGrammar GrammarCompiler::Impl::CompileGrammar( + const std::string& ebnf_str, std::string root_rule_name +) { + if (!cache_enabled_) { + return no_cache_compiler_.CompileGrammar(ebnf_str, root_rule_name); + } + return grammar_level_cache_.Get(GrammarKey{ebnf_str, root_rule_name}); +} + +void GrammarCompiler::Impl::ClearCache() { + grammar_level_cache_.Clear(); + if (rule_level_cache_.has_value()) { + rule_level_cache_->ClearCache(); + } +} + +int64_t GrammarCompiler::Impl::GetCacheSizeBytes() const { + return static_cast(grammar_level_cache_.MemorySize()) + + static_cast(MemorySize(rule_level_cache_)); +} + +int64_t GrammarCompiler::Impl::CacheLimitBytes() const { + const auto size = grammar_level_cache_.MaxMemorySize(); + if (size == grammar_level_cache_.kUnlimitedSize) return -1; + return static_cast(size) + (rule_level_cache_.has_value() + ? static_cast(rule_level_cache_->GetMaxSize()) + : 0); +} + +/******************* GrammarCompiler *******************/ + +GrammarCompiler::GrammarCompiler( + const TokenizerInfo& tokenizer_info, + int max_threads, + bool cache_enabled, + int64_t max_memory_bytes +) + : pimpl_(std::make_shared(tokenizer_info, max_threads, cache_enabled, max_memory_bytes)) { +} + +CompiledGrammar GrammarCompiler::CompileJSONSchema( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order +) { + return pimpl_->CompileJSONSchema( + schema, any_whitespace, indent, separators, strict_mode, max_whitespace_cnt, any_order + ); +} + +CompiledGrammar GrammarCompiler::CompileBuiltinJSONGrammar() { + return pimpl_->CompileBuiltinJSONGrammar(); +} + +CompiledGrammar GrammarCompiler::CompileStructuralTag(const std::string& structural_tag_json) { + return pimpl_->CompileStructuralTag(structural_tag_json); +} + +CompiledGrammar GrammarCompiler::CompileRegex(const std::string& regex) { + return pimpl_->CompileRegex(regex); +} + +CompiledGrammar GrammarCompiler::CompileLark( + const std::string& lark_string, const std::vector& named_grammars +) { + return pimpl_->CompileLark(lark_string, named_grammars); +} + +CompiledGrammar GrammarCompiler::CompileGrammar(const Grammar& grammar) { + return pimpl_->CompileGrammar(grammar); +} + +CompiledGrammar GrammarCompiler::CompileGrammar( + const std::string& ebnf_str, const std::string& root_rule_name +) { + return pimpl_->CompileGrammar(ebnf_str, root_rule_name); +} + +void GrammarCompiler::ClearCache() { pimpl_->ClearCache(); } + +int64_t GrammarCompiler::GetCacheSizeBytes() const { return pimpl_->GetCacheSizeBytes(); } + +int64_t GrammarCompiler::CacheLimitBytes() const { return pimpl_->CacheLimitBytes(); } + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_functor.cc b/zig/vendor/xgrammar/cpp/grammar_functor.cc new file mode 100644 index 0000000000..d014bdc9a2 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_functor.cc @@ -0,0 +1,4016 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_functor.cc + */ + +#include "grammar_functor.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "fsm.h" +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "suffix_automata.h" +#include "support/container.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +using GrammarExpr = Grammar::Impl::GrammarExpr; +using ExprType = Grammar::Impl::GrammarExprType; + +/*************************** Impl of grammar constructors ***************************/ + +/*! + * \brief Base class for grammar mutators that add subgrammars. + * + * Provides functionality to visit a subgrammar and add its rules to the builder + * while maintaining proper rule references and names. + */ +class SubGrammarAdderImpl : public GrammarMutator { + public: + SubGrammarAdderImpl() = default; + + /*! + * \brief Visit a subgrammar and add the rules to the builder. + * \param grammar The subgrammar to visit. + * \return The new id of the root rule of this subgrammar. + */ + int32_t ApplyWithBuilder(GrammarBuilder* builder, const Grammar& sub_grammar) { + InitGrammar(sub_grammar); + InitBuilder(builder); + new_rule_ids_names.reserve(base_grammar_->NumRules()); + new_rule_ids_names.clear(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto new_name = builder_->GetNewRuleName(base_grammar_->GetRule(i).name); + auto new_id = builder_->AddEmptyRule(new_name); + new_rule_ids_names.emplace_back(new_id, new_name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = new_rule_ids_names[i].second; + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(new_rule_ids_names[i].first, new_body_expr_id); + auto new_lookahead_assertion_id = VisitLookaheadAssertion(rule.lookahead_assertion_id); + builder_->UpdateLookaheadAssertion(new_rule_ids_names[i].first, new_lookahead_assertion_id); + builder_->UpdateMaxTokens(new_rule_ids_names[i].first, rule.max_tokens); + builder_->UpdateMaxChars(new_rule_ids_names[i].first, rule.max_chars); + builder_->UpdateCaptureName(new_rule_ids_names[i].first, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + auto remapped_info = *suffix_stop_info; + if (remapped_info.body_rule_id >= 0) { + remapped_info.body_rule_id = new_rule_ids_names[remapped_info.body_rule_id].first; + remapped_info.marker_rule_id = new_rule_ids_names[remapped_info.marker_rule_id].first; + } + builder_->UpdateSuffixStopInfo(new_rule_ids_names[i].first, remapped_info); + } + builder_->UpdateLazy(new_rule_ids_names[i].first, rule.is_lazy); + builder_->UpdateRuleTemperature(new_rule_ids_names[i].first, rule.temperature); + } + return new_rule_ids_names[base_grammar_->GetRootRuleId()].first; + } + + int32_t VisitRuleRef(const GrammarExpr& grammar_expr) final { + return builder_->AddRuleRef(new_rule_ids_names[grammar_expr[0]].first); + } + + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + return builder_->AddRepeat( + new_rule_ids_names[grammar_expr[0]].first, grammar_expr[1], grammar_expr[2] + ); + } + + int32_t VisitTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TagDispatch old_tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + Grammar::Impl::TagDispatch new_tag_dispatch; + for (const auto& [trigger, rule_id] : old_tag_dispatch.tag_rule_pairs) { + new_tag_dispatch.tag_rule_pairs.emplace_back(trigger, new_rule_ids_names[rule_id].first); + } + new_tag_dispatch.loop_after_dispatch = old_tag_dispatch.loop_after_dispatch; + new_tag_dispatch.excludes = old_tag_dispatch.excludes; + return builder_->AddTagDispatch(new_tag_dispatch); + } + + int32_t VisitTokenTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TokenTagDispatch old_ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + Grammar::Impl::TokenTagDispatch new_ttd; + for (const auto& [token_id, rule_id] : old_ttd.trigger_rule_pairs) { + new_ttd.trigger_rule_pairs.emplace_back(token_id, new_rule_ids_names[rule_id].first); + } + new_ttd.loop_after_dispatch = old_ttd.loop_after_dispatch; + new_ttd.excludes = old_ttd.excludes; + return builder_->AddTokenTagDispatch(new_ttd); + } + + std::vector> new_rule_ids_names; +}; + +/*! + * \brief Implementation of grammar union operation. + * + * Creates a new grammar that accepts strings from any of the input grammars. + * The resulting grammar has a new root rule that chooses between the root rules + * of all input grammars. + */ +class GrammarUnionFunctorImpl : public GrammarMutator { + public: + GrammarUnionFunctorImpl() = default; + + Grammar Apply(const std::vector& grammars) { + InitGrammar(); + InitBuilder(); + auto root_rule_id = builder_->AddEmptyRule("root"); + + std::vector new_root_choices; + new_root_choices.reserve(grammars.size()); + + for (const auto& grammar : grammars) { + auto new_root_id_for_grammar = SubGrammarAdderImpl().ApplyWithBuilder(builder_, grammar); + auto new_rule_ref = builder_->AddRuleRef(new_root_id_for_grammar); + auto new_rule_ref_seq = builder_->AddSequence({new_rule_ref}); + new_root_choices.push_back(new_rule_ref_seq); + } + + builder_->UpdateRuleBody(root_rule_id, builder_->AddChoices(new_root_choices)); + return builder_->Get(root_rule_id); + } + + // Avoid hiding the original Apply(const Grammar&) + Grammar Apply(const Grammar& grammar) final { + XGRAMMAR_LOG(FATAL) << "Should not be called"; + XGRAMMAR_UNREACHABLE(); + } +}; + +/*! + * \brief Implementation of grammar concatenation operation. + * + * Creates a new grammar that accepts strings that are concatenations of strings + * from the input grammars in order. The resulting grammar has a new root rule + * that concatenates the root rules of all input grammars. + */ +class GrammarConcatFunctorImpl : public GrammarMutator { + public: + GrammarConcatFunctorImpl() = default; + + Grammar Apply(const std::vector& grammars) { + InitGrammar(); + InitBuilder(); + auto root_rule_id = builder_->AddEmptyRule("root"); + + std::vector new_root_sequence; + new_root_sequence.reserve(grammars.size()); + + for (const auto& grammar : grammars) { + auto new_root_id_for_grammar = SubGrammarAdderImpl().ApplyWithBuilder(builder_, grammar); + auto new_rule_ref = builder_->AddRuleRef(new_root_id_for_grammar); + new_root_sequence.push_back(new_rule_ref); + } + + auto new_root_seq = builder_->AddSequence(new_root_sequence); + builder_->UpdateRuleBody(root_rule_id, builder_->AddChoices({new_root_seq})); + + return builder_->Get(root_rule_id); + } + + // Avoid hiding the original Apply(const Grammar&) + Grammar Apply(const Grammar& grammar) final { + XGRAMMAR_LOG(FATAL) << "Should not be called"; + XGRAMMAR_UNREACHABLE(); + } +}; + +/*************************** Impl of grammar normalizers ***************************/ + +/*! + * \brief Eliminates single-element sequence or choice or character class in the grammar. + * \example `A ::= choices("a")` --> `A ::= "a"` (the body is a string) + * \example `A ::= sequence("a")` --> `A ::= "a"` (the body is a string) + * \example `A ::= [a-a]` --> `A ::= "a"` (the body is a string) + */ +class SingleElementExprEliminator : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + private: + int32_t VisitSequence(const GrammarExpr& grammar_expr) final { + std::vector sequence_ids; + for (int32_t i : grammar_expr) { + sequence_ids.push_back(VisitExpr(i)); + } + if (sequence_ids.size() == 1) { + return sequence_ids[0]; + } + return builder_->AddSequence(sequence_ids); + } + + int32_t VisitChoices(const GrammarExpr& grammar_expr) final { + std::vector choice_ids; + for (int32_t i : grammar_expr) { + choice_ids.push_back(VisitExpr(i)); + } + if (choice_ids.size() == 1) { + return choice_ids[0]; + } + return builder_->AddChoices(choice_ids); + } + + int32_t VisitCharacterClass(const GrammarExpr& grammar_expr) final { + if (grammar_expr.data_len == 3 && grammar_expr[0] == 0 && grammar_expr[1] == grammar_expr[2]) { + std::string str = CharToUTF8(grammar_expr[1]); + std::vector bytes; + bytes.reserve(str.size()); + for (char c : str) { + bytes.push_back(static_cast(c)); + } + return builder_->AddByteString(bytes); + } + return builder_->AddGrammarExpr(grammar_expr); + } +}; + +/*! + * \brief Take a grammar from SingleElementExprEliminator and normalize the structure of the + * grammar. + * + * \note The normalized form: + * Each rule should be either: + * - A sequence of choices, each choice is a sequence of elements. Elements can be a character + * class, a byte string, or a rule reference. Only the first choice can be an empty string, + * indicating the rule can be empty. E.g. + * `rule_name ::= ("" | (element1_1 element1_2 ...) | (element2_1 element2_2 ...) | ...)` + * - A macro. Now only TagDispatch is supported. + * + * The lookahead assertion should be a sequence. + * + * New rules may be created to make every rule fit the normalized form. + * + * \example `A ::= ((a) (((b)) (c)) "")` -> `A ::= ((a b c))` + * \example `A ::= (a | (b | (c | "")))` -> `A ::= ("" | (a) | (b) | (c))` + * \example `A ::= (a | (b (c | d)))` -> `A ::= ((a) | (b A_1)), A_1 ::= ((c) | (d))` + * \example `A ::= (a | TagDispatch((tag1, rule1)))` -> `A ::= ((a) | (A_1)), A_1 ::= + * TagDispatch((tag1, rule1))` + */ +class StructureNormalizerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + auto grammar_new = SingleElementExprEliminator().Apply(grammar); + InitGrammar(grammar_new); + InitBuilder(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + cur_rule_name_ = rule.name; + auto new_body_expr_id = VisitRuleBody(grammar_expr); + builder_->UpdateRuleBody(i, new_body_expr_id); + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } + + private: + int32_t VisitLookaheadAssertion(int32_t lookahead_assertion_id) final { + if (lookahead_assertion_id == -1) { + return -1; + } + auto assertion_expr = base_grammar_->GetGrammarExpr(lookahead_assertion_id); + switch (assertion_expr.type) { + case GrammarExprType::kSequence: + return builder_->AddSequence(VisitSequence_(assertion_expr)); + case GrammarExprType::kChoices: + XGRAMMAR_LOG(FATAL) << "Choices in lookahead assertion are not supported yet"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kEmptyStr: + XGRAMMAR_LOG(FATAL) << "Empty string should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kTagDispatch: + XGRAMMAR_LOG(FATAL) << "TagDispatch should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kRegex: + XGRAMMAR_LOG(FATAL) << "Regex should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kSubstring: + XGRAMMAR_LOG(FATAL) << "Substring should not be in lookahead assertion"; + XGRAMMAR_UNREACHABLE(); + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + case GrammarExprType::kTokenTagDispatch: + return builder_->AddSequence({builder_->AddGrammarExpr(assertion_expr)}); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected lookahead assertion type: " + << static_cast(assertion_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! \brief Visit a GrammarExpr as a rule body. */ + int32_t VisitRuleBody(const GrammarExpr& grammar_expr) { + switch (grammar_expr.type) { + case GrammarExprType::kSequence: + return builder_->AddChoices({builder_->AddSequence(VisitSequence_(grammar_expr))}); + case GrammarExprType::kChoices: + return builder_->AddChoices(VisitChoices_(grammar_expr)); + case GrammarExprType::kEmptyStr: + return builder_->AddChoices({builder_->AddEmptyStr()}); + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + return builder_->AddChoices({builder_->AddSequence({builder_->AddGrammarExpr(grammar_expr)}) + }); + case GrammarExprType::kTagDispatch: + return VisitTagDispatch(grammar_expr); + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(grammar_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + return builder_->AddChoices({builder_->AddSequence({builder_->AddRuleRef(new_rule_id)})}); + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: + // A regex or substring is kept as the direct body of the rule, like a tag dispatch. + return builder_->AddGrammarExpr(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! + * \brief Visit a GrammarExpr containing choices. + * \returns A list of new choice GrammarExpr ids. + */ + std::vector VisitChoices_(const GrammarExpr& grammar_expr) { + std::vector new_choice_ids; + bool found_empty = false; + for (auto i : grammar_expr) { + auto choice_expr = base_grammar_->GetGrammarExpr(i); + switch (choice_expr.type) { + case GrammarExprType::kSequence: + VisitSequenceInChoices(choice_expr, &new_choice_ids, &found_empty); + break; + case GrammarExprType::kChoices: + VisitChoicesInChoices(choice_expr, &new_choice_ids, &found_empty); + break; + case GrammarExprType::kEmptyStr: + found_empty = true; + break; + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + VisitElementInChoices(choice_expr, &new_choice_ids); + break; + case GrammarExprType::kTagDispatch: { + auto tag_dispatch_expr_id = VisitTagDispatch(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, tag_dispatch_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: { + auto leaf_expr_id = builder_->AddGrammarExpr(choice_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, leaf_expr_id); + auto new_sequence_id = builder_->AddSequence({builder_->AddRuleRef(new_rule_id)}); + new_choice_ids.push_back(new_sequence_id); + break; + } + default: + XGRAMMAR_LOG(FATAL) << "Unexpected choice type: " << static_cast(choice_expr.type); + } + } + if (found_empty) { + new_choice_ids.insert(new_choice_ids.begin(), builder_->AddEmptyStr()); + } + XGRAMMAR_ICHECK(new_choice_ids.size() >= 1); + return new_choice_ids; + } + + /*! \brief Visit a sequence GrammarExpr that is one of a list of choices. */ + void VisitSequenceInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids, bool* found_empty + ) { + auto sub_sequence_ids = VisitSequence_(grammar_expr); + if (sub_sequence_ids.size() == 0) { + *found_empty = true; + } else { + new_choice_ids->push_back(builder_->AddSequence(sub_sequence_ids)); + } + } + + /*! \brief Visit a choice GrammarExpr that is one of a list of choices. */ + void VisitChoicesInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids, bool* found_empty + ) { + auto sub_choice_ids = VisitChoices_(grammar_expr); + bool contains_empty = + builder_->GetGrammarExpr(sub_choice_ids[0]).type == GrammarExprType::kEmptyStr; + if (contains_empty) { + *found_empty = true; + new_choice_ids->insert( + new_choice_ids->end(), sub_choice_ids.begin() + 1, sub_choice_ids.end() + ); + } else { + new_choice_ids->insert(new_choice_ids->end(), sub_choice_ids.begin(), sub_choice_ids.end()); + } + } + + /*! \brief Visit an atom element GrammarExpr that is one of a list of choices. */ + void VisitElementInChoices( + const GrammarExpr& grammar_expr, std::vector* new_choice_ids + ) { + auto sub_expr_id = builder_->AddGrammarExpr(grammar_expr); + new_choice_ids->push_back(builder_->AddSequence({sub_expr_id})); + } + + /*! + * \brief Visit a GrammarExpr containing a sequence. + * \returns A list of new sequence GrammarExpr ids. + */ + std::vector VisitSequence_(const GrammarExpr& grammar_expr) { + std::vector new_sequence_ids; + for (auto i : grammar_expr) { + auto element_expr = base_grammar_->GetGrammarExpr(i); + switch (element_expr.type) { + case GrammarExprType::kSequence: + VisitSequenceInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kChoices: + VisitChoiceInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kEmptyStr: + break; + case GrammarExprType::kByteString: + case GrammarExprType::kCharacterClass: + case GrammarExprType::kCharacterClassStar: + case GrammarExprType::kRuleRef: + case GrammarExprType::kRepeat: + case GrammarExprType::kToken: + case GrammarExprType::kExcludeToken: + VisitElementInSequence(element_expr, &new_sequence_ids); + break; + case GrammarExprType::kTagDispatch: { + auto tag_dispatch_expr_id = VisitTagDispatch(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, tag_dispatch_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + case GrammarExprType::kTokenTagDispatch: { + auto ttd_expr_id = VisitTokenTagDispatch(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, ttd_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + case GrammarExprType::kRegex: + case GrammarExprType::kSubstring: { + auto leaf_expr_id = builder_->AddGrammarExpr(element_expr); + auto new_rule_id = builder_->AddRuleWithHint(cur_rule_name_, leaf_expr_id); + new_sequence_ids.push_back(builder_->AddRuleRef(new_rule_id)); + break; + } + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " + << static_cast(element_expr.type); + } + } + return new_sequence_ids; + } + + /*! \brief Visit a sequence GrammarExpr that is one element in another sequence. */ + void VisitSequenceInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + auto sub_sequence_ids = VisitSequence_(grammar_expr); + new_sequence_ids->insert( + new_sequence_ids->end(), sub_sequence_ids.begin(), sub_sequence_ids.end() + ); + } + + /*! \brief Visit a choice GrammarExpr that is one element in a sequence. */ + void VisitChoiceInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + auto sub_choice_ids = VisitChoices_(grammar_expr); + if (sub_choice_ids.size() == 1) { + auto choice_element_expr = builder_->GetGrammarExpr(sub_choice_ids[0]); + if (choice_element_expr.type != GrammarExprType::kEmptyStr) { + new_sequence_ids->insert( + new_sequence_ids->end(), choice_element_expr.begin(), choice_element_expr.end() + ); + } + } else { + auto new_choice_id = builder_->AddChoices(sub_choice_ids); + auto new_choice_rule_id = builder_->AddRuleWithHint(cur_rule_name_, new_choice_id); + new_sequence_ids->push_back(builder_->AddRuleRef(new_choice_rule_id)); + } + } + + /*! \brief Visit an atom element GrammarExpr that is in a sequence. */ + void VisitElementInSequence( + const GrammarExpr& grammar_expr, std::vector* new_sequence_ids + ) { + new_sequence_ids->push_back(builder_->AddGrammarExpr(grammar_expr)); + } +}; + +/*! + * \brief A class that normalizes a grammar by applying a series of transformations. + * + * The normalizer applies the following transformations in order: + * 1. SingleElementExprEliminator - Eliminates single element expressions + * 2. NestedRuleUnwrapper - Unwraps nested rules + */ +class GrammarNormalizerImpl { + public: + GrammarNormalizerImpl() = default; + + Grammar Apply(const Grammar& grammar) { + auto renamed_grammar = RootRuleRenamer::Apply(grammar); + return StructureNormalizerImpl().Apply(renamed_grammar); + } +}; + +/*************************** Impl of grammar optimizers ***************************/ + +/*! + * \brief Base for optimizer passes that rewrite a grammar in place. It binds a GrammarBuilder + * directly to the target grammar (no copy) and walks every rule body and lookahead assertion. + * \details An expr whose subtree is unchanged keeps its original id; a new expr is appended only + * where a rewrite actually happens, so a pass with nothing to do writes nothing. Changed ids + * propagate upward: a container is rebuilt when any child changed, and a rule body or lookahead is + * updated when its expr changed. Stale exprs left behind are removed later by DeadCodeEliminator. + */ +class InPlaceGrammarRewriter { + public: + using GrammarExpr = Grammar::Impl::GrammarExpr; + using GrammarExprType = Grammar::Impl::GrammarExprType; + + void Apply(Grammar* grammar) { + grammar_ = grammar; + builder_ = GrammarBuilder::FromMutableGrammar(grammar); + // Expr ids are dense, so the memo is a plain array over the original exprs. Appended exprs + // are never visited again, so they need no memo slots. + memo_.assign(builder_.NumGrammarExprs(), -1); + Prepare(); + int32_t num_rules = builder_.NumRules(); + for (int32_t rule_id = 0; rule_id < num_rules; ++rule_id) { + int32_t body_expr_id = builder_.GetRule(rule_id).body_expr_id; + int32_t new_body_expr_id = VisitExpr(body_expr_id); + if (new_body_expr_id != body_expr_id) { + builder_.UpdateRuleBody(rule_id, new_body_expr_id); + } + int32_t lookahead_id = builder_.GetRule(rule_id).lookahead_assertion_id; + if (lookahead_id != -1) { + int32_t new_lookahead_id = VisitExpr(lookahead_id); + if (new_lookahead_id != lookahead_id) { + builder_.UpdateLookaheadAssertion(rule_id, new_lookahead_id); + } + } + } + } + + virtual ~InPlaceGrammarRewriter() = default; + + protected: + /*! \brief Hook run after the builder is bound and before the walk. */ + virtual void Prepare() {} + + /*! \brief Visit an expr; return its rewritten id, or the original id when nothing changed. + * May append new exprs to the arena, invalidating outstanding GrammarExpr views. */ + int32_t VisitExpr(int32_t expr_id) { + XGRAMMAR_DCHECK(expr_id < static_cast(memo_.size())); + if (memo_[expr_id] != -1) { + return memo_[expr_id]; + } + int32_t result; + switch (builder_.GetGrammarExpr(expr_id).type) { + case GrammarExprType::kSequence: + result = VisitSequence(expr_id); + break; + case GrammarExprType::kChoices: + result = VisitChoices(expr_id); + break; + default: + result = expr_id; + } + memo_[expr_id] = result; + return result; + } + + /*! \brief Rebuild a sequence, recursing into each element. Overridden to add rewriting. + * No memory is allocated when no element changes. */ + virtual int32_t VisitSequence(int32_t expr_id) { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_element_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t element_id = expr[i]; + int32_t new_element_id = VisitExpr(element_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_element_id != element_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_element_ids.reserve(size); + new_element_ids.insert(new_element_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_element_ids.push_back(new_element_id); + } + } + if (!changed) { + return expr_id; + } + return builder_.AddSequence(new_element_ids); + } + + /*! \brief Rebuild a choices, recursing into each choice. Overridden to add rewriting. + * No memory is allocated when no choice changes. */ + virtual int32_t VisitChoices(int32_t expr_id) { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_choice_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t choice_id = expr[i]; + int32_t new_choice_id = VisitExpr(choice_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_choice_id != choice_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_choice_ids.push_back(new_choice_id); + } + } + if (!changed) { + return expr_id; + } + return builder_.AddChoices(new_choice_ids); + } + + GrammarBuilder builder_; + Grammar* grammar_; + // Maps an original expr id to its rewritten id (or itself when unchanged); -1 means unvisited. + std::vector memo_; +}; + +/*! + * \brief Inline rules that can be inlined. + * + * Now we only inline rule references that: + * 1. at the beginning of a sequence + * 2. The rule should be a sequence of choices, cannot be empty, cannot refer to other rules + * + * \details Rewrites the grammar in place: only choices that actually have an inlinable leading + * rule reference are rebuilt, the rest keep their original ids. Inlinability is judged on the + * original grammar so the result does not depend on the order rules are rewritten. Stale exprs are + * removed later by DeadCodeEliminator. + */ +class RuleInlinerImpl : public InPlaceGrammarRewriter { + protected: + void Prepare() override { + int32_t num_rules = builder_.NumRules(); + original_body_expr_ids_.reserve(num_rules); + for (int32_t rule_id = 0; rule_id < num_rules; ++rule_id) { + original_body_expr_ids_.push_back(builder_.GetRule(rule_id).body_expr_id); + } + can_rule_be_inlined_.assign(num_rules, -1); + } + + int32_t VisitChoices(int32_t expr_id) override { + auto expr = builder_.GetGrammarExpr(expr_id); + int32_t size = expr.size(); + std::vector new_choice_ids; + bool changed = false; + for (int32_t i = 0; i < size; ++i) { + int32_t choice_id = expr[i]; + auto choice_expr = builder_.GetGrammarExpr(choice_id); + int32_t inline_rule_id = -1; + if (choice_expr.type == GrammarExprType::kSequence && choice_expr.size() > 0) { + auto first_element = builder_.GetGrammarExpr(choice_expr[0]); + if (first_element.type == GrammarExprType::kRuleRef && CanRuleBeInlined(first_element[0])) { + inline_rule_id = first_element[0]; + } + } + if (inline_rule_id == -1) { + int32_t new_choice_id = VisitExpr(choice_id); + // The visit may have appended exprs to the arena and invalidated the view, so re-fetch. + expr = builder_.GetGrammarExpr(expr_id); + if (!changed && new_choice_id != choice_id) { + // Materialize the already scanned, unchanged prefix on the first change. + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + if (changed) { + new_choice_ids.push_back(new_choice_id); + } + continue; + } + if (!changed) { + changed = true; + new_choice_ids.reserve(size); + new_choice_ids.insert(new_choice_ids.end(), expr.begin(), expr.begin() + i); + } + // Do inlining: prepend each choice of the referenced rule to the rest of this sequence. + // Copy and visit the needed element ids before appending anything. + std::vector rest_element_ids(choice_expr.begin() + 1, choice_expr.end()); + for (int32_t& element_id : rest_element_ids) { + element_id = VisitExpr(element_id); + } + auto ref_body = builder_.GetGrammarExpr(original_body_expr_ids_[inline_rule_id]); + std::vector ref_choice_ids(ref_body.begin(), ref_body.end()); + for (int32_t ref_choice_id : ref_choice_ids) { + auto ref_choice_expr = builder_.GetGrammarExpr(ref_choice_id); + XGRAMMAR_ICHECK(ref_choice_expr.type == GrammarExprType::kSequence); + std::vector new_sequence(ref_choice_expr.begin(), ref_choice_expr.end()); + for (int32_t& element_id : new_sequence) { + element_id = VisitExpr(element_id); + } + new_sequence.insert(new_sequence.end(), rest_element_ids.begin(), rest_element_ids.end()); + new_choice_ids.push_back(builder_.AddSequence(new_sequence)); + } + // The appended sequences invalidated the view, so re-fetch it for the next iteration. + expr = builder_.GetGrammarExpr(expr_id); + } + if (!changed) { + return expr_id; + } + return builder_.AddChoices(new_choice_ids); + } + + private: + /*! \brief A rule can be inlined iff its body is a non-empty choices of sequences that contain no + * rule references. Judged on the original grammar via original_body_expr_ids_. */ + bool CanRuleBeInlined(int32_t rule_id) { + if (can_rule_be_inlined_[rule_id] != -1) { + return can_rule_be_inlined_[rule_id] != 0; + } + const auto& rule = (*grammar_)->GetRule(rule_id); + // Inlining a budgeted rule would erase the rule its length budget applies to. Inlining a + // capture-relevant rule would eliminate its completion events, so its capture or hidden span + // would never be recorded. Inlining a lazy rule would erase its committed-shortest semantics. + // Inlining a temperature rule would erase the rule its sampling temperature applies to. + if (rule.max_tokens >= 0 || rule.max_chars >= 0 || !rule.capture_name.empty() || + (*grammar_)->GetSuffixStopInfo(rule_id) != nullptr || rule.is_lazy || + rule.temperature.has_value()) { + can_rule_be_inlined_[rule_id] = false; + return false; + } + bool result = true; + auto body = builder_.GetGrammarExpr(original_body_expr_ids_[rule_id]); + if (body.type != GrammarExprType::kChoices || body.size() == 0) { + result = false; + } else { + for (int32_t choice_id : body) { + auto choice_expr = builder_.GetGrammarExpr(choice_id); + if (choice_expr.type == GrammarExprType::kEmptyStr) { + result = false; + break; + } + XGRAMMAR_ICHECK(choice_expr.type == GrammarExprType::kSequence); + bool has_rule_ref = false; + for (int32_t element_id : choice_expr) { + if (builder_.GetGrammarExpr(element_id).type == GrammarExprType::kRuleRef) { + has_rule_ref = true; + break; + } + } + if (has_rule_ref) { + result = false; + break; + } + } + } + can_rule_be_inlined_[rule_id] = result ? 1 : 0; + return result; + } + + // Rule body expr ids captured before any rewriting, for order-independent inlinability checks. + std::vector original_body_expr_ids_; + // Per-rule cache of CanRuleBeInlined: -1 unknown, 0 false, 1 true. + std::vector can_rule_be_inlined_; +}; + +/*! + * \brief Analyze all referenced rules or the main rule. Return a list of all referenced rule ids. + * This is useful for dead code elimination. + */ +class UsedRulesAnalyzer : public GrammarVisitor> { + public: + UsedRulesAnalyzer() = default; + + std::vector Apply(const Grammar& grammar) final { + InitGrammar(grammar); + + std::set visited; + + std::queue().swap(visit_queue_); + + visit_queue_.push(base_grammar_->GetRootRuleId()); + while (!visit_queue_.empty()) { + auto rule_id = visit_queue_.front(); + visit_queue_.pop(); + if (visited.count(rule_id)) { + continue; + } + visited.insert(rule_id); + auto rule = base_grammar_->GetRule(rule_id); + VisitExpr(rule.body_expr_id); + if (rule.lookahead_assertion_id != -1) { + VisitExpr(rule.lookahead_assertion_id); + } + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(rule_id); + suffix_stop_info != nullptr && suffix_stop_info->body_rule_id != -1) { + visit_queue_.push(suffix_stop_info->body_rule_id); + visit_queue_.push(suffix_stop_info->marker_rule_id); + } + } + + return std::vector(visited.begin(), visited.end()); + } + + void VisitTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + visit_queue_.push(rule_id); + } + } + + void VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + visit_queue_.push(rule_id); + } + } + + void VisitRuleRef(const GrammarExpr& grammar_expr) { visit_queue_.push(grammar_expr[0]); } + + void VisitRepeat(const GrammarExpr& grammar_expr) { visit_queue_.push(grammar_expr[0]); } + + private: + std::queue visit_queue_; +}; + +class DeadCodeEliminatorImpl : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + InitGrammar(grammar); + InitBuilder(); + auto used_rules = UsedRulesAnalyzer().Apply(grammar); + rule_id_map_.clear(); + for (auto rule_id : used_rules) { + rule_id_map_[rule_id] = builder_->AddEmptyRule(grammar->GetRule(rule_id).name); + } + for (auto rule_id : used_rules) { + auto rule = grammar->GetRule(rule_id); + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(rule_id_map_[rule_id], new_body_expr_id); + builder_->UpdateLookaheadAssertion( + rule_id_map_[rule_id], VisitLookaheadAssertion(rule.lookahead_assertion_id) + ); + builder_->UpdateMaxTokens(rule_id_map_[rule_id], rule.max_tokens); + builder_->UpdateMaxChars(rule_id_map_[rule_id], rule.max_chars); + builder_->UpdateCaptureName(rule_id_map_[rule_id], rule.capture_name); + if (const auto* suffix_stop_info = grammar->GetSuffixStopInfo(rule_id)) { + auto remapped_info = *suffix_stop_info; + if (remapped_info.body_rule_id >= 0) { + remapped_info.body_rule_id = rule_id_map_.at(remapped_info.body_rule_id); + remapped_info.marker_rule_id = rule_id_map_.at(remapped_info.marker_rule_id); + } + builder_->UpdateSuffixStopInfo(rule_id_map_[rule_id], remapped_info); + } + builder_->UpdateLazy(rule_id_map_[rule_id], rule.is_lazy); + builder_->UpdateRuleTemperature(rule_id_map_[rule_id], rule.temperature); + } + XGRAMMAR_CHECK(rule_id_map_.count(grammar->GetRootRuleId()) > 0); + return builder_->Get(rule_id_map_[grammar->GetRootRuleId()]); + } + + int32_t VisitTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TagDispatch tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + XGRAMMAR_DCHECK(rule_id_map_.count(rule_id) > 0); + rule_id = rule_id_map_[rule_id]; + } + return builder_->AddTagDispatch(tag_dispatch); + } + + int32_t VisitTokenTagDispatch(const GrammarExpr& grammar_expr) final { + Grammar::Impl::TokenTagDispatch ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + XGRAMMAR_DCHECK(rule_id_map_.count(rule_id) > 0); + rule_id = rule_id_map_[rule_id]; + } + return builder_->AddTokenTagDispatch(ttd); + } + + int32_t VisitRuleRef(const GrammarExpr& grammar_expr) final { + XGRAMMAR_DCHECK(rule_id_map_.count(grammar_expr[0]) > 0); + auto new_rule_id = rule_id_map_[grammar_expr[0]]; + return builder_->AddRuleRef(new_rule_id); + } + + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + XGRAMMAR_DCHECK(rule_id_map_.count(grammar_expr[0]) > 0); + auto new_rule_id = rule_id_map_[grammar_expr[0]]; + return builder_->AddRepeat(new_rule_id, grammar_expr[1], grammar_expr[2]); + } + + private: + std::unordered_map rule_id_map_; +}; + +class LookaheadAssertionAnalyzerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + InitGrammar(grammar); + InitBuilder(grammar); + auto root_rule = grammar->GetRootRule(); + auto root_grammar_expr = base_grammar_->GetGrammarExpr(root_rule.body_expr_id); + if (root_grammar_expr.type == GrammarExprType::kTagDispatch || + root_grammar_expr.type == GrammarExprType::kTokenTagDispatch || + root_grammar_expr.type == GrammarExprType::kRegex || + root_grammar_expr.type == GrammarExprType::kSubstring) { + return grammar; + } + BuildRuleLookaheadInfo(); + for (int i = 0; i < static_cast(grammar->NumRules()); ++i) { + auto rule = grammar->GetRule(i); + if (i == grammar->GetRootRuleId()) { + continue; + } + if (rule.lookahead_assertion_id != -1) { + builder_->UpdateLookaheadExact(i, IsExactLookaheadAssertion(i)); + continue; + } + auto look_head_assertion_id = DetectLookaheadAssertion(i); + if (look_head_assertion_id != -1) { + builder_->UpdateLookaheadAssertion(i, look_head_assertion_id); + builder_->UpdateLookaheadExact(i); + } + } + return builder_->Get(grammar->GetRootRuleId()); + } + + bool IsExactLookaheadAssertion(int32_t rule_id) { + XGRAMMAR_DCHECK(base_grammar_->GetRule(rule_id).lookahead_assertion_id != -1); + return CanUseDerivedLookahead(rule_id); + } + + int32_t DetectLookaheadAssertion(int32_t rule_id) { + if (!CanUseDerivedLookahead(rule_id)) { + return -1; + } + return builder_->AddSequence(rule_lookahead_infos_[rule_id].suffix_after_first_occurrence); + } + + private: + struct RuleLookaheadInfo { + bool is_triggered_by_dispatch = false; + bool appears_as_last_in_other_rule = false; + int non_last_occurrence_count = 0; + std::vector suffix_after_first_occurrence; + }; + + bool CanUseDerivedLookahead(int32_t rule_id) const { + const auto& info = rule_lookahead_infos_[rule_id]; + return !info.is_triggered_by_dispatch && !info.appears_as_last_in_other_rule && + info.non_last_occurrence_count == 1; + } + + void BuildRuleLookaheadInfo() { + rule_lookahead_infos_.assign(base_grammar_->NumRules(), RuleLookaheadInfo{}); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (grammar_expr.type == GrammarExprType::kTagDispatch) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + rule_lookahead_infos_[rule_id].is_triggered_by_dispatch = true; + } + continue; + } + if (grammar_expr.type == GrammarExprType::kTokenTagDispatch) { + auto token_tag_dispatch = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : token_tag_dispatch.trigger_rule_pairs) { + rule_lookahead_infos_[rule_id].is_triggered_by_dispatch = true; + } + continue; + } + if (grammar_expr.type == GrammarExprType::kRegex || + grammar_expr.type == GrammarExprType::kSubstring) { + // A regex or substring rule is a leaf: it references no other rules. + continue; + } + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kChoices); + for (auto sequence_id : grammar_expr) { + auto sequence_expr = base_grammar_->GetGrammarExpr(sequence_id); + if (sequence_expr.type != GrammarExprType::kSequence || sequence_expr.size() == 0) { + continue; + } + auto last_element = base_grammar_->GetGrammarExpr(sequence_expr.end()[-1]); + if (last_element.type == GrammarExprType::kRuleRef && i != last_element[0]) { + rule_lookahead_infos_[last_element[0]].appears_as_last_in_other_rule = true; + } + for (int j = 0; j < sequence_expr.size() - 1; ++j) { + auto element_expr = base_grammar_->GetGrammarExpr(sequence_expr[j]); + if (element_expr.type != GrammarExprType::kRuleRef) { + continue; + } + auto& info = rule_lookahead_infos_[element_expr[0]]; + if (info.non_last_occurrence_count == 0) { + info.suffix_after_first_occurrence.assign( + sequence_expr.begin() + j + 1, sequence_expr.end() + ); + } + ++info.non_last_occurrence_count; + } + } + } + } + + std::vector rule_lookahead_infos_; +}; + +/*! + * \brief Finds the rule reference graph of a grammar. + * + * The rule reference graph shows which rules reference which other rules. + * The returned graph is inverted: it points from referee to referer. + */ +class RuleRefGraphFinder : public GrammarVisitor>> { + public: + RuleRefGraphFinder() = default; + + std::vector> Apply(const Grammar& grammar) { + InitGrammar(grammar); + rule_visit_graph_ = std::vector>(base_grammar_->NumRules()); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + cur_rule_id_ = i; + VisitExpr(grammar_expr); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + std::sort(rule_visit_graph_[i].begin(), rule_visit_graph_[i].end()); + auto end_it = std::unique(rule_visit_graph_[i].begin(), rule_visit_graph_[i].end()); + rule_visit_graph_[i].erase(end_it, rule_visit_graph_[i].end()); + } + return std::move(rule_visit_graph_); + } + + private: + void VisitRuleRef(const GrammarExpr& grammar_expr) { + rule_visit_graph_[grammar_expr[0]].push_back(cur_rule_id_); + } + + void VisitRepeat(const GrammarExpr& grammar_expr) { + rule_visit_graph_[grammar_expr[0]].push_back(cur_rule_id_); + } + + void VisitTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + rule_visit_graph_[rule_id].push_back(cur_rule_id_); + } + } + + void VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = base_grammar_->GetTokenTagDispatch(grammar_expr); + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + rule_visit_graph_[rule_id].push_back(cur_rule_id_); + } + } + + // Inversed reference graph: pointing from referee to referer + std::vector> rule_visit_graph_; + int32_t cur_rule_id_; +}; + +/*! + * \brief Analyzes which rules in a grammar can match the empty string. + */ +class AllowEmptyRuleAnalyzerImpl : public GrammarVisitor> { + public: + AllowEmptyRuleAnalyzerImpl() = default; + + std::vector Apply(const Grammar& grammar) final { + InitGrammar(grammar); + + // Step 1: Find rules that explicitly allow empty string + std::unordered_set empty_rule_id_set; + FindExplicitEmptyRules(&empty_rule_id_set); + + // Step 2: Find rules that indirectly allow empty string. Using the Bellman-Ford algorithm + // on the rule reference graph. + std::vector> rule_ref_graph = RuleRefGraphFinder().Apply(grammar); + FindIndirectEmptyRules(&empty_rule_id_set, rule_ref_graph); + + auto result = std::vector(empty_rule_id_set.begin(), empty_rule_id_set.end()); + std::sort(result.begin(), result.end()); + return result; + } + + void FindExplicitEmptyRules(std::unordered_set* empty_rule_id_set) { + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (grammar_expr.type == GrammarExprType::kTagDispatch || + grammar_expr.type == GrammarExprType::kTokenTagDispatch) { + empty_rule_id_set->insert(i); + continue; + } + + if (grammar_expr.type == GrammarExprType::kSubstring) { + // A substring automaton always accepts the empty string: every state is accepting. + empty_rule_id_set->insert(i); + continue; + } + + if (grammar_expr.type == GrammarExprType::kRegex) { + // Nullability is checked syntactically on the parsed regex, so no FSM is built here. + // Parse errors are reported by GrammarFSMBuilder later. + auto matches_empty_result = + RegexFSMBuilder::MatchesEmpty(base_grammar_->GetRegexString(grammar_expr)); + if (matches_empty_result.IsOk() && std::move(matches_empty_result).Unwrap()) { + empty_rule_id_set->insert(i); + } + continue; + } + + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kChoices); + if (base_grammar_->GetGrammarExpr(grammar_expr[0]).type == GrammarExprType::kEmptyStr) { + empty_rule_id_set->insert(i); + continue; + } + + for (auto seq_id : grammar_expr) { + auto seq_expr = base_grammar_->GetGrammarExpr(seq_id); + if (std::all_of(seq_expr.begin(), seq_expr.end(), [&](int32_t i) { + return base_grammar_->GetGrammarExpr(i).type == GrammarExprType::kCharacterClassStar; + })) { + empty_rule_id_set->insert(i); + break; + } + } + } + } + + bool SeqExprIsEpsilon( + const GrammarExpr& seq_expr, const std::unordered_set& empty_rule_id_set + ) { + if (seq_expr.type == GrammarExprType::kEmptyStr) { + return true; + } + XGRAMMAR_DCHECK(seq_expr.type == GrammarExprType::kSequence); + + return std::all_of(seq_expr.begin(), seq_expr.end(), [&](int32_t i) { + auto element_expr = base_grammar_->GetGrammarExpr(i); + return (element_expr.type == GrammarExprType::kRuleRef && + empty_rule_id_set.count(element_expr[0])) || + element_expr.type == GrammarExprType::kCharacterClassStar || + (element_expr.type == GrammarExprType::kRepeat && + (empty_rule_id_set.count(element_expr[0]) || element_expr[1] == 0)); + }); + } + + void FindIndirectEmptyRules( + std::unordered_set* empty_rule_id_set, + const std::vector>& rule_ref_graph + ) { + std::queue queue; + for (auto i : *empty_rule_id_set) { + queue.push(i); + } + + while (!queue.empty()) { + auto rule_id = queue.front(); + queue.pop(); + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rule_ref_graph.size())); + for (auto referer_rule_id : rule_ref_graph[rule_id]) { + if (empty_rule_id_set->count(referer_rule_id)) { + continue; + } + auto rule = base_grammar_->GetRule(referer_rule_id); + auto grammar_expr = base_grammar_->GetGrammarExpr(rule.body_expr_id); + + XGRAMMAR_DCHECK( + grammar_expr.type != GrammarExprType::kTagDispatch && + grammar_expr.type != GrammarExprType::kTokenTagDispatch + ) << "TagDispatch rules should already exist in empty_rule_id_set"; + + bool is_epsilon = std::any_of(grammar_expr.begin(), grammar_expr.end(), [&](int32_t i) { + auto seq_expr = base_grammar_->GetGrammarExpr(i); + return SeqExprIsEpsilon(seq_expr, *empty_rule_id_set); + }); + + if (is_epsilon) { + empty_rule_id_set->insert(referer_rule_id); + queue.push(referer_rule_id); + } + } + } + } +}; + +class GrammarFSMBuilderImpl { + public: + explicit GrammarFSMBuilderImpl( + FSM& target_fsm, + const std::string* rule_name = nullptr, + GrammarBuilder* grammar_builder = nullptr + ) + : target_fsm_(target_fsm), rule_name_(rule_name), grammar_builder_(grammar_builder) {} + + static void Apply(Grammar* grammar) { + FSM complete_fsm; + std::vector> per_rule_fsms; + std::vector state_mapping; + + // Compiling a kRegex rule may append new rules through this builder: large bounded + // repetitions in a regex become kRepeatRef edges referencing new rules. NumRules() is + // re-read every iteration so the new rules get their FSMs built as well. + int32_t num_original_rules = (*grammar)->NumRules(); + GrammarBuilder grammar_builder = GrammarBuilder::FromMutableGrammar(grammar); + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + auto rule_fsm = BuildRuleFSM(*grammar, i, &grammar_builder); + per_rule_fsms.push_back(rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping)); + } + + // The rules created during FSM building missed the AllowEmptyRuleAnalyzer pass; complete + // allow_empty_rule_ids for them. New rule ids are larger than all existing ids, so + // appending keeps the list sorted. + for (int i = num_original_rules; i < (*grammar)->NumRules(); ++i) { + const auto& rule = (*grammar)->GetRule(i); + const auto& body_expr = (*grammar)->GetGrammarExpr(rule.body_expr_id); + XGRAMMAR_DCHECK(body_expr.type == Grammar::Impl::GrammarExprType::kRegex); + auto matches_empty_result = + RegexFSMBuilder::MatchesEmpty((*grammar)->GetRegexString(body_expr)); + if (matches_empty_result.IsOk() && std::move(matches_empty_result).Unwrap()) { + (*grammar)->allow_empty_rule_ids.push_back(i); + } + } + + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + XGRAMMAR_DCHECK(per_rule_fsms[i].has_value()) + << "Rule " << i << " (" << (*grammar)->GetRule(i).name + << ") does not have an FSM after optimization"; + } + + // Compress to compact fsm + CompactFSM compact_complete_fsm = complete_fsm.ToCompact(); + std::vector> compact_per_rule_fsms( + (*grammar)->NumRules() + ); + for (int i = 0; i < (*grammar)->NumRules(); ++i) { + if (per_rule_fsms[i]) { + auto compact_fsm_with_se = CompactFSMWithStartEnd( + compact_complete_fsm, + per_rule_fsms[i]->GetFsm().GetStart(), + per_rule_fsms[i]->GetFsm().GetEnds() + ); + compact_per_rule_fsms[i] = CompactFSMWithStartEndWithSize( + compact_fsm_with_se, per_rule_fsms[i]->GetEdgeNum(), per_rule_fsms[i]->GetNodeNum() + ); + } + } + + (*grammar)->complete_fsm = std::move(compact_complete_fsm); + (*grammar)->per_rule_fsms = std::move(compact_per_rule_fsms); + } + + /* Basic Building functions.*/ + static FSMWithStartEnd RuleRef(const GrammarExpr& expr); + static FSMWithStartEnd CharacterClass(const GrammarExpr& expr); + static FSMWithStartEnd ByteString(const GrammarExpr& expr); + static FSMWithStartEnd Token(const GrammarExpr& expr); + static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); + static std::optional TokenTagDispatch(const Grammar::Impl::TokenTagDispatch& ttd + ); + static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); + static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); + static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + static Result Regex(const std::string& regex, bool json_string = false); + /* Building tool functions.*/ + static std::optional BuildTagDispatchFSM( + const std::vector>& string_trigger_rules, + bool loop_after_dispatch, + const std::vector& excluded_strings + ); + + private: + static FSMWithStartEnd BuildRuleFSM( + const Grammar& grammar, int rule_id, GrammarBuilder* grammar_builder = nullptr + ); + static FSMWithStartEnd BuildExpressionFSM( + const GrammarExpr& expr, + const Grammar& grammar, + const std::string* rule_name = nullptr, + GrammarBuilder* grammar_builder = nullptr + ); + void BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void BuildEmptyString(int start_state, std::vector* end_states); + void BuildByteString(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildRuleRef(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRepeat(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildToken(const GrammarExpr& expr, int start_state, std::vector* end_states); + void BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states + ); + void BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states + ); + void BuildSubstring( + const std::vector& chunks, int start_state, std::vector* end_states + ); + void BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states + ); + void BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states + ); + void AddCharacterClassTransitions(const GrammarExpr& expr, int start_state, int end_state); + void BuildNegativeCharacterClass(const GrammarExpr& expr, int start_state, int end_state); + void AppendFSM(FSMWithStartEnd fsm, int start_state, std::vector* end_states); + void AddCharacterRange(int from, int to, uint32_t min, uint32_t max); + void AddCodepointRange(int from, int to, uint32_t low, uint32_t high); + + FSM& target_fsm_; + const std::string* rule_name_; + // If not null, kRegex expressions may add new rules through this builder; see Apply(). + GrammarBuilder* grammar_builder_ = nullptr; +}; + +// This function will add a range [min, max] of unicode characters to the FSM. +void GrammarFSMBuilderImpl::AddCharacterRange(int from, int to, uint32_t min, uint32_t max) { + AddPackedUTF8RangeEdges(target_fsm_, from, to, min, max); +} + +void GrammarFSMBuilderImpl::AddCodepointRange(int from, int to, uint32_t low, uint32_t high) { + // Do not bridge UTF-8 widths with the packed-byte range helper: its historical minimum + // byte sequences include overlong encodings (e.g. C0 80), which are not codepoints. + constexpr uint32_t width_ends[] = {0x7F, 0x7FF, 0xFFFF, 0x10FFFF}; + for (uint32_t width_end : width_ends) { + if (low <= high && low <= width_end) { + auto end = std::min(high, width_end); + AddCharacterRange(from, to, CodepointToPackedUTF8(low), CodepointToPackedUTF8(end)); + low = end + 1; + } + } +} + +void GrammarFSMBuilderImpl::BuildNegativeCharacterClass( + const GrammarExpr& expr, int start_state, int end_state +) { + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); + XGRAMMAR_DCHECK(expr[0]); // Negative character class should be true. + // Complement codepoint ranges before encoding them. Truncating their endpoints to bytes + // incorrectly made, for example, [^\x00-\U0010ffff] accept every multi-byte character. + constexpr uint32_t kMaxCodepoint = 0x10FFFF; + std::vector> excluded_ranges; + for (int i = 1; i < static_cast(expr.size()); i += 2) { + excluded_ranges.emplace_back(expr[i], expr[i + 1]); + } + std::sort(excluded_ranges.begin(), excluded_ranges.end()); + uint32_t next = 0; + for (const auto& [low, high] : excluded_ranges) { + if (low > next) { + AddCodepointRange(start_state, end_state, next, std::min(low - 1, kMaxCodepoint)); + } + if (high >= kMaxCodepoint) { + return; + } + next = std::max(next, high + 1); + } + if (next <= kMaxCodepoint) { + AddCodepointRange(start_state, end_state, next, kMaxCodepoint); + } +} + +void GrammarFSMBuilderImpl::AddCharacterClassTransitions( + const GrammarExpr& expr, int start_state, int end_state +) { + XGRAMMAR_DCHECK( + expr.type == ExprType::kCharacterClass || expr.type == ExprType::kCharacterClassStar + ); + bool is_negative = expr[0]; + if (is_negative) { + BuildNegativeCharacterClass(expr, start_state, end_state); + } else { + for (int i = 1; i < static_cast(expr.size()); i += 2) { + uint32_t codepoint_min = static_cast(expr[i]); + uint32_t codepoint_max = static_cast(expr[i + 1]); + AddCodepointRange(start_state, end_state, codepoint_min, codepoint_max); + } + } +} + +void GrammarFSMBuilderImpl::BuildCharacterClass( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClass); + end_states->clear(); + int end_state = target_fsm_.AddState(); + AddCharacterClassTransitions(expr, start_state, end_state); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildCharacterClassStar( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kCharacterClassStar); + end_states->clear(); + AddCharacterClassTransitions(expr, start_state, start_state); + end_states->push_back(start_state); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::BuildRuleFSM( + const Grammar& grammar, int rule_id, GrammarBuilder* grammar_builder +) { + const auto& rule = grammar->GetRule(rule_id); + return BuildExpressionFSM( + grammar->GetGrammarExpr(rule.body_expr_id), grammar, &rule.name, grammar_builder + ); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::BuildExpressionFSM( + const GrammarExpr& expr, + const Grammar& grammar, + const std::string* rule_name, + GrammarBuilder* grammar_builder +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm, rule_name, grammar_builder); + builder.BuildExpression(expr, grammar, start_state, &end_states); + FSMWithStartEnd result(result_fsm, start_state, std::move(end_states)); + if (expr.type != ExprType::kTagDispatch && expr.type != ExprType::kTokenTagDispatch) { + result = result.SimplifyEpsilon(); + result = result.MergeEquivalentStates(); + } + return result; +} + +FSMWithStartEnd GrammarFSMBuilderImpl::RuleRef(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildRuleRef(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::CharacterClass(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + if (expr.type == ExprType::kCharacterClassStar) { + builder.BuildCharacterClassStar(expr, start_state, &end_states); + } else { + builder.BuildCharacterClass(expr, start_state, &end_states); + } + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::ByteString(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildByteString(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::Token(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +FSMWithStartEnd GrammarFSMBuilderImpl::ExcludeToken(const GrammarExpr& expr) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildExcludeToken(expr, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +std::optional GrammarFSMBuilderImpl::TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildTokenTagDispatch(token_tag_dispatch, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +void GrammarFSMBuilderImpl::BuildExpression( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + switch (expr.type) { + case ExprType::kEmptyStr: + return BuildEmptyString(start_state, end_states); + case ExprType::kByteString: + return BuildByteString(expr, start_state, end_states); + case ExprType::kCharacterClass: + return BuildCharacterClass(expr, start_state, end_states); + case ExprType::kCharacterClassStar: + return BuildCharacterClassStar(expr, start_state, end_states); + case ExprType::kRuleRef: + return BuildRuleRef(expr, start_state, end_states); + case ExprType::kRepeat: + return BuildRepeat(expr, start_state, end_states); + case ExprType::kToken: + return BuildToken(expr, start_state, end_states); + case ExprType::kExcludeToken: + return BuildExcludeToken(expr, start_state, end_states); + case ExprType::kSequence: + return BuildSequence(expr, grammar, start_state, end_states); + case ExprType::kChoices: + return BuildChoices(expr, grammar, start_state, end_states); + case ExprType::kRegex: + return BuildRegex( + grammar->GetRegexString(expr), + grammar->GetRegexIsJSONString(expr), + start_state, + end_states + ); + case ExprType::kSubstring: + return BuildSubstring(grammar->GetSubstringChunks(expr), start_state, end_states); + case ExprType::kTagDispatch: + return BuildTagDispatch(grammar->GetTagDispatch(expr), start_state, end_states); + case ExprType::kTokenTagDispatch: + return BuildTokenTagDispatch(grammar->GetTokenTagDispatch(expr), start_state, end_states); + } + XGRAMMAR_UNREACHABLE(); +} + +void GrammarFSMBuilderImpl::BuildEmptyString(int start_state, std::vector* end_states) { + end_states->clear(); + end_states->push_back(start_state); +} + +void GrammarFSMBuilderImpl::BuildByteString( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kByteString); + end_states->clear(); + int current_state = start_state; + for (int32_t byte : expr) { + int next_state = target_fsm_.AddState(); + target_fsm_.AddEdge( + current_state, next_state, static_cast(byte), static_cast(byte) + ); + current_state = next_state; + } + end_states->push_back(current_state); +} + +void GrammarFSMBuilderImpl::BuildRuleRef( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRuleRef); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRuleEdge(start_state, end_state, expr[0]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildRepeat( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kRepeat); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddRepeatEdge(start_state, end_state, expr[0], expr[1], expr[2]); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddTokenEdge(start_state, end_state, std::vector(expr.begin(), expr.end())); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildExcludeToken( + const GrammarExpr& expr, int start_state, std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kExcludeToken); + end_states->clear(); + int end_state = target_fsm_.AddState(); + target_fsm_.AddExcludeTokenEdge( + start_state, end_state, std::vector(expr.begin(), expr.end()) + ); + end_states->push_back(end_state); +} + +void GrammarFSMBuilderImpl::BuildSequence( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + if (expr.size() == 0) { + end_states->push_back(start_state); + return; + } + + BuildExpression(grammar->GetGrammarExpr(expr[0]), grammar, start_state, end_states); + + std::vector next_end_states; + for (int index = 1; index < static_cast(expr.size()); ++index) { + int element_start = target_fsm_.AddState(); + BuildExpression(grammar->GetGrammarExpr(expr[index]), grammar, element_start, &next_end_states); + for (int32_t previous_end_state : *end_states) { + target_fsm_.AddEpsilonEdge(previous_end_state, element_start); + } + end_states->swap(next_end_states); + } +} + +std::optional GrammarFSMBuilderImpl::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { + FSM result_fsm; + int start_state = result_fsm.AddState(); + std::vector end_states; + GrammarFSMBuilderImpl builder(result_fsm); + builder.BuildSequence(expr, grammar, start_state, &end_states); + return FSMWithStartEnd(result_fsm, start_state, std::move(end_states)); +} + +void GrammarFSMBuilderImpl::BuildChoices( + const GrammarExpr& expr, + const Grammar& grammar, + int start_state, + std::vector* end_states +) { + XGRAMMAR_DCHECK(expr.type == ExprType::kChoices); + XGRAMMAR_DCHECK(start_state >= 0 && start_state < target_fsm_.NumStates()); + end_states->clear(); + + int non_empty_choice_count = 0; + bool nullable = false; + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type == ExprType::kEmptyStr) { + nullable = true; + } else { + ++non_empty_choice_count; + } + } + + if (non_empty_choice_count == 0) { + end_states->push_back(start_state); + return; + } + + if (non_empty_choice_count == 1 && !nullable) { + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type != ExprType::kEmptyStr) { + BuildExpression(choice_expr, grammar, start_state, end_states); + return; + } + } + XGRAMMAR_UNREACHABLE(); + } + + std::vector branch_end_states; + for (int32_t choice_id : expr) { + const auto& choice_expr = grammar->GetGrammarExpr(choice_id); + if (choice_expr.type == ExprType::kEmptyStr) { + continue; + } + int branch_start_state = target_fsm_.AddState(); + BuildExpression(choice_expr, grammar, branch_start_state, &branch_end_states); + target_fsm_.AddEpsilonEdge(start_state, branch_start_state); + end_states->insert(end_states->end(), branch_end_states.begin(), branch_end_states.end()); + } + + if (nullable) { + int nullable_branch_state = target_fsm_.AddState(); + target_fsm_.AddEpsilonEdge(start_state, nullable_branch_state); + end_states->push_back(nullable_branch_state); + } +} + +std::optional GrammarFSMBuilderImpl::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + return BuildExpressionFSM(expr, grammar); +} + +void GrammarFSMBuilderImpl::AppendFSM( + FSMWithStartEnd fsm, int start_state, std::vector* end_states +) { + const bool target_is_empty = target_fsm_.NumStates() == 1 && start_state == 0 && + target_fsm_.GetEdges(0).empty() && + target_fsm_.GetEdgeAuxData().empty() && fsm.GetStart() == 0; + if (target_is_empty) { + *end_states = fsm.GetEnds(); + target_fsm_ = std::move(fsm.GetFsm()); + return; + } + + std::vector state_mapping; + target_fsm_.AddFSM(fsm.GetFsm(), &state_mapping); + target_fsm_.AddEpsilonEdge(start_state, state_mapping[fsm.GetStart()]); + end_states->clear(); + end_states->reserve(fsm.GetEnds().size()); + for (int end_state : fsm.GetEnds()) { + end_states->push_back(state_mapping[end_state]); + } +} + +void GrammarFSMBuilderImpl::BuildRegex( + const std::string& regex, bool json_string, int start_state, std::vector* end_states +) { + const std::string rule_hint = rule_name_ != nullptr ? *rule_name_ : ""; + auto build_result = + json_string + ? RegexFSMBuilder::BuildWithForbiddenChars( + regex, GrammarFSMBuilder::JSONStringForbiddenChars(), grammar_builder_, rule_hint + ) + : RegexFSMBuilder::Build(regex, grammar_builder_, rule_hint); + if (build_result.IsErr()) { + auto error = std::move(build_result).UnwrapErr(); + if (rule_name_ != nullptr) { + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for rule " << *rule_name_ + << " with regex " << regex << ": " << error.what(); + } + XGRAMMAR_LOG(FATAL) << "Failed to build the automaton for regex " << regex << ": " + << error.what(); + } + AppendFSM(std::move(build_result).Unwrap(), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildSubstring( + const std::vector& chunks, int start_state, std::vector* end_states +) { + AppendFSM(SuffixAutomata::Build(chunks), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch, + int start_state, + std::vector* end_states +) { + auto build_result = TagDispatch(tag_dispatch); + XGRAMMAR_CHECK(build_result.has_value()) << "Failed to build tag dispatch FSM"; + AppendFSM(std::move(*build_result), start_state, end_states); +} + +void GrammarFSMBuilderImpl::BuildTokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch, + int start_state, + std::vector* end_states +) { + int trigger_count = static_cast(token_tag_dispatch.trigger_rule_pairs.size()); + std::vector dispatch_states; + dispatch_states.reserve(trigger_count); + for (int index = 0; index < trigger_count; ++index) { + dispatch_states.push_back(target_fsm_.AddState()); + } + + int end_state = -1; + end_states->clear(); + end_states->push_back(start_state); + if (!token_tag_dispatch.loop_after_dispatch) { + end_state = target_fsm_.AddState(); + end_states->push_back(end_state); + } + + std::vector excluded_token_ids; + excluded_token_ids.reserve( + token_tag_dispatch.trigger_rule_pairs.size() + token_tag_dispatch.excludes.size() + ); + for (const auto& trigger_rule_pair : token_tag_dispatch.trigger_rule_pairs) { + excluded_token_ids.push_back(trigger_rule_pair.first); + } + excluded_token_ids.insert( + excluded_token_ids.end(), + token_tag_dispatch.excludes.begin(), + token_tag_dispatch.excludes.end() + ); + std::sort(excluded_token_ids.begin(), excluded_token_ids.end()); + excluded_token_ids.erase( + std::unique(excluded_token_ids.begin(), excluded_token_ids.end()), excluded_token_ids.end() + ); + + for (int index = 0; index < trigger_count; ++index) { + auto [token_id, rule_id] = token_tag_dispatch.trigger_rule_pairs[index]; + int dispatch_target = token_tag_dispatch.loop_after_dispatch ? start_state : end_state; + target_fsm_.AddTokenEdge(start_state, dispatch_states[index], {token_id}); + target_fsm_.AddRuleEdge(dispatch_states[index], dispatch_target, rule_id); + } + target_fsm_.AddExcludeTokenEdge(start_state, start_state, excluded_token_ids); +} + +std::optional GrammarFSMBuilderImpl::BuildTagDispatchFSM( + const std::vector>& string_trigger_rules, + bool loop_after_dispatch, + const std::vector& excluded_strings +) { + std::vector tag_names; + tag_names.reserve(string_trigger_rules.size()); + for (const auto& [tag_name, tag_id] : string_trigger_rules) { + tag_names.push_back(tag_name); + } + std::vector end_states; + auto trie_result = TrieFSMBuilder::Build(tag_names, excluded_strings, &end_states, true, true); + if (!trie_result.has_value()) { + return std::nullopt; + } + auto trie_fsm = trie_result->GetFsm(); + auto start = trie_result->GetStart(); + + // The final end states are all but the trie's original end states. + std::vector ends; + for (int i = 0; i < trie_fsm.NumStates(); i++) { + if (!trie_result->IsEndState(i)) { + ends.push_back(i); + } + } + + // Add rule ref edges for string triggers + for (int i = 0; i < static_cast(string_trigger_rules.size()); i++) { + int next_state; + if (loop_after_dispatch) { + next_state = start; + } else { + next_state = trie_fsm.AddState(); + ends.push_back(next_state); + } + trie_fsm.AddRuleEdge(end_states[i], next_state, string_trigger_rules[i].second); + } + + return FSMWithStartEnd(trie_fsm, start, std::move(ends)); +} + +std::optional GrammarFSMBuilderImpl::TagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch +) { + std::vector> string_trigger_rules( + tag_dispatch.tag_rule_pairs.begin(), tag_dispatch.tag_rule_pairs.end() + ); + + return BuildTagDispatchFSM( + string_trigger_rules, tag_dispatch.loop_after_dispatch, tag_dispatch.excludes + ); +} + +Result GrammarFSMBuilderImpl::Regex(const std::string& regex, bool json_string) { + auto build_result = json_string ? RegexFSMBuilder::BuildWithForbiddenChars( + regex, GrammarFSMBuilder::JSONStringForbiddenChars() + ) + : RegexFSMBuilder::Build(regex); + if (build_result.IsErr()) { + return build_result; + } + auto result = std::move(build_result).Unwrap(); + result = result.SimplifyEpsilon(); + result = result.MergeEquivalentStates(); + return ResultOk(std::move(result)); +} + +/*! + * \brief Lower the regular subset used by RegexToEBNF without changing terminal semantics. + * + * A shared callee entry/exit would let one call return to another call's continuation. Inline + * each finite call instead, and connect only direct tail self-references to that call's own + * entry. The entry is private even when the caller's start also has an optional exit. + */ +class RegularGrammarFSMBuilder { + public: + RegularGrammarFSMBuilder(const Grammar& grammar, int max_num_states) + : grammar_(grammar), max_num_states_(max_num_states), active_(grammar->NumRules(), false) {} + + Result Build() { + if (max_num_states_ < 2) { + return ResultErr("Regular grammar FSM state limit must be at least 2"); + } + int start = fsm_.AddState(); + int end = fsm_.AddState(); + if (!BuildRule(grammar_->GetRootRuleId(), start, end)) { + return ResultErr(error_); + } + auto result = FSMWithStartEnd(fsm_, start, {end}).SimplifyEpsilon(); + return ResultOk(result.MergeEquivalentStates()); + } + + private: + bool Fail(const std::string& message) { + error_ = "Cannot flatten regular grammar: " + message; + return false; + } + + int AddState() { + if (fsm_.NumStates() >= max_num_states_) { + Fail("intermediate FSM exceeds " + std::to_string(max_num_states_) + " states"); + return -1; + } + return fsm_.AddState(); + } + + bool Append(const FSMWithStartEnd& leaf, int start, int end) { + if (leaf.GetFsm().NumStates() > max_num_states_ - fsm_.NumStates()) { + return Fail("intermediate FSM exceeds " + std::to_string(max_num_states_) + " states"); + } + std::vector mapping; + fsm_.AddFSM(leaf.GetFsm(), &mapping); + fsm_.AddEpsilonEdge(start, mapping[leaf.GetStart()]); + for (int leaf_end : leaf.GetEnds()) { + fsm_.AddEpsilonEdge(mapping[leaf_end], end); + } + return true; + } + + bool BuildRule(int rule_id, int start, int end) { + const auto& rule = grammar_->GetRule(rule_id); + if (active_[rule_id]) { + return Fail("mutual or non-tail recursion involving rule " + rule.name); + } + if (depth_ >= 256) { + return Fail("rule call depth exceeds 256"); + } + if (rule.lookahead_assertion_id >= 0 || rule.max_tokens >= 0 || rule.max_chars >= 0 || + !rule.capture_name.empty() || rule.is_lazy || rule.temperature.has_value() || + grammar_->GetSuffixStopInfo(rule_id) != nullptr) { + return Fail("annotated rule " + rule.name + " is not a plain regular rule"); + } + int entry = AddState(); + if (entry < 0) { + return false; + } + fsm_.AddEpsilonEdge(start, entry); + active_[rule_id] = true; + ++depth_; + bool ok = BuildExpression( + grammar_->GetGrammarExpr(rule.body_expr_id), entry, end, rule_id, entry, true + ); + --depth_; + active_[rule_id] = false; + return ok; + } + + bool BuildExpression( + const GrammarExpr& expr, int start, int end, int rule_id, int rule_start, bool tail + ) { + switch (expr.type) { + case ExprType::kEmptyStr: + fsm_.AddEpsilonEdge(start, end); + return true; + case ExprType::kByteString: + if (expr.size() >= max_num_states_ - fsm_.NumStates()) { + return Fail("byte string exceeds the intermediate FSM state limit"); + } + return Append(GrammarFSMBuilderImpl::ByteString(expr), start, end); + case ExprType::kCharacterClass: + case ExprType::kCharacterClassStar: + return Append(GrammarFSMBuilderImpl::CharacterClass(expr), start, end); + case ExprType::kRuleRef: + if (expr[0] == rule_id) { + if (!tail) { + return Fail("non-tail self-recursion in rule " + grammar_->GetRule(rule_id).name); + } + fsm_.AddEpsilonEdge(start, rule_start); + return true; + } + return BuildRule(expr[0], start, end); + case ExprType::kChoices: + for (int child : expr) { + if (!BuildExpression( + grammar_->GetGrammarExpr(child), start, end, rule_id, rule_start, tail + )) { + return false; + } + } + return true; + case ExprType::kSequence: { + int current = start; + for (int i = 0; i < expr.size(); ++i) { + int next = i + 1 == expr.size() ? end : AddState(); + if (next < 0 || !BuildExpression( + grammar_->GetGrammarExpr(expr[i]), + current, + next, + rule_id, + rule_start, + tail && i + 1 == expr.size() + )) { + return false; + } + current = next; + } + if (expr.size() == 0) { + fsm_.AddEpsilonEdge(start, end); + } + return true; + } + case ExprType::kRepeat: { + int64_t lower = expr[1], upper = expr[2]; + if (lower < 0 || (upper != -1 && upper < lower)) { + return Fail("invalid repetition bounds"); + } + // Even empty callees should not make an enormous explicit repetition consume + // unbounded compilation time. Large bounded patterns need a different representation. + int64_t count = upper == -1 ? lower : upper; + if (count > max_num_states_) { + return Fail("repetition exceeds the intermediate FSM state limit"); + } + int current = start; + for (int64_t i = 0; i <= count; ++i) { + if (i >= lower) { + fsm_.AddEpsilonEdge(current, end); + } + if (i == count) { + break; + } + int next = AddState(); + if (next < 0 || !BuildRule(expr[0], current, next)) { + return false; + } + current = next; + } + return upper != -1 || BuildRule(expr[0], current, current); + } + default: + return Fail( + "unsupported terminal or expression in rule " + grammar_->GetRule(rule_id).name + ); + } + } + + const Grammar& grammar_; + int max_num_states_; + FSM fsm_; + std::vector active_; + int depth_ = 0; + std::string error_; +}; + +class RepetitionRangeExpanderImpl : public GrammarMutator { + public: + using GrammarMutator::Apply; + using GrammarMutator::GrammarMutator; + + private: + int32_t VisitRepeat(const GrammarExpr& grammar_expr) final { + int32_t ref_rule_id = grammar_expr[0]; + int64_t lower = grammar_expr[1]; + int64_t upper = grammar_expr[2]; + return HandleRepetitionRange(cur_rule_name_, ref_rule_id, lower, upper); + } + + /*! + * \brief Handle repetition range by unzipping into explicit sequence/choice (for small bounds). + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t LegacyHandleRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Handle repetition range {lower, upper}, using unzip for small bounds or kRepeat for + * large. Identical repetitions are expanded only once and shared via memoization. + * \param cur_rule_name Name hint for generated rules. + * \param rule_id The rule to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t HandleRepetitionRange( + const std::string& cur_rule_name, int32_t rule_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Expand a repetition range into rules. Called by HandleRepetitionRange on cache miss. + * \param cur_rule_name Name hint for generated rules. + * \param grammar_expr_id The expression to repeat. + * \param lower Minimum count (inclusive). + * \param upper Maximum count (inclusive), or -1 for unbounded. + * \return grammar_expr_id of the repetition result. + */ + int32_t ExpandRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper + ); + + /*! + * \brief Memoization of expanded repetitions, mapping (content of the repeated expr, lower, + * upper) to the resulting grammar_expr_id. + * + * Grammars may contain a large number of identical repetitions. E.g. a JSON schema converted + * with max_whitespace_cnt emits one [ \n\r\t]{0,n} repetition per whitespace position, so a + * schema with 50k properties produces 200k+ identical repetitions. Expanding each occurrence + * into its own chain of rules multiplies the rule count by more than an order of magnitude, + * which blows up all downstream compilation stages (FSM building, token mask cache) in both + * time and memory. Sharing one expansion among identical repetitions keeps the rule count + * linear in the schema size. + */ + std::map, int32_t> repetition_cache_; +}; + +/****************** Repetition range helpers ******************/ + +int32_t RepetitionRangeExpanderImpl::LegacyHandleRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper +) { + // Construct expr expr ... expr (l times) + + std::vector elements; + for (int64_t i = 0; i < lower; ++i) { + elements.push_back(grammar_expr_id); + } + + // Case 1: {l}: + // expr expr ... expr (l times) + if (upper == lower) { + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); + } + + // Case 2: {l,}: + // expr expr ... expr (l times) rest + // rest ::= "" | expr rest + if (upper == -1) { + auto new_rule_name = builder_->GetNewRuleName(cur_rule_name); + auto new_rule_id = builder_->AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_->AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_->AddChoices( + {builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id, ref_to_new_rule})} + ); + builder_->UpdateRuleBody(new_rule_id, new_grammar_expr_id); + elements.push_back(builder_->AddRuleRef(new_rule_id)); + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); + } + + // Case 3: {l, r} (r - l >= 1) + // expr expr ... expr (l times) rest1 + // rest1 ::= "" | expr rest2 + // rest2 ::= "" | expr rest3 + // ... + // rest(r - l) ::= "" | expr + std::vector rest_rule_ids; + + for (int64_t i = 0; i < upper - lower; ++i) { + auto new_rule_name = builder_->GetNewRuleName(cur_rule_name); + rest_rule_ids.push_back(builder_->AddEmptyRule(new_rule_name)); + } + for (int64_t i = 0; i < upper - lower - 1; ++i) { + auto ref_to_next_rule = builder_->AddRuleRef(rest_rule_ids[i + 1]); + auto new_grammar_expr_id = builder_->AddChoices( + {builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id, ref_to_next_rule})} + ); + builder_->UpdateRuleBody(rest_rule_ids[i], new_grammar_expr_id); + } + auto last_grammar_expr_id = + builder_->AddChoices({builder_->AddEmptyStr(), builder_->AddSequence({grammar_expr_id})}); + builder_->UpdateRuleBody(rest_rule_ids.back(), last_grammar_expr_id); + + elements.push_back(builder_->AddRuleRef(rest_rule_ids[0])); + auto result_rule_id = builder_->AddRuleWithHint( + cur_rule_name, builder_->AddChoices({builder_->AddSequence(elements)}) + ); + return builder_->AddRuleRef(result_rule_id); +} + +int32_t RepetitionRangeExpanderImpl::HandleRepetitionRange( + const std::string& cur_rule_name, int32_t rule_id, int64_t lower, int64_t upper +) { + // Check if the referred rule is only one single element. If so, we can directly use the element + // for further optimization. + int32_t grammar_expr_id = builder_->AddRuleRef(rule_id); + const auto& ref_rule = base_grammar_->GetRule(rule_id); + const auto& ref_rule_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + // Keep the reference to budgeted, suffix/stop, lazy, and temperature rules: replacing it with + // the rule's content would erase the rule that the runtime semantics apply to. + if (ref_rule.max_tokens < 0 && ref_rule.max_chars < 0 && + base_grammar_->GetSuffixStopInfo(rule_id) == nullptr && !ref_rule.is_lazy && + !ref_rule.temperature.has_value() && + ref_rule_body.type == GrammarBuilder::GrammarExprType::kChoices && + ref_rule_body.size() == 1) { + const auto& ref_choice = base_grammar_->GetGrammarExpr(ref_rule_body[0]); + if (ref_choice.size() == 1) { + // Visit the element instead of copying it verbatim: a nested repetition such as + // b ::= a{2,} must be expanded too, or an unbounded kRepeat reaches the parser. + grammar_expr_id = VisitExpr(ref_choice[0]); + } + } + + // Memoize on (content of the repeated expr, lower, upper) so that identical repetitions share + // one expansion instead of each producing its own chain of rules. + const auto repeated_expr = builder_->GetGrammarExpr(grammar_expr_id); + std::vector cache_key; + cache_key.reserve(repeated_expr.size() + 3); + cache_key.push_back(static_cast(repeated_expr.type)); + cache_key.insert(cache_key.end(), repeated_expr.begin(), repeated_expr.end()); + cache_key.push_back(lower); + cache_key.push_back(upper); + auto it = repetition_cache_.find(cache_key); + if (it != repetition_cache_.end()) { + return it->second; + } + + int32_t result = ExpandRepetitionRange(cur_rule_name, grammar_expr_id, lower, upper); + repetition_cache_.emplace(std::move(cache_key), result); + return result; +} + +int32_t RepetitionRangeExpanderImpl::ExpandRepetitionRange( + const std::string& cur_rule_name, int32_t grammar_expr_id, int64_t lower, int64_t upper +) { + static const int64_t kUnzipThreshold = 128; + XGRAMMAR_CHECK(lower >= 0 && (upper == -1 || upper >= lower)) + << "Invalid repetition range {" << lower << ", " << upper << "}"; + + // Case 1.1 small upper (<=threshold), unzip the repetition. + // Case 1.2 unbounded upper, and lower is also small (<=threshold), unzip the lower part. + if ((upper != -1 && upper <= kUnzipThreshold) || (upper == -1 && lower <= kUnzipThreshold)) { + return LegacyHandleRepetitionRange(cur_rule_name, grammar_expr_id, lower, upper); + } + + // Case 2. upper is unbounded, and lower is large (>threshold). + // Or upper is bounded, but upper > threshold. + + // Case 2.1.1. lower is smaller than threshold, and upper is large. Transform {lower, upper} into: + // {threshold, upper} | {lower, threshold} + std::vector choices; + if (lower < kUnzipThreshold) { + choices.push_back(builder_->AddSequence( + {LegacyHandleRepetitionRange(cur_rule_name, grammar_expr_id, lower, kUnzipThreshold - 1)} + )); + lower = kUnzipThreshold; + } + + std::optional infinite_repetition_id = std::nullopt; + std::vector repeated_sequence; + // Now, we transform {lower, upper} into {max{threshold, lower}, upper}. + // Case 2.2 upper is unbounded. We will transform it into {lower} {0, inf}. + if (upper == -1) { + const auto& rule_expr = builder_->GetGrammarExpr(grammar_expr_id); + if (rule_expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + std::vector character_ranges; + bool is_negative = rule_expr[0]; + for (int i = 1; i < static_cast(rule_expr.size()); i += 2) { + character_ranges.push_back({rule_expr[i], rule_expr[i + 1]}); + } + infinite_repetition_id = builder_->AddCharacterClassStar(character_ranges, is_negative); + } else { + const auto unbounded_rule_id = + builder_->AddEmptyRule(builder_->GetNewRuleName(cur_rule_name + "_repeat_inf")); + int recursion_sequence = + builder_->AddSequence({grammar_expr_id, builder_->AddRuleRef(unbounded_rule_id)}); + int recursion_choice = builder_->AddChoices({builder_->AddEmptyStr(), recursion_sequence}); + builder_->UpdateRuleBody(unbounded_rule_id, recursion_choice); + infinite_repetition_id = builder_->AddRuleRef(unbounded_rule_id); + } + upper = lower; + } + + // Handle the {lower, upper} part, where threshold <= lower <= upper. + const auto repeat_name = cur_rule_name + "_repeat_1"; + XGRAMMAR_DCHECK(lower >= kUnzipThreshold && upper >= lower); + + // If we have infinite repetition part, add it to the sequence. + if (infinite_repetition_id.has_value()) { + repeated_sequence.push_back(infinite_repetition_id.value()); + } + + // The repetition body. + if (upper != kUnzipThreshold) { + XGRAMMAR_DCHECK(upper > kUnzipThreshold); + auto new_grammar_expr_id = builder_->AddChoices({builder_->AddSequence({grammar_expr_id})}); + auto new_rule_id = builder_->AddRuleWithHint(repeat_name, new_grammar_expr_id); + auto new_repeated_ref_rule_expr = builder_->AddChoices({builder_->AddSequence( + {builder_->AddRepeat(new_rule_id, lower - kUnzipThreshold, upper - kUnzipThreshold)} + )}); + auto new_repeated_rule_id = + builder_->AddRuleWithHint(repeat_name + "_inner", new_repeated_ref_rule_expr); + repeated_sequence.push_back(builder_->AddRuleRef(new_repeated_rule_id)); + std::vector repetition_lookahead(kUnzipThreshold, grammar_expr_id); + builder_->UpdateLookaheadAssertion(new_rule_id, builder_->AddSequence(repetition_lookahead)); + } + + // Add the last threshold grammar_expr_id to the sequence. + for (int i = 0; i < kUnzipThreshold; ++i) { + repeated_sequence.push_back(grammar_expr_id); + } + + // Add the sequence to choices. + choices.push_back(builder_->AddSequence(repeated_sequence)); + auto result_rule_id = builder_->AddRuleWithHint(cur_rule_name, builder_->AddChoices(choices)); + return builder_->AddRuleRef(result_rule_id); +} + +class RepetitionNormalizerImpl { + public: + void Apply(Grammar* grammar) { + auto& grammar_ref = *grammar; + for (int i = 0; i < grammar_ref->NumGrammarExprs(); ++i) { + auto expr = grammar_ref->GetGrammarExpr(i); + if (expr.type != Grammar::Impl::GrammarExprType::kRepeat) { + continue; + } + int repeat_rule_id = expr[0]; + grammar_ref->GetRule(repeat_rule_id).is_exact_lookahead = true; + if (std::binary_search( + grammar_ref->allow_empty_rule_ids.begin(), + grammar_ref->allow_empty_rule_ids.end(), + repeat_rule_id + )) { + // The repeated rule can be empty, so we need to normalize it. + expr.SetData(1, 0); // Set min repeat to 0 + } + } + } +}; + +/*! + * \brief Rewrite lazy rule bodies into their terminal-like form where possible: unwrap the + * single-reference chains produced by regex conversion, and flatten the right-recursive plus + * pattern (x ::= cc x | cc) and star pattern (x ::= cc x | "") produced by regex conversion and + * repetition expansion into (cc cc*) and (cc*). Grammars without lazy rules are returned + * unchanged. + */ +class LazyBodyFlattenerImpl : public GrammarMutator { + public: + using GrammarMutator::GrammarMutator; + + Grammar Apply(const Grammar& grammar) final { + bool has_lazy_rule = false; + for (int i = 0; i < grammar->NumRules(); ++i) { + has_lazy_rule = has_lazy_rule || grammar->GetRule(i).is_lazy; + } + if (!has_lazy_rule) { + return grammar; + } + InitGrammar(grammar); + InitBuilder(); + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + int32_t new_body_expr_id = + rule.is_lazy ? BuildFlattenedLazyBody(rule.body_expr_id) : VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(i, new_body_expr_id); + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } + + private: + int32_t BuildFlattenedLazyBody(int32_t body_expr_id) { + auto body_type = base_grammar_->GetGrammarExpr(body_expr_id).type; + if (body_type == GrammarExprType::kRegex || body_type == GrammarExprType::kSubstring) { + return VisitExpr(body_expr_id); + } + // Unwrap chains of single rule references (r ::= (x), x ::= (y), ...) produced by regex + // conversion, and detect the plus-desugar pattern at the top level. + int32_t cur_body_id = body_expr_id; + int32_t cur_rule_id = -1; + for (int depth = 0; depth < 64; ++depth) { + const auto& body = base_grammar_->GetGrammarExpr(cur_body_id); + if (body.type != GrammarExprType::kChoices || body.size() != 1) { + break; + } + const auto& choice = base_grammar_->GetGrammarExpr(body[0]); + if (choice.type != GrammarExprType::kSequence || choice.size() != 1) { + break; + } + const auto& element = base_grammar_->GetGrammarExpr(choice[0]); + if (element.type != GrammarExprType::kRuleRef || base_grammar_->GetRule(element[0]).is_lazy) { + break; + } + cur_rule_id = element[0]; + cur_body_id = base_grammar_->GetRule(cur_rule_id).body_expr_id; + } + + const auto& body = base_grammar_->GetGrammarExpr(cur_body_id); + if (body.type != GrammarExprType::kChoices) { + XGRAMMAR_LOG(WARNING) << "The body of the lazy rule '" << cur_rule_name_ + << "' cannot be flattened into a terminal-like form"; + return VisitExpr(cur_body_id); + } + std::vector repeat_elements; + if (cur_rule_id != -1 && TryEmitRepeatPattern(body, cur_rule_id, &repeat_elements)) { + return builder_->AddChoices({builder_->AddSequence(repeat_elements)}); + } + std::vector new_choice_ids; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence) { + new_choice_ids.push_back(VisitExpr(choice_id)); + continue; + } + std::vector elements; + if (!FlattenSequenceInto(choice, &elements, 0)) { + // Not flattenable; copy as is and let the terminal-like validation report the error. + XGRAMMAR_LOG(WARNING) << "The body of the lazy rule '" << cur_rule_name_ + << "' cannot be flattened into a terminal-like form"; + return VisitExpr(cur_body_id); + } + new_choice_ids.push_back(builder_->AddSequence(elements)); + } + return builder_->AddChoices(new_choice_ids); + } + + /*! \brief Append the flattened elements of the sequence, splicing rule references whose body + * is a single terminal-like sequence or the plus-desugar pattern, and coalescing references + * to single-character alternations into character classes. Returns false if some element + * cannot be flattened. */ + bool FlattenSequenceInto(const GrammarExpr& seq, std::vector* elements, int depth) { + if (depth > 64) { + return false; + } + for (auto element_id : seq) { + const auto& element = base_grammar_->GetGrammarExpr(element_id); + if (element.type == GrammarExprType::kByteString || + element.type == GrammarExprType::kCharacterClass || + element.type == GrammarExprType::kCharacterClassStar) { + elements->push_back(builder_->AddGrammarExpr(element)); + continue; + } + if (element.type != GrammarExprType::kRuleRef) { + return false; + } + const auto& ref_rule = base_grammar_->GetRule(element[0]); + if (ref_rule.is_lazy) { + return false; + } + const auto& ref_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + if (TryEmitRepeatPattern(ref_body, element[0], elements)) { + continue; + } + if (ref_body.type == GrammarExprType::kChoices && ref_body.size() == 1) { + const auto& only_choice = base_grammar_->GetGrammarExpr(ref_body[0]); + if (only_choice.type == GrammarExprType::kEmptyStr) { + continue; + } + if (only_choice.type == GrammarExprType::kSequence && + FlattenSequenceInto(only_choice, elements, depth + 1)) { + continue; + } + } + std::vector ranges; + if (CollectSingleCharRanges(element_id, &ranges, 0)) { + elements->push_back(builder_->AddCharacterClass(UnionRanges(std::move(ranges)), false)); + continue; + } + return false; + } + return true; + } + + /*! \brief Resolve an expr matching exactly one character into the set of codepoint ranges it + * accepts, appending them to ranges. Accepts character classes, single-byte strings, and + * references to non-lazy rules that are alternations of such elements. Returns false + * otherwise. */ + bool CollectSingleCharRanges( + int32_t expr_id, std::vector* ranges, int depth + ) { + if (depth > 64) { + return false; + } + const auto& expr = base_grammar_->GetGrammarExpr(expr_id); + if (expr.type == GrammarExprType::kCharacterClass) { + AppendPositiveRanges(expr, ranges); + return true; + } + if (expr.type == GrammarExprType::kByteString && expr.size() == 1) { + ranges->push_back({expr[0], expr[0]}); + return true; + } + if (expr.type != GrammarExprType::kRuleRef) { + return false; + } + const auto& rule = base_grammar_->GetRule(expr[0]); + if (rule.is_lazy) { + return false; + } + const auto& body = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (body.type != GrammarExprType::kChoices) { + return false; + } + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence || choice.size() != 1 || + !CollectSingleCharRanges(choice[0], ranges, depth + 1)) { + return false; + } + } + return true; + } + + /*! \brief If the body matches the plus-desugar pattern (self ::= e self | e) or the + * (generalized) star-desugar pattern (self ::= "" | e1 self | e2 self | ...), append the + * equivalent (e e*) or ((e1|e2|...)*) to the elements and return true. */ + bool TryEmitRepeatPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + return TryEmitPlusPattern(body, self_rule_id, elements) || + TryEmitGeneralizedPlusPattern(body, self_rule_id, elements) || + TryEmitStarPattern(body, self_rule_id, elements); + } + + /*! \brief If the body matches the generalized plus-desugar pattern (self ::= e1 self | ... | + * e1 | ...) with single-character elements whose recursive and base unions are equal, append + * the equivalent (cc cc*) over the union to the elements and return true. */ + bool TryEmitGeneralizedPlusPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices) { + return false; + } + std::vector recursive_ranges; + std::vector base_ranges; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type != GrammarExprType::kSequence) { + return false; + } + if (choice.size() == 1) { + if (!CollectSingleCharRanges(choice[0], &base_ranges, 0)) { + return false; + } + continue; + } + if (choice.size() == 2) { + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type == GrammarExprType::kRuleRef && tail[0] == self_rule_id && + CollectSingleCharRanges(choice[0], &recursive_ranges, 0)) { + continue; + } + } + return false; + } + recursive_ranges = UnionRanges(std::move(recursive_ranges)); + base_ranges = UnionRanges(std::move(base_ranges)); + // The unions must coincide: with differing sets (e.g. self ::= a self | b, which is a*b), + // the language is not (a|b)+. + if (recursive_ranges.empty() || !RangesEqual(recursive_ranges, base_ranges)) { + return false; + } + elements->push_back(builder_->AddCharacterClass(recursive_ranges, false)); + elements->push_back(builder_->AddCharacterClassStar(recursive_ranges, false)); + return true; + } + + /*! \brief If the body matches the plus-desugar pattern (self ::= e self | e) with e resolving + * to a star-expressible expr, append the equivalent (e e*) (or (e*) when e itself resolves to + * a star) to the elements and return true. */ + bool TryEmitPlusPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices || body.size() != 2) { + return false; + } + for (int recursive_pos = 0; recursive_pos < 2; ++recursive_pos) { + const auto& recursive = base_grammar_->GetGrammarExpr(body[recursive_pos]); + const auto& base = base_grammar_->GetGrammarExpr(body[1 - recursive_pos]); + if (recursive.type != GrammarExprType::kSequence || recursive.size() != 2 || + base.type != GrammarExprType::kSequence || base.size() != 1) { + continue; + } + const auto& element = base_grammar_->GetGrammarExpr(recursive[0]); + const auto& tail = base_grammar_->GetGrammarExpr(recursive[1]); + const auto& base_element = base_grammar_->GetGrammarExpr(base[0]); + if (tail.type != GrammarExprType::kRuleRef || tail[0] != self_rule_id || + element.type != base_element.type || element.size() != base_element.size() || + !std::equal(element.begin(), element.end(), base_element.begin())) { + continue; + } + int32_t resolved_id = ResolveStarExpressible(recursive[0]); + if (resolved_id == -1) { + // Not a single terminal; e may still be a single-character alternation, giving + // (e1|e2|...)+ = cc cc* over the union of the ranges. + std::vector ranges; + if (!CollectSingleCharRanges(recursive[0], &ranges, 0)) { + continue; + } + ranges = UnionRanges(std::move(ranges)); + elements->push_back(builder_->AddCharacterClass(ranges, false)); + elements->push_back(builder_->AddCharacterClassStar(ranges, false)); + return true; + } + const auto& resolved = base_grammar_->GetGrammarExpr(resolved_id); + if (resolved.type == GrammarExprType::kCharacterClassStar) { + // (e*)+ is e*. + elements->push_back(builder_->AddGrammarExpr(resolved)); + return true; + } + std::vector character_ranges; + bool is_negative = false; + if (resolved.type == GrammarExprType::kCharacterClass) { + is_negative = static_cast(resolved[0]); + for (int i = 1; i < static_cast(resolved.size()); i += 2) { + character_ranges.push_back({resolved[i], resolved[i + 1]}); + } + } else { // single-byte kByteString + character_ranges.push_back({resolved[0], resolved[0]}); + } + elements->push_back(builder_->AddGrammarExpr(resolved)); + elements->push_back(builder_->AddCharacterClassStar(character_ranges, is_negative)); + return true; + } + return false; + } + + /*! \brief If the body matches the generalized star-desugar pattern (self ::= "" | e1 self | + * e2 self | ...) with each e resolving to character ranges, append the equivalent single + * character class star ((e1|e2|...)*) to the elements and return true. */ + bool TryEmitStarPattern( + const GrammarExpr& body, int32_t self_rule_id, std::vector* elements + ) { + if (body.type != GrammarExprType::kChoices) { + return false; + } + bool has_empty = false; + std::vector ranges; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type == GrammarExprType::kEmptyStr) { + has_empty = true; + continue; + } + if (choice.type != GrammarExprType::kSequence || choice.size() != 2) { + return false; + } + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type != GrammarExprType::kRuleRef || tail[0] != self_rule_id) { + return false; + } + if (!CollectStarRanges(choice[0], &ranges, 0)) { + return false; + } + } + if (!has_empty || ranges.empty()) { + return false; + } + elements->push_back(builder_->AddCharacterClassStar(UnionRanges(std::move(ranges)), false)); + return true; + } + + /*! \brief Resolve an expr repeated under an enclosing star into the set of codepoint ranges it + * repeats over, appending them to ranges. Accepts character classes, character class stars, + * single-byte strings, and references to non-lazy rules that are alternations of such + * elements, or star/plus recursions over them — under an enclosing star, all of these are + * equivalent to the union of their character ranges. Returns false otherwise. */ + bool CollectStarRanges( + int32_t expr_id, std::vector* ranges, int depth + ) { + if (depth > 64) { + return false; + } + const auto& expr = base_grammar_->GetGrammarExpr(expr_id); + if (expr.type == GrammarExprType::kCharacterClass || + expr.type == GrammarExprType::kCharacterClassStar) { + AppendPositiveRanges(expr, ranges); + return true; + } + if (expr.type == GrammarExprType::kByteString && expr.size() == 1) { + ranges->push_back({expr[0], expr[0]}); + return true; + } + if (expr.type != GrammarExprType::kRuleRef) { + return false; + } + int32_t rule_id = expr[0]; + const auto& rule = base_grammar_->GetRule(rule_id); + if (rule.is_lazy) { + return false; + } + const auto& body = base_grammar_->GetGrammarExpr(rule.body_expr_id); + if (body.type != GrammarExprType::kChoices) { + return false; + } + bool has_empty = false; + std::vector base_elements; + std::vector recursive_elements; + for (auto choice_id : body) { + const auto& choice = base_grammar_->GetGrammarExpr(choice_id); + if (choice.type == GrammarExprType::kEmptyStr) { + has_empty = true; + continue; + } + if (choice.type != GrammarExprType::kSequence) { + return false; + } + if (choice.size() == 1) { + base_elements.push_back(choice[0]); + continue; + } + if (choice.size() == 2) { + const auto& tail = base_grammar_->GetGrammarExpr(choice[1]); + if (tail.type == GrammarExprType::kRuleRef && tail[0] == rule_id) { + recursive_elements.push_back(choice[0]); + continue; + } + } + return false; + } + // The safe shapes: an alternation (a | b | ...), a star ("" | e1 self | ...), and a plus + // (e self | e, or single-character alternated forms with equal recursive/base unions). + // Mixed shapes like (a self | b) are a*b, whose star is not the union, so they are rejected. + std::vector* collect = nullptr; + if (recursive_elements.empty()) { + collect = &base_elements; + } else if (base_elements.empty() && has_empty) { + collect = &recursive_elements; + } else if (recursive_elements.size() == 1 && base_elements.size() == 1 && !has_empty && + ExprsEqual(recursive_elements[0], base_elements[0])) { + collect = &recursive_elements; + } else if (!has_empty) { + std::vector recursive_ranges; + std::vector base_ranges; + for (auto element_id : recursive_elements) { + if (!CollectSingleCharRanges(element_id, &recursive_ranges, depth + 1)) { + return false; + } + } + for (auto element_id : base_elements) { + if (!CollectSingleCharRanges(element_id, &base_ranges, depth + 1)) { + return false; + } + } + recursive_ranges = UnionRanges(std::move(recursive_ranges)); + if (!RangesEqual(recursive_ranges, UnionRanges(std::move(base_ranges)))) { + return false; + } + ranges->insert(ranges->end(), recursive_ranges.begin(), recursive_ranges.end()); + return true; + } else { + return false; + } + for (auto element_id : *collect) { + if (!CollectStarRanges(element_id, ranges, depth + 1)) { + return false; + } + } + return true; + } + + /*! \brief Whether two exprs have identical type and content. */ + bool ExprsEqual(int32_t lhs_id, int32_t rhs_id) { + const auto& lhs = base_grammar_->GetGrammarExpr(lhs_id); + const auto& rhs = base_grammar_->GetGrammarExpr(rhs_id); + return lhs.type == rhs.type && lhs.size() == rhs.size() && + std::equal(lhs.begin(), lhs.end(), rhs.begin()); + } + + /*! \brief Whether two normalized range vectors are identical. */ + static bool RangesEqual( + const std::vector& lhs, + const std::vector& rhs + ) { + if (lhs.size() != rhs.size()) { + return false; + } + for (size_t i = 0; i < lhs.size(); ++i) { + if (lhs[i].lower != rhs[i].lower || lhs[i].upper != rhs[i].upper) { + return false; + } + } + return true; + } + + /*! \brief Append the positive codepoint ranges of a character class or character class star, + * complementing negated classes over [0, 0x10FFFF]. */ + void AppendPositiveRanges( + const GrammarExpr& expr, std::vector* ranges + ) { + std::vector class_ranges; + for (int i = 1; i < static_cast(expr.size()); i += 2) { + class_ranges.push_back({expr[i], expr[i + 1]}); + } + if (!static_cast(expr[0])) { + ranges->insert(ranges->end(), class_ranges.begin(), class_ranges.end()); + return; + } + class_ranges = UnionRanges(std::move(class_ranges)); + int32_t next = 0; + for (const auto& range : class_ranges) { + if (range.lower > next) { + ranges->push_back({next, range.lower - 1}); + } + next = std::max(next, range.upper + 1); + } + if (next <= 0x10FFFF) { + ranges->push_back({next, 0x10FFFF}); + } + } + + /*! \brief Sort the ranges and merge overlapping or adjacent ones. */ + static std::vector UnionRanges( + std::vector ranges + ) { + std::sort(ranges.begin(), ranges.end(), [](const auto& a, const auto& b) { + return a.lower < b.lower; + }); + std::vector result; + for (const auto& range : ranges) { + if (!result.empty() && range.lower <= result.back().upper + 1) { + result.back().upper = std::max(result.back().upper, range.upper); + } else { + result.push_back(range); + } + } + return result; + } + + /*! \brief Resolve an expr through chains of non-lazy single-reference rules to a + * star-expressible expr: a character class, a single-byte string, or a character class star. + * Returns the resolved expr id, or -1 if it does not resolve to one. */ + int32_t ResolveStarExpressible(int32_t expr_id) { + int32_t cur_id = expr_id; + for (int depth = 0; depth < 64; ++depth) { + const auto& cur = base_grammar_->GetGrammarExpr(cur_id); + if (cur.type == GrammarExprType::kCharacterClass || + cur.type == GrammarExprType::kCharacterClassStar || + (cur.type == GrammarExprType::kByteString && cur.size() == 1)) { + return cur_id; + } + if (cur.type != GrammarExprType::kRuleRef) { + return -1; + } + const auto& ref_rule = base_grammar_->GetRule(cur[0]); + if (ref_rule.is_lazy) { + return -1; + } + const auto& ref_body = base_grammar_->GetGrammarExpr(ref_rule.body_expr_id); + if (ref_body.type != GrammarExprType::kChoices || ref_body.size() != 1) { + return -1; + } + const auto& only_choice = base_grammar_->GetGrammarExpr(ref_body[0]); + if (only_choice.type != GrammarExprType::kSequence || only_choice.size() != 1) { + return -1; + } + cur_id = only_choice[0]; + } + return -1; + } +}; + +class GrammarOptimizerImpl { + public: + static Grammar Apply(const Grammar& grammar) { + // ByteStringFuser and RuleInliner rewrite the grammar in place, so work on a private copy: the + // input grammar may be shared (e.g. a cached grammar) and must not be mutated. Copy the impl + // directly (contiguous vector copies) instead of going through GrammarBuilder, which would + // also build the unneeded rule name map. + Grammar result(std::make_shared(*grammar.operator->())); + ByteStringFuser::Apply(&result); + RuleInliner::Apply(&result); + result = RepetitionRangeExpander::Apply(result); + result = LazyBodyFlattenerImpl().Apply(result); + result = DeadCodeEliminator::Apply(result); + result = LookaheadAssertionAnalyzer::Apply(result); + result->allow_empty_rule_ids = AllowEmptyRuleAnalyzer::Apply(result); + ValidateLazyRules(result); + RepetitionNormalizer::Apply(&result); + GrammarFSMBuilder::Apply(&result); + result->optimized = true; + return result; + } + + private: + /*! + * \brief Committed-shortest (lazy) matching requires the whole rule body to compile into a + * single per-rule FSM without rule references, so that the states of one occurrence are exactly + * the states with the rule's id. + */ + static void ValidateLazyRules(const Grammar& grammar) { + for (int32_t i = 0; i < grammar->NumRules(); ++i) { + const auto& rule = grammar->GetRule(i); + if (!rule.is_lazy) { + continue; + } + const auto& body = grammar->GetGrammarExpr(rule.body_expr_id); + if (body.type == Grammar::Impl::GrammarExprType::kRegex || + body.type == Grammar::Impl::GrammarExprType::kSubstring) { + continue; + } + XGRAMMAR_CHECK(body.type == Grammar::Impl::GrammarExprType::kChoices) + << "lazy rule '" << rule.name << "' must have a terminal-like body"; + for (auto choice_id : body) { + const auto& choice = grammar->GetGrammarExpr(choice_id); + if (choice.type == Grammar::Impl::GrammarExprType::kEmptyStr) { + continue; + } + for (auto element_id : choice) { + const auto& element = grammar->GetGrammarExpr(element_id); + XGRAMMAR_CHECK( + element.type == Grammar::Impl::GrammarExprType::kByteString || + element.type == Grammar::Impl::GrammarExprType::kCharacterClass || + element.type == Grammar::Impl::GrammarExprType::kCharacterClassStar + ) << "lazy rule '" + << rule.name + << "' must have a terminal-like body (strings, character classes, and regexes that " + "compile to a single FSM); rule references and repetition ranges are not supported"; + } + } + } + } +}; + +/*! + * \brief Fuse adjacent byte string elements in sequences. + * \details Rewrites the grammar in place: only sequences that actually contain a run of adjacent + * byte strings are rebuilt, the rest keep their original ids. Stale exprs are removed later by + * DeadCodeEliminator. + */ +class ByteStringFuserImpl : public InPlaceGrammarRewriter { + protected: + int32_t VisitSequence(int32_t expr_id) override { + // Read-only probe: rewrite only if the sequence contains an empty byte string or two adjacent + // byte strings. The probe appends nothing, so no memory is allocated for the common unchanged + // case. + bool previous_is_byte_string = false; + bool needs_rewrite = false; + { + auto expr = builder_.GetGrammarExpr(expr_id); + for (int32_t element_id : expr) { + auto element = builder_.GetGrammarExpr(element_id); + bool is_byte_string = element.type == GrammarExprType::kByteString; + if (is_byte_string && (previous_is_byte_string || element.size() == 0)) { + needs_rewrite = true; + break; + } + previous_is_byte_string = is_byte_string; + } + } + if (!needs_rewrite) { + return expr_id; + } + // Copy the element ids first: fusing appends to the arena and invalidates expr views. + auto expr = builder_.GetGrammarExpr(expr_id); + std::vector element_ids(expr.begin(), expr.end()); + std::vector new_element_ids; + for (size_t i = 0; i < element_ids.size();) { + if (builder_.GetGrammarExpr(element_ids[i]).type != GrammarExprType::kByteString) { + new_element_ids.push_back(element_ids[i]); + ++i; + continue; + } + // Gather the run of adjacent byte strings starting at i. + std::vector fused_bytes; + size_t run_end = i; + while (run_end < element_ids.size()) { + auto element = builder_.GetGrammarExpr(element_ids[run_end]); + if (element.type != GrammarExprType::kByteString) { + break; + } + fused_bytes.insert(fused_bytes.end(), element.begin(), element.end()); + ++run_end; + } + // Empty byte strings are epsilon and can be omitted from a sequence. + if (!fused_bytes.empty()) { + if (run_end - i == 1) { + new_element_ids.push_back(element_ids[i]); + } else { + new_element_ids.push_back(builder_.AddByteString(fused_bytes)); + } + } + i = run_end; + } + return builder_.AddSequence(new_element_ids); + } +}; + +class RootRuleRenamerImpl { + public: + static Grammar Apply(const Grammar& grammar) { + // If the root name is "root", return directly. + if (grammar->GetRootRule().name == "root") { + return grammar; + } + + // Collect all the rule names. + std::unordered_set rule_names; + int root_name_rule_id = -1; + for (int i = 0; i < grammar->NumRules(); i++) { + const auto& rule_name = grammar->GetRule(i).name; + if (rule_name == "root") { + root_name_rule_id = i; + } + rule_names.insert(rule_name); + } + + // Rename the rules. + Grammar grammar_copy = grammar; + grammar_copy->GetRule(grammar_copy->GetRootRuleId()).name = "root"; + if (root_name_rule_id != -1) { + std::string rule_prefix = "root_"; + bool renamed = false; + for (int i = 0; i <= grammar_copy->NumRules(); i++) { + std::string new_rule_name = rule_prefix + std::to_string(i); + if (rule_names.find(new_rule_name) == rule_names.end()) { + grammar_copy->GetRule(root_name_rule_id).name = new_rule_name; + renamed = true; + break; + } + } + XGRAMMAR_DCHECK(renamed) << "Rule renaming must succeed within (n + 1) attempts."; + } + return grammar_copy; + } +}; + +class GrammarFSMHasherImpl { + public: + void Apply(Grammar* grammar); + static std::optional HashSequence(const Grammar& grammar, int32_t sequence_id); + + static constexpr int16_t kNotEndStateFlag = -0x100; + static constexpr int16_t kEndStateFlag = -0x200; + static constexpr int16_t kSelfRecursionFlag = -0x300; + static constexpr int16_t kSimpleCycleFlag = -0x400; + static constexpr int16_t kUnKnownFlag = -0x500; + + private: + Grammar* grammar_; + std::vector visited_; + std::vector> ref_graph_from_referrer_to_referee_; + std::vector> ref_graph_from_referee_to_referrer_; + std::vector> sorted_edges_; + std::vector has_inward_edges_; + + /*! + * \brief The worklist of fsms that are ready to be hashed: fsms whose references are all + * hashed (except possibly a self-recursion). Maintained incrementally so that the main hashing + * loop is O(V + E) instead of rescanning all rules after each hashed fsm. + */ + std::queue ready_queue_; + + /*! + * \brief Get the hash value of a fsm, with a given grammar. + */ + uint64_t HashFsm(int fsm_index); + + /*! + * \brief Find a simple cycle in the reference graph, And hash the + * fsms in the simple cycle. + */ + bool FindSimpleCycle(); + + /*! + * \brief Hash the fsms in the simple cycle. + */ + void HashSimpleCycle(const std::vector& simple_cycle); + + /*! + * \brief Check if a fsm is ready to be hashed: it is not hashed yet, and it references no + * unhashed fsms other than itself. + */ + bool IsReadyToHash(int32_t fsm_index) const { + if (visited_[fsm_index]) { + return false; + } + const auto& referees = ref_graph_from_referrer_to_referee_[fsm_index]; + return referees.empty() || (referees.size() == 1 && referees[0] == fsm_index); + } + + /*! + * \brief Remove the hashed fsm from the reference graph, and push the referrers that become + * ready to hash into the ready queue. + */ + void RemoveHashedFsmFromRefGraph(int32_t fsm_index); + + std::pair IsPartialHashable(int fsm_index); +}; + +bool GrammarFSMHasherImpl::FindSimpleCycle() { + // Try to find a simple cycle. + std::vector not_simple_cycle = visited_; + // Allocated once and cleaned up after each walk, to avoid an O(num_rules) allocation per + // outer iteration. + std::vector in_stack(ref_graph_from_referee_to_referrer_.size(), false); + for (size_t i = 0; i < ref_graph_from_referee_to_referrer_.size(); i++) { + if (not_simple_cycle[i]) { + continue; + } + // Not a simple cycle if it has more than one referee. + std::stack dfs_stack; + std::vector simple_cycle; + std::vector walked_states; + dfs_stack.push(static_cast(i)); + int32_t current_fsm_index = i; + in_stack[current_fsm_index] = true; + walked_states.push_back(current_fsm_index); + while ((ref_graph_from_referrer_to_referee_[current_fsm_index].size() == 1) && + !not_simple_cycle[current_fsm_index]) { + XGRAMMAR_CHECK(current_fsm_index != ref_graph_from_referrer_to_referee_[current_fsm_index][0]) + << "Self-recursion cycle found in the reference graph, which is not allowed."; + not_simple_cycle[current_fsm_index] = true; + current_fsm_index = ref_graph_from_referrer_to_referee_[current_fsm_index][0]; + if (in_stack[current_fsm_index]) { + simple_cycle.push_back(current_fsm_index); + while (dfs_stack.top() != current_fsm_index) { + simple_cycle.push_back(dfs_stack.top()); + dfs_stack.pop(); + } + // Found a simple cycle. + break; + } else { + dfs_stack.push(current_fsm_index); + in_stack[current_fsm_index] = true; + walked_states.push_back(current_fsm_index); + } + } + if (!simple_cycle.empty()) { + HashSimpleCycle(simple_cycle); + return true; + } + for (auto state : walked_states) { + in_stack[state] = false; + } + } + return false; +} + +void GrammarFSMHasherImpl::HashSimpleCycle(const std::vector& simple_cycle) { + // Initialize the cycle hash. + for (const auto& cycle_id : simple_cycle) { + visited_[cycle_id] = true; + grammar_->ImplPtr()->per_rule_fsm_hashes[cycle_id] = kSimpleCycleFlag; + } + + std::vector local_cycle_hash; + local_cycle_hash.reserve(simple_cycle.size()); + for (const auto& cycle_id : simple_cycle) { + local_cycle_hash.push_back(HashFsm(cycle_id)); + } + std::vector local_cycle_hash_copy = local_cycle_hash; + for (int i = 0; i < static_cast(local_cycle_hash.size()); i++) { + uint64_t current_hash = 0; + for (int j = 0; j < static_cast(local_cycle_hash.size()); j++) { + current_hash = + HashCombine(current_hash, local_cycle_hash_copy[(i + j) % local_cycle_hash.size()]); + } + local_cycle_hash[i] = current_hash; + } + + for (int i = 0; i < static_cast(simple_cycle.size()); i++) { + grammar_->ImplPtr()->per_rule_fsm_hashes[simple_cycle[i]] = local_cycle_hash[i]; + RemoveHashedFsmFromRefGraph(simple_cycle[i]); + } +} + +void GrammarFSMHasherImpl::RemoveHashedFsmFromRefGraph(int32_t fsm_index) { + for (const auto& referer : ref_graph_from_referee_to_referrer_[fsm_index]) { + auto& referees = ref_graph_from_referrer_to_referee_[referer]; + auto it = std::find(referees.begin(), referees.end(), fsm_index); + if (it != referees.end()) { + referees.erase(it); + } + if (IsReadyToHash(referer)) { + ready_queue_.push(referer); + } + } +} + +void GrammarFSMHasherImpl::Apply(Grammar* grammar) { + grammar_ = grammar; + grammar->ImplPtr()->per_rule_fsm_hashes = + std::vector>((*grammar)->NumRules()); + grammar->ImplPtr()->per_rule_fsm_new_state_ids.resize((*grammar)->NumRules()); + ref_graph_from_referee_to_referrer_.clear(); + ref_graph_from_referrer_to_referee_.clear(); + sorted_edges_.clear(); + visited_ = std::vector((*grammar)->NumRules(), false); + has_inward_edges_ = std::vector((*grammar)->complete_fsm.NumStates(), false); + for (int i = 0; i < grammar_->ImplPtr()->complete_fsm.NumStates(); i++) { + for (const auto& edge : grammar->ImplPtr()->complete_fsm.GetEdges(i)) { + has_inward_edges_[edge.target] = true; + } + } + + // Get the reference graph. + ref_graph_from_referee_to_referrer_ = RuleRefGraphFinder().Apply(*grammar); + ref_graph_from_referrer_to_referee_ = std::vector>((*grammar)->NumRules()); + for (int referee = 0; referee < static_cast(ref_graph_from_referee_to_referrer_.size()); + ++referee) { + for (int referer : ref_graph_from_referee_to_referrer_[referee]) { + ref_graph_from_referrer_to_referee_[referer].push_back(referee); + } + } + + // Sort the edges. + const auto& complete_fsm = grammar->ImplPtr()->complete_fsm; + sorted_edges_.reserve(complete_fsm.NumStates()); + for (int i = 0; i < complete_fsm.NumStates(); i++) { + const auto& edges = complete_fsm.GetEdges(i); + sorted_edges_.emplace_back(); + sorted_edges_.back().reserve(edges.size()); + for (const auto& edge : edges) { + sorted_edges_.back().emplace_back(edge); + } + std::sort(sorted_edges_.back().begin(), sorted_edges_.back().end()); + } + + // Disable non-fsms. + for (size_t i = 0; i < grammar->ImplPtr()->per_rule_fsms.size(); i++) { + if (!grammar->ImplPtr()->per_rule_fsms[i].has_value()) { + visited_[i] = true; + } + } + + // Hash the fsms which can be hashed: terminal fsms, or self-recursion fsms. The ready queue + // is seeded with all currently hashable fsms and maintained incrementally as fsms are hashed, + // so the whole loop is O(V + E) over the reference graph. When no fsm is ready, try to break + // a simple cycle in the reference graph and continue. + ready_queue_ = {}; + for (int i = 0; i < (*grammar)->NumRules(); i++) { + if (IsReadyToHash(i)) { + ready_queue_.push(i); + } + } + while (true) { + if (ready_queue_.empty()) { + // Try to find a simple cycle. We must ensure there are not self-recursion cycles. + if (!FindSimpleCycle()) { + break; + } + continue; + } + int32_t current_operating_index = ready_queue_.front(); + ready_queue_.pop(); + // Skip stale entries: an fsm may be pushed multiple times before it is processed. + if (!IsReadyToHash(current_operating_index)) { + continue; + } + visited_[current_operating_index] = true; + grammar->ImplPtr()->per_rule_fsm_hashes[current_operating_index] = + HashFsm(current_operating_index); + RemoveHashedFsmFromRefGraph(current_operating_index); + } + + // Try to hash the remaining fsms: they must contain something can't be hashed, like repetition. + // We can do this: if the fsm's start state has no inward edges, and all the ref edges are hashed + // except the edges at the start state, we can hash it. + std::vector> partial_hashed_list; + for (int i = 0; i < (*grammar)->NumRules(); i++) { + if (grammar->ImplPtr()->per_rule_fsm_hashes[i].has_value()) { + continue; + } + if (!grammar->ImplPtr()->per_rule_fsms[i].has_value()) { + continue; + } + if (has_inward_edges_[grammar->ImplPtr()->per_rule_fsms[i]->GetFsm().GetStart()]) { + continue; + } + const auto& [can_be_hashed, hash_value] = IsPartialHashable(i); + if (can_be_hashed) { + partial_hashed_list.emplace_back(i, hash_value); + } + } + for (const auto& [rule_id, hash_value] : partial_hashed_list) { + grammar->ImplPtr()->per_rule_fsm_hashes[rule_id] = hash_value; + } +} + +std::pair GrammarFSMHasherImpl::IsPartialHashable(int fsm_index) { + uint64_t hash_result = 0; + XGRAMMAR_DCHECK(fsm_index >= 0 && fsm_index < (*grammar_)->NumRules()) + << "Invalid fsm index: " << fsm_index << " num_rules: " << (*grammar_)->NumRules(); + XGRAMMAR_DCHECK(grammar_->ImplPtr()->per_rule_fsms[fsm_index].has_value()); + const auto& fsm = grammar_->ImplPtr()->per_rule_fsms[fsm_index].value().GetFsm(); + std::map original_state_id_to_new_id; + original_state_id_to_new_id[fsm.GetStart()] = 0; + std::queue bfs_queue; + std::set> hash_and_target; + bfs_queue.push(fsm.GetStart()); + // Perform a bfs to hash all the edges. + while (!bfs_queue.empty()) { + int current_old_state_id = bfs_queue.front(); + bool is_start = current_old_state_id == fsm.GetStart(); + int current_new_state_id = original_state_id_to_new_id[current_old_state_id]; + bfs_queue.pop(); + + // Check if the current state is an end state. + if (fsm.IsEndState(current_old_state_id)) { + hash_result = HashCombine( + hash_result, current_new_state_id, kEndStateFlag, kEndStateFlag, current_new_state_id + ); + } else { + hash_result = HashCombine( + hash_result, + current_new_state_id, + kNotEndStateFlag, + kNotEndStateFlag, + current_new_state_id + ); + } + + // Hash the edges. + + // First, check the edges which are rule references (including repeat refs). + // To keep consistent, we need to sort them with hashes. + int32_t unhashed_rules_count = 0; + auto hash_rule_like_edge = [&](int32_t ref_rule_id, int32_t target) { + if (ref_rule_id == fsm_index) { + hash_and_target.insert({kSelfRecursionFlag, target}); + return true; + } + if (!grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()) { + if (!is_start) { + return false; + } else { + unhashed_rules_count++; + if (unhashed_rules_count > 1) { + return false; + } + hash_and_target.insert({kUnKnownFlag, target}); + } + return true; + } + hash_and_target.insert({grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(), target} + ); + return true; + }; + + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (edge.IsRuleRef()) { + if (!hash_rule_like_edge(edge.GetRefRuleId(), edge.target)) { + return {false, 0}; + } + } else if (edge.IsRepeatRef()) { + auto info = grammar_->ImplPtr()->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + if (!hash_rule_like_edge(info.RuleId(), edge.target)) { + return {false, 0}; + } + } + } + + // Hash them. + for (const auto& [hash, target] : hash_and_target) { + if (original_state_id_to_new_id.find(target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(target); + } + int32_t target_new_id = original_state_id_to_new_id[target]; + hash_result = HashCombine(hash_result, current_new_state_id, hash, target_new_id); + } + + // Then, check the edges which are not rule/repeat references. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (original_state_id_to_new_id.find(edge.target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[edge.target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(edge.target); + } + int32_t target_new_id = original_state_id_to_new_id[edge.target]; + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + continue; + } + hash_result = HashCombine( + hash_result, + current_new_state_id, + static_cast(edge.min), + static_cast(edge.max), + target_new_id + ); + } + } + std::vector> new_id_mapping; + new_id_mapping.reserve(original_state_id_to_new_id.size()); + for (const auto& [original_state_id, new_state_id] : original_state_id_to_new_id) { + new_id_mapping.emplace_back(original_state_id, new_state_id); + } + grammar_->ImplPtr()->per_rule_fsm_new_state_ids[fsm_index] = new_id_mapping; + return {true, hash_result}; +} + +uint64_t GrammarFSMHasherImpl::HashFsm(int fsm_index) { + uint64_t hash_result = 0; + XGRAMMAR_DCHECK(fsm_index >= 0 && fsm_index < (*grammar_)->NumRules()) + << "Invalid fsm index: " << fsm_index << " num_rules: " << (*grammar_)->NumRules(); + XGRAMMAR_DCHECK(grammar_->ImplPtr()->per_rule_fsms[fsm_index].has_value()); + const auto& fsm = grammar_->ImplPtr()->per_rule_fsms[fsm_index].value().GetFsm(); + std::map original_state_id_to_new_id; + original_state_id_to_new_id[fsm.GetStart()] = 0; + std::queue bfs_queue; + std::set> hash_and_target; + bfs_queue.push(fsm.GetStart()); + + // Perform a bfs to hash all the edges. + while (!bfs_queue.empty()) { + int current_old_state_id = bfs_queue.front(); + int current_new_state_id = original_state_id_to_new_id[current_old_state_id]; + bfs_queue.pop(); + + // Check if the current state is an end state. + if (fsm.IsEndState(current_old_state_id)) { + hash_result = HashCombine( + hash_result, current_new_state_id, kEndStateFlag, kEndStateFlag, current_new_state_id + ); + } else { + hash_result = HashCombine( + hash_result, + current_new_state_id, + kNotEndStateFlag, + kNotEndStateFlag, + current_new_state_id + ); + } + + // Hash the edges. + + // First, check the edges which are rule references (including repeat refs). + // To keep consistent, we need to sort them with hashes. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (edge.IsRuleRef()) { + int32_t ref_rule_id = edge.GetRefRuleId(); + if (ref_rule_id == fsm_index) { + hash_and_target.insert({kSelfRecursionFlag, edge.target}); + } else { + XGRAMMAR_CHECK(grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()); + hash_and_target.insert( + {grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(), edge.target} + ); + } + } else if (edge.IsRepeatRef()) { + auto info = grammar_->ImplPtr()->complete_fsm.GetRepeatEdgeInfo(edge.GetAuxIndex()); + int32_t ref_rule_id = info.RuleId(); + if (ref_rule_id == fsm_index) { + uint64_t base_hash = kSelfRecursionFlag; + uint64_t repeat_hash = HashCombine(base_hash, info.Lower(), info.Upper()); + hash_and_target.insert({repeat_hash, edge.target}); + } else { + XGRAMMAR_CHECK(grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].has_value()); + uint64_t base_hash = grammar_->ImplPtr()->per_rule_fsm_hashes[ref_rule_id].value(); + uint64_t repeat_hash = HashCombine(base_hash, info.Lower(), info.Upper()); + hash_and_target.insert({static_cast(repeat_hash), edge.target}); + } + } + } + + // Hash them. + for (const auto& [hash, target] : hash_and_target) { + if (original_state_id_to_new_id.find(target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(target); + } + int32_t target_new_id = original_state_id_to_new_id[target]; + hash_result = HashCombine(hash_result, current_new_state_id, hash, target_new_id); + } + + // Then, check the edges which are not rule/repeat references. + for (const auto& edge : sorted_edges_[current_old_state_id]) { + if (original_state_id_to_new_id.find(edge.target) == original_state_id_to_new_id.end()) { + original_state_id_to_new_id[edge.target] = + static_cast(original_state_id_to_new_id.size()); + bfs_queue.push(edge.target); + } + int32_t target_new_id = original_state_id_to_new_id[edge.target]; + if (edge.IsRuleRef() || edge.IsRepeatRef()) { + continue; + } + hash_result = HashCombine( + hash_result, + current_new_state_id, + static_cast(edge.min), + static_cast(edge.max), + target_new_id + ); + } + } + std::vector> new_id_mapping; + new_id_mapping.reserve(original_state_id_to_new_id.size()); + for (const auto& [original_state_id, new_state_id] : original_state_id_to_new_id) { + new_id_mapping.emplace_back(original_state_id, new_state_id); + } + grammar_->ImplPtr()->per_rule_fsm_new_state_ids[fsm_index] = new_id_mapping; + return hash_result; +} + +std::optional GrammarFSMHasherImpl::HashSequence( + const Grammar& grammar, int32_t sequence_id +) { + using GrammarExprType = Grammar::Impl::GrammarExprType; + if (sequence_id == -1) { + return std::nullopt; + } + uint64_t hash_result = 0; + const auto& sequence_expr = grammar->GetGrammarExpr(sequence_id); + XGRAMMAR_DCHECK(sequence_expr.type == GrammarExprType::kSequence) + << "GrammarExpr is not a sequence"; + for (const auto& expr_id : sequence_expr) { + const auto& expr = grammar->GetGrammarExpr(expr_id); + hash_result = HashCombine(hash_result, static_cast(expr.type)); + switch (expr.type) { + case (GrammarExprType::kByteString): + case (GrammarExprType::kCharacterClass): + case (GrammarExprType::kCharacterClassStar): + case (GrammarExprType::kEmptyStr): { + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + case (GrammarExprType::kRuleRef): { + if (grammar->per_rule_fsm_hashes[expr[0]].has_value()) { + hash_result = HashCombine(hash_result, grammar->per_rule_fsm_hashes[expr[0]].value()); + } else { + return std::nullopt; + } + break; + } + case (GrammarExprType::kRepeat): { + if (grammar->per_rule_fsm_hashes[expr[0]].has_value()) { + hash_result = HashCombine(hash_result, grammar->per_rule_fsm_hashes[expr[0]].value()); + } else { + return std::nullopt; + } + hash_result = HashCombine(hash_result, expr[1]); + hash_result = HashCombine(hash_result, expr[2]); + break; + } + case (GrammarExprType::kSequence): + case (GrammarExprType::kChoices): { + return std::nullopt; + } + case (GrammarExprType::kTagDispatch): + case (GrammarExprType::kTokenTagDispatch): { + return std::nullopt; + } + case (GrammarExprType::kRegex): + case (GrammarExprType::kSubstring): { + // Hash the content, like a byte string. + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + case (GrammarExprType::kToken): + case (GrammarExprType::kExcludeToken): { + for (const auto& element : expr) { + hash_result = HashCombine(hash_result, element); + } + break; + } + } + } + return hash_result; +} + +class RuleLevelCache::Impl { + public: + using NodeKey = std::tuple< + uint64_t /*The hash value of the FSM*/, + int32_t /* The normalized node id*/, + int32_t /*The number of states*/, + int32_t /* The number of edges*/>; + using NodeType = std::pair; + + explicit Impl(size_t max_cache_memory_size) : max_cache_memory_size_(max_cache_memory_size) {} + + std::optional GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt + ); + + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask + ); + + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask + ); + + void ClearCache(); + + friend size_t MemorySize(const Impl* impl) { + int64_t total = 0; + for (const auto& shard : impl->shards_) { + total += shard.current_cache_memory_size; + } + return total; + } + + size_t GetMaxSize() const { return max_cache_memory_size_; } + + private: + /*! + * \brief The cache is sharded to reduce lock contention: the token mask cache generation + * queries and inserts from all compilation threads, and a single global mutex would serialize + * them (large grammars issue millions of cache operations). + */ + static constexpr size_t kNumShards = 16; + + struct Shard { + std::mutex mutex; + int64_t current_cache_memory_size = 0; + // The cache map: (fsm_hash, node_id, ...) -> index in cache_list + List cache_list; + std::unordered_map cache; + }; + + Shard& GetShard(const NodeKey& key) { + return shards_[HashCombine(std::get<0>(key), std::get<1>(key)) % kNumShards]; + } + + /*! \brief The memory budget of one shard. Eviction is performed per shard. */ + size_t ShardMaxSize() const { + return max_cache_memory_size_ == kUnlimitedSize ? kUnlimitedSize + : max_cache_memory_size_ / kNumShards; + } + + const size_t max_cache_memory_size_; + std::array shards_; +}; + +std::optional RuleLevelCache::GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt +) { + return pimpl_->GetCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); +} + +bool RuleLevelCache::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask +) { + return pimpl_->AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, token_mask); +} + +bool RuleLevelCache::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask +) { + return pimpl_->AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, std::move(token_mask)); +} + +void RuleLevelCache::ClearCache() { pimpl_->ClearCache(); } + +size_t RuleLevelCache::GetMaxSize() const { return pimpl_->GetMaxSize(); } + +std::optional RuleLevelCache::Impl::GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt +) { + // Find in the cache. + NodeKey key = std::make_tuple(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); + Shard& shard = GetShard(key); + std::lock_guard lock(shard.mutex); + auto it = shard.cache.find(key); + if (it == shard.cache.end()) { + return std::nullopt; + } + + // Move the node to the back of the list. + shard.cache_list.MoveBack(it->second); + return List::iterator(it->second, shard.cache_list)->second; +} + +bool RuleLevelCache::Impl::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask +) { + return AddCache(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt, AdaptiveTokenMask(token_mask)); +} + +bool RuleLevelCache::Impl::AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask +) { + // Check if we can add to the cache. + NodeKey key = std::make_tuple(fsm_hash, fsm_new_node_id, state_cnt, edge_cnt); + Shard& shard = GetShard(key); + const size_t shard_max_size = ShardMaxSize(); + std::lock_guard lock(shard.mutex); + if (shard_max_size != kUnlimitedSize && MemorySize(token_mask) > shard_max_size) { + // The token mask is too large to be cached. + return false; + } + if (shard.cache.find(key) != shard.cache.end()) { + // Already exists. + return false; + } + + // Evict old entries if needed. + if (shard_max_size != kUnlimitedSize) { + size_t new_item_size = MemorySize(token_mask); + while ((shard.current_cache_memory_size) > static_cast(shard_max_size - new_item_size) + ) { + auto oldest_it = shard.cache_list.begin(); + if (oldest_it == shard.cache_list.end()) { + // This should not happen if the size of the new item is smaller than + // the shard budget, but this is a safeguard. + break; + } + shard.current_cache_memory_size -= MemorySize(oldest_it->second); + shard.cache.erase(oldest_it->first); + shard.cache_list.Erase(oldest_it); + } + } + + // Add to the cache. + auto new_it = shard.cache_list.PushBack(NodeType(key, std::move(token_mask))); + shard.current_cache_memory_size += MemorySize(new_it->second); + shard.cache[key] = new_it.Index(); + return true; +} + +RuleLevelCache::RuleLevelCache(size_t max_cache_memory_size) + : pimpl_(std::make_shared(max_cache_memory_size)) {} + +void RuleLevelCache::Impl::ClearCache() { + for (auto& shard : shards_) { + std::lock_guard lock(shard.mutex); + shard.cache_list.Clear(); + shard.cache.clear(); + shard.current_cache_memory_size = 0; + } +} + +size_t MemorySize(const RuleLevelCache& manager) { return MemorySize(manager.ImplPtr()); } + +/*************************** Forward grammar constructors to their impl ***************************/ + +Grammar GrammarUnionFunctor::Apply(const std::vector& grammars) { + return GrammarUnionFunctorImpl().Apply(grammars); +} + +Grammar GrammarConcatFunctor::Apply(const std::vector& grammars) { + return GrammarConcatFunctorImpl().Apply(grammars); +} + +int32_t SubGrammarAdder::Apply(GrammarBuilder* builder, const Grammar& sub_grammar) { + return SubGrammarAdderImpl().ApplyWithBuilder(builder, sub_grammar); +} + +/*************************** Forward grammar Normalizers to their impl ***************************/ + +Grammar GrammarNormalizer::Apply(const Grammar& grammar) { + return GrammarNormalizerImpl().Apply(grammar); +} + +Grammar StructureNormalizer::Apply(const Grammar& grammar) { + return StructureNormalizerImpl().Apply(grammar); +} + +/*************************** Forward grammar optimizers to their impl ***************************/ + +void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl::Apply(grammar); } + +void RepetitionNormalizer::Apply(Grammar* grammar) { RepetitionNormalizerImpl().Apply(grammar); } + +void GrammarFSMHasher::Apply(Grammar* grammar) { GrammarFSMHasherImpl().Apply(grammar); } + +std::optional GrammarFSMHasher::HashSequence( + const Grammar& grammar, int32_t sequence_id +) { + return GrammarFSMHasherImpl().HashSequence(grammar, sequence_id); +} + +FSMWithStartEnd GrammarFSMBuilder::RuleRef(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::RuleRef(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::CharacterClass(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::CharacterClass(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::ByteString(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::ByteString(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::Token(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::Token(expr); +} + +FSMWithStartEnd GrammarFSMBuilder::ExcludeToken(const GrammarExpr& expr) { + return GrammarFSMBuilderImpl::ExcludeToken(expr); +} + +std::optional GrammarFSMBuilder::TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& ttd +) { + return GrammarFSMBuilderImpl::TokenTagDispatch(ttd); +} + +std::optional GrammarFSMBuilder::Sequence( + const GrammarExpr& expr, const Grammar& grammar +) { + return GrammarFSMBuilderImpl::Sequence(expr, grammar); +} + +std::optional GrammarFSMBuilder::Choices( + const GrammarExpr& expr, const Grammar& grammar +) { + return GrammarFSMBuilderImpl::Choices(expr, grammar); +} + +Result GrammarFSMBuilder::Regex(const std::string& regex, bool json_string) { + return GrammarFSMBuilderImpl::Regex(regex, json_string); +} + +Result GrammarFSMBuilder::FromRegularGrammar( + const Grammar& grammar, int max_num_states +) { + return RegularGrammarFSMBuilder(grammar, max_num_states).Build(); +} + +const std::bitset<256>& GrammarFSMBuilder::JSONStringForbiddenChars() { + static const std::bitset<256> forbidden_chars = [] { + std::bitset<256> chars; + for (int c = 0x00; c <= 0x1F; ++c) { + chars.set(c); + } + chars.set('"'); + chars.set('\\'); + return chars; + }(); + return forbidden_chars; +} + +std::optional GrammarFSMBuilder::TagDispatch( + const Grammar::Impl::TagDispatch& tag_dispatch +) { + return GrammarFSMBuilderImpl::TagDispatch(tag_dispatch); +} + +std::vector AllowEmptyRuleAnalyzer::Apply(const Grammar& grammar) { + return AllowEmptyRuleAnalyzerImpl().Apply(grammar); +} + +void RuleInliner::Apply(Grammar* grammar) { RuleInlinerImpl().Apply(grammar); } + +Grammar DeadCodeEliminator::Apply(const Grammar& grammar) { + return DeadCodeEliminatorImpl().Apply(grammar); +} + +Grammar LookaheadAssertionAnalyzer::Apply(const Grammar& grammar) { + return LookaheadAssertionAnalyzerImpl().Apply(grammar); +} + +Grammar RepetitionRangeExpander::Apply(const Grammar& grammar) { + return RepetitionRangeExpanderImpl().Apply(grammar); +} + +Grammar GrammarOptimizer::Apply(const Grammar& grammar) { + return GrammarOptimizerImpl::Apply(grammar); +} + +void ByteStringFuser::Apply(Grammar* grammar) { ByteStringFuserImpl().Apply(grammar); } + +Grammar RootRuleRenamer::Apply(const Grammar& grammar) { + return RootRuleRenamerImpl().Apply(grammar); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_functor.h b/zig/vendor/xgrammar/cpp/grammar_functor.h new file mode 100644 index 0000000000..e7871cab9b --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_functor.h @@ -0,0 +1,520 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_functor.h + * \brief The header for the simplification of the BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_FUNCTOR_H_ +#define XGRAMMAR_GRAMMAR_FUNCTOR_H_ + +#include + +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "support/utils.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief Base class for visitors and mutators of the BNF grammar. + * \tparam T The type of the return value of visitor functions. Typical values: + * - int32_t: the id of the new grammar_expr + * - void: no return value + * \tparam ReturnType The type of the return value of the transform function Apply(). Typical values + * are void (for visitor) and Grammar (for mutator). + */ +template +class GrammarFunctor { + public: + /*! + * \brief Constructor. + * \param grammar The grammar to visit or mutate. + */ + explicit GrammarFunctor() {} + + /*! + * \brief Apply the transformation to the grammar, or visit the grammar. + * \return The transformed grammar, or the visiting result, or void. + */ + virtual ReturnType Apply(const Grammar& grammar) { + // The initializer MUST be called at first when overriding the Apply() function. + InitGrammar(grammar); + if constexpr (std::is_same::value) { + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + VisitExpr(rule.body_expr_id); + VisitLookaheadAssertion(rule.lookahead_assertion_id); + } + return ReturnType(); + } else if constexpr (std::is_same::value && + std::is_same::value) { + InitBuilder(); + // First add empty rules to ensure the new rule ids the same as the old ones, then update + // the rule bodies + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + builder_->AddEmptyRule(base_grammar_->GetRule(i).name); + } + for (int i = 0; i < static_cast(base_grammar_->NumRules()); ++i) { + auto rule = base_grammar_->GetRule(i); + cur_rule_name_ = rule.name; + auto new_body_expr_id = VisitExpr(rule.body_expr_id); + builder_->UpdateRuleBody(i, new_body_expr_id); + // Handle lookahead assertion + builder_->UpdateLookaheadAssertion(i, VisitLookaheadAssertion(rule.lookahead_assertion_id)); + builder_->UpdateMaxTokens(i, rule.max_tokens); + builder_->UpdateMaxChars(i, rule.max_chars); + builder_->UpdateCaptureName(i, rule.capture_name); + if (const auto* suffix_stop_info = base_grammar_->GetSuffixStopInfo(i)) { + builder_->UpdateSuffixStopInfo(i, *suffix_stop_info); + } + builder_->UpdateLazy(i, rule.is_lazy); + builder_->UpdateRuleTemperature(i, rule.temperature); + } + return builder_->Get(base_grammar_->GetRootRule().name); + } else { + return ReturnType(); + } + } + + /*! \brief Virtual destructor. */ + virtual ~GrammarFunctor() = default; + + protected: + using Rule = Grammar::Impl::Rule; + using GrammarExpr = Grammar::Impl::GrammarExpr; + using GrammarExprType = Grammar::Impl::GrammarExprType; + + /*! \brief Initialize the functor. Should be called at the beginning of Apply(). */ + virtual void InitGrammar() {} + + virtual void InitGrammar(const Grammar& grammar) { base_grammar_ = grammar; } + + virtual void InitBuilder() { + owned_builder_ = GrammarBuilder(); + builder_ = &owned_builder_; + } + + virtual void InitBuilder(const Grammar& grammar) { + owned_builder_ = GrammarBuilder(grammar); + builder_ = &owned_builder_; + } + + virtual void InitBuilder(GrammarBuilder* builder) { builder_ = builder; } + + /*! \brief Visit a lookahead assertion expr referred by id. */ + virtual T VisitLookaheadAssertion(int32_t lookahead_assertion_id) { + if (lookahead_assertion_id == -1) { + if constexpr (std::is_same::value) { + return -1; + } else { + return T(); + } + } + return VisitExpr(lookahead_assertion_id); + } + + /*! \brief Visit a GrammarExpr by id. */ + virtual T VisitExpr(int32_t old_grammar_expr_id) { + return VisitExpr(base_grammar_->GetGrammarExpr(old_grammar_expr_id)); + } + + /*! \brief Visit a GrammarExpr. Dispatch to the corresponding Visit function. */ + virtual T VisitExpr(const GrammarExpr& grammar_expr) { + switch (grammar_expr.type) { + case GrammarExprType::kSequence: + return VisitSequence(grammar_expr); + case GrammarExprType::kChoices: + return VisitChoices(grammar_expr); + case GrammarExprType::kEmptyStr: + return VisitEmptyStr(grammar_expr); + case GrammarExprType::kByteString: + return VisitByteString(grammar_expr); + case GrammarExprType::kCharacterClass: + return VisitCharacterClass(grammar_expr); + case GrammarExprType::kCharacterClassStar: + return VisitCharacterClassStar(grammar_expr); + case GrammarExprType::kRuleRef: + return VisitRuleRef(grammar_expr); + case GrammarExprType::kTagDispatch: + return VisitTagDispatch(grammar_expr); + case GrammarExprType::kRepeat: + return VisitRepeat(grammar_expr); + case GrammarExprType::kToken: + return VisitToken(grammar_expr); + case GrammarExprType::kExcludeToken: + return VisitExcludeToken(grammar_expr); + case GrammarExprType::kTokenTagDispatch: + return VisitTokenTagDispatch(grammar_expr); + case GrammarExprType::kRegex: + return VisitRegex(grammar_expr); + case GrammarExprType::kSubstring: + return VisitSubstring(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected sequence type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } + } + + /*! \brief Visit a choices GrammarExpr. */ + virtual T VisitChoices(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + for (auto i : grammar_expr) { + VisitExpr(i); + } + } else if constexpr (std::is_same::value) { + std::vector choice_ids; + for (int32_t i : grammar_expr) { + choice_ids.push_back(VisitExpr(i)); + } + return builder_->AddChoices(choice_ids); + } else { + return T(); + } + } + + /*! \brief Visit a sequence GrammarExpr. */ + virtual T VisitSequence(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + for (auto i : grammar_expr) { + VisitExpr(i); + } + } else if constexpr (std::is_same::value) { + std::vector sequence_ids; + for (int32_t i : grammar_expr) { + sequence_ids.push_back(VisitExpr(i)); + } + return builder_->AddSequence(sequence_ids); + } else { + return T(); + } + } + + virtual T VisitTagDispatch(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + return; + } else if constexpr (std::is_same::value) { + Grammar::Impl::TagDispatch tag_dispatch = base_grammar_->GetTagDispatch(grammar_expr); + return builder_->AddTagDispatch(tag_dispatch); + } else { + return T(); + } + } + + /*! \brief Visit an element GrammarExpr, including empty string, character class, and rule ref. */ + virtual T VisitElement(const GrammarExpr& grammar_expr) { + if constexpr (std::is_same::value) { + return; + } else if constexpr (std::is_same::value) { + return builder_->AddGrammarExpr(grammar_expr); + } else { + return T(); + } + } + + /*! \brief Visit an empty string GrammarExpr. */ + virtual T VisitEmptyStr(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a character class GrammarExpr. */ + virtual T VisitByteString(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a character class GrammarExpr. */ + virtual T VisitCharacterClass(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a star quantifier GrammarExpr. */ + virtual T VisitCharacterClassStar(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a rule reference GrammarExpr. */ + virtual T VisitRuleRef(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a repeat GrammarExpr. */ + virtual T VisitRepeat(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + virtual T VisitToken(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + virtual T VisitExcludeToken(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + virtual T VisitTokenTagDispatch(const GrammarExpr& grammar_expr) { + return VisitElement(grammar_expr); + } + + /*! \brief Visit a regex GrammarExpr. It is a leaf: the pattern string is carried as-is. */ + virtual T VisitRegex(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief Visit a substring GrammarExpr. It is a leaf: the chunk list is carried as-is. */ + virtual T VisitSubstring(const GrammarExpr& grammar_expr) { return VisitElement(grammar_expr); } + + /*! \brief The grammar to visit or mutate. */ + Grammar base_grammar_{NullObj{}}; + + /*! + * \brief The builder to build the new grammar. It is empty when the mutator is constructed, and + * can be used to build a new grammar in subclasses. + */ + GrammarBuilder* builder_ = nullptr; + + GrammarBuilder owned_builder_; + + /*! \brief The name of the current rule being visited. */ + std::string cur_rule_name_; +}; + +/*! + * \brief Visitor of Grammar. + * \tparam ReturnType The return type of the Apply() function. Denotes the collected information. + */ +template +using GrammarVisitor = GrammarFunctor; + +/*! + * \brief Mutator of Grammar. The Apply() function returns the updated grammar. + */ +using GrammarMutator = GrammarFunctor; + +/****** All below methods are implemented as functor to hide the implementation ******/ + +/*************************** Grammar Constructor ***************************/ +/*! + * \brief Find the union of multiple grammars as a new grammar. + */ +class GrammarUnionFunctor { + public: + static Grammar Apply(const std::vector& grammars); +}; + +/*! + * \brief Find the concatenation of multiple grammars as a new grammar. + */ +class GrammarConcatFunctor { + public: + static Grammar Apply(const std::vector& grammars); +}; + +/*! + * \brief Add a sub grammar to the current builder. The return value + * of Apply is the new rule id of the sub grammar's root rule. + */ +class SubGrammarAdder { + public: + static int32_t Apply(GrammarBuilder* builder, const Grammar& sub_grammar); +}; + +/*************************** Grammar Normalizer ***************************/ + +/*! + * \brief Normalize a Grammar: expand the nested rules, combine consequent sequences and strings, + * etc. + */ +class GrammarNormalizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Normalize the structure of the grammar. It will ensure each rule is a choices of + * sequences of elements, or a tag dispatch. The expanded context will be a sequence of elements. + */ +class StructureNormalizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*************************** Grammar Optimizer ***************************/ + +/*! + * \brief Fuse adjacent byte string elements in sequences. + * \details Rewrites *grammar in place. Only sequences that actually contain adjacent byte strings + * are rewritten; if nothing needs fusing, the grammar is left untouched. The caller must own the + * grammar, as the input is mutated directly. + */ +class ByteStringFuser { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Analyze the grammar to find the rules that are allowed to be empty. + */ +class AllowEmptyRuleAnalyzer { + public: + static std::vector Apply(const Grammar& grammar); +}; + +/*! + * \brief Inline the rule references in the grammar. + * \details Rewrites *grammar in place. Only choices with an inlinable leading rule reference are + * rewritten; if nothing can be inlined, the grammar is left untouched. The caller must own the + * grammar, as the input is mutated directly. + */ +class RuleInliner { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Eliminate the not referenced rules in the grammar. + */ +class DeadCodeEliminator { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Analyze and add lookahead assertions in the grammar. + */ +class LookaheadAssertionAnalyzer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Build the FSMs of the grammar. + */ +class GrammarFSMBuilder { + using GrammarExpr = Grammar::Impl::GrammarExpr; + + public: + static void Apply(Grammar* grammar); + static FSMWithStartEnd RuleRef(const GrammarExpr& expr); + static FSMWithStartEnd CharacterClass(const GrammarExpr& expr); + static FSMWithStartEnd ByteString(const GrammarExpr& expr); + static FSMWithStartEnd Token(const GrammarExpr& expr); + static FSMWithStartEnd ExcludeToken(const GrammarExpr& expr); + static std::optional TokenTagDispatch( + const Grammar::Impl::TokenTagDispatch& token_tag_dispatch + ); + static std::optional Sequence(const GrammarExpr& expr, const Grammar& grammar); + static std::optional Choices(const GrammarExpr& expr, const Grammar& grammar); + static std::optional TagDispatch(const Grammar::Impl::TagDispatch& tag_dispatch); + /*! + * \brief Build the automaton of a regex pattern. Returns the error message on failure. + * \param regex The regex pattern string. + * \param json_string Whether the regex matches the body of a JSON string literal. If true, + * the characters in JSONStringForbiddenChars() are excluded from every character match. + */ + static Result Regex(const std::string& regex, bool json_string = false); + /*! + * \brief Flatten a regular EBNF grammar without reinterpreting its terminal expressions. + * Supports finite rule calls, direct tail self-recursion, and repetition. Each call gets + * separate return states. Other recursion, non-byte terminals, and rule annotations return + * an error. Intended for the existing RegexToEBNF fallback, not arbitrary CFGs. + * \param max_num_states Bound on the intermediate automaton size. + */ + static Result FromRegularGrammar( + const Grammar& grammar, int max_num_states = 100000 + ); + /*! \brief The characters that must be escaped inside a JSON string literal: the control + * characters 0x00-0x1F, the quote '"' and the backslash '\\'. */ + static const std::bitset<256>& JSONStringForbiddenChars(); +}; + +/*! + * \brief Normalize the repetition expression. If the context of + * repetition expression is nullable, then the repetition range will be + * normalized from {m, n} to {0, n} to reduce uncertainty. + */ +class RepetitionNormalizer { + public: + static void Apply(Grammar* grammar); +}; + +/*! + * \brief Expand kRepeat grammar expressions using HandleRepetitionRange logic. + * Transforms repetition structures into explicit sequences and choices. + */ +class RepetitionRangeExpander { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Optimize the grammar when compiling. + * \note No matter whether the grammar is optimized, grammar optimizer will + * return a new grammar. The following optimization will be applied: + * 1. Byte fuser. + * 2. Rule inliner. + * 3. Dead code eliminator. + * 4. Lookahead assertion analyzer. + * 5. Allow-empty rule analyzer. + * 6. Repetition normalizer. + * 7. FSM builder. + */ +class GrammarOptimizer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Rename the root rule of the grammar to "root". + */ +class RootRuleRenamer { + public: + static Grammar Apply(const Grammar& grammar); +}; + +/*! + * \brief Hash the fsms in the grammar, + * and get the new state ids of each fsm's states. + */ +class GrammarFSMHasher { + public: + static void Apply(Grammar* grammar); + static std::optional HashSequence(const Grammar& grammar, int32_t sequence_id); +}; + +/*! + * \brief Store the crossing cache for different grammars. + * \param max_cache_size The maximum size of the cache numbers. + * \details LRU algorithm is implemented. + */ +class RuleLevelCache { + public: + static const size_t kUnlimitedSize = static_cast(-1); + + std::optional GetCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt + ); + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + const AdaptiveTokenMask& token_mask + ); + bool AddCache( + const uint64_t& fsm_hash, + int32_t fsm_new_node_id, + const int32_t& state_cnt, + const int32_t edge_cnt, + AdaptiveTokenMask&& token_mask + ); + RuleLevelCache(size_t max_cache_memory_size = kUnlimitedSize); + + void ClearCache(); + + size_t GetMaxSize() const; + + friend size_t MemorySize(const RuleLevelCache& manager); + + XGRAMMAR_DEFINE_PIMPL_METHODS(RuleLevelCache); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_FUNCTOR_H_ diff --git a/zig/vendor/xgrammar/cpp/grammar_impl.h b/zig/vendor/xgrammar/cpp/grammar_impl.h new file mode 100644 index 0000000000..7b36a954cd --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_impl.h @@ -0,0 +1,487 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.h + * \brief The header for the support of grammar-guided generation. + */ + +#ifndef XGRAMMAR_GRAMMAR_IMPL_H_ +#define XGRAMMAR_GRAMMAR_IMPL_H_ + +#include + +#include +#include +#include +#include + +#include "fsm.h" +#include "support/logging.h" +#include "support/reflection.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +/*! + * \brief This class stores the abstract syntax tree (AST) of the Backus-Naur Form (BNF) grammar. + * The BNF definition here is standard BNF, and the characters are represented using regex-style + * character classes (e.g. [a-z], [^a-z]). + * + * \details + * ### Rules + * The BNF grammar AST consists of a set of rules. Each rule contains a name and a definition, and + * corresponds to a production in the grammar. The definition of a rule is a GrammarExpr. Each rule + * has a rule_id for reference. + * + * ### GrammarExprs + * GrammarExpr is the definition of a rule or part of the definition of a rule. It can contain + * elements, empty string, reference to other GrammarExprs, or reference to other rules. Each + * GrammarExpr corresponds to a grammar_expr_id for reference. + * + * For example, in the following rule: rule ::= ("a" "b") | "c" + * ("a" "b"), "c", ("a" "b") | "c" are all GrammarExprs. + * + * #### Types of GrammarExprs + * Every GrammarExpr is represented by a type as well as a variable-length array containing its + * data. GrammarExpr has several types: + * - Byte string: a string of bytes (0~255). Supports UTF-8 strings. + * - Character class: a range of characters (each character is a unicode codepoint), e.g. [a-z], + * [ac-z]. Can be negated: [^a-z], [^ac-z]. Now only ascii chars is allowed in [], but this + * expression can accept/reject unicode chars. + * - Character class star: a star quantifier of a character class. e.g. [a-z]*, [^a-z]*. + * - EmptyStr: an empty string, i.e. "" + * - Rule reference: a reference to another rule + * - Sequence: a sequence of grammar_exprs, e.g. ("a" "b"). These grammar_exprs are concatenated + * together. + * - Choices: a choice of grammar_exprs, e.g. ("a" "b") | "c". Each grammar_expr can be matched. + * + * #### Storage of GrammarExprs + * Each type of GrammarExpr has a different data format. For the format of each type of GrammarExpr, + * see docs in Grammar::Impl::GrammarExprType. + * + * We store all GrammarExprs in csr_matrix style. That is, they are stored consecutively in one + * vector (data vector) and the starting position of each GrammarExpr is recorded in the indptr + * vector. + * + * \remark The character class star GrammarExpr is for the special support for elements like [a-z]* + * in the grammar. We add it to make the matching more efficient, as we can avoid recursion into + * rules when matching a sequence of characters. It should be used like: + * rule1 ::= ((element1 element2 rule2 ...) | ...) + * rule2 ::= character_class_star_grammar_expr(id_of_a_character_class_grammar_expr) + */ +class Grammar::Impl { + public: + /*! \brief A rule with name. */ + struct Rule { + /*! \brief The name of the rule. */ + std::string name; + /*! \brief The GrammarExpr id of the body of the rule. */ + int32_t body_expr_id; + /*! \brief The id of the associated lookahead assertion expr. For now it must be a id of a + * sequence GrammarExpr. -1 if not exists. */ + int32_t lookahead_assertion_id = -1; + /*! \brief Whether the lookahead assertion is exact. */ + bool is_exact_lookahead = false; + /*! \brief The token budget of the rule. When non-negative, the matcher bounds each + * occurrence of this rule to at most max_tokens LLM tokens, forcing it to end at the + * earliest possible position once the budget is exhausted. -1 means no budget. */ + int32_t max_tokens = -1; + /*! \brief The Unicode codepoint budget of the rule. When non-negative, the matcher bounds + * each occurrence of this rule to at most max_chars codepoints, forcing it to end at the + * earliest possible position once the budget is exhausted. -1 means no budget. */ + int32_t max_chars = -1; + /*! \brief The capture group name of the rule. When non-empty, the matcher records the input + * span matched by this rule on every completion, retrievable via GrammarMatcher::GetCaptures. + * Empty means no capture. */ + std::string capture_name = {}; + /*! \brief Whether the rule matches with committed-shortest (lazy) semantics: at the first + * position where the body can complete, it must complete. */ + bool is_lazy = false; + /*! \brief The sampling temperature to use while matching this rule. */ + std::optional temperature = std::nullopt; + }; + + /*! \brief Sparse per-rule metadata used to materialize suffix and stop captures. */ + struct SuffixStopInfo { + /*! \brief The rule carrying this metadata. */ + int32_t rule_id = -1; + /*! \brief Trailing bytes hidden only from this rule's own capture. */ + int32_t hidden_suffix_bytes = 0; + /*! \brief Trailing bytes hidden from this rule and every enclosing capture. */ + int32_t hidden_stop_bytes = 0; + /*! \brief Helper rule matching the body before a variable-length marker. A self-reference + * marks a zero-width event immediately following a fixed dynamic-dispatch marker. */ + int32_t body_rule_id = -1; + /*! \brief Helper rule matching a variable-length marker. */ + int32_t marker_rule_id = -1; + /*! \brief Capture name for the marker bytes. */ + std::string stop_capture_name = {}; + + bool IsEmpty() const { + return hidden_suffix_bytes == 0 && hidden_stop_bytes == 0 && body_rule_id == -1 && + marker_rule_id == -1 && stop_capture_name.empty(); + } + }; + + /*! \brief Get the number of rules. */ + int32_t NumRules() const { return rules_.size(); } + /*! \brief Get the rule with the given id. */ + const Rule& GetRule(int32_t rule_id) const { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rules_.size())) + << "rule_id " << rule_id << " is out of bound"; + return rules_[rule_id]; + } + Rule& GetRule(int32_t rule_id) { + XGRAMMAR_DCHECK(rule_id >= 0 && rule_id < static_cast(rules_.size())) + << "rule_id " << rule_id << " is out of bound"; + return rules_[rule_id]; + } + /*! \brief Get sparse suffix/stop metadata for a rule, or nullptr when none exists. */ + const SuffixStopInfo* GetSuffixStopInfo(int32_t rule_id) const { + auto it = std::lower_bound( + suffix_stop_infos_.begin(), + suffix_stop_infos_.end(), + rule_id, + [](const SuffixStopInfo& info, int32_t id) { return info.rule_id < id; } + ); + return it != suffix_stop_infos_.end() && it->rule_id == rule_id ? &*it : nullptr; + } + /*! \brief Get the root rule id of the grammar. */ + int32_t GetRootRuleId() const { return root_rule_id_; } + /*! \brief Get the root rule of the grammar. */ + const Rule& GetRootRule() const { + XGRAMMAR_DCHECK(root_rule_id_ >= 0 && root_rule_id_ < static_cast(rules_.size())) + << "root_rule_id " << root_rule_id_ << " is out of bound"; + return rules_[root_rule_id_]; + } + + /*! + * \brief Check that every id and offset stored in the grammar is in range. Used after + * deserialization, where the fields are restored verbatim and none of the builder checks run. + * \return An error message if the grammar is malformed. + */ + std::optional Validate() const; + + /*! \brief The type of the grammar expr. */ + enum class GrammarExprType : int32_t { + // data format: [byte0, byte1, ...] + kByteString, + // data format: [is_negative, lower0, upper0, lower1, upper1, ...] + kCharacterClass, + kCharacterClassStar, + // data format: [] + kEmptyStr, + // data format: [rule_id] + kRuleRef, + // data format: [grammar_expr_id0, grammar_expr_id1, ...] + kSequence, + // data format: [grammar_expr_id0, grammar_expr_id1, ...] + kChoices, + // data format: [tag_expr0, rule_id0, tag_expr1, rule_id1, ..., loop_after_dispatch, + // excluded_str_expr_id] + kTagDispatch, + // data format: [rule_id, min_repeat_count, max_repeat_count] + kRepeat, + // data format: [token_id_0, token_id_1, ...] + kToken, + // data format: [token_id_0, token_id_1, ...] + kExcludeToken, + // data format: [trigger_cnt, (token_id, rule_id) × N, + // loop_after_dispatch, + // exclude_cnt, token_id × M] + kTokenTagDispatch, + // data format: [json_string, byte0, byte1, ...] + // The bytes are the regex pattern string. Like kTagDispatch, it can only be the body of a + // rule. The pattern is carried through the grammar passes as-is; when GrammarFSMBuilder + // runs, the pattern is compiled into an automaton, so every regex rule always has a + // per-rule FSM after optimization. If json_string is 1, the regex matches the body of a + // JSON string literal: the characters that must be escaped in a JSON string (the control + // characters, '"' and '\\') are excluded from every character match of the automaton. + kRegex, + // data format: [chunk0_len, byte0_0, byte0_1, ..., chunk1_len, byte1_0, ...] + // A list of length-prefixed byte string chunks. Matches every contiguous subsequence of + // the chunk list, including the empty one. Like kRegex, it can only be the body of a rule + // and is carried through the grammar passes as-is; when GrammarFSMBuilder runs, it is + // compiled into an automaton via a chunk-level suffix automaton (see SuffixAutomata). + kSubstring, + }; + + /*! \brief The object representing a grammar expr. */ + struct GrammarExpr { + /*! \brief The type of the grammar expr. */ + GrammarExprType type; + /*! \brief The data of the GrammarExpr. A variable-length array. */ + const int32_t* data; + /*! \brief The length of the data array. */ + int32_t data_len; + + int32_t size() const { return data_len; } + /*! \brief Get the i-th element of the data array. */ + const int32_t& operator[](int i) const { + XGRAMMAR_DCHECK(i >= 0 && i < static_cast(data_len)) + << "Index " << i << " is out of bound"; + return data[i]; + } + const int32_t* begin() const { return data; } + const int32_t* end() const { return data + data_len; } + void SetData(int index, int value) { const_cast(data)[index] = value; } + }; + + /*! \brief Get the number of grammar_exprs. */ + int32_t NumGrammarExprs() const { return grammar_expr_indptr_.size(); } + + /*! \brief Get the grammar_expr with the given id. */ + GrammarExpr GetGrammarExpr(int32_t grammar_expr_id) const { + XGRAMMAR_DCHECK( + grammar_expr_id >= 0 && grammar_expr_id < static_cast(grammar_expr_indptr_.size()) + ) << "grammar_expr_id " + << grammar_expr_id << " is out of bound"; + int start_index = grammar_expr_indptr_[grammar_expr_id]; + auto start_ptr = grammar_expr_data_.data() + start_index; + auto type = static_cast(start_ptr[0]); + auto data_ptr = start_ptr + 2; + auto data_len = start_ptr[1]; + return {type, data_ptr, data_len}; + } + + /******************* GrammarExpr Getters *******************/ + + /*! \brief Get the string of the byte string grammar expr. */ + std::string GetByteString(const GrammarExpr& grammar_expr) const { + std::string str; + str.reserve(grammar_expr.size()); + for (int i = 0; i < grammar_expr.size(); ++i) { + str.push_back(static_cast(static_cast(grammar_expr[i]))); + } + return str; + } + + /*! \brief Get the string of the byte string grammar expr. */ + std::string GetByteString(int32_t grammar_expr_id) const { + return GetByteString(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief Get the regex pattern string of the regex grammar expr. */ + std::string GetRegexString(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kRegex) << "GrammarExpr is not a regex"; + XGRAMMAR_DCHECK(grammar_expr.size() >= 1) << "Regex expr must contain the json_string flag"; + std::string str; + str.reserve(grammar_expr.size() - 1); + for (int i = 1; i < grammar_expr.size(); ++i) { + str.push_back(static_cast(static_cast(grammar_expr[i]))); + } + return str; + } + + /*! \brief Get the regex pattern string of the regex grammar expr with the given id. */ + std::string GetRegexString(int32_t grammar_expr_id) const { + return GetRegexString(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief Get whether the regex grammar expr matches the body of a JSON string literal. */ + bool GetRegexIsJSONString(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kRegex) << "GrammarExpr is not a regex"; + XGRAMMAR_DCHECK(grammar_expr.size() >= 1) << "Regex expr must contain the json_string flag"; + return grammar_expr[0] != 0; + } + + /*! \brief Get the chunk list of the substring grammar expr. */ + std::vector GetSubstringChunks(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kSubstring) + << "GrammarExpr is not a substring"; + std::vector chunks; + for (int i = 0; i < grammar_expr.size();) { + int32_t chunk_len = grammar_expr[i++]; + XGRAMMAR_DCHECK(chunk_len >= 0 && i + chunk_len <= grammar_expr.size()) + << "Invalid substring chunk length"; + std::string chunk; + chunk.reserve(chunk_len); + for (int32_t j = 0; j < chunk_len; ++j) { + chunk.push_back(static_cast(static_cast(grammar_expr[i++]))); + } + chunks.push_back(std::move(chunk)); + } + return chunks; + } + + /*! \brief Get the chunk list of the substring grammar expr with the given id. */ + std::vector GetSubstringChunks(int32_t grammar_expr_id) const { + return GetSubstringChunks(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief The object representing a tag dispatch. */ + struct TagDispatch { + /*! \brief The tag and rule id pairs. */ + std::vector> tag_rule_pairs; + /*! \brief If true, the tag dispatch will loop after dispatching. */ + bool loop_after_dispatch; + /*! \brief The strings that are excluded by the tag dispatch. */ + std::vector excludes; + static const int kTagDispatchExtraParameter = 2; + }; + + /*! \brief Get the tag dispatch from the grammar expr. */ + TagDispatch GetTagDispatch(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTagDispatch) + << "GrammarExpr is not a tag dispatch"; + + TagDispatch result; + XGRAMMAR_DCHECK(grammar_expr.size() >= TagDispatch::kTagDispatchExtraParameter); + result.tag_rule_pairs.reserve( + (grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter) / 2 + ); + + for (int i = 0; i < grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter; i += 2) { + auto tag_expr_id = grammar_expr[i]; + auto rule_id = grammar_expr[i + 1]; + result.tag_rule_pairs.push_back({GetByteString(tag_expr_id), rule_id}); + } + + result.loop_after_dispatch = static_cast( + grammar_expr[grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter] + ); + + auto exclude_str_expr = GetGrammarExpr( + grammar_expr[grammar_expr.size() - TagDispatch::kTagDispatchExtraParameter + 1] + ); + XGRAMMAR_DCHECK(exclude_str_expr.type == GrammarExprType::kChoices); + result.excludes.reserve(exclude_str_expr.size()); + for (int j = 0; j < exclude_str_expr.size(); j++) { + result.excludes.push_back(GetByteString(exclude_str_expr[j])); + } + return result; + } + + /*! \brief Get the tag dispatch from the grammar expr with the given id. */ + TagDispatch GetTagDispatch(int32_t grammar_expr_id) const { + return GetTagDispatch(GetGrammarExpr(grammar_expr_id)); + } + + /*! \brief The object representing a token tag dispatch. */ + struct TokenTagDispatch { + std::vector> trigger_rule_pairs; // token_id → rule_id + bool loop_after_dispatch; + std::vector excludes; + }; + + /*! \brief Decode a kTokenTagDispatch expr into the TokenTagDispatch struct. */ + TokenTagDispatch GetTokenTagDispatch(const GrammarExpr& grammar_expr) const { + XGRAMMAR_DCHECK(grammar_expr.type == GrammarExprType::kTokenTagDispatch); + TokenTagDispatch result; + int pos = 0; + int32_t trigger_count = grammar_expr[pos++]; + for (int i = 0; i < trigger_count; ++i) { + auto token_id = grammar_expr[pos++]; + auto rule_id = grammar_expr[pos++]; + result.trigger_rule_pairs.push_back({token_id, rule_id}); + } + result.loop_after_dispatch = static_cast(grammar_expr[pos++]); + int32_t exclude_count = grammar_expr[pos++]; + for (int i = 0; i < exclude_count; ++i) { + result.excludes.push_back(grammar_expr[pos++]); + } + XGRAMMAR_DCHECK(pos == grammar_expr.size()); + return result; + } + + /*! \brief Get the token tag dispatch from the grammar expr with the given id. */ + TokenTagDispatch GetTokenTagDispatch(int32_t grammar_expr_id) const { + return GetTokenTagDispatch(GetGrammarExpr(grammar_expr_id)); + } + + private: + /*! \brief The rules of the grammar. rule_id corresponds the index of this vector. */ + std::vector rules_; + /*! \brief Suffix/stop metadata, sorted by rule_id and omitted for ordinary rules. */ + std::vector suffix_stop_infos_; + /*! \brief The data of all grammar_exprs. */ + std::vector grammar_expr_data_; + /*! \brief The start index of every grammar_expr in grammar_expr_data_. grammar_expr_id is the + * index to the elements in this vector. */ + std::vector grammar_expr_indptr_; + /*! \brief The id of the root rule. */ + int32_t root_rule_id_ = -1; + + public: + /******************* Aux information for matching *******************/ + + /*! \brief The complete FSM for the grammar. It contains the FSMs for all rules. */ + CompactFSM complete_fsm{NullObj{}}; + + /*! + * \brief The FSM for each rule. + * \details The FSM will be used in matching if it exists. If it does not exist (std::nullopt), + * the rule will be used in matching, and the rule's body must be a kChoices expr. + */ + std::vector> per_rule_fsms; + + /*! + * \brief The hash value for each rule's FSM. + */ + std::vector> per_rule_fsm_hashes; + + /*! + * \brief The new state ids of each FSM's states. + */ + std::vector>> per_rule_fsm_new_state_ids; + + /*! \brief The ids of the rules that are allowed to be empty. */ + std::vector allow_empty_rule_ids; + + /*! \brief Whether the grammar is optimized. */ + bool optimized = false; + + friend class GrammarBuilder; + friend class GrammarCompiler; + + friend std::size_t MemorySize(const Impl& impl); + friend struct member_trait; +}; + +XGRAMMAR_MEMBER_ARRAY( + Grammar::Impl::Rule, + &Grammar::Impl::Rule::name, + &Grammar::Impl::Rule::body_expr_id, + &Grammar::Impl::Rule::lookahead_assertion_id, + &Grammar::Impl::Rule::is_exact_lookahead, + &Grammar::Impl::Rule::max_tokens, + &Grammar::Impl::Rule::max_chars, + &Grammar::Impl::Rule::capture_name, + &Grammar::Impl::Rule::is_lazy, + &Grammar::Impl::Rule::temperature +); + +XGRAMMAR_MEMBER_ARRAY( + Grammar::Impl::SuffixStopInfo, + &Grammar::Impl::SuffixStopInfo::rule_id, + &Grammar::Impl::SuffixStopInfo::hidden_suffix_bytes, + &Grammar::Impl::SuffixStopInfo::hidden_stop_bytes, + &Grammar::Impl::SuffixStopInfo::body_rule_id, + &Grammar::Impl::SuffixStopInfo::marker_rule_id, + &Grammar::Impl::SuffixStopInfo::stop_capture_name +); + +XGRAMMAR_MEMBER_TABLE( + Grammar::Impl, + "rules", + &Grammar::Impl::rules_, + "suffix_stop_infos", + &Grammar::Impl::suffix_stop_infos_, + "grammar_expr_data", + &Grammar::Impl::grammar_expr_data_, + "grammar_expr_indptr", + &Grammar::Impl::grammar_expr_indptr_, + "root_rule_id", + &Grammar::Impl::root_rule_id_, + "complete_fsm", + &Grammar::Impl::complete_fsm, + "per_rule_fsms", + &Grammar::Impl::per_rule_fsms, + "allow_empty_rule_ids", + &Grammar::Impl::allow_empty_rule_ids, + "optimized", + &Grammar::Impl::optimized +); + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_IMPL_H_ diff --git a/zig/vendor/xgrammar/cpp/grammar_matcher.cc b/zig/vendor/xgrammar/cpp/grammar_matcher.cc new file mode 100644 index 0000000000..b31ec0ce66 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_matcher.cc @@ -0,0 +1,2856 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_matcher.cc + * \brief This source file implement the matcher class, especially the logic related to LLM tokens, + * like accepting tokens, leveraging the token mask cache to generate the mask, etc. matcher_base.cc + * implements the basic matching algorithm from strings to grammar. + */ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compiled_grammar_impl.h" +#include "earley_parser.h" +#include "grammar_impl.h" +#include "support/dynamic_bitset.h" +#include "support/encoding.h" +#include "support/int_set.h" +#include "support/logging.h" +#include "support/thread_pool.h" +#include "testing.h" +#include "tokenizer_info_impl.h" + +namespace xgrammar { + +/******************* Tool functions for token mask *******************/ +using GrammarExprType = Grammar::Impl::GrammarExprType; + +int32_t GetBitmaskSize(int vocab_size) { return DynamicBitset::GetBufferSize(vocab_size); } + +DLDataType GetBitmaskDLType() { return DLDataType{kDLInt, 32, 1}; } + +namespace details { + +using Clock = std::chrono::steady_clock; +using TimePoint = Clock::time_point; + +void ClearTokenBitmaskRow(int32_t* bitmask_data, int32_t bitmask_size, int32_t position) { + std::fill_n(bitmask_data + position * bitmask_size, bitmask_size, 0); +} + +bool TraverseDraftTreeRecursive( + int32_t current_position, + int32_t parent_position, + const int64_t* retrieve_next_token, + const int64_t* retrieve_next_sibling, + const int64_t* draft_tokens, + GrammarMatcher& matcher, + DLTensor* token_bitmask, + float* temperatures, + double time_threshold, + const TimePoint& start_time +) { + int32_t* bitmask_data = reinterpret_cast(token_bitmask->data); + int32_t bitmask_size = static_cast(token_bitmask->shape[1]); + + bool accepted; + if (current_position == 0) { + // The first token generated by the target model is always accepted. + accepted = true; + } else { + XGRAMMAR_CHECK(parent_position >= 0) + << "Non-root draft tree nodes must have a valid parent position"; + int64_t current_token_id = draft_tokens[current_position]; + if (current_token_id < 0 || current_token_id >= static_cast(bitmask_size) * 32) { + accepted = false; + } else { + int32_t* parent_bitmask = bitmask_data + parent_position * bitmask_size; + // 32 boolean bitmask values are packed into 32-bit integers. + uint32_t token_mask = uint32_t{1} << static_cast(current_token_id % 32); + accepted = (static_cast(parent_bitmask[current_token_id / 32]) & token_mask) != 0; + } + + // Check timeout for non-root nodes so the root token mask is still computed. + if (accepted && time_threshold > 0) { + auto elapsed = std::chrono::duration(Clock::now() - start_time).count(); + if (elapsed > time_threshold) { + return false; + } + } + } + + if (accepted) { + bool token_accepted = true; + if (current_position != 0) { + token_accepted = matcher.AcceptToken(static_cast(draft_tokens[current_position])); + } + + if (token_accepted) { + if (!matcher.IsTerminated()) { + matcher.FillNextTokenBitmask(token_bitmask, current_position); + if (temperatures != nullptr) { + temperatures[current_position] = matcher.GetTemperature().value_or(-1.0f); + } + + if (retrieve_next_token[current_position] != -1) { + bool success = TraverseDraftTreeRecursive( + retrieve_next_token[current_position], + current_position, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + matcher, + token_bitmask, + temperatures, + time_threshold, + start_time + ); + if (!success) { + if (current_position != 0) { + matcher.Rollback(1); + } + return false; + } + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + + if (current_position != 0) { + matcher.Rollback(1); + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + } else { + ClearTokenBitmaskRow(bitmask_data, bitmask_size, current_position); + } + + if (retrieve_next_sibling[current_position] != -1) { + bool success = TraverseDraftTreeRecursive( + retrieve_next_sibling[current_position], + parent_position, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + matcher, + token_bitmask, + temperatures, + time_threshold, + start_time + ); + if (!success) { + return false; + } + } + + return true; +} + +} // namespace details + +int32_t* CheckAndGetBitmaskPtr(const DLTensor& token_bitmask, int vocab_size, int index) { + XGRAMMAR_CHECK(token_bitmask.dtype.code == kDLInt && token_bitmask.dtype.bits == 32) + << "The provied bitmask's dtype is not valid: should be int32"; + + int32_t buffer_size = GetBitmaskSize(vocab_size); + if (token_bitmask.ndim == 1) { + XGRAMMAR_CHECK(token_bitmask.shape[0] == buffer_size) + << "The provided bitmask's shape is not valid: should be (" << buffer_size << ", )"; + XGRAMMAR_CHECK(index == 0) << "The index should be 0 when the bitmask is 1D"; + } else { + XGRAMMAR_CHECK(token_bitmask.ndim == 2) + << "The provided bitmask's shape is not valid: should be (batch_size, " << buffer_size + << ")"; + XGRAMMAR_CHECK(token_bitmask.shape[1] == buffer_size) + << "The provided bitmask's shape is not valid: should be (batch_size, " << buffer_size + << ")"; + XGRAMMAR_CHECK(index >= 0 && index < token_bitmask.shape[0]) + << "The provided index is out of bounds"; + } + + XGRAMMAR_CHECK( + token_bitmask.device.device_type == kDLCPU || + token_bitmask.device.device_type == kDLCUDAHost || + token_bitmask.device.device_type == kDLROCMHost + ) << "The provided bitmask's device is not valid: should be CPU"; + + return reinterpret_cast(token_bitmask.data) + index * buffer_size; +} + +void _DebugGetMaskedTokensFromBitmask( + std::vector* rejected_tokens, const DLTensor& token_bitmask, int vocab_size, int index +) { + int32_t* data_ptr = CheckAndGetBitmaskPtr(token_bitmask, vocab_size, index); + DynamicBitset bitset(vocab_size, reinterpret_cast(data_ptr)); + rejected_tokens->clear(); + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + rejected_tokens->push_back(i); + } +} + +std::pair _IsSingleTokenBitmask(const DLTensor& bitmask, int vocab_size, int index) { + int32_t* data_ptr = CheckAndGetBitmaskPtr(bitmask, vocab_size, index); + DynamicBitset bitset(vocab_size, reinterpret_cast(data_ptr)); + if (bitset.Count() == 1) { + return std::make_pair(true, bitset.FindFirstOne()); + } else { + return std::make_pair(false, -1); + } +} + +void ApplyMask32Bits( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + XGRAMMAR_CHECK(logits->dtype.code == kDLFloat && logits->dtype.bits == 32) + << "The provided logits's dtype is not valid: should be float32"; + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + int logits_stride0 = logits->strides[0]; + int bitmask_stride0 = bitmask.strides[0]; + if (indices.has_value()) { + for (auto idx : indices.value()) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = -std::numeric_limits::infinity(); + } + } + } else { + for (int idx = 0; idx < logits_shape.first; ++idx) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = -std::numeric_limits::infinity(); + } + } + } +} + +void ApplyMask16Bits( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + XGRAMMAR_CHECK(logits->dtype.bits == 16) + << "The provided logits's dtype is not valid: should be bfloat16 or float16"; + uint16_t kMinusInfinity; + const uint16_t kMinusInfinityBf16 = 0xff80; + const uint16_t kMinusInfinityFp16 = 0xfc00; + switch (logits->dtype.code) { + case kDLBfloat: + kMinusInfinity = kMinusInfinityBf16; + break; + case kDLFloat: + kMinusInfinity = kMinusInfinityFp16; + break; + default: + XGRAMMAR_LOG(FATAL + ) << "The provided logits's dtype is not valid: should be bfloat16 or float16"; + } + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + int logits_stride0 = logits->strides[0]; + int bitmask_stride0 = bitmask.strides[0]; + if (indices.has_value()) { + for (auto idx : indices.value()) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = kMinusInfinity; + } + } + } else { + for (int idx = 0; idx < logits_shape.first; ++idx) { + uint32_t* data_ptr = reinterpret_cast(bitmask.data) + idx * bitmask_stride0; + DynamicBitset bitset(vocab_size, data_ptr); + auto logits_ptr = reinterpret_cast(logits->data) + idx * logits_stride0; + for (int i = bitset.FindFirstZero(); i != -1; i = bitset.FindNextZero(i)) { + logits_ptr[i] = kMinusInfinity; + } + } + } +} + +void ApplyTokenBitmaskInplaceCPU( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size, + std::optional> indices +) { + // Check device and dim + XGRAMMAR_CHECK( + logits->device.device_type == kDLCPU || logits->device.device_type == kDLCUDAHost || + logits->device.device_type == kDLROCMHost + ) << "The provided logits's device is not valid: should be CPU"; + XGRAMMAR_CHECK( + bitmask.device.device_type == kDLCPU || bitmask.device.device_type == kDLCUDAHost || + bitmask.device.device_type == kDLROCMHost + ) << "The provided bitmask's device is not valid: should be CPU"; + XGRAMMAR_CHECK(logits->ndim == 2 || logits->ndim == 1) + << "The provided logits's shape is not valid: should be 2D or 1D"; + XGRAMMAR_CHECK(bitmask.ndim == 2 || bitmask.ndim == 1) + << "The provided bitmask's shape is not valid: should be 2D or 1D"; + + // Check type + XGRAMMAR_CHECK(logits->dtype.lanes == 1) + << "The provided logits's dtype is not valid: lanes should be 1"; + XGRAMMAR_CHECK( + bitmask.dtype.code == kDLInt && bitmask.dtype.bits == 32 && bitmask.dtype.lanes == 1 + ) << "The provided bitmask's dtype is not valid: should be int32"; + + // Check shape + std::pair logits_shape = + logits->ndim == 2 + ? std::make_pair(static_cast(logits->shape[0]), static_cast(logits->shape[1])) + : std::make_pair(1, static_cast(logits->shape[0])); + std::pair bitmask_shape = + bitmask.ndim == 2 + ? std::make_pair(static_cast(bitmask.shape[0]), static_cast(bitmask.shape[1])) + : std::make_pair(1, static_cast(bitmask.shape[0])); + + XGRAMMAR_CHECK( + vocab_size <= bitmask_shape.second * DynamicBitset::BITS_PER_BLOCK && + vocab_size <= logits_shape.second + ); + + if (!indices.has_value()) { + XGRAMMAR_CHECK(logits_shape.first == bitmask_shape.first) + << "When indices is not provided, the logits's batch size should be equal to the " + "bitmask's batch size, but got " + << logits_shape.first << " vs " << bitmask_shape.first; + } else { + // Each index selects a row of both logits and bitmask; an out-of-range index would lead to + // an out-of-bounds access, so validate every index against both batch sizes here. + for (int idx : indices.value()) { + XGRAMMAR_CHECK(idx >= 0 && idx < logits_shape.first && idx < bitmask_shape.first) + << "The provided index " << idx << " is out of bounds: it should be in [0, " + << std::min(logits_shape.first, bitmask_shape.first) << ")."; + } + } + + // Apply mask + if (logits->dtype.bits == 32) { + ApplyMask32Bits(logits, bitmask, vocab_size, indices); + } else if (logits->dtype.bits == 16) { + ApplyMask16Bits(logits, bitmask, vocab_size, indices); + } else { + XGRAMMAR_LOG(FATAL + ) << "The provided logits's dtype is not valid: should be float32 or float16/bfloat16"; + } +} + +/******************* Grammar Matcher with Adaptive Token Mask *******************/ + +/* + * Note on the matching algorithm (this is the old description for the matching algorithm, please + * refer to https://arxiv.org/pdf/2411.15100 for the latest description) + * + * Given a context-free grammar, we match the characters in a string one by one. + * + * We adopt a non-deterministic pushdown automata (NPDA) in matching. To be specific, we maintain + * several stacks, each of which represents a possible path in the NPDA, and update the stacks + * during matching. + * + * ## Stack Structure (see grammar_matcher_state.h) + * The element of every stack is a StackElement object, referring a position in the grammar. If a + * StackElement points to a RuleRef element (referring to another rule), the next element of the + * stack will be a position in this rule. If a StackElement is a CharacterClass element, it will be + * the last in the stack, meaning *the next* character to match. + * + * ## Matching Process (see grammar_matcher_base.h) + * When accepting a new character and it is accepted by a stack, the last element of the stack will + * be advanced to the next position in the grammar. If it gets to the end of the rule, several + * elements at the end may be popped out, and the last element of the stack will be advanced. + * + * One stack may split since there may be multiple possible next positions. In this case, similar + * stacks with different top elements will be added. When one stack cannot accept the new character, + * it will be removed from the stacks. + * + * ## Storage of Stacks (see grammar_matcher_state.h) + * Note these stacks form a tree structure as when splitting, the new stacks share the same prefix. + * We store all StackElements as a tree, where every path from tree root to a node represents a + * stack. To represent stack tops, we attach additional pointers pointing the stack top nodes. + * Also, We maintain a history of the stack top pointers, so we can rollback to the previous state. + * + * All tree nodes are maintained by a buffer, and utilize reference counting to recycle. If a node + * is neither pointed by a stack top pointer, not pointed by some child nodes, it will be freed. + * + * ## Example + * ### Grammar + * root ::= [a] R + * R ::= [b] S [c] | [b] [c] T + * S ::= "" | [c] [d] + * T ::= [e] + * + * ### The previous step + * Previous accepted string: ab + * Previous stack tree: + * A------ + * | \ \ + * B D< E< + * | + * C< + * + * A: (rule root, choice 0, element 1) + * B: (rule R, choice 0, element 1) + * C: (rule S, choice 1, element 0) + * D: (rule R, choice 0, element 2) + * E: (rule R, choice 1, element 1) + * < means the stack top pointers in the previous step. + * The stacks in the previous step is: (A, B, C), (A, D), (A, E) + * + * ### The current step + * Current accepted string: abc + * Current stack tree: + * A----------------- G<< + * | \ \ \ + * B--- D< E< H + * | \ | + * C< F<< I<< + * + * F: (rule S, choice 1, element 1) + * G: (rule root, choice 0, element 2) (means the matching process has finished, and will be deleted + * when the next char comes) + * H: (rule R, choice 1, element 2) + * I: (rule T, choice 0, element 0) + * << means the stack top pointers in the current step. + * The stacks in the current step is: (A, B, F), (A, H, I), (G,) + * + * ## Preprocess (see grammar_matcher_preproc.h) + * We will store all information about tokens that needed in matching in a CompiledGrammar + * object. Tokens are sorted by codepoint, allowing us to reuse the repeated prefixes between + * different tokens. + * + * For a given position in a rule, if we only consider this rule and its sub-rules during matching, + * without considering its parent rules (in actual matching, we also need to consider its parent + * rules), we can already determine that some tokens are acceptable while others are definitely + * rejected. Therefore, for a position in a rule, we can divide the token set into three categories: + * - accepted_indices: If a token is accepted by this rule + * - rejected_indices: If a token is rejected by this rule + * - uncertain_indices: Whether it can be accepted depends on the information from the parent + * level during actual matching. To be specific, If this token has a prefix that has not been + * rejected and has reached the end of this rule, then it is possible for it to be further accepted + * by the parent rule. + * + * During actual matching, we will directly accept or reject the tokens in accepted_indices and + * rejected_indices, and only consider the tokens in uncertain_indices. That speeds up the matching + * process. + */ + +/* \brief The concrete implementation of GrammarMatcherNode. */ +class GrammarMatcher::Impl : public EarleyParser { + public: + Impl( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens = std::nullopt, + bool terminate_without_stop_token = false, + // max_rollback_tokens_ is deprecated and not used. + int max_rollback_tokens = -1, + std::optional default_temperature = std::nullopt + ) + : EarleyParser(compiled_grammar->grammar), + compiled_grammar_(compiled_grammar), + tokenizer_info_(compiled_grammar->tokenizer_info), + stop_token_ids_(override_stop_tokens.value_or(tokenizer_info_.GetStopTokenIds())), + terminate_without_stop_token_(terminate_without_stop_token), + default_temperature_(default_temperature), + tmp_accepted_bitset_(tokenizer_info_.GetVocabSize()) { + if (override_stop_tokens.has_value()) { + XGRAMMAR_CHECK(!override_stop_tokens->empty()) + << "The override_stop_tokens should not be empty"; + // Stop token ids are written into the token bitmask, so they must be in range. + for (int id : *override_stop_tokens) { + XGRAMMAR_CHECK(id >= 0 && id < tokenizer_info_.GetVocabSize()) + << "The override stop token id " << id << " is out of the vocabulary range [0, " + << tokenizer_info_.GetVocabSize() << ")"; + } + } + if (has_budget_rules_ || has_char_budget_rules_) { + for (int32_t rule_id = 0; rule_id < grammar_->NumRules(); ++rule_id) { + const auto& rule = grammar_->GetRule(rule_id); + if (rule.max_tokens < 0 && rule.max_chars < 0) { + continue; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(rule_id); + if (suffix_stop_info != nullptr && suffix_stop_info->body_rule_id >= 0) { + has_budget_marker_rules_ = true; + break; + } + } + } + if (has_char_budget_rules_) { + has_token_edges_ = HasTokenEdges(); + } + XGRAMMAR_CHECK( + !default_temperature_.has_value() || + (std::isfinite(default_temperature_.value()) && default_temperature_.value() >= 0) + ) << "The default_temperature must be a finite non-negative number"; + } + + bool AcceptToken(int32_t token_id, bool debug_print = false); + + bool AcceptString(const std::string& input_str, bool debug_print = false); + + bool FillNextTokenBitmask(DLTensor* next_token_bitmask, int index, bool debug_print = false); + + std::optional GetTemperature() const; + + std::string FindJumpForwardString(); + + void Rollback(int num_tokens); + + bool IsTerminated() const; + + void Reset() { + token_length_history.clear(); + current_token_index_ = -1; + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + budget_exceeded_warned_ = false; + char_budget_exceeded_warned_ = false; + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + budget_body_match_cache_.clear(); + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + accepted_bytes_.clear(); + row_byte_end_.assign(1, 0); + EarleyParser::Reset(); + } + + std::vector> GetCaptures(bool deduplicate) const; + + int GetMaxRollbackTokens() const { return -1; } + + const std::vector& GetStopTokenIds() const { return stop_token_ids_; } + + std::string _DebugPrintInternalState() const { return PrintStates(); } + + private: + using StoreType = AdaptiveTokenMask::StoreType; + + /*! + * \brief If is_uncertain_saved is true, find the next token in uncertain_indices. Otherwise, + * find the next token that is set to true in uncertain_tokens_bitset. + * \param iterator_uncertain The helper iterator to iterate over uncertain_indices or + * uncertain_tokens_bitset. + * \returns The index of the next token, or -1 if no more token. + */ + int GetNextUncertainToken( + bool is_uncertain_saved, + int* iterator_uncertain, + const std::vector& uncertain_indices, + const std::vector& uncertain_tokens_bitset + ); + + /*! \brief Set the acceptable next token in next_token_bitmask. */ + void SetTokenBitmask( + int32_t* bitmask_data_ptr, + const DynamicBitset& accepted_bitset, + const std::vector& rejected_indices, + bool can_reach_end, + bool allow_special_token = false + ); + + /*! + * \brief Accept the stop token and terminates the matcher. + * \returns Whether the stop token can be accepted. + */ + bool AcceptStopToken(); + + bool IsStopTokenAccepted() const; + + /*! \brief Check if the token bitmask is all-true. */ + bool IsTokenBitmaskAllTrue(int32_t* bitmask_data_ptr); + + std::string PrintBitmask(int32_t* bitmask_data_ptr, const TokenizerInfo& tokenizer_info); + + /*! \brief Restores skip_expired_states_ to false when the accept call returns. */ + class SkipExpiredGuard { + public: + explicit SkipExpiredGuard(Impl* impl, bool enable) : impl_(impl) { + impl_->skip_expired_states_ = enable; + } + ~SkipExpiredGuard() { impl_->skip_expired_states_ = false; } + SkipExpiredGuard(const SkipExpiredGuard&) = delete; + SkipExpiredGuard& operator=(const SkipExpiredGuard&) = delete; + + private: + Impl* impl_; + }; + + /*! + * \brief RAII guard that enables capture-event recording in the Earley parser for its + * lifetime. Used in the definitive accept paths (AcceptToken / AcceptString) only, so that + * speculative advances (mask computation, jump-forward search) never record capture events. + */ + class CaptureRecordingGuard { + public: + explicit CaptureRecordingGuard(Impl* impl) : impl_(impl) { + if (impl_->capture_tracking_) { + impl_->capture_recording_ = true; + } + } + ~CaptureRecordingGuard() { impl_->capture_recording_ = false; } + CaptureRecordingGuard(const CaptureRecordingGuard&) = delete; + CaptureRecordingGuard& operator=(const CaptureRecordingGuard&) = delete; + + private: + Impl* impl_; + }; + + /*! \brief Fill the bitmask from the current states, optionally excluding the states whose + * budget deadline has passed. */ + void FillBitmaskForStates( + int32_t* bitmask_data_ptr, int index, bool skip_expired, bool debug_print + ); + + void FillBitmaskForCharBudgetBoundary( + const AdaptiveTokenMask& adaptive_token_mask, int32_t remaining_chars + ); + + bool AdvanceWithCharacterBudget(uint8_t byte, bool debug_print = false); + + bool AdvanceAtomicTokenWithCharacterBudget( + int32_t token_id, int32_t token_char_count, bool debug_print = false + ); + + /*! \brief Whether byte offsets are needed for captures or budgeted suffix/stop rules. */ + bool ShouldTrackAcceptedBytes() const { + return IsCaptureTrackingEnabled() || has_budget_marker_rules_; + } + + /*! \brief Whether the bytes consumed by this expired suffix/stop occurrence form a complete + * match of its body, so max_tokens may end it without consuming the marker. */ + bool CanForceCompleteWithoutMarker(const ParserState& state, bool use_char_budget); + + /*! \brief Commit the current budget boundary: discard expired states and complete eligible + * suffix/stop occurrences without their marker. Returns whether a derivation remains viable. */ + bool ApplyBudgetEnforcement(bool debug_print = false); + + /*! \brief Enforce exhausted character budgets before consuming a new codepoint. */ + bool ApplyCharacterBudgetEnforcement(bool debug_print = false); + + /*! \brief Reapply character-budget enforcement for newly exposed parent occurrences. */ + bool ApplyCharacterBudgetEnforcementToFixedPoint(bool debug_print = false); + + /*! \brief Consume the pending budget decision and warn once when a token budget was + * exceeded. */ + void FinishCharacterBudgetAccept() { + if (char_budget_relaxed_ && !char_budget_exceeded_warned_) { + char_budget_exceeded_warned_ = true; + XGRAMMAR_LOG(WARNING + ) << "The character budget (max_chars) of a rule was exceeded: the rule could not end at " + "the position where its budget ran out, so the budget is relaxed until the rule " + "can end. This warning is reported once per matcher."; + } + char_budget_relaxed_ = false; + } + + bool FinishAccept(bool consumed_past_deadline) { + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + if (consumed_past_deadline && !budget_exceeded_warned_) { + budget_exceeded_warned_ = true; + XGRAMMAR_LOG(WARNING + ) << "The token budget (max_tokens) of a rule was exceeded: the rule could not end at " + "the position where its budget ran out, so the budget is relaxed until the rule " + "can end. This warning is reported once per matcher."; + } + FinishCharacterBudgetAccept(); + return true; + } + + bool BitmaskHasAnyToken(int32_t* bitmask_data_ptr) const { + auto bitset = DynamicBitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + return bitset.Any(); + } + + /*! \brief Record that num_rows new input positions were created, each consuming one byte of + * bytes in order. Used for the byte-by-byte advance paths. */ + void AppendPerByteRows(const std::string& bytes) { + for (char c : bytes) { + accepted_bytes_.push_back(static_cast(c)); + row_byte_end_.push_back(static_cast(accepted_bytes_.size())); + } + } + + /*! \brief Record that one new input position was created, consuming all bytes of the token. + * Used for the atomic token advance path. */ + void AppendAtomicRow(const std::string& bytes) { + accepted_bytes_.insert(accepted_bytes_.end(), bytes.begin(), bytes.end()); + row_byte_end_.push_back(static_cast(accepted_bytes_.size())); + } + + CompiledGrammar compiled_grammar_; + TokenizerInfo tokenizer_info_; + std::vector stop_token_ids_; + bool terminate_without_stop_token_; + std::optional default_temperature_; + mutable bool has_warned_temperature_conflict_ = false; + std::deque token_length_history; + + /*! \brief Set by the last mask computation when an exhausted budget could be enforced: the + * next accept commits the same state transition used to compute that mask. */ + bool budget_enforce_pending_ = false; + /*! \brief Whether the pending decision includes a suffix/stop completion without its marker. */ + bool budget_force_close_pending_ = false; + /*! \brief Whether the one-time budget-exceeded warning has been reported. */ + bool budget_exceeded_warned_ = false; + /*! \brief Whether the one-time character-budget warning has been reported. */ + bool char_budget_exceeded_warned_ = false; + /*! \brief Whether the current definitive accept relaxed a character budget. */ + bool char_budget_relaxed_ = false; + /*! \brief Whether speculative character-budget relaxation should be recorded. */ + bool record_char_budget_relaxation_ = false; + /*! \brief Whether byte history is needed to recognize a budgeted suffix/stop body boundary. */ + bool has_budget_marker_rules_ = false; + /*! \brief Whether the grammar has Token/ExcludeToken edges, the only edges that can accept a + * token the byte-level walk rejected. */ + bool has_token_edges_ = false; + /*! \brief Whether the complete FSM has any Token/ExcludeToken edge. */ + bool HasTokenEdges() const { + const auto& fsm = grammar_->complete_fsm; + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetEdges(state)) { + if (edge.IsToken() || edge.IsExcludeToken()) { + return true; + } + } + } + return false; + } + + struct BudgetBodyMatchProgress { + int64_t begin_byte; + int64_t end_byte; + std::unordered_set states; + std::string temporary_input_snapshot; + }; + /*! \brief Incremental body-FSM progress keyed by (rule id, occurrence start row). */ + std::unordered_map budget_body_match_cache_; + + /*! \brief The bytes accepted so far. Maintained for captures and budgeted marker rules. */ + std::vector accepted_bytes_; + /*! \brief row_byte_end_[i] is the number of accepted bytes after input position i. Aligned + * with the parser's state history whenever accepted_bytes_ is maintained. */ + std::vector row_byte_end_{0}; + /*! \brief Bytes tentatively consumed by the current token or string advance. */ + std::string temporary_input_bytes_; + /*! \brief Parser row immediately before temporary_input_bytes_. */ + int32_t temporary_input_start_row_ = -1; + + // Temporary data for FillNextTokenBitmask. They are stored here to avoid repeated allocation. + DynamicBitset tmp_accepted_bitset_; + std::vector tmp_rejected_indices_; + std::vector tmp_rejected_indices_delta_; +}; + +class BatchGrammarMatcher::Impl { + public: + Impl(std::variant max_threads) { + if (std::holds_alternative(max_threads)) { + int32_t num_threads = std::get(max_threads); + XGRAMMAR_CHECK(num_threads >= 1) + << "The num_threads should be at least 1, but got " << num_threads; + if (num_threads > 1) { + if (num_threads > static_cast(std::thread::hardware_concurrency())) { + XGRAMMAR_LOG(WARNING) << "The num_threads " << num_threads << " is larger than the " + << "number of hardware threads. Using " + << static_cast(std::thread::hardware_concurrency()) + << " instead."; + } + max_threads_ = + std::min(num_threads, static_cast(std::thread::hardware_concurrency())); + } + } else { + std::string str = std::get(max_threads); + XGRAMMAR_CHECK(str == "auto"); + max_threads_ = std::thread::hardware_concurrency() / 2; + } + } + + void BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print + ); + + static std::vector BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print + ); + + static std::vector BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print + ); + + static void BatchRollback( + std::vector* matchers, const std::vector& num_tokens + ); + + private: + std::optional thread_pool_ = std::nullopt; + int32_t max_threads_ = 1; +}; + +bool GrammarMatcher::Impl::AcceptStopToken() { + if (terminate_without_stop_token_) { + return false; + } + if (!IsCompleted()) { + return false; + } + XGRAMMAR_DCHECK(!stop_token_is_accepted_); + token_length_history.push_back(0); + stop_token_is_accepted_ = true; + return true; +} + +bool GrammarMatcher::Impl::IsTerminated() const { + if (terminate_without_stop_token_) { + return IsCompleted(); + } + return IsStopTokenAccepted(); +} + +bool GrammarMatcher::Impl::IsStopTokenAccepted() const { return stop_token_is_accepted_; } + +bool GrammarMatcher::Impl::CanForceCompleteWithoutMarker( + const ParserState& state, bool use_char_budget +) { + bool expired = use_char_budget ? IsCharExpiredState(state) : IsExpiredState(state); + if (!expired || state.rule_id < 0) { + return false; + } + const auto& rule = grammar_->GetRule(state.rule_id); + if (use_char_budget ? rule.max_chars < 0 : rule.max_tokens < 0) { + return false; + } + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(state.rule_id); + if (suffix_stop_info == nullptr || suffix_stop_info->body_rule_id < 0) { + return false; + } + XGRAMMAR_DCHECK(ShouldTrackAcceptedBytes()); + int32_t start_row = + state.rule_start_pos == ParserState::kNoPrevInputPos ? 0 : state.rule_start_pos; + auto byte_position_for_row = [&](int32_t row) { + if (row >= 0 && row < static_cast(row_byte_end_.size())) { + return row_byte_end_[row]; + } + XGRAMMAR_DCHECK(temporary_input_start_row_ >= 0); + XGRAMMAR_DCHECK(row >= temporary_input_start_row_); + XGRAMMAR_DCHECK( + row <= temporary_input_start_row_ + static_cast(temporary_input_bytes_.size()) + ); + return static_cast(accepted_bytes_.size() + row - temporary_input_start_row_); + }; + int64_t begin = byte_position_for_row(start_row); + int64_t end = accepted_bytes_.size() + temporary_input_bytes_.size(); + + XGRAMMAR_DCHECK(grammar_->per_rule_fsms[suffix_stop_info->body_rule_id].has_value()); + const auto& body_fsm = grammar_->per_rule_fsms[suffix_stop_info->body_rule_id]->GetFsm(); + XGRAMMAR_DCHECK(body_fsm.IsLeaf()) << "A suffix/stop body helper must compile to a leaf FSM"; + + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + auto [it, inserted] = budget_body_match_cache_.try_emplace( + occurrence, BudgetBodyMatchProgress{begin, begin, {body_fsm.GetStart()}, {}} + ); + auto& progress = it->second; + if (inserted) { + body_fsm.GetFsm().GetEpsilonClosure(&progress.states); + } else if (progress.begin_byte != begin || progress.end_byte > end || + (progress.end_byte > static_cast(accepted_bytes_.size()) && + progress.temporary_input_snapshot != temporary_input_bytes_)) { + // This is only expected after restoring external matcher state. Rollback and reset clear the + // cache eagerly, but reinitialize defensively if a caller supplies an equivalent history. + progress = BudgetBodyMatchProgress{begin, begin, {body_fsm.GetStart()}, {}}; + body_fsm.GetFsm().GetEpsilonClosure(&progress.states); + } + + std::unordered_set next_states; + for (int64_t offset = progress.end_byte; offset < end && !progress.states.empty(); ++offset) { + uint8_t byte = + offset < static_cast(accepted_bytes_.size()) + ? accepted_bytes_[offset] + : static_cast(temporary_input_bytes_[offset - accepted_bytes_.size()]); + body_fsm.GetFsm().Advance( + progress.states, byte, &next_states, FSMEdge::EdgeType::kCharRange, true + ); + progress.states = next_states; + } + progress.end_byte = end; + progress.temporary_input_snapshot = + end > static_cast(accepted_bytes_.size()) ? temporary_input_bytes_ : std::string{}; + return std::any_of(progress.states.begin(), progress.states.end(), [&](int state_id) { + return body_fsm.IsEndState(state_id); + }); +} + +bool GrammarMatcher::Impl::ApplyBudgetEnforcement(bool debug_print) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + const auto latest_row = scanable_state_history_[scanable_state_history_.size() - 1]; + std::vector latest_states(latest_row.begin(), latest_row.end()); + std::vector force_completed_states; + std::unordered_set force_completed_occurrences; + + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + tmp_accept_stop_token_ = IsCompleted(); + + for (const auto& state : latest_states) { + if (!IsExpiredState(state)) { + EnqueueWithoutProcessing(state); + continue; + } + if (!CanForceCompleteWithoutMarker(state, false)) { + continue; + } + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + if (force_completed_occurrences.insert(occurrence).second) { + force_completed_states.push_back(state); + } + } + + for (const auto& state : force_completed_states) { + Complete(state, debug_print, /*marker_present=*/false); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + + bool any_expired = false; + bool any_alive = false; + for (const auto& state : tmp_states_to_be_added_) { + if (IsExpiredState(state)) { + any_expired = true; + } else { + any_alive = true; + } + } + if (any_expired && (any_alive || tmp_accept_stop_token_)) { + // Completing a budgeted suffix/stop rule can expose an already-expired parent alternative. + // Apply the same preference as the ordinary max_tokens path: once another derivation can + // continue or finish at this boundary, expired derivations may not consume another token. + tmp_states_to_be_added_.erase( + std::remove_if( + tmp_states_to_be_added_.begin(), + tmp_states_to_be_added_.end(), + [&](const ParserState& state) { return IsExpiredState(state); } + ), + tmp_states_to_be_added_.end() + ); + } + + bool viable = tmp_accept_stop_token_ || !tmp_states_to_be_added_.empty(); + if (!viable) { + return false; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + is_completed_.back() = tmp_accept_stop_token_; + return true; +} + +bool GrammarMatcher::Impl::ApplyCharacterBudgetEnforcement(bool debug_print) { + XGRAMMAR_DCHECK(tmp_process_state_queue_.empty()); + const auto latest_row = scanable_state_history_[scanable_state_history_.size() - 1]; + std::vector latest_states(latest_row.begin(), latest_row.end()); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + std::vector> previous_completable_states( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + std::vector previous_capture_events; + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + std::vector force_completed_states; + std::unordered_set force_completed_occurrences; + + tmp_states_visited_in_queue_.Clear(); + tmp_states_to_be_added_.clear(); + tmp_completed_lazy_occurrences_.clear(); + tmp_accept_stop_token_ = IsCompleted(); + + for (const auto& state : latest_states) { + if (!IsCharExpiredState(state)) { + EnqueueWithoutProcessing(state); + continue; + } + if (!CanForceCompleteWithoutMarker(state, true)) { + continue; + } + int64_t occurrence = + (static_cast(state.rule_id) << 32) | static_cast(state.rule_start_pos); + if (force_completed_occurrences.insert(occurrence).second) { + force_completed_states.push_back(state); + } + } + + for (const auto& state : force_completed_states) { + Complete(state, debug_print, /*marker_present=*/false); + } + while (!tmp_process_state_queue_.empty()) { + const auto state = std::move(tmp_process_state_queue_.front()); + tmp_process_state_queue_.pop(); + auto [scanable, completable] = Predict(state, debug_print); + if (completable) { + Complete(state, debug_print); + } + if (scanable) { + tmp_states_to_be_added_.push_back(state); + } + } + if (!tmp_completed_lazy_occurrences_.empty()) { + RemoveCommittedLazyStates(); + } + + bool any_expired = false; + bool any_alive = false; + for (const auto& state : tmp_states_to_be_added_) { + if (IsCharExpiredState(state)) { + any_expired = true; + } else { + any_alive = true; + } + } + if (any_expired && (any_alive || tmp_accept_stop_token_)) { + tmp_states_to_be_added_.erase( + std::remove_if( + tmp_states_to_be_added_.begin(), + tmp_states_to_be_added_.end(), + [&](const ParserState& state) { return IsCharExpiredState(state); } + ), + tmp_states_to_be_added_.end() + ); + } + + bool viable = tmp_accept_stop_token_ || !tmp_states_to_be_added_.empty(); + if (!viable) { + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + return false; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(tmp_states_to_be_added_); + is_completed_.back() = tmp_accept_stop_token_; + return true; +} + +bool GrammarMatcher::Impl::ApplyCharacterBudgetEnforcementToFixedPoint(bool debug_print) { + bool applied_any = false; + for (int32_t iteration = 0; iteration <= grammar_->NumRules(); ++iteration) { + bool has_expired_state = false; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (IsCharExpiredState(state)) { + has_expired_state = true; + break; + } + } + if (!has_expired_state) { + return applied_any; + } + if (!ApplyCharacterBudgetEnforcement(debug_print)) { + return applied_any; + } + applied_any = true; + } + return applied_any; +} + +bool GrammarMatcher::Impl::AdvanceWithCharacterBudget(uint8_t byte, bool debug_print) { + if (!has_char_budget_rules_ || !StartsUTF8Codepoint(byte)) { + return Advance(byte, debug_print); + } + + bool has_expired_state = false; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (IsCharExpiredState(state)) { + has_expired_state = true; + break; + } + } + if (!has_expired_state) { + return Advance(byte, debug_print); + } + + std::vector previous_states = GetLatestScanableStates(); + bool previous_completed = IsCompleted(); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + std::vector> previous_completable_states( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + std::vector previous_capture_events; + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + + bool enforced = ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + if (!enforced && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + bool accepted = Advance(byte, debug_print); + if (!accepted && enforced) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(previous_states); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + is_completed_.back() = previous_completed; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + } + return accepted; +} + +bool GrammarMatcher::Impl::AdvanceAtomicTokenWithCharacterBudget( + int32_t token_id, int32_t token_char_count, bool debug_print +) { + bool has_expired_state = false; + bool crosses_budget = false; + if (has_char_budget_rules_ && token_char_count > 0) { + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + has_expired_state = has_expired_state || IsCharExpiredState(state); + crosses_budget = + crosses_budget || (state.char_budget_deadline >= 0 && + GetCurrentCharIndex() + token_char_count > state.char_budget_deadline); + } + } + + std::vector previous_states; + std::vector> previous_completable_states; + bool previous_completed = false; + std::vector previous_capture_events; + bool enforced = false; + if (has_expired_state) { + previous_states = GetLatestScanableStates(); + auto previous_completable_row = rule_id_to_completable_states_.Back(); + previous_completable_states.assign( + previous_completable_row.data, + previous_completable_row.data + previous_completable_row.data_len + ); + previous_completed = IsCompleted(); + if (capture_tracking_) { + previous_capture_events = CopyLastCaptureRow(); + } + enforced = ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + if (!enforced && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + } + + bool accepted = AdvanceAtomicToken(token_id, debug_print, token_char_count); + if (!accepted && enforced) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(previous_states); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(previous_completable_states); + is_completed_.back() = previous_completed; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(previous_capture_events); + } + } + if (accepted && crosses_budget && record_char_budget_relaxation_) { + char_budget_relaxed_ = true; + } + return accepted; +} + +std::optional GrammarMatcher::Impl::GetTemperature() const { + std::optional syntax_temperature = std::nullopt; + bool has_temperature_conflict = false; + for (const auto& state : GetLatestScanableStates()) { + if (state.active_temperature_rule_id == -1) { + continue; + } + const auto& temperature = grammar_->GetRule(state.active_temperature_rule_id).temperature; + XGRAMMAR_DCHECK(temperature.has_value()); + if (syntax_temperature.has_value() && temperature.value() != syntax_temperature.value()) { + has_temperature_conflict = true; + } + if (!syntax_temperature.has_value() || temperature.value() > syntax_temperature.value()) { + syntax_temperature = temperature; + } + } + if (has_temperature_conflict && !has_warned_temperature_conflict_) { + XGRAMMAR_LOG(WARNING) << "Multiple active grammar paths specify different temperatures. " + "Using the maximum temperature " + << syntax_temperature.value() << "."; + has_warned_temperature_conflict_ = true; + } + return syntax_temperature.has_value() ? syntax_temperature : default_temperature_; +} + +// TODO(yixin): Polish verbose logging +bool GrammarMatcher::Impl::AcceptToken(int32_t token_id, bool debug_print) { + if (IsStopTokenAccepted()) { + XGRAMMAR_LOG(WARNING) << "The matcher has terminated after accepting the stop token, but is " + << "trying to accept new token with id " << token_id << "."; + return false; + } + + if (token_id < 0 || token_id >= tokenizer_info_.GetVocabSize()) { + XGRAMMAR_LOG(WARNING) << "The token id " << token_id << " is out of range [0, " + << tokenizer_info_.GetVocabSize() << "). Rejecting the token."; + return false; + } + + bool is_stop_token = + std::find(stop_token_ids_.begin(), stop_token_ids_.end(), token_id) != stop_token_ids_.end(); + const auto& special_token_ids = tokenizer_info_.GetSpecialTokenIds(); + if (!is_stop_token && std::find(special_token_ids.begin(), special_token_ids.end(), token_id) != + special_token_ids.end()) { + // Padding ids, i.e. ids in [decoded_vocab.size(), vocab_size), are registered as special ids + // too, but they have no entry in decoded_vocab, so only decode ids backed by a real token. + const auto& decoded_vocab = tokenizer_info_.GetDecodedVocab(); + XGRAMMAR_LOG(WARNING) << "GrammarMatcher cannot accept special token id " << token_id << ": " + << (token_id < static_cast(decoded_vocab.size()) + ? decoded_vocab[token_id] + : "") + << ". Rejecting the token."; + return false; + } + + current_token_index_ = static_cast(token_length_history.size()); + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + bool enforce_budget = budget_enforce_pending_; + bool force_budget_close = budget_force_close_pending_; + if (force_budget_close) { + // Validate on a copy first: a rejected token must not commit the budget-close transition. + Impl trial(*this); + trial.capture_recording_ = false; + bool applied = trial.ApplyBudgetEnforcement(); + XGRAMMAR_DCHECK(applied); + if (!applied) { + return false; + } + trial.budget_enforce_pending_ = false; + trial.budget_force_close_pending_ = false; + if (!trial.AcceptToken(token_id, false)) { + return false; + } + } + + // Capture events are only recorded on the definitive accept path. + CaptureRecordingGuard capture_guard(this); + if (force_budget_close) { + bool applied = ApplyBudgetEnforcement(debug_print); + XGRAMMAR_DCHECK(applied); + if (!applied) { + return false; + } + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + } + SkipExpiredGuard skip_expired_guard(this, enforce_budget && !force_budget_close); + + // The token extends a derivation past its budget iff expired states are allowed to scan + // (no enforcement) while some exist. + bool consumed_past_deadline = false; + if (has_budget_rules_ && !enforce_budget) { + const auto& row = scanable_state_history_[scanable_state_history_.size() - 1]; + for (const auto& state : row) { + if (IsExpiredState(state)) { + consumed_past_deadline = true; + break; + } + } + } + + if (debug_print) { + std::string states_str; + for (const auto& state : GetLatestScanableStates()) { + states_str += " " + state.ToString() + "\n"; + } + XGRAMMAR_LOG(INFO) << "Accepting token id " << token_id << ", string: \"" + << EscapeString(tokenizer_info_.GetDecodedVocab()[token_id]) + << "\", current state:\n" + << states_str; + } + // Handle the stop token + if (is_stop_token) { + bool accepted = AcceptStopToken(); + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The token is an end token. Is accepted: " << accepted; + } + return accepted ? FinishAccept(consumed_past_deadline) : false; + } + + const auto& token = tokenizer_info_.GetDecodedVocab()[token_id]; + int32_t token_char_count = 0; + if (has_char_budget_rules_) { + for (uint8_t byte : token) { + token_char_count += StartsUTF8Codepoint(byte); + } + } + + const int32_t size_before_token = rule_id_to_completable_states_.size(); + std::vector states_before_token; + std::vector> completable_before_token; + std::vector capture_row_before_token; + bool completed_before_token = false; + if (has_char_budget_rules_) { + states_before_token = GetLatestScanableStates(); + auto completable_row_before_token = rule_id_to_completable_states_.Back(); + completable_before_token.assign( + completable_row_before_token.data, + completable_row_before_token.data + completable_row_before_token.data_len + ); + capture_row_before_token = CopyLastCaptureRow(); + completed_before_token = is_completed_.back(); + } + auto restore_row_before_token = [&]() { + if (!has_char_budget_rules_) { + return; + } + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(states_before_token); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(completable_before_token); + is_completed_.back() = completed_before_token; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(capture_row_before_token); + } + }; + + // Phase 1: Try atomic token path (from current state, before byte path) + std::vector atomic_states; + std::vector> atomic_completable; + std::vector atomic_capture_row; + bool atomic_completed = false; + bool atomic_success = + has_char_budget_rules_ + ? AdvanceAtomicTokenWithCharacterBudget(token_id, token_char_count, debug_print) + : AdvanceAtomicToken(token_id, debug_print); + if (atomic_success) { + atomic_states = GetLatestScanableStates(); + auto row = rule_id_to_completable_states_.Back(); + atomic_completable.assign(row.data, row.data + row.data_len); + atomic_capture_row = CopyLastCaptureRow(); + atomic_completed = is_completed_.back(); + PopLastStates(1); + } + restore_row_before_token(); + + // Phase 2: Try byte-by-byte path (from the same original state) + record_char_budget_relaxation_ = true; + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + if (track_temporary_input) { + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + int pos = 0; + bool byte_path_success = true; + for (auto char_value : token) { + bool accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(char_value), debug_print) + : Advance(static_cast(char_value), debug_print); + if (!accepted) { + byte_path_success = false; + break; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(char_value); + } + ++pos; + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + + // Phase 3: Combine results (no priority — merge with deduplication) + if (!byte_path_success && !atomic_success) { + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Token #" << token_id << "<" << EscapeString(token) + << "> rejected at position " << pos; + } + PopLastStates(pos); + restore_row_before_token(); + record_char_budget_relaxation_ = false; + char_budget_relaxed_ = false; + return false; + } + + if (atomic_success && !byte_path_success) { + PopLastStates(pos); + restore_row_before_token(); + char_budget_relaxed_ = false; + bool accepted = + has_char_budget_rules_ + ? AdvanceAtomicTokenWithCharacterBudget(token_id, token_char_count, debug_print) + : AdvanceAtomicToken(token_id, debug_print); + XGRAMMAR_DCHECK(accepted); + token_length_history.push_back(1); + if (ShouldTrackAcceptedBytes()) { + AppendAtomicRow(token); + } + } else if (byte_path_success && !atomic_success) { + token_length_history.push_back(token.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(token); + } + } else { + // Both paths succeeded — merge atomic token states into byte path + if (token.empty()) { + // Zero-length token: byte path created 0 timepoints, just push atomic states + scanable_state_history_.PushBack(atomic_states); + rule_id_to_completable_states_.PushBack(atomic_completable); + is_completed_.push_back(atomic_completed); + PushCaptureRow(atomic_capture_row); + PushCharCountRow(GetCurrentCharIndex(), HasEnteredCharBudget()); + token_length_history.push_back(1); + if (ShouldTrackAcceptedBytes()) { + AppendAtomicRow(token); + } + } else { + auto byte_states = GetLatestScanableStates(); + std::vector merged = byte_states; + StateEqualForParsing state_eq; + for (const auto& s : atomic_states) { + if (std::find_if(merged.begin(), merged.end(), [&](const auto& m) { + return state_eq(m, s); + }) == merged.end()) { + merged.push_back(s); + } + } + + auto byte_row = rule_id_to_completable_states_.Back(); + std::vector> merged_completable( + byte_row.data, byte_row.data + byte_row.data_len + ); + std::vector merged_capture_row = CopyLastCaptureRow(); + bool byte_completed = is_completed_.back(); + int32_t final_char_index = GetCurrentCharIndex(); + bool final_char_budget_entered = HasEnteredCharBudget(); + PopLastStates(1); + + for (const auto& cs : atomic_completable) { + if (std::find_if(merged_completable.begin(), merged_completable.end(), [&](const auto& m) { + return m.first == cs.first && state_eq(m.second, cs.second); + }) == merged_completable.end()) { + merged_completable.push_back(cs); + } + } + + for (auto event : atomic_capture_row) { + // In the atomic path, the after-token position was size_before_token; in the byte path + // it is the row being re-pushed. Remap events that started at the after-token position + // (empty spans) so that they stay empty in the byte-path numbering. + if (event.start_pos == size_before_token) { + event.start_pos = rule_id_to_completable_states_.size(); + } + if (event.occurrence_start_pos == size_before_token) { + event.occurrence_start_pos = rule_id_to_completable_states_.size(); + } + for (auto& target : event.stop_capture_targets) { + if (target.start_pos == size_before_token) { + target.start_pos = rule_id_to_completable_states_.size(); + } + } + auto existing = std::find_if( + merged_capture_row.begin(), + merged_capture_row.end(), + [&](const CaptureEvent& e) { + return e.rule_id == event.rule_id && e.start_pos == event.start_pos && + e.occurrence_start_pos == event.occurrence_start_pos; + } + ); + if (existing == merged_capture_row.end()) { + merged_capture_row.push_back(event); + } else { + existing->hidden_suffix_bytes = + std::max(existing->hidden_suffix_bytes, event.hidden_suffix_bytes); + existing->hidden_stop_bytes = + std::max(existing->hidden_stop_bytes, event.hidden_stop_bytes); + for (const auto& target : event.stop_capture_targets) { + if (std::find( + existing->stop_capture_targets.begin(), + existing->stop_capture_targets.end(), + target + ) == existing->stop_capture_targets.end()) { + existing->stop_capture_targets.push_back(target); + } + } + } + } + + scanable_state_history_.PushBack(merged); + rule_id_to_completable_states_.PushBack(merged_completable); + is_completed_.push_back(byte_completed || atomic_completed); + PushCaptureRow(merged_capture_row); + PushCharCountRow(final_char_index, final_char_budget_entered); + token_length_history.push_back(token.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(token); + } + } + } + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Token #" << token_id << "<" << EscapeString(token) << "> accepted."; + } + if (has_char_budget_rules_) { + // Close exhausted occurrences at the token boundary, including after zero-byte atomic tokens. + ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + } + record_char_budget_relaxation_ = false; + return FinishAccept(consumed_past_deadline); +} + +bool GrammarMatcher::Impl::AcceptString(const std::string& input_str, bool debug_print) { + if (IsStopTokenAccepted()) { + XGRAMMAR_LOG(WARNING) << "The matcher has terminated after accepting the stop token, but is " + << "trying to accept new string \"" << EscapeString(input_str) << "\"."; + return false; + } + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Trying to accept string \"" << EscapeString(input_str) + << "\". Current state:\n" + << PrintStates(); + } + + current_token_index_ = static_cast(token_length_history.size()); + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = true; + + // Capture events are only recorded on the definitive accept path. + CaptureRecordingGuard capture_guard(this); + + std::vector states_before_input; + std::vector> completable_before_input; + std::vector capture_row_before_input; + bool completed_before_input = false; + if (has_char_budget_rules_) { + states_before_input = GetLatestScanableStates(); + auto completable_row_before_input = rule_id_to_completable_states_.Back(); + completable_before_input.assign( + completable_row_before_input.data, + completable_row_before_input.data + completable_row_before_input.data_len + ); + capture_row_before_input = CopyLastCaptureRow(); + completed_before_input = is_completed_.back(); + } + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + if (track_temporary_input) { + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + int accepted_cnt = 0; + for (auto char_value : input_str) { + bool accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(char_value), debug_print) + : Advance(static_cast(char_value), debug_print); + if (!accepted) { + if (debug_print) { + XGRAMMAR_LOG(INFO) << "String \"" << EscapeString(input_str) << "\" is rejected at " + << "position " << accepted_cnt << ", char " << EscapeString(char_value); + } + PopLastStates(accepted_cnt); + if (has_char_budget_rules_) { + scanable_state_history_.PopBack(1); + scanable_state_history_.PushBack(states_before_input); + rule_id_to_completable_states_.PopBack(1); + rule_id_to_completable_states_.PushBack(completable_before_input); + is_completed_.back() = completed_before_input; + if (capture_tracking_) { + capture_event_history_.PopBack(1); + capture_event_history_.PushBack(capture_row_before_input); + } + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + record_char_budget_relaxation_ = false; + char_budget_relaxed_ = false; + return false; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(char_value); + } + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Char " << EscapeString(char_value) << " is accepted. Current state:\n" + << PrintStates(); + } + ++accepted_cnt; + } + token_length_history.push_back(input_str.size()); + if (ShouldTrackAcceptedBytes()) { + AppendPerByteRows(input_str); + } + if (has_char_budget_rules_ && !input_str.empty()) { + // Leave the parser at the enforced boundary even when no later input is supplied. + ApplyCharacterBudgetEnforcementToFixedPoint(debug_print); + } + if (track_temporary_input) { + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + } + record_char_budget_relaxation_ = false; + FinishCharacterBudgetAccept(); + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "String \"" << EscapeString(input_str) << "\" is accepted."; + } + return true; +} + +std::string GrammarMatcher::Impl::PrintBitmask( + int32_t* bitmask_data_ptr, const TokenizerInfo& tokenizer_info +) { + constexpr int kMaxPrintTokens = 100; + std::vector accepted_ids; + std::vector rejected_ids; + auto bitset = + DynamicBitset(tokenizer_info.GetVocabSize(), reinterpret_cast(bitmask_data_ptr)); + for (int i = 0; i < tokenizer_info.GetVocabSize(); ++i) { + if (bitset[i]) { + accepted_ids.push_back(i); + } else { + rejected_ids.push_back(i); + } + } + std::stringstream ss; + ss << "TokenBitmask(num_tokens=" << tokenizer_info.GetVocabSize() + << ", accepted_num=" << accepted_ids.size() << ", rejected_num=" << rejected_ids.size() + << ",\naccepted_ids=" << PrintTokenByIds(accepted_ids, tokenizer_info, kMaxPrintTokens) + << ",\nrejected_ids=" << PrintTokenByIds(rejected_ids, tokenizer_info, kMaxPrintTokens) << ")"; + return ss.str(); +} + +bool GrammarMatcher::Impl::IsTokenBitmaskAllTrue(int32_t* bitmask_data_ptr) { + DynamicBitset next_token_bitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + return next_token_bitset.All(); +} + +bool GrammarMatcher::Impl::FillNextTokenBitmask( + DLTensor* next_token_bitmask, int index, bool debug_print +) { + XGRAMMAR_CHECK(!IsStopTokenAccepted()) + << "GrammarMatcher has terminated after accepting the stop token, but is trying to " + "find the next token mask"; + int32_t* bitmask_data_ptr = + CheckAndGetBitmaskPtr(*next_token_bitmask, tokenizer_info_.GetVocabSize(), index); + current_token_index_ = static_cast(token_length_history.size()); + if (has_budget_rules_) { + const auto& row = scanable_state_history_[scanable_state_history_.size() - 1]; + bool any_expired = false; + bool any_alive = false; + bool can_force_close_without_marker = false; + for (const auto& state : row) { + if (IsExpiredState(state)) { + any_expired = true; + can_force_close_without_marker = + can_force_close_without_marker || + (has_budget_marker_rules_ && CanForceCompleteWithoutMarker(state, false)); + } else { + any_alive = true; + } + } + if (can_force_close_without_marker) { + // Compute the enforcing state on a copy. In addition to dropping expired derivations, + // this lets a budgeted suffix/stop rule end through its body when the marker has not + // appeared yet. The real state is committed only after an allowed token is accepted. + Impl trial(*this); + trial.capture_recording_ = false; + bool applied = trial.ApplyBudgetEnforcement(debug_print); + if (applied) { + trial.FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/false, debug_print); + } + if (applied && (trial.IsCompleted() || trial.BitmaskHasAnyToken(bitmask_data_ptr))) { + budget_enforce_pending_ = true; + budget_force_close_pending_ = true; + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); + } + } + if (any_expired && (any_alive || IsCompleted())) { + // No marker-less completion is needed: retain the original max_tokens fast path. + FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/true, debug_print); + if (BitmaskHasAnyToken(bitmask_data_ptr) || IsCompleted()) { + budget_enforce_pending_ = true; + budget_force_close_pending_ = false; + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); + } + } + // No enforcement at this position: either no budget has expired, or the expired + // derivations cannot end here — relax the budget for one step and report it on the next + // accept. + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + } + FillBitmaskForStates(bitmask_data_ptr, index, /*skip_expired=*/false, debug_print); + return !IsTokenBitmaskAllTrue(bitmask_data_ptr); +} + +void GrammarMatcher::Impl::FillBitmaskForCharBudgetBoundary( + const AdaptiveTokenMask& adaptive_token_mask, int32_t remaining_chars +) { + const auto& token_char_counts = tokenizer_info_.ImplPtr()->GetTokenCharCounts(); + const auto& vocab = tokenizer_info_.GetSortedDecodedVocab(); + + std::vector tokens_to_check = adaptive_token_mask.uncertain_indices; + switch (adaptive_token_mask.store_type) { + case StoreType::kAccepted: + for (int32_t index : adaptive_token_mask.accepted_indices) { + if (token_char_counts[index] <= remaining_chars) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } else { + tokens_to_check.push_back(index); + } + } + break; + case StoreType::kAcceptedBitset: + for (int32_t index = 0; index < static_cast(vocab.size()); ++index) { + int32_t token_id = vocab[index].first; + if (!adaptive_token_mask.accepted_bitset[token_id]) { + continue; + } + if (token_char_counts[index] <= remaining_chars) { + tmp_accepted_bitset_.Set(token_id, true); + } else { + tokens_to_check.push_back(index); + } + } + break; + case StoreType::kRejected: { + std::vector blocked = adaptive_token_mask.rejected_indices; + blocked.insert( + blocked.end(), + adaptive_token_mask.uncertain_indices.begin(), + adaptive_token_mask.uncertain_indices.end() + ); + std::sort(blocked.begin(), blocked.end()); + for (int32_t index = 0; index < static_cast(vocab.size()); ++index) { + if (token_char_counts[index] <= remaining_chars) { + if (!std::binary_search(blocked.begin(), blocked.end(), index)) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } + } else if (!std::binary_search( + adaptive_token_mask.rejected_indices.begin(), + adaptive_token_mask.rejected_indices.end(), + index + )) { + tokens_to_check.push_back(index); + } + } + break; + } + } + + std::sort(tokens_to_check.begin(), tokens_to_check.end()); + tokens_to_check.erase( + std::unique(tokens_to_check.begin(), tokens_to_check.end()), tokens_to_check.end() + ); + + int32_t saved_temporary_input_start_row = temporary_input_start_row_; + std::string saved_temporary_input_bytes = std::move(temporary_input_bytes_); + std::vector all_latest_states = GetLatestScanableStates(); + bool all_latest_completed = IsCompleted(); + + PushStatesToCheck(all_latest_states, all_latest_completed); + int32_t initial_state_history_size = scanable_state_history_.size(); + int32_t previous_index = -1; + int32_t previous_matched_size = 0; + temporary_input_start_row_ = initial_state_history_size - 1; + temporary_input_bytes_.clear(); + std::vector byte_rejected_tokens; + + for (int32_t index : tokens_to_check) { + const std::string& token = vocab[index].second; + int32_t common_prefix_length = 0; + if (previous_index >= 0) { + common_prefix_length = std::mismatch( + vocab[previous_index].second.begin(), + vocab[previous_index].second.end(), + token.begin(), + token.end() + ) + .first - + vocab[previous_index].second.begin(); + common_prefix_length = std::min(common_prefix_length, previous_matched_size); + } + PopLastStates( + scanable_state_history_.size() - initial_state_history_size - common_prefix_length + ); + temporary_input_bytes_.resize(common_prefix_length); + + bool accepted = !token.empty(); + int32_t matched_size = common_prefix_length; + for (; matched_size < static_cast(token.size()); ++matched_size) { + if (!AdvanceWithCharacterBudget(static_cast(token[matched_size]))) { + accepted = false; + break; + } + temporary_input_bytes_.push_back(token[matched_size]); + } + if (accepted) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + } else { + byte_rejected_tokens.push_back(index); + } + previous_index = index; + previous_matched_size = matched_size; + } + PopLastStates(scanable_state_history_.size() - initial_state_history_size + 1); + + for (int32_t index : byte_rejected_tokens) { + PushStatesToCheck(all_latest_states, all_latest_completed); + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + if (AdvanceAtomicTokenWithCharacterBudget(vocab[index].first, token_char_counts[index])) { + tmp_accepted_bitset_.Set(vocab[index].first, true); + PopLastStates(1); + } + PopLastStates(1); + } + + temporary_input_start_row_ = saved_temporary_input_start_row; + temporary_input_bytes_ = std::move(saved_temporary_input_bytes); +} + +void GrammarMatcher::Impl::FillBitmaskForStates( + int32_t* bitmask_data_ptr, int index, bool skip_expired, bool debug_print +) { + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + const auto& subtree_range = tokenizer_info_.GetTrieSubtreeNodesRange(); + const auto& adaptive_token_mask_cache = compiled_grammar_->adaptive_token_mask_cache; + // We need to have a copy, because scanable_state_history_ will be modified during the + // FillNextTokenBitmask process, which can lead to undefined behavior. + std::vector latest_states; + for (const auto& state : scanable_state_history_[scanable_state_history_.size() - 1]) { + if (skip_expired && IsExpiredState(state)) { + continue; + } + latest_states.push_back(state); + } + + // We check all the latest states of the earley parser, and check all the masks of the leaf + // states. The final accepted token set is the union of the accepted token sets of all leaf + // states. The final rejected token set is the intersection of the rejected token sets of all leaf + // states. + + // Note these indices store the indices in sorted_decoded_vocab, instead of the token ids. + tmp_accepted_bitset_.Reset(); + // {-1} means the universal set, i.e. all tokens initially + tmp_rejected_indices_.assign({-1}); + + if (debug_print) { + XGRAMMAR_LOG(INFO) << "FillNextTokenBitmask: index=" << index + << ", num of states=" << latest_states.size(); + } + + std::vector> + latest_states_with_masks; + + for (const auto& state : latest_states) { + auto adaptive_token_mask_it = adaptive_token_mask_cache.find(state); + XGRAMMAR_CHECK(adaptive_token_mask_it != adaptive_token_mask_cache.end()) << state; + const auto& adaptive_token_mask = adaptive_token_mask_it->second; + if (state.char_budget_deadline >= 0) { + int32_t remaining_chars = state.char_budget_deadline - GetCurrentCharIndex(); + if (remaining_chars <= tokenizer_info_.ImplPtr()->GetMaxTokenChars()) { + FillBitmaskForCharBudgetBoundary(adaptive_token_mask, std::max(remaining_chars, 0)); + continue; + } + } + latest_states_with_masks.push_back(std::make_pair(state, adaptive_token_mask_it)); + if (adaptive_token_mask.store_type == StoreType::kAcceptedBitset) { + tmp_accepted_bitset_ |= adaptive_token_mask.accepted_bitset; + } else if (adaptive_token_mask.store_type == StoreType::kAccepted) { + for (auto idx : adaptive_token_mask.accepted_indices) { + tmp_accepted_bitset_.Set(sorted_decoded_vocab[idx].first, true); + } + } + } + + const auto& token_char_counts = tokenizer_info_.ImplPtr()->GetTokenCharCounts(); + for (const auto& [state, adaptive_token_mask_it] : latest_states_with_masks) { + const auto& adaptive_token_mask = adaptive_token_mask_it->second; + + // For each ParserState, we will check every uncertain token and put them into the accepted or + // rejected list. + + // Step 2. Update the accepted tokens in accepted_indices_delta, or the rejected tokens in + // rejected_indices_delta. + + // If the accepted tokens are saved, it means it is likely to be smaller than the rejected + // tokens, so we will just find the accepted tokens, and vice versa. + + tmp_rejected_indices_delta_.clear(); + + // Examine only the current one ParserState + PushOneStateToCheck(state); + bool track_temporary_input = has_char_budget_rules_ && has_budget_marker_rules_; + int32_t saved_temporary_input_start_row = -1; + std::string saved_temporary_input_bytes; + if (track_temporary_input) { + saved_temporary_input_start_row = temporary_input_start_row_; + saved_temporary_input_bytes = std::move(temporary_input_bytes_); + temporary_input_start_row_ = scanable_state_history_.size() - 1; + temporary_input_bytes_.clear(); + } + + const std::string* prev_token = nullptr; + int prev_matched_size = 0; + if (debug_print) { + XGRAMMAR_LOG(INFO) << "The ParserState is " << state << ", the mask is " + << adaptive_token_mask.Print(tokenizer_info_); + } + int last_rejected_uncertain_range = 0; + const auto& uncertain_indices = adaptive_token_mask.uncertain_indices; + for (auto token_it = uncertain_indices.begin(); token_it != uncertain_indices.end();) { + const auto cur_token_idx = *token_it++; + // Check if the current token is already accepted. If it is, we can skip it. + if (tmp_accepted_bitset_[sorted_decoded_vocab[cur_token_idx].first]) { + continue; + } + + // Check if the current token is in the rejected range. i.e. check if the current token + // is on the subtree of the rejected token. + if (cur_token_idx < last_rejected_uncertain_range) { + if (adaptive_token_mask.store_type == StoreType::kRejected) { + tmp_rejected_indices_delta_.push_back(cur_token_idx); + } else { + // Accepted-list/bitset masks need no per-token record for a rejected subtree. + // Sorted, distinct indices bound the number of remaining tokens in this subtree. + const auto search_size = std::min( + uncertain_indices.end() - token_it, last_rejected_uncertain_range - cur_token_idx - 1 + ); + token_it = + std::lower_bound(token_it, token_it + search_size, last_rejected_uncertain_range); + } + continue; + } + + const auto& cur_token = sorted_decoded_vocab[cur_token_idx].second; + bool accepted = !cur_token.empty() || !has_char_budget_rules_; + + // Step 2.1. Find the longest common prefix with the accepted part of the previous token. + // We can reuse the previous matched size to avoid unnecessary matching. + if (prev_token) { + int lcp_len = std::mismatch( + cur_token.begin(), cur_token.end(), prev_token->begin(), prev_token->end() + ) + .first - + cur_token.begin(); + if (lcp_len > prev_matched_size) { + last_rejected_uncertain_range = subtree_range[cur_token_idx]; + accepted = false; + } else if (lcp_len < prev_matched_size) { + PopLastStates(prev_matched_size - lcp_len); + if (track_temporary_input) { + temporary_input_bytes_.resize(lcp_len); + } + } + prev_matched_size = std::min(prev_matched_size, lcp_len); + } + + // Step 2.2. Find if the current token is accepted or rejected. + if (accepted) { + for (int j = prev_matched_size; j < static_cast(cur_token.size()); ++j) { + bool byte_accepted = has_char_budget_rules_ + ? AdvanceWithCharacterBudget(static_cast(cur_token[j])) + : Advance(static_cast(cur_token[j])); + if (!byte_accepted) { + last_rejected_uncertain_range = subtree_range[cur_token_idx]; + accepted = false; + break; + } + if (track_temporary_input) { + temporary_input_bytes_.push_back(cur_token[j]); + } + prev_matched_size = j + 1; + } + } + + bool retried_atomically = false; + if (!accepted && has_char_budget_rules_ && has_token_edges_) { + // Retry through Token/ExcludeToken edges in place: rebuild the single-state row, run the + // atomic path, then rebuild it again so trial-time enforcement does not leak. A full + // matcher copy here made every mask O(generated length). + retried_atomically = true; + PopLastStates(prev_matched_size + 1); + if (track_temporary_input) { + temporary_input_bytes_.clear(); + } + PushOneStateToCheck(state); + const bool saved_capture_recording = capture_recording_; + const bool saved_record_relaxation = record_char_budget_relaxation_; + capture_recording_ = false; + record_char_budget_relaxation_ = false; + accepted = AdvanceAtomicTokenWithCharacterBudget( + sorted_decoded_vocab[cur_token_idx].first, token_char_counts[cur_token_idx] + ); + capture_recording_ = saved_capture_recording; + record_char_budget_relaxation_ = saved_record_relaxation; + PopLastStates(accepted ? 2 : 1); + PushOneStateToCheck(state); + prev_matched_size = 0; + if (accepted) { + last_rejected_uncertain_range = cur_token_idx + 1; + } + } + + // Step 2.3. Push the result to the delta list. + if (adaptive_token_mask.store_type == StoreType::kAcceptedBitset || + adaptive_token_mask.store_type == StoreType::kAccepted) { + if (accepted) { + tmp_accepted_bitset_.Set(sorted_decoded_vocab[cur_token_idx].first, true); + } + } else { + if (!accepted) { + tmp_rejected_indices_delta_.push_back(cur_token_idx); + } + } + + // The retry dropped this token's rows, so its prefix cannot be reused. + prev_token = retried_atomically ? nullptr : &cur_token; + } + + PopLastStates(prev_matched_size + 1); + if (track_temporary_input) { + temporary_input_start_row_ = saved_temporary_input_start_row; + temporary_input_bytes_ = std::move(saved_temporary_input_bytes); + } + // Step 3. Update the accepted_indices or rejected_indices + if (adaptive_token_mask.store_type == StoreType::kRejected) { + // rejected_indices = Intersect( + // rejected_indices, + // adaptive_token_mask.rejected_indices + rejected_indices_delta) + IntsetUnion(&tmp_rejected_indices_delta_, adaptive_token_mask.rejected_indices); + IntsetIntersection(&tmp_rejected_indices_, tmp_rejected_indices_delta_); + } + } + + // Finally update the rejected_ids bitset + bool can_reach_end = IsCompleted(); + SetTokenBitmask( + bitmask_data_ptr, tmp_accepted_bitset_, tmp_rejected_indices_, can_reach_end, false + ); + if (debug_print) { + XGRAMMAR_LOG(INFO) << "Filled bitmask: " << PrintBitmask(bitmask_data_ptr, tokenizer_info_); + } +} + +std::string GrammarMatcher::Impl::FindJumpForwardString() { + XGRAMMAR_CHECK(!IsStopTokenAccepted()) + << "GrammarMatcher has terminated after accepting the stop token, but is trying to " + "get the jump forward string"; + + current_token_index_ = static_cast(token_length_history.size()); + if (budget_force_close_pending_) { + Impl trial(*this); + trial.capture_recording_ = false; + if (trial.ApplyBudgetEnforcement()) { + trial.budget_enforce_pending_ = false; + trial.budget_force_close_pending_ = false; + return trial.FindJumpForwardString(); + } + } + + std::string result; + int num_accepted_chars = 0; + bool can_find_next_char = true; + + while (can_find_next_char) { + const auto& states = scanable_state_history_[scanable_state_history_.size() - 1]; + + // The state comes to the end of the grammar + if (IsCompleted()) { + can_find_next_char = false; + break; + } + + // 1. Check that for every leaf ParserState, the next possible char is unique and the same + // -1 means not found yet; 0~255 means the next char + int next_char = -1; + for (const auto& state : states) { + if (budget_enforce_pending_ && IsExpiredState(state)) { + continue; + } + XGRAMMAR_DCHECK(state.rule_id != -1 && grammar_->per_rule_fsms[state.rule_id].has_value()); + const auto& fsm = grammar_->per_rule_fsms[state.rule_id].value(); + const auto& current_edges = fsm.GetFsm().GetFsm().GetEdges(state.element_id); + for (const auto& edge : current_edges) { + if (!edge.IsCharRange()) { + continue; + } + if (edge.min != edge.max) { + can_find_next_char = false; + break; + } + if (next_char == -1) { + next_char = edge.min; + } else if (next_char != edge.min) { + can_find_next_char = false; + break; + } + } + } + + if (next_char == -1) { + can_find_next_char = false; + } + + // 2. If found, accept the char and iterate to the next position + if (can_find_next_char) { + if (StartsUTF8Codepoint(static_cast(next_char)) && + std::any_of(states.begin(), states.end(), [&](const ParserState& state) { + return IsCharExpiredState(state); + })) { + break; + } + result += static_cast(next_char); + Advance(next_char); + ++num_accepted_chars; + } + } + + // Rollback all chars accepted + PopLastStates(num_accepted_chars); + return result; +} + +void GrammarMatcher::Impl::Rollback(int num_tokens) { + XGRAMMAR_CHECK(num_tokens <= static_cast(token_length_history.size())) + << "Intended to rollback " << num_tokens << " tokens, but only the last " + << token_length_history.size() << " steps of history are saved"; + while (num_tokens > 0) { + int steps = token_length_history.back(); + PopLastStates(steps); + token_length_history.pop_back(); + if (ShouldTrackAcceptedBytes() && steps > 0) { + row_byte_end_.resize(row_byte_end_.size() - steps); + accepted_bytes_.resize(row_byte_end_.back()); + } + --num_tokens; + } + budget_enforce_pending_ = false; + budget_force_close_pending_ = false; + char_budget_relaxed_ = false; + record_char_budget_relaxation_ = false; + temporary_input_start_row_ = -1; + temporary_input_bytes_.clear(); + budget_body_match_cache_.clear(); +} + +std::vector> GrammarMatcher::Impl::GetCaptures(bool deduplicate +) const { + std::vector> result; + if (!IsCaptureTrackingEnabled()) { + return result; + } + XGRAMMAR_DCHECK(capture_event_history_.size() == static_cast(row_byte_end_.size())) + << "The capture history is not aligned with the byte history: " + << capture_event_history_.size() << " vs " << row_byte_end_.size(); + + // Flatten the event history. Events are ordered by completion position, and by completion + // order within a position. + struct FlatEvent { + int32_t rule_id; + int32_t start_row; + int32_t occurrence_start_pos; + int32_t end_row; + int32_t hidden_suffix_bytes; + int32_t hidden_stop_bytes; + std::vector stop_capture_targets; + bool marker_present = false; + int64_t marker_start_byte = -1; + }; + std::vector events; + for (int32_t row = 0; row < capture_event_history_.size(); ++row) { + for (const auto& event : capture_event_history_[row]) { + int32_t start_row = event.start_pos == ParserState::kNoPrevInputPos ? 0 : event.start_pos; + XGRAMMAR_DCHECK(start_row <= row); + events.push_back( + {event.rule_id, + start_row, + event.occurrence_start_pos, + row, + event.hidden_suffix_bytes, + event.hidden_stop_bytes, + event.stop_capture_targets, + false, + -1} + ); + } + } + + // The Earley parser explores all parse hypotheses in parallel, so one occurrence of a rule + // (identified by its rule id and start position) may complete at several candidate end + // positions before the following input decides the real one. When deduplicate is true, we + // keep only the last (longest) completion of each occurrence. + std::vector keep(events.size(), true); + if (deduplicate) { + std::unordered_map last_index; + for (size_t i = 0; i < events.size(); ++i) { + int64_t key = (static_cast(events[i].rule_id) << 32) | + static_cast(events[i].start_row); + auto it = last_index.find(key); + if (it != last_index.end()) { + keep[it->second] = false; + it->second = i; + } else { + last_index[key] = i; + } + } + } + + auto contains_end_state = [](const CompactFSMWithStartEnd& fsm, + const std::unordered_set& states) { + return std::any_of(states.begin(), states.end(), [&](int state) { + return fsm.IsEndState(state); + }); + }; + + auto accepted_prefixes = [&](const CompactFSMWithStartEnd& fsm, int64_t begin, int64_t end) { + XGRAMMAR_DCHECK(fsm.IsLeaf()) << "A suffix/stop capture helper must compile to a leaf FSM"; + std::vector result(static_cast(end - begin + 1), false); + std::unordered_set states{fsm.GetStart()}; + fsm.GetFsm().GetEpsilonClosure(&states); + result[0] = contains_end_state(fsm, states); + std::unordered_set next_states; + for (int64_t offset = 0; offset < end - begin; ++offset) { + fsm.GetFsm().Advance( + states, accepted_bytes_[begin + offset], &next_states, FSMEdge::EdgeType::kCharRange, true + ); + states = next_states; + result[static_cast(offset + 1)] = contains_end_state(fsm, states); + } + return result; + }; + + // Compute which suffixes of [begin, end) are accepted by marker_fsm. This is a reverse NFA + // traversal, so locating a delimiter remains linear in the captured byte length even when the + // body (for example /.*/) can end at every byte. + auto accepted_suffixes = [&](const CompactFSMWithStartEnd& marker_fsm, int64_t begin, int64_t end + ) { + XGRAMMAR_DCHECK(marker_fsm.IsLeaf()) + << "A suffix/stop marker helper must compile to a leaf FSM"; + const auto& fsm = marker_fsm.GetFsm(); + std::unordered_set reachable; + marker_fsm.GetReachableStates(&reachable); + struct ReverseCharEdge { + int32_t source; + int32_t min; + int32_t max; + }; + std::vector> reverse_epsilon(fsm.NumStates()); + std::vector> reverse_char(fsm.NumStates()); + for (int32_t source : reachable) { + for (const auto& edge : fsm.GetEdges(source)) { + if (!reachable.count(edge.target)) { + continue; + } + if (edge.IsEpsilon()) { + reverse_epsilon[edge.target].push_back(source); + } else if (edge.IsCharRange()) { + reverse_char[edge.target].push_back({source, edge.min, edge.max}); + } + } + } + auto add_reverse_epsilon_closure = [&](std::unordered_set* states) { + std::vector queue(states->begin(), states->end()); + for (size_t i = 0; i < queue.size(); ++i) { + for (int32_t predecessor : reverse_epsilon[queue[i]]) { + if (states->insert(predecessor).second) { + queue.push_back(predecessor); + } + } + } + }; + + std::vector result(static_cast(end - begin + 1), false); + std::unordered_set states(marker_fsm.GetEnds().begin(), marker_fsm.GetEnds().end()); + add_reverse_epsilon_closure(&states); + result[static_cast(end - begin)] = states.count(marker_fsm.GetStart()) != 0; + std::unordered_set previous_states; + for (int64_t offset = end - begin; offset > 0; --offset) { + previous_states.clear(); + int32_t byte = accepted_bytes_[begin + offset - 1]; + for (int32_t target : states) { + for (const auto& edge : reverse_char[target]) { + if (edge.min <= byte && byte <= edge.max) { + previous_states.insert(edge.source); + } + } + } + add_reverse_epsilon_closure(&previous_states); + states = previous_states; + result[static_cast(offset - 1)] = states.count(marker_fsm.GetStart()) != 0; + } + return result; + }; + + // Resolve every marker to a concrete byte span once. Fixed strings use their recorded length; + // regexes and named terminals use helper FSMs for both sides of the original body/marker split. + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i]) { + continue; + } + auto& event = events[i]; + int32_t hidden_bytes = std::max(event.hidden_suffix_bytes, event.hidden_stop_bytes); + if (hidden_bytes <= 0) { + continue; + } + const auto& rule = grammar_->GetRule(event.rule_id); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(event.rule_id); + XGRAMMAR_DCHECK(suffix_stop_info != nullptr); + int64_t event_start = row_byte_end_[event.start_row]; + int64_t event_end = row_byte_end_[event.end_row]; + event.marker_present = true; + if (suffix_stop_info->body_rule_id == -1) { + event.marker_start_byte = + std::max(event_start, event_end - static_cast(hidden_bytes)); + continue; + } + + XGRAMMAR_DCHECK(suffix_stop_info->body_rule_id >= 0 && suffix_stop_info->marker_rule_id >= 0); + XGRAMMAR_DCHECK( + grammar_->per_rule_fsms[suffix_stop_info->body_rule_id].has_value() && + grammar_->per_rule_fsms[suffix_stop_info->marker_rule_id].has_value() + ); + const auto& body_fsm = grammar_->per_rule_fsms[suffix_stop_info->body_rule_id]->GetFsm(); + const auto& marker_fsm = grammar_->per_rule_fsms[suffix_stop_info->marker_rule_id]->GetFsm(); + std::vector body_ends = accepted_prefixes(body_fsm, event_start, event_end); + std::vector marker_starts = accepted_suffixes(marker_fsm, event_start, event_end); + bool found = false; + for (size_t offset = 0; offset < body_ends.size(); ++offset) { + if (body_ends[offset] && marker_starts[offset]) { + event.marker_start_byte = event_start + static_cast(offset); + found = true; + break; + } + } + XGRAMMAR_DCHECK(found) << "Could not recover the suffix/stop marker boundary for rule " + << rule.name; + if (!found) { + event.marker_present = false; + event.marker_start_byte = -1; + } + } + + // Associate each stop marker only with captured occurrences reached through its concrete + // Earley parent chain. Failed or unrelated branches may cover identical byte ranges, so byte + // overlap alone is not a valid indication that a marker belongs to a capture. + std::unordered_map>> + hidden_stop_spans_by_capture; + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i] || events[i].hidden_stop_bytes <= 0 || !events[i].marker_present) { + continue; + } + int64_t event_end = row_byte_end_[events[i].end_row]; + int64_t hidden_start = events[i].marker_start_byte; + if (hidden_start < event_end) { + for (const auto& target : events[i].stop_capture_targets) { + int64_t target_key = + (static_cast(target.rule_id) << 32) | static_cast(target.start_pos); + hidden_stop_spans_by_capture[target_key].emplace_back(hidden_start, event_end); + } + } + } + for (auto& [_, spans] : hidden_stop_spans_by_capture) { + std::sort(spans.begin(), spans.end()); + } + + for (size_t i = 0; i < events.size(); ++i) { + if (!keep[i]) { + continue; + } + const auto& event = events[i]; + const auto& rule = grammar_->GetRule(event.rule_id); + const auto* suffix_stop_info = grammar_->GetSuffixStopInfo(event.rule_id); + int64_t capture_start = row_byte_end_[event.start_row]; + int64_t capture_end = row_byte_end_[event.end_row]; + if (event.marker_present && suffix_stop_info != nullptr && + !suffix_stop_info->stop_capture_name.empty()) { + std::string marker_capture; + if (event.marker_start_byte < capture_end) { + marker_capture.assign( + reinterpret_cast(accepted_bytes_.data() + event.marker_start_byte), + static_cast(capture_end - event.marker_start_byte) + ); + } + result.emplace_back(suffix_stop_info->stop_capture_name, std::move(marker_capture)); + } + if (rule.capture_name.empty()) { + continue; + } + if (event.hidden_suffix_bytes > 0 && event.marker_present) { + capture_end = event.marker_start_byte; + } + + std::string capture; + int64_t cursor = capture_start; + auto append_bytes = [&](int64_t begin, int64_t end) { + if (begin < end) { + capture.append( + reinterpret_cast(accepted_bytes_.data() + begin), + static_cast(end - begin) + ); + } + }; + int64_t capture_key = (static_cast(event.rule_id) << 32) | + static_cast(event.occurrence_start_pos); + auto hidden_spans_it = hidden_stop_spans_by_capture.find(capture_key); + if (hidden_spans_it != hidden_stop_spans_by_capture.end()) { + for (const auto& hidden : hidden_spans_it->second) { + if (hidden.second <= capture_start) { + continue; + } + if (hidden.first >= capture_end) { + break; + } + int64_t hidden_begin = std::max(hidden.first, capture_start); + int64_t hidden_end = std::min(hidden.second, capture_end); + append_bytes(cursor, hidden_begin); + cursor = std::max(cursor, hidden_end); + } + } + append_bytes(cursor, capture_end); + result.emplace_back(rule.capture_name, std::move(capture)); + } + return result; +} + +void GrammarMatcher::Impl::SetTokenBitmask( + int32_t* bitmask_data_ptr, + const DynamicBitset& accepted_bitset, + const std::vector& rejected_indices, + bool can_reach_end, + bool allow_special_token +) { + // next_token_bitmask = set(all accepted tokens) = + // 1. all_tokens - (rejected_ids / accepted_ids) + // (when rejected_ids != {-1}, i.e. rejected_ids is not the universal set) + // 2. accepted_ids + // (otherwise, when rejected_ids is the universal set) + DynamicBitset next_token_bitset( + tokenizer_info_.GetVocabSize(), reinterpret_cast(bitmask_data_ptr) + ); + const auto& sorted_decoded_vocab = tokenizer_info_.GetSortedDecodedVocab(); + + if (rejected_indices.size() == 1 && rejected_indices[0] == -1) { + // If rejected_indices is the universal set, the final accepted token set is just + // accepted_indices + next_token_bitset = accepted_bitset; + + if (allow_special_token) { + for (int id : tokenizer_info_.GetSpecialTokenIds()) { + next_token_bitset.Set(id, true); + } + } + } else { + // Otherwise, the final rejected token set is (rejected_indices \ accepted_indices) + next_token_bitset.Set(); + + for (auto i : rejected_indices) { + auto id = sorted_decoded_vocab[i].first; + if (!accepted_bitset[id]) { + next_token_bitset.Set(id, false); + } + } + if (!allow_special_token) { + for (int id : tokenizer_info_.GetSpecialTokenIds()) { + next_token_bitset.Set(id, false); + } + } + } + // Overridden stop tokens can occur in the compiler's accepted token set. + // Apply their termination policy after constructing either mask representation. + for (int id : stop_token_ids_) { + next_token_bitset.Set(id, can_reach_end); + } +} + +int GrammarMatcher::Impl::GetNextUncertainToken( + bool is_uncertain_saved, + int* iterator_uncertain, + const std::vector& uncertain_indices, + const std::vector& uncertain_tokens_bitset +) { + if (is_uncertain_saved) { + ++*iterator_uncertain; + if (*iterator_uncertain == static_cast(uncertain_indices.size())) { + return -1; + } + return uncertain_indices[*iterator_uncertain]; + } else { + ++*iterator_uncertain; + while (*iterator_uncertain < static_cast(uncertain_tokens_bitset.size()) && + !uncertain_tokens_bitset[*iterator_uncertain]) { + ++*iterator_uncertain; + } + if (*iterator_uncertain == static_cast(uncertain_tokens_bitset.size())) { + return -1; + } + return *iterator_uncertain; + } +} + +void BatchGrammarMatcher::Impl::BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print +) { + XGRAMMAR_CHECK(!indices.has_value() || indices->size() == matchers->size()) + << "The size of indices (" << (indices.has_value() ? indices->size() : 0) + << ") should be the same as the size of matchers (" << matchers->size() << ")."; + // Initialize the thread pool if needed. It should be initialized each time, + // because ThreadPool cannot be reused after Join(). + if (max_threads_ > 1) { + thread_pool_.emplace(max_threads_); + } + if (!thread_pool_.has_value()) { + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + int index = indices.has_value() ? (*indices)[i] : i; + XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0]) + << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0] + << ") for batch_id " << i << "."; + matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print); + } + } else { + auto fill_next_token_mask = [&](int32_t batch_id) { + auto& matcher = (*matchers)[batch_id]; + int index = indices.has_value() ? (*indices)[batch_id] : batch_id; + XGRAMMAR_CHECK(index >= 0 && index < next_token_bitmask->shape[0]) + << "The index " << index << " is out of range [0, " << next_token_bitmask->shape[0] + << ") for batch_id " << batch_id << "."; + matcher->FillNextTokenBitmask(next_token_bitmask, index, debug_print); + }; + for (int i = 0; i < static_cast(matchers->size()); i++) { + thread_pool_->Execute([fill_next_token_mask, i]() { fill_next_token_mask(i); }); + } + thread_pool_->Join(); + } +} + +std::vector BatchGrammarMatcher::Impl::BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print +) { + XGRAMMAR_CHECK(matchers->size() == input_strs.size()) + << "The size of matchers (" << matchers->size() << ") and input_strs (" << input_strs.size() + << ") should be the same."; + std::vector accepted(matchers->size()); + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + accepted[i] = matcher->AcceptString(input_strs[i], debug_print); + } + return accepted; +} + +std::vector BatchGrammarMatcher::Impl::BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print +) { + XGRAMMAR_CHECK(matchers->size() == token_ids.size()) + << "The size of matchers (" << matchers->size() << ") and token_ids (" << token_ids.size() + << ") should be the same."; + std::vector accepted(matchers->size()); + for (int i = 0; i < static_cast(matchers->size()); i++) { + auto& matcher = (*matchers)[i]; + accepted[i] = matcher->AcceptToken(token_ids[i], debug_print); + } + return accepted; +} + +void BatchGrammarMatcher::Impl::BatchRollback( + std::vector* matchers, const std::vector& num_tokens +) { + XGRAMMAR_CHECK(matchers->size() == num_tokens.size()) + << "The size of matchers (" << matchers->size() << ") and num_tokens (" << num_tokens.size() + << ") should be the same."; + for (int i = 0; i < static_cast(matchers->size()); i++) { + (*matchers)[i].Rollback(num_tokens[i]); + } +} + +GrammarMatcher::GrammarMatcher( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens, + bool terminate_without_stop_token, + int max_rollback_tokens, + std::optional default_temperature +) + : pimpl_(std::make_shared( + compiled_grammar, + override_stop_tokens, + terminate_without_stop_token, + max_rollback_tokens, + default_temperature + )) {} + +bool GrammarMatcher::AcceptToken(int32_t token_id, bool debug_print) { + return pimpl_->AcceptToken(token_id, debug_print); +} + +bool GrammarMatcher::AcceptString(const std::string& input_str, bool debug_print) { + return pimpl_->AcceptString(input_str, debug_print); +} + +bool GrammarMatcher::FillNextTokenBitmask( + DLTensor* next_token_bitmask, int index, bool debug_print +) { + return pimpl_->FillNextTokenBitmask(next_token_bitmask, index, debug_print); +} + +bool GrammarMatcher::TraverseDraftTree( + const DLTensor* retrieve_next_token, + const DLTensor* retrieve_next_sibling, + const DLTensor* draft_tokens, + DLTensor* token_bitmask, + double time_threshold, + DLTensor* temperatures +) { + auto check_cpu = [](const DLTensor* tensor, const char* name) { + XGRAMMAR_CHECK( + tensor->device.device_type == kDLCPU || tensor->device.device_type == kDLCUDAHost || + tensor->device.device_type == kDLROCMHost + ) << "The " + << name << " tensor must be on CPU"; + }; + + XGRAMMAR_CHECK( + retrieve_next_token->ndim == 1 && retrieve_next_token->dtype.code == kDLInt && + retrieve_next_token->dtype.bits == 64 + ) << "The retrieve_next_token tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + retrieve_next_sibling->ndim == 1 && retrieve_next_sibling->dtype.code == kDLInt && + retrieve_next_sibling->dtype.bits == 64 + ) << "The retrieve_next_sibling tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + draft_tokens->ndim == 1 && draft_tokens->dtype.code == kDLInt && + draft_tokens->dtype.bits == 64 + ) << "The draft_tokens tensor must be a 1D int64 tensor"; + XGRAMMAR_CHECK( + token_bitmask->ndim == 2 && token_bitmask->dtype.code == kDLInt && + token_bitmask->dtype.bits == 32 + ) << "The token_bitmask tensor must be a 2D int32 tensor"; + if (temperatures != nullptr) { + XGRAMMAR_CHECK( + temperatures->ndim == 1 && temperatures->dtype.code == kDLFloat && + temperatures->dtype.bits == 32 + ) << "The temperatures tensor must be a 1D float32 tensor"; + } + + check_cpu(retrieve_next_token, "retrieve_next_token"); + check_cpu(retrieve_next_sibling, "retrieve_next_sibling"); + check_cpu(draft_tokens, "draft_tokens"); + check_cpu(token_bitmask, "token_bitmask"); + if (temperatures != nullptr) { + check_cpu(temperatures, "temperatures"); + } + + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == retrieve_next_sibling->shape[0]) + << "The retrieve_next_token and retrieve_next_sibling tensors must have the same length"; + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == draft_tokens->shape[0]) + << "The retrieve_next_token and draft_tokens tensors must have the same length"; + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == token_bitmask->shape[0]) + << "The token_bitmask batch size must match the number of nodes in the tree"; + if (temperatures != nullptr) { + XGRAMMAR_CHECK(retrieve_next_token->shape[0] == temperatures->shape[0]) + << "The temperatures size must match the number of nodes in the tree"; + std::fill_n(reinterpret_cast(temperatures->data), temperatures->shape[0], -1.0f); + } + XGRAMMAR_CHECK(retrieve_next_sibling->shape[0] > 0 && retrieve_next_sibling->data != nullptr) + << "The draft tree must not be empty"; + XGRAMMAR_CHECK(reinterpret_cast(retrieve_next_sibling->data)[0] == -1) + << "The root node must not have siblings"; + + // The traversal follows retrieve_next_token / retrieve_next_sibling as node indices; a value + // outside [-1, num_nodes) would recurse into an out-of-bounds position, so validate them upfront. + // The edges must also form a tree: the root is never a target, and every other node is the + // child or next sibling of at most one node. A node targeted twice (a cycle or a shared node) + // would otherwise make the traversal recurse without bound. + int64_t num_nodes = retrieve_next_token->shape[0]; + const int64_t* next_token_data = reinterpret_cast(retrieve_next_token->data); + const int64_t* next_sibling_data = reinterpret_cast(retrieve_next_sibling->data); + std::vector in_degree(num_nodes, 0); + for (int64_t i = 0; i < num_nodes; ++i) { + XGRAMMAR_CHECK(next_token_data[i] >= -1 && next_token_data[i] < num_nodes) + << "retrieve_next_token[" << i << "] = " << next_token_data[i] + << " is out of bounds: it should be in [-1, " << num_nodes << ")."; + XGRAMMAR_CHECK(next_sibling_data[i] >= -1 && next_sibling_data[i] < num_nodes) + << "retrieve_next_sibling[" << i << "] = " << next_sibling_data[i] + << " is out of bounds: it should be in [-1, " << num_nodes << ")."; + if (next_token_data[i] != -1) { + ++in_degree[next_token_data[i]]; + } + if (next_sibling_data[i] != -1) { + ++in_degree[next_sibling_data[i]]; + } + } + XGRAMMAR_CHECK(in_degree[0] == 0) << "The root node must not be the child or sibling of a node"; + for (int64_t i = 1; i < num_nodes; ++i) { + XGRAMMAR_CHECK(in_degree[i] <= 1) + << "Node " << i << " is the target of " << in_degree[i] + << " retrieve_next_token / retrieve_next_sibling entries; the draft tree must not contain " + "cycles or shared nodes."; + } + + return details::TraverseDraftTreeRecursive( + 0, + -1, + reinterpret_cast(retrieve_next_token->data), + reinterpret_cast(retrieve_next_sibling->data), + reinterpret_cast(draft_tokens->data), + *this, + token_bitmask, + temperatures == nullptr ? nullptr : reinterpret_cast(temperatures->data), + time_threshold, + details::Clock::now() + ); +} + +std::string GrammarMatcher::FindJumpForwardString() { return pimpl_->FindJumpForwardString(); } + +void GrammarMatcher::Rollback(int num_tokens) { pimpl_->Rollback(num_tokens); } + +std::vector> GrammarMatcher::GetCaptures(bool deduplicate +) const { + return pimpl_->GetCaptures(deduplicate); +} + +bool GrammarMatcher::IsTerminated() const { return pimpl_->IsTerminated(); } + +bool GrammarMatcher::IsCompleted() const { return pimpl_->IsCompleted(); } + +void GrammarMatcher::Reset() { pimpl_->Reset(); } + +GrammarMatcher GrammarMatcher::Fork() const { + return GrammarMatcher(std::make_shared(*pimpl_)); +} + +int GrammarMatcher::GetMaxRollbackTokens() const { return pimpl_->GetMaxRollbackTokens(); } + +std::optional GrammarMatcher::GetTemperature() const { return pimpl_->GetTemperature(); } + +const std::vector& GrammarMatcher::GetStopTokenIds() const { + return pimpl_->GetStopTokenIds(); +} + +std::string GrammarMatcher::_DebugPrintInternalState() const { + return pimpl_->_DebugPrintInternalState(); +} + +void BatchGrammarMatcher::BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices, + bool debug_print +) { + pimpl_->BatchFillNextTokenBitmask(matchers, next_token_bitmask, indices, debug_print); +} + +void BatchGrammarMatcher::BatchFillTemperature( + const std::vector& matchers, + DLTensor* temperatures, + const std::optional>& indices +) { + XGRAMMAR_CHECK( + temperatures->ndim == 1 && temperatures->dtype.code == kDLFloat && + temperatures->dtype.bits == 32 + ) << "The temperatures tensor must be a 1D float32 tensor"; + XGRAMMAR_CHECK( + temperatures->device.device_type == kDLCPU || + temperatures->device.device_type == kDLCUDAHost || + temperatures->device.device_type == kDLROCMHost + ) << "The temperatures tensor must be on CPU"; + if (indices.has_value()) { + XGRAMMAR_CHECK(indices->size() == matchers.size()) + << "The indices size must match the number of matchers"; + } + auto* temperatures_data = reinterpret_cast(temperatures->data); + for (int64_t i = 0; i < static_cast(matchers.size()); ++i) { + int64_t position = indices.has_value() ? (*indices)[i] : i; + XGRAMMAR_CHECK(position >= 0 && position < temperatures->shape[0]) + << "The index " << position << " is out of bounds for the temperatures tensor"; + temperatures_data[position] = matchers[i].GetTemperature().value_or(-1.0f); + } +} + +std::vector BatchGrammarMatcher::BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print +) { + return Impl::BatchAcceptString(matchers, input_strs, debug_print); +} + +std::vector BatchGrammarMatcher::BatchAcceptToken( + std::vector* matchers, const std::vector& token_ids, bool debug_print +) { + return Impl::BatchAcceptToken(matchers, token_ids, debug_print); +} + +void BatchGrammarMatcher::BatchRollback( + std::vector* matchers, const std::vector& num_tokens +) { + Impl::BatchRollback(matchers, num_tokens); +} + +BatchGrammarMatcher::BatchGrammarMatcher(std::variant max_threads) + : pimpl_(std::make_shared(max_threads)) {} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_parser.cc b/zig/vendor/xgrammar/cpp/grammar_parser.cc new file mode 100644 index 0000000000..1d00ec4092 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_parser.cc @@ -0,0 +1,1593 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_parser.cc + */ + +#include "grammar_parser.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_impl.h" +#include "support/encoding.h" +#include "support/logging.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +class EBNFLexer::Impl { + public: + using Token = EBNFLexer::Token; + using TokenType = EBNFLexer::TokenType; + + std::vector Tokenize(const std::string& input); + + private: + std::string input_; + const char* cur_ = nullptr; + int cur_line_ = 1; + int cur_column_ = 1; + + constexpr static int64_t kMaxIntegerInGrammar = 1e15; + + // Helper functions + + /*! + * \brief Consume a character sequence and return the next token. Return a token if it's a + * single token, or a vector of tokens if it's a sequence of tokens. + * + * \return std::variant> + */ + std::variant> NextToken(); + Token ParseIdentifierOrBooleanToken(); + Token ParseStringToken(); + std::vector ParseCharClassToken(); + Token ParseIntegerToken(); + [[noreturn]] void ReportLexerError(const std::string& msg, int line = -1, int column = -1); + char Peek(int delta = 0) const; + void Consume(int cnt = 1); + void ConsumeSpace(); + std::string ParseIdentifierToken(); + void ConvertIdentifierToRuleName(std::vector* tokens); + static bool IsNameChar(char c, bool is_first = false); +}; + +// Look at the next character +inline char EBNFLexer::Impl::Peek(int delta) const { return *(cur_ + delta); } + +// Consume characters and update position information +inline void EBNFLexer::Impl::Consume(int cnt) { + for (int i = 0; i < cnt; ++i) { + // Newline\n \r \r\n + if (*cur_ == '\n' || (*cur_ == '\r' && *(cur_ + 1) != '\n')) { + ++cur_line_; + cur_column_ = 1; + } else { + ++cur_column_; + } + ++cur_; + } +} + +// Skip whitespace and comments +void EBNFLexer::Impl::ConsumeSpace() { + while (Peek() && + (Peek() == ' ' || Peek() == '\t' || Peek() == '#' || Peek() == '\n' || Peek() == '\r')) { + Consume(); + if (Peek(-1) == '#') { + while (Peek() && Peek() != '\n' && Peek() != '\r') { + Consume(); + } + if (!Peek()) { + return; + } + Consume(); + if (Peek(-1) == '\r' && Peek() == '\n') { + Consume(); + } + } + } +} + +// Report parsing error +void EBNFLexer::Impl::ReportLexerError(const std::string& msg, int line, int column) { + int line_to_print = line == -1 ? cur_line_ : line; + int column_to_print = column == -1 ? cur_column_ : column; + XGRAMMAR_LOG(FATAL) << "EBNF lexer error at line " + std::to_string(line_to_print) + ", column " + + std::to_string(column_to_print) + ": " + msg; + XGRAMMAR_UNREACHABLE(); +} + +// Check if a character can be part of an identifier +bool EBNFLexer::Impl::IsNameChar(char c, bool is_first) { + return c == '_' || c == '-' || c == '.' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (!is_first && c >= '0' && c <= '9'); +} + +// Parse identifier +std::string EBNFLexer::Impl::ParseIdentifierToken() { + const char* start = cur_; + bool first_char = true; + while (*cur_ && IsNameChar(*cur_, first_char)) { + Consume(); + first_char = false; + } + if (start == cur_) { + ReportLexerError("Expect identifier"); + } + return std::string(start, cur_ - start); +} + +// Parse identifier or boolean value +EBNFLexer::Token EBNFLexer::Impl::ParseIdentifierOrBooleanToken() { + int start_line = cur_line_; + int start_column = cur_column_; + + std::string identifier = ParseIdentifierToken(); + + // Check if it's a boolean value + if (identifier == "true" || identifier == "false") { + return { + TokenType::BooleanLiteral, + identifier, + identifier == "true" ? true : false, + start_line, + start_column + }; + } + + // A rule definition may carry an attribute block before ::=, e.g. + // name[max_tokens=10] ::= ..., name[max_chars=10] ::= ..., name[capture] ::= ..., + // name[capture="x"] ::= ..., + // name[capture_hidden_suffix_bytes=3] ::= ..., name[capture_hidden_stop_bytes=3] ::= ..., + // name[capture_hidden_body_rule_id=1, capture_hidden_marker_rule_id=2] ::= ..., + // name[stop_capture="marker"] ::= ..., name[lazy] ::= ..., name[temperature=0.7] ::= ..., or a + // comma-separated combination. + // The bracket group is treated as an attribute block only when it is followed by "::="; + // otherwise it is left to be lexed as a character class. + if (*cur_ == '[') { + int delta = 1; + auto skip_space = [&]() { + while (Peek(delta) == ' ' || Peek(delta) == '\t') { + ++delta; + } + }; + // Match the keyword at the current position and advance delta past it on success. + auto match_keyword = [&](const char* keyword) { + int len = 0; + while (keyword[len] != '\0') { + if (Peek(delta + len) != keyword[len]) { + return false; + } + ++len; + } + delta += len; + return true; + }; + bool matched = true; + bool has_max_tokens = false; + bool has_max_chars = false; + bool has_capture = false; + bool has_capture_hidden_suffix_bytes = false; + bool has_capture_hidden_stop_bytes = false; + bool has_capture_hidden_body_rule_id = false; + bool has_capture_hidden_marker_rule_id = false; + bool has_stop_capture = false; + bool has_lazy = false; + bool has_temperature = false; + double temperature_value = 0; + int64_t max_tokens_value = -1; + int64_t max_chars_value = -1; + int64_t capture_hidden_suffix_bytes_value = 0; + int64_t capture_hidden_stop_bytes_value = 0; + int64_t capture_hidden_body_rule_id_value = -1; + int64_t capture_hidden_marker_rule_id_value = -1; + std::string capture_value; + std::string stop_capture_value; + auto parse_integer_value = [&](int64_t* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + *value = 0; + int digits = 0; + while (Peek(delta) >= '0' && Peek(delta) <= '9' && digits < 10) { + *value = *value * 10 + (Peek(delta) - '0'); + ++delta; + ++digits; + } + return digits > 0 && !(Peek(delta) >= '0' && Peek(delta) <= '9'); + }; + auto parse_string_value = [&](std::string* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + if (Peek(delta) != '"') { + return false; + } + ++delta; + while (Peek(delta) != '"') { + char c = Peek(delta); + if (c == '\0' || c == '\n' || c == '\r' || c == '\\') { + return false; + } + value->push_back(c); + ++delta; + } + ++delta; + return true; + }; + auto parse_float_value = [&](double* value) { + skip_space(); + if (Peek(delta) != '=') { + return false; + } + ++delta; + skip_space(); + std::string text; + while (true) { + char c = Peek(delta); + if ((c >= '0' && c <= '9') || c == '.' || c == 'e' || c == 'E' || c == '+' || c == '-') { + text.push_back(c); + ++delta; + } else { + break; + } + } + if (text.empty()) { + return false; + } + try { + *value = std::stod(text); + } catch (const std::out_of_range&) { + *value = std::numeric_limits::infinity(); + } catch (const std::exception&) { + return false; + } + return true; + }; + // Parse a comma-separated attribute list. Each attribute may appear at most once. + while (matched) { + skip_space(); + if (!has_max_tokens && match_keyword("max_tokens")) { + has_max_tokens = true; + matched = parse_integer_value(&max_tokens_value); + } else if (!has_max_chars && match_keyword("max_chars")) { + has_max_chars = true; + matched = parse_integer_value(&max_chars_value); + } else if (!has_capture_hidden_suffix_bytes && match_keyword("capture_hidden_suffix_bytes")) { + has_capture_hidden_suffix_bytes = true; + matched = parse_integer_value(&capture_hidden_suffix_bytes_value); + } else if (!has_capture_hidden_stop_bytes && match_keyword("capture_hidden_stop_bytes")) { + has_capture_hidden_stop_bytes = true; + matched = parse_integer_value(&capture_hidden_stop_bytes_value); + } else if (!has_capture_hidden_body_rule_id && match_keyword("capture_hidden_body_rule_id")) { + has_capture_hidden_body_rule_id = true; + matched = parse_integer_value(&capture_hidden_body_rule_id_value); + } else if (!has_capture_hidden_marker_rule_id && + match_keyword("capture_hidden_marker_rule_id")) { + has_capture_hidden_marker_rule_id = true; + matched = parse_integer_value(&capture_hidden_marker_rule_id_value); + } else if (!has_stop_capture && match_keyword("stop_capture")) { + has_stop_capture = true; + matched = parse_string_value(&stop_capture_value); + } else if (!has_capture && match_keyword("capture")) { + has_capture = true; + skip_space(); + if (Peek(delta) == '=') { + matched = parse_string_value(&capture_value); + } else { + capture_value = identifier; + } + } else if (!has_lazy && match_keyword("lazy")) { + has_lazy = true; + } else if (!has_temperature && match_keyword("temperature")) { + has_temperature = true; + matched = parse_float_value(&temperature_value); + } else { + matched = false; + } + if (!matched) { + break; + } + skip_space(); + if (Peek(delta) == ',') { + ++delta; + continue; + } + break; + } + if (matched) { + skip_space(); + if (Peek(delta) == ']') { + ++delta; + } else { + matched = false; + } + } + if (matched) { + int after_bracket = delta; + while (Peek(after_bracket) == ' ' || Peek(after_bracket) == '\t') { + ++after_bracket; + } + if (!(Peek(after_bracket) == ':' && Peek(after_bracket + 1) == ':' && + Peek(after_bracket + 2) == '=')) { + matched = false; + } + } + if (matched) { + if (has_capture && capture_value.empty()) { + ReportLexerError("The capture name must not be empty", start_line, start_column); + } + if (has_stop_capture && stop_capture_value.empty()) { + ReportLexerError("The stop capture name must not be empty", start_line, start_column); + } + if (has_capture_hidden_body_rule_id != has_capture_hidden_marker_rule_id) { + ReportLexerError( + "The capture-hidden body and marker rule ids must be specified together", + start_line, + start_column + ); + } + if ((has_capture_hidden_suffix_bytes && capture_hidden_suffix_bytes_value <= 0) || + (has_capture_hidden_stop_bytes && capture_hidden_stop_bytes_value <= 0)) { + ReportLexerError( + "The number of capture-hidden bytes must be positive", start_line, start_column + ); + } + constexpr int64_t kMaxInt32 = std::numeric_limits::max(); + if (has_max_chars && max_chars_value < 0) { + ReportLexerError( + "The max_chars rule attribute must be non-negative", start_line, start_column + ); + } + if ((has_max_chars && max_chars_value > kMaxInt32) || + (has_capture_hidden_suffix_bytes && capture_hidden_suffix_bytes_value > kMaxInt32) || + (has_capture_hidden_stop_bytes && capture_hidden_stop_bytes_value > kMaxInt32) || + (has_capture_hidden_body_rule_id && capture_hidden_body_rule_id_value > kMaxInt32) || + (has_capture_hidden_marker_rule_id && capture_hidden_marker_rule_id_value > kMaxInt32)) { + ReportLexerError("The rule attribute value is too large", start_line, start_column); + } + if (has_temperature && !(std::isfinite(temperature_value) && temperature_value >= 0 && + temperature_value <= std::numeric_limits::max())) { + ReportLexerError( + "The temperature must be a finite non-negative number", start_line, start_column + ); + } + Consume(delta); + Token token{TokenType::Identifier, identifier, identifier, start_line, start_column}; + token.max_tokens = static_cast(max_tokens_value); + token.max_chars = static_cast(max_chars_value); + token.capture_name = capture_value; + token.capture_hidden_suffix_bytes = static_cast(capture_hidden_suffix_bytes_value); + token.capture_hidden_stop_bytes = static_cast(capture_hidden_stop_bytes_value); + token.capture_hidden_body_rule_id = static_cast(capture_hidden_body_rule_id_value); + token.capture_hidden_marker_rule_id = + static_cast(capture_hidden_marker_rule_id_value); + token.stop_capture_name = stop_capture_value; + token.is_lazy = has_lazy; + if (has_temperature) { + token.temperature = static_cast(temperature_value); + } + return token; + } + } + + // Otherwise it's an identifier + return {TokenType::Identifier, identifier, identifier, start_line, start_column}; +} + +// Parse string literal +EBNFLexer::Token EBNFLexer::Impl::ParseStringToken() { + int start_line = cur_line_; + int start_column = cur_column_; + const char* start_pos = cur_; + + Consume(); // Skip opening quote + + std::vector codepoints; + while (Peek() && Peek() != '"' && Peek() != '\n' && Peek() != '\r') { + auto [codepoint, len] = ParseNextUTF8OrEscaped(cur_); + if (codepoint == CharHandlingError::kInvalidUTF8) { + ReportLexerError("Invalid UTF8 sequence"); + } + if (codepoint == CharHandlingError::kInvalidEscape) { + ReportLexerError("Invalid escape sequence"); + } + Consume(len); + codepoints.push_back(codepoint); + } + + if (Peek() != '"') { + ReportLexerError("Expect \" in string literal"); + } + Consume(); // Skip closing quote + + // Extract original lexeme + std::string lexeme(start_pos, cur_ - start_pos); + + // Convert codepoints to UTF-8 string value + std::string value; + for (auto codepoint : codepoints) { + value += CharToUTF8(codepoint); + } + + return {TokenType::StringLiteral, lexeme, value, start_line, start_column}; +} + +// Parse character class. +std::vector EBNFLexer::Impl::ParseCharClassToken() { + std::vector tokens; + + tokens.push_back({TokenType::LBracket, "[", "", cur_line_, cur_column_}); + Consume(); // Skip '[' + + if (Peek() == '^') { + tokens.push_back({TokenType::Caret, "^", "", cur_line_, cur_column_}); + Consume(); + } + + static const std::unordered_map kRegexEscapeChars = { + // clang-format off + {'^', '^'}, {'$', '$'}, {'\\', '\\'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'}, {'-', '-'} // clang-format on + }; + + static const std::unordered_set kRegexSpecialEscapes = {'d', 'D', 's', 'S', 'w', 'W'}; + + while (Peek() && Peek() != ']') { + if (Peek() == '\r' || Peek() == '\n') { + ReportLexerError("Character class should not contain newline"); + } else if (Peek() == '-') { + // Handle dash; this dash could be a range expression or a normal dash. + // It will further be handled in EBNFParser::ParseCharClass. + tokens.push_back({TokenType::Dash, "-", "", cur_line_, cur_column_}); + Consume(); + } else if (Peek() == '\\' && kRegexSpecialEscapes.count(Peek(1))) { + // Handle escaped characters with special function + tokens.push_back( + {TokenType::EscapeInCharClass, + std::string(cur_, cur_ + 2), + std::string(cur_ + 1, cur_ + 2), + cur_line_, + cur_column_} + ); + Consume(2); + } else { + // Handle normal characters + auto [codepoint, len] = ParseNextUTF8OrEscaped(cur_, kRegexEscapeChars); + if (codepoint == CharHandlingError::kInvalidUTF8) { + ReportLexerError("Invalid UTF8 sequence"); + } + + if (codepoint == CharHandlingError::kInvalidEscape) { + ReportLexerError("Invalid escape sequence" + std::string(cur_, cur_ + 2)); + } + + tokens.push_back( + {TokenType::CharInCharClass, + std::string(cur_, cur_ + len), + codepoint, + cur_line_, + cur_column_} + ); + Consume(len); + } + } + + if (!Peek()) { + ReportLexerError("Unterminated character class"); + } + + tokens.push_back({TokenType::RBracket, "]", "", cur_line_, cur_column_}); + Consume(); // Skip ']' + + return tokens; +} + +// Parse integer +EBNFLexer::Token EBNFLexer::Impl::ParseIntegerToken() { + int start_line = cur_line_; + int start_column = cur_column_; + const char* start_pos = cur_; + bool is_negative = false; + + if (Peek() == '-') { + is_negative = true; + Consume(); + } else if (Peek() == '+') { + Consume(); + } + + int64_t num = 0; + while (Peek() && isdigit(Peek())) { + num = num * 10 + (Peek() - '0'); + Consume(); + if (num > kMaxIntegerInGrammar) { + ReportLexerError( + "Integer is too large: parsed " + std::to_string(num) + ", max allowed is " + + std::to_string(kMaxIntegerInGrammar) + ); + } + } + + std::string lexeme(start_pos, cur_ - start_pos); + return {TokenType::IntegerLiteral, lexeme, is_negative ? -num : num, start_line, start_column}; +} + +// Get the next token +std::variant> EBNFLexer::Impl::NextToken() { + ConsumeSpace(); // Skip whitespace and comments + + auto start_line = cur_line_; + auto start_column = cur_column_; + + if (!Peek()) { + return EBNFLexer::Token{TokenType::EndOfFile, "", "", start_line, start_column}; + } + + // Determine token type based on current character + switch (Peek()) { + case '(': + if (Peek(1) == '=') { + Consume(2); + return EBNFLexer::Token{TokenType::LookaheadLParen, "(=", "", start_line, start_column}; + } else { + Consume(); + return EBNFLexer::Token{TokenType::LParen, "(", "", start_line, start_column}; + } + case ')': + Consume(); + return EBNFLexer::Token{TokenType::RParen, ")", "", start_line, start_column}; + case '{': + Consume(); + return EBNFLexer::Token{TokenType::LBrace, "{", "", start_line, start_column}; + case '}': + Consume(); + return EBNFLexer::Token{TokenType::RBrace, "}", "", start_line, start_column}; + case '|': + Consume(); + return EBNFLexer::Token{TokenType::Pipe, "|", "", start_line, start_column}; + case ',': + Consume(); + return EBNFLexer::Token{TokenType::Comma, ",", "", start_line, start_column}; + case '*': + Consume(); + return EBNFLexer::Token{TokenType::Star, "*", "", start_line, start_column}; + case '+': + Consume(); + return EBNFLexer::Token{TokenType::Plus, "+", "", start_line, start_column}; + case '?': + Consume(); + return EBNFLexer::Token{TokenType::Question, "?", "", start_line, start_column}; + case '=': + Consume(); + return EBNFLexer::Token{TokenType::Equal, "=", "", start_line, start_column}; + case ':': + if (Peek(1) == ':' && Peek(2) == '=') { + Consume(3); + return EBNFLexer::Token{TokenType::Assign, "::=", "", start_line, start_column}; + } + ReportLexerError("Unexpected character: ':'"); + break; + case '"': + return ParseStringToken(); + case '[': + return ParseCharClassToken(); + default: + if (IsNameChar(*cur_, true)) { + return ParseIdentifierOrBooleanToken(); + } else if (isdigit(*cur_) || *cur_ == '-' || *cur_ == '+') { + return ParseIntegerToken(); + } + + // Unrecognized character, report error + ReportLexerError("Unexpected character: " + std::string(1, *cur_)); + } + + // Should not reach here + XGRAMMAR_UNREACHABLE(); +} + +void EBNFLexer::Impl::ConvertIdentifierToRuleName(std::vector* tokens) { + for (int i = 0; i < static_cast(tokens->size()); ++i) { + if (tokens->at(i).type == TokenType::Assign) { + if (i == 0) { + ReportLexerError( + "Assign should not be the first token", tokens->at(i).line, tokens->at(i).column + ); + } + if (tokens->at(i - 1).type != TokenType::Identifier) { + ReportLexerError( + "Assign should be preceded by an identifier", + tokens->at(i - 1).line, + tokens->at(i - 1).column + ); + } + if (i >= 2 && tokens->at(i - 2).line == tokens->at(i - 1).line) { + ReportLexerError( + "The rule name should be at the beginning of the line", + tokens->at(i - 1).line, + tokens->at(i - 1).column + ); + } + tokens->at(i - 1).type = TokenType::RuleName; + } + } +} + +// Tokenize the entire input and return a vector of tokens +std::vector EBNFLexer::Impl::Tokenize(const std::string& input) { + // Reset position to the beginning + input_ = input; + cur_ = input_.c_str(); + cur_line_ = 1; + cur_column_ = 1; + + // Collect all tokens + std::vector tokens; + + while (true) { + auto token = NextToken(); + + if (auto* token_value = std::get_if(&token)) { + tokens.push_back(*token_value); + // Stop when we reach the end of file + if (token_value->type == TokenType::EndOfFile) { + break; + } + } else { + auto vec = std::get_if>(&token); + XGRAMMAR_DCHECK(vec != nullptr); + tokens.insert(tokens.end(), vec->begin(), vec->end()); + } + } + + ConvertIdentifierToRuleName(&tokens); + + return tokens; +} + +EBNFLexer::EBNFLexer() : pimpl_(std::make_shared()) {} + +std::vector EBNFLexer::Tokenize(const std::string& input) { + return pimpl_->Tokenize(input); +} + +class EBNFParser { + public: + /*! \brief The logic of parsing the grammar string. */ + Grammar Parse( + const std::vector& tokens, + const std::string& root_rule_name, + const int& max_nest_layer = 1000 + ); + + private: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using Token = EBNFLexer::Token; + using TokenType = EBNFLexer::TokenType; + + struct ParsedRule { + Rule rule; + SuffixStopInfo suffix_stop_info; + }; + + // Parsing different parts of the grammar + std::string ParseIdentifier(); + int32_t ParseCharClass(); + int32_t ParseString(); + int32_t ParseRuleRef(); + int32_t ParseElement(); + int64_t ParseInteger(); + std::pair ParseRepetitionRange(); + int32_t ParseElementWithQuantifier(); + int32_t ParseLookaheadAssertion(); + int32_t ParseSequence(); + int32_t ParseChoices(); + ParsedRule ParseRule(); + + // Parser for macro + class MacroIR { + public: + struct StringNode; + struct IntegerNode; + struct BooleanNode; + struct IdentifierNode; + struct TupleNode; + + using Node = std::variant; + using NodePtr = std::unique_ptr; + + struct StringNode { + std::string value; + }; + struct IntegerNode { + int64_t value; + }; + struct BooleanNode { + bool value; + }; + struct IdentifierNode { + std::string name; + }; + struct TupleNode { + std::vector elements; + }; + + struct Arguments { + std::vector arguments; + std::unordered_map named_arguments; + }; + }; + MacroIR::Arguments ParseMacroArguments(); + MacroIR::NodePtr ParseMacroValue(); + + int32_t ParseTagDispatch(); + int32_t ParseTokenSet(); + int32_t ParseExcludeToken(); + int32_t ParseTokenTagDispatch(); + int32_t ParseRegexMacro(); + int32_t ParseSubstringMacro(); + + // Helper functions + + // Helper for ParseElementWithQuantifier + int32_t HandleStarQuantifier(int32_t grammar_expr_id); + int32_t HandlePlusQuantifier(int32_t grammar_expr_id); + int32_t HandleQuestionQuantifier(int32_t grammar_expr_id); + + // When parsing, we first find the names of all rules, and build the mapping from name to rule id. + void InitRuleNames(); + + // Consume a token and advance to the next + void Consume(int cnt = 1); + + // Peek at the current token with optional offset + const Token& Peek(int delta = 0) const; + + // Consume token if it matches expected type, otherwise report error + void PeekAndConsume(TokenType type, const std::string& message); + + // Report a parsing error with the given message + [[noreturn]] void ReportParseError(const std::string& msg, int delta_element = 0); + + // The grammar builder + GrammarBuilder builder_; + + // The current token pointer + const Token* current_token_ = nullptr; + + // Tokens from lexer + std::vector tokens_; + + // The current rule name. Help to generate a name for a new rule. + std::string cur_rule_name_; + + // The name of the root rule + std::string root_rule_name_; + + int nest_layer_guard_ = 0; + + int max_nest_layer_ = 1000; // Max nest layer of the grammar + + static const std::unordered_map> kMacroFunctions; +}; + +const std::unordered_map> + EBNFParser::kMacroFunctions = { + {"TagDispatch", [](EBNFParser* parser) { return parser->ParseTagDispatch(); }}, + {"Token", [](EBNFParser* parser) { return parser->ParseTokenSet(); }}, + {"ExcludeToken", [](EBNFParser* parser) { return parser->ParseExcludeToken(); }}, + {"TokenTagDispatch", [](EBNFParser* parser) { return parser->ParseTokenTagDispatch(); }}, + {"Regex", [](EBNFParser* parser) { return parser->ParseRegexMacro(); }}, + {"Substring", [](EBNFParser* parser) { return parser->ParseSubstringMacro(); }}, +}; + +const EBNFParser::Token& EBNFParser::Peek(int delta) const { return *(current_token_ + delta); } + +void EBNFParser::Consume(int cnt) { current_token_ += cnt; } + +void EBNFParser::PeekAndConsume(TokenType type, const std::string& message) { + if (Peek().type != type) { + ReportParseError(message); + } + Consume(); +} + +void EBNFParser::ReportParseError(const std::string& msg, int delta_element) { + XGRAMMAR_DCHECK(current_token_ + delta_element < tokens_.data() + tokens_.size()); + int line_to_print = Peek(delta_element).line; + int column_to_print = Peek(delta_element).column; + XGRAMMAR_LOG(FATAL) << "EBNF parser error at line " + std::to_string(line_to_print) + + ", column " + std::to_string(column_to_print) + ": " + msg; + XGRAMMAR_UNREACHABLE(); +} + +std::string EBNFParser::ParseIdentifier() { + if (Peek().type != TokenType::Identifier) { + ReportParseError("Expect identifier"); + } + std::string identifier = std::any_cast(Peek().value); + Consume(); + return identifier; +} + +int32_t EBNFParser::ParseCharClass() { + PeekAndConsume(TokenType::LBracket, "Expect [ in character class"); + + std::vector elements; + bool is_negated = false; + + if (Peek().type == TokenType::Caret) { + is_negated = true; + Consume(); + } + + while (Peek().type != TokenType::RBracket && Peek().type != TokenType::EndOfFile) { + if (Peek().type == TokenType::EscapeInCharClass) { + ReportParseError("Character class escape is not supported yet in EBNF"); + } + + TCodepoint codepoint; + if (Peek().type == TokenType::CharInCharClass) { + codepoint = std::any_cast(Peek().value); + } else if (Peek().type == TokenType::Dash) { + codepoint = static_cast(static_cast('-')); + } else { + ReportParseError("Unexpected character in character class: " + Peek().lexeme); + } + Consume(); + + if (Peek().type == TokenType::Dash && + (Peek(1).type == TokenType::CharInCharClass || Peek(1).type == TokenType::Dash)) { + // Range expression + TCodepoint codepoint2; + if (Peek(1).type == TokenType::CharInCharClass) { + codepoint2 = std::any_cast(Peek(1).value); + } else { + XGRAMMAR_DCHECK(Peek(1).type == TokenType::Dash); + codepoint2 = static_cast(static_cast('-')); + } + + if (codepoint > codepoint2) { + ReportParseError("Invalid character class: lower bound is larger than upper bound", -1); + } + elements.push_back({codepoint, codepoint2}); + Consume(2); + } else { + // Single character + elements.push_back({codepoint, codepoint}); + } + } + + PeekAndConsume(TokenType::RBracket, "Expect ] in character class"); + + return builder_.AddCharacterClass(elements, is_negated); +} + +int32_t EBNFParser::ParseString() { + if (Peek().type != TokenType::StringLiteral) { + ReportParseError("Expect string literal"); + } + + std::string str_value = std::any_cast(Peek().value); + Consume(); + + if (str_value.empty()) { + return builder_.AddEmptyStr(); + } + + return builder_.AddByteString(str_value); +} + +int32_t EBNFParser::ParseRuleRef() { + std::string name = ParseIdentifier(); + auto rule_id = builder_.GetRuleId(name); + if (rule_id == -1) { + ReportParseError("Rule \"" + name + "\" is not defined", -1); + } + return builder_.AddRuleRef(rule_id); +} + +int32_t EBNFParser::ParseElement() { + if (Peek().type == TokenType::LParen) { + nest_layer_guard_++; + if (nest_layer_guard_ > max_nest_layer_) { + ReportParseError("Nest layer exceeded the maximum limit", -1); + } + Consume(); + if (Peek().type == TokenType::RParen) { + // Special case: ( ) + Consume(); + nest_layer_guard_--; + return builder_.AddEmptyStr(); + } + auto grammar_expr_id = ParseChoices(); + PeekAndConsume(TokenType::RParen, "Expect )"); + nest_layer_guard_--; + return grammar_expr_id; + } else if (Peek().type == TokenType::LBracket) { + return ParseCharClass(); + } else if (Peek().type == TokenType::StringLiteral) { + return ParseString(); + } else if (Peek().type == TokenType::Identifier) { + auto id = std::any_cast(Peek().value); + if (kMacroFunctions.count(id)) { + return kMacroFunctions.at(id)(this); + } else { + return ParseRuleRef(); + } + } else { + ReportParseError("Expect element, but got " + Peek().lexeme); + } +} + +int64_t EBNFParser::ParseInteger() { + if (Peek().type != TokenType::IntegerLiteral) { + ReportParseError("Expect integer, but got " + Peek().lexeme); + } + int64_t num = std::any_cast(Peek().value); + Consume(); + return num; +} + +std::pair EBNFParser::ParseRepetitionRange() { + PeekAndConsume(TokenType::LBrace, "Expect {"); + + int64_t lower = ParseInteger(); + + if (lower < 0) { + ReportParseError("Lower bound cannot be negative", -1); + } + + if (Peek().type == TokenType::Comma) { + Consume(); + if (Peek().type == TokenType::RBrace) { + Consume(); + return {lower, -1}; + } + // The grammar printer emits {n, -1} for unbounded upper bounds, and + // '-' is a valid identifier-start char (IsNameChar), so the lexer + // produces Identifier("-1") rather than IntegerLiteral. Accept it + // as equivalent to {n,}. + if (Peek().type == TokenType::Identifier && Peek().lexeme == "-1") { + Consume(); + PeekAndConsume(TokenType::RBrace, "Expect }"); + return {lower, -1}; + } + int64_t upper = ParseInteger(); + if (upper < lower) { + ReportParseError( + "Lower bound is larger than upper bound: " + std::to_string(lower) + " > " + + std::to_string(upper), + -1 + ); + } + PeekAndConsume(TokenType::RBrace, "Expect }"); + return {lower, upper}; + } else if (Peek().type == TokenType::RBrace) { + Consume(); + return {lower, lower}; + } + + ReportParseError("Expect ',' or '}' in repetition range"); +} + +int32_t EBNFParser::HandleStarQuantifier(int32_t grammar_expr_id) { + Grammar::Impl::GrammarExpr grammar_expr = builder_.GetGrammarExpr(grammar_expr_id); + if (grammar_expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + // We have special handling for character class star, e.g. [a-z]* + grammar_expr.type = GrammarBuilder::GrammarExprType::kCharacterClassStar; + // Copy grammar expr because the grammar may change during insertion, and grammar_expr is in the + // grammar, so it may become invalid + std::vector grammar_expr_data(grammar_expr.begin(), grammar_expr.end()); + return builder_.AddGrammarExpr( + {grammar_expr.type, grammar_expr_data.data(), grammar_expr.data_len} + ); + } else { + // For other star quantifiers, we transform it into a rule: + // a* --> rule ::= a rule | "" + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_rule_id = builder_.AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_.AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_.AddChoices( + {builder_.AddEmptyStr(), builder_.AddSequence({grammar_expr_id, ref_to_new_rule})} + ); + builder_.UpdateRuleBody(new_rule_id, new_grammar_expr_id); + + // Return the reference to the new rule + return builder_.AddRuleRef(new_rule_id); + } +} + +int32_t EBNFParser::HandlePlusQuantifier(int32_t grammar_expr_id) { + // a+ --> rule ::= a rule | a + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_rule_id = builder_.AddEmptyRule(new_rule_name); + auto ref_to_new_rule = builder_.AddRuleRef(new_rule_id); + auto new_grammar_expr_id = builder_.AddChoices( + {builder_.AddSequence({grammar_expr_id, ref_to_new_rule}), grammar_expr_id} + ); + builder_.UpdateRuleBody(new_rule_id, new_grammar_expr_id); + + // Return the reference to the new rule + return builder_.AddRuleRef(new_rule_id); +} + +int32_t EBNFParser::HandleQuestionQuantifier(int32_t grammar_expr_id) { + // a? --> rule ::= a | empty + auto new_rule_name = builder_.GetNewRuleName(cur_rule_name_); + auto new_grammar_expr_id = builder_.AddChoices({builder_.AddEmptyStr(), grammar_expr_id}); + auto new_rule_id = builder_.AddRule({new_rule_name, new_grammar_expr_id}); + return builder_.AddRuleRef(new_rule_id); +} + +int32_t EBNFParser::ParseElementWithQuantifier() { + int32_t grammar_expr_id = ParseElement(); + + if (Peek().type == TokenType::Star) { + Consume(); + return HandleStarQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::Plus) { + Consume(); + return HandlePlusQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::Question) { + Consume(); + return HandleQuestionQuantifier(grammar_expr_id); + } else if (Peek().type == TokenType::LBrace) { + auto [lower, upper] = ParseRepetitionRange(); + return builder_.AddRepeatFromExpr( + cur_rule_name_, + grammar_expr_id, + static_cast(lower), + upper == -1 ? -1 : static_cast(upper) + ); + } + + return grammar_expr_id; +} + +int32_t EBNFParser::ParseSequence() { + std::vector elements; + + do { + elements.push_back(ParseElementWithQuantifier()); + } while (Peek().type != TokenType::Pipe && Peek().type != TokenType::RParen && + Peek().type != TokenType::LookaheadLParen && Peek().type != TokenType::RuleName && + Peek().type != TokenType::EndOfFile); + + return builder_.AddSequence(elements); +} + +int32_t EBNFParser::ParseChoices() { + std::vector choices; + + choices.push_back(ParseSequence()); + + while (Peek().type == TokenType::Pipe) { + Consume(); + choices.push_back(ParseSequence()); + } + + return builder_.AddChoices(choices); +} + +// Parse macro arguments and return a MacroIR::Arguments structure +EBNFParser::MacroIR::Arguments EBNFParser::ParseMacroArguments() { + MacroIR::Arguments args; + + PeekAndConsume(TokenType::LParen, "Expect ( after macro function name"); + + // Parse arguments + if (Peek().type != TokenType::RParen) { + while (true) { + // Check if it's a named argument (identifier = value) + if (Peek().type == TokenType::Identifier && Peek(1).type == TokenType::Equal) { + std::string name = std::any_cast(Peek().value); + Consume(); // Consume identifier + Consume(); // Consume = + + // Parse the value + args.named_arguments[name] = ParseMacroValue(); + } else { + // Regular positional argument + args.arguments.push_back(ParseMacroValue()); + } + + // Check for comma or end of arguments + if (Peek().type == TokenType::Comma) { + Consume(); + } else if (Peek().type == TokenType::RParen) { + break; + } else { + ReportParseError("Expect , or ) in macro arguments"); + } + } + } + + PeekAndConsume(TokenType::RParen, "Expect ) after macro arguments"); + return args; +} + +// Parse a single macro value (string, integer, boolean, or tuple) +EBNFParser::MacroIR::NodePtr EBNFParser::ParseMacroValue() { + if (Peek().type == TokenType::StringLiteral) { + // String value + std::string value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::StringNode{value}); + } else if (Peek().type == TokenType::IntegerLiteral) { + // Integer value + int64_t value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::IntegerNode{value}); + } else if (Peek().type == TokenType::BooleanLiteral) { + // Boolean value + bool value = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::BooleanNode{value}); + } else if (Peek().type == TokenType::Identifier) { + // Identifier value + std::string name = std::any_cast(Peek().value); + Consume(); + return std::make_unique(MacroIR::IdentifierNode{name}); + } else if (Peek().type == TokenType::LParen) { + // Tuple value. Nested tuples recurse once per layer, so guard the depth the same way the + // ordinary parenthesis path does to avoid a stack overflow on deeply nested input. + nest_layer_guard_++; + if (nest_layer_guard_ > max_nest_layer_) { + ReportParseError("Nest layer exceeded the maximum limit", -1); + } + Consume(); // Consume ( + + MacroIR::TupleNode tuple; + + // Parse tuple elements (supports trailing comma) + if (Peek().type != TokenType::RParen) { + while (true) { + tuple.elements.push_back(ParseMacroValue()); + + if (Peek().type == TokenType::Comma) { + Consume(); + if (Peek().type == TokenType::RParen) { + break; + } + } else if (Peek().type == TokenType::RParen) { + break; + } else { + ReportParseError("Expect , or ) in tuple"); + } + } + } + + Consume(); // Consume ) + nest_layer_guard_--; + return std::make_unique(std::move(tuple)); + } else { + ReportParseError("Expect string, integer, boolean, or tuple in macro argument"); + } +} + +int32_t EBNFParser::ParseTagDispatch() { + Consume(); // Consume TagDispatch operator + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; // Used to report parse errors + + Grammar::Impl::TagDispatch tag_dispatch; + + static const std::unordered_set kValidNamedArgs = { + "loop_after_dispatch", "excludes" + }; + for (const auto& [name, _] : args.named_arguments) { + if (kValidNamedArgs.count(name) == 0) { + ReportParseError("Unknown named argument for TagDispatch: " + name, delta_element); + } + } + + // Positional parameters: ("tag_string", rule_name) — string triggers only + for (const auto& arg : args.arguments) { + auto tuple_node = std::get_if(arg.get()); + if (tuple_node == nullptr) { + ReportParseError("Each tag dispatch element must be a tuple", delta_element); + } + + if (tuple_node->elements.size() != 2) { + ReportParseError("Each tag dispatch element must be a pair (tag, rule)", delta_element); + } + + auto tag_str_node = std::get_if(tuple_node->elements[0].get()); + if (tag_str_node == nullptr || tag_str_node->value.empty()) { + ReportParseError("Tag must be a non-empty string literal", delta_element); + } + + auto rule_name_node = std::get_if(tuple_node->elements[1].get()); + if (rule_name_node == nullptr) { + ReportParseError("Rule reference must be an identifier", delta_element); + } + + auto rule_id = builder_.GetRuleId(rule_name_node->name); + if (rule_id == -1) { + ReportParseError("Rule \"" + rule_name_node->name + "\" is not defined", delta_element); + } + tag_dispatch.tag_rule_pairs.push_back({tag_str_node->value, rule_id}); + } + + // loop_after_dispatch + tag_dispatch.loop_after_dispatch = true; + if (auto it = args.named_arguments.find("loop_after_dispatch"); + it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("loop_after_dispatch must be a boolean literal", delta_element); + } + tag_dispatch.loop_after_dispatch = bool_node->value; + } + + // excludes — string only + if (auto it = args.named_arguments.find("excludes"); it != args.named_arguments.end()) { + auto tuple_node = std::get_if(it->second.get()); + if (tuple_node == nullptr) { + ReportParseError("excludes must be a tuple", delta_element); + } + for (const auto& element : tuple_node->elements) { + auto str_node = std::get_if(element.get()); + if (str_node == nullptr || str_node->value.empty()) { + ReportParseError("Exclude must be a non-empty string literal", delta_element); + } + tag_dispatch.excludes.push_back(str_node->value); + } + } + + // Well-formedness checks: string excludes vs string triggers + for (const auto& excl_str : tag_dispatch.excludes) { + for (const auto& [trigger_str, _] : tag_dispatch.tag_rule_pairs) { + if (trigger_str.rfind(excl_str, 0) == 0) { + ReportParseError( + "Exclude string must not be a prefix of trigger string: " + excl_str, delta_element + ); + } + } + } + + return builder_.AddTagDispatch(tag_dispatch); +} + +int32_t EBNFParser::ParseRegexMacro() { + Consume(); // Consume Regex operator + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; // Used to report parse errors + + if (args.arguments.size() != 1) { + ReportParseError("Regex expects exactly one string argument", delta_element); + } + auto pattern_node = std::get_if(args.arguments[0].get()); + if (pattern_node == nullptr) { + ReportParseError("Regex pattern must be a string literal", delta_element); + } + + bool json_string = false; + for (const auto& [name, _] : args.named_arguments) { + if (name != "json_string" && name != "flags") { + ReportParseError("Regex does not support the named argument " + name, delta_element); + } + } + if (auto it = args.named_arguments.find("json_string"); it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("json_string must be a boolean", delta_element); + } + json_string = bool_node->value; + } + std::string pattern = pattern_node->value; + if (auto it = args.named_arguments.find("flags"); it != args.named_arguments.end()) { + auto flags_node = std::get_if(it->second.get()); + if (flags_node == nullptr) { + ReportParseError("flags must be a string", delta_element); + } + bool case_insensitive = false; + bool dot_all = false; + for (char flag : flags_node->value) { + if (flag == 'i') { + case_insensitive = true; + } else if (flag == 's') { + dot_all = true; + } else if (flag == 'u') { + // XGrammar regular expressions use Unicode codepoint semantics by default. + } else { + ReportParseError( + "regular-expression flag '" + std::string(1, flag) + "' is not supported", delta_element + ); + } + } + // The flags argument opts into the standard dot semantics: '.' does not match '\n' unless + // the 's' flag is given. Without the argument the pattern is stored verbatim, where the + // engine's '.' matches every codepoint. + pattern = RewriteRegexDots(pattern, dot_all); + if (case_insensitive && pattern.compare(0, 4, "(?i)") != 0) { + pattern = "(?i)" + pattern; + } + } + return builder_.AddRegex(pattern, json_string); +} + +int32_t EBNFParser::ParseSubstringMacro() { + Consume(); // Consume Substring identifier + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("Substring() does not accept named arguments", delta_element); + } + + std::vector chunks; + chunks.reserve(args.arguments.size()); + for (const auto& arg : args.arguments) { + auto string_node = std::get_if(arg.get()); + if (string_node == nullptr) { + ReportParseError("Substring() arguments must be strings", delta_element); + } + chunks.push_back(string_node->value); + } + + return builder_.AddSubstring(chunks); +} + +int32_t EBNFParser::ParseTokenSet() { + Consume(); // Consume Token identifier + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("Token() does not accept named arguments", delta_element); + } + + if (args.arguments.empty()) { + ReportParseError("Token() requires at least one integer argument", delta_element); + } + + std::vector token_ids; + for (const auto& arg : args.arguments) { + auto int_node = std::get_if(arg.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("Token() arguments must be non-negative integers", delta_element); + } + token_ids.push_back(static_cast(int_node->value)); + } + + std::sort(token_ids.begin(), token_ids.end()); + token_ids.erase(std::unique(token_ids.begin(), token_ids.end()), token_ids.end()); + + return builder_.AddTokenSet(token_ids); +} + +int32_t EBNFParser::ParseExcludeToken() { + Consume(); + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + if (!args.named_arguments.empty()) { + ReportParseError("ExcludeToken() does not accept named arguments", delta_element); + } + if (args.arguments.empty()) { + ReportParseError("ExcludeToken() requires at least one integer argument", delta_element); + } + + std::vector token_ids; + for (const auto& arg : args.arguments) { + auto int_node = std::get_if(arg.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("ExcludeToken() arguments must be non-negative integers", delta_element); + } + token_ids.push_back(static_cast(int_node->value)); + } + std::sort(token_ids.begin(), token_ids.end()); + token_ids.erase(std::unique(token_ids.begin(), token_ids.end()), token_ids.end()); + + return builder_.AddExcludeTokenSet(token_ids); +} + +int32_t EBNFParser::ParseTokenTagDispatch() { + Consume(); + auto start = current_token_; + auto args = ParseMacroArguments(); + auto delta_element = start - current_token_; + + Grammar::Impl::TokenTagDispatch ttd; + + static const std::unordered_set kValidNamedArgs = { + "loop_after_dispatch", "excludes" + }; + for (const auto& [name, _] : args.named_arguments) { + if (kValidNamedArgs.count(name) == 0) { + ReportParseError("Unknown named argument for TokenTagDispatch: " + name, delta_element); + } + } + + for (const auto& arg : args.arguments) { + auto tuple_node = std::get_if(arg.get()); + if (tuple_node == nullptr || tuple_node->elements.size() != 2) { + ReportParseError( + "Each TokenTagDispatch element must be a pair (token_id, rule)", delta_element + ); + } + auto id_node = std::get_if(tuple_node->elements[0].get()); + if (id_node == nullptr || id_node->value < 0) { + ReportParseError("Token trigger ID must be a non-negative integer", delta_element); + } + auto rule_node = std::get_if(tuple_node->elements[1].get()); + if (rule_node == nullptr) { + ReportParseError("Rule reference must be an identifier", delta_element); + } + auto rule_id = builder_.GetRuleId(rule_node->name); + if (rule_id == -1) { + ReportParseError("Rule \"" + rule_node->name + "\" is not defined", delta_element); + } + ttd.trigger_rule_pairs.push_back({static_cast(id_node->value), rule_id}); + } + + ttd.loop_after_dispatch = true; + if (auto it = args.named_arguments.find("loop_after_dispatch"); + it != args.named_arguments.end()) { + auto bool_node = std::get_if(it->second.get()); + if (bool_node == nullptr) { + ReportParseError("loop_after_dispatch must be a boolean", delta_element); + } + ttd.loop_after_dispatch = bool_node->value; + } + + if (auto it = args.named_arguments.find("excludes"); it != args.named_arguments.end()) { + auto tuple_node = std::get_if(it->second.get()); + if (tuple_node == nullptr) { + ReportParseError("excludes must be a tuple", delta_element); + } + for (const auto& element : tuple_node->elements) { + auto int_node = std::get_if(element.get()); + if (int_node == nullptr || int_node->value < 0) { + ReportParseError("Exclude token ID must be a non-negative integer", delta_element); + } + ttd.excludes.push_back(static_cast(int_node->value)); + } + } + + for (auto excl_id : ttd.excludes) { + for (const auto& [tid, _] : ttd.trigger_rule_pairs) { + if (tid == excl_id) { + ReportParseError( + "Token trigger ID " + std::to_string(tid) + " must not overlap with exclude token ID", + delta_element + ); + } + } + } + + return builder_.AddTokenTagDispatch(ttd); +} + +int32_t EBNFParser::ParseLookaheadAssertion() { + PeekAndConsume(TokenType::LookaheadLParen, "Expect (= in lookahead assertion"); + auto result = ParseChoices(); + PeekAndConsume(TokenType::RParen, "Expect )"); + return result; +} + +EBNFParser::ParsedRule EBNFParser::ParseRule() { + if (Peek().type != TokenType::RuleName) { + ReportParseError("Expect rule name"); + } + cur_rule_name_ = std::any_cast(Peek().value); + int32_t max_tokens = Peek().max_tokens; + int32_t max_chars = Peek().max_chars; + std::string capture_name = Peek().capture_name; + int32_t capture_hidden_suffix_bytes = Peek().capture_hidden_suffix_bytes; + int32_t capture_hidden_stop_bytes = Peek().capture_hidden_stop_bytes; + int32_t capture_hidden_body_rule_id = Peek().capture_hidden_body_rule_id; + int32_t capture_hidden_marker_rule_id = Peek().capture_hidden_marker_rule_id; + std::string stop_capture_name = Peek().stop_capture_name; + bool is_lazy = Peek().is_lazy; + std::optional temperature = Peek().temperature; + Consume(); + + PeekAndConsume(TokenType::Assign, "Expect ::="); + + auto body_id = ParseChoices(); + + int32_t lookahead_id = -1; + if (Peek().type == TokenType::LookaheadLParen) { + lookahead_id = ParseLookaheadAssertion(); + } + + ParsedRule result; + result.rule = Rule{cur_rule_name_, body_id, lookahead_id}; + result.rule.max_tokens = max_tokens; + result.rule.max_chars = max_chars; + result.rule.capture_name = capture_name; + result.rule.is_lazy = is_lazy; + result.rule.temperature = temperature; + result.suffix_stop_info.hidden_suffix_bytes = capture_hidden_suffix_bytes; + result.suffix_stop_info.hidden_stop_bytes = capture_hidden_stop_bytes; + result.suffix_stop_info.body_rule_id = capture_hidden_body_rule_id; + result.suffix_stop_info.marker_rule_id = capture_hidden_marker_rule_id; + result.suffix_stop_info.stop_capture_name = std::move(stop_capture_name); + return result; +} + +void EBNFParser::InitRuleNames() { + int delta_element = 0; + for (auto& token : tokens_) { + if (token.type == TokenType::RuleName) { + auto name = std::any_cast(token.value); + if (builder_.GetRuleId(name) != -1) { + ReportParseError("Rule \"" + name + "\" is defined multiple times", delta_element); + } + builder_.AddEmptyRule(name); + } + ++delta_element; + } + if (builder_.GetRuleId(root_rule_name_) == -1) { + ReportParseError("The root rule with name \"" + root_rule_name_ + "\" is not found", 0); + } +} + +Grammar EBNFParser::Parse( + const std::vector& tokens, + const std::string& root_rule_name, + const int& max_nest_layer +) { + max_nest_layer_ = max_nest_layer; + nest_layer_guard_ = 0; + tokens_ = tokens; + current_token_ = tokens_.data(); + root_rule_name_ = root_rule_name; + + // First collect rule names + InitRuleNames(); + + // Then parse all the rules + while (Peek().type != TokenType::EndOfFile) { + auto parsed_rule = ParseRule(); + const auto& rule = parsed_rule.rule; + builder_.UpdateRuleBody(rule.name, rule.body_expr_id); + builder_.UpdateLookaheadAssertion(rule.name, rule.lookahead_assertion_id); + builder_.UpdateMaxTokens(rule.name, rule.max_tokens); + builder_.UpdateMaxChars(rule.name, rule.max_chars); + builder_.UpdateCaptureName(rule.name, rule.capture_name); + builder_.UpdateSuffixStopInfo(rule.name, parsed_rule.suffix_stop_info); + builder_.UpdateLazy(rule.name, rule.is_lazy); + builder_.UpdateRuleTemperature(builder_.GetRuleId(rule.name), rule.temperature); + } + + return builder_.Get(root_rule_name); +} + +Grammar ParseEBNF(const std::string& ebnf_string, const std::string& root_rule_name) { + EBNFLexer lexer; + auto tokens = lexer.Tokenize(ebnf_string); + EBNFParser parser; + return parser.Parse(std::move(tokens), root_rule_name); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_parser.h b/zig/vendor/xgrammar/cpp/grammar_parser.h new file mode 100644 index 0000000000..bfb3d8727e --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_parser.h @@ -0,0 +1,107 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_parser.h + * \brief The header for the parser of BNF/EBNF grammar into BNF AST. + */ + +#ifndef XGRAMMAR_GRAMMAR_PARSER_H_ +#define XGRAMMAR_GRAMMAR_PARSER_H_ + +#include + +#include +#include + +namespace xgrammar { + +class EBNFLexer { + public: + // Token types + enum class TokenType { + RuleName, // the name of a rule definition, e.g.: root, rule1 + Identifier, // reference to a rule, or a Macro name, e.g.: root, rule1, TagDispatch + StringLiteral, // e.g.: "tag1", "hello" + BooleanLiteral, // true, false + IntegerLiteral, // 123 + LParen, // ( + RParen, // ) + LBrace, // { + RBrace, // } + Pipe, // | + Comma, // , + EndOfFile, // End of file + + // Symbols and quantifiers + Assign, // ::= + Equal, // = + Star, // * + Plus, // + + Question, // ? + + // Character class + LBracket, // [ + RBracket, // ] + Dash, // - + Caret, // ^ + CharInCharClass, // a character in a character class, e.g. a and z in [a-z]; escaped chars + // with no special meaning are also included, e.g. . in [a\.z] + EscapeInCharClass, // Escaped sequence with special function, e.g. \S in [\S] + + // Special structures + LookaheadLParen, // (= + }; + + // Token structure + struct Token { + TokenType type; + std::string lexeme; // original text + std::any value; // The processed value. Can be a int for integer literal, a string for string + // literal, etc. + int line; + int column; + // The token budget attached to a rule-definition identifier via name[max_tokens=N], or -1. + int32_t max_tokens = -1; + // The character budget attached to a rule-definition identifier via name[max_chars=N], or -1. + int32_t max_chars = -1; + // The capture name attached to a rule-definition identifier via name[capture="x"], or empty. + std::string capture_name = {}; + // Trailing bytes hidden only from the rule's own capture. + int32_t capture_hidden_suffix_bytes = 0; + // Trailing bytes hidden from the rule and every enclosing capture. + int32_t capture_hidden_stop_bytes = 0; + // Helper rule ids used to recover a variable-length suffix/stop marker boundary. + int32_t capture_hidden_body_rule_id = -1; + int32_t capture_hidden_marker_rule_id = -1; + // Capture name for the bytes matched by a suffix/stop marker. + std::string stop_capture_name = {}; + // Whether the identifier is a rule name carrying the [lazy] attribute, e.g. r[lazy] ::= ... + bool is_lazy = false; + // The sampling temperature attached to a rule-definition identifier via name[temperature=T]. + std::optional temperature = std::nullopt; + }; + + EBNFLexer(); + std::vector Tokenize(const std::string& input); + + XGRAMMAR_DEFINE_PIMPL_METHODS(EBNFLexer); +}; + +/*! + * \brief This class parses a BNF/EBNF grammar string into an BNF abstract syntax tree (AST). + * \details This function accepts the EBNF notation defined in the W3C XML Specification + * (https://www.w3.org/TR/xml/#sec-notation), which is a popular standard, with the following + * changes: + * - Using # as comment mark instead of C-style comments + * - Accept C-style unicode escape sequence \u01AB, \U000001AB, \xAB instead of #x0123 + * - Rule A-B (match A and not match B) is not supported yet + * + * See tests/python/serve/json.ebnf for an example. + * \param ebnf_string The grammar string. + * \param root_rule_name The name of the root rule. Default is "root". + * \return The parsed grammar. + */ +Grammar ParseEBNF(const std::string& ebnf_string, const std::string& root_rule_name = "root"); + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_PARSER_H_ diff --git a/zig/vendor/xgrammar/cpp/grammar_printer.cc b/zig/vendor/xgrammar/cpp/grammar_printer.cc new file mode 100644 index 0000000000..a598dfb9af --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_printer.cc @@ -0,0 +1,315 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_printer.cc + */ + +#include "grammar_printer.h" + +#include + +#include +#include +#include + +#include "support/encoding.h" + +namespace xgrammar { + +std::string GrammarPrinter::PrintRule(const Rule& rule, const SuffixStopInfo* suffix_stop_info) { + std::string res = rule.name; + // Print the attributes as one comma-separated bracket group, re-parseable by the EBNF lexer. + if (rule.max_tokens >= 0 || rule.max_chars >= 0 || !rule.capture_name.empty() || + suffix_stop_info != nullptr || rule.is_lazy || rule.temperature.has_value()) { + std::string attributes; + auto append_attribute = [&](const std::string& attribute) { + if (!attributes.empty()) { + attributes += ", "; + } + attributes += attribute; + }; + if (rule.max_tokens >= 0) { + append_attribute("max_tokens=" + std::to_string(rule.max_tokens)); + } + if (rule.max_chars >= 0) { + append_attribute("max_chars=" + std::to_string(rule.max_chars)); + } + if (!rule.capture_name.empty()) { + append_attribute("capture=\"" + rule.capture_name + "\""); + } + if (suffix_stop_info != nullptr && suffix_stop_info->hidden_suffix_bytes > 0) { + append_attribute( + "capture_hidden_suffix_bytes=" + std::to_string(suffix_stop_info->hidden_suffix_bytes) + ); + } + if (suffix_stop_info != nullptr && suffix_stop_info->hidden_stop_bytes > 0) { + append_attribute( + "capture_hidden_stop_bytes=" + std::to_string(suffix_stop_info->hidden_stop_bytes) + ); + } + if (suffix_stop_info != nullptr && suffix_stop_info->body_rule_id >= 0) { + append_attribute( + "capture_hidden_body_rule_id=" + std::to_string(suffix_stop_info->body_rule_id) + ); + append_attribute( + "capture_hidden_marker_rule_id=" + std::to_string(suffix_stop_info->marker_rule_id) + ); + } + if (suffix_stop_info != nullptr && !suffix_stop_info->stop_capture_name.empty()) { + append_attribute("stop_capture=\"" + suffix_stop_info->stop_capture_name + "\""); + } + if (rule.is_lazy) { + append_attribute("lazy"); + } + if (rule.temperature.has_value()) { + std::ostringstream temperature; + temperature << std::setprecision(std::numeric_limits::max_digits10) + << rule.temperature.value(); + append_attribute("temperature=" + temperature.str()); + } + res += "[" + attributes + "]"; + } + res += " ::= " + PrintGrammarExpr(rule.body_expr_id); + if (rule.lookahead_assertion_id != -1) { + res += " (=" + PrintGrammarExpr(rule.lookahead_assertion_id) + ")"; + } + return res; +} + +std::string GrammarPrinter::PrintRule(int32_t rule_id) { + return PrintRule(grammar_->GetRule(rule_id), grammar_->GetSuffixStopInfo(rule_id)); +} + +std::string GrammarPrinter::PrintGrammarExpr(const GrammarExpr& grammar_expr) { + std::string result; + switch (grammar_expr.type) { + case GrammarExprType::kByteString: + return PrintByteString(grammar_expr); + case GrammarExprType::kCharacterClass: + return PrintCharacterClass(grammar_expr); + case GrammarExprType::kCharacterClassStar: + return PrintCharacterClassStar(grammar_expr); + case GrammarExprType::kEmptyStr: + return PrintEmptyStr(grammar_expr); + case GrammarExprType::kRuleRef: + return PrintRuleRef(grammar_expr); + case GrammarExprType::kSequence: + return PrintSequence(grammar_expr); + case GrammarExprType::kChoices: + return PrintChoices(grammar_expr); + case GrammarExprType::kTagDispatch: + return PrintTagDispatch(grammar_expr); + case GrammarExprType::kRepeat: + return PrintRepeat(grammar_expr); + case GrammarExprType::kToken: + return PrintToken(grammar_expr); + case GrammarExprType::kExcludeToken: + return PrintExcludeToken(grammar_expr); + case GrammarExprType::kTokenTagDispatch: + return PrintTokenTagDispatch(grammar_expr); + case GrammarExprType::kRegex: + return PrintRegex(grammar_expr); + case GrammarExprType::kSubstring: + return PrintSubstring(grammar_expr); + default: + XGRAMMAR_LOG(FATAL) << "Unexpected GrammarExpr type: " << static_cast(grammar_expr.type); + XGRAMMAR_UNREACHABLE(); + } +} + +std::string GrammarPrinter::PrintGrammarExpr(int32_t grammar_expr_id) { + return PrintGrammarExpr(grammar_->GetGrammarExpr(grammar_expr_id)); +} + +std::string GrammarPrinter::PrintByteString(const GrammarExpr& grammar_expr) { + std::string internal_str; + internal_str.reserve(grammar_expr.data_len); + for (int i = 0; i < grammar_expr.data_len; ++i) { + internal_str += static_cast(grammar_expr[i]); + } + return "\"" + EscapeString(internal_str) + "\""; +} + +std::string GrammarPrinter::PrintCharacterClass(const GrammarExpr& grammar_expr) { + static const std::unordered_map kCustomEscapeMap = { + {'-', "\\-"}, {']', "\\]"} + }; + std::string result = "["; + bool is_negative = static_cast(grammar_expr[0]); + if (is_negative) { + result += "^"; + } + for (auto i = 1; i < grammar_expr.data_len; i += 2) { + result += EscapeString(grammar_expr[i], kCustomEscapeMap); + if (grammar_expr[i] == grammar_expr[i + 1]) { + continue; + } + result += "-"; + result += EscapeString(grammar_expr[i + 1], kCustomEscapeMap); + } + result += "]"; + return result; +} + +std::string GrammarPrinter::PrintCharacterClassStar(const GrammarExpr& grammar_expr) { + return PrintCharacterClass(grammar_expr) + "*"; +} + +std::string GrammarPrinter::PrintEmptyStr(const GrammarExpr& grammar_expr) { return "\"\""; } + +std::string GrammarPrinter::PrintRuleRef(const GrammarExpr& grammar_expr) { + return grammar_->GetRule(grammar_expr[0]).name; +} + +std::string GrammarPrinter::PrintSequence(const GrammarExpr& grammar_expr) { + std::string result; + result += "("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + result += PrintGrammarExpr(grammar_expr[i]); + if (i + 1 != grammar_expr.data_len) { + result += " "; + } + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintChoices(const GrammarExpr& grammar_expr) { + std::string result; + + result += "("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + result += PrintGrammarExpr(grammar_expr[i]); + if (i + 1 != grammar_expr.data_len) { + result += " | "; + } + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintRegex(const GrammarExpr& grammar_expr) { + std::string result = "Regex(" + PrintString(grammar_->GetRegexString(grammar_expr)); + if (grammar_->GetRegexIsJSONString(grammar_expr)) { + result += ", json_string=true"; + } + return result + ")"; +} + +std::string GrammarPrinter::PrintSubstring(const GrammarExpr& grammar_expr) { + // EscapeString(std::string) stops at embedded NUL bytes, so escape codepoint by codepoint to + // keep NUL chunks (allowed by substring expressions) re-parseable. + auto escape_chunk = [](const std::string& chunk) { + std::string result = "\""; + size_t offset = 0; + while (offset < chunk.size()) { + if (chunk[offset] == '\0') { + result += "\\0"; + ++offset; + continue; + } + auto [codepoint, length] = ParseNextUTF8(chunk.c_str() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8) { + result += EscapeString(static_cast(chunk[offset])); + ++offset; + continue; + } + result += EscapeString(codepoint); + offset += static_cast(length); + } + return result + "\""; + }; + + auto chunks = grammar_->GetSubstringChunks(grammar_expr); + std::string result = "Substring("; + for (size_t i = 0; i < chunks.size(); ++i) { + if (i > 0) { + result += ", "; + } + result += escape_chunk(chunks[i]); + } + return result + ")"; +} + +std::string GrammarPrinter::PrintString(const std::string& str) { + return "\"" + EscapeString(str) + "\""; +} + +std::string GrammarPrinter::PrintBoolean(bool value) { return value ? "true" : "false"; } + +std::string GrammarPrinter::PrintTagDispatch(const GrammarExpr& grammar_expr) { + auto tag_dispatch = grammar_->GetTagDispatch(grammar_expr); + std::string result = "TagDispatch(\n"; + std::string indent = " "; + for (const auto& [trigger, rule_id] : tag_dispatch.tag_rule_pairs) { + result += indent + "(" + PrintString(trigger) + ", " + grammar_->GetRule(rule_id).name + "),\n"; + } + result += + indent + "loop_after_dispatch=" + PrintBoolean(tag_dispatch.loop_after_dispatch) + ",\n"; + result += indent + "excludes=("; + for (int i = 0; i < static_cast(tag_dispatch.excludes.size()); ++i) { + if (i > 0) result += ", "; + result += PrintString(tag_dispatch.excludes[i]); + } + result += ")\n)"; + return result; +} + +std::string GrammarPrinter::PrintRepeat(const GrammarExpr& grammar_expr) { + int32_t lower_bound = grammar_expr[1]; + int32_t upper_bound = grammar_expr[2]; + std::string result = grammar_->GetRule(grammar_expr[0]).name + "{"; + result += std::to_string(lower_bound); + result += ", "; + result += std::to_string(upper_bound); + result += "}"; + return result; +} + +std::string GrammarPrinter::PrintToken(const GrammarExpr& grammar_expr) { + std::string result = "Token("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + if (i > 0) result += ", "; + result += std::to_string(grammar_expr[i]); + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintExcludeToken(const GrammarExpr& grammar_expr) { + std::string result = "ExcludeToken("; + for (int i = 0; i < grammar_expr.data_len; ++i) { + if (i > 0) result += ", "; + result += std::to_string(grammar_expr[i]); + } + result += ")"; + return result; +} + +std::string GrammarPrinter::PrintTokenTagDispatch(const GrammarExpr& grammar_expr) { + auto ttd = grammar_->GetTokenTagDispatch(grammar_expr); + std::string result = "TokenTagDispatch(\n"; + std::string indent = " "; + for (const auto& [token_id, rule_id] : ttd.trigger_rule_pairs) { + result += + indent + "(" + std::to_string(token_id) + ", " + grammar_->GetRule(rule_id).name + "),\n"; + } + result += indent + "loop_after_dispatch=" + PrintBoolean(ttd.loop_after_dispatch) + ",\n"; + result += indent + "excludes=("; + for (int i = 0; i < static_cast(ttd.excludes.size()); ++i) { + if (i > 0) result += ", "; + result += std::to_string(ttd.excludes[i]); + } + result += ")\n)"; + return result; +} + +std::string GrammarPrinter::ToString() { + std::string result; + int num_rules = grammar_->NumRules(); + for (auto i = 0; i < num_rules; ++i) { + result += PrintRule(i) + "\n"; + } + return result; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/grammar_printer.h b/zig/vendor/xgrammar/cpp/grammar_printer.h new file mode 100644 index 0000000000..4d5b1ef844 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/grammar_printer.h @@ -0,0 +1,86 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar_printer.h + * \brief The header for printing the AST of a BNF grammar. + */ + +#ifndef XGRAMMAR_GRAMMAR_PRINTER_H_ +#define XGRAMMAR_GRAMMAR_PRINTER_H_ + +#include + +#include + +#include "grammar_impl.h" + +namespace xgrammar { + +/*! + * \brief Prints the BNF AST with standard BNF format. + */ +class GrammarPrinter { + private: + using Rule = Grammar::Impl::Rule; + using SuffixStopInfo = Grammar::Impl::SuffixStopInfo; + using GrammarExprType = Grammar::Impl::GrammarExprType; + using GrammarExpr = Grammar::Impl::GrammarExpr; + + public: + /*! + * \brief Constructor. + * \param grammar The grammar to print. + */ + explicit GrammarPrinter(const Grammar& grammar) : grammar_(grammar) {} + + /*! \brief Print the complete grammar. */ + std::string ToString(); + + /*! \brief Print a rule. */ + std::string PrintRule(const Rule& rule, const SuffixStopInfo* suffix_stop_info); + /*! \brief Print a rule corresponding to the given id. */ + std::string PrintRule(int32_t rule_id); + /*! \brief Print a GrammarExpr. */ + std::string PrintGrammarExpr(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr corresponding to the given id. */ + std::string PrintGrammarExpr(int32_t grammar_expr_id); + + private: + /*! \brief Print a GrammarExpr for byte string. */ + std::string PrintByteString(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for character class. */ + std::string PrintCharacterClass(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for a star quantifier of a character class. */ + std::string PrintCharacterClassStar(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for empty string. */ + std::string PrintEmptyStr(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for rule reference. */ + std::string PrintRuleRef(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for grammar_expr sequence. */ + std::string PrintSequence(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for grammar_expr choices. */ + std::string PrintChoices(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for tag dispatch. */ + std::string PrintTagDispatch(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for repeat. */ + std::string PrintRepeat(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for token. */ + std::string PrintToken(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for exclude token. */ + std::string PrintExcludeToken(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for token tag dispatch. */ + std::string PrintTokenTagDispatch(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for regex. */ + std::string PrintRegex(const GrammarExpr& grammar_expr); + /*! \brief Print a GrammarExpr for substring. */ + std::string PrintSubstring(const GrammarExpr& grammar_expr); + /*! \brief Print a string. */ + std::string PrintString(const std::string& str); + /*! \brief Print a boolean. */ + std::string PrintBoolean(bool value); + + Grammar grammar_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_PRINTER_H_ diff --git a/zig/vendor/xgrammar/cpp/json_schema_converter.cc b/zig/vendor/xgrammar/cpp/json_schema_converter.cc new file mode 100644 index 0000000000..9261cdfd8a --- /dev/null +++ b/zig/vendor/xgrammar/cpp/json_schema_converter.cc @@ -0,0 +1,5126 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter.cc + * \brief Implementation of JSONSchemaConverter and related utilities. + */ +#include "json_schema_converter.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "json_schema_converter_ext.h" +#include "regex_converter.h" +#include "support/json_parse.h" +#include "support/logging.h" + +namespace xgrammar { + +// ==================== Spec ToString implementations ==================== + +std::string IntegerSpec::ToString() const { + return "IntegerSpec{minimum=" + (minimum.has_value() ? std::to_string(*minimum) : "null") + + ", maximum=" + (maximum.has_value() ? std::to_string(*maximum) : "null") + + ", exclusive_minimum=" + + (exclusive_minimum.has_value() ? std::to_string(*exclusive_minimum) : "null") + + ", exclusive_maximum=" + + (exclusive_maximum.has_value() ? std::to_string(*exclusive_maximum) : "null") + + ", multiple_of=" + (multiple_of.has_value() ? std::to_string(*multiple_of) : "null") + "}"; +} + +std::string NumberSpec::ToString() const { + return "NumberSpec{minimum=" + (minimum.has_value() ? std::to_string(*minimum) : "null") + + ", maximum=" + (maximum.has_value() ? std::to_string(*maximum) : "null") + + ", exclusive_minimum=" + + (exclusive_minimum.has_value() ? std::to_string(*exclusive_minimum) : "null") + + ", exclusive_maximum=" + + (exclusive_maximum.has_value() ? std::to_string(*exclusive_maximum) : "null") + "}"; +} + +std::string StringSpec::ToString() const { + return "StringSpec{pattern=" + (pattern.has_value() ? "\"" + *pattern + "\"" : "null") + + ", format=" + (format.has_value() ? "\"" + *format + "\"" : "null") + + ", min_length=" + std::to_string(min_length) + + ", max_length=" + std::to_string(max_length) + "}"; +} + +std::string BooleanSpec::ToString() const { return "BooleanSpec{}"; } + +std::string NullSpec::ToString() const { return "NullSpec{}"; } + +std::string AnySpec::ToString() const { return "AnySpec{}"; } + +std::string ArraySpec::ToString() const { + return "ArraySpec{prefix_items.size()=" + std::to_string(prefix_items.size()) + + ", allow_additional_items=" + (allow_additional_items ? "true" : "false") + + ", additional_items=" + (additional_items ? "SchemaSpec" : "null") + + ", min_items=" + std::to_string(min_items) + ", max_items=" + std::to_string(max_items) + + "}"; +} + +std::string ObjectSpec::ToString() const { + std::string s = + "ObjectSpec{properties.size()=" + std::to_string(properties.size()) + ", properties=["; + for (size_t i = 0; i < properties.size(); ++i) { + if (i != 0) s += ", "; + s += properties[i].name; + } + s += "], pattern_properties.size()=" + std::to_string(pattern_properties.size()) + ", required=["; + bool first = true; + for (const auto& r : required) { + if (!first) s += ", "; + s += r; + first = false; + } + s += + std::string("], allow_additional_properties=") + + (allow_additional_properties ? "true" : "false") + + ", additional_properties_schema=" + (additional_properties_schema ? "SchemaSpec" : "null") + + ", allow_unevaluated_properties=" + (allow_unevaluated_properties ? "true" : "false") + + ", unevaluated_properties_schema=" + (unevaluated_properties_schema ? "SchemaSpec" : "null") + + ", property_names=" + (property_names ? "SchemaSpec" : "null") + + ", min_properties=" + std::to_string(min_properties) + + ", max_properties=" + std::to_string(max_properties) + "}"; + return s; +} + +std::string ConstSpec::ToString() const { return "ConstSpec{json_value=\"" + json_value + "\"}"; } + +std::string EnumSpec::ToString() const { + std::string s = + "EnumSpec{json_values.size()=" + std::to_string(json_values.size()) + ", json_values=["; + for (size_t i = 0; i < json_values.size(); ++i) { + if (i != 0) s += ", "; + s += "\"" + json_values[i] + "\""; + } + s += "]}"; + return s; +} + +std::string RefSpec::ToString() const { return "RefSpec{uri=\"" + uri + "\"}"; } + +std::string AnyOfSpec::ToString() const { + return "AnyOfSpec{options.size()=" + std::to_string(options.size()) + "}"; +} + +std::string OneOfSpec::ToString() const { + return "OneOfSpec{options.size()=" + std::to_string(options.size()) + "}"; +} + +std::string AllOfSpec::ToString() const { + return "AllOfSpec{schemas.size()=" + std::to_string(schemas.size()) + "}"; +} + +std::string TypeArraySpec::ToString() const { + return "TypeArraySpec{type_schemas.size()=" + std::to_string(type_schemas.size()) + "}"; +} + +std::string SchemaSpec::ToString() const { + std::string spec_str; + std::visit([&spec_str](const auto& s) { spec_str = s.ToString(); }, spec); + return "SchemaSpec{spec=" + spec_str + ", cache_key=\"" + cache_key + "\", rule_name_hint=\"" + + rule_name_hint + "\"}"; +} + +// ==================== SchemaParser (Internal) ==================== + +namespace { + +enum class SchemaErrorType : int { + kInvalidSchema = 0, + kUnsatisfiableSchema = 1, + kUnsupportedSchema = 2, +}; + +using SchemaError = TypedError; + +// Unbounded integer multipleOf emits a modulo DFA: states ~= N, transitions ~= 10N. +// Fail closed above the cap to keep generated grammars bounded. +constexpr int64_t kIntegerMultipleOfMax = 1024; +constexpr int64_t kIntegerMultipleOfRangeWidthMax = 10000; + +bool IsMultipleOf(int64_t value, int64_t multiple_of) { return (value % multiple_of) == 0; } + +bool HasMultipleInRange(int64_t start, int64_t end, int64_t multiple_of) { + for (int64_t value = start; value <= end; ++value) { + if (IsMultipleOf(value, multiple_of)) return true; + if (value == std::numeric_limits::max()) break; + } + return false; +} + +constexpr const char* kUnsupportedOneOfMessage = + "oneOf with overlapping or non-provably-disjoint branches cannot be represented exactly; " + "falling back to anyOf semantics"; + +bool IsSchemaAnnotationKey(const std::string& key) { + static const std::unordered_set kAnnotationKeys = { + "title", + "default", + "description", + "examples", + "deprecated", + "readOnly", + "writeOnly", + "$comment", + "$schema", + }; + return kAnnotationKeys.count(key) != 0; +} + +bool HasOnlyKeys( + const picojson::object& schema, const std::unordered_set& allowed_keys +) { + for (const auto& [key, _] : schema) { + if (allowed_keys.count(key) == 0 && !IsSchemaAnnotationKey(key)) { + return false; + } + } + return true; +} + +bool IsSupportedJSONType(const std::string& type) { + static const std::unordered_set kTypes = { + "null", + "boolean", + "object", + "array", + "number", + "string", + "integer", + }; + return kTypes.count(type) != 0; +} + +bool NormalizeTypeSet( + const picojson::value& type_value, std::unordered_set* type_set +) { + if (type_value.is()) { + const auto& type = type_value.get(); + if (!IsSupportedJSONType(type)) { + return false; + } + type_set->insert(type); + return true; + } + if (!type_value.is()) { + return false; + } + + const auto& type_array = type_value.get(); + if (type_array.empty()) { + return false; + } + for (const auto& item : type_array) { + if (!item.is()) { + return false; + } + const auto& type = item.get(); + if (!IsSupportedJSONType(type)) { + return false; + } + type_set->insert(type); + } + return true; +} + +bool IsNumericValue(const picojson::value& value) { + return value.is() || value.is(); +} + +bool IsIntegerValue(const picojson::value& value) { + if (value.is()) { + return true; + } + if (!value.is()) { + return false; + } + double number = value.get(); + return std::isfinite(number) && std::floor(number) == number; +} + +bool JSONValuesMayOverlap(const picojson::value& lhs, const picojson::value& rhs) { + if (IsNumericValue(lhs) || IsNumericValue(rhs)) { + if (!IsNumericValue(lhs) || !IsNumericValue(rhs)) { + return false; + } + if (lhs.is() && rhs.is()) { + return lhs.get() == rhs.get(); + } + return true; + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is(); + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is() && lhs.get() == rhs.get(); + } + if (lhs.is() || rhs.is()) { + return lhs.is() && rhs.is() && + lhs.get() == rhs.get(); + } + if (lhs.is() || rhs.is()) { + if (!lhs.is() || !rhs.is()) { + return false; + } + const auto& lhs_array = lhs.get(); + const auto& rhs_array = rhs.get(); + if (lhs_array.size() != rhs_array.size()) { + return false; + } + for (size_t i = 0; i < lhs_array.size(); ++i) { + if (!JSONValuesMayOverlap(lhs_array[i], rhs_array[i])) { + return false; + } + } + return true; + } + if (lhs.is() || rhs.is()) { + if (!lhs.is() || !rhs.is()) { + return false; + } + const auto& lhs_object = lhs.get(); + const auto& rhs_object = rhs.get(); + if (lhs_object.size() != rhs_object.size()) { + return false; + } + for (const auto& [key, lhs_value] : lhs_object) { + auto rhs_it = rhs_object.find(key); + if (rhs_it == rhs_object.end() || !JSONValuesMayOverlap(lhs_value, rhs_it->second)) { + return false; + } + } + return true; + } + return lhs.serialize() == rhs.serialize(); +} + +bool ValueMatchesType(const picojson::value& value, const std::string& type) { + if (type == "null") { + return value.is(); + } + if (type == "boolean") { + return value.is(); + } + if (type == "string") { + return value.is(); + } + if (type == "integer") { + return IsIntegerValue(value); + } + if (type == "number") { + return IsNumericValue(value); + } + if (type == "array") { + return value.is(); + } + if (type == "object") { + return value.is(); + } + return false; +} + +bool IsRangeWidthOverCap(int64_t start, int64_t end, int64_t cap) { + uint64_t cap_u = static_cast(cap); + if (start <= 0 && end >= 0) { + // Count [start, end] inclusively without evaluating -INT64_MIN or overflowing the sum. + uint64_t negative_count = start < 0 ? static_cast(-(start + 1)) + 1 : 0; + if (negative_count > cap_u) return true; + uint64_t remaining = cap_u - negative_count; + if (remaining == 0) return true; + --remaining; // zero + uint64_t positive_count = end > 0 ? static_cast(end) : 0; + return positive_count > remaining; + } + + uint64_t value_count = static_cast(end - start) + 1; + return value_count > cap_u; +} + +// Effective inclusive integer range after folding exclusive bounds into minimum/maximum. A nullopt +// side means that side is unbounded. +struct EffectiveIntegerRange { + std::optional start; + std::optional end; +}; + +// Fold the inclusive [minimum, maximum] bounds together with any exclusive bounds so the stricter +// bound wins on each side. Shared by ParseInteger (range validation) and GenerateInteger (grammar +// emission) so the two can never disagree about the effective range. Precondition: +// exclusive_minimum != INT64_MAX and exclusive_maximum != INT64_MIN (ParseInteger rejects those +// before building the spec), so the +1/-1 below cannot overflow. +EffectiveIntegerRange ComputeEffectiveIntegerRange(const IntegerSpec& spec) { + EffectiveIntegerRange range; + if (spec.minimum.has_value()) { + range.start = spec.minimum; + } + if (spec.exclusive_minimum.has_value()) { + // Smallest integer strictly greater than exclusive_minimum; the larger lower bound wins. + int64_t excl_start = *spec.exclusive_minimum + 1; + range.start = range.start.has_value() ? std::max(*range.start, excl_start) : excl_start; + } + if (spec.maximum.has_value()) { + range.end = spec.maximum; + } + if (spec.exclusive_maximum.has_value()) { + // Largest integer strictly less than exclusive_maximum; the smaller upper bound wins. + int64_t excl_end = *spec.exclusive_maximum - 1; + range.end = range.end.has_value() ? std::min(*range.end, excl_end) : excl_end; + } + return range; +} + +bool TypeSetsOverlap( + const std::unordered_set& lhs, const std::unordered_set& rhs +) { + for (const auto& lhs_type : lhs) { + for (const auto& rhs_type : rhs) { + if (lhs_type == rhs_type) { + return true; + } + if ((lhs_type == "integer" || lhs_type == "number") && + (rhs_type == "integer" || rhs_type == "number")) { + return true; + } + } + } + return false; +} + +bool FiniteValuesOverlap( + const std::vector& lhs, const std::vector& rhs +) { + for (const auto& lhs_value : lhs) { + for (const auto& rhs_value : rhs) { + if (JSONValuesMayOverlap(lhs_value, rhs_value)) { + return true; + } + } + } + return false; +} + +bool FiniteValuesOverlapTypeSet( + const std::vector& values, const std::unordered_set& type_set +) { + for (const auto& value : values) { + if (IsNumericValue(value) && (type_set.count("integer") || type_set.count("number"))) { + return true; + } + for (const auto& type : type_set) { + if (ValueMatchesType(value, type)) { + return true; + } + } + } + return false; +} + +bool TryGetFiniteValues(const picojson::object& schema, std::vector* values) { + if (schema.count("const")) { + values->push_back(schema.at("const")); + return true; + } + if (schema.count("enum")) { + if (!schema.at("enum").is()) { + return false; + } + const auto& enum_values = schema.at("enum").get(); + if (enum_values.empty()) { + return false; + } + values->insert(values->end(), enum_values.begin(), enum_values.end()); + return true; + } + return false; +} + +struct OneOfArmProof { + enum class Kind { kTypeSet, kFiniteValues }; + + Kind kind; + std::unordered_set type_set; + std::vector finite_values; +}; + +std::optional ClassifyTypeOrFiniteOneOfArm(const picojson::value& option) { + if (!option.is()) { + return std::nullopt; + } + const auto& schema = option.get(); + + if (schema.count("$ref") || schema.count("anyOf") || schema.count("allOf") || + schema.count("oneOf")) { + return std::nullopt; + } + + std::vector finite_values; + if (TryGetFiniteValues(schema, &finite_values)) { + OneOfArmProof proof; + proof.kind = OneOfArmProof::Kind::kFiniteValues; + proof.finite_values = std::move(finite_values); + return proof; + } + + if (!schema.count("type") || !HasOnlyKeys(schema, {"type"})) { + return std::nullopt; + } + + std::unordered_set type_set; + if (!NormalizeTypeSet(schema.at("type"), &type_set)) { + return std::nullopt; + } + if (type_set.count("object")) { + return std::nullopt; + } + + OneOfArmProof proof; + proof.kind = OneOfArmProof::Kind::kTypeSet; + proof.type_set = std::move(type_set); + return proof; +} + +bool OneOfArmProofsAreDisjoint(const OneOfArmProof& lhs, const OneOfArmProof& rhs) { + if (lhs.kind == OneOfArmProof::Kind::kTypeSet && rhs.kind == OneOfArmProof::Kind::kTypeSet) { + return !TypeSetsOverlap(lhs.type_set, rhs.type_set); + } + if (lhs.kind == OneOfArmProof::Kind::kFiniteValues && + rhs.kind == OneOfArmProof::Kind::kFiniteValues) { + return !FiniteValuesOverlap(lhs.finite_values, rhs.finite_values); + } + if (lhs.kind == OneOfArmProof::Kind::kFiniteValues && rhs.kind == OneOfArmProof::Kind::kTypeSet) { + return !FiniteValuesOverlapTypeSet(lhs.finite_values, rhs.type_set); + } + return !FiniteValuesOverlapTypeSet(rhs.finite_values, lhs.type_set); +} + +std::optional> GetDiscriminatorValues( + const picojson::value& option, const std::string& discriminator_key +) { + if (!option.is()) { + return std::nullopt; + } + const auto& schema = option.get(); + if (schema.count("$ref") || schema.count("anyOf") || schema.count("allOf") || + schema.count("oneOf")) { + return std::nullopt; + } + if (!schema.count("type") || !schema.at("type").is() || + schema.at("type").get() != "object") { + return std::nullopt; + } + if (!schema.count("required") || !schema.at("required").is()) { + return std::nullopt; + } + + bool requires_discriminator = false; + for (const auto& required_key : schema.at("required").get()) { + if (!required_key.is()) { + return std::nullopt; + } + if (required_key.get() == discriminator_key) { + requires_discriminator = true; + } + } + if (!requires_discriminator) { + return std::nullopt; + } + + if (!schema.count("properties") || !schema.at("properties").is()) { + return std::nullopt; + } + const auto& properties = schema.at("properties").get(); + auto property_it = properties.find(discriminator_key); + if (property_it == properties.end() || !property_it->second.is()) { + return std::nullopt; + } + + std::vector values; + if (!TryGetFiniteValues(property_it->second.get(), &values)) { + return std::nullopt; + } + return values; +} + +std::vector GetDiscriminatorCandidates(const picojson::value& option) { + std::vector candidates; + if (!option.is()) { + return candidates; + } + const auto& schema = option.get(); + if (!schema.count("required") || !schema.at("required").is() || + !schema.count("properties") || !schema.at("properties").is()) { + return candidates; + } + const auto& properties = schema.at("properties").get(); + for (const auto& required_key : schema.at("required").get()) { + if (!required_key.is()) { + continue; + } + const auto& key = required_key.get(); + auto property_it = properties.find(key); + if (property_it == properties.end() || !property_it->second.is()) { + continue; + } + std::vector values; + if (TryGetFiniteValues(property_it->second.get(), &values)) { + candidates.push_back(key); + } + } + return candidates; +} + +bool TryProveStrictDiscriminatorOneOf(const picojson::array& options) { + if (options.empty()) { + return false; + } + + for (const auto& discriminator_key : GetDiscriminatorCandidates(options.front())) { + std::vector> branch_values; + bool all_branches_have_key = true; + for (const auto& option : options) { + auto values = GetDiscriminatorValues(option, discriminator_key); + if (!values.has_value()) { + all_branches_have_key = false; + break; + } + branch_values.push_back(std::move(values.value())); + } + if (!all_branches_have_key) { + continue; + } + + bool pairwise_disjoint = true; + for (size_t i = 0; i < branch_values.size() && pairwise_disjoint; ++i) { + for (size_t j = i + 1; j < branch_values.size(); ++j) { + if (FiniteValuesOverlap(branch_values[i], branch_values[j])) { + pairwise_disjoint = false; + break; + } + } + } + if (pairwise_disjoint) { + return true; + } + } + return false; +} + +bool TryProveTypeOrFiniteOneOf(const picojson::array& options) { + std::vector proofs; + proofs.reserve(options.size()); + for (const auto& option : options) { + auto proof = ClassifyTypeOrFiniteOneOfArm(option); + if (!proof.has_value()) { + return false; + } + proofs.push_back(std::move(proof.value())); + } + + for (size_t i = 0; i < proofs.size(); ++i) { + for (size_t j = i + 1; j < proofs.size(); ++j) { + if (!OneOfArmProofsAreDisjoint(proofs[i], proofs[j])) { + return false; + } + } + } + return true; +} + +bool TryProvePairwiseDisjointOneOf(const picojson::array& options) { + return TryProveStrictDiscriminatorOneOf(options) || TryProveTypeOrFiniteOneOf(options); +} + +/*! + * \brief Parser for JSON Schema, converts JSON Schema to SchemaSpec intermediate representation. + */ +class SchemaParser { + public: + struct Config { + bool strict_mode = false; + JSONFormat json_format; + }; + + explicit SchemaParser(const picojson::value& root_schema, const Config& config) + : config_(config), root_schema_(root_schema) {} + + Result Parse( + const picojson::value& schema, + const std::string& rule_name_hint = "root", + std::optional default_type = std::nullopt + ); + + const picojson::value& GetRootSchema() const { return root_schema_; } + bool IsStrictMode() const { return config_.strict_mode; } + + Result ResolveRef( + const std::string& uri, const std::string& rule_name_hint + ); + + private: + Result ParseInteger(const picojson::object& schema); + Result ParseNumber(const picojson::object& schema); + Result ParseString(const picojson::object& schema); + Result ParseBoolean(const picojson::object& schema); + Result ParseNull(const picojson::object& schema); + Result ParseArray(const picojson::object& schema); + Result ParseObject(const picojson::object& schema); + Result ParseConst(const picojson::object& schema); + Result ParseEnum(const picojson::object& schema); + Result ParseRef(const picojson::object& schema); + Result ParseAnyOf( + const picojson::object& schema, const std::string& keyword + ); + Result ParseOneOf(const picojson::object& schema); + Result ParseAllOf(const picojson::object& schema); + Result ParseTypeArray( + const picojson::object& schema, const std::string& rule_name_hint + ); + + std::string ComputeCacheKey(const picojson::value& schema); + + static void WarnUnsupportedKeywords( + const picojson::object& schema, const std::vector& keywords, bool verbose = false + ); + + Config config_; + picojson::value root_schema_; + std::unordered_map ref_cache_; + std::unordered_map schema_cache_; +}; + +std::string SchemaParser::ComputeCacheKey(const picojson::value& schema) { + static const std::unordered_set kSkippedKeys = { + "title", + "default", + "description", + "examples", + "deprecated", + "readOnly", + "writeOnly", + "$comment", + "$schema", + }; + + if (schema.is()) { + std::string result = "{"; + std::vector> sorted_kv; + for (const auto& kv : schema.get()) { + if (kSkippedKeys.count(kv.first) == 0) { + sorted_kv.push_back(kv); + } + } + std::sort(sorted_kv.begin(), sorted_kv.end(), [](const auto& lhs, const auto& rhs) { + return lhs.first < rhs.first; + }); + int64_t idx = 0; + for (const auto& [key, value] : sorted_kv) { + if (idx != 0) { + result += ","; + } + ++idx; + result += "\"" + key + "\":" + ComputeCacheKey(value); + } + return result + "}"; + } else if (schema.is()) { + std::string result = "["; + int64_t idx = 0; + for (const auto& item : schema.get()) { + if (idx != 0) { + result += ","; + } + ++idx; + result += ComputeCacheKey(item); + } + return result + "]"; + } + return schema.serialize(false); +} + +void SchemaParser::WarnUnsupportedKeywords( + const picojson::object& schema, const std::vector& keywords, bool verbose +) { + if (!verbose) { + return; + } + for (const auto& keyword : keywords) { + if (schema.find(keyword) != schema.end()) { + XGRAMMAR_LOG(WARNING) << "Keyword " << keyword << " is not supported"; + } + } +} + +Result SchemaParser::Parse( + const picojson::value& schema, + const std::string& rule_name_hint, + std::optional default_type +) { + std::string cache_key = ComputeCacheKey(schema); + if (schema_cache_.count(cache_key)) { + return ResultOk(schema_cache_[cache_key]); + } + + if (schema.is()) { + if (!schema.get()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Schema 'false' cannot accept any value" + ); + } + auto spec = SchemaSpec::Make(AnySpec{}, cache_key, rule_name_hint); + schema_cache_[cache_key] = spec; + return ResultOk(spec); + } + + if (!schema.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "Schema should be an object or bool, but got " + schema.serialize(false) + ); + } + + const auto& schema_obj = schema.get(); + WarnUnsupportedKeywords( + schema_obj, {"not", "if", "then", "else", "dependentRequired", "dependentSchemas"} + ); + + SchemaSpecPtr result; + + if (schema_obj.count("$ref")) { + auto ref_result = ParseRef(schema_obj); + if (ref_result.IsErr()) return ResultErr(std::move(ref_result).UnwrapErr()); + auto ref_spec = std::move(ref_result).Unwrap(); + result = SchemaSpec::Make(std::move(ref_spec), cache_key, rule_name_hint); + } else if (schema_obj.count("const")) { + auto const_result = ParseConst(schema_obj); + if (const_result.IsErr()) return ResultErr(std::move(const_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(const_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("enum")) { + auto enum_result = ParseEnum(schema_obj); + if (enum_result.IsErr()) return ResultErr(std::move(enum_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(enum_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("anyOf")) { + auto anyof_result = ParseAnyOf(schema_obj, "anyOf"); + if (anyof_result.IsErr()) return ResultErr(std::move(anyof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(anyof_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("oneOf")) { + auto oneof_result = ParseOneOf(schema_obj); + if (oneof_result.IsErr()) { + if (oneof_result.ErrRef().Type() != SchemaErrorType::kUnsupportedSchema) { + return ResultErr(std::move(oneof_result).UnwrapErr()); + } + XGRAMMAR_LOG(WARNING) << oneof_result.ErrRef().what(); + auto anyof_result = ParseAnyOf(schema_obj, "oneOf"); + if (anyof_result.IsErr()) return ResultErr(std::move(anyof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(anyof_result).Unwrap(), cache_key, rule_name_hint); + } else { + result = SchemaSpec::Make(std::move(oneof_result).Unwrap(), cache_key, rule_name_hint); + } + } else if (schema_obj.count("allOf")) { + auto allof_result = ParseAllOf(schema_obj); + if (allof_result.IsErr()) return ResultErr(std::move(allof_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(allof_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("type") || default_type.has_value()) { + if (schema_obj.count("type") && schema_obj.at("type").is()) { + auto type_array_result = ParseTypeArray(schema_obj, rule_name_hint); + if (type_array_result.IsErr()) return ResultErr(std::move(type_array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(type_array_result).Unwrap(), cache_key, rule_name_hint); + } else { + if (schema_obj.count("type") && !schema_obj.at("type").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Type should be a string"); + } + const std::string& type = schema_obj.count("type") ? schema_obj.at("type").get() + : default_type.value(); + if (type == "integer") { + auto int_result = ParseInteger(schema_obj); + if (int_result.IsErr()) return ResultErr(std::move(int_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(int_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "number") { + auto num_result = ParseNumber(schema_obj); + if (num_result.IsErr()) return ResultErr(std::move(num_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(num_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "string") { + auto str_result = ParseString(schema_obj); + if (str_result.IsErr()) return ResultErr(std::move(str_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(str_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "boolean") { + auto bool_result = ParseBoolean(schema_obj); + if (bool_result.IsErr()) return ResultErr(std::move(bool_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(bool_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "null") { + auto null_result = ParseNull(schema_obj); + if (null_result.IsErr()) return ResultErr(std::move(null_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(null_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "array") { + auto array_result = ParseArray(schema_obj); + if (array_result.IsErr()) return ResultErr(std::move(array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(array_result).Unwrap(), cache_key, rule_name_hint); + } else if (type == "object") { + auto obj_result = ParseObject(schema_obj); + if (obj_result.IsErr()) return ResultErr(std::move(obj_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(obj_result).Unwrap(), cache_key, rule_name_hint); + } else { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Unsupported type \"" + type + "\"" + ); + } + } + } else if (schema_obj.count("properties") || schema_obj.count("additionalProperties") || + schema_obj.count("unevaluatedProperties")) { + auto obj_result = ParseObject(schema_obj); + if (obj_result.IsErr()) return ResultErr(std::move(obj_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(obj_result).Unwrap(), cache_key, rule_name_hint); + } else if (schema_obj.count("items") || schema_obj.count("prefixItems") || + schema_obj.count("unevaluatedItems")) { + auto array_result = ParseArray(schema_obj); + if (array_result.IsErr()) return ResultErr(std::move(array_result).UnwrapErr()); + result = SchemaSpec::Make(std::move(array_result).Unwrap(), cache_key, rule_name_hint); + } else { + result = SchemaSpec::Make(AnySpec{}, cache_key, rule_name_hint); + } + + schema_cache_[cache_key] = result; + return ResultOk(result); +} + +Result SchemaParser::ParseInteger(const picojson::object& schema) { + IntegerSpec spec; + + auto checkAndConvertIntegerBound = [](const picojson::value& value + ) -> Result { + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + if (value.is()) return ResultOk(value.get()); + double val = value.get(); + if (val != std::floor(val)) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Integer constraint must be a whole number" + ); + } + static const double PROBLEMATIC_MIN = -9223372036854776000.0; + static const double PROBLEMATIC_MAX = 9223372036854776000.0; + if (val == PROBLEMATIC_MIN) { + XGRAMMAR_CHECK(false + ) << "Integer exceeds minimum limit due to precision loss at 64-bit boundary"; + } + + if (val == PROBLEMATIC_MAX) { + XGRAMMAR_CHECK(false + ) << "Integer exceeds maximum limit due to precision loss at 64-bit boundary"; + } + static const double MAX_INT64_AS_DOUBLE = + static_cast(std::numeric_limits::max()); + static const double MIN_INT64_AS_DOUBLE = + static_cast(std::numeric_limits::min()); + XGRAMMAR_CHECK(val <= MAX_INT64_AS_DOUBLE) << "Integer exceeds maximum limit"; + XGRAMMAR_CHECK(val >= MIN_INT64_AS_DOUBLE) << "Integer exceeds minimum limit"; + return ResultOk(static_cast(val)); + }; + + auto checkAndConvertMultipleOf = [](const picojson::value& value + ) -> Result { + double val; + if (value.is()) { + val = static_cast(value.get()); + } else if (value.is()) { + val = value.get(); + } else { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + if (val <= 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "multipleOf must be greater than 0" + ); + } + if (val != std::floor(val)) { + return ResultErr( + SchemaErrorType::kUnsupportedSchema, "multipleOf for type:integer must be an integer" + ); + } + if (val > static_cast(kIntegerMultipleOfMax)) { + return ResultErr( + SchemaErrorType::kUnsupportedSchema, + "multipleOf for type:integer must be > 0 and <= " + std::to_string(kIntegerMultipleOfMax) + ); + } + return ResultOk(static_cast(val)); + }; + + if (schema.count("multipleOf")) { + auto result = checkAndConvertMultipleOf(schema.at("multipleOf")); + if (result.IsErr()) { + if (result.ErrRef().Type() != SchemaErrorType::kUnsupportedSchema) { + return ResultErr(std::move(result).UnwrapErr()); + } + XGRAMMAR_LOG(WARNING) << result.ErrRef().what() << "; ignoring multipleOf"; + } else { + spec.multiple_of = std::move(result).Unwrap(); + } + } + if (schema.count("minimum")) { + auto result = checkAndConvertIntegerBound(schema.at("minimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.minimum = std::move(result).Unwrap(); + } + if (schema.count("maximum")) { + auto result = checkAndConvertIntegerBound(schema.at("maximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.maximum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMinimum")) { + auto result = checkAndConvertIntegerBound(schema.at("exclusiveMinimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + int64_t val = std::move(result).Unwrap(); + if (val == std::numeric_limits::max()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "exclusiveMinimum would cause integer overflow" + ); + } + spec.exclusive_minimum = val; + } + if (schema.count("exclusiveMaximum")) { + auto result = checkAndConvertIntegerBound(schema.at("exclusiveMaximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + int64_t val = std::move(result).Unwrap(); + if (val == std::numeric_limits::min()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "exclusiveMaximum would cause integer underflow" + ); + } + spec.exclusive_maximum = val; + } + + EffectiveIntegerRange effective_range = ComputeEffectiveIntegerRange(spec); + int64_t effective_min = effective_range.start.value_or(std::numeric_limits::min()); + int64_t effective_max = effective_range.end.value_or(std::numeric_limits::max()); + if (effective_min > effective_max) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Invalid range: minimum greater than maximum" + ); + } + if (spec.multiple_of.has_value()) { + bool has_lower_bound = spec.minimum.has_value() || spec.exclusive_minimum.has_value(); + bool has_upper_bound = spec.maximum.has_value() || spec.exclusive_maximum.has_value(); + if (has_lower_bound || has_upper_bound) { + if (!has_lower_bound || !has_upper_bound || + IsRangeWidthOverCap(effective_min, effective_max, kIntegerMultipleOfRangeWidthMax)) { + XGRAMMAR_LOG(WARNING + ) << "range + multipleOf combination not yet supported; ignoring multipleOf"; + spec.multiple_of.reset(); + return ResultOk(std::move(spec)); + } + if (!HasMultipleInRange(effective_min, effective_max, *spec.multiple_of)) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "range contains no multipleOf value" + ); + } + } + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseNumber(const picojson::object& schema) { + if (schema.count("multipleOf")) { + const auto& value = schema.at("multipleOf"); + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + double multiple_of = + value.is() ? static_cast(value.get()) : value.get(); + if (multiple_of <= 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "multipleOf must be greater than 0" + ); + } + XGRAMMAR_LOG(WARNING) << "multipleOf is not supported for type:number; ignoring multipleOf"; + } + NumberSpec spec; + + auto getDouble = [](const picojson::value& value) -> Result { + if (!value.is() && !value.is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "Value must be a number"); + } + return ResultOk(value.get()); + }; + + if (schema.count("minimum")) { + auto result = getDouble(schema.at("minimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.minimum = std::move(result).Unwrap(); + } + if (schema.count("maximum")) { + auto result = getDouble(schema.at("maximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.maximum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMinimum")) { + auto result = getDouble(schema.at("exclusiveMinimum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.exclusive_minimum = std::move(result).Unwrap(); + } + if (schema.count("exclusiveMaximum")) { + auto result = getDouble(schema.at("exclusiveMaximum")); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + spec.exclusive_maximum = std::move(result).Unwrap(); + } + + // The range is empty if any lower bound conflicts with any upper bound. An + // exclusive bound also rules out equality, so it uses ">=" instead of ">". + auto empty = []() { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "Invalid range: empty range" + ); + }; + + // minimum (x >= min) vs maximum (x <= max). + if (spec.minimum && spec.maximum && *spec.minimum > *spec.maximum) { + return empty(); + } + // minimum (x >= min) vs exclusiveMaximum (x < exclMax). + if (spec.minimum && spec.exclusive_maximum && *spec.minimum >= *spec.exclusive_maximum) { + return empty(); + } + // exclusiveMinimum (x > exclMin) vs maximum (x <= max). + if (spec.exclusive_minimum && spec.maximum && *spec.exclusive_minimum >= *spec.maximum) { + return empty(); + } + // exclusiveMinimum (x > exclMin) vs exclusiveMaximum (x < exclMax). + if (spec.exclusive_minimum && spec.exclusive_maximum && + *spec.exclusive_minimum >= *spec.exclusive_maximum) { + return empty(); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseString(const picojson::object& schema) { + StringSpec spec; + if (schema.count("format")) spec.format = schema.at("format").get(); + if (schema.count("pattern")) spec.pattern = schema.at("pattern").get(); + // Lengths become int32 repetition bounds. A minimum beyond int32 can never be satisfied; a + // maximum beyond it is unbounded in practice. Neither may wrap around when converted. + constexpr int64_t kMaxBound = std::numeric_limits::max(); + if (schema.count("minLength")) { + if (!schema.at("minLength").is() || + schema.at("minLength").get() > kMaxBound) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "minLength must be an integer not exceeding " + std::to_string(kMaxBound) + ); + } + spec.min_length = static_cast(schema.at("minLength").get()); + } + if (schema.count("maxLength")) { + if (!schema.at("maxLength").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxLength must be an integer" + ); + } + if (schema.at("maxLength").get() <= kMaxBound) { + spec.max_length = static_cast(schema.at("maxLength").get()); + } + } + if (spec.max_length != -1 && spec.min_length > spec.max_length) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minLength " + std::to_string(spec.min_length) + " is greater than maxLength " + + std::to_string(spec.max_length) + ); + } + // A pattern or built-in format takes the whole GenerateString branch, so minLength/maxLength + // would be dropped silently (issue #749). Warn here rather than in GenerateString: the XML + // tool-calling converter overrides that method, but every converter goes through ParseString. + if (spec.min_length != 0 || spec.max_length != -1) { + const char* generative = nullptr; + if (spec.pattern.has_value()) { + generative = "pattern"; + } else if (spec.format.has_value() && JSONSchemaConverter::IsBuiltinFormat(*spec.format)) { + generative = "format"; + } + if (generative != nullptr) { + XGRAMMAR_LOG(WARNING) << generative + << " combined with minLength/maxLength is not supported; ignoring " + "minLength/maxLength"; + } + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseBoolean(const picojson::object&) { + return ResultOk(BooleanSpec{}); +} + +Result SchemaParser::ParseNull(const picojson::object&) { + return ResultOk(NullSpec{}); +} + +Result SchemaParser::ParseArray(const picojson::object& schema) { + WarnUnsupportedKeywords(schema, {"uniqueItems", "contains", "minContains", "maxContains"}); + ArraySpec spec; + + if (schema.count("prefixItems")) { + if (!schema.at("prefixItems").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "prefixItems must be an array" + ); + } + for (const auto& item : schema.at("prefixItems").get()) { + if (item.is() && !item.get()) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "prefixItems contains false" + ); + } else if (!item.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "prefixItems must be an array of objects or booleans" + ); + } + auto item_result = Parse(item, "prefix_item"); + if (item_result.IsErr()) return ResultErr(std::move(item_result).UnwrapErr()); + spec.prefix_items.push_back(std::move(item_result).Unwrap()); + } + } + + if (schema.count("items")) { + auto items_value = schema.at("items"); + if (!items_value.is() && !items_value.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "items must be a boolean or an object" + ); + } + if (items_value.is() && !items_value.get()) { + spec.allow_additional_items = false; + } else { + spec.allow_additional_items = true; + auto items_result = Parse(items_value, "item"); + if (items_result.IsErr()) return ResultErr(std::move(items_result).UnwrapErr()); + spec.additional_items = std::move(items_result).Unwrap(); + } + } else if (schema.count("unevaluatedItems")) { + auto unevaluated_items_value = schema.at("unevaluatedItems"); + if (!unevaluated_items_value.is() && !unevaluated_items_value.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "unevaluatedItems must be a boolean or an object" + ); + } + if (unevaluated_items_value.is() && !unevaluated_items_value.get()) { + spec.allow_additional_items = false; + } else { + spec.allow_additional_items = true; + auto items_result = Parse(unevaluated_items_value, "unevaluated_item"); + if (items_result.IsErr()) return ResultErr(std::move(items_result).UnwrapErr()); + spec.additional_items = std::move(items_result).Unwrap(); + } + } else if (!config_.strict_mode) { + spec.allow_additional_items = true; + spec.additional_items = SchemaSpec::Make(AnySpec{}, "", "any"); + } else { + spec.allow_additional_items = false; + } + + if (schema.count("minItems")) { + if (!schema.at("minItems").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "minItems must be an integer"); + } + spec.min_items = std::max(static_cast(0), schema.at("minItems").get()); + } + if (schema.count("minContains")) { + if (!schema.at("minContains").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "minContains must be an integer" + ); + } + spec.min_items = std::max(spec.min_items, schema.at("minContains").get()); + } + if (schema.count("maxItems")) { + if (!schema.at("maxItems").is() || schema.at("maxItems").get() < 0) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxItems must be a non-negative integer" + ); + } + spec.max_items = schema.at("maxItems").get(); + } + // Item counts become int32 repetition bounds, see ParseString for the rationale. + constexpr int64_t kMaxBound = std::numeric_limits::max(); + if (spec.min_items > kMaxBound) { + return ResultErr( + SchemaErrorType::kInvalidSchema, + "minItems and minContains must not exceed " + std::to_string(kMaxBound) + ); + } + if (spec.max_items > kMaxBound) { + spec.max_items = -1; + } + + if (spec.max_items != -1 && spec.min_items > spec.max_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minItems is greater than maxItems: " + std::to_string(spec.min_items) + " > " + + std::to_string(spec.max_items) + ); + } + if (spec.max_items != -1 && spec.max_items < static_cast(spec.prefix_items.size())) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxItems is less than the number of prefixItems: " + std::to_string(spec.max_items) + + " < " + std::to_string(spec.prefix_items.size()) + ); + } + if (!spec.allow_additional_items) { + int64_t prefix_size = static_cast(spec.prefix_items.size()); + if (prefix_size < spec.min_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minItems is greater than the number of prefixItems, but additional items are not " + "allowed: " + + std::to_string(spec.min_items) + " > " + std::to_string(prefix_size) + ); + } + if (spec.max_items != -1 && prefix_size > spec.max_items) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxItems is less than the number of prefixItems, but additional items are not " + "allowed: " + + std::to_string(spec.max_items) + " < " + std::to_string(prefix_size) + ); + } + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseObject(const picojson::object& schema) { + ObjectSpec spec; + + if (schema.count("properties")) { + if (!schema.at("properties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "properties must be an object" + ); + } + auto properties_obj = schema.at("properties").get(); + for (const auto& key : properties_obj.ordered_keys()) { + auto prop_result = Parse(properties_obj.at(key), key); + if (prop_result.IsErr()) return ResultErr(std::move(prop_result).UnwrapErr()); + spec.properties.push_back({key, std::move(prop_result).Unwrap()}); + } + } + + if (schema.count("required")) { + if (!schema.at("required").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "required must be an array"); + } + for (const auto& req : schema.at("required").get()) { + spec.required.insert(req.get()); + } + } + + if (schema.count("patternProperties")) { + if (!schema.at("patternProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "patternProperties must be an object" + ); + } + auto pattern_props = schema.at("patternProperties").get(); + for (const auto& key : pattern_props.ordered_keys()) { + auto prop_result = Parse(pattern_props.at(key), "pattern_prop"); + if (prop_result.IsErr()) return ResultErr(std::move(prop_result).UnwrapErr()); + spec.pattern_properties.push_back({key, std::move(prop_result).Unwrap()}); + } + } + + if (schema.count("propertyNames")) { + if (!schema.at("propertyNames").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "propertyNames must be an object" + ); + } + auto property_names_obj = schema.at("propertyNames").get(); + if (property_names_obj.count("type") && property_names_obj.at("type").is() && + property_names_obj.at("type").get() != "string") { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "propertyNames must be an object that validates string" + ); + } + auto prop_names_result = Parse(schema.at("propertyNames"), "property_name", "string"); + if (prop_names_result.IsErr()) return ResultErr(std::move(prop_names_result).UnwrapErr()); + spec.property_names = std::move(prop_names_result).Unwrap(); + } + + spec.allow_additional_properties = !config_.strict_mode; + if (schema.count("additionalProperties")) { + auto add_props = schema.at("additionalProperties"); + if (add_props.is()) { + spec.allow_additional_properties = add_props.get(); + } else { + spec.allow_additional_properties = true; + auto add_props_result = Parse(add_props, "additional"); + if (add_props_result.IsErr()) return ResultErr(std::move(add_props_result).UnwrapErr()); + spec.additional_properties_schema = std::move(add_props_result).Unwrap(); + } + } + + spec.allow_unevaluated_properties = true; + if (schema.count("additionalProperties")) { + spec.allow_unevaluated_properties = spec.allow_additional_properties; + } else if (schema.count("unevaluatedProperties")) { + auto uneval_props = schema.at("unevaluatedProperties"); + if (uneval_props.is()) { + spec.allow_unevaluated_properties = uneval_props.get(); + } else { + spec.allow_unevaluated_properties = true; + auto uneval_result = Parse(uneval_props, "unevaluated"); + if (uneval_result.IsErr()) return ResultErr(std::move(uneval_result).UnwrapErr()); + spec.unevaluated_properties_schema = std::move(uneval_result).Unwrap(); + } + } else if (config_.strict_mode) { + spec.allow_unevaluated_properties = false; + } + + if (schema.count("minProperties")) { + if (!schema.at("minProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "minProperties must be an integer" + ); + } + spec.min_properties = static_cast(schema.at("minProperties").get()); + if (spec.min_properties < 0) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "minProperties must be a non-negative integer" + ); + } + } + if (schema.count("maxProperties")) { + if (!schema.at("maxProperties").is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "maxProperties must be an integer" + ); + } + spec.max_properties = static_cast(schema.at("maxProperties").get()); + if (spec.max_properties < 0) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, "maxProperties must be a non-negative integer" + ); + } + } + + if (spec.max_properties != -1 && spec.min_properties > spec.max_properties) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minProperties is greater than maxProperties: " + std::to_string(spec.min_properties) + + " > " + std::to_string(spec.max_properties) + ); + } + if (spec.max_properties != -1 && static_cast(spec.required.size()) > spec.max_properties) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "maxProperties is less than the number of required properties: " + + std::to_string(spec.max_properties) + " < " + std::to_string(spec.required.size()) + ); + } + if (spec.pattern_properties.empty() && !spec.property_names && + !spec.allow_additional_properties && !spec.allow_unevaluated_properties && + spec.min_properties > static_cast(spec.properties.size())) { + return ResultErr( + SchemaErrorType::kUnsatisfiableSchema, + "minProperties is greater than the number of properties, but additional properties aren't " + "allowed: " + + std::to_string(spec.min_properties) + " > " + std::to_string(spec.properties.size()) + ); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseConst(const picojson::object& schema) { + ConstSpec spec; + spec.json_value = schema.at("const").serialize(); + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseEnum(const picojson::object& schema) { + EnumSpec spec; + if (!schema.at("enum").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "enum must be an array"); + } + const auto& enum_array = schema.at("enum").get(); + if (enum_array.empty()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "enum array must not be empty"); + } + for (const auto& value : enum_array) { + spec.json_values.push_back(value.serialize()); + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseRef(const picojson::object& schema) { + if (!schema.at("$ref").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "$ref must be a string"); + } + RefSpec spec; + spec.uri = schema.at("$ref").get(); + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ResolveRef( + const std::string& uri, const std::string& rule_name_hint +) { + if (ref_cache_.count(uri)) return ResultOk(ref_cache_[uri]); + + if (uri == "#") { + auto placeholder = SchemaSpec::Make(AnySpec{}, "", "root"); + ref_cache_[uri] = placeholder; + auto result = Parse(root_schema_, "root"); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + auto resolved = std::move(result).Unwrap(); + ref_cache_[uri] = resolved; + return ResultOk(resolved); + } + + if (uri.size() < 2 || uri[0] != '#' || uri[1] != '/') { + XGRAMMAR_LOG(WARNING) << "URI should either be '#' or start with '#/' but got " << uri; + return ResultOk(SchemaSpec::Make(AnySpec{}, "", "any")); + } + + std::vector parts; + std::stringstream ss(uri.substr(2)); + std::string part; + std::string new_rule_name_prefix; + while (std::getline(ss, part, '/')) { + if (!part.empty()) parts.push_back(part); + if (!new_rule_name_prefix.empty()) new_rule_name_prefix += "_"; + for (const auto& c : part) { + if (std::isalpha(c) || c == '_' || c == '-' || c == '.') new_rule_name_prefix += c; + } + } + + auto current = std::cref(root_schema_); + for (const auto& p : parts) { + if (!current.get().is() || !current.get().contains(p)) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "Cannot find field " + p + " in " + uri + ); + } + current = current.get().get(p); + } + + auto result = Parse(current, new_rule_name_prefix); + if (result.IsErr()) return ResultErr(std::move(result).UnwrapErr()); + auto resolved = std::move(result).Unwrap(); + ref_cache_[uri] = resolved; + return ResultOk(resolved); +} + +Result SchemaParser::ParseAnyOf( + const picojson::object& schema, const std::string& keyword +) { + AnyOfSpec spec; + if (!schema.at(keyword).is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, keyword + " must be an array"); + } + int idx = 0; + for (const auto& option : schema.at(keyword).get()) { + auto option_result = Parse(option, "case_" + std::to_string(idx)); + if (option_result.IsErr()) return ResultErr(std::move(option_result).UnwrapErr()); + spec.options.push_back(std::move(option_result).Unwrap()); + ++idx; + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseOneOf(const picojson::object& schema) { + OneOfSpec spec; + if (!schema.at("oneOf").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "oneOf must be an array"); + } + + const auto& options = schema.at("oneOf").get(); + if (options.empty()) { + return ResultErr(SchemaErrorType::kUnsupportedSchema, kUnsupportedOneOfMessage); + } + + int idx = 0; + for (const auto& option : options) { + auto option_result = Parse(option, "case_" + std::to_string(idx)); + if (option_result.IsErr()) return ResultErr(std::move(option_result).UnwrapErr()); + spec.options.push_back(std::move(option_result).Unwrap()); + ++idx; + } + + if (!TryProvePairwiseDisjointOneOf(options)) { + return ResultErr(SchemaErrorType::kUnsupportedSchema, kUnsupportedOneOfMessage); + } + + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseAllOf(const picojson::object& schema) { + AllOfSpec spec; + if (!schema.at("allOf").is()) { + return ResultErr(SchemaErrorType::kInvalidSchema, "allOf must be an array"); + } + int idx = 0; + for (const auto& sub_schema : schema.at("allOf").get()) { + auto sub_result = Parse(sub_schema, "all_" + std::to_string(idx)); + if (sub_result.IsErr()) return ResultErr(std::move(sub_result).UnwrapErr()); + spec.schemas.push_back(std::move(sub_result).Unwrap()); + ++idx; + } + return ResultOk(std::move(spec)); +} + +Result SchemaParser::ParseTypeArray( + const picojson::object& schema, const std::string& rule_name_hint +) { + TypeArraySpec spec; + auto type_array = schema.at("type").get(); + picojson::object schema_copy = schema; + if (type_array.empty()) { + schema_copy.erase("type"); + auto any_result = Parse(picojson::value(schema_copy), rule_name_hint); + if (any_result.IsErr()) return ResultErr(std::move(any_result).UnwrapErr()); + spec.type_schemas.push_back(std::move(any_result).Unwrap()); + return ResultOk(std::move(spec)); + } + for (const auto& type : type_array) { + if (!type.is()) { + return ResultErr( + SchemaErrorType::kInvalidSchema, "type must be a string or an array of strings" + ); + } + schema_copy["type"] = type; + auto type_result = + Parse(picojson::value(schema_copy), rule_name_hint + "_" + type.get()); + if (type_result.IsErr()) return ResultErr(std::move(type_result).UnwrapErr()); + spec.type_schemas.push_back(std::move(type_result).Unwrap()); + } + return ResultOk(std::move(spec)); +} + +} // namespace + +// ==================== IndentManager Implementation ==================== + +IndentManager::IndentManager( + std::optional indent, + const std::string& separator, + bool any_whitespace, + std::optional max_whitespace_cnt +) + : any_whitespace_(any_whitespace), + enable_newline_(indent.has_value()), + indent_(indent.value_or(0)), + separator_(separator), + total_indent_(0), + is_first_({true}), + max_whitespace_cnt_(max_whitespace_cnt) { + if (max_whitespace_cnt.has_value() && max_whitespace_cnt.value() <= 0) { + XGRAMMAR_LOG(FATAL) << "max_whitespace_cnt must be positive."; + } +} + +void IndentManager::StartIndent() { + total_indent_ += indent_; + is_first_.push_back(true); +} + +void IndentManager::EndIndent() { + total_indent_ -= indent_; + is_first_.pop_back(); +} + +std::string IndentManager::StartSeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + if (!enable_newline_) { + return "\"\""; + } + return "\"\\n" + std::string(total_indent_, ' ') + "\""; +} + +std::string IndentManager::MiddleSeparator() { + if (any_whitespace_) { + std::string whitespace_part; + if (!max_whitespace_cnt_.has_value()) { + whitespace_part = "[ \\n\\r\\t]*"; + } else { + whitespace_part = "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + return whitespace_part + " \"" + separator_ + "\" " + whitespace_part; + } + if (!enable_newline_) { + return "\"" + separator_ + "\""; + } + return "\"" + separator_ + "\\n" + std::string(total_indent_, ' ') + "\""; +} + +std::string IndentManager::EndSeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + if (!enable_newline_) { + return "\"\""; + } + return "\"\\n" + std::string(total_indent_ - indent_, ' ') + "\""; +} + +std::string IndentManager::EmptySeparator() { + if (any_whitespace_) { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } + return "\"\""; +} + +std::string IndentManager::NextSeparator(bool is_end) { + if (any_whitespace_) { + if (is_first_.back() || is_end) { + is_first_.back() = false; + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } else { + return "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + } else { + std::string whitespace_part; + if (!max_whitespace_cnt_.has_value()) { + whitespace_part = "[ \\n\\r\\t]*"; + } else { + whitespace_part = "[ \\n\\r\\t]{0," + std::to_string(max_whitespace_cnt_.value()) + "}"; + } + return whitespace_part + " \"" + separator_ + "\" " + whitespace_part; + } + } + + std::string res = ""; + if (!is_first_.back() && !is_end) { + res += separator_; + } + is_first_.back() = false; + + if (enable_newline_) { + res += "\\n"; + } + + if (!is_end) { + res += std::string(total_indent_, ' '); + } else { + res += std::string(total_indent_ - indent_, ' '); + } + + return "\"" + res + "\""; +} + +// ==================== Static Constants ==================== + +const std::string JSONSchemaConverter::kBasicAny = "basic_any"; +const std::string JSONSchemaConverter::kBasicInteger = "basic_integer"; +const std::string JSONSchemaConverter::kBasicNumber = "basic_number"; +const std::string JSONSchemaConverter::kBasicString = "basic_string"; +const std::string JSONSchemaConverter::kBasicBoolean = "basic_boolean"; +const std::string JSONSchemaConverter::kBasicNull = "basic_null"; +const std::string JSONSchemaConverter::kBasicArray = "basic_array"; +const std::string JSONSchemaConverter::kBasicObject = "basic_object"; +const std::string JSONSchemaConverter::kBasicEscape = "basic_escape"; +const std::string JSONSchemaConverter::kBasicStringSub = "basic_string_sub"; + +// ==================== JSONSchemaConverter Implementation ==================== + +JSONSchemaConverter::JSONSchemaConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + bool any_order, + std::vector excludes +) + : indent_manager_( + indent, + separators.has_value() ? separators->first + : (any_whitespace ? "," : (indent.has_value() ? "," : ", ")), + any_whitespace, + max_whitespace_cnt + ), + any_whitespace_(any_whitespace), + max_whitespace_cnt_(max_whitespace_cnt), + any_order_(any_order), + excludes_(std::move(excludes)), + ref_resolver_(std::move(ref_resolver)) { + std::string colon_sep = + separators.has_value() ? separators->second : (any_whitespace ? ":" : ": "); + std::string whitespace = GetWhitespacePattern(); + colon_expr_id_ = FormattingExpression( + any_whitespace ? whitespace + " \"" + colon_sep + "\" " + whitespace : "\"" + colon_sep + "\"" + ); +} + +Grammar JSONSchemaConverter::Convert(const SchemaSpecPtr& spec) { + AddBasicRules(); + + // Register the root rule for circular reference handling + // This allows $ref: "#" to resolve to "root" + int32_t root_rule_id = builder_.AddEmptyRuleWithHint("root"); + std::string root_rule_name = builder_.GetRule(root_rule_id).name; + uri_to_rule_id_[RefCacheKey("#")] = root_rule_id; + + // Check if the spec can be directly mapped to an existing rule + auto cached_rule = GetCache(spec->cache_key); + if (cached_rule.has_value()) { + // Root schema matches a basic type, just reference it + builder_.UpdateRuleBody(root_rule_id, RuleRef(*cached_rule)); + } else { + // Generate the rule body + if (!spec->cache_key.empty()) { + AddCache(spec->cache_key, root_rule_id); + } + builder_.UpdateRuleBody(root_rule_id, GenerateFromSpec(spec, root_rule_name)); + } + return builder_.Get(root_rule_id); +} + +void JSONSchemaConverter::AddBasicRules() { AddBasicRules({}); } + +void JSONSchemaConverter::AddBasicRules(const std::vector& additional_rule_names) { + std::vector basic_rule_names = { + kBasicEscape, + kBasicStringSub, + kBasicAny, + kBasicInteger, + kBasicNumber, + kBasicString, + kBasicBoolean, + kBasicNull, + kBasicArray, + kBasicObject, + }; + basic_rule_names.insert( + basic_rule_names.end(), additional_rule_names.begin(), additional_rule_names.end() + ); + for (const auto& name : basic_rule_names) { + builder_.AddEmptyRule(name); + } + AddHelperRules(); + + // Create basic rules with a temporary indent manager for compact format + auto saved_indent_manager = indent_manager_; + indent_manager_ = IndentManager( + std::nullopt, + any_whitespace_ ? "," : ", ", + any_whitespace_, + any_whitespace_ ? max_whitespace_cnt_ : std::nullopt + ); + + // basic_any - use "{}" as the cache key for empty schema + auto any_spec = SchemaSpec::Make(AnySpec{}, "{}", kBasicAny); + builder_.UpdateRuleBody(kBasicAny, GenerateAny(std::get(any_spec->spec), kBasicAny)); + AddCache("{}", builder_.GetRuleId(kBasicAny)); + + // basic_integer - cache_key matches SchemaParser::ComputeCacheKey for {"type": "integer"} + constexpr const char* kIntegerCacheKey = "{\"type\":\"integer\"}"; + builder_.UpdateRuleBody(kBasicInteger, GenerateInteger(IntegerSpec{}, kBasicInteger)); + AddCache(kIntegerCacheKey, builder_.GetRuleId(kBasicInteger)); + + // basic_number - cache_key matches SchemaParser::ComputeCacheKey for {"type": "number"} + constexpr const char* kNumberCacheKey = "{\"type\":\"number\"}"; + builder_.UpdateRuleBody(kBasicNumber, GenerateNumber(NumberSpec{}, kBasicNumber)); + AddCache(kNumberCacheKey, builder_.GetRuleId(kBasicNumber)); + + constexpr const char* kStringCacheKey = "{\"type\":\"string\"}"; + builder_.UpdateRuleBody(kBasicString, Sequence({ByteString("\""), RuleRef(kBasicStringSub)})); + AddCache(kStringCacheKey, builder_.GetRuleId(kBasicString)); + + // basic_boolean - cache_key matches SchemaParser::ComputeCacheKey for {"type": "boolean"} + constexpr const char* kBooleanCacheKey = "{\"type\":\"boolean\"}"; + builder_.UpdateRuleBody(kBasicBoolean, GenerateBoolean(BooleanSpec{}, kBasicBoolean)); + AddCache(kBooleanCacheKey, builder_.GetRuleId(kBasicBoolean)); + + // basic_null - cache_key matches SchemaParser::ComputeCacheKey for {"type": "null"} + constexpr const char* kNullCacheKey = "{\"type\":\"null\"}"; + builder_.UpdateRuleBody(kBasicNull, GenerateNull(NullSpec{}, kBasicNull)); + AddCache(kNullCacheKey, builder_.GetRuleId(kBasicNull)); + + // basic_array - cache_key matches SchemaParser::ComputeCacheKey for {"type": "array"} + constexpr const char* kArrayCacheKey = "{\"type\":\"array\"}"; + ArraySpec array_spec_val; + array_spec_val.allow_additional_items = true; + array_spec_val.additional_items = any_spec; + builder_.UpdateRuleBody(kBasicArray, GenerateArray(array_spec_val, kBasicArray)); + AddCache(kArrayCacheKey, builder_.GetRuleId(kBasicArray)); + + // basic_object - cache_key matches SchemaParser::ComputeCacheKey for {"type": "object"} + constexpr const char* kObjectCacheKey = "{\"type\":\"object\"}"; + ObjectSpec obj_spec_val; + obj_spec_val.allow_additional_properties = true; + obj_spec_val.additional_properties_schema = any_spec; + builder_.UpdateRuleBody(kBasicObject, GenerateObject(obj_spec_val, kBasicObject)); + AddCache(kObjectCacheKey, builder_.GetRuleId(kBasicObject)); + + indent_manager_ = saved_indent_manager; +} + +// The content of a JSON string, escapes included, as RegexExpression parses it. +static constexpr const char kJSONStringBodyRegex[] = + R"(([^"\\\x00-\x1f]|\\(["\\/bfnrt]|u[0-9a-fA-F]{4}))*)"; + +void JSONSchemaConverter::AddHelperRules() { + if (max_whitespace_cnt_.has_value()) { + // Preserve historical helper-rule numbering after grammar optimization. The text parser + // allocated one initial bounded-repetition helper that dead-code elimination later removed. + builder_.AddRuleWithHint(kBasicStringSub, Empty()); + } + int32_t escaped_character = builder_.AddCharacterClass( + {{'"', '"'}, + {'\\', '\\'}, + {'/', '/'}, + {'b', 'b'}, + {'f', 'f'}, + {'n', 'n'}, + {'r', 'r'}, + {'t', 't'}} + ); + int32_t hexadecimal_character = builder_.AddCharacterClass({{'A', 'F'}, {'a', 'f'}, {'0', '9'}}); + int32_t unicode_escape = Sequence( + {ByteString("u"), + hexadecimal_character, + hexadecimal_character, + hexadecimal_character, + hexadecimal_character} + ); + builder_.UpdateRuleBody(kBasicEscape, Choice({escaped_character, unicode_escape})); + + int32_t normal_character = builder_.AddCharacterClass( + {{0, 0x1f}, {'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'}}, true + ); + int32_t string_sub_ref = RuleRef(kBasicStringSub); + int32_t string_sub_body = Choice( + {ByteString("\""), + Sequence({normal_character, string_sub_ref}), + Sequence({ByteString("\\"), RuleRef(kBasicEscape), string_sub_ref})} + ); + builder_.UpdateRuleBody(kBasicStringSub, string_sub_body); + if (!excludes_.empty()) { + builder_.UpdateRuleBody( + kBasicStringSub, ExcludingString(kJSONStringBodyRegex, kBasicStringSub, true) + ); + } + int32_t closing_context = + builder_.AddCharacterClass({{',', ','}, {'}', '}'}, {']', ']'}, {':', ':'}}); + builder_.UpdateLookaheadAssertion( + kBasicStringSub, Sequence({WhitespaceExpression(), closing_context}) + ); +} + +// Keep converter-specific node reuse local; GrammarBuilder creates all AST nodes. +int32_t JSONSchemaConverter::Empty() { + if (!empty_expr_id_.has_value()) { + empty_expr_id_ = builder_.AddEmptyStr(); + } + return *empty_expr_id_; +} + +int32_t JSONSchemaConverter::Unsatisfiable() { + if (!unsatisfiable_expr_id_.has_value()) { + unsatisfiable_expr_id_ = builder_.AddCharacterClass({{0, 0x10ffff}}, true); + } + return *unsatisfiable_expr_id_; +} + +int32_t JSONSchemaConverter::ByteString(const std::string& value) { + auto it = byte_string_expr_ids_.find(value); + if (it != byte_string_expr_ids_.end()) { + return it->second; + } + int32_t expr_id = value.empty() ? Empty() : builder_.AddByteString(value); + byte_string_expr_ids_[value] = expr_id; + return expr_id; +} + +int32_t JSONSchemaConverter::TagDispatch( + bool loop_after_dispatch, std::vector excludes +) { + return builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{{}, loop_after_dispatch, std::move(excludes)} + ); +} + +int32_t JSONSchemaConverter::RuleRef(int32_t rule_id) { + auto it = rule_ref_expr_ids_.find(rule_id); + if (it != rule_ref_expr_ids_.end()) { + return it->second; + } + int32_t expr_id = builder_.AddRuleRef(rule_id); + rule_ref_expr_ids_[rule_id] = expr_id; + return expr_id; +} + +int32_t JSONSchemaConverter::RuleRef(const std::string& rule_name) { + int32_t rule_id = builder_.GetRuleId(rule_name); + XGRAMMAR_CHECK(rule_id != -1) << "Rule " << rule_name << " is not allocated"; + return RuleRef(rule_id); +} + +int32_t JSONSchemaConverter::Sequence(const std::vector& elements) { + if (elements.empty()) { + return Empty(); + } + if (elements.size() == 1) { + return elements[0]; + } + return builder_.AddSequence(elements); +} + +int32_t JSONSchemaConverter::Choice(const std::vector& choices) { + if (choices.empty()) { + return Empty(); + } + if (choices.size() == 1) { + return choices[0]; + } + return builder_.AddChoices(choices); +} + +int32_t JSONSchemaConverter::Repeat( + const std::string& rule_name_hint, int32_t expr_id, int32_t min_count, int32_t max_count +) { + if (min_count == 0 && max_count == 0) { + return Empty(); + } + if (min_count == 1 && max_count == 1) { + return expr_id; + } + if (min_count == 0 && max_count == 1) { + return Choice({Empty(), expr_id}); + } + if (min_count == 0 && max_count == -1) { + auto expr = builder_.GetGrammarExpr(expr_id); + if (expr.type == GrammarBuilder::GrammarExprType::kCharacterClass) { + std::vector data(expr.begin(), expr.end()); + return builder_.AddGrammarExpr( + {GrammarBuilder::GrammarExprType::kCharacterClassStar, + data.data(), + static_cast(data.size())} + ); + } + } + return builder_.AddRepeatFromExpr(rule_name_hint, expr_id, min_count, max_count); +} + +int32_t JSONSchemaConverter::AddSubGrammar(const Grammar& grammar) { + int32_t rule_id = SubGrammarAdder::Apply(&builder_, grammar); + return RuleRef(rule_id); +} + +std::string JSONSchemaConverter::GetWhitespacePattern() const { + if (!max_whitespace_cnt_.has_value()) { + return "[ \\n\\r\\t]*"; + } + return "[ \\n\\r\\t]{0," + std::to_string(*max_whitespace_cnt_) + "}"; +} + +int32_t JSONSchemaConverter::WhitespaceExpression() { + std::vector elements = { + {' ', ' '}, {'\n', '\n'}, {'\r', '\r'}, {'\t', '\t'} + }; + if (!max_whitespace_cnt_.has_value()) { + if (!whitespace_expr_id_.has_value()) { + whitespace_expr_id_ = builder_.AddCharacterClassStar(elements); + } + return *whitespace_expr_id_; + } + // Bounded whitespace occurrences intentionally remain distinct, matching the historical + // parser-produced rule shape after normalization. + return Repeat( + "whitespace", + builder_.AddCharacterClass(elements), + 0, + static_cast(*max_whitespace_cnt_) + ); +} + +int32_t JSONSchemaConverter::FormattingExpression(const std::string& expression) { + const std::string whitespace = GetWhitespacePattern(); + if (expression == whitespace) { + return WhitespaceExpression(); + } + + const std::string prefix = whitespace + " "; + const std::string suffix = " " + whitespace; + if (expression.size() >= prefix.size() + suffix.size() && + expression.compare(0, prefix.size(), prefix) == 0 && + expression.compare(expression.size() - suffix.size(), suffix.size(), suffix) == 0) { + return Sequence( + {WhitespaceExpression(), + FormattingExpression( + expression.substr(prefix.size(), expression.size() - prefix.size() - suffix.size()) + ), + WhitespaceExpression()} + ); + } + + picojson::value value; + std::string error = ParseJSON(value, expression); + XGRAMMAR_CHECK(error.empty() && value.is()) + << "Unsupported indentation expression: " << expression; + return ByteString(value.get()); +} + +std::string JSONSchemaConverter::NextSeparator(bool is_end) { + return indent_manager_.NextSeparator(is_end); +} + +int32_t JSONSchemaConverter::NextSeparatorExpression(bool is_end) { + return FormattingExpression(NextSeparator(is_end)); +} + +std::string JSONSchemaConverter::GetKeyPattern() const { return kBasicString; } + +int32_t JSONSchemaConverter::KeyPatternExpression() { return RuleRef(GetKeyPattern()); } + +int32_t JSONSchemaConverter::BuildTrieBody(const TrieNode& node, const std::string& rule_name) { + std::vector choices; + if (!node.is_terminal) { + choices.push_back(ByteString("\"")); + } + + std::vector excluded = { + {0, 0x1f}, {'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'} + }; + for (const auto& [character, child] : node.children) { + static_cast(child); + excluded.push_back({character, character}); + } + choices.push_back(Sequence({builder_.AddCharacterClass(excluded, true), RuleRef(kBasicStringSub)}) + ); + choices.push_back(Sequence({ByteString("\\"), RuleRef(kBasicEscape), RuleRef(kBasicStringSub)})); + for (const auto& [character, child] : node.children) { + choices.push_back(Sequence( + {ByteString(std::string(1, static_cast(character))), BuildTrieBody(child, rule_name)} + )); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (properties.empty()) { + return KeyPatternExpression(); + } + if (!excludes_.empty()) { + std::vector keys; + for (const auto& property : properties) { + auto encoded = picojson::value(property.name).serialize(false); + keys.push_back(encoded.substr(1, encoded.size() - 2)); + } + return Sequence( + {ByteString("\""), + ExcludingString(kJSONStringBodyRegex, rule_name + "_addl_key", true, keys)} + ); + } + + // Build trie from property names + // TODO(linzhang): The trie only excludes the literal unescaped spelling of each property name. + TrieNode root; + for (const auto& prop : properties) { + TrieNode* cur = &root; + for (unsigned char c : prop.name) { + cur = &cur->children[c]; + } + cur->is_terminal = true; + } + + int32_t key_rule_id = builder_.AddEmptyRuleWithHint(rule_name + "_addl_key"); + std::string key_rule_name = builder_.GetRule(key_rule_id).name; + builder_.UpdateRuleBody( + key_rule_id, Sequence({ByteString("\""), BuildTrieBody(root, key_rule_name)}) + ); + builder_.UpdateLookaheadAssertion( + key_rule_id, + Sequence( + {WhitespaceExpression(), + builder_.AddCharacterClass({{',', ','}, {'}', '}'}, {']', ']'}, {':', ':'}})} + ) + ); + return RuleRef(key_rule_id); +} + +std::string JSONSchemaConverter::GetBasicAnyRuleName() const { return kBasicAny; } + +void JSONSchemaConverter::AddCache(const std::string& key, int32_t rule_id) { + if (!key.empty()) { + rule_cache_manager_.AddCache(key, true, rule_id); + } +} + +std::optional JSONSchemaConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + return rule_cache_manager_.GetCache(key, true); +} + +int32_t JSONSchemaConverter::CreateRule( + const SchemaSpecPtr& spec, const std::string& rule_name_hint +) { + // Only check cache for basic rules (pre-populated in AddBasicRules) + // Don't cache other rules to match original behavior + auto cached = GetCache(spec->cache_key); + if (cached.has_value()) { + return cached.value(); + } + int32_t rule_id = builder_.AddEmptyRuleWithHint(rule_name_hint); + // Copy the name before generating: GenerateFromSpec may add rules and reallocate the + // builder's rule storage, invalidating references into it. + std::string rule_name = builder_.GetRule(rule_id).name; + builder_.UpdateRuleBody(rule_id, GenerateFromSpec(spec, rule_name)); + return rule_id; +} + +int32_t JSONSchemaConverter::GenerateFromSpec( + const SchemaSpecPtr& spec, const std::string& rule_name_hint +) { + return std::visit( + [this, &rule_name_hint](const auto& s) -> int32_t { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return GenerateInteger(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateNumber(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateString(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateBoolean(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateNull(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateArray(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateObject(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAny(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateConst(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateEnum(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateRef(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAnyOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateOneOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateAllOf(s, rule_name_hint); + } else if constexpr (std::is_same_v) { + return GenerateTypeArray(s, rule_name_hint); + } else { + XGRAMMAR_LOG(FATAL) << "Unknown spec type"; + } + }, + spec->spec + ); +} + +/*! + * \brief Emit the grammar expression matching a regex. Prefer the Regex node with + * json_string=true so the pattern is compiled into a single automaton by GrammarFSMBuilder; + * json_string=true excludes the characters that must be escaped in a JSON string ('"', '\\' + * and the control characters) from every character match, so classes like \S cannot emit an + * unescaped quote. Fall back to the CFG expansion when the FSM regex engine does not support + * the pattern, or when the exclusion makes the pattern unmatchable (e.g. a pattern requiring + * a literal '"'). + */ +int32_t JSONSchemaConverter::RegexExpression( + const std::string& regex, bool json_string, bool force_cfg_expansion +) { + bool can_use_fsm = !force_cfg_expansion; + if (json_string) { + can_use_fsm = + can_use_fsm && std::all_of(regex.begin(), regex.end(), [](unsigned char character) { + return character >= 0x20 && character <= 0x7e; + }); + } + if (can_use_fsm) { + auto fsm_result = GrammarFSMBuilder::Regex(regex, json_string); + if (fsm_result.IsOk()) { + auto fsm = std::move(fsm_result).Unwrap(); + std::unordered_set reachable_states; + fsm.GetReachableStates(&reachable_states); + bool language_is_empty = + std::none_of(reachable_states.begin(), reachable_states.end(), [&](int state) { + return fsm.IsEndState(state); + }); + if (!language_is_empty) { + return builder_.AddRegex(regex, json_string); + } + } + } + + // Keep regex conversion independent. Only the uncommon fallback path converts its existing + // EBNF result to a subgrammar; the JSON Schema rule graph itself is still built directly. + return AddSubGrammar(Grammar::FromEBNF(RegexToEBNF(regex))); +} + +// ==================== Generate Methods ==================== + +void JSONSchemaConverter::WarnDroppedLengthConstraints( + const StringSpec& spec, const std::string& rule_name +) const { + XGRAMMAR_LOG(WARNING) << "Ignoring the length constraints of string " << rule_name + << " (minLength=" << spec.min_length << ", maxLength=" << spec.max_length + << "): they are not applied together with JSONSchemaFormat.excludes"; +} + +bool JSONSchemaConverter::IsAllowedString(const std::string& text) const { + return std::none_of(excludes_.begin(), excludes_.end(), [&](const auto& excluded) { + return text.find(excluded) != std::string::npos; + }); +} + +bool JSONSchemaConverter::IsAllowedLiteral(const picojson::value& value, bool raw_string) const { + if (excludes_.empty()) return true; + if (value.is()) { + if (raw_string) return IsAllowedString(value.get()); + auto encoded = value.serialize(false); + return IsAllowedString(encoded.substr(1, encoded.size() - 2)); + } + if (value.is()) { + for (const auto& item : value.get()) { + if (!IsAllowedLiteral(item)) return false; + } + } + if (value.is()) { + for (const auto& [key, item] : value.get()) { + if (!IsAllowedLiteral(picojson::value(key)) || !IsAllowedLiteral(item)) return false; + } + } + return true; +} + +bool JSONSchemaConverter::IsAllowedJSONLiteral(const std::string& json_value, bool raw_string) + const { + if (excludes_.empty()) return true; + picojson::value value; + XGRAMMAR_CHECK(picojson::parse(value, json_value).empty()); + return IsAllowedLiteral(value, raw_string); +} + +int32_t JSONSchemaConverter::ExcludingString( + const std::string& regex, + const std::string& rule_name, + bool close_json_string, + const std::vector& excluded_keys +) { + auto parsed = GrammarFSMBuilder::Regex(regex, false); + XGRAMMAR_CHECK(parsed.IsOk()) << "Cannot build the FSM of " << regex << ": " + << std::move(parsed).UnwrapErr().what(); + auto exclusion = GrammarFSMBuilder::TagDispatch({{}, false, excludes_}); + XGRAMMAR_CHECK(exclusion.has_value()) << "Invalid JSONSchemaFormat.excludes"; + auto intersected = FSMWithStartEnd::Intersect(std::move(parsed).Unwrap(), *exclusion); + XGRAMMAR_CHECK(intersected.IsOk()) << "Cannot intersect JSONSchemaFormat.excludes: " + << std::move(intersected).UnwrapErr().what(); + auto fsm = std::move(intersected).Unwrap(); + if (!excluded_keys.empty()) { + auto keys = TrieFSMBuilder::Build(excluded_keys, {}); + XGRAMMAR_CHECK(keys.has_value()); + auto other_keys = keys->Not(); + XGRAMMAR_CHECK(other_keys.IsOk()); + auto filtered = FSMWithStartEnd::Intersect(fsm, std::move(other_keys).Unwrap()); + XGRAMMAR_CHECK(filtered.IsOk()); + fsm = std::move(filtered).Unwrap(); + } + // Do not emit paths that can never complete after exclusions. Otherwise a matcher + // could accept a forbidden alternative's prefix and reach an all-rejected mask later. + std::vector> predecessors(fsm.NumStates()); + for (int state = 0; state < fsm.NumStates(); ++state) { + for (const auto& edge : fsm.GetFsm().GetEdges(state)) { + predecessors[edge.target].push_back(state); + } + } + std::vector productive(fsm.NumStates(), false); + std::vector worklist(fsm.GetEnds().begin(), fsm.GetEnds().end()); + for (int state : worklist) productive[state] = true; + for (size_t index = 0; index < worklist.size(); ++index) { + for (int state : predecessors[worklist[index]]) { + if (!productive[state]) { + productive[state] = true; + worklist.push_back(state); + } + } + } + if (!productive[fsm.GetStart()]) { + return Unsatisfiable(); + } + + // Collapse UTF-8 paths into codepoint transitions before emitting character classes. + // Emitting individual continuation bytes as string literals would not round-trip through + // EBNF, whose escaped string literals represent Unicode codepoints rather than raw bytes. + struct CodepointRange { + int min, max, target; + }; + using Ranges = std::vector; + auto merge_ranges = [](Ranges ranges) { + std::sort(ranges.begin(), ranges.end(), [](const auto& a, const auto& b) { + return std::tie(a.target, a.min, a.max) < std::tie(b.target, b.min, b.max); + }); + Ranges merged; + for (const auto& range : ranges) { + if (!merged.empty() && merged.back().target == range.target && + range.min <= merged.back().max + 1) { + merged.back().max = std::max(merged.back().max, range.max); + } else { + merged.push_back(range); + } + } + return merged; + }; + auto append_ranges = [](Ranges* ranges, int min, int max, int shift, const CodepointRange& suffix + ) { + if (suffix.min == 0 && suffix.max == (1 << shift) - 1) { + ranges->push_back({min << shift, (max << shift) | suffix.max, suffix.target}); + } else { + for (int prefix = min; prefix <= max; ++prefix) { + ranges->push_back( + {(prefix << shift) | suffix.min, (prefix << shift) | suffix.max, suffix.target} + ); + } + } + }; + // Memoize suffixes so wide Unicode classes do not enumerate every codepoint. + std::map, Ranges> suffix_cache; + std::function suffix_ranges = [&](int state, + int remaining) -> const Ranges& { + auto [it, inserted] = suffix_cache.emplace(std::make_pair(state, remaining), Ranges{}); + if (!inserted) return it->second; + Ranges ranges; + if (remaining == 0) { + ranges.push_back({0, 0, state}); + } else { + for (const auto& edge : fsm.GetFsm().GetEdges(state)) { + if (!productive[edge.target]) continue; + XGRAMMAR_CHECK(edge.IsCharRange()); + int min = std::max(edge.min, 0x80), max = std::min(edge.max, 0xbf); + if (min > max) continue; + for (const auto& suffix : suffix_ranges(edge.target, remaining - 1)) { + append_ranges(&ranges, min & 0x3f, max & 0x3f, 6 * (remaining - 1), suffix); + } + } + } + it->second = merge_ranges(std::move(ranges)); + return it->second; + }; + + std::vector rules(fsm.NumStates(), -1); + std::vector pending{fsm.GetStart()}; + // Share complete fallback branches across key prefixes to reuse their token masks. + // The key contains the target state followed by every codepoint interval. + std::map, int32_t> shared_key_continuations; + rules[fsm.GetStart()] = builder_.AddEmptyRuleWithHint(rule_name + "_exclude"); + for (size_t index = 0; index < pending.size(); ++index) { + int state = pending[index]; + Ranges ranges; + for (const auto& edge : fsm.GetFsm().GetEdges(state)) { + if (!productive[edge.target]) continue; + XGRAMMAR_CHECK(edge.IsCharRange()); + if (edge.min < 128) { + ranges.push_back({edge.min, std::min(edge.max, 127), edge.target}); + } + const int leading_min[] = {0, 0xc2, 0xe0, 0xf0}; + const int leading_max[] = {0, 0xdf, 0xef, 0xf4}; + for (int remaining = 1; remaining <= 3; ++remaining) { + int min = std::max(edge.min, leading_min[remaining]); + int max = std::min(edge.max, leading_max[remaining]); + if (min > max) continue; + int mask = (1 << (6 - remaining)) - 1; + Ranges unicode_ranges; + for (const auto& suffix : suffix_ranges(edge.target, remaining)) { + append_ranges(&unicode_ranges, min & mask, max & mask, 6 * remaining, suffix); + } + // The byte FSM can contain non-canonical UTF-8 paths. Never turn an overlong + // encoding into a second transition for an ASCII character (bypassing exclusions). + const int codepoint_min[] = {0, 0x80, 0x800, 0x10000}; + const int codepoint_max[] = {0, 0x7ff, 0xffff, 0x10ffff}; + for (auto range : unicode_ranges) { + range.min = std::max(range.min, codepoint_min[remaining]); + range.max = std::min(range.max, codepoint_max[remaining]); + if (range.min <= range.max) ranges.push_back(range); + } + } + } + std::map> transitions; + for (const auto& range : merge_ranges(std::move(ranges))) { + transitions[range.target].push_back({range.min, range.max}); + } + std::vector choices; + // Keep the closing quote on the accepting body states. A nullable body rule + // would otherwise finish at every character and force token masks to speculate + // across its parent rules. The quote is appended AFTER the exclusion intersection: + // it is JSON syntax, not string content subject to excludes. + if (fsm.IsEndState(state)) { + choices.push_back(close_json_string ? ByteString("\"") : Empty()); + } + for (const auto& [target, codepoints] : transitions) { + if (rules[target] == -1) { + rules[target] = builder_.AddEmptyRuleWithHint(rule_name + "_exclude"); + pending.push_back(target); + } + bool single_ascii = codepoints.size() == 1 && codepoints[0].lower == codepoints[0].upper && + codepoints[0].upper <= 0x7f; + // Keep self-loops local for the compiler's speculative string-mask fast path. + if (!excluded_keys.empty() && !single_ascii && target != state) { + std::vector key{target}; + key.reserve(1 + codepoints.size() * 2); + for (const auto& range : codepoints) { + key.push_back(range.lower); + key.push_back(range.upper); + } + auto [it, inserted] = shared_key_continuations.emplace(std::move(key), -1); + if (inserted) { + it->second = builder_.AddRuleWithHint( + rule_name + "_exclude_continuation", + Sequence({builder_.AddCharacterClass(codepoints), RuleRef(rules[target])}) + ); + } + choices.push_back(RuleRef(it->second)); + } else { + choices.push_back(Sequence({builder_.AddCharacterClass(codepoints), RuleRef(rules[target])}) + ); + } + } + if (choices.empty()) { + choices.push_back(Unsatisfiable()); + } + builder_.UpdateRuleBody(rules[state], Choice(choices)); + } + return RuleRef(rules[fsm.GetStart()]); +} + +int32_t JSONSchemaConverter::GenerateInteger( + const IntegerSpec& spec, const std::string& rule_name +) { + // Shared with ParseInteger's range validation so emission and validation agree on the effective + // range; a nullopt side means that side is unbounded. + const EffectiveIntegerRange range = ComputeEffectiveIntegerRange(spec); + std::optional start = range.start; + std::optional end = range.end; + + if (spec.multiple_of.has_value()) { + // ParseInteger keeps multiple_of only when the range is fully bounded (enumerate the + // multiples) or fully unbounded (emit a modulo DFA); the half-bounded case is dropped there. + if (start.has_value() && end.has_value()) { + std::vector multiples; + for (int64_t value = *start; value <= *end; ++value) { + if (IsMultipleOf(value, *spec.multiple_of)) { + multiples.push_back(ByteString(std::to_string(value))); + } + if (value == std::numeric_limits::max()) { + break; + } + } + return Choice(multiples); + } + return GenerateIntegerMultipleOfDFA(*spec.multiple_of, rule_name); + } + if (start.has_value() || end.has_value()) { + return RegexExpression( + GenerateRangeRegex(start, end), + false, + /*force_cfg_expansion=*/true + ); + } + int32_t optional_minus = Choice({Empty(), ByteString("-")}); + return Choice( + {ByteString("0"), + Sequence( + {optional_minus, + builder_.AddCharacterClass({{'1', '9'}}), + builder_.AddCharacterClassStar({{'0', '9'}})} + )} + ); +} + +int32_t JSONSchemaConverter::GenerateIntegerMultipleOfDFA( + int64_t multiple_of, const std::string& rule_name +) { + std::vector states(multiple_of); + for (int64_t state = 0; state < multiple_of; ++state) { + states[state] = builder_.AddEmptyRuleWithHint( + rule_name + "_multiple_of_" + std::to_string(multiple_of) + "_mod_" + std::to_string(state) + ); + } + for (int64_t state = 0; state < multiple_of; ++state) { + std::vector transitions; + if (state == 0) { + transitions.push_back(Empty()); + } + for (int64_t digit = 0; digit <= 9; ++digit) { + int64_t next_state = (state * 10 + digit) % multiple_of; + transitions.push_back( + Sequence({ByteString(std::to_string(digit)), RuleRef(states[next_state])}) + ); + } + builder_.UpdateRuleBody(states[state], Choice(transitions)); + } + + std::vector non_zero_starts; + for (int64_t digit = 1; digit <= 9; ++digit) { + non_zero_starts.push_back( + Sequence({ByteString(std::to_string(digit)), RuleRef(states[digit % multiple_of])}) + ); + } + return Choice( + {ByteString("0"), Sequence({Choice({Empty(), ByteString("-")}), Choice(non_zero_starts)})} + ); +} + +int32_t JSONSchemaConverter::GenerateNumber(const NumberSpec& spec, const std::string& rule_name) { + std::optional start = spec.minimum; + std::optional end = spec.maximum; + bool exclusive_start = false; + bool exclusive_end = false; + // When both bounds are present the larger lower bound wins; on a tie the + // exclusive one is stricter. + if (spec.exclusive_minimum.has_value() && + (!start.has_value() || *spec.exclusive_minimum >= *start)) { + start = spec.exclusive_minimum; + exclusive_start = true; + } + if (spec.exclusive_maximum.has_value() && (!end.has_value() || *spec.exclusive_maximum <= *end)) { + end = spec.exclusive_maximum; + exclusive_end = true; + } + if (start.has_value() || end.has_value()) { + return RegexExpression( + GenerateFloatRangeRegex(start, end, /*precision=*/6, exclusive_start, exclusive_end), + false, + /*force_cfg_expansion=*/true + ); + } + + int32_t optional_minus = Choice({Empty(), ByteString("-")}); + int32_t integer_part = Choice( + {ByteString("0"), + Sequence( + {builder_.AddCharacterClass({{'1', '9'}}), builder_.AddCharacterClassStar({{'0', '9'}})} + )} + ); + int32_t one_or_more_digits = + Repeat(rule_name + "_digits", builder_.AddCharacterClass({{'0', '9'}}), 1, -1); + int32_t fraction = Choice({Empty(), Sequence({ByteString("."), one_or_more_digits})}); + int32_t exponent = Choice( + {Empty(), + Sequence( + {builder_.AddCharacterClass({{'e', 'e'}, {'E', 'E'}}), + Choice({Empty(), builder_.AddCharacterClass({{'+', '+'}, {'-', '-'}})}), + one_or_more_digits} + )} + ); + // Note: The format must be "-"? ("0" | ...) not ("0" | "-"? ...) + // The first allows -0, -123, 0, 123 + // The second allows 0, -123, 123 but not -0 + return Sequence({optional_minus, integer_part, fraction, exponent}); +} + +int32_t JSONSchemaConverter::GenerateString(const StringSpec& spec, const std::string& rule_name) { + // Check for format + if (spec.format.has_value()) { + auto regex = JSONFormatToRegexPattern(*spec.format); + if (regex.has_value()) { + // The built-in format regexes use constructs that the FSM regex engine does not fully + // support yet (e.g. quoted email local parts), so they keep the CFG expansion. + return Sequence({ByteString("\""), RegexExpression(*regex, false, true), ByteString("\"")}); + } + } + // Check for pattern + if (spec.pattern.has_value()) { + return Sequence( + {ByteString("\""), RegexExpression(*spec.pattern, /*json_string=*/true), ByteString("\"")} + ); + } + // Check for length constraints. They are dropped when there are exclusions: intersecting the + // unrolled bound with the exclusion automaton emits one rule per position and automaton state + // (about 18 rules per character for three markers), so the string keeps only the exclusions. + // Exclusions apply only to the strings without pattern and format above: those constraints + // are the schema's own contract for the string. + if (spec.min_length != 0 || spec.max_length != -1) { + if (!excludes_.empty()) { + WarnDroppedLengthConstraints(spec, rule_name); + } else { + int32_t character = + builder_.AddCharacterClass({{'"', '"'}, {'\\', '\\'}, {'\r', '\r'}, {'\n', '\n'}}, true); + int32_t body = Repeat(rule_name + "_characters", character, spec.min_length, spec.max_length); + return Sequence({ByteString("\""), body, ByteString("\"")}); + } + } + // Default string + return Sequence({ByteString("\""), RuleRef(kBasicStringSub)}); +} + +int32_t JSONSchemaConverter::GenerateBoolean( + const BooleanSpec& spec, const std::string& rule_name +) { + return Choice({ByteString("true"), ByteString("false")}); +} + +int32_t JSONSchemaConverter::GenerateNull(const NullSpec& spec, const std::string& rule_name) { + return ByteString("null"); +} + +int32_t JSONSchemaConverter::GenerateArray(const ArraySpec& spec, const std::string& rule_name) { + indent_manager_.StartIndent(); + int32_t start_separator = FormattingExpression(indent_manager_.StartSeparator()); + int32_t middle_separator = FormattingExpression(indent_manager_.MiddleSeparator()); + int32_t end_separator = FormattingExpression(indent_manager_.EndSeparator()); + int32_t empty_separator = FormattingExpression(indent_manager_.EmptySeparator()); + + std::vector item_rule_ids; + for (size_t index = 0; index < spec.prefix_items.size(); ++index) { + item_rule_ids.push_back( + CreateRule(spec.prefix_items[index], rule_name + "_item_" + std::to_string(index)) + ); + } + int32_t additional_rule_id = -1; + if (spec.allow_additional_items && spec.additional_items) { + additional_rule_id = CreateRule(spec.additional_items, rule_name + "_additional"); + } + indent_manager_.EndIndent(); + + int32_t left_bracket = ByteString("["); + int32_t right_bracket = ByteString("]"); + int32_t empty_array = Sequence({left_bracket, empty_separator, right_bracket}); + + if (item_rule_ids.empty()) { + if (!spec.allow_additional_items || spec.max_items == 0) { + return empty_array; + } + int32_t additional = RuleRef(additional_rule_id); + int32_t tail = Repeat( + rule_name + "_items", + Sequence({middle_separator, additional}), + spec.min_items == 0 ? 0 : static_cast(spec.min_items - 1), + spec.max_items == -1 ? -1 : static_cast(spec.max_items - 1) + ); + int32_t nonempty = + Sequence({left_bracket, start_separator, additional, tail, end_separator, right_bracket}); + return spec.min_items == 0 ? Choice({nonempty, empty_array}) : nonempty; + } + + // Per Draft 2020-12, prefixItems entries are positional: the instance may + // end after any prefix position (subject to minItems), and additional items + // are only allowed after the full prefix (issue #824). + size_t mandatory_count = static_cast(std::min( + std::max(0, spec.min_items), static_cast(item_rule_ids.size()) + )); + + // Mandatory head: the first min(minItems, n) items, separated. + std::vector prefix_elements; + for (size_t index = 0; index < mandatory_count; ++index) { + if (index != 0) { + prefix_elements.push_back(middle_separator); + } + prefix_elements.push_back(RuleRef(item_rule_ids[index])); + } + + // Suffix after the mandatory head, flattened into a right-recursive chain + // of rules suffix_k ::= "" | sep item_k suffix_{k+1} so each position + // is encoded once instead of once per truncation length. The chain ends + // with the additional-items tail. Positions from index 1 on are separated + // by middle_separator; position 0, when it is not part of the mandatory + // head, gets its own rule without the separator. + int32_t suffix = Empty(); + if (spec.allow_additional_items && spec.additional_items) { + int64_t minimum_additional = + std::max(int64_t{0}, spec.min_items - static_cast(item_rule_ids.size())); + suffix = Repeat( + rule_name + "_additional_items", + Sequence({middle_separator, RuleRef(additional_rule_id)}), + static_cast(minimum_additional), + spec.max_items == -1 + ? -1 + : static_cast(spec.max_items - static_cast(item_rule_ids.size())) + ); + } + size_t chain_start = std::max(mandatory_count, 1); + for (size_t k = item_rule_ids.size(); k-- > chain_start;) { + int32_t with_item = Sequence({middle_separator, RuleRef(item_rule_ids[k]), suffix}); + int32_t suffix_rule_id = builder_.AddRuleWithHint( + rule_name + "_suffix_" + std::to_string(k), Choice({Empty(), with_item}) + ); + suffix = RuleRef(suffix_rule_id); + } + if (mandatory_count == 0) { + int32_t with_first = Sequence({RuleRef(item_rule_ids[0]), suffix}); + int32_t suffix_rule_id = + builder_.AddRuleWithHint(rule_name + "_suffix_0", Choice({Empty(), with_first})); + suffix = RuleRef(suffix_rule_id); + } + + std::vector content_elements = prefix_elements; + content_elements.push_back(suffix); + int32_t prefix = Sequence(content_elements); + return Sequence({left_bracket, start_separator, prefix, end_separator, right_bracket}); +} + +int32_t JSONSchemaConverter::FormatPropertyKey( + const std::string& key, const SchemaSpecPtr& schema +) { + if (!IsAllowedLiteral(picojson::value(key))) { + return Unsatisfiable(); + } + return ByteString(picojson::value(key).serialize()); +} + +int32_t JSONSchemaConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + return Sequence({FormatPropertyKey(key, schema), colon_expr_id_, RuleRef(value_rule_id)}); +} + +int32_t JSONSchemaConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + return Sequence({key_pattern_expr, colon_expr_id_, RuleRef(value_rule_id)}); +} + +int32_t JSONSchemaConverter::GetPropertyWithNumberConstraints( + int32_t pattern, + int min_properties, + int max_properties, + int already_repeated_times, + const std::string& rule_name +) { + if (max_properties != -1 && max_properties == already_repeated_times) { + return Empty(); + } + int lower = std::max(0, min_properties - already_repeated_times); + int upper = max_properties == -1 ? -1 : std::max(-1, max_properties - already_repeated_times); + return Repeat(rule_name + "_properties", pattern, lower, upper); +} + +int32_t JSONSchemaConverter::GetAnyOrderRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override +) { + int32_t first_separator = NextSeparatorExpression(); + int32_t middle_separator = NextSeparatorExpression(); + int32_t last_separator = NextSeparatorExpression(true); + + // Build one "item" alternation over every property (any required/optional key) plus any + // additional/pattern key; any_order does not care which key goes where. + std::vector items; + for (size_t index = 0; index < properties.size(); ++index) { + const auto& property = properties[index]; + int32_t value_rule_id = + CreateRule(property.schema, rule_name + "_prop_" + std::to_string(index)); + items.push_back(FormatProperty(property.name, value_rule_id, rule_name, index, property.schema) + ); + } + if (additional != nullptr) { + if (additional_property_override.has_value()) { + items.push_back(*additional_property_override); + } else { + int32_t value_rule_id = CreateRule(additional, rule_name + "_" + additional_suffix); + items.push_back(FormatOtherProperty( + GetKeyPatternExcluding(properties, rule_name), + value_rule_id, + rule_name, + additional_suffix, + additional + )); + } + } + + int32_t item_rule_id = builder_.AddRuleWithHint(rule_name + "_item", Choice(items)); + + // Repeat `item` between n = max(minProperties, #required) and m = maxProperties times; only the + // count is constrained, not which keys appear. + int minimum_count = std::max(min_properties, static_cast(required.size())); + int32_t repeated_items = GetPropertyWithNumberConstraints( + Sequence({middle_separator, RuleRef(item_rule_id)}), + minimum_count, + max_properties, + 1, + rule_name + ); + return Sequence({first_separator, RuleRef(item_rule_id), repeated_items, last_separator}); +} + +int32_t JSONSchemaConverter::GetPartialRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override +) { + if (max_properties == 0) { + return Empty(); + } + if (any_order_) { + return GetAnyOrderRuleForProperties( + properties, + required, + additional, + rule_name, + additional_suffix, + min_properties, + max_properties, + additional_property_override + ); + } + + int32_t first_separator = NextSeparatorExpression(); + int32_t middle_separator = NextSeparatorExpression(); + int32_t last_separator = NextSeparatorExpression(true); + + std::vector property_patterns; + for (size_t index = 0; index < properties.size(); ++index) { + int32_t value_rule_id = + CreateRule(properties[index].schema, rule_name + "_prop_" + std::to_string(index)); + property_patterns.push_back(FormatProperty( + properties[index].name, value_rule_id, rule_name, index, properties[index].schema + )); + } + + bool allow_additional = additional != nullptr; + std::optional additional_pattern; + auto get_additional_pattern = [&]() -> int32_t { + if (!additional_pattern.has_value()) { + if (additional_property_override.has_value()) { + additional_pattern = *additional_property_override; + } else { + int32_t value_rule_id = CreateRule(additional, rule_name + "_" + additional_suffix); + additional_pattern = FormatOtherProperty( + GetKeyPatternExcluding(properties, rule_name), + value_rule_id, + rule_name, + additional_suffix, + additional + ); + } + } + return *additional_pattern; + }; + + if (min_properties == 0 && max_properties == -1) { + // Case 1: No property number constraints + std::vector tails(properties.size(), Empty()); + std::vector is_required(properties.size(), false); + + if (allow_additional) { + int32_t repeated_additional = Repeat( + rule_name + "_additional_properties", + Sequence({middle_separator, get_additional_pattern()}), + 0, + -1 + ); + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(static_cast(properties.size()) - 1), + repeated_additional + ); + tails.back() = RuleRef(tail_rule_id); + } + + for (int index = static_cast(properties.size()) - 2; index >= 0; --index) { + int32_t with_property = + Sequence({middle_separator, property_patterns[index + 1], tails[index + 1]}); + int32_t body = with_property; + if (!required.count(properties[index + 1].name)) { + body = Choice({tails[index + 1], with_property}); + } else { + is_required[index + 1] = true; + } + int32_t tail_rule_id = + builder_.AddRuleWithHint(rule_name + "_part_" + std::to_string(index), body); + tails[index] = RuleRef(tail_rule_id); + } + if (required.count(properties[0].name)) { + is_required[0] = true; + } + + std::vector choices; + for (size_t index = 0; index < properties.size(); ++index) { + choices.push_back(Sequence({property_patterns[index], tails[index]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence({get_additional_pattern(), tails.back()})); + } + return Sequence({first_separator, Choice(choices), last_separator}); + } + + const int property_count = static_cast(properties.size()); + std::vector is_required(property_count, false); + std::vector matched_min(property_count, 0); + bool found_required = required.count(properties[0].name); + matched_min[0] = 1; + for (int index = 1; index < property_count; ++index) { + if (required.count(properties[index].name)) { + is_required[index] = true; + matched_min[index] = matched_min[index - 1] + 1; + } else { + matched_min[index] = matched_min[index - 1]; + } + if (!found_required) { + matched_min[index] = 1; + } + if (is_required[index]) { + found_required = true; + } + } + if (required.count(properties[0].name)) { + is_required[0] = true; + } + + if (max_properties == -1) { + // Case 2: With constraint on the lower bound of the properties number + std::vector> tails(property_count); + matched_min.back() = allow_additional ? std::max(1, matched_min.back()) + : std::max(min_properties, matched_min.back()); + for (int index = property_count - 2; index >= 0; --index) { + matched_min[index] = std::max(matched_min[index], matched_min[index + 1] - 1); + } + + for (int matched = matched_min.back(); matched <= property_count; ++matched) { + int32_t body = allow_additional ? GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + matched, + rule_name + ) + : Empty(); + if (allow_additional) { + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(property_count - 1) + "_" + + std::to_string(matched), + body + ); + tails.back().push_back(RuleRef(tail_rule_id)); + } else { + tails.back().push_back(body); + } + } + + for (int index = property_count - 2; index >= 0; --index) { + for (int matched = matched_min[index]; matched <= index + 1; ++matched) { + int32_t with_property = Sequence( + {middle_separator, + property_patterns[index + 1], + tails[index + 1][matched + 1 - matched_min[index + 1]]} + ); + int32_t body = + (is_required[index + 1] || matched == matched_min[index + 1] - 1) + ? with_property + : Choice({tails[index + 1][matched - matched_min[index + 1]], with_property}); + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(index) + "_" + std::to_string(matched), body + ); + tails[index].push_back(RuleRef(tail_rule_id)); + } + } + + std::vector choices; + for (int index = 0; index < property_count; ++index) { + if (matched_min[index] > 1) { + break; + } + choices.push_back(Sequence({property_patterns[index], tails[index][1 - matched_min[index]]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence( + {get_additional_pattern(), + GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + 1, + rule_name + )} + )); + } + return Sequence({first_separator, Choice(choices), last_separator}); + } + + // Case 3: With constraints on both lower & upper bound of the properties number + std::vector> tails(property_count); + std::vector matched_max(property_count, property_count); + matched_max[0] = 1; + for (int index = 1; index < property_count; ++index) { + matched_max[index] = matched_max[index - 1] + 1; + } + matched_min.back() = allow_additional ? std::max(1, matched_min.back()) + : std::max(min_properties, matched_min.back()); + matched_max.back() = std::min(max_properties, matched_max.back()); + for (int index = property_count - 2; index >= 0; --index) { + matched_min[index] = std::max(matched_min[index], matched_min[index + 1] - 1); + matched_max[index] = is_required[index + 1] + ? std::min(matched_max[index], matched_max[index + 1] - 1) + : std::min(matched_max[index], matched_max[index + 1]); + } + + for (int matched = matched_min.back(); matched <= matched_max.back(); ++matched) { + int32_t body = allow_additional ? GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + matched, + rule_name + ) + : Empty(); + if (allow_additional) { + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(property_count - 1) + "_" + std::to_string(matched), + body + ); + tails.back().push_back(RuleRef(tail_rule_id)); + } else { + tails.back().push_back(body); + } + } + + for (int index = property_count - 2; index >= 0; --index) { + for (int matched = matched_min[index]; matched <= matched_max[index]; ++matched) { + int32_t body; + if (matched == matched_max[index + 1]) { + body = tails[index + 1][matched - matched_min[index + 1]]; + } else { + int32_t with_property = Sequence( + {middle_separator, + property_patterns[index + 1], + tails[index + 1][matched + 1 - matched_min[index + 1]]} + ); + body = (is_required[index + 1] || matched == matched_min[index + 1] - 1) + ? with_property + : Choice({tails[index + 1][matched - matched_min[index + 1]], with_property}); + } + int32_t tail_rule_id = builder_.AddRuleWithHint( + rule_name + "_part_" + std::to_string(index) + "_" + std::to_string(matched), body + ); + tails[index].push_back(RuleRef(tail_rule_id)); + } + } + + std::vector choices; + for (int index = 0; index < property_count; ++index) { + if (matched_max[index] < matched_min[index]) { + continue; + } + if (matched_min[index] > 1) { + break; + } + choices.push_back(Sequence({property_patterns[index], tails[index][1 - matched_min[index]]})); + if (is_required[index]) { + break; + } + } + if (allow_additional && required.empty()) { + choices.push_back(Sequence( + {get_additional_pattern(), + GetPropertyWithNumberConstraints( + Sequence({middle_separator, get_additional_pattern()}), + min_properties, + max_properties, + 1, + rule_name + )} + )); + } + return Sequence({first_separator, Choice(choices), last_separator}); +} + +int32_t JSONSchemaConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool need_braces +) { + // Determine additional property handling + std::string additional_suffix; + SchemaSpecPtr additional_property; + if (spec.allow_additional_properties && spec.additional_properties_schema) { + additional_suffix = "addl"; + additional_property = spec.additional_properties_schema; + } else if (spec.allow_unevaluated_properties && spec.unevaluated_properties_schema) { + additional_suffix = "uneval"; + additional_property = spec.unevaluated_properties_schema; + } else if (spec.allow_additional_properties || spec.allow_unevaluated_properties) { + additional_suffix = "addl"; + additional_property = SchemaSpec::Make(AnySpec{}, "", "any"); + } + + indent_manager_.StartIndent(); + bool has_content = false; + bool could_be_empty = false; + int32_t content = Empty(); + + // Build a key rule through GenerateString rather than spelling out a JSON string here. At the + // JSON root this still produces `"key"`, while XML-style converters override GenerateString to + // produce the unquoted key body expected inside their parameter wrappers. + auto create_pattern_key_rule = [&](const std::string& pattern, + const std::string& rule_name_hint) -> int32_t { + StringSpec key_spec; + key_spec.pattern = pattern; + return CreateRule( + SchemaSpec::Make(std::move(key_spec), /*cache_key=*/"", rule_name_hint), rule_name_hint + ); + }; + + if (!spec.properties.empty() && (!spec.pattern_properties.empty() || spec.property_names)) { + // Case 1a: properties coexist with patternProperties and/or propertyNames. + // Use GetPartialRuleForProperties for named properties, and build + // patternProperties/propertyNames as the additional property pattern override. + SchemaSpecPtr effective_additional = additional_property; + std::string effective_suffix = additional_suffix; + std::optional additional_override; + + if (!spec.pattern_properties.empty()) { + // Build patternProperties as additional property alternatives + std::vector patterns; + for (size_t index = 0; index < spec.pattern_properties.size(); ++index) { + const auto& pattern_property = spec.pattern_properties[index]; + std::string pattern_suffix = "pp_" + std::to_string(index); + int32_t key_rule_id = create_pattern_key_rule( + pattern_property.pattern, rule_name + "_" + pattern_suffix + "_key" + ); + int32_t value_rule_id = + CreateRule(pattern_property.schema, rule_name + "_" + pattern_suffix); + patterns.push_back(FormatOtherProperty( + RuleRef(key_rule_id), value_rule_id, rule_name, pattern_suffix, pattern_property.schema + )); + } + // Merge with existing additionalProperties if present + if (effective_additional) { + int32_t value_rule_id = + CreateRule(effective_additional, rule_name + "_" + effective_suffix); + patterns.push_back(FormatOtherProperty( + KeyPatternExpression(), value_rule_id, rule_name, effective_suffix, effective_additional + )); + } + additional_override = Choice(patterns); + if (!effective_additional) { + effective_additional = SchemaSpec::Make(AnySpec{}, "", "any"); + } + effective_suffix = "pp"; + } else if (spec.property_names && effective_additional) { + // propertyNames constrains keys of additional properties. + // Only apply when additional properties are allowed - when additionalProperties + // is false, no extra keys beyond named properties should be permitted. + int32_t key_rule_id = CreateRule(spec.property_names, rule_name + "_name"); + int32_t value_rule_id = CreateRule(effective_additional, rule_name + "_" + effective_suffix); + additional_override = FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + /*rule_name_suffix=*/"pn", + effective_additional + ); + effective_suffix = "pn"; + } + + content = GetPartialRuleForProperties( + spec.properties, + spec.required, + effective_additional, + rule_name, + effective_suffix, + spec.min_properties, + spec.max_properties, + additional_override + ); + has_content = spec.max_properties != 0; + could_be_empty = spec.required.empty() && spec.min_properties == 0; + } else if (!spec.pattern_properties.empty() || spec.property_names) { + // Case 1b: patternProperties or propertyNames without named properties + if (spec.max_properties != 0) { + int32_t beginning_separator = NextSeparatorExpression(); + std::vector property_choices; + if (!spec.pattern_properties.empty()) { + for (size_t index = 0; index < spec.pattern_properties.size(); ++index) { + const auto& pattern_property = spec.pattern_properties[index]; + std::string pattern_suffix = "prop_" + std::to_string(index); + int32_t key_rule_id = create_pattern_key_rule( + pattern_property.pattern, rule_name + "_" + pattern_suffix + "_key" + ); + int32_t value_rule_id = + CreateRule(pattern_property.schema, rule_name + "_" + pattern_suffix); + property_choices.push_back(Sequence( + {beginning_separator, + FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + pattern_suffix, + pattern_property.schema + )} + )); + } + } else { + int32_t key_rule_id = CreateRule(spec.property_names, rule_name + "_name"); + // propertyNames constrains only the key, so a typed additionalProperties + // schema still applies to the value (issue #826). + int32_t value_rule_id; + if (additional_property) { + value_rule_id = CreateRule(additional_property, rule_name + "_" + additional_suffix); + } else { + value_rule_id = builder_.GetRuleId(GetBasicAnyRuleName()); + XGRAMMAR_DCHECK(value_rule_id != -1); + } + property_choices.push_back(Sequence( + {beginning_separator, + FormatOtherProperty( + RuleRef(key_rule_id), + value_rule_id, + rule_name, + /*rule_name_suffix=*/"pn", + additional_property + )} + )); + } + + int32_t property_rule_id = + builder_.AddRuleWithHint(rule_name + "_prop", Choice(property_choices)); + int32_t subsequent_property = + Sequence({NextSeparatorExpression(), RuleRef(property_rule_id)}); + content = Sequence( + {RuleRef(property_rule_id), + GetPropertyWithNumberConstraints( + subsequent_property, spec.min_properties, spec.max_properties, 1, rule_name + ), + NextSeparatorExpression(true)} + ); + has_content = true; + could_be_empty = spec.min_properties == 0; + } else { + could_be_empty = true; + } + } else if (!spec.properties.empty()) { + // Case 2: properties defined (no patternProperties/propertyNames) + content = GetPartialRuleForProperties( + spec.properties, + spec.required, + additional_property, + rule_name, + additional_suffix, + spec.min_properties, + spec.max_properties + ); + has_content = spec.max_properties != 0; + could_be_empty = spec.required.empty() && spec.min_properties == 0; + } else if (additional_property) { + // Case 3: no properties defined, additional properties allowed + if (spec.max_properties != 0) { + int32_t value_rule_id = CreateRule(additional_property, rule_name + "_" + additional_suffix); + int32_t property = FormatOtherProperty( + KeyPatternExpression(), value_rule_id, rule_name, additional_suffix, additional_property + ); + content = Sequence( + {NextSeparatorExpression(), + property, + GetPropertyWithNumberConstraints( + Sequence({NextSeparatorExpression(), property}), + spec.min_properties, + spec.max_properties, + 1, + rule_name + ), + NextSeparatorExpression(true)} + ); + has_content = true; + } + could_be_empty = spec.min_properties == 0; + } else { + // Case 4: no properties, no additional properties, no pattern properties + // The object is unconditionally empty. + could_be_empty = true; + } + + indent_manager_.EndIndent(); + + int32_t result = need_braces ? Sequence({ByteString("{"), content, ByteString("}")}) : content; + if (could_be_empty) { + int32_t empty_content = any_whitespace_ ? WhitespaceExpression() : Empty(); + int32_t empty_result = + need_braces ? Sequence({ByteString("{"), empty_content, ByteString("}")}) : empty_content; + return has_content ? Choice({result, empty_result}) : empty_result; + } + return result; +} + +int32_t JSONSchemaConverter::GenerateAny(const AnySpec& spec, const std::string& rule_name) { + return Choice( + {RuleRef(kBasicNumber), + RuleRef(kBasicString), + RuleRef(kBasicBoolean), + RuleRef(kBasicNull), + RuleRef(kBasicArray), + RuleRef(kBasicObject)} + ); +} + +int32_t JSONSchemaConverter::GenerateConst(const ConstSpec& spec, const std::string& rule_name) { + if (!IsAllowedJSONLiteral(spec.json_value)) { + return Unsatisfiable(); + } + return ByteString(spec.json_value); +} + +int32_t JSONSchemaConverter::GenerateEnum(const EnumSpec& spec, const std::string& rule_name) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + if (IsAllowedJSONLiteral(value)) { + values.push_back(ByteString(value)); + } + } + if (values.empty()) { + return Unsatisfiable(); + } + return Choice(values); +} + +SchemaSpecPtr JSONSchemaConverter::ResolveRefSchema( + const RefSpec& spec, const std::string& rule_name_hint +) { + if (!ref_resolver_) { + XGRAMMAR_LOG(FATAL) << "Ref resolver not set; cannot resolve $ref: " << spec.uri; + } + return ref_resolver_(spec.uri, rule_name_hint); +} + +std::string JSONSchemaConverter::RefCacheKey(const std::string& uri) const { return uri; } + +int32_t JSONSchemaConverter::GenerateRef(const RefSpec& spec, const std::string& rule_name) { + const std::string cache_key = RefCacheKey(spec.uri); + // First check if we have a direct URI mapping (for circular references) + if (uri_to_rule_id_.count(cache_key)) { + return RuleRef(uri_to_rule_id_[cache_key]); + } + + // Derive rule name from URI path (like original URIToRule) so that the same + // $ref always gets the same rule name, and allocate before resolving to prevent + // dead recursion when the ref target contains a ref back. + std::string rule_name_hint = "ref"; + if (spec.uri.size() >= 2 && spec.uri[0] == '#' && spec.uri[1] == '/') { + std::string new_rule_name_prefix; + std::stringstream ss(spec.uri.substr(2)); + std::string part; + while (std::getline(ss, part, '/')) { + if (!part.empty()) { + if (!new_rule_name_prefix.empty()) { + new_rule_name_prefix += "_"; + } + for (char c : part) { + if (std::isalpha(static_cast(c)) || c == '_' || c == '-' || c == '.') { + new_rule_name_prefix += c; + } + } + } + } + if (!new_rule_name_prefix.empty()) { + rule_name_hint = std::move(new_rule_name_prefix); + } + } + + int32_t allocated_rule_id = builder_.AddEmptyRuleWithHint(rule_name_hint); + std::string allocated_rule_name = builder_.GetRule(allocated_rule_id).name; + uri_to_rule_id_[cache_key] = allocated_rule_id; + SchemaSpecPtr resolved = ResolveRefSchema(spec, allocated_rule_name); + builder_.UpdateRuleBody(allocated_rule_id, GenerateFromSpec(resolved, allocated_rule_name)); + if (!resolved->cache_key.empty()) { + AddCache(resolved->cache_key, allocated_rule_id); + } + return RuleRef(allocated_rule_id); +} + +int32_t JSONSchemaConverter::GenerateAnyOf(const AnyOfSpec& spec, const std::string& rule_name) { + std::vector choices; + for (size_t index = 0; index < spec.options.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.options[index], rule_name + "_case_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GenerateOneOf(const OneOfSpec& spec, const std::string& rule_name) { + std::vector choices; + for (size_t index = 0; index < spec.options.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.options[index], rule_name + "_case_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +int32_t JSONSchemaConverter::GenerateAllOf(const AllOfSpec& spec, const std::string& rule_name) { + if (spec.schemas.size() == 1) { + return GenerateFromSpec(spec.schemas[0], rule_name + "_case_0"); + } + XGRAMMAR_LOG(WARNING) << "Support for allOf with multiple options is still ongoing"; + return GenerateFromSpec(SchemaSpec::Make(AnySpec{}, "", "any"), rule_name); +} + +int32_t JSONSchemaConverter::GenerateTypeArray( + const TypeArraySpec& spec, const std::string& rule_name +) { + std::vector choices; + for (size_t index = 0; index < spec.type_schemas.size(); ++index) { + choices.push_back( + RuleRef(CreateRule(spec.type_schemas[index], rule_name + "_type_" + std::to_string(index))) + ); + } + return Choice(choices); +} + +// ==================== Static Helper Methods ==================== + +std::optional JSONSchemaConverter::JSONFormatToRegexPattern(const std::string& format +) { + static const auto regex_map = []() -> std::unordered_map { + std::unordered_map m; + + std::string atext = "[\\w!#$%&'*+/=?^`{|}~-]"; + std::string dot_string = "(" + atext + "+(\\." + atext + "+)*)"; + std::string quoted_string = + "\\\\\"(\\\\[\\x20-\\x7E]|[\\x20\\x21\\x23-\\x5B\\x5D-\\x7E])*\\\\\""; + std::string domain = + "([A-Za-z0-9]([\\-A-Za-z0-9]*[A-Za-z0-9])?)((\\.[A-Za-z0-9][\\-A-Za-z0-9]*[A-Za-z0-9])*" + ")"; + m["email"] = "^(" + dot_string + "|" + quoted_string + ")@" + domain + "$"; + + m["date"] = "^(\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2]\\d|3[01]))$"; + m["time"] = + "^([01]\\d|2[0-3]):[0-5]\\d:([0-5]\\d|60)(\\.\\d+)?(Z|[+-]([01]\\d|2[0-3]):[0-5]\\d)$"; + m["date-time"] = + "^(\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[1-2]\\d|3[01]))T([01]\\d|2[0-3]):[0-5]\\d:([0-5]\\d|60)(" + "\\.\\d+)?(Z|[+-]([01]\\d|2[0-3]):[0-5]\\d)$"; + m["duration"] = + "^P((\\d+D|\\d+M(\\d+D)?|\\d+Y(\\d+M(\\d+D)?)?)(T(\\d+S|\\d+M(\\d+S)?|\\d+H(\\d+M(\\d+" + "S)?" + ")?))?|T(\\d+S|\\d+M(\\d+S)?|\\d+H(\\d+M(\\d+S)?)?)|\\d+W)$"; + + std::string decbyte = "(25[0-5]|2[0-4]\\d|[0-1]?\\d?\\d)"; + m["ipv4"] = "^(" + decbyte + "\\.){3}" + decbyte + "$"; + + m["ipv6"] = + "(" + "([0-9a-fA-F]{1,4}:){7,7}[0-9a-fA-F]{1,4}|" + "([0-9a-fA-F]{1,4}:){1,7}:|" + "([0-9a-fA-F]{1,4}:){1,6}:[0-9a-fA-F]{1,4}|" + "([0-9a-fA-F]{1,4}:){1,5}(:[0-9a-fA-F]{1,4}){1,2}|" + "([0-9a-fA-F]{1,4}:){1,4}(:[0-9a-fA-F]{1,4}){1,3}|" + "([0-9a-fA-F]{1,4}:){1,3}(:[0-9a-fA-F]{1,4}){1,4}|" + "([0-9a-fA-F]{1,4}:){1,2}(:[0-9a-fA-F]{1,4}){1,5}|" + "[0-9a-fA-F]{1,4}:((:[0-9a-fA-F]{1,4}){1,6})|" + ":((:[0-9a-fA-F]{1,4}){1,7}|:)|" + "::(ffff(:0{1,4}){0,1}:){0,1}" + "((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\\.){3,3}" + "(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])|" + "([0-9a-fA-F]{1,4}:){1,4}:" + "((25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])\\.){3,3}" + "(25[0-5]|(2[0-4]|1{0,1}[0-9]){0,1}[0-9])" + ")"; + + m["hostname"] = "^([a-z0-9]([a-z0-9-]*[a-z0-9])?)(\\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)*$"; + m["uuid"] = "^[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}$"; + + std::string schema_pat = "[a-zA-Z][a-zA-Z+\\.-]*"; + std::string pchar = "([\\w\\.~!$&'()*+,;=:@-]|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string query_fragment_char = "([\\w\\.~!$&'()*+,;=:@/\\?-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string query = "(\\?" + query_fragment_char + ")?"; + std::string fragment = "(#" + query_fragment_char + ")?"; + std::string path_abempty = "(/" + pchar + "*)*"; + std::string path_absolute_rootless_empty = "/?(" + pchar + "+(/" + pchar + "*)*)?"; + std::string userinfo = "([\\w\\.~!$&'()*+,;=:-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string host = "([\\w\\.~!$&'()*+,;=-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + std::string authority = "(" + userinfo + "@)?" + host + "(:\\d*)?"; + std::string hier_part = + "(//" + authority + path_abempty + "|" + path_absolute_rootless_empty + ")"; + m["uri"] = "^" + schema_pat + ":" + hier_part + query + fragment + "$"; + + pchar = "([\\w\\.~!$&'()*+,;=:@-]|%[0-9A-Fa-f][0-9A-Fa-f])"; + query_fragment_char = "([\\w\\.~!$&'()*+,;=:@/\\?-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + query = "(\\?" + query_fragment_char + ")?"; + fragment = "(#" + query_fragment_char + ")?"; + path_abempty = "(/" + pchar + "*)*"; + std::string path_absolute = "/(" + pchar + "+(/" + pchar + "*)*)?"; + std::string segment_nz_nc = "([\\w\\.~!$&'()*+,;=@-]|%[0-9A-Fa-f][0-9A-Fa-f])+"; + std::string path_noscheme = segment_nz_nc + "(/" + pchar + "*)*"; + userinfo = "([\\w\\.~!$&'()*+,;=:-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + host = "([\\w\\.~!$&'()*+,;=-]|%[0-9A-Fa-f][0-9A-Fa-f])*"; + authority = "(" + userinfo + "@)?" + host + "(:\\d*)?"; + std::string relative_part = + "(//" + authority + path_abempty + "|" + path_absolute + "|" + path_noscheme + ")?"; + m["uri-reference"] = "^" + relative_part + query + fragment + "$"; + + std::string literals = + "([\\x21\\x23-\\x24\\x26\\x28-\\x3B\\x3D\\x3F-\\x5B\\x5D\\x5F\\x61-\\x7A\\x7E]" + "|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string op = "[+#\\./;\\?&=,!@|]"; + std::string varchar = "(\\w|%[0-9A-Fa-f][0-9A-Fa-f])"; + std::string varname = varchar + "(\\.?" + varchar + ")*"; + std::string varspec = varname + "(:[1-9]\\d?\\d?\\d?|\\*)?"; + std::string variable_list = varspec + "(," + varspec + ")*"; + std::string expression = "\\{(" + op + ")?" + variable_list + "\\}"; + m["uri-template"] = "^(" + literals + "|" + expression + ")*$"; + + m["json-pointer"] = "^(/([\\x00-\\x2E]|[\\x30-\\x7D]|[\\x7F-\\U0010FFFF]|~[01])*)*$"; + m["relative-json-pointer"] = + "^(0|[1-9][0-9]*)(#|(/([\\x00-\\x2E]|[\\x30-\\x7D]|[\\x7F-\\U0010FFFF]|~[01])*)*)$"; + + return m; + }(); + + auto it = regex_map.find(format); + if (it == regex_map.end()) { + return std::nullopt; + } + return it->second; +} + +// ==================== XMLToolCallingConverter Implementation ==================== + +namespace { + +constexpr const char* kStringCacheKey = "{\"type\":\"string\"}"; +constexpr const char* kObjectCacheKey = "{\"type\":\"object\"}"; + +} // namespace + +const std::string XMLToolCallingConverter::kXMLString = "xml_string"; +const std::string XMLToolCallingConverter::kXMLAny = "xml_any"; +const std::string XMLToolCallingConverter::kXMLObject = "xml_object"; +const std::string XMLToolCallingConverter::kXMLVariableName = "xml_variable_name"; + +const std::unordered_map + XMLToolCallingConverter::kKeyWrapperMap = { + {JSONFormat::kQwenXML, converter_ext::GetQwenXMLWrapper()}, + {JSONFormat::kMiniMaxXML, converter_ext::GetMiniMaxXMLWrapper()}, + {JSONFormat::kDeepSeekXML, converter_ext::GetDeepSeekXMLWrapper()}, + {JSONFormat::kDeepSeekV41XML, converter_ext::GetDeepSeekV41XMLWrapper()}, + {JSONFormat::kGlmXML, converter_ext::GetGLMXMLWrapper()}, + {JSONFormat::kCohereXML, converter_ext::GetCohereXMLWrapper()}, + {JSONFormat::kKimiK3XML, converter_ext::GetKimiK3XMLWrapper()}, +}; + +XMLToolCallingConverter::XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver, + JSONFormat json_format, + bool any_order, + std::vector excludes +) + : JSONSchemaConverter( + indent, + separators, + any_whitespace, + max_whitespace_cnt, + ref_resolver, + any_order, + std::move(excludes) + ), + json_format_(json_format), + nested_object_level_(0), + xml_wrapper_(kKeyWrapperMap.at(json_format)) { + // XML formatting can add whitespace outside a constrained raw string's rule. + // Reject exclusions that could straddle that boundary instead of silently bypassing them. + auto is_padding = [](char c) { return c == ' ' || c == '\t' || c == '\r' || c == '\n'; }; + for (const auto& excluded : excludes_) { + XGRAMMAR_CHECK( + excluded.empty() || (!is_padding(excluded.front()) && !is_padding(excluded.back())) + ) << "XML JSONSchemaFormat.excludes must not start or end with formatting whitespace"; + } +} + +Grammar XMLToolCallingConverter::Convert(const SchemaSpecPtr& spec) { + nested_object_level_ = 0; + return JSONSchemaConverter::Convert(spec); +} + +std::optional XMLToolCallingConverter::GetRenderedJSONType(const SchemaSpecPtr& spec) { + if (spec == nullptr) { + return std::nullopt; + } + auto type_of_json_value = [](const std::string& json_value) -> std::optional { + picojson::value value; + if (!ParseJSON(value, json_value).empty()) { + return std::nullopt; + } + if (value.is()) return "string"; + if (value.is()) return "boolean"; + if (value.is()) return "number"; + if (value.is()) return "null"; + if (value.is()) return "object"; + if (value.is()) return "array"; + return std::nullopt; + }; + + return std::visit( + [&](auto&& arg) -> std::optional { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return "string"; + } else if constexpr (std::is_same_v || std::is_same_v) { + // Both integer and floating-point values have the JSON type "number". + return "number"; + } else if constexpr (std::is_same_v) { + return "boolean"; + } else if constexpr (std::is_same_v) { + return "null"; + } else if constexpr (std::is_same_v) { + return "array"; + } else if constexpr (std::is_same_v) { + return "object"; + } else if constexpr (std::is_same_v) { + return type_of_json_value(arg.json_value); + } else if constexpr (std::is_same_v) { + // Only pin the attribute when every alternative renders with the same type. + std::optional common; + for (const auto& json_value : arg.json_values) { + auto type_name = type_of_json_value(json_value); + if (!type_name.has_value()) return std::nullopt; + if (!common.has_value()) { + common = type_name; + } else if (*common != *type_name) { + return std::nullopt; + } + } + return common; + } else { + // Any, $ref and the combinators may render as more than one type; keep them open. + return std::nullopt; + } + }, + spec->spec + ); +} + +std::string XMLToolCallingConverter::XMLValue(const std::string& json_value) const { + picojson::value value; + std::string error = ParseJSON(value, json_value); + if (error.empty() && value.is()) { + return value.get(); + } + return json_value; +} + +int32_t XMLToolCallingConverter::XMLKeySuffix(const std::optional& pinned_type) { + auto value_choices = [this](const std::vector& values) { + std::vector choices; + choices.reserve(values.size()); + for (const auto* value : values) { + choices.push_back(ByteString(value)); + } + return Choice(choices); + }; + if (json_format_ == JSONFormat::kKimiK3XML) { + const auto& suffix = converter_ext::GetKimiK3XMLKeySuffix(); + // A declared property carries exactly the type its value grammar is rendered with, so the + // parser decodes the value back to the schema's type. Free-form keys have no single schema + // type, so they keep the full set. + int32_t type_expr = + pinned_type.has_value() ? ByteString(*pinned_type) : value_choices(suffix.values); + return Sequence({ByteString(suffix.prefix), type_expr, ByteString(suffix.suffix)}); + } + return ByteString(xml_wrapper_.key_wrapper_suffix); +} + +void XMLToolCallingConverter::AddBasicRules() { + // First add JSON basic rules. These should be in the inner layer of the XML format. + XGRAMMAR_DCHECK(nested_object_level_ == 0); + // The nested part, true json format, is at level 2. + nested_object_level_ = 2; + JSONSchemaConverter::AddBasicRules({kXMLString, kXMLAny, kXMLObject, kXMLVariableName}); + + auto any_spec = SchemaSpec::Make(AnySpec{}, "{}", kBasicAny); + + // The outer part, xml format, is at level 1. + nested_object_level_ = 1; + // Keep the unrestricted raw body as a single TagDispatch. The argument suffix is matched + // by the enclosing property rule, outside this body's exclusion scope. + auto string_excludes = excludes_; + string_excludes.push_back(xml_wrapper_.parameter_suffix); + builder_.UpdateRuleBody(kXMLString, TagDispatch(false, std::move(string_excludes))); + AddCache(kStringCacheKey, builder_.GetRuleId(kXMLString)); + + // Add XML any rule + builder_.UpdateRuleBody(kXMLAny, GenerateAny(AnySpec{}, kXMLAny)); + AddCache("{}", builder_.GetRuleId(kXMLAny)); + + // Reset the nested object level to 0, which is the root level. + nested_object_level_ = 0; + + // Add XML object rule + ObjectSpec xml_object_spec; + xml_object_spec.allow_additional_properties = true; + xml_object_spec.additional_properties_schema = any_spec; + builder_.UpdateRuleBody(kXMLObject, GenerateObject(xml_object_spec, kXMLObject)); + AddCache(kObjectCacheKey, builder_.GetRuleId(kXMLObject)); + + // Add XML variable name rule + builder_.UpdateRuleBody( + kXMLVariableName, + !excludes_.empty() + ? ExcludingString("[a-zA-Z_][a-zA-Z0-9_]*", kXMLVariableName, false) + : Sequence( + {builder_.AddCharacterClass({{'a', 'z'}, {'A', 'Z'}, {'_', '_'}}), + builder_.AddCharacterClassStar({{'a', 'z'}, {'A', 'Z'}, {'0', '9'}, {'_', '_'}})} + ) + ); +} + +std::string XMLToolCallingConverter::GetKeyPattern() const { + if (nested_object_level_ <= 1) { + return kXMLVariableName; + } + return kBasicString; +} + +std::string XMLToolCallingConverter::GetBasicAnyRuleName() const { + if (nested_object_level_ <= 1) { + return kXMLAny; + } + return kBasicAny; +} + +int32_t XMLToolCallingConverter::GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name +) { + if (nested_object_level_ <= 1) { + return RuleRef(GetKeyPattern()); + } + return JSONSchemaConverter::GetKeyPatternExcluding(properties, rule_name); +} + +std::string XMLToolCallingConverter::NextSeparator(bool is_end) { + if (nested_object_level_ <= 1) { + return GetWhitespacePattern(); + } + return JSONSchemaConverter::NextSeparator(is_end); +} + +int32_t XMLToolCallingConverter::GenerateString( + const StringSpec& spec, const std::string& rule_name +) { + if (nested_object_level_ <= 1) { + if (spec.format.has_value()) { + auto regex = JSONFormatToRegexPattern(*spec.format); + if (regex.has_value()) { + return RegexExpression(*regex, false, true); + } + } + if (spec.pattern.has_value()) { + return RegexExpression(*spec.pattern, false, /*force_cfg_expansion=*/true); + } + const bool bounded = spec.min_length != 0 || spec.max_length != -1; + // Without exclusions, a length bound or an unrecognized format keeps the plain repetition. + if (excludes_.empty() && (bounded || spec.format.has_value())) { + return Repeat( + rule_name + "_characters", + builder_.AddCharacterClass({{0, 0x10ffff}}), + spec.min_length, + spec.max_length + ); + } + // With exclusions the raw string keeps only the exclusions: length constraints are dropped + // (see JSONSchemaConverter::GenerateString) and an unrecognized format is unconstrained. + if (bounded) { + WarnDroppedLengthConstraints(spec, rule_name); + } + return RuleRef(kXMLString); + } + return JSONSchemaConverter::GenerateString(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateAny(const AnySpec& spec, const std::string& rule_name) { + if (nested_object_level_ == 0) { + return RuleRef(kXMLObject); + } + if (nested_object_level_ == 1) { + return Choice({RuleRef(kXMLString), RuleRef(kBasicArray), RuleRef(kBasicObject)}); + } + return JSONSchemaConverter::GenerateAny(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateArray( + const ArraySpec& spec, const std::string& rule_name +) { + nested_object_level_++; + auto result = JSONSchemaConverter::GenerateArray(spec, rule_name); + nested_object_level_--; + return result; +} + +int32_t XMLToolCallingConverter::GenerateConst( + const ConstSpec& spec, const std::string& rule_name +) { + if (nested_object_level_ == 0) { + picojson::value value; + XGRAMMAR_CHECK(ParseJSON(value, spec.json_value).empty()); + if (value.is()) { + // A root object is a parameter list, including when all its values are fixed. + // Nested object constants still use the JSON representation below. + ObjectSpec object; + object.allow_unevaluated_properties = false; + const auto& properties = value.get(); + for (const auto& key : properties.ordered_keys()) { + object.properties.push_back( + {key, SchemaSpec::Make(ConstSpec{properties.at(key).serialize()})} + ); + object.required.insert(key); + } + // As with JSON literals, keep a fixed order even when any_order is enabled. + // The general any-order object rule permits repeated keys and is not exact for const. + bool saved_any_order = any_order_; + any_order_ = false; + int32_t result = GenerateObject(object, rule_name); + any_order_ = saved_any_order; + return result; + } + } + if (nested_object_level_ <= 1) { + if (!IsAllowedJSONLiteral(spec.json_value, /*raw_string=*/true)) { + return Unsatisfiable(); + } + return ByteString(XMLValue(spec.json_value)); + } + return JSONSchemaConverter::GenerateConst(spec, rule_name); +} + +int32_t XMLToolCallingConverter::GenerateEnum(const EnumSpec& spec, const std::string& rule_name) { + XGRAMMAR_DCHECK(!spec.json_values.empty()) + << "GenerateEnum called with empty enum spec for rule: " << rule_name; + if (nested_object_level_ <= 1) { + std::vector values; + values.reserve(spec.json_values.size()); + for (const auto& value : spec.json_values) { + values.push_back(GenerateConst(ConstSpec{value}, rule_name)); + } + return Choice(values); + } + return JSONSchemaConverter::GenerateEnum(spec, rule_name); +} + +int32_t XMLToolCallingConverter::FormatPropertyKey( + const std::string& key, const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + // Only kimi_k3_xml encodes the value's type next to the key; the other formats would + // discard the result, so don't walk the schema for them. + std::optional pinned_type; + if (json_format_ == JSONFormat::kKimiK3XML) { + pinned_type = GetRenderedJSONType(schema); + } + if (!IsAllowedString(EscapeAttrValue(key))) { + return Unsatisfiable(); + } + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix + EscapeAttrValue(key)), + XMLKeySuffix(pinned_type)} + ); + } + return JSONSchemaConverter::FormatPropertyKey(key, schema); +} + +int32_t XMLToolCallingConverter::FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + if (json_format_ == JSONFormat::kDeepSeekXML || json_format_ == JSONFormat::kDeepSeekV41XML) { + if (!IsAllowedString(key)) { + return Unsatisfiable(); + } + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix + key), + FormatDeepSeekParamSuffix(schema, value_rule_id)} + ); + } + std::vector elements = {FormatPropertyKey(key, schema)}; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + // xml_string already accepts whitespace. Adding whitespace repetitions around it preserves the + // language but creates one Earley state for every possible split with the string body. + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + elements.push_back(RuleRef(value_rule_id)); + } else { + elements.push_back(WhitespaceExpression()); + elements.push_back(RuleRef(value_rule_id)); + elements.push_back(WhitespaceExpression()); + } + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + } + return JSONSchemaConverter::FormatProperty(key, value_rule_id, rule_name, idx, schema); +} + +int32_t XMLToolCallingConverter::FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema +) { + if (nested_object_level_ <= 1) { + if (json_format_ == JSONFormat::kDeepSeekXML || json_format_ == JSONFormat::kDeepSeekV41XML) { + return Sequence( + {ByteString(xml_wrapper_.key_wrapper_prefix), + key_pattern_expr, + FormatDeepSeekParamSuffix(schema, value_rule_id)} + ); + } + std::vector elements = { + ByteString(xml_wrapper_.key_wrapper_prefix), + key_pattern_expr, + XMLKeySuffix( + json_format_ == JSONFormat::kKimiK3XML ? GetRenderedJSONType(schema) : std::nullopt + ) + }; + if (!xml_wrapper_.value_wrapper_prefix.empty()) { + elements.push_back(WhitespaceExpression()); + elements.push_back(ByteString(xml_wrapper_.value_wrapper_prefix)); + } + if (value_rule_id == builder_.GetRuleId(kXMLString)) { + elements.push_back(RuleRef(value_rule_id)); + } else { + elements.push_back(WhitespaceExpression()); + elements.push_back(RuleRef(value_rule_id)); + elements.push_back(WhitespaceExpression()); + } + elements.push_back(ByteString(xml_wrapper_.parameter_suffix)); + return Sequence(elements); + } + return JSONSchemaConverter::FormatOtherProperty( + key_pattern_expr, value_rule_id, rule_name, rule_name_suffix, schema + ); +} + +int32_t XMLToolCallingConverter::GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces +) { + nested_object_level_++; + bool need_brace = nested_object_level_ > 1; + auto result = JSONSchemaConverter::GenerateObject(spec, rule_name, need_brace); + nested_object_level_--; + return result; +} + +void XMLToolCallingConverter::AddCache(const std::string& key, int32_t rule_id) { + if (key.empty()) { + return; + } + rule_cache_manager_.AddCache(key, EncodingContext(), rule_id); +} + +std::optional XMLToolCallingConverter::GetCache(const std::string& key) const { + if (key.empty()) { + return std::nullopt; + } + if ((json_format_ == JSONFormat::kQwenXML || json_format_ == JSONFormat::kDeepSeekXML || + json_format_ == JSONFormat::kDeepSeekV41XML) && + nested_object_level_ == 0 && key == "{}") { + // Unconstrained tool arguments are an XML parameter list, not one parameter's raw value. + return rule_cache_manager_.GetCache(kObjectCacheKey, 0); + } + // At level 0, {"type":"object"} is the root tool-arguments object and uses XML parameter + // tags. At level 1 it is the value of one such parameter and must use the inner JSON object + // rule, including braces. Without this distinction, the outer XML object cache is reused for + // the value before GenerateObject() can advance nested_object_level_. + if (nested_object_level_ == 1 && key == kObjectCacheKey) { + return rule_cache_manager_.GetCache(key, 2); + } + return rule_cache_manager_.GetCache(key, EncodingContext()); +} + +std::string XMLToolCallingConverter::RefCacheKey(const std::string& uri) const { + return std::to_string(EncodingContext()) + ":" + uri; +} + +// ==================== Range Regex Generation ==================== + +// Stateless utility that turns a numeric range into an anchored regex matching +// exactly the JSON integers / numbers inside it. Every method is static; the +// class exists only to group the helpers and keep the internal ones private. +class NumberGenerator { + public: + // Anchored regex matching every integer x with start <= x <= end. Either bound + // may be std::nullopt for an open side; an empty range yields "^()$". Bounds + // span the whole int64 range (|INT64_MIN| is handled without negation overflow). + static std::string IntegerRangeRegex(std::optional start, std::optional end); + + // Anchored regex matching every number in the range, written with up to + // `precision` fraction digits. `exclusive_start` / `exclusive_end` exclude the + // boundary value itself (turning >= / <= into > / <). Either bound may be + // std::nullopt for an open side; an empty range yields "^()$". + static std::string FloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end + ); + + private: + // Regex alternatives for the fraction digits following a decimal point. + struct FracPatternSet { + // Each pattern matches a non-empty fraction digit string. + std::vector parts; + // Whether having no fraction digits at all also satisfies the bound. + bool include_empty = false; + }; + + // --- Regex fragment primitives --- + static std::string DigitClass(char lo, char hi); // one digit in [lo, hi] (or \d) + static std::string ExactDigits(int k); // exactly k free digits: \d{k} + static std::string FreeDigits(int max_count); // 0..max_count free digits: \d{0,n} + static std::string OptionalZeros(int max_count); // 0..max_count zeros: 0{0,n} + static std::string SomeZeros(int max_count); // 1..max_count zeros: 0{1,n} + static bool AllChar(const std::string& s, char c); + + // --- Integer range (operate on non-negative decimal magnitude strings) --- + static std::string AbsDigits(int64_t v); + static int CompareDigitStr(const std::string& a, const std::string& b); + static std::vector IntSameLen(const std::string& a, const std::string& b); + static std::vector NumberPatternsStr(const std::string& lo, const std::string& hi); + static std::string SubRangeRegexStr(const std::string& lo, const std::string& hi); + static std::vector AtLeastPositivePatternsStr(const std::string& v_str); + + // --- Float range --- + static std::string FormatFloat(double value, int precision); + // Snaps a non-negative bound to the precision grid in the direction that keeps + // the range sound: a lower bound rounds up, an upper bound rounds down, so no + // out-of-range value is ever admitted. Returns the canonical grid string and, + // via strict_out, whether the boundary value must still be excluded. + static std::string RoundBoundToGrid( + double value, int precision, bool is_lower, bool strict_in, bool* strict_out + ); + // Adds (inc) or subtracts (!inc) one grid step (10^-precision) to a canonical + // non-negative decimal string, returning the canonical result. + static std::string AdjustGrid(const std::string& s, int precision, bool inc); + static void SplitDecimal(const std::string& s, std::string* int_part, std::string* frac_part); + static int CompareDecimal( + const std::string& int_a, + const std::string& frac_a, + const std::string& int_b, + const std::string& frac_b + ); + static std::string StripAnchors(const std::string& regex); + static int64_t ParseIntCapped(const std::string& digits); + static FracPatternSet FracGreaterPatterns(const std::string& s, bool strict, int max_len); + static FracPatternSet FracLessPatterns(const std::string& s, bool strict, int max_len); + static FracPatternSet FracBetweenPatterns( + const std::string& a, bool strict_a, const std::string& b, bool strict_b, int max_len + ); + static std::vector PositiveRangeParts( + const std::string& low, + bool strict_low, + const std::optional& high, + bool strict_high, + int precision + ); +}; + +// Helpers for integer range regex generation. They operate purely on +// fixed-length decimal digit strings (suffixes may carry leading zeros), so the +// patterns are correct by construction regardless of digit position. + +// A regex fragment matching a single digit in [lo, hi]. +std::string NumberGenerator::DigitClass(char lo, char hi) { + if (lo == hi) { + return std::string(1, lo); + } + if (lo == '0' && hi == '9') { + return "\\d"; + } + return "[" + std::string(1, lo) + "-" + std::string(1, hi) + "]"; +} + +// A regex fragment matching k free digits (each 0-9). Empty when k <= 0. +std::string NumberGenerator::ExactDigits(int k) { + if (k <= 0) { + return ""; + } + if (k == 1) { + return "\\d"; + } + return "\\d{" + std::to_string(k) + "}"; +} + +bool NumberGenerator::AllChar(const std::string& s, char c) { + return std::all_of(s.begin(), s.end(), [c](char ch) { return ch == c; }); +} + +// Patterns matching every equal-length digit string t with +// value(a) <= value(t) <= value(b). Requires a.size() == b.size() and +// value(a) <= value(b). Partitions t by its first digit: +// * first digit == a[0]: the suffix must be >= a's suffix (<= 99..9); +// * first digit strictly between a[0] and b[0]: the suffix is unconstrained; +// * first digit == b[0]: the suffix must be <= b's suffix (>= 00..0). +// The partition is exact and non-overlapping, so the union is sound and +// complete for [a, b]. +std::vector NumberGenerator::IntSameLen(const std::string& a, const std::string& b) { + int n = static_cast(a.size()); + if (a == b) { + return {a}; + } + if (n == 1) { + return {DigitClass(a[0], b[0])}; + } + if (a[0] == b[0]) { + std::vector res; + for (auto& p : IntSameLen(a.substr(1), b.substr(1))) { + res.push_back(std::string(1, a[0]) + p); + } + return res; + } + // a[0] < b[0] + std::string a_suf = a.substr(1); + std::string b_suf = b.substr(1); + if (AllChar(a_suf, '0') && AllChar(b_suf, '9')) { + // The whole suffix space is free: collapse to one box pattern. + if (a[0] == '0' && b[0] == '9') { + return {ExactDigits(n)}; + } + return {DigitClass(a[0], b[0]) + ExactDigits(n - 1)}; + } + std::vector res; + std::string nines(n - 1, '9'); + std::string zeros(n - 1, '0'); + for (auto& p : IntSameLen(a_suf, nines)) { + res.push_back(std::string(1, a[0]) + p); + } + if (b[0] - a[0] >= 2) { + res.push_back( + DigitClass(static_cast(a[0] + 1), static_cast(b[0] - 1)) + ExactDigits(n - 1) + ); + } + for (auto& p : IntSameLen(zeros, b_suf)) { + res.push_back(std::string(1, b[0]) + p); + } + return res; +} + +// Compares two non-negative decimal magnitude strings (no leading zeros except +// "0") by value. +int NumberGenerator::CompareDigitStr(const std::string& a, const std::string& b) { + if (a.size() != b.size()) { + return a.size() < b.size() ? -1 : 1; + } + if (a < b) { + return -1; + } + return a > b ? 1 : 0; +} + +// Patterns matching every integer whose magnitude has value in [lo, hi], where +// lo and hi are non-negative decimal magnitude strings (no leading zeros except +// "0"). An empty range (value(lo) > value(hi)) yields no patterns. Operating on +// strings keeps the whole int64 range representable, including +// |INT64_MIN| = 9223372036854775808, which does not fit in int64. +std::vector NumberGenerator::NumberPatternsStr( + const std::string& lo, const std::string& hi +) { + std::vector patterns; + if (CompareDigitStr(lo, hi) > 0) { + return patterns; + } + int lo_len = static_cast(lo.size()); + int hi_len = static_cast(hi.size()); + // Split [lo, hi] by digit length; each length yields a same-length segment + // handled exactly by IntSameLen. + for (int len = lo_len; len <= hi_len; ++len) { + std::string a_str = (len == lo_len) ? lo : ("1" + std::string(len - 1, '0')); + std::string b_str = (len == hi_len) ? hi : std::string(len, '9'); + for (auto& p : IntSameLen(a_str, b_str)) { + patterns.push_back(p); + } + } + return patterns; +} + +// Joins NumberPatternsStr alternatives into a parenthesised regex group. +std::string NumberGenerator::SubRangeRegexStr(const std::string& lo, const std::string& hi) { + std::vector patterns = NumberPatternsStr(lo, hi); + std::string joined; + for (size_t i = 0; i < patterns.size(); ++i) { + if (i > 0) { + joined += "|"; + } + joined += patterns[i]; + } + return "(" + joined + ")"; +} + +// Patterns matching every integer in [value(v_str), +infinity) for v_str a +// positive magnitude string (no leading zeros). Same-length values come from +// IntSameLen(v_str, 99..9); strictly longer values are any non-zero-led number. +std::vector NumberGenerator::AtLeastPositivePatternsStr(const std::string& v_str) { + int len = static_cast(v_str.size()); + std::vector res = IntSameLen(v_str, std::string(len, '9')); + res.push_back("[1-9]\\d{" + std::to_string(len) + ",}"); + return res; +} + +// The magnitude (absolute value) of v as a decimal string. Derived from the +// signed text rather than by negating v, so INT64_MIN is handled correctly. +std::string NumberGenerator::AbsDigits(int64_t v) { + std::string s = std::to_string(v); + return (!s.empty() && s[0] == '-') ? s.substr(1) : s; +} + +std::string NumberGenerator::IntegerRangeRegex( + std::optional start, std::optional end +) { + std::vector parts; + std::ostringstream result; + + if (!start && !end) { + return "^-?\\d+$"; + } + + if (start && !end) { + if (start.value() <= 0) { + if (start.value() < 0) { + // Negatives in [start, -1] are the magnitudes [1, |start|], negated. + parts.push_back("-" + SubRangeRegexStr("1", AbsDigits(start.value()))); + } + parts.push_back("0"); + parts.push_back("[1-9]\\d*"); + } else { + // x >= start with start > 0: same-length values >= start, plus every + // value with strictly more digits. + for (auto& p : AtLeastPositivePatternsStr(std::to_string(start.value()))) { + parts.push_back(p); + } + } + } + + if (!start && end) { + if (end.value() >= 0) { + parts.push_back("-[1-9]\\d*"); + parts.push_back("0"); + if (end.value() > 0) { + parts.push_back(SubRangeRegexStr("1", std::to_string(end.value()))); + } + } else { + // x <= end with end < 0: x = -a where a >= |end| > 0, so negate every + // pattern for the range [|end|, +infinity). + for (auto& p : AtLeastPositivePatternsStr(AbsDigits(end.value()))) { + parts.push_back("-" + p); + } + } + } + + if (start && end) { + int64_t range_start = start.value(); + int64_t range_end = end.value(); + + if (range_start > range_end) { + return "^()$"; + } + + if (range_start < 0) { + int64_t neg_start = range_start; + int64_t neg_end = std::min(static_cast(-1), range_end); + // Negatives in [neg_start, neg_end] are the magnitudes + // [|neg_end|, |neg_start|], negated. + parts.push_back("-" + SubRangeRegexStr(AbsDigits(neg_end), AbsDigits(neg_start))); + } + + if (range_start <= 0 && range_end >= 0) { + parts.push_back("0"); + } + + if (range_end > 0) { + int64_t pos_start = std::max(static_cast(1), range_start); + parts.push_back(SubRangeRegexStr(std::to_string(pos_start), std::to_string(range_end))); + } + } + + result << "^("; + for (size_t i = 0; i < parts.size(); ++i) { + if (i > 0) { + result << "|"; + } + result << parts[i]; + } + result << ")$"; + + return result.str(); +} + +std::string NumberGenerator::FormatFloat(double value, int precision) { + // Casting a double outside [INT64_MIN, INT64_MAX] (or NaN/Inf) to int64_t is + // undefined behavior, so range-check before the integer fast path. 2^63 == + // 9223372036854775808.0 is exactly representable and one past INT64_MAX, so the + // upper comparison must be strict. + if (value >= -9223372036854775808.0 && value < 9223372036854775808.0 && + value == static_cast(value)) { + return std::to_string(static_cast(value)); + } + + std::ostringstream oss; + oss << std::fixed << std::setprecision(precision) << value; + std::string result = oss.str(); + + size_t decimalPos = result.find('.'); + if (decimalPos != std::string::npos) { + size_t lastNonZero = result.find_last_not_of('0'); + if (lastNonZero != std::string::npos && lastNonZero > decimalPos) { + result.erase(lastNonZero + 1); + } else if (lastNonZero == decimalPos) { + result.erase(decimalPos); + } + } + + return result; +} + +std::string NumberGenerator::AdjustGrid(const std::string& s, int precision, bool inc) { + std::string int_part, frac_part; + SplitDecimal(s, &int_part, &frac_part); + // Build the scaled-integer numerator (value * 10^precision) as a digit string. + // Callers only pass FormatFloat output (<= precision fraction digits); guard + // the count so a longer string can never wrap the unsigned append count. + frac_part.append(std::max(0, precision - static_cast(frac_part.size())), '0'); + std::string num = int_part + frac_part; + + if (inc) { + int i = static_cast(num.size()) - 1; + for (; i >= 0 && num[i] == '9'; --i) { + num[i] = '0'; + } + if (i < 0) { + num.insert(num.begin(), '1'); + } else { + num[i]++; + } + } else { + int i = static_cast(num.size()) - 1; + for (; i >= 0 && num[i] == '0'; --i) { + num[i] = '9'; + } + if (i < 0) { + // Underflow below zero; clamp to zero (does not occur for the bounds the + // float pipeline feeds in, which are all >= one grid step when decremented). + num.assign(num.size(), '0'); + } else { + num[i]--; + } + } + + // Re-split into integer and `precision`-digit fraction, then canonicalize. + while (static_cast(num.size()) <= precision) { + num.insert(num.begin(), '0'); + } + std::string new_int = num.substr(0, num.size() - precision); + std::string new_frac = num.substr(num.size() - precision); + size_t nz = new_int.find_first_not_of('0'); + new_int = (nz == std::string::npos) ? "0" : new_int.substr(nz); + size_t lnz = new_frac.find_last_not_of('0'); + new_frac = (lnz == std::string::npos) ? "" : new_frac.substr(0, lnz + 1); + return new_frac.empty() ? new_int : new_int + "." + new_frac; +} + +std::string NumberGenerator::RoundBoundToGrid( + double value, int precision, bool is_lower, bool strict_in, bool* strict_out +) { + // FormatFloat rounds to the nearest grid point; if that lands exactly on the + // bound, keep the original strictness. Otherwise step to the grid point just + // inside the range so no out-of-range value is admitted, and the boundary is + // now strictly interior, so it becomes inclusive. + std::string r = FormatFloat(value, precision); + double rv = std::stod(r); + if (rv == value) { + *strict_out = strict_in; + return r; + } + *strict_out = false; + if (is_lower && rv < value) { + // Rounded below a lower bound: move up to the smallest grid point >= value. + r = AdjustGrid(r, precision, /*inc=*/true); + } else if (!is_lower && rv > value) { + // Rounded above an upper bound: move down to the largest grid point <= value. + r = AdjustGrid(r, precision, /*inc=*/false); + } + return r; +} + +// Helpers for GenerateFloatRangeRegex. Fraction patterns operate on the +// digit string after the decimal point, compared against a canonical bound +// fraction (canonical: produced by FormatFloat, so no trailing zeros). + +// Matches 0 to max_count free digits. +std::string NumberGenerator::FreeDigits(int max_count) { + if (max_count <= 0) { + return ""; + } + return "\\d{0," + std::to_string(max_count) + "}"; +} + +// Matches 0 to max_count zeros. +std::string NumberGenerator::OptionalZeros(int max_count) { + if (max_count <= 0) { + return ""; + } + return "0{0," + std::to_string(max_count) + "}"; +} + +// Matches 1 to max_count zeros. +std::string NumberGenerator::SomeZeros(int max_count) { + return "0{1," + std::to_string(max_count) + "}"; +} + +// Patterns for fraction strings t (1 <= |t| <= max_len) whose value 0.t is +// greater than 0.s (or equal when !strict). |s| <= max_len. +NumberGenerator::FracPatternSet NumberGenerator::FracGreaterPatterns( + const std::string& s, bool strict, int max_len +) { + FracPatternSet result; + int n = static_cast(s.size()); + // t agrees with s up to position i, then has a larger digit + for (int i = 0; i < n; ++i) { + if (s[i] < '9') { + result.parts.push_back( + s.substr(0, i) + DigitClass(s[i] + 1, '9') + FreeDigits(max_len - i - 1) + ); + } + } + // t extends s with a nonzero digit (after optional zeros) + for (int k = 0; n + k + 1 <= max_len; ++k) { + result.parts.push_back(s + std::string(k, '0') + "[1-9]" + FreeDigits(max_len - n - k - 1)); + } + if (!strict) { + // t has the same value as s: s plus optional trailing zeros + if (n > 0) { + result.parts.push_back(s + OptionalZeros(max_len - n)); + } else { + result.include_empty = true; + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + } + return result; +} + +// Patterns for fraction strings t (1 <= |t| <= max_len) whose value 0.t is +// less than 0.s (or equal when !strict). |s| <= max_len. +NumberGenerator::FracPatternSet NumberGenerator::FracLessPatterns( + const std::string& s, bool strict, int max_len +) { + FracPatternSet result; + int n = static_cast(s.size()); + // t agrees with s up to position i, then has a smaller digit + for (int i = 0; i < n; ++i) { + if (s[i] > '0') { + result.parts.push_back( + s.substr(0, i) + DigitClass('0', s[i] - 1) + FreeDigits(max_len - i - 1) + ); + } + } + // t is a proper prefix of s plus optional trailing zeros: strictly smaller, + // since the remaining digits of s contain a nonzero one + for (int i = 0; i < n; ++i) { + if (i == 0) { + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } else { + result.parts.push_back(s.substr(0, i) + OptionalZeros(max_len - i)); + } + } + if (!strict) { + // t has the same value as s + if (n > 0) { + result.parts.push_back(s + OptionalZeros(max_len - n)); + } else if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + result.include_empty = n > 0 || !strict; + return result; +} + +// Patterns for fraction strings t whose value 0.t lies between 0.a and 0.b. +// Requires value(0.a) < value(0.b) and b non-empty. +NumberGenerator::FracPatternSet NumberGenerator::FracBetweenPatterns( + const std::string& a, bool strict_a, const std::string& b, bool strict_b, int max_len +) { + FracPatternSet result; + // Longest common prefix of b and zero-padded a. Always stops before |b|: + // value(0.a) < value(0.b) implies b is not a prefix of padded a. + int common_len = 0; + while (common_len < static_cast(b.size()) && + (common_len < static_cast(a.size()) ? a[common_len] : '0') == b[common_len]) { + ++common_len; + } + std::string common = b.substr(0, common_len); + char digit_a = common_len < static_cast(a.size()) ? a[common_len] : '0'; + char digit_b = b[common_len]; + + // a digit strictly between the bounds' digits, then anything + if (digit_b - digit_a >= 2) { + result.parts.push_back( + common + DigitClass(digit_a + 1, digit_b - 1) + FreeDigits(max_len - common_len - 1) + ); + } + // lower boundary: t continues with digit_a, the rest must exceed a's suffix + if (common_len < static_cast(a.size())) { + FracPatternSet sub_lower = + FracGreaterPatterns(a.substr(common_len + 1), strict_a, max_len - common_len - 1); + for (auto& part : sub_lower.parts) { + result.parts.push_back(common + digit_a + std::move(part)); + } + if (sub_lower.include_empty) { + result.parts.push_back(common + std::string(1, digit_a)); + } + } else { + // a's value equals value(0.common): only nonzero extensions of + // common + digit_a ('0') are strictly greater + FracPatternSet sub_lower = FracGreaterPatterns("", true, max_len - common_len - 1); + for (auto& part : sub_lower.parts) { + result.parts.push_back(common + digit_a + std::move(part)); + } + if (!strict_a) { + // t has the same value as a + if (!a.empty()) { + result.parts.push_back(a + OptionalZeros(max_len - static_cast(a.size()))); + } else { + result.include_empty = true; + if (max_len >= 1) { + result.parts.push_back(SomeZeros(max_len)); + } + } + } + } + // upper boundary: t continues with digit_b, the rest must stay below b's suffix + FracPatternSet sub_upper = + FracLessPatterns(b.substr(common_len + 1), strict_b, max_len - common_len - 1); + for (auto& part : sub_upper.parts) { + result.parts.push_back(common + digit_b + std::move(part)); + } + if (sub_upper.include_empty) { + result.parts.push_back(common + std::string(1, digit_b)); + } + return result; +} + +// Splits a canonical decimal string from FormatFloat ("12" or "12.34") into +// integer and fraction parts. +void NumberGenerator::SplitDecimal( + const std::string& s, std::string* int_part, std::string* frac_part +) { + size_t dot = s.find('.'); + if (dot == std::string::npos) { + *int_part = s; + frac_part->clear(); + } else { + *int_part = s.substr(0, dot); + *frac_part = s.substr(dot + 1); + } +} + +// Compares the values of two canonical non-negative decimals. +int NumberGenerator::CompareDecimal( + const std::string& int_a, + const std::string& frac_a, + const std::string& int_b, + const std::string& frac_b +) { + if (int_a.size() != int_b.size()) { + return int_a.size() < int_b.size() ? -1 : 1; + } + if (int_a != int_b) { + return int_a < int_b ? -1 : 1; + } + size_t max_frac = std::max(frac_a.size(), frac_b.size()); + for (size_t i = 0; i < max_frac; ++i) { + char da = i < frac_a.size() ? frac_a[i] : '0'; + char db = i < frac_b.size() ? frac_b[i] : '0'; + if (da != db) { + return da < db ? -1 : 1; + } + } + return 0; +} + +// Strips the ^( )$ anchors added by IntegerRangeRegex, keeping the group. +std::string NumberGenerator::StripAnchors(const std::string& regex) { + return regex.substr(1, regex.size() - 2); +} + +int64_t NumberGenerator::ParseIntCapped(const std::string& digits) { + // `digits` is a canonical non-negative integer string (no leading zeros). + // Parse it exactly when it fits in int64; clamp to INT64_MAX otherwise (such + // magnitudes are beyond practical float bounds and double integer precision). + static const std::string kMaxInt64 = std::to_string(std::numeric_limits::max()); + if (digits.size() > kMaxInt64.size() || + (digits.size() == kMaxInt64.size() && digits > kMaxInt64)) { + return std::numeric_limits::max(); + } + return std::stoll(digits); +} + +// Patterns for unsigned decimals (integer part plus optional fraction of up +// to `precision` digits) within the given bounds. `low` is required and +// non-negative; `high` is optional. Patterns for the value 0 are never +// produced: when low's value is 0 the bound is treated as strict, and the +// caller emits the zero pattern itself. +std::vector NumberGenerator::PositiveRangeParts( + const std::string& low, + bool strict_low, + const std::optional& high, + bool strict_high, + int precision +) { + std::vector parts; + std::string int_low, frac_low; + SplitDecimal(low, &int_low, &frac_low); + if (int_low == "0" && frac_low.empty()) { + strict_low = true; + } + int64_t int_low_value = ParseIntCapped(int_low); + std::string opt_any_frac = "(\\.\\d{1," + std::to_string(precision) + "})?"; + + auto add_with_int_part = [&](const std::string& int_part, const FracPatternSet& set) { + for (const auto& part : set.parts) { + parts.push_back(int_part + "\\." + part); + } + if (set.include_empty) { + parts.push_back(int_part); + } + }; + + if (!high.has_value()) { + add_with_int_part(int_low, FracGreaterPatterns(frac_low, strict_low, precision)); + // Guard the +1 against int64 overflow (int_low_value may be clamped to + // INT64_MAX for very large bounds). + if (int_low_value < std::numeric_limits::max()) { + parts.push_back( + StripAnchors(IntegerRangeRegex(int_low_value + 1, std::nullopt)) + opt_any_frac + ); + } + return parts; + } + + std::string int_high, frac_high; + SplitDecimal(*high, &int_high, &frac_high); + int64_t int_high_value = ParseIntCapped(int_high); + int cmp = CompareDecimal(int_low, frac_low, int_high, frac_high); + if (cmp > 0 || (cmp == 0 && (strict_low || strict_high))) { + return parts; + } + if (cmp == 0) { + // single representable value, with optional redundant trailing zeros + if (frac_low.empty()) { + parts.push_back(int_low + "(\\." + SomeZeros(precision) + ")?"); + } else { + parts.push_back( + int_low + "\\." + frac_low + OptionalZeros(precision - static_cast(frac_low.size())) + ); + } + return parts; + } + if (int_low == int_high) { + add_with_int_part( + int_low, FracBetweenPatterns(frac_low, strict_low, frac_high, strict_high, precision) + ); + } else { + add_with_int_part(int_low, FracGreaterPatterns(frac_low, strict_low, precision)); + if (int_high_value - int_low_value >= 2) { + parts.push_back( + StripAnchors(IntegerRangeRegex(int_low_value + 1, int_high_value - 1)) + opt_any_frac + ); + } + add_with_int_part(int_high, FracLessPatterns(frac_high, strict_high, precision)); + } + return parts; +} + +std::string NumberGenerator::FloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end +) { + if (start && end) { + if (start.value() > end.value() || + (start.value() == end.value() && (exclusive_start || exclusive_end))) { + return "^()$"; + } + } + + if (!start && !end) { + return "^-?\\d+(\\.\\d{1," + std::to_string(precision) + "})?$"; + } + + std::vector parts; + + // Negative values: x is in [start, end] iff -x is in [-end, -start], so the + // positive-range patterns are reused on the negated bounds and prefixed + // with '-'. + bool negatives_in_range = !start.has_value() || start.value() < 0; + if (negatives_in_range) { + std::string low = "0"; + bool strict_low = true; + if (end.has_value() && end.value() < 0) { + low = + RoundBoundToGrid(-end.value(), precision, /*is_lower=*/true, exclusive_end, &strict_low); + } + std::optional high; + bool strict_high = false; + if (start.has_value()) { + high = RoundBoundToGrid( + -start.value(), precision, /*is_lower=*/false, exclusive_start, &strict_high + ); + } + for (auto& part : PositiveRangeParts(low, strict_low, high, strict_high, precision)) { + parts.push_back("-" + std::move(part)); + } + } + + bool zero_allowed = + (!start.has_value() || start.value() < 0 || (start.value() == 0 && !exclusive_start)) && + (!end.has_value() || end.value() > 0 || (end.value() == 0 && !exclusive_end)); + if (zero_allowed) { + parts.push_back("0(\\." + SomeZeros(precision) + ")?"); + // Negative zero written with an all-zero fraction ("-0.0".."-0.000000") also + // denotes 0. PositiveRangeParts never emits magnitude 0, so add these forms + // explicitly when the range covers the negative side. + if (negatives_in_range) { + parts.push_back("-0(\\." + SomeZeros(precision) + ")"); + } + } + + // Positive values + if (!end.has_value() || end.value() > 0) { + std::string low = "0"; + bool strict_low = true; + if (start.has_value() && start.value() > 0) { + low = RoundBoundToGrid( + start.value(), precision, /*is_lower=*/true, exclusive_start, &strict_low + ); + } + std::optional high; + bool strict_high = false; + if (end.has_value()) { + high = + RoundBoundToGrid(end.value(), precision, /*is_lower=*/false, exclusive_end, &strict_high); + } + for (auto& part : PositiveRangeParts(low, strict_low, high, strict_high, precision)) { + parts.push_back(std::move(part)); + } + } + + std::ostringstream result; + result << "^("; + for (size_t i = 0; i < parts.size(); ++i) { + if (i > 0) { + result << "|"; + } + result << parts[i]; + } + result << ")$"; + + return result.str(); +} + +std::string JSONSchemaConverter::GenerateRangeRegex( + std::optional start, std::optional end +) { + return NumberGenerator::IntegerRangeRegex(start, end); +} + +std::string JSONSchemaConverter::GenerateFloatRangeRegex( + std::optional start, + std::optional end, + int precision, + bool exclusive_start, + bool exclusive_end +) { + return NumberGenerator::FloatRangeRegex(start, end, precision, exclusive_start, exclusive_end); +} + +// ==================== Public API Functions ==================== + +std::optional JSONFormatFromString(const std::string& format) { + static const std::unordered_map kNameToFormat = { + {"json", JSONFormat::kJSON}, + {"qwen_xml", JSONFormat::kQwenXML}, + {"minimax_xml", JSONFormat::kMiniMaxXML}, + {"minimax_m3_xml", JSONFormat::kMiniMaxM3XML}, + {"deepseek_xml", JSONFormat::kDeepSeekXML}, + {"deepseek_v4_1_xml", JSONFormat::kDeepSeekV41XML}, + {"glm_xml", JSONFormat::kGlmXML}, + {"cohere_xml", JSONFormat::kCohereXML}, + {"kimi_k3_xml", JSONFormat::kKimiK3XML}, + }; + auto it = kNameToFormat.find(format); + if (it == kNameToFormat.end()) { + return std::nullopt; + } + return it->second; +} + +Grammar JSONSchemaToGrammar( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + bool any_order, + JSONFormat json_format, + std::vector excludes +) { + picojson::value schema_value; + std::string error = ParseJSON(schema_value, schema); + XGRAMMAR_CHECK(error.empty()) << "Failed to parse JSON: " << error + << ". The JSON string is:" << schema; + SchemaParser parser(schema_value, {strict_mode, json_format}); + auto spec_result = parser.Parse(schema_value, "root"); + if (spec_result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(spec_result).UnwrapErr().what(); + } + auto spec = std::move(spec_result).Unwrap(); + auto ref_resolver = [&parser](const std::string& uri, const std::string& rule_name_hint) { + auto result = parser.ResolveRef(uri, rule_name_hint); + if (result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(result).UnwrapErr().what(); + } + return std::move(result).Unwrap(); + }; + + switch (json_format) { + case JSONFormat::kJSON: { + JSONSchemaConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order, + std::move(excludes) + ); + return converter.Convert(spec); + } + case JSONFormat::kQwenXML: + case JSONFormat::kMiniMaxXML: + case JSONFormat::kDeepSeekXML: + case JSONFormat::kDeepSeekV41XML: + case JSONFormat::kGlmXML: + case JSONFormat::kKimiK3XML: { + XMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + json_format, + any_order, + std::move(excludes) + ); + return converter.Convert(spec); + } + case JSONFormat::kMiniMaxM3XML: { + XGRAMMAR_CHECK(excludes.empty()) + << "JSONSchemaFormat.excludes is not supported for minimax_m3_xml"; + MiniMaxM3XMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order + ); + return converter.Convert(spec); + } + case JSONFormat::kCohereXML: { + XGRAMMAR_CHECK(excludes.empty()) + << "JSONSchemaFormat.excludes is not supported for cohere_xml"; + CohereXMLToolCallingConverter converter( + indent, + std::move(separators), + any_whitespace, + max_whitespace_cnt, + std::move(ref_resolver), + any_order + ); + return converter.Convert(spec); + } + default: + XGRAMMAR_LOG(FATAL) << "Invalid JSON format: " << static_cast(json_format); + } + XGRAMMAR_UNREACHABLE(); +} + +std::string JSONSchemaToEBNF( + const std::string& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + JSONFormat json_format, + bool any_order +) { + picojson::value schema_value; + std::string err = ParseJSON(schema_value, schema); + XGRAMMAR_CHECK(err.empty()) << "Failed to parse JSON: " << err + << ". The JSON string is:" << schema; + return JSONSchemaToEBNF( + schema_value, + any_whitespace, + indent, + separators, + strict_mode, + max_whitespace_cnt, + json_format, + any_order + ); +} + +std::string JSONSchemaToEBNF( + const picojson::value& schema, + bool any_whitespace, + std::optional indent, + std::optional> separators, + bool strict_mode, + std::optional max_whitespace_cnt, + JSONFormat json_format, + bool any_order +) { + // Parse JSON Schema to SchemaSpec + SchemaParser parser(schema, {strict_mode, json_format}); + auto spec_result = parser.Parse(schema, "root"); + if (spec_result.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(spec_result).UnwrapErr().what(); + } + auto spec = std::move(spec_result).Unwrap(); + + auto ref_resolver = [&parser](const std::string& uri, const std::string& rule_name_hint) { + auto r = parser.ResolveRef(uri, rule_name_hint); + if (r.IsErr()) { + XGRAMMAR_LOG(FATAL) << std::move(r).UnwrapErr().what(); + } + return std::move(r).Unwrap(); + }; + + // Create converter based on format + switch (json_format) { + case JSONFormat::kJSON: { + JSONSchemaConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kQwenXML: + case JSONFormat::kMiniMaxXML: + case JSONFormat::kDeepSeekXML: + case JSONFormat::kDeepSeekV41XML: + case JSONFormat::kGlmXML: + case JSONFormat::kKimiK3XML: { + XMLToolCallingConverter converter( + indent, + separators, + any_whitespace, + max_whitespace_cnt, + ref_resolver, + json_format, + any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kMiniMaxM3XML: { + MiniMaxM3XMLToolCallingConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + case JSONFormat::kCohereXML: { + CohereXMLToolCallingConverter converter( + indent, separators, any_whitespace, max_whitespace_cnt, ref_resolver, any_order + ); + return GrammarNormalizer::Apply(converter.Convert(spec)).ToString(); + } + default: + XGRAMMAR_LOG(FATAL) << "Invalid JSON format: " << static_cast(json_format); + } + XGRAMMAR_UNREACHABLE(); +} + +// Wrapper functions for testing +std::string GenerateRangeRegex(std::optional start, std::optional end) { + return JSONSchemaConverter::GenerateRangeRegex(start, end); +} + +std::string GenerateFloatRangeRegex( + std::optional start, std::optional end, bool exclusive_start, bool exclusive_end +) { + return JSONSchemaConverter::GenerateFloatRangeRegex( + start, end, 6, exclusive_start, exclusive_end + ); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/json_schema_converter.h b/zig/vendor/xgrammar/cpp/json_schema_converter.h new file mode 100644 index 0000000000..8ca48963af --- /dev/null +++ b/zig/vendor/xgrammar/cpp/json_schema_converter.h @@ -0,0 +1,666 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter.h + * \brief Convert a JSON Schema directly to a grammar AST. + */ + +#ifndef XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ +#define XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_builder.h" +#include "support/utils.h" + +namespace xgrammar { + +// ==================== SchemaSpec: Intermediate Representation for JSON Schema ==================== + +// Forward declaration +struct SchemaSpec; +using SchemaSpecPtr = std::shared_ptr; + +// Basic Type Specs +struct IntegerSpec { + std::optional minimum; + std::optional maximum; + std::optional exclusive_minimum; + std::optional exclusive_maximum; + std::optional multiple_of; + + std::string ToString() const; +}; + +struct NumberSpec { + std::optional minimum; + std::optional maximum; + std::optional exclusive_minimum; + std::optional exclusive_maximum; + + std::string ToString() const; +}; + +struct StringSpec { + std::optional pattern; + std::optional format; + int min_length = 0; + int max_length = -1; // -1 means no limit + + std::string ToString() const; +}; + +struct BooleanSpec { + std::string ToString() const; +}; + +struct NullSpec { + std::string ToString() const; +}; + +struct AnySpec { + std::string ToString() const; +}; + +// Complex Type Specs +struct ArraySpec { + std::vector prefix_items; + bool allow_additional_items = true; + SchemaSpecPtr additional_items; // nullptr means not allowed + int64_t min_items = 0; + int64_t max_items = -1; // -1 means no limit + + std::string ToString() const; +}; + +struct ObjectSpec { + struct Property { + std::string name; + SchemaSpecPtr schema; + }; + + struct PatternProperty { + std::string pattern; // regex pattern for key + SchemaSpecPtr schema; + }; + + std::vector properties; + std::vector pattern_properties; + std::unordered_set required; + + bool allow_additional_properties = false; + SchemaSpecPtr additional_properties_schema; + bool allow_unevaluated_properties = true; + SchemaSpecPtr unevaluated_properties_schema; + SchemaSpecPtr property_names; + + int min_properties = 0; + int max_properties = -1; // -1 means no limit + + std::string ToString() const; +}; + +// Composite Type Specs +struct ConstSpec { + std::string json_value; // JSON serialized value + + std::string ToString() const; +}; + +struct EnumSpec { + std::vector json_values; // JSON serialized values + + std::string ToString() const; +}; + +struct RefSpec { + std::string uri; + + std::string ToString() const; +}; + +struct AnyOfSpec { + std::vector options; + + std::string ToString() const; +}; + +struct OneOfSpec { + std::vector options; + + std::string ToString() const; +}; + +struct AllOfSpec { + std::vector schemas; + + std::string ToString() const; +}; + +struct TypeArraySpec { + // Handle "type": ["string", "integer"] cases + std::vector type_schemas; + + std::string ToString() const; +}; + +// Unified SchemaSpec +using SchemaSpecVariant = std::variant< + IntegerSpec, + NumberSpec, + StringSpec, + BooleanSpec, + NullSpec, + ArraySpec, + ObjectSpec, + AnySpec, + ConstSpec, + EnumSpec, + RefSpec, + AnyOfSpec, + OneOfSpec, + AllOfSpec, + TypeArraySpec>; + +struct SchemaSpec { + SchemaSpecVariant spec; + std::string cache_key; // for deduplication + std::string rule_name_hint; // suggested rule name + + std::string ToString() const; + + // Helper method to create SchemaSpec + template + static SchemaSpecPtr Make(T&& spec_value, std::string cache_key = "", std::string hint = "") { + auto ptr = std::make_shared(); + ptr->spec = std::forward(spec_value); + ptr->cache_key = std::move(cache_key); + ptr->rule_name_hint = std::move(hint); + return ptr; + } +}; + +// ==================== JSONFormat Enum ==================== + +enum class JSONFormat : int { + kJSON = 0, + kQwenXML = 1, + kMiniMaxXML = 2, + kDeepSeekXML = 3, + kGlmXML = 4, + kCohereXML = 5, + kKimiK3XML = 6, + kMiniMaxM3XML = 7, + kDeepSeekV41XML = 8, +}; + +/*! + * \brief Convert a format name to JSONFormat. + * \param format One of "json", "qwen_xml", "minimax_xml", "minimax_m3_xml", "deepseek_xml", + * "glm_xml", "cohere_xml", "kimi_k3_xml", or "deepseek_v4_1_xml". + * \return The corresponding JSONFormat, or std::nullopt if the name is not recognized. + */ +std::optional JSONFormatFromString(const std::string& format); + +/*! + * \brief Manage the rule generation cache. Wraps key-value cache for schema deduplication. + * The cached value is the rule id in the grammar builder. + */ +class GenerateCacheManager { + public: + /*! \brief Add a key-value pair to the cache. */ + void AddCache(const std::string& key, int context, int32_t rule_id) { + cache_[{key, context}] = rule_id; + } + + /*! \brief Get cached rule id by key. Returns std::nullopt if not found. */ + std::optional GetCache(const std::string& key, int context) const { + auto it = cache_.find({key, context}); + if (it != cache_.end()) { + return it->second; + } + return std::nullopt; + } + + private: + std::unordered_map, int32_t> cache_; +}; + +/*! + * \brief Manage the indent and separator for the generation of EBNF grammar. + */ +class IndentManager { + public: + IndentManager( + std::optional indent, + const std::string& separator, + bool any_whitespace, + std::optional max_whitespace_cnt + ); + + void StartIndent(); + void EndIndent(); + std::string StartSeparator(); + std::string MiddleSeparator(); + std::string EndSeparator(); + std::string EmptySeparator(); + std::string NextSeparator(bool is_end = false); + + private: + bool any_whitespace_; + bool enable_newline_; + int64_t indent_; + std::string separator_; + int64_t total_indent_; + std::vector is_first_; + std::optional max_whitespace_cnt_; + + friend class JSONSchemaConverter; +}; + +/*! + * \brief Convert SchemaSpec directly to a grammar AST. + * + * This is the base class for grammar generation. It generates JSON-format grammar by default. + * Subclasses can override virtual methods to generate different formats (e.g., XML). + */ +class JSONSchemaConverter { + public: + using RefResolver = + std::function; + + JSONSchemaConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false, + std::vector excludes = {} + ); + + virtual ~JSONSchemaConverter() = default; + + /*! + * \brief Convert SchemaSpec directly to a grammar AST. + * \param spec The SchemaSpec to convert. + * \return The grammar AST. + */ + Grammar Convert(const SchemaSpecPtr& spec); + + /*! \brief Whether \p format is compiled to a regex, which shadows minLength/maxLength. */ + static bool IsBuiltinFormat(const std::string& format) { + return JSONFormatToRegexPattern(format).has_value(); + } + + protected: + using CharacterClassElement = GrammarBuilder::CharacterClassElement; + + // ==================== Virtual methods for generation ==================== + // Subclasses can override these to customize output format + + virtual int32_t GenerateInteger(const IntegerSpec& spec, const std::string& rule_name); + virtual int32_t GenerateNumber(const NumberSpec& spec, const std::string& rule_name); + virtual int32_t GenerateString(const StringSpec& spec, const std::string& rule_name); + virtual int32_t GenerateBoolean(const BooleanSpec& spec, const std::string& rule_name); + virtual int32_t GenerateNull(const NullSpec& spec, const std::string& rule_name); + virtual int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name); + virtual int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool need_brace = true + ); + virtual int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name); + virtual int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name); + virtual int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name); + virtual int32_t GenerateRef(const RefSpec& spec, const std::string& rule_name); + /*! \brief Key reference rules by their output encoding context. */ + virtual std::string RefCacheKey(const std::string& uri) const; + virtual int32_t GenerateAnyOf(const AnyOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateOneOf(const OneOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateAllOf(const AllOfSpec& spec, const std::string& rule_name); + virtual int32_t GenerateTypeArray(const TypeArraySpec& spec, const std::string& rule_name); + + // ==================== Hooks for customization ==================== + + /*! + * \brief Format a property key. Override for different formats. + * \param schema The schema of the property's value. Formats that encode the value's type + * next to the key (e.g. the Kimi-K3 `type` attribute) need it; the JSON format ignores it. + */ + virtual int32_t FormatPropertyKey(const std::string& key, const SchemaSpecPtr& schema); + + /*! \brief Format a property (key + value). Override for different formats. */ + virtual int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ); + + /*! + * \brief Format an "other" property (additional/unevaluated). Override for different formats. + * \param schema The schema selected for the property's value, or nullptr when the dynamic key + * has no single schema. Formats that encode the value's type next to the key need it. + */ + virtual int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ); + + /*! \brief Get the basic string rule name. Override for different formats. */ + virtual std::string GetKeyPattern() const; + + /*! \brief Get a key pattern that excludes specific property names. */ + virtual int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ); + + /*! \brief Get the basic any rule name. Override for different formats. */ + virtual std::string GetBasicAnyRuleName() const; + + /*! \brief Add basic rules for the format. Override for different formats. */ + virtual void AddBasicRules(); + void AddBasicRules(const std::vector& additional_rule_names); + + /*! \brief Add a key-value pair to the generation cache. Override for custom cache behavior. */ + virtual void AddCache(const std::string& key, int32_t rule_id); + + /*! \brief Get cached value by key. Returns std::nullopt if not found. */ + virtual std::optional GetCache(const std::string& key) const; + + // ==================== Helper methods (for subclasses to use) ==================== + + /*! \brief Dispatch to the appropriate Generate method based on spec type. */ + int32_t GenerateFromSpec(const SchemaSpecPtr& spec, const std::string& rule_name_hint); + + /*! \brief Create a rule and return the rule id (handles caching). */ + int32_t CreateRule(const SchemaSpecPtr& spec, const std::string& rule_name_hint); + + /*! \brief Resolve a reference to its parsed schema. */ + SchemaSpecPtr ResolveRefSchema(const RefSpec& spec, const std::string& rule_name_hint); + + /*! \brief Get next separator from indent manager. */ + virtual std::string NextSeparator(bool is_end = false); + + /*! \brief Get whitespace pattern. */ + std::string GetWhitespacePattern() const; + + int32_t Empty(); + /*! \brief An expression that matches nothing, for alternatives ruled out by excludes_. */ + int32_t Unsatisfiable(); + int32_t ByteString(const std::string& value); + int32_t TagDispatch(bool loop_after_dispatch, std::vector excludes); + int32_t RuleRef(int32_t rule_id); + int32_t RuleRef(const std::string& rule_name); + int32_t Sequence(const std::vector& elements); + int32_t Choice(const std::vector& choices); + int32_t Repeat( + const std::string& rule_name_hint, int32_t expr_id, int32_t min_count, int32_t max_count + ); + int32_t AddSubGrammar(const Grammar& grammar); + + int32_t WhitespaceExpression(); + int32_t FormattingExpression(const std::string& expression); + int32_t NextSeparatorExpression(bool is_end = false); + int32_t KeyPatternExpression(); + + int32_t RegexExpression( + const std::string& regex, bool json_string = false, bool force_cfg_expansion = false + ); + /*! + * \brief Rules matching the regex (one of the converter's own ASCII string bodies) minus + * every string containing one of excludes_, and minus excluded_keys as whole strings. With + * close_json_string the closing quote is appended after the filtering, so it is never part of + * an exclusion. + */ + int32_t ExcludingString( + const std::string& regex, + const std::string& rule_name, + bool close_json_string, + const std::vector& excluded_keys = {} + ); + bool IsAllowedString(const std::string& text) const; + bool IsAllowedLiteral(const picojson::value& value, bool raw_string = false) const; + /*! \brief Whether the JSON literal contains none of excludes_ (see IsAllowedLiteral). */ + bool IsAllowedJSONLiteral(const std::string& json_value, bool raw_string = false) const; + /*! \brief Log that the string's minLength/maxLength are ignored because excludes_ is set. */ + void WarnDroppedLengthConstraints(const StringSpec& spec, const std::string& rule_name) const; + + /*! \brief Helper to create rule with repetition constraints. */ + int32_t GetPropertyWithNumberConstraints( + int32_t pattern, + int min_properties, + int max_properties, + int already_repeated_times, + const std::string& rule_name + ); + + /*! \brief Generate partial rule for object properties. + * \param additional_property_override When set, used as the additional property + * pattern instead of the default GetKeyPattern() : value. This supports patternProperties + * and propertyNames constraints on additional keys. + */ + int32_t GetPartialRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override = std::nullopt + ); + + /*! \brief Generate the object rule in "any order" mode: an "item" alternation over all property + * keys, repeated between max(min_properties, required.size()) and max_properties times. Only the + * entry count is bounded, not which keys appear. + */ + int32_t GetAnyOrderRuleForProperties( + const std::vector& properties, + const std::unordered_set& required, + const SchemaSpecPtr& additional, + const std::string& rule_name, + const std::string& additional_suffix, + int min_properties, + int max_properties, + const std::optional& additional_property_override = std::nullopt + ); + + // ==================== Protected members ==================== + + GrammarBuilder builder_; + IndentManager indent_manager_; + int32_t colon_expr_id_; + bool any_whitespace_; + std::optional max_whitespace_cnt_; + // When true, object properties may appear in any order (see GetAnyOrderRuleForProperties). + // Applies to all objects (including nested ones). Default false preserves the fixed-order + // behavior. + bool any_order_ = false; + std::vector excludes_; + + public: + // Basic rule names + static const std::string kBasicAny; + static const std::string kBasicInteger; + static const std::string kBasicNumber; + static const std::string kBasicString; + static const std::string kBasicBoolean; + static const std::string kBasicNull; + static const std::string kBasicArray; + static const std::string kBasicObject; + static const std::string kBasicEscape; + static const std::string kBasicStringSub; + + protected: + GenerateCacheManager rule_cache_manager_; + + private: + void AddHelperRules(); + + std::unordered_map uri_to_rule_id_; // For circular reference handling + RefResolver ref_resolver_; // Resolves $ref URI to SchemaSpecPtr at generate time + + // Trie over property names, for key patterns that exclude specific properties + struct TrieNode { + bool is_terminal = false; + std::map children; + }; + int32_t BuildTrieBody(const TrieNode& node, const std::string& rule_name); + + // Reused grammar expression ids + std::optional empty_expr_id_; + std::optional unsatisfiable_expr_id_; + std::unordered_map byte_string_expr_ids_; + std::unordered_map rule_ref_expr_ids_; + std::optional whitespace_expr_id_; + + // Helper for integer/number range regex generation + static std::string GenerateRangeRegex(std::optional start, std::optional end); + int32_t GenerateIntegerMultipleOfDFA(int64_t multiple_of, const std::string& rule_name); + static std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + int precision = 6, + bool exclusive_start = false, + bool exclusive_end = false + ); + + protected: + static std::optional JSONFormatToRegexPattern(const std::string& format); + + // Expose for testing + friend std::string GenerateRangeRegex(std::optional start, std::optional end); + friend std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + bool exclusive_start, + bool exclusive_end + ); +}; + +/*! + * \brief Convert a JSON Schema string directly to an unnormalized grammar AST. + * + * Callers that need normalized grammar should apply GrammarNormalizer after composing any + * subgrammars. + */ +Grammar JSONSchemaToGrammar( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool any_order = false, + JSONFormat json_format = JSONFormat::kJSON, + std::vector excludes = {} +); + +// ==================== Public API functions (backward compatible) ==================== + +/*! + * \brief Convert JSON schema string to EBNF grammar string. + * \param schema The JSON schema string. + * \param any_whitespace Whether to ignore the indentation restrictions, and allow any whitespace. + * Default: true. + * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. If std::nullopt, the default separators will be used: {",", ": "} when the + * indent is not -1, and {", ", ": "} otherwise. This follows the convention in python + * json.dumps(). Default: std::nullopt. + * \param strict_mode Whether to use strict mode. In strict + * mode, the generated grammar will not allow properties and items that is not specified in the + * schema. This is equivalent to setting unevaluatedProperties and unevaluatedItems to false. + * This helps LLM to generate accurate output in the grammar-guided generation with JSON + * schema. Default: true. + * \param max_whitespace_cnt The maximum number of whitespace characters for the whitespace + * which is used for indentation or JSON elements separation when any_whitespace is True. If + * std::nullopt, it means unlimited. Default: std::nullopt. + * \param json_format Define the root + * format of the object. If it's JSONFormat::kJSON, then it will generate a fully JSON-style + * grammar. If it's JSONFormat::kXML, then it will generate a grammar with the root format is + * XML-style, while the inner format is JSON-style. Default: JSONFormat::kJSON. + * \returns The EBNF grammar string. + */ + +std::string JSONSchemaToEBNF( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + JSONFormat json_format = JSONFormat::kJSON, + bool any_order = false +); + +/*! + * \brief Convert JSON schema string to EBNF grammar string. + * \param schema The JSON schema object. + * \param any_whitespace Whether to ignore the indentation restrictions, and allow any whitespace. + * Default: true. + * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. If std::nullopt, the default separators will be used: {",", ": "} when the + * indent is not -1, and {", ", ": "} otherwise. This follows the convention in python + * json.dumps(). Default: std::nullopt. + * \param strict_mode Whether to use strict mode. In strict + * mode, the generated grammar will not allow properties and items that is not specified in the + * schema. This is equivalent to setting unevaluatedProperties and unevaluatedItems to false. + * This helps LLM to generate accurate output in the grammar-guided generation with JSON + * schema. Default: true. + * \param max_whitespace_cnt The maximum number of whitespace characters for the whitespace + * which is used for indentation or JSON elements separation when any_whitespace is True. If + * std::nullopt, it means unlimited. Default: std::nullopt. + * \param json_format Define the root format of the object. If it's JSONFormat::kJSON, + * then it will generate a fully JSON-style grammar. If it's JSONFormat::kXML, then it will + * generate a grammar with the root format is XML-style, while the inner format is JSON-style. + * Default: JSONFormat::kJSON. + * \returns The EBNF grammar string. + */ +std::string JSONSchemaToEBNF( + const picojson::value& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + JSONFormat json_format = JSONFormat::kJSON, + bool any_order = false +); + +/*! + * \brief Generate regex pattern for integer/float range. + * \param start The start of the range (inclusive). If null assume negative infinity. + * \param end The end of the range (inclusive). If null assume infinity. + * \returns The regex pattern that matches integers/floats in the given range. + */ +std::string GenerateRangeRegex(std::optional start, std::optional end); + +std::string GenerateFloatRangeRegex( + std::optional start, + std::optional end, + bool exclusive_start = false, + bool exclusive_end = false +); + +} // namespace xgrammar + +#endif // XGRAMMAR_JSON_SCHEMA_CONVERTER_H_ diff --git a/zig/vendor/xgrammar/cpp/json_schema_converter_ext.h b/zig/vendor/xgrammar/cpp/json_schema_converter_ext.h new file mode 100644 index 0000000000..221bb8cb36 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/json_schema_converter_ext.h @@ -0,0 +1,322 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/json_schema_converter_ext.h + * \brief Extended format converters for JSON Schema, including XML Tool Calling format. + */ + +#ifndef XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ +#define XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "json_schema_converter.h" + +namespace xgrammar { + +namespace converter_ext { + +// Wrapper strings for XML parameter tags. +struct XMLWrapper { + std::string key_wrapper_prefix; + std::string key_wrapper_suffix; + std::string value_wrapper_prefix; + std::string parameter_suffix; +}; + +} // namespace converter_ext + +/*! + * \brief Converter for MiniMax M3's recursive namespace-prefixed XML format. + * + * This initial implementation supports schemas whose object property names are + * known when the grammar is built. Schemas requiring runtime element names are + * rejected explicitly. + */ +class MiniMaxM3XMLToolCallingConverter : public JSONSchemaConverter { + public: + MiniMaxM3XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false + ); + + protected: + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + std::string NextSeparator(bool is_end = false) override; + void AddBasicRules() override; + + private: + int32_t FormatElement(const std::string& name, int32_t value_rule_id); + int32_t GenerateLiteral(const picojson::value& value); + void ValidateObject(const ObjectSpec& spec) const; + static void ValidateElementName(const std::string& name); +}; + +/*! + * \brief Converter for XML Tool Calling format (e.g., Qwen style). + * + * This converter generates a grammar where: + * - The outermost object uses XML format: value + * - Inner values use standard JSON format + */ +class XMLToolCallingConverter : public JSONSchemaConverter { + public: + XMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + JSONFormat json_format = JSONFormat::kQwenXML, + bool any_order = false, + std::vector excludes = {} + ); + + /*! \brief Convert SchemaSpec to grammar with XML format for root object. Note that this function + * is not thread-safe.*/ + Grammar Convert(const SchemaSpecPtr& spec); + + protected: + // Override methods for XML format + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + // Override format hooks + int32_t FormatPropertyKey(const std::string& key, const SchemaSpecPtr& schema) override; + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ) override; + + std::string GetKeyPattern() const override; + std::string GetBasicAnyRuleName() const override; + int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ) override; + + std::string NextSeparator(bool is_end = false) override; + + void AddBasicRules() override; + + void AddCache(const std::string& key, int32_t rule_id) override; + std::optional GetCache(const std::string& key) const override; + std::string RefCacheKey(const std::string& uri) const override; + + protected: + using XMLWrapper = converter_ext::XMLWrapper; + + static const std::unordered_map kKeyWrapperMap; + static const std::string kXMLString; + static const std::string kXMLAny; + static const std::string kXMLObject; + static const std::string kXMLVariableName; + + std::string XMLValue(const std::string& json_value) const; + std::string EscapeAttrValue(const std::string& value) const; + + /*! + * \brief Return a single rendered JSON type when it can be determined without resolving + * references or combinators, or std::nullopt otherwise (\p spec may be nullptr). + * + * Shared by Kimi-K3 type attributes and DeepSeek string attributes. Both integer and + * number schemas render as "number"; string values are rendered as raw text. + */ + static std::optional GetRenderedJSONType(const SchemaSpecPtr& spec); + + /*! + * \brief Build the expression between the property key and its value. + * \param pinned_type For kimi_k3_xml, the single type attribute this property must carry. + * std::nullopt keeps every type allowed, which is what free-form keys + * (additionalProperties / patternProperties) need. + */ + int32_t XMLKeySuffix(const std::optional& pinned_type = std::nullopt); + + /*! + * \brief Build a DeepSeek XML parameter's string attribute, value and closing tag. + * string="true" wraps raw strings, string="false" wraps JSON values. Unions and mixed enums + * produce one alternative per option; GetRenderedJSONType supplies the type classification. + * A negative value_rule_id defers value-rule creation until a typed leaf is reached. + */ + int32_t FormatDeepSeekParamSuffix( + const SchemaSpecPtr& schema, int32_t value_rule_id, const std::string& rule_name_hint = "" + ); + + // Parameter suffix rules are independent of the key, so references can be shared across + // named and dynamic parameters. Allocate them before resolving refs to handle cycles. + std::unordered_map deepseek_param_ref_rules_; + + JSONFormat json_format_; + // Root parameter lists, raw parameter values, and nested JSON have distinct grammars. + int EncodingContext() const { return std::min(nested_object_level_, 2); } + // Track if we're at the root object level + int nested_object_level_ = 0; + const XMLWrapper xml_wrapper_; +}; + +/*! + * \brief Converter for Cohere XML Tool Calling format. + * + * This converter generates recursive Cohere value tags: + * value. + * Object properties use named value tags. Array items use unnamed value tags. + */ +class CohereXMLToolCallingConverter : public XMLToolCallingConverter { + public: + CohereXMLToolCallingConverter( + std::optional indent, + std::optional> separators, + bool any_whitespace, + std::optional max_whitespace_cnt, + RefResolver ref_resolver = nullptr, + bool any_order = false + ); + + protected: + int32_t GenerateString(const StringSpec& spec, const std::string& rule_name) override; + int32_t GenerateObject( + const ObjectSpec& spec, const std::string& rule_name, bool dummy_need_braces = false + ) override; + int32_t GenerateAny(const AnySpec& spec, const std::string& rule_name) override; + int32_t GenerateArray(const ArraySpec& spec, const std::string& rule_name) override; + int32_t GenerateConst(const ConstSpec& spec, const std::string& rule_name) override; + int32_t GenerateEnum(const EnumSpec& spec, const std::string& rule_name) override; + + int32_t FormatProperty( + const std::string& key, + int32_t value_rule_id, + const std::string& rule_name, + int64_t idx, + const SchemaSpecPtr& schema + ) override; + int32_t FormatOtherProperty( + int32_t key_pattern_expr, + int32_t value_rule_id, + const std::string& rule_name, + const std::string& rule_name_suffix, + const SchemaSpecPtr& schema + ) override; + + std::string GetKeyPattern() const override; + int32_t GetKeyPatternExcluding( + const std::vector& properties, const std::string& rule_name + ) override; + std::string NextSeparator(bool is_end = false) override; + + void AddBasicRules() override; + void AddCache(const std::string& key, int32_t rule_id) override; + std::optional GetCache(const std::string& key) const override; + std::string RefCacheKey(const std::string& uri) const override; + + private: + struct CohereKeyTrieNode { + bool is_terminal = false; + std::map children; + }; + + static const std::string kCohereKey; + static const std::string kCohereAnyScalar; + static const std::string kCohereAnyList; + + /*! \brief `& name, const std::optional& key_pattern_expr + ); + /*! + * \brief The type attribute, value and closing tag of one parameter, correlated with + * \p schema. The suffix does not depend on the parameter name, so `$ref` schemas get one + * memoized rule per URI that is registered before the reference is resolved; recursive + * references (directly, or through nested dict/list items) then point back at that rule. + */ + int32_t FormatCohereParamSuffix(const SchemaSpecPtr& schema, int32_t value_rule_id); + int32_t FormatCohereSuffixWithType(int32_t type_expression, int32_t value_rule_id); + int32_t FormatAnyCohereSuffix(); + int32_t FormatCohereParam( + const std::optional& name, + const std::optional& key_pattern_expr, + const SchemaSpecPtr& schema, + int32_t value_rule_id + ); + int32_t FormatAnyCohereParam( + const std::optional& name, const std::optional& key_pattern_expr + ); + int32_t FormatCohereValue(int32_t value_rule_id); + int32_t GetCohereTypePattern(const SchemaSpecPtr& schema); + static std::string CohereTypeForJSONLiteral(const std::string& json_value); + static std::optional CommonCohereTypeForJSONLiterals( + const std::vector& json_values + ); + std::optional> GetCohereCompositeOptions(const SchemaSpecPtr& schema + ) const; + int32_t BuildCohereKeyExcludingBody(const CohereKeyTrieNode& node, int depth); + bool AtCohereRoot() const; + bool InCohereValueContext() const; + + std::vector object_stack_; + std::vector additional_property_stack_; + // Parameter suffix rules keyed by `$ref` URI; see FormatCohereParamSuffix. + std::unordered_map cohere_param_ref_rules_; + int cohere_array_level_ = 0; +}; + +namespace converter_ext { + +XMLWrapper GetQwenXMLWrapper(); +XMLWrapper GetMiniMaxXMLWrapper(); +XMLWrapper GetDeepSeekXMLWrapper(); +XMLWrapper GetDeepSeekV41XMLWrapper(); +XMLWrapper GetGLMXMLWrapper(); +XMLWrapper GetCohereXMLWrapper(); +XMLWrapper GetKimiK3XMLWrapper(); + +struct XMLKeySuffix { + const char* prefix; + std::vector values; + const char* suffix; +}; + +const XMLKeySuffix& GetKimiK3XMLKeySuffix(); + +} // namespace converter_ext +} // namespace xgrammar + +#endif // XGRAMMAR_JSON_SCHEMA_CONVERTER_EXT_H_ diff --git a/zig/vendor/xgrammar/cpp/lark_converter.cc b/zig/vendor/xgrammar/cpp/lark_converter.cc new file mode 100644 index 0000000000..e4408d4cae --- /dev/null +++ b/zig/vendor/xgrammar/cpp/lark_converter.cc @@ -0,0 +1,2714 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/lark_converter.cc + */ + +#include "lark_converter.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "fsm_builder.h" +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "support/encoding.h" +#include "support/json_parse.h" +#include "support/logging.h" + +namespace xgrammar { +namespace { + +struct Location { + int line = 1; + int column = 1; +}; + +[[noreturn]] void RaiseLarkError( + const std::string& source, const Location& location, const std::string& message +) { + size_t line_start = 0; + int current_line = 1; + while (current_line < location.line && line_start < source.size()) { + size_t newline = source.find('\n', line_start); + if (newline == std::string::npos) { + line_start = source.size(); + break; + } + line_start = newline + 1; + ++current_line; + } + size_t line_end = source.find('\n', line_start); + if (line_end == std::string::npos) { + line_end = source.size(); + } + std::string line_text = source.substr(line_start, line_end - line_start); + std::ostringstream os; + os << "Lark error at line " << location.line << ", column " << location.column << ": " << message; + if (!line_text.empty()) { + os << "\n" << line_text << "\n" << std::string(std::max(0, location.column - 1), ' ') << "^"; + } + throw XGrammarError(os.str()); +} + +enum class TokenType { + kName, + kString, + kRegex, + kNumber, + kSpecialToken, + kGrammarRef, + kJson, + kRegexExt, + kGrammarOptions, + kImport, + kIgnore, + kLark, + kIf, + kUnsupportedDirective, + kColon, + kDoubleColon, + kComma, + kDot, + kDotDot, + kArrow, + kEquals, + kLParen, + kRParen, + kLBracket, + kRBracket, + kLBrace, + kRBrace, + kPipe, + kAnd, + kTilde, + kQuestion, + kStar, + kPlus, + kNewline, + kEnd, +}; + +struct Token { + TokenType type; + std::string text; + std::string flags; + Location location; +}; + +class LarkLexer { + public: + explicit LarkLexer(const std::string& source) : source_(source) {} + + std::vector Tokenize() { + std::vector result; + while (position_ < source_.size()) { + char c = source_[position_]; + if (c == ' ' || c == '\t' || c == '\f') { + Advance(); + continue; + } + if (c == '\r' || c == '\n') { + Location location = CurrentLocation(); + if (c == '\r') { + Advance(); + if (position_ < source_.size() && source_[position_] == '\n') { + Advance(); + } + } else { + Advance(); + } + result.push_back({TokenType::kNewline, "\n", "", location}); + continue; + } + if (c == '#') { + SkipComment(); + continue; + } + if (c == '/' && PeekChar(1) == '/') { + SkipComment(); + continue; + } + + Location location = CurrentLocation(); + switch (c) { + case ':': + if (PeekChar(1) == ':') { + result.push_back(SimpleToken(TokenType::kDoubleColon, 2)); + } else { + result.push_back(SimpleToken(TokenType::kColon, 1)); + } + break; + case ',': + result.push_back(SimpleToken(TokenType::kComma, 1)); + break; + case '.': + if (PeekChar(1) == '.') { + result.push_back(SimpleToken(TokenType::kDotDot, 2)); + } else { + result.push_back(SimpleToken(TokenType::kDot, 1)); + } + break; + case '-': + if (PeekChar(1) == '>') { + result.push_back(SimpleToken(TokenType::kArrow, 2)); + } else if (std::isdigit(static_cast(PeekChar(1)))) { + result.push_back(LexNumber()); + } else { + RaiseLarkError(source_, location, "unexpected '-' character"); + } + break; + case '+': + if (std::isdigit(static_cast(PeekChar(1)))) { + result.push_back(LexNumber()); + } else { + result.push_back(SimpleToken(TokenType::kPlus, 1)); + } + break; + case '=': + result.push_back(SimpleToken(TokenType::kEquals, 1)); + break; + case '(': + result.push_back(SimpleToken(TokenType::kLParen, 1)); + break; + case ')': + result.push_back(SimpleToken(TokenType::kRParen, 1)); + break; + case '[': + result.push_back(SimpleToken(TokenType::kLBracket, 1)); + break; + case ']': + result.push_back(SimpleToken(TokenType::kRBracket, 1)); + break; + case '{': + result.push_back(SimpleToken(TokenType::kLBrace, 1)); + break; + case '}': + result.push_back(SimpleToken(TokenType::kRBrace, 1)); + break; + case '|': + result.push_back(SimpleToken(TokenType::kPipe, 1)); + break; + case '&': + result.push_back(SimpleToken(TokenType::kAnd, 1)); + break; + case '~': + result.push_back(SimpleToken(TokenType::kTilde, 1)); + break; + case '?': + if (std::isalpha(static_cast(PeekChar(1))) || PeekChar(1) == '_') { + result.push_back(LexName()); + } else { + result.push_back(SimpleToken(TokenType::kQuestion, 1)); + } + break; + case '*': + result.push_back(SimpleToken(TokenType::kStar, 1)); + break; + case '!': + if (std::isalpha(static_cast(PeekChar(1))) || PeekChar(1) == '_') { + result.push_back(LexName()); + } else { + RaiseLarkError(source_, location, "unexpected '!' character"); + } + break; + case '"': + result.push_back(LexString()); + break; + case '/': + result.push_back(LexRegex()); + break; + case '<': + result.push_back(LexSpecialToken()); + break; + case '@': + result.push_back(LexGrammarRef()); + break; + case '%': + result.push_back(LexDirective()); + break; + default: + if (std::isalpha(static_cast(c)) || c == '_') { + result.push_back(LexName()); + } else if (std::isdigit(static_cast(c))) { + result.push_back(LexNumber()); + } else { + RaiseLarkError(source_, location, std::string("unexpected character '") + c + "'"); + } + } + } + result.push_back({TokenType::kEnd, "", "", CurrentLocation()}); + return result; + } + + private: + char PeekChar(size_t offset) const { + size_t index = position_ + offset; + return index < source_.size() ? source_[index] : '\0'; + } + + Location CurrentLocation() const { return {line_, column_}; } + + void Advance() { + if (position_ >= source_.size()) { + return; + } + char c = source_[position_++]; + if (c == '\n') { + ++line_; + column_ = 1; + } else { + ++column_; + } + } + + void AdvanceTo(size_t end_position) { + while (position_ < end_position) { + Advance(); + } + } + + Token SimpleToken(TokenType type, size_t length) { + Location location = CurrentLocation(); + std::string text = source_.substr(position_, length); + AdvanceTo(position_ + length); + return {type, std::move(text), "", location}; + } + + void SkipComment() { + while (position_ < source_.size() && source_[position_] != '\n' && source_[position_] != '\r') { + Advance(); + } + } + + Token LexName() { + Location location = CurrentLocation(); + size_t start = position_; + if (source_[position_] == '!' || source_[position_] == '?') { + Advance(); + } + while (position_ < source_.size()) { + char c = source_[position_]; + if (!std::isalnum(static_cast(c)) && c != '_' && c != '-') { + break; + } + Advance(); + } + return {TokenType::kName, source_.substr(start, position_ - start), "", location}; + } + + Token LexNumber() { + Location location = CurrentLocation(); + size_t start = position_; + if (source_[position_] == '+' || source_[position_] == '-') { + Advance(); + } + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + if (PeekChar(0) == '.' && PeekChar(1) != '.') { + Advance(); + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + } + if (PeekChar(0) == 'e' || PeekChar(0) == 'E') { + Advance(); + if (PeekChar(0) == '+' || PeekChar(0) == '-') { + Advance(); + } + while (std::isdigit(static_cast(PeekChar(0)))) { + Advance(); + } + } + return {TokenType::kNumber, source_.substr(start, position_ - start), "", location}; + } + + Token LexString() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + bool escaped = false; + while (position_ < source_.size()) { + char c = source_[position_]; + if (!escaped && c == '"') { + Advance(); + if (PeekChar(0) == 'i') { + Advance(); + } + return {TokenType::kString, source_.substr(start, position_ - start), "", location}; + } + if (c == '\n' || c == '\r') { + RaiseLarkError(source_, location, "unterminated string literal"); + } + if (!escaped && c == '\\') { + escaped = true; + } else { + escaped = false; + } + Advance(); + } + RaiseLarkError(source_, location, "unterminated string literal"); + } + + Token LexRegex() { + Location location = CurrentLocation(); + Advance(); + size_t pattern_start = position_; + bool escaped = false; + while (position_ < source_.size()) { + char c = source_[position_]; + if (!escaped && c == '/') { + std::string pattern = source_.substr(pattern_start, position_ - pattern_start); + Advance(); + size_t flags_start = position_; + while (std::isalpha(static_cast(PeekChar(0)))) { + Advance(); + } + return { + TokenType::kRegex, + std::move(pattern), + source_.substr(flags_start, position_ - flags_start), + location + }; + } + if (!escaped && c == '\\') { + escaped = true; + } else { + escaped = false; + } + Advance(); + } + RaiseLarkError(source_, location, "unterminated regular expression"); + } + + Token LexSpecialToken() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (position_ < source_.size() && source_[position_] != '>') { + char c = source_[position_]; + if (std::isspace(static_cast(c)) || c == '<') { + RaiseLarkError(source_, location, "invalid special token"); + } + Advance(); + } + if (position_ == source_.size()) { + RaiseLarkError(source_, location, "unterminated special token"); + } + Advance(); + return {TokenType::kSpecialToken, source_.substr(start, position_ - start), "", location}; + } + + Token LexGrammarRef() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (std::isalnum(static_cast(PeekChar(0))) || PeekChar(0) == '_' || + PeekChar(0) == '-') { + Advance(); + } + if (position_ == start + 1) { + RaiseLarkError(source_, location, "empty grammar reference"); + } + return {TokenType::kGrammarRef, source_.substr(start, position_ - start), "", location}; + } + + Token LexDirective() { + Location location = CurrentLocation(); + size_t start = position_; + Advance(); + while (std::isalpha(static_cast(PeekChar(0))) || PeekChar(0) == '_') { + Advance(); + } + std::string directive = source_.substr(start, position_ - start); + if (directive == "%json") { + return LexJSONValue(TokenType::kJson, location, directive); + } + if (directive == "%regex") { + return LexJSONValue(TokenType::kRegexExt, location, directive); + } + if (directive == "%grammar_options") { + return LexJSONValue(TokenType::kGrammarOptions, location, directive); + } + if (directive == "%import") { + return {TokenType::kImport, directive, "", location}; + } + if (directive == "%ignore") { + return {TokenType::kIgnore, directive, "", location}; + } + if (directive == "%lark") { + return {TokenType::kLark, directive, "", location}; + } + if (directive == "%if") { + return {TokenType::kIf, directive, "", location}; + } + return {TokenType::kUnsupportedDirective, directive, "", location}; + } + + Token LexJSONValue(TokenType type, const Location& location, const std::string& directive) { + while (position_ < source_.size() && + std::isspace(static_cast(source_[position_]))) { + Advance(); + } + auto begin = source_.begin() + static_cast(position_); + auto end = source_.end(); + picojson::value value; + std::string error; + auto parsed_end = ParseJSON(value, begin, end, &error); + if (!error.empty() || parsed_end == begin) { + RaiseLarkError( + source_, location, "failed to parse JSON value after " + directive + ": " + error + ); + } + size_t new_position = static_cast(parsed_end - source_.begin()); + AdvanceTo(new_position); + return {type, value.serialize(), "", location}; + } + + const std::string& source_; + size_t position_ = 0; + int line_ = 1; + int column_ = 1; +}; + +struct Document; + +struct Node { + enum class Kind { + kSequence, + kChoice, + kRepeat, + kString, + kRegex, + kRange, + kName, + kJson, + kRegexExt, + kNestedLark, + kSpecialToken, + kGrammarRef, + kNot, + }; + + Kind kind = Kind::kSequence; + Location location; + std::string text; + std::string text2; + std::string flags; + int32_t min_repeat = 0; + int32_t max_repeat = 0; + std::vector children; + std::shared_ptr nested; +}; + +struct Definition { + std::string name; + bool is_terminal = false; + bool lazy = false; + std::optional suffix; + std::optional temperature; + Location suffix_location; + std::optional stop; + Location stop_location; + std::optional stop_capture_name; + Location stop_capture_location; + std::optional max_tokens; + Location max_tokens_location; + std::optional max_chars; + Location max_chars_location; + std::optional capture_name; + Location capture_location; + Node body; + Location location; +}; + +struct Import { + std::string path; + std::string local_name; + Location location; +}; + +struct Document { + std::vector definitions; + std::vector ignores; + std::vector imports; + std::vector> options; +}; + +class LarkParser { + public: + LarkParser(const std::string& source, std::vector tokens) + : source_(source), tokens_(std::move(tokens)) {} + + Document Parse() { return ParseDocument(false); } + + private: + const Token& Peek(size_t offset = 0) const { + size_t index = std::min(position_ + offset, tokens_.size() - 1); + return tokens_[index]; + } + + bool Match(TokenType type) { + if (Peek().type != type) { + return false; + } + ++position_; + return true; + } + + Token Consume(TokenType type, const std::string& message) { + if (Peek().type != type) { + RaiseLarkError(source_, Peek().location, message); + } + return tokens_[position_++]; + } + + void ConsumeNewlines() { + while (Match(TokenType::kNewline)) { + } + } + + Document ParseDocument(bool stop_at_rbrace) { + Document document; + ConsumeNewlines(); + while (Peek().type != TokenType::kEnd && !(stop_at_rbrace && Peek().type == TokenType::kRBrace) + ) { + switch (Peek().type) { + case TokenType::kImport: + ParseImport(&document); + break; + case TokenType::kIgnore: + ParseIgnore(&document); + break; + case TokenType::kGrammarOptions: + ParseOptions(&document); + break; + case TokenType::kUnsupportedDirective: + RaiseLarkError( + source_, Peek().location, "directive " + Peek().text + " is not supported" + ); + default: + document.definitions.push_back(ParseDefinition()); + break; + } + if (Peek().type != TokenType::kNewline && Peek().type != TokenType::kEnd && + !(stop_at_rbrace && Peek().type == TokenType::kRBrace)) { + RaiseLarkError(source_, Peek().location, "expected end of grammar item"); + } + ConsumeNewlines(); + } + return document; + } + + static bool IsTerminalName(const std::string& raw_name) { + size_t index = 0; + if (!raw_name.empty() && (raw_name[0] == '!' || raw_name[0] == '?')) { + index = 1; + } + if (index < raw_name.size() && raw_name[index] == '_') { + ++index; + } + return index < raw_name.size() && std::isupper(static_cast(raw_name[index])); + } + + static std::string NormalizeRuleName(std::string name) { + if (!name.empty() && (name[0] == '!' || name[0] == '?')) { + name.erase(name.begin()); + } + return name; + } + + void ValidateCaptureName(const std::string& capture_name, const Location& location) const { + if (capture_name.empty()) { + RaiseLarkError(source_, location, "capture name must not be empty"); + } + for (char c : capture_name) { + bool valid = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || + c == '_' || c == '-' || c == '.'; + if (!valid) { + RaiseLarkError( + source_, location, "capture name must only contain letters, digits, '_', '-' and '.'" + ); + } + } + } + + Node ParseStopLikeValue(const std::string& attribute_name) { + Token token = Peek(); + if (Match(TokenType::kString)) { + return ParseStringNode(token); + } + if (Match(TokenType::kRegex)) { + Node result; + result.kind = Node::Kind::kRegex; + result.location = token.location; + result.text = token.text; + result.flags = token.flags; + return result; + } + if (Match(TokenType::kName)) { + if (!IsTerminalName(token.text)) { + RaiseLarkError( + source_, token.location, attribute_name + " terminal name must be uppercase" + ); + } + Node result; + result.kind = Node::Kind::kName; + result.location = token.location; + result.text = NormalizeRuleName(token.text); + return result; + } + RaiseLarkError( + source_, + token.location, + "expected string literal, regular expression, or uppercase terminal name after " + + attribute_name + "=" + ); + } + + void ParseImport(Document* document) { + Location location = Consume(TokenType::kImport, "expected %import").location; + Token first = Consume(TokenType::kName, "expected import path"); + std::string path = first.text; + while (Match(TokenType::kDot)) { + path += "." + Consume(TokenType::kName, "expected name after '.'").text; + } + + if (Match(TokenType::kLParen)) { + do { + Token name = Consume(TokenType::kName, "expected imported terminal name"); + document->imports.push_back({path + "." + name.text, name.text, location}); + } while (Match(TokenType::kComma)); + Consume(TokenType::kRParen, "expected ')' after import list"); + return; + } + + std::string local_name = path.substr(path.find_last_of('.') + 1); + if (Match(TokenType::kArrow)) { + local_name = Consume(TokenType::kName, "expected import alias").text; + } + document->imports.push_back({path, local_name, location}); + } + + void ParseIgnore(Document* document) { + Consume(TokenType::kIgnore, "expected %ignore"); + document->ignores.push_back(ParseChoice()); + } + + void ParseOptions(Document* document) { + Token token = Consume(TokenType::kGrammarOptions, "expected %grammar_options"); + picojson::value value; + std::string error = ParseJSON(value, token.text); + if (!error.empty()) { + RaiseLarkError(source_, token.location, "invalid %grammar_options value: " + error); + } + document->options.push_back({std::move(value), token.location}); + } + + Definition ParseDefinition() { + Token name_token = Consume(TokenType::kName, "expected rule or terminal name"); + Definition result; + result.name = NormalizeRuleName(name_token.text); + result.is_terminal = IsTerminalName(name_token.text); + result.location = name_token.location; + + if (Peek().type == TokenType::kLBracket) { + if (result.is_terminal) { + RaiseLarkError(source_, Peek().location, "attributes are only supported on rules"); + } + ParseAttributes(&result); + } + if (Peek().type == TokenType::kDot) { + RaiseLarkError(source_, Peek().location, "rule and terminal priorities are not supported"); + } + if (Peek().type == TokenType::kDoubleColon) { + RaiseLarkError(source_, Peek().location, "parametric grammar is not supported"); + } + if (Peek().type == TokenType::kLBrace) { + RaiseLarkError(source_, Peek().location, "Lark templates are not supported"); + } + Consume(TokenType::kColon, "expected ':' after rule name"); + result.body = ParseChoice(); + return result; + } + + void ParseAttributes(Definition* definition) { + Consume(TokenType::kLBracket, "expected '['"); + while (Peek().type != TokenType::kRBracket) { + Token key = Consume(TokenType::kName, "expected rule attribute"); + if (key.text == "lazy" && Peek().type != TokenType::kEquals) { + definition->lazy = true; + } else if (key.text == "max_tokens") { + Consume(TokenType::kEquals, "expected '=' after max_tokens attribute"); + Location value_location = Peek().location; + int32_t value = ParseInteger(); + if (value <= 0) { + RaiseLarkError(source_, value_location, "max_tokens must be positive"); + } + if (value > 1'000'000) { + RaiseLarkError(source_, value_location, "max_tokens is too large"); + } + if (definition->max_tokens.has_value()) { + RaiseLarkError(source_, key.location, "max_tokens attribute is specified more than once"); + } + definition->max_tokens = value; + definition->max_tokens_location = key.location; + } else if (key.text == "max_chars") { + Consume(TokenType::kEquals, "expected '=' after max_chars attribute"); + int32_t value = ParseInteger(); + if (definition->max_chars.has_value()) { + RaiseLarkError(source_, key.location, "max_chars attribute is specified more than once"); + } + definition->max_chars = value; + definition->max_chars_location = key.location; + } else if (key.text == "capture") { + std::string capture_name; + Location capture_location = key.location; + if (Match(TokenType::kEquals)) { + Token name_token = Consume(TokenType::kString, "expected string literal after capture="); + Node name_node = ParseStringNode(name_token); + if (!name_node.flags.empty()) { + RaiseLarkError( + source_, name_node.location, "case-insensitive flags are not supported on capture" + ); + } + capture_name = std::move(name_node.text); + capture_location = name_node.location; + } else { + capture_name = definition->name; + } + ValidateCaptureName(capture_name, capture_location); + if (definition->capture_name.has_value()) { + RaiseLarkError(source_, key.location, "capture attribute is specified more than once"); + } + definition->capture_name = std::move(capture_name); + definition->capture_location = capture_location; + } else if (key.text == "suffix") { + Consume(TokenType::kEquals, "expected '=' after suffix attribute"); + Node suffix = ParseStopLikeValue("suffix"); + if (suffix.kind == Node::Kind::kString && suffix.text.empty()) { + RaiseLarkError(source_, suffix.location, "suffix must not be empty"); + } + if (definition->suffix.has_value()) { + RaiseLarkError(source_, key.location, "suffix attribute is specified more than once"); + } + if (definition->stop.has_value()) { + RaiseLarkError(source_, key.location, "suffix cannot be combined with stop"); + } + Location suffix_location = suffix.location; + definition->suffix = std::move(suffix); + definition->suffix_location = suffix_location; + } else if (key.text == "stop") { + Consume(TokenType::kEquals, "expected '=' after stop attribute"); + Node stop = ParseStopLikeValue("stop"); + if (stop.kind == Node::Kind::kString && stop.text.empty()) { + RaiseLarkError(source_, stop.location, "stop must not be empty"); + } + if (definition->stop.has_value()) { + RaiseLarkError(source_, key.location, "stop attribute is specified more than once"); + } + if (definition->suffix.has_value()) { + RaiseLarkError(source_, key.location, "stop cannot be combined with suffix"); + } + Location stop_location = stop.location; + definition->stop = std::move(stop); + definition->stop_location = stop_location; + } else if (key.text == "stop_capture") { + Consume(TokenType::kEquals, "expected '=' after stop_capture attribute"); + Token name_token = + Consume(TokenType::kString, "expected string literal after stop_capture="); + Node name_node = ParseStringNode(name_token); + if (!name_node.flags.empty()) { + RaiseLarkError( + source_, + name_node.location, + "case-insensitive flags are not supported on stop_capture" + ); + } + ValidateCaptureName(name_node.text, name_node.location); + if (definition->stop_capture_name.has_value()) { + RaiseLarkError( + source_, key.location, "stop_capture attribute is specified more than once" + ); + } + Location stop_capture_location = name_node.location; + definition->stop_capture_name = std::move(name_node.text); + definition->stop_capture_location = stop_capture_location; + } else if (key.text == "temperature") { + Consume(TokenType::kEquals, "expected '=' after temperature attribute"); + Token value = Consume(TokenType::kNumber, "expected number after temperature="); + if (definition->temperature.has_value()) { + RaiseLarkError( + source_, key.location, "temperature attribute is specified more than once" + ); + } + try { + size_t parsed_length = 0; + float temperature = std::stof(value.text, &parsed_length); + if (parsed_length != value.text.size() || !std::isfinite(temperature) || + temperature < 0) { + RaiseLarkError( + source_, value.location, "temperature must be a finite non-negative number" + ); + } + definition->temperature = temperature; + } catch (const std::exception&) { + RaiseLarkError( + source_, value.location, "temperature must be a finite non-negative number" + ); + } + } else { + RaiseLarkError( + source_, + key.location, + "rule attribute '" + key.text + "' is not supported by XGrammar Lark" + ); + } + if (!Match(TokenType::kComma)) { + break; + } + } + Consume(TokenType::kRBracket, "expected ']' after rule attributes"); + if (definition->stop_capture_name.has_value() && !definition->suffix.has_value() && + !definition->stop.has_value()) { + RaiseLarkError( + source_, definition->stop_capture_location, "stop_capture requires stop or suffix" + ); + } + } + + Node ParseChoice() { + Location location = Peek().location; + std::vector alternatives; + alternatives.push_back(ParseSequence()); + while (MatchAlternativeSeparator()) { + alternatives.push_back(ParseSequence()); + } + if (alternatives.size() == 1) { + return std::move(alternatives[0]); + } + Node result; + result.kind = Node::Kind::kChoice; + result.location = location; + result.children = std::move(alternatives); + return result; + } + + bool MatchAlternativeSeparator() { + if (Match(TokenType::kPipe)) { + return true; + } + size_t saved_position = position_; + ConsumeNewlines(); + if (Match(TokenType::kPipe)) { + return true; + } + if (Peek().type == TokenType::kRParen || Peek().type == TokenType::kRBracket || + Peek().type == TokenType::kRBrace) { + return false; + } + position_ = saved_position; + return false; + } + + Node ParseSequence() { + Location location = Peek().location; + std::vector elements; + while (!IsSequenceEnd(Peek().type)) { + elements.push_back(ParseExpr()); + } + if (Match(TokenType::kArrow)) { + Consume(TokenType::kName, "expected alias name after '->'"); + } + if (Peek().type == TokenType::kAnd) { + RaiseLarkError(source_, Peek().location, "terminal intersection '&' is not supported"); + } + if (Peek().type == TokenType::kIf) { + RaiseLarkError(source_, Peek().location, "parametric %if conditions are not supported"); + } + Node result; + result.kind = Node::Kind::kSequence; + result.location = location; + result.children = std::move(elements); + return result; + } + + static bool IsSequenceEnd(TokenType type) { + return type == TokenType::kNewline || type == TokenType::kPipe || type == TokenType::kRParen || + type == TokenType::kRBracket || type == TokenType::kRBrace || type == TokenType::kEnd || + type == TokenType::kArrow || type == TokenType::kAnd || type == TokenType::kIf; + } + + int32_t ParseInteger() { + Token token = Consume(TokenType::kNumber, "expected integer"); + try { + size_t parsed = 0; + long long value = std::stoll(token.text, &parsed); + if (parsed != token.text.size() || value < 0 || value > std::numeric_limits::max()) { + RaiseLarkError(source_, token.location, "invalid non-negative repetition count"); + } + return static_cast(value); + } catch (const std::exception&) { + RaiseLarkError(source_, token.location, "invalid repetition count"); + } + } + + Node ParseExpr() { + Location location = Peek().location; + bool negated = Match(TokenType::kTilde); + Node atom = ParseAtom(); + if (negated) { + Node not_node; + not_node.kind = Node::Kind::kNot; + not_node.location = location; + not_node.children.push_back(std::move(atom)); + atom = std::move(not_node); + } + + int32_t min_repeat = -1; + int32_t max_repeat = -1; + if (Match(TokenType::kQuestion)) { + min_repeat = 0; + max_repeat = 1; + } else if (Match(TokenType::kStar)) { + min_repeat = 0; + max_repeat = -1; + } else if (Match(TokenType::kPlus)) { + min_repeat = 1; + max_repeat = -1; + } else if (Match(TokenType::kTilde)) { + min_repeat = ParseInteger(); + max_repeat = min_repeat; + if (Match(TokenType::kDotDot)) { + max_repeat = ParseInteger(); + } + } else if (Match(TokenType::kLBrace)) { + min_repeat = Peek().type == TokenType::kComma ? 0 : ParseInteger(); + if (Match(TokenType::kComma)) { + max_repeat = Peek().type == TokenType::kRBrace ? -1 : ParseInteger(); + } else { + max_repeat = min_repeat; + } + Consume(TokenType::kRBrace, "expected '}' after repetition range"); + } + + if (min_repeat == -1) { + return atom; + } + if (max_repeat != -1 && max_repeat < min_repeat) { + RaiseLarkError(source_, location, "repetition end must be greater than or equal to start"); + } + Node repeat; + repeat.kind = Node::Kind::kRepeat; + repeat.location = location; + repeat.min_repeat = min_repeat; + repeat.max_repeat = max_repeat; + repeat.children.push_back(std::move(atom)); + return repeat; + } + + // Groups recurse ParseChoice -> ParseSequence -> ParseExpr -> ParseAtom once per level, so bound + // the depth instead of overflowing the stack. Each level uses a few KB of stack; the limit keeps + // a wide margin below the 1 MB main-thread stack of Windows. + Node ParseGroupBody(const Token& open) { + if (group_depth_ >= kMaxGroupDepth) { + RaiseLarkError( + source_, + open.location, + "groups are nested deeper than " + std::to_string(kMaxGroupDepth) + " levels" + ); + } + ++group_depth_; + Node result = ParseChoice(); + --group_depth_; + return result; + } + + Node ParseAtom() { + Token token = Peek(); + if (Match(TokenType::kLParen)) { + Node result = ParseGroupBody(token); + Consume(TokenType::kRParen, "expected ')' after group"); + return result; + } + if (Match(TokenType::kLBracket)) { + Node inner = ParseGroupBody(token); + Consume(TokenType::kRBracket, "expected ']' after optional group"); + Node result; + result.kind = Node::Kind::kRepeat; + result.location = token.location; + result.min_repeat = 0; + result.max_repeat = 1; + result.children.push_back(std::move(inner)); + return result; + } + if (Match(TokenType::kString)) { + Node result = ParseStringNode(token); + if (Match(TokenType::kDotDot)) { + Token end = Consume(TokenType::kString, "expected string after '..'"); + Node end_node = ParseStringNode(end); + if (!result.flags.empty()) { + RaiseLarkError(source_, token.location, "flags are not allowed on character ranges"); + } + if (!end_node.flags.empty()) { + RaiseLarkError(source_, end.location, "flags are not allowed on character ranges"); + } + Node range; + range.kind = Node::Kind::kRange; + range.location = token.location; + range.text = result.text; + range.text2 = end_node.text; + return range; + } + return result; + } + if (Match(TokenType::kRegex)) { + Node result; + result.kind = Node::Kind::kRegex; + result.location = token.location; + result.text = token.text; + result.flags = token.flags; + return result; + } + if (Match(TokenType::kName)) { + Node result; + result.kind = Node::Kind::kName; + result.location = token.location; + result.text = NormalizeRuleName(token.text); + if (Peek().type == TokenType::kDoubleColon) { + RaiseLarkError(source_, Peek().location, "parametric grammar is not supported"); + } + if (Peek().type == TokenType::kLBrace && Peek(1).type != TokenType::kComma && + Peek(1).type != TokenType::kNumber) { + RaiseLarkError(source_, Peek().location, "Lark templates are not supported"); + } + return result; + } + if (Match(TokenType::kJson)) { + Node result; + result.kind = Node::Kind::kJson; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kRegexExt)) { + Node result; + result.kind = Node::Kind::kRegexExt; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kSpecialToken)) { + Node result; + result.kind = Node::Kind::kSpecialToken; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kGrammarRef)) { + Node result; + result.kind = Node::Kind::kGrammarRef; + result.location = token.location; + result.text = token.text; + return result; + } + if (Match(TokenType::kLark)) { + Consume(TokenType::kLBrace, "expected '{' after %lark"); + Node result; + result.kind = Node::Kind::kNestedLark; + result.location = token.location; + result.nested = std::make_shared(ParseDocument(true)); + Consume(TokenType::kRBrace, "expected '}' after nested Lark grammar"); + return result; + } + if (token.type == TokenType::kUnsupportedDirective) { + RaiseLarkError(source_, token.location, "directive " + token.text + " is not supported"); + } + RaiseLarkError(source_, token.location, "expected grammar expression"); + } + + Node ParseStringNode(const Token& token) { + std::string json_string = token.text; + std::string flags; + if (!json_string.empty() && json_string.back() == 'i') { + flags = "i"; + json_string.pop_back(); + } + picojson::value value; + std::string error = ParseJSON(value, json_string); + if (!error.empty() || !value.is()) { + RaiseLarkError(source_, token.location, "invalid string literal: " + error); + } + Node result; + result.kind = Node::Kind::kString; + result.location = token.location; + result.text = value.get(); + result.flags = std::move(flags); + return result; + } + + const std::string& source_; + std::vector tokens_; + size_t position_ = 0; + int group_depth_ = 0; + static constexpr int kMaxGroupDepth = 200; +}; + +const std::unordered_map& CommonRegexes() { + static const std::unordered_map regexes = { + {"common.DIGIT", "[0-9]"}, + {"common.HEXDIGIT", "[a-fA-F0-9]"}, + {"common.INT", "[0-9]+"}, + {"common.SIGNED_INT", "(\\+|-)?[0-9]+"}, + {"common.DECIMAL", "([0-9]+\\.[0-9]*)|(\\.[0-9]+)"}, + {"common._EXP", "[eE](\\+|-)?[0-9]+"}, + {"common.FLOAT", "([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+"}, + {"common.SIGNED_FLOAT", + "(\\+|-)?(([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+)"}, + {"common.NUMBER", + "([0-9]+)|([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+"}, + {"common.SIGNED_NUMBER", + "(\\+|-)?(([0-9]+)|([0-9]+\\.[0-9]*|\\.[0-9]+)([eE](\\+|-)?[0-9]+)?|[0-9]+[eE](\\+|-)?[0-9]+" + ")"}, + {"common.ESCAPED_STRING", "\\\"([^\\\"\\\\]|\\\\.)*\\\""}, + {"common.LCASE_LETTER", "[a-z]"}, + {"common.UCASE_LETTER", "[A-Z]"}, + {"common.LETTER", "[A-Za-z]"}, + {"common.WORD", "[A-Za-z]+"}, + {"common.CNAME", "[_A-Za-z][_A-Za-z0-9]*"}, + {"common.WS_INLINE", "[ \\t]+"}, + {"common.WS", "[ \\t\\f\\r\\n]+"}, + {"common.CR", "\\r"}, + {"common.LF", "\\n"}, + {"common.NEWLINE", "(\\r?\\n)+"}, + {"common.SH_COMMENT", "#[^\\n]*"}, + {"common.CPP_COMMENT", "//[^\\n]*"}, + {"common.C_COMMENT", "\\/\\*[^*]*\\*+(?:[^/*][^*]*\\*+)*\\/"}, + {"common.SQL_COMMENT", "--[^\\n]*"}, + }; + return regexes; +} + +std::string Trim(std::string value) { + size_t begin = 0; + while (begin < value.size() && std::isspace(static_cast(value[begin]))) { + ++begin; + } + size_t end = value.size(); + while (end > begin && std::isspace(static_cast(value[end - 1]))) { + --end; + } + return value.substr(begin, end - begin); +} + +std::optional ParseFixedRegexLiteral(const std::string& pattern) { + std::string result; + for (size_t i = 0; i < pattern.size();) { + char c = pattern[i++]; + if (c != '\\') { + if (std::string(".^$*+?()[]{}|").find(c) != std::string::npos) { + return std::nullopt; + } + result.push_back(c); + continue; + } + if (i == pattern.size()) { + return std::nullopt; + } + char escaped = pattern[i++]; + switch (escaped) { + case 'n': + result.push_back('\n'); + break; + case 'r': + result.push_back('\r'); + break; + case 't': + result.push_back('\t'); + break; + case 'f': + result.push_back('\f'); + break; + case 'v': + result.push_back('\v'); + break; + case '0': + result.push_back('\0'); + break; + case '^': + case '$': + case '.': + case '*': + case '+': + case '?': + case '(': + case ')': + case '[': + case ']': + case '{': + case '}': + case '|': + case '\\': + case '/': + case '-': + result.push_back(escaped); + break; + case 'x': + case 'u': { + bool braced = escaped == 'u' && i < pattern.size() && pattern[i] == '{'; + TCodepoint codepoint = 0; + if (braced) { + ++i; + int digit_count = 0; + while (i < pattern.size() && HexCharToInt(pattern[i]) != -1 && digit_count < 6) { + codepoint = codepoint * 16 + HexCharToInt(pattern[i++]); + ++digit_count; + } + if (digit_count == 0 || i >= pattern.size() || pattern[i++] != '}') { + return std::nullopt; + } + } else { + int digit_count = escaped == 'x' ? 2 : 4; + if (i + static_cast(digit_count) > pattern.size()) { + return std::nullopt; + } + for (int digit = 0; digit < digit_count; ++digit) { + int value = HexCharToInt(pattern[i++]); + if (value == -1) { + return std::nullopt; + } + codepoint = codepoint * 16 + value; + } + } + if (codepoint > 0x10FFFF || (codepoint >= 0xD800 && codepoint <= 0xDFFF)) { + return std::nullopt; + } + result += CharToUTF8(codepoint); + break; + } + default: + return std::nullopt; + } + } + return result; +} + +struct NamedGrammarRegistry { + std::unordered_map> inputs; + std::unordered_map compiled; + std::vector active; +}; + +class LarkCompiler { + public: + LarkCompiler( + const std::string& source, + Document document, + const std::optional& tokenizer_info, + NamedGrammarRegistry& named_grammars + ) + : source_(source), + document_(std::move(document)), + tokenizer_info_(tokenizer_info), + named_grammars_(named_grammars) {} + + Grammar Compile() { + ExpandImports(); + ParseOptions(); + IndexDefinitions(); + ValidateTerminalCycles(); + + for (const auto& definition : document_.definitions) { + rule_ids_[definition.name] = builder_.AddEmptyRule(definition.name); + } + + for (const auto& definition : document_.definitions) { + if (definition.is_terminal) { + builder_.UpdateRuleBody( + rule_ids_.at(definition.name), CompileNode(definition.body, definition.name, true) + ); + } + } + + CompileIgnore(); + + const Definition& start_definition = *definition_by_name_.at("start"); + std::optional dynamic_start_body = start_definition.temperature.has_value() + ? std::nullopt + : CompileDynamicStart(start_definition); + + for (const auto& definition : document_.definitions) { + if (definition.is_terminal) { + continue; + } + if (dynamic_unused_rules_.count(definition.name)) { + if (definition.max_tokens.has_value()) { + RaiseLarkError( + source_, + definition.max_tokens_location, + "max_tokens is not supported on rules consumed by dynamic dispatch" + ); + } + if (definition.max_chars.has_value()) { + XGRAMMAR_LOG(WARNING) << "Ignoring max_chars on rule '" << definition.name + << "' because it is consumed by dynamic dispatch."; + } + if (definition.capture_name.has_value()) { + RaiseLarkError( + source_, + definition.capture_location, + "capture is not supported on rules consumed by dynamic dispatch" + ); + } + builder_.UpdateRuleBody(rule_ids_.at(definition.name), builder_.AddEmptyStr()); + continue; + } + int32_t body_expr_id; + if (definition.temperature.has_value()) { + body_expr_id = CompileTemperatureRule(definition); + if (definition.max_chars.has_value()) { + builder_.UpdateMaxChars(rule_ids_.at(definition.name), definition.max_chars.value()); + } + } else if (definition.name == "start") { + if (dynamic_start_body.has_value()) { + if (definition.max_tokens.has_value()) { + RaiseLarkError( + source_, + definition.max_tokens_location, + "max_tokens is not supported on a dynamic dispatch start rule" + ); + } + if (definition.max_chars.has_value()) { + XGRAMMAR_LOG(WARNING) << "Ignoring max_chars on dynamic dispatch start rule '" + << definition.name << "'."; + } + body_expr_id = dynamic_start_body.value(); + } else if (definition.max_tokens.has_value() || definition.max_chars.has_value()) { + body_expr_id = CompileBudgetRule(definition); + } else if (HasLazySemantics(definition)) { + body_expr_id = CompileLazyRule(definition); + } else { + body_expr_id = CompileNode(definition.body, definition.name, false); + } + if (allow_initial_skip_ && skip_rule_id_ != -1) { + body_expr_id = builder_.AddSequence({builder_.AddRuleRef(skip_rule_id_), body_expr_id}); + } + } else if (definition.max_tokens.has_value() || definition.max_chars.has_value()) { + body_expr_id = CompileBudgetRule(definition); + } else if (HasLazySemantics(definition)) { + body_expr_id = CompileLazyRule(definition); + } else { + body_expr_id = CompileNode(definition.body, definition.name, false); + } + int32_t rule_id = rule_ids_.at(definition.name); + builder_.UpdateRuleBody(rule_id, body_expr_id); + if (definition.capture_name.has_value()) { + builder_.UpdateCaptureName(rule_id, definition.capture_name.value()); + } + builder_.UpdateRuleTemperature(rule_id, definition.temperature); + } + + auto start_it = rule_ids_.find("start"); + if (start_it == rule_ids_.end()) { + RaiseLarkError(source_, {1, 1}, "no start rule found"); + } + int32_t root_rule_id = start_it->second; + if (start_definition.temperature.has_value() && skip_rule_id_ != -1) { + std::vector elements; + if (allow_initial_skip_) { + elements.push_back(builder_.AddRuleRef(skip_rule_id_)); + } + elements.push_back(builder_.AddRuleRef(root_rule_id)); + elements.push_back(builder_.AddRuleRef(skip_rule_id_)); + root_rule_id = + builder_.AddRuleWithHint("start_with_skip", builder_.AddSequence(std::move(elements))); + } + return DeadCodeEliminator::Apply(GrammarNormalizer().Apply(builder_.Get(root_rule_id))); + } + + private: + struct SpecialTokenSet { + bool excluded = false; + std::vector token_ids; + }; + + struct Trigger { + enum class Level { kString, kToken } level; + std::string string; + std::vector token_ids; + Location location; + }; + + struct DynamicAlternative { + Trigger trigger; + Node remainder; + int32_t marker_event_rule_id = -1; + }; + + static bool HasLazySemantics(const Definition& definition) { + return definition.lazy || definition.suffix.has_value() || definition.stop.has_value(); + } + + int32_t CompileTemperatureRule(const Definition& definition) { + const Node* body = UnwrapSingle(&definition.body); + if (body->kind == Node::Kind::kJson || body->kind == Node::Kind::kNestedLark || + body->kind == Node::Kind::kGrammarRef) { + return CompileNode(*body, definition.name, false, false); + } + try { + return CompileNode(definition.body, definition.name, true, false); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + definition.location, + std::string(error.what()) + "; temperature is only supported on terminals and subgrammars" + ); + } + } + + void ExpandImports() { + for (const auto& import : document_.imports) { + auto it = CommonRegexes().find(import.path); + if (it == CommonRegexes().end()) { + RaiseLarkError(source_, import.location, "unknown common import '" + import.path + "'"); + } + Node regex; + regex.kind = Node::Kind::kRegex; + regex.location = import.location; + regex.text = it->second; + Definition definition; + definition.name = import.local_name; + definition.is_terminal = true; + definition.body = std::move(regex); + definition.location = import.location; + document_.definitions.push_back(std::move(definition)); + } + } + + void ParseOptions() { + for (const auto& [value, location] : document_.options) { + if (!value.is()) { + RaiseLarkError(source_, location, "%grammar_options value must be an object"); + } + for (const auto& [key, option] : value.get()) { + if (key == "allow_initial_skip") { + if (!option.is()) { + RaiseLarkError(source_, location, "allow_initial_skip must be a boolean"); + } + allow_initial_skip_ = allow_initial_skip_ || option.get(); + } else if (key == "no_forcing" || key == "allow_invalid_utf8") { + if (!option.is()) { + RaiseLarkError(source_, location, key + " must be a boolean"); + } + if (option.get()) { + RaiseLarkError( + source_, location, "%grammar_options option '" + key + "' is not supported" + ); + } + } else { + RaiseLarkError(source_, location, "unknown %grammar_options option '" + key + "'"); + } + } + } + } + + void IndexDefinitions() { + for (auto& definition : document_.definitions) { + if (definition_by_name_.count(definition.name)) { + RaiseLarkError( + source_, definition.location, "duplicate rule or terminal '" + definition.name + "'" + ); + } + definition_by_name_[definition.name] = &definition; + } + if (!definition_by_name_.count("start")) { + RaiseLarkError(source_, {1, 1}, "no start rule found"); + } + if (definition_by_name_.at("start")->is_terminal) { + RaiseLarkError(source_, definition_by_name_.at("start")->location, "start must be a rule"); + } + } + + void CollectReferencedNames(const Node& node, std::vector* names) const { + if (node.kind == Node::Kind::kName) { + names->push_back(node.text); + } + for (const Node& child : node.children) { + CollectReferencedNames(child, names); + } + } + + void ValidateTerminalCycles() { + std::unordered_map states; + for (const auto& definition : document_.definitions) { + if (definition.is_terminal && states[definition.name] == 0) { + VisitTerminal(definition, &states); + } + } + } + + void VisitTerminal(const Definition& definition, std::unordered_map* states) { + (*states)[definition.name] = 1; + std::vector names; + CollectReferencedNames(definition.body, &names); + for (const std::string& name : names) { + auto it = definition_by_name_.find(name); + if (it == definition_by_name_.end()) { + RaiseLarkError(source_, definition.location, "unknown name '" + name + "'"); + } + if (!it->second->is_terminal) { + RaiseLarkError( + source_, + definition.location, + "terminal '" + definition.name + "' cannot reference rule '" + name + "'" + ); + } + if ((*states)[name] == 1) { + RaiseLarkError( + source_, definition.location, "circular reference in terminal '" + name + "'" + ); + } + if ((*states)[name] == 0) { + VisitTerminal(*it->second, states); + } + } + (*states)[definition.name] = 2; + } + + void CompileIgnore() { + if (document_.ignores.empty()) { + return; + } + std::vector ignore_choices; + for (const Node& ignore : document_.ignores) { + ignore_choices.push_back(CompileNode(ignore, "lark_ignore", true)); + } + int32_t ignore_body = + ignore_choices.size() == 1 ? ignore_choices[0] : builder_.AddChoices(ignore_choices); + int32_t ignore_item_rule = builder_.AddRuleWithHint("lark_ignore_item", ignore_body); + int32_t ignore_repeat = builder_.AddRepeat(ignore_item_rule, 0, -1); + skip_rule_id_ = builder_.AddRuleWithHint("lark_ignore", ignore_repeat); + } + + int32_t CompileStringLiteral(const Node& node) { + if (node.flags.empty()) { + return node.text.empty() ? builder_.AddEmptyStr() : builder_.AddByteString(node.text); + } + if (node.flags != "i") { + RaiseLarkError( + source_, node.location, "unsupported string literal flags '" + node.flags + "'" + ); + } + std::vector codepoints = ParseUTF8(node.text.c_str()); + if (!node.text.empty() && + (codepoints.empty() || codepoints[0] == CharHandlingError::kInvalidUTF8)) { + RaiseLarkError(source_, node.location, "case-insensitive string is not valid UTF-8"); + } + std::vector elements; + elements.reserve(codepoints.size()); + for (TCodepoint codepoint : codepoints) { + if (codepoint > 0x7F) { + RaiseLarkError( + source_, + node.location, + "case-insensitive string literals currently support ASCII characters only" + ); + } + if ((codepoint >= 'a' && codepoint <= 'z') || (codepoint >= 'A' && codepoint <= 'Z')) { + TCodepoint lowercase = + static_cast(std::tolower(static_cast(codepoint))); + TCodepoint uppercase = + static_cast(std::toupper(static_cast(codepoint))); + elements.push_back( + builder_.AddCharacterClass({{lowercase, lowercase}, {uppercase, uppercase}}) + ); + } else { + elements.push_back(builder_.AddByteString(CharToUTF8(codepoint))); + } + } + if (elements.empty()) { + return builder_.AddEmptyStr(); + } + return elements.size() == 1 ? elements[0] : builder_.AddSequence(elements); + } + + struct RegexFlags { + bool case_insensitive = false; + bool dot_all = false; + }; + + RegexFlags ParseRegexFlags(const Node& node) const { + RegexFlags result; + for (char flag : node.flags) { + if (flag == 'i') { + result.case_insensitive = true; + } else if (flag == 's') { + result.dot_all = true; + } else if (flag == 'u') { + // XGrammar regular expressions use Unicode codepoint semantics by default. + } else if (flag == 'l') { + RaiseLarkError(source_, node.location, "regular-expression flag 'l' is not supported"); + } else { + RaiseLarkError( + source_, + node.location, + "regular-expression flag '" + std::string(1, flag) + "' is not supported" + ); + } + } + return result; + } + + std::string PrepareRegexPattern(const Node& node) const { + return RewriteRegexDots(node.text, ParseRegexFlags(node).dot_all); + } + + static std::string EscapeRegexLiteral(const std::string& value) { + static const std::string kRegexMeta = R"(\.^$|()[]{}*+?)"; + static constexpr char kHex[] = "0123456789ABCDEF"; + std::string result; + for (unsigned char byte : value) { + if (byte == '\n') { + result += "\\n"; + } else if (byte == '\r') { + result += "\\r"; + } else if (byte == '\t') { + result += "\\t"; + } else if (byte < 0x20 || byte == 0x7F) { + result += "\\x"; + result += kHex[byte >> 4]; + result += kHex[byte & 0x0F]; + } else { + char character = static_cast(byte); + if (kRegexMeta.find(character) != std::string::npos) { + result += '\\'; + } + result += character; + } + } + return result; + } + + std::string StringLiteralToRegex(const Node& node) { + if (node.flags.empty()) { + return EscapeRegexLiteral(node.text); + } + if (node.flags != "i") { + RaiseLarkError( + source_, node.location, "unsupported string literal flags '" + node.flags + "'" + ); + } + std::vector codepoints = ParseUTF8(node.text.c_str()); + if (!node.text.empty() && + (codepoints.empty() || codepoints[0] == CharHandlingError::kInvalidUTF8)) { + RaiseLarkError(source_, node.location, "case-insensitive string is not valid UTF-8"); + } + std::string result; + for (TCodepoint codepoint : codepoints) { + if (codepoint > 0x7F) { + RaiseLarkError( + source_, + node.location, + "case-insensitive string literals currently support ASCII characters only" + ); + } + char character = static_cast(codepoint); + if ((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z')) { + char lower = static_cast(std::tolower(static_cast(character))); + char upper = static_cast(std::toupper(static_cast(character))); + result += "["; + result += lower; + result += upper; + result += "]"; + } else { + result += EscapeRegexLiteral(std::string(1, character)); + } + } + return result; + } + + std::string TerminalNodeToRegex( + const Node& node, std::unordered_set* visiting = nullptr + ) { + auto wrap = [](const std::string& pattern) { + return pattern.empty() ? std::string() : "(?:" + pattern + ")"; + }; + switch (node.kind) { + case Node::Kind::kSequence: { + std::string result; + for (const Node& child : node.children) { + result += wrap(TerminalNodeToRegex(child, visiting)); + } + return result; + } + case Node::Kind::kChoice: { + std::string result = "(?:"; + for (size_t i = 0; i < node.children.size(); ++i) { + if (i != 0) { + result += "|"; + } + result += TerminalNodeToRegex(node.children[i], visiting); + } + return result + ")"; + } + case Node::Kind::kRepeat: { + std::string child = TerminalNodeToRegex(node.children[0], visiting); + if (child.empty()) { + return ""; + } + std::string result = "(?:" + child + ")"; + if (node.min_repeat == 0 && node.max_repeat == -1) { + return result + "*"; + } + if (node.min_repeat == 1 && node.max_repeat == -1) { + return result + "+"; + } + if (node.min_repeat == 0 && node.max_repeat == 1) { + return result + "?"; + } + result += "{" + std::to_string(node.min_repeat); + if (node.max_repeat != node.min_repeat) { + result += ","; + if (node.max_repeat != -1) { + result += std::to_string(node.max_repeat); + } + } + return result + "}"; + } + case Node::Kind::kString: + return StringLiteralToRegex(node); + case Node::Kind::kRegex: + if (ParseRegexFlags(node).case_insensitive) { + RaiseLarkError( + source_, + node.location, + "regular-expression flag 'i' is not supported with suffix or stop attributes" + ); + } + return "(?:" + PrepareRegexPattern(node) + ")"; + case Node::Kind::kRange: { + std::vector begin = ParseUTF8(node.text.c_str()); + std::vector end = ParseUTF8(node.text2.c_str()); + if (begin.size() != 1 || end.size() != 1 || begin[0] == CharHandlingError::kInvalidUTF8 || + end[0] == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "character range endpoints must be one character"); + } + if (begin[0] > end[0]) { + RaiseLarkError(source_, node.location, "character range start must not exceed end"); + } + auto escape_class_character = [](TCodepoint codepoint) { + std::string value = CharToUTF8(codepoint); + if (value == "\\" || value == "]" || value == "-" || value == "^") { + return "\\" + value; + } + return value; + }; + return "[" + escape_class_character(begin[0]) + "-" + escape_class_character(end[0]) + "]"; + } + case Node::Kind::kName: { + auto definition_it = definition_by_name_.find(node.text); + if (definition_it == definition_by_name_.end()) { + RaiseLarkError(source_, node.location, "unknown name '" + node.text + "'"); + } + if (!definition_it->second->is_terminal) { + RaiseLarkError( + source_, node.location, "terminal cannot reference rule '" + node.text + "'" + ); + } + std::unordered_set local_visiting; + if (visiting == nullptr) { + visiting = &local_visiting; + } + if (!visiting->insert(node.text).second) { + RaiseLarkError( + source_, node.location, "recursive terminal '" + node.text + "' is not supported" + ); + } + std::string result = TerminalNodeToRegex(definition_it->second->body, visiting); + visiting->erase(node.text); + return result; + } + case Node::Kind::kSpecialToken: + RaiseLarkError(source_, node.location, "special tokens cannot be used in terminals"); + case Node::Kind::kJson: + RaiseLarkError(source_, node.location, "%json cannot be used in terminals"); + case Node::Kind::kNestedLark: + RaiseLarkError(source_, node.location, "nested %lark cannot be used in terminals"); + case Node::Kind::kRegexExt: + RaiseLarkError( + source_, node.location, "structured %regex cannot be used with suffix or stop" + ); + case Node::Kind::kGrammarRef: + RaiseLarkError(source_, node.location, "named grammars cannot be used in terminals"); + case Node::Kind::kNot: + RaiseLarkError( + source_, node.location, "regular-expression complement '~' is not supported" + ); + } + RaiseLarkError(source_, node.location, "unsupported terminal node"); + } + + std::vector ParseStructuredRegexChunks(const Node& node) const { + picojson::value value; + std::string error = ParseJSON(value, node.text); + if (!error.empty()) { + RaiseLarkError(source_, node.location, "failed to parse %regex: " + error); + } + if (!value.is()) { + RaiseLarkError(source_, node.location, "%regex value must be an object"); + } + + const auto& object = value.get(); + std::vector fields; + for (const auto& [key, field_value] : object) { + if (key != "substring_chunks" && key != "substring_chars" && key != "substring_words") { + RaiseLarkError(source_, node.location, "unknown field '" + key + "' in %regex"); + } + if (!field_value.is()) { + fields.push_back(key); + } + } + if (fields.empty()) { + RaiseLarkError(source_, node.location, "no fields set on %regex"); + } + if (fields.size() != 1) { + RaiseLarkError(source_, node.location, "only one field can be set on %regex"); + } + + const std::string& field = fields[0]; + const picojson::value& field_value = object.at(field); + if (field == "substring_words") { + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_words must be a string"); + } + RaiseLarkError(source_, node.location, "substring_words is not supported yet"); + } + if (field == "substring_chars") { + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_chars must be a string"); + } + const std::string& text = field_value.get(); + std::vector chunks; + for (size_t offset = 0; offset < text.size();) { + if (text[offset] == '\0') { + chunks.emplace_back(1, '\0'); + ++offset; + continue; + } + auto [codepoint, length] = ParseNextUTF8(text.c_str() + offset); + if (codepoint == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "substring_chars must be valid UTF-8"); + } + chunks.push_back(text.substr(offset, length)); + offset += length; + } + return chunks; + } + + if (!field_value.is()) { + RaiseLarkError(source_, node.location, "substring_chunks must be an array of strings"); + } + std::vector chunks; + const auto& array = field_value.get(); + chunks.reserve(array.size()); + for (const picojson::value& chunk : array) { + if (!chunk.is()) { + RaiseLarkError(source_, node.location, "substring_chunks must be an array of strings"); + } + chunks.push_back(chunk.get()); + } + return chunks; + } + + int32_t CompileStructuredRegex(const Node& node, const std::string& rule_hint) { + std::vector chunks = ParseStructuredRegexChunks(node); + // A substring expr can only be the body of a rule, so wrap it and return a reference. + int32_t substring_expr_id = builder_.AddSubstring(chunks); + int32_t rule_id = builder_.AddRuleWithHint(rule_hint + "_substring", substring_expr_id); + return builder_.AddRuleRef(rule_id); + } + + const Grammar& ResolveNamedGrammar(const std::string& name, const Location& location) { + auto input_it = named_grammars_.inputs.find(name); + if (input_it == named_grammars_.inputs.end()) { + RaiseLarkError(source_, location, "unknown named grammar '@" + name + "'"); + } + if (std::holds_alternative(input_it->second)) { + return std::get(input_it->second); + } + auto compiled_it = named_grammars_.compiled.find(name); + if (compiled_it != named_grammars_.compiled.end()) { + return compiled_it->second; + } + + auto active_it = std::find(named_grammars_.active.begin(), named_grammars_.active.end(), name); + if (active_it != named_grammars_.active.end()) { + std::ostringstream cycle; + for (auto it = active_it; it != named_grammars_.active.end(); ++it) { + if (it != active_it) { + cycle << " -> "; + } + cycle << "@" << *it; + } + cycle << " -> @" << name; + RaiseLarkError(source_, location, "circular named grammar reference: " + cycle.str()); + } + + named_grammars_.active.push_back(name); + try { + const std::string& named_source = std::get(input_it->second); + auto tokens = LarkLexer(named_source).Tokenize(); + auto document = LarkParser(named_source, std::move(tokens)).Parse(); + Grammar compiled = + LarkCompiler(named_source, std::move(document), tokenizer_info_, named_grammars_) + .Compile(); + auto compiled_it = named_grammars_.compiled.emplace(name, std::move(compiled)).first; + named_grammars_.active.pop_back(); + return compiled_it->second; + } catch (const std::exception& error) { + named_grammars_.active.pop_back(); + RaiseLarkError( + source_, + location, + "failed to compile named grammar '@" + name + "': " + std::string(error.what()) + ); + } + } + + int32_t CompileNode( + const Node& node, const std::string& rule_hint, bool terminal_mode, bool append_skip = true + ) { + switch (node.kind) { + case Node::Kind::kSequence: { + if (node.children.empty()) { + return builder_.AddEmptyStr(); + } + std::vector elements; + elements.reserve(node.children.size()); + for (const Node& child : node.children) { + elements.push_back(CompileNode(child, rule_hint, terminal_mode, append_skip)); + } + return elements.size() == 1 ? elements[0] : builder_.AddSequence(elements); + } + case Node::Kind::kChoice: { + std::vector choices; + choices.reserve(node.children.size()); + for (const Node& child : node.children) { + choices.push_back(CompileNode(child, rule_hint, terminal_mode, append_skip)); + } + return choices.size() == 1 ? choices[0] : builder_.AddChoices(choices); + } + case Node::Kind::kRepeat: { + int32_t child = + CompileNode(node.children[0], rule_hint + "_repeat", terminal_mode, append_skip); + return builder_.AddRepeatFromExpr( + rule_hint + "_repeat", child, node.min_repeat, node.max_repeat + ); + } + case Node::Kind::kName: { + auto definition_it = definition_by_name_.find(node.text); + if (definition_it == definition_by_name_.end()) { + RaiseLarkError(source_, node.location, "unknown name '" + node.text + "'"); + } + if (terminal_mode && !definition_it->second->is_terminal) { + RaiseLarkError( + source_, node.location, "terminal cannot reference rule '" + node.text + "'" + ); + } + int32_t result = builder_.AddRuleRef(rule_ids_.at(node.text)); + // Lazy rules are compiled like terminals (lexemes), so they also take a trailing skip. + // Temperature rules are also compiled like terminals. + bool is_lexeme = definition_it->second->is_terminal || + HasLazySemantics(*definition_it->second) || + definition_it->second->temperature.has_value(); + return !terminal_mode && append_skip && is_lexeme ? AppendSkip(result) : result; + } + case Node::Kind::kString: { + int32_t result = CompileStringLiteral(node); + return !terminal_mode && append_skip && !node.text.empty() ? AppendSkip(result) : result; + } + case Node::Kind::kRange: { + auto begin = ParseUTF8(node.text.c_str()); + auto end = ParseUTF8(node.text2.c_str()); + if (begin.size() != 1 || end.size() != 1 || begin[0] == CharHandlingError::kInvalidUTF8 || + end[0] == CharHandlingError::kInvalidUTF8) { + RaiseLarkError(source_, node.location, "character range endpoints must be one character"); + } + if (begin[0] > end[0]) { + RaiseLarkError(source_, node.location, "character range start must not exceed end"); + } + int32_t result = builder_.AddCharacterClass({{begin[0], end[0]}}); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + case Node::Kind::kRegex: { + RegexFlags flags = ParseRegexFlags(node); + std::string pattern = RewriteRegexDots(node.text, flags.dot_all); + if (flags.case_insensitive) { + // The FSM regex engine handles the (?i) prefix with ASCII case folding. Validate the + // pattern eagerly so that errors carry the source location. + std::string flagged_pattern = "(?i)" + pattern; + auto matches_empty = RegexFSMBuilder::MatchesEmpty(flagged_pattern); + if (matches_empty.IsErr()) { + RaiseLarkError( + source_, + node.location, + "failed to compile regular expression: " + + std::string(std::move(matches_empty).UnwrapErr().what()) + ); + } + int32_t regex_rule_id = + builder_.AddRuleWithHint(rule_hint + "_regex", builder_.AddRegex(flagged_pattern)); + int32_t result = builder_.AddRuleRef(regex_rule_id); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + try { + int32_t root = SubGrammarAdder::Apply(&builder_, Grammar::FromRegex(pattern)); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile regular expression: " + std::string(error.what()) + ); + } + } + case Node::Kind::kJson: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "%json cannot be used in terminals"); + } + try { + int32_t root = SubGrammarAdder::Apply(&builder_, Grammar::FromJSONSchema(node.text)); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile inline JSON schema: " + std::string(error.what()) + ); + } + } + case Node::Kind::kNestedLark: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "nested %lark cannot be used in terminals"); + } + try { + LarkCompiler compiler(source_, *node.nested, tokenizer_info_, named_grammars_); + int32_t root = SubGrammarAdder::Apply(&builder_, compiler.Compile()); + int32_t result = builder_.AddRuleRef(root); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } catch (const std::exception& error) { + RaiseLarkError( + source_, + node.location, + "failed to compile nested Lark grammar: " + std::string(error.what()) + ); + } + } + case Node::Kind::kSpecialToken: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "special tokens cannot be used in terminals"); + } + SpecialTokenSet token_set = ResolveSpecialToken(node.text, node.location); + int32_t result = token_set.excluded ? builder_.AddExcludeTokenSet(token_set.token_ids) + : builder_.AddTokenSet(token_set.token_ids); + return append_skip ? AppendSkip(result) : result; + } + case Node::Kind::kRegexExt: { + int32_t result = CompileStructuredRegex(node, rule_hint); + return terminal_mode || !append_skip ? result : AppendSkip(result); + } + case Node::Kind::kGrammarRef: { + if (terminal_mode) { + RaiseLarkError(source_, node.location, "named grammars cannot be used in terminals"); + } + std::string name = node.text.substr(1); + auto root_it = named_grammar_roots_.find(name); + if (root_it == named_grammar_roots_.end()) { + int32_t root = + SubGrammarAdder::Apply(&builder_, ResolveNamedGrammar(name, node.location)); + root_it = named_grammar_roots_.emplace(name, root).first; + } + int32_t result = builder_.AddRuleRef(root_it->second); + return append_skip ? AppendSkip(result) : result; + } + case Node::Kind::kNot: + RaiseLarkError( + source_, node.location, "regular-expression complement '~' is not supported" + ); + } + RaiseLarkError(source_, node.location, "unsupported grammar node"); + } + + int32_t AppendSkip(int32_t expression) { + if (skip_rule_id_ == -1) { + return expression; + } + return builder_.AddSequence({expression, builder_.AddRuleRef(skip_rule_id_)}); + } + + /*! \brief Compile a rule with a token or character budget. */ + int32_t CompileBudgetRule(const Definition& definition) { + int32_t rule_id = rule_ids_.at(definition.name); + if (definition.max_tokens.has_value()) { + builder_.UpdateMaxTokens(rule_id, definition.max_tokens.value()); + } + if (definition.max_chars.has_value()) { + builder_.UpdateMaxChars(rule_id, definition.max_chars.value()); + } + if (HasLazySemantics(definition)) { + return CompileLazyRule(definition); + } + return CompileNode(definition.body, definition.name, false); + } + + SpecialTokenSet ResolveSpecialToken(const std::string& token, const Location& location) const { + if (token.size() >= 4 && token.substr(0, 2) == "<[" && token.substr(token.size() - 2) == "]>") { + std::string contents = token.substr(2, token.size() - 4); + SpecialTokenSet result; + if (!contents.empty() && contents[0] == '^') { + result.excluded = true; + contents.erase(contents.begin()); + } + if (contents == "*") { + if (result.excluded) { + RaiseLarkError(source_, location, "negated wildcard special token is not supported"); + } + if (!tokenizer_info_.has_value()) { + RaiseLarkError(source_, location, "wildcard special token requires tokenizer_info"); + } + result.token_ids.reserve(tokenizer_info_->GetVocabSize()); + for (int32_t token_id = 0; token_id < tokenizer_info_->GetVocabSize(); ++token_id) { + result.token_ids.push_back(token_id); + } + return result; + } + if (contents.find('*') != std::string::npos) { + RaiseLarkError(source_, location, "wildcard cannot be mixed with token ranges"); + } + size_t offset = 0; + while (offset <= contents.size()) { + size_t comma = contents.find(',', offset); + std::string range = Trim(contents.substr(offset, comma - offset)); + if (!range.empty()) { + size_t dash = range.find('-'); + if (dash != std::string::npos && range.find('-', dash + 1) != std::string::npos) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + int64_t first; + int64_t last; + try { + auto parse_token_id = [](const std::string& value) { + std::string trimmed = Trim(value); + size_t parsed = 0; + int64_t result = std::stoll(trimmed, &parsed); + if (parsed != trimmed.size()) { + throw std::invalid_argument("trailing characters"); + } + return result; + }; + first = parse_token_id(range.substr(0, dash)); + last = dash == std::string::npos ? first : parse_token_id(range.substr(dash + 1)); + } catch (const std::exception&) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + if (first < 0 || last < first || last > std::numeric_limits::max()) { + RaiseLarkError( + source_, location, "invalid numeric special-token range '" + range + "'" + ); + } + if (last - first > 1'000'000) { + RaiseLarkError(source_, location, "special-token range is too large"); + } + for (int64_t token_id = first; token_id <= last; ++token_id) { + result.token_ids.push_back(static_cast(token_id)); + } + } + if (comma == std::string::npos) { + break; + } + offset = comma + 1; + } + if (result.token_ids.empty()) { + RaiseLarkError(source_, location, "empty numeric special-token range"); + } + std::sort(result.token_ids.begin(), result.token_ids.end()); + result.token_ids.erase( + std::unique(result.token_ids.begin(), result.token_ids.end()), result.token_ids.end() + ); + return result; + } + + if (!tokenizer_info_.has_value()) { + RaiseLarkError( + source_, location, "named special token " + token + " requires tokenizer_info" + ); + } + SpecialTokenSet result; + const auto& decoded_vocab = tokenizer_info_->GetDecodedVocab(); + for (int32_t token_id = 0; token_id < static_cast(decoded_vocab.size()); ++token_id) { + if (decoded_vocab[token_id] == token) { + result.token_ids.push_back(token_id); + } + } + if (result.token_ids.empty()) { + RaiseLarkError(source_, location, "unknown special token " + token); + } + return result; + } + + static const Node* UnwrapSingle(const Node* node) { + while (node->kind == Node::Kind::kSequence && node->children.size() == 1) { + node = &node->children[0]; + } + return node; + } + + bool IsAnyText(const Node& node, std::unordered_set* visiting = nullptr) const { + if (node.kind == Node::Kind::kRegex) { + std::string pattern; + for (char c : node.text) { + if (c != ' ' && c != '\t' && c != '\r') { + pattern.push_back(c); + } + } + if (node.flags.find_first_not_of("isu") != std::string::npos) { + return false; + } + if (node.flags.find('s') != std::string::npos) { + return pattern == ".*"; + } + return pattern == "(.|\\n)*" || pattern == "(\\n|.)*" || pattern == "(?s:.*)" || + pattern == "(?:.|\\n)*" || pattern == "(?:\\n|.)*" || pattern == "[\\s\\S]*"; + } + if (node.kind == Node::Kind::kSequence && node.children.size() == 1) { + return IsAnyText(node.children[0], visiting); + } + if (node.kind == Node::Kind::kName) { + std::unordered_set local_visiting; + if (visiting == nullptr) { + visiting = &local_visiting; + } + if (visiting->count(node.text)) { + return false; + } + auto it = definition_by_name_.find(node.text); + if (it == definition_by_name_.end()) { + return false; + } + visiting->insert(node.text); + bool result = IsAnyText(it->second->body, visiting); + visiting->erase(node.text); + return result; + } + return false; + } + + std::optional ExtractLazyRegexTrigger(const Node& node) const { + if (node.kind != Node::Kind::kRegex) { + return std::nullopt; + } + if (node.flags.find('i') != std::string::npos || + node.flags.find_first_not_of("su") != std::string::npos) { + return std::nullopt; + } + std::vector prefixes; + if (node.flags.find('s') != std::string::npos) { + prefixes = {".*"}; + } else { + prefixes = {"(.|\\n)*", "(\\n|.)*", "(?:.|\\n)*", "(?:\\n|.)*", "[\\s\\S]*", "(?s:.*)"}; + } + for (const std::string& prefix : prefixes) { + if (node.text.size() <= prefix.size() || node.text.compare(0, prefix.size(), prefix) != 0) { + continue; + } + auto trigger = ParseFixedRegexLiteral(node.text.substr(prefix.size())); + if (trigger.has_value() && !trigger->empty()) { + return Trigger{Trigger::Level::kString, std::move(trigger.value()), {}, node.location}; + } + } + return std::nullopt; + } + + std::optional ExtractLazyTrigger(const Definition& definition) const { + if (definition.stop.has_value()) { + const Node& marker = definition.stop.value(); + if (!IsAnyText(definition.body) || marker.kind != Node::Kind::kString || + !marker.flags.empty()) { + return std::nullopt; + } + return Trigger{Trigger::Level::kString, marker.text, {}, definition.stop_location}; + } + if (definition.suffix.has_value()) { + const Node& marker = definition.suffix.value(); + if (!IsAnyText(definition.body) || marker.kind != Node::Kind::kString || + !marker.flags.empty()) { + return std::nullopt; + } + return Trigger{Trigger::Level::kString, marker.text, {}, definition.suffix_location}; + } + if (!definition.lazy) { + return std::nullopt; + } + const Node* body = UnwrapSingle(&definition.body); + if (body->kind == Node::Kind::kRegex) { + auto regex_trigger = ExtractLazyRegexTrigger(*body); + if (regex_trigger.has_value()) { + return regex_trigger; + } + } + if (definition.body.kind != Node::Kind::kSequence || definition.body.children.size() != 2 || + !IsAnyText(definition.body.children[0])) { + return std::nullopt; + } + const Node& trigger = definition.body.children[1]; + if (trigger.kind == Node::Kind::kString && !trigger.text.empty() && trigger.flags.empty()) { + return Trigger{Trigger::Level::kString, trigger.text, {}, trigger.location}; + } + if (trigger.kind == Node::Kind::kSpecialToken) { + SpecialTokenSet token_set = ResolveSpecialToken(trigger.text, trigger.location); + if (token_set.excluded) { + RaiseLarkError(source_, trigger.location, "lazy special-token trigger cannot be negated"); + } + return Trigger{Trigger::Level::kToken, "", token_set.token_ids, trigger.location}; + } + return std::nullopt; + } + + int32_t CompileLazyRule(const Definition& definition) { + int32_t rule_id = rule_ids_.at(definition.name); + const Node* marker = definition.suffix.has_value() + ? &definition.suffix.value() + : (definition.stop.has_value() ? &definition.stop.value() : nullptr); + bool marker_has_fixed_byte_length = marker != nullptr && marker->kind == Node::Kind::kString; + int32_t hidden_bytes = + marker_has_fixed_byte_length ? static_cast(marker->text.size()) : 1; + Grammar::Impl::SuffixStopInfo suffix_stop_info; + if (definition.suffix.has_value()) { + suffix_stop_info.hidden_suffix_bytes = hidden_bytes; + } else if (definition.stop.has_value()) { + suffix_stop_info.hidden_stop_bytes = hidden_bytes; + } + if (definition.stop_capture_name.has_value()) { + suffix_stop_info.stop_capture_name = definition.stop_capture_name.value(); + } + const Node* body = UnwrapSingle(&definition.body); + if (!definition.suffix.has_value() && !definition.stop.has_value() && + body->kind == Node::Kind::kRegex && ExtractLazyRegexTrigger(*body).has_value()) { + RaiseLarkError( + source_, + definition.location, + "lazy regex suffix is only supported on a head used by dynamic dispatch" + ); + } + std::optional body_pattern; + std::optional marker_pattern; + if (marker != nullptr && (!marker_has_fixed_byte_length || definition.max_tokens.has_value() || + definition.max_chars.has_value())) { + body_pattern = TerminalNodeToRegex(definition.body); + marker_pattern = TerminalNodeToRegex(*marker); + int32_t body_helper_expr = builder_.AddRegex(body_pattern.value()); + int32_t body_helper_rule = + builder_.AddRuleWithHint(definition.name + "_stop_body", body_helper_expr); + int32_t marker_helper_expr = builder_.AddRegex(marker_pattern.value()); + int32_t marker_helper_rule = + builder_.AddRuleWithHint(definition.name + "_stop_marker", marker_helper_expr); + suffix_stop_info.body_rule_id = body_helper_rule; + suffix_stop_info.marker_rule_id = marker_helper_rule; + } + builder_.UpdateSuffixStopInfo(rule_id, suffix_stop_info); + auto trigger = ExtractLazyTrigger(definition); + if (!trigger.has_value()) { + // General committed-shortest lazy rule: compiled like a terminal (no skip insertion); + // the terminal-like requirement is validated after grammar optimization. suffix="s" and + // stop="s" both desugar to the lazy rule over (body "s"); their only difference is capture + // scope, represented by the metadata set above. + builder_.UpdateLazy(rule_id, true); + if (marker != nullptr) { + if (!body_pattern.has_value()) { + body_pattern = TerminalNodeToRegex(definition.body); + marker_pattern = TerminalNodeToRegex(*marker); + } + return builder_.AddRegex( + "(?:" + body_pattern.value() + ")(?:" + marker_pattern.value() + ")" + ); + } + return CompileNode(definition.body, definition.name, true); + } + int32_t empty_rule = builder_.AddRuleWithHint("lark_lazy_end", builder_.AddEmptyStr()); + int32_t result; + if (trigger->level == Trigger::Level::kString) { + result = builder_.AddTagDispatch({{{trigger->string, empty_rule}}, false, {}}); + } else { + Grammar::Impl::TokenTagDispatch dispatch; + for (int32_t token_id : trigger->token_ids) { + dispatch.trigger_rule_pairs.push_back({token_id, empty_rule}); + } + dispatch.loop_after_dispatch = false; + result = builder_.AddTokenTagDispatch(dispatch); + } + return AppendSkip(result); + } + + static std::vector FlattenSequence(const Node& node) { + if (node.kind == Node::Kind::kSequence) { + return node.children; + } + return {node}; + } + + std::optional CompileDynamicStart(const Definition& start) { + std::unordered_set unused_rules; + std::vector start_elements = FlattenSequence(start.body); + if (start_elements.size() != 2) { + return std::nullopt; + } + const Node* loop = UnwrapSingle(&start_elements[0]); + if (loop->kind != Node::Kind::kRepeat || loop->min_repeat != 0 || loop->max_repeat != -1) { + return std::nullopt; + } + const Node* loop_body = UnwrapSingle(&loop->children[0]); + std::vector tool_names; + if (loop_body->kind == Node::Kind::kChoice) { + for (const Node& alternative : loop_body->children) { + const Node* name = UnwrapSingle(&alternative); + if (name->kind != Node::Kind::kName) { + return std::nullopt; + } + tool_names.push_back(name->text); + } + } else if (loop_body->kind == Node::Kind::kName) { + tool_names.push_back(loop_body->text); + } else { + return std::nullopt; + } + + const Node* tail_name = UnwrapSingle(&start_elements[1]); + if (tail_name->kind != Node::Kind::kName) { + return std::nullopt; + } + auto tail_it = definition_by_name_.find(tail_name->text); + if (tail_it == definition_by_name_.end() || !IsAnyText(tail_it->second->body)) { + return std::nullopt; + } + unused_rules.insert(tail_name->text); + + std::vector alternatives; + for (const std::string& tool_name : tool_names) { + auto tool_it = definition_by_name_.find(tool_name); + if (tool_it == definition_by_name_.end() || tool_it->second->is_terminal) { + return std::nullopt; + } + unused_rules.insert(tool_name); + std::vector tool_elements = FlattenSequence(tool_it->second->body); + if (tool_elements.empty()) { + return std::nullopt; + } + + std::optional trigger; + int32_t marker_event_rule_id = -1; + size_t remainder_begin = 0; + const Node* first = UnwrapSingle(&tool_elements[0]); + if (first->kind == Node::Kind::kName) { + auto head_it = definition_by_name_.find(first->text); + if (head_it != definition_by_name_.end()) { + trigger = ExtractLazyTrigger(*head_it->second); + if (trigger.has_value()) { + unused_rules.insert(first->text); + remainder_begin = 1; + const Definition& head = *head_it->second; + if (head.stop.has_value() || head.stop_capture_name.has_value()) { + // The dispatch FSM consumes the trigger before entering the remainder. Insert a + // zero-width rule there so capture materialization can recover that preceding + // marker without giving up the deterministic dispatch path. + int32_t event_expr = builder_.AddEmptyStr(); + marker_event_rule_id = + builder_.AddRuleWithHint(head.name + "_dynamic_marker", event_expr); + int32_t marker_expr = builder_.AddByteString(trigger->string); + int32_t marker_rule_id = + builder_.AddRuleWithHint(head.name + "_dynamic_marker_text", marker_expr); + Grammar::Impl::SuffixStopInfo suffix_stop_info; + suffix_stop_info.body_rule_id = marker_event_rule_id; + suffix_stop_info.marker_rule_id = marker_rule_id; + int32_t hidden_bytes = static_cast(trigger->string.size()); + if (head.stop.has_value()) { + suffix_stop_info.hidden_stop_bytes = hidden_bytes; + } else { + suffix_stop_info.hidden_suffix_bytes = hidden_bytes; + } + if (head.stop_capture_name.has_value()) { + suffix_stop_info.stop_capture_name = head.stop_capture_name.value(); + } + builder_.UpdateSuffixStopInfo(marker_event_rule_id, suffix_stop_info); + } + } + } + } + if (!trigger.has_value() && tool_elements.size() >= 2 && IsAnyText(tool_elements[0])) { + const Node* token_trigger = UnwrapSingle(&tool_elements[1]); + if (token_trigger->kind == Node::Kind::kSpecialToken) { + SpecialTokenSet token_set = + ResolveSpecialToken(token_trigger->text, token_trigger->location); + if (token_set.excluded) { + RaiseLarkError( + source_, token_trigger->location, "dynamic special-token trigger cannot be negated" + ); + } + trigger = + Trigger{Trigger::Level::kToken, "", token_set.token_ids, token_trigger->location}; + remainder_begin = 2; + } + } + if (!trigger.has_value()) { + return std::nullopt; + } + + Node remainder; + remainder.kind = Node::Kind::kSequence; + remainder.location = tool_it->second->location; + remainder.children.assign( + tool_elements.begin() + static_cast(remainder_begin), tool_elements.end() + ); + alternatives.push_back( + {std::move(trigger.value()), std::move(remainder), marker_event_rule_id} + ); + } + + if (alternatives.empty()) { + return std::nullopt; + } + Trigger::Level level = alternatives[0].trigger.level; + for (const auto& alternative : alternatives) { + if (alternative.trigger.level != level) { + RaiseLarkError( + source_, + start.location, + "a dynamic Lark start rule cannot mix string and token triggers" + ); + } + } + + if (level == Trigger::Level::kString) { + std::unordered_map> grouped; + std::vector trigger_order; + for (const auto& alternative : alternatives) { + if (!grouped.count(alternative.trigger.string)) { + trigger_order.push_back(alternative.trigger.string); + } + grouped[alternative.trigger.string].push_back(&alternative); + } + Grammar::Impl::TagDispatch dispatch; + dispatch.loop_after_dispatch = true; + for (const std::string& trigger : trigger_order) { + std::vector remainder_choices; + for (const DynamicAlternative* alternative : grouped.at(trigger)) { + int32_t remainder = CompileNode(alternative->remainder, "lark_dynamic_body", false); + if (alternative->marker_event_rule_id >= 0) { + remainder = builder_.AddSequence( + {builder_.AddRuleRef(alternative->marker_event_rule_id), remainder} + ); + } + remainder_choices.push_back(remainder); + } + int32_t body = remainder_choices.size() == 1 ? remainder_choices[0] + : builder_.AddChoices(remainder_choices); + int32_t body_rule = builder_.AddRuleWithHint("lark_dynamic_body", body); + dispatch.tag_rule_pairs.push_back({trigger, body_rule}); + } + dynamic_unused_rules_ = std::move(unused_rules); + return builder_.AddTagDispatch(dispatch); + } + + std::unordered_map> grouped; + std::vector token_order; + for (const auto& alternative : alternatives) { + for (int32_t token_id : alternative.trigger.token_ids) { + if (!grouped.count(token_id)) { + token_order.push_back(token_id); + } + grouped[token_id].push_back(&alternative); + } + } + Grammar::Impl::TokenTagDispatch dispatch; + dispatch.loop_after_dispatch = true; + for (int32_t token_id : token_order) { + std::vector remainder_choices; + for (const DynamicAlternative* alternative : grouped.at(token_id)) { + remainder_choices.push_back( + CompileNode(alternative->remainder, "lark_dynamic_token_body", false) + ); + } + int32_t body = remainder_choices.size() == 1 ? remainder_choices[0] + : builder_.AddChoices(remainder_choices); + int32_t body_rule = builder_.AddRuleWithHint("lark_dynamic_token_body", body); + dispatch.trigger_rule_pairs.push_back({token_id, body_rule}); + } + dynamic_unused_rules_ = std::move(unused_rules); + return builder_.AddTokenTagDispatch(dispatch); + } + + const std::string& source_; + Document document_; + const std::optional& tokenizer_info_; + NamedGrammarRegistry& named_grammars_; + GrammarBuilder builder_; + std::unordered_map definition_by_name_; + std::unordered_map rule_ids_; + std::unordered_map named_grammar_roots_; + int32_t skip_rule_id_ = -1; + bool allow_initial_skip_ = false; + std::unordered_set dynamic_unused_rules_; +}; + +} // namespace + +Grammar LarkToGrammar( + const std::string& lark_string, + const std::optional& tokenizer_info, + const std::vector& named_grammars +) { + NamedGrammarRegistry named_grammar_registry; + for (const auto& [name, grammar_or_source] : named_grammars) { + if (name.empty()) { + throw XGrammarError("Named grammar names must not be empty"); + } + if (!std::all_of(name.begin(), name.end(), [](unsigned char character) { + return std::isalnum(character) || character == '_' || character == '-'; + })) { + throw XGrammarError( + "Invalid named grammar name '" + name + + "': names may contain only letters, digits, underscores, and hyphens" + ); + } + if (!named_grammar_registry.inputs.emplace(name, grammar_or_source).second) { + throw XGrammarError("Duplicate named grammar '" + name + "'"); + } + } + auto tokens = LarkLexer(lark_string).Tokenize(); + auto document = LarkParser(lark_string, std::move(tokens)).Parse(); + return LarkCompiler(lark_string, std::move(document), tokenizer_info, named_grammar_registry) + .Compile(); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/lark_converter.h b/zig/vendor/xgrammar/cpp/lark_converter.h new file mode 100644 index 0000000000..411652047f --- /dev/null +++ b/zig/vendor/xgrammar/cpp/lark_converter.h @@ -0,0 +1,26 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/lark_converter.h + * \brief Convert Lark syntax to XGrammar Grammar IR. + */ + +#ifndef XGRAMMAR_LARK_CONVERTER_H_ +#define XGRAMMAR_LARK_CONVERTER_H_ + +#include + +#include +#include +#include + +namespace xgrammar { + +Grammar LarkToGrammar( + const std::string& lark_string, + const std::optional& tokenizer_info = std::nullopt, + const std::vector& named_grammars = {} +); + +} // namespace xgrammar + +#endif // XGRAMMAR_LARK_CONVERTER_H_ diff --git a/zig/vendor/xgrammar/cpp/regex_converter.cc b/zig/vendor/xgrammar/cpp/regex_converter.cc new file mode 100644 index 0000000000..b46bd76ada --- /dev/null +++ b/zig/vendor/xgrammar/cpp/regex_converter.cc @@ -0,0 +1,413 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/regex_converter.cc + */ +#include "regex_converter.h" + +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/logging.h" +#include "support/utils.h" + +namespace xgrammar { + +/*! + * \brief Convert a regex to EBNF. + * \details The implementation refers to the regex described in + * https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Regular_expressions + */ +class RegexConverter { + public: + explicit RegexConverter(const std::string& regex) : regex_(regex) { + if (!regex.empty()) { + // ParseUTF8 takes a C string and stops at the first NUL, so a regex + // containing one would either yield an empty codepoint vector (leading + // NUL -- an out-of-bounds read below) or be silently truncated to the + // prefix. A NUL is never meaningful in a regex, so reject it here. + if (regex.find('\0') != std::string::npos) { + XGRAMMAR_LOG(FATAL) << "The regex must not contain null characters."; + XGRAMMAR_UNREACHABLE(); + } + regex_codepoints_ = ParseUTF8(regex_.c_str(), false); + if (regex_codepoints_[0] == kInvalidUTF8) { + XGRAMMAR_LOG(FATAL) << "The regex is not a valid UTF-8 string."; + XGRAMMAR_UNREACHABLE(); + } + } + regex_codepoints_.push_back(0); // Add a null terminator + } + std::string Convert(); + + private: + /** + * \brief Add a segment string to the result EBNF string. It especially adds a space if needed + * and add_space is true. + */ + void AddEBNFSegment(const std::string& element); + + [[noreturn]] void RaiseError(const std::string& message); + void RaiseWarning(const std::string& message); + + std::string HandleCharacterClass(); + std::string HandleRepetitionRange(); + std::string HandleCharEscape(); + std::string HandleEscape(); + std::string HandleEscapeInCharClass(); + /** + * \brief Handle group modifier. The general format is "(?" + modifier + content + ")". E.g. + * "(?:abc)" is a non-capturing group. + */ + void HandleGroupModifier(); + + std::string regex_; + std::vector regex_codepoints_; + TCodepoint* start_; + TCodepoint* current_; + TCodepoint* end_; + std::string result_ebnf_; + int parenthesis_level_ = 0; +}; + +void RegexConverter::AddEBNFSegment(const std::string& element) { + if (!result_ebnf_.empty()) { + result_ebnf_ += ' '; + } + result_ebnf_ += element; +} + +void RegexConverter::RaiseError(const std::string& message) { + XGRAMMAR_LOG(FATAL) << "Regex parsing error at position " << current_ - start_ + 1 << ": " + << message; + XGRAMMAR_UNREACHABLE(); +} + +void RegexConverter::RaiseWarning(const std::string& message) { + XGRAMMAR_LOG(WARNING) << "Regex parsing warning at position " << current_ - start_ + 1 << ": " + << message; +} + +std::string RegexConverter::HandleCharacterClass() { + std::string char_class = "["; + ++current_; + if (*current_ == ']') { + RaiseError("Empty character class is not allowed in regex."); + } + while (*current_ != ']' && current_ != end_) { + if (*current_ == '\\') { + char_class += HandleEscapeInCharClass(); + } else { + char_class += CharToUTF8(*current_); + ++current_; + } + } + if (current_ == end_) { + RaiseError("Unclosed '['"); + } + char_class += ']'; + ++current_; + return char_class; +} + +// {x}: Match exactly x occurrences of the preceding regular expression. +// {x,} +// {x,y} +std::string RegexConverter::HandleRepetitionRange() { + std::string result = "{"; + ++current_; + if (!isdigit(*current_)) { + RaiseError("Invalid repetition count."); + } + while (isdigit(*current_)) { + result += static_cast(*current_); + ++current_; + } + if (*current_ != ',' && *current_ != '}') { + RaiseError("Invalid repetition count."); + } + result += static_cast(*current_); + ++current_; + if (current_[-1] == '}') { + // Matches {x} + return result; + } + if (!isdigit(*current_) && *current_ != '}') { + RaiseError("Invalid repetition count."); + } + while (isdigit(*current_)) { + result += static_cast(*current_); + ++current_; + } + if (*current_ != '}') { + RaiseError("Invalid repetition count."); + } + result += '}'; + ++current_; + return result; +} + +std::string RegexConverter::HandleCharEscape() { + // clang-format off + static const std::unordered_map CUSTOM_ESCAPE_MAP = { + {'^', '^'}, {'$', '$'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, {'\\', '\\'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'}, {'-', '-'} + }; + // clang-format on + if (end_ - current_ < 2 || (current_[1] == 'u' && end_ - current_ < 5) || + (current_[1] == 'x' && end_ - current_ < 4) || (current_[1] == 'c' && end_ - current_ < 3)) { + RaiseError("Escape sequence is not finished."); + } + auto [codepoint, len] = ParseNextEscaped(current_, CUSTOM_ESCAPE_MAP); + if (codepoint != CharHandlingError::kInvalidEscape) { + current_ += len; + return EscapeString(codepoint); + } else if (current_[1] == 'u' && current_[2] == '{') { + current_ += 3; + int len = 0; + TCodepoint value = 0; + while (HexCharToInt(current_[len]) != -1 && len <= 6) { + value = value * 16 + HexCharToInt(current_[len]); + ++len; + } + if (len == 0 || len > 6 || current_[len] != '}') { + RaiseError("Invalid Unicode escape sequence."); + } + current_ += len + 1; + return EscapeString(value); + } else if (current_[1] == 'c') { + current_ += 2; + if (!std::isalpha(*current_)) { + RaiseError("Invalid control character escape sequence."); + } + ++current_; + return EscapeString((*(current_ - 1)) % 32); + } else { + RaiseWarning( + "Escape sequence '\\" + EscapeString(current_[1]) + + "' is not recognized. The character itself will be matched" + ); + current_ += 2; + return EscapeString(current_[-1]); + } +} + +std::string RegexConverter::HandleEscapeInCharClass() { + if (end_ - current_ < 2) { + RaiseError("Escape sequence is not finished."); + } + if (current_[1] == 'd') { + current_ += 2; + return "0-9"; + } else if (current_[1] == 'D') { + current_ += 2; + return R"(\x00-\x2F\x3A-\U0010FFFF)"; + } else if (current_[1] == 'w') { + current_ += 2; + return "a-zA-Z0-9_"; + } else if (current_[1] == 'W') { + current_ += 2; + return R"(\x00-\x2F\x3A-\x40\x5B-\x5E\x60\x7B-\U0010FFFF)"; + } else if (current_[1] == 's') { + current_ += 2; + return R"(\f\n\r\t\v\u0020\u00a0)"; + } else if (current_[1] == 'S') { + current_ += 2; + return R"(\x00-\x08\x0E-\x1F\x21-\x9F\xA1-\U0010FFFF)"; + } else { + auto res = HandleCharEscape(); + if (res == "]" || res == "-") { + return "\\" + res; + } else { + return res; + } + } +} + +std::string RegexConverter::HandleEscape() { + // clang-format off + static const std::unordered_map CUSTOM_ESCAPE_MAP = { + {'^', '^'}, {'$', '$'}, {'.', '.'}, {'*', '*'}, {'+', '+'}, {'?', '?'}, {'\\', '\\'}, + {'(', '('}, {')', ')'}, {'[', '['}, {']', ']'}, {'{', '{'}, {'}', '}'}, {'|', '|'}, + {'/', '/'} + }; + // clang-format on + if (end_ - current_ < 2) { + RaiseError("Escape sequence is not finished."); + } + if (current_[1] == 'd') { + current_ += 2; + return "[0-9]"; + } else if (current_[1] == 'D') { + current_ += 2; + return "[^0-9]"; + } else if (current_[1] == 'w') { + current_ += 2; + return "[a-zA-Z0-9_]"; + } else if (current_[1] == 'W') { + current_ += 2; + return "[^a-zA-Z0-9_]"; + } else if (current_[1] == 's') { + current_ += 2; + return R"([\f\n\r\t\v\u0020\u00a0])"; + } else if (current_[1] == 'S') { + current_ += 2; + return R"([^\f\n\r\t\v\u0020\u00a0])"; + } else if ((current_[1] >= '1' && current_[1] <= '9') || current_[1] == 'k') { + RaiseError("Backreference is not supported yet."); + } else if (current_[1] == 'p' || current_[1] == 'P') { + RaiseError("Unicode character class escape sequence is not supported yet."); + } else if (current_[1] == 'b' || current_[1] == 'B') { + RaiseError("Word boundary is not supported yet."); + } else { + return "\"" + HandleCharEscape() + "\""; + } +} + +void RegexConverter::HandleGroupModifier() { + if (current_ == end_) { + RaiseError("Group modifier is not finished."); + } + if (*current_ == ':') { + // Non-capturing group. + ++current_; + } else if (*current_ == '=' || *current_ == '!') { + // Positive or negative lookahead. + RaiseError("Lookahead is not supported yet."); + } else if (*current_ == '<' && current_ + 1 != end_ && + (current_[1] == '=' || current_[1] == '!')) { + // Positive or negative lookbehind. + RaiseError("Lookbehind is not supported yet."); + } else if (*current_ == '<') { + ++current_; + while (current_ != end_ && isalpha(*current_)) { + ++current_; + } + if (current_ == end_ || *current_ != '>') { + RaiseError("Invalid named capturing group."); + } + // Just ignore the named of the group. + ++current_; + } else { + // Group modifier flag. + RaiseError("Group modifier flag is not supported yet."); + } +} + +std::string RegexConverter::Convert() { + start_ = regex_codepoints_.data(); + current_ = start_; + end_ = start_ + regex_codepoints_.size() - 1; + bool is_empty = true; + while (current_ != end_) { + if (*current_ == '^') { + if (current_ != start_) { + RaiseWarning( + "'^' should be at the start of the regex, but found in the middle. It is ignored." + ); + } + ++current_; + } else if (*current_ == '$') { + if (current_ != end_ - 1) { + RaiseWarning( + "'$' should be at the end of the regex, but found in the middle. It is ignored." + ); + } + ++current_; + } else if (*current_ == '[') { + is_empty = false; + AddEBNFSegment(HandleCharacterClass()); + } else if (*current_ == '(') { + is_empty = false; + ++current_; + ++parenthesis_level_; + AddEBNFSegment("("); + if (current_ != end_ && *current_ == '?') { + ++current_; + HandleGroupModifier(); + } + } else if (*current_ == ')') { + is_empty = false; + if (parenthesis_level_ == 0) { + RaiseError("Unmatched ')'"); + } + // Empty alternative before ')' (e.g. "(a|)" or "(a|$)"): emit "" so it isn't a bare '|'. + if (!result_ebnf_.empty() && result_ebnf_.back() == '|') { + AddEBNFSegment("\"\""); + } + --parenthesis_level_; + AddEBNFSegment(")"); + ++current_; + } else if (*current_ == '*' || *current_ == '+' || *current_ == '?') { + is_empty = false; + result_ebnf_ += static_cast(*current_); + ++current_; + if (current_ != end_ && *current_ == '?') { + // Ignore the non-greedy modifier because our grammar handles all repetition numbers + // non-deterministically. + ++current_; + } + if (current_ != end_ && + (*current_ == '{' || *current_ == '*' || *current_ == '+' || *current_ == '?')) { + RaiseError("Two consecutive repetition modifiers are not allowed."); + } + } else if (*current_ == '{') { + is_empty = false; + result_ebnf_ += HandleRepetitionRange(); + if (current_ != end_ && *current_ == '?') { + // Still ignore the non-greedy modifier. + ++current_; + } + if (current_ != end_ && + (*current_ == '{' || *current_ == '*' || *current_ == '+' || *current_ == '?')) { + RaiseError("Two consecutive repetition modifiers are not allowed."); + } + } else if (*current_ == '|') { + is_empty = false; + // Empty alternative before '|': emit "" so there's no bare '|' on the left. + // Covers leading ("^$|abc"), consecutive ("a||b") and group-start ("(|a)") cases. + if (result_ebnf_.empty() || result_ebnf_.back() == '|' || result_ebnf_.back() == '(') { + AddEBNFSegment("\"\""); + } + AddEBNFSegment("|"); + ++current_; + } else if (*current_ == '\\') { + is_empty = false; + AddEBNFSegment(HandleEscape()); + } else if (*current_ == '.') { + is_empty = false; + AddEBNFSegment(R"([\u0000-\U0010FFFF])"); + ++current_; + } else { + is_empty = false; + // Non-special characters are matched literally. + AddEBNFSegment("\"" + EscapeString(*current_) + "\""); + ++current_; + } + } + if (parenthesis_level_ != 0) { + RaiseError("The parenthesis is not closed."); + } + // Trailing empty alternative, e.g. "abc|": emit "" so it doesn't end with a bare '|'. + if (!result_ebnf_.empty() && result_ebnf_.back() == '|') { + AddEBNFSegment("\"\""); + } + if (is_empty) { + AddEBNFSegment("\"\""); + } + return result_ebnf_; +} + +std::string RegexToEBNF(const std::string& regex, bool with_rule_name) { + RegexConverter converter(regex); + if (with_rule_name) { + return "root ::= " + converter.Convert() + "\n"; + } else { + return converter.Convert(); + } +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/regex_converter.h b/zig/vendor/xgrammar/cpp/regex_converter.h new file mode 100644 index 0000000000..fcb3442ac1 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/regex_converter.h @@ -0,0 +1,21 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/regex_converter.h + * \brief Convert a regex string to EBNF grammar string. + */ + +#ifndef XGRAMMAR_REGEX_CONVERTER_H_ +#define XGRAMMAR_REGEX_CONVERTER_H_ + +#include + +namespace xgrammar { + +/*! + * \brief Convert a regex string to EBNF grammar string. + */ +std::string RegexToEBNF(const std::string& regex, bool with_rule_name = true); + +} // namespace xgrammar + +#endif // XGRAMMAR_REGEX_CONVERTER_H_ diff --git a/zig/vendor/xgrammar/cpp/structural_tag.cc b/zig/vendor/xgrammar/cpp/structural_tag.cc new file mode 100644 index 0000000000..a92840cf0c --- /dev/null +++ b/zig/vendor/xgrammar/cpp/structural_tag.cc @@ -0,0 +1,2546 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/structural_tag.cc + */ +#include "structural_tag.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "grammar_builder.h" +#include "grammar_functor.h" +#include "grammar_impl.h" +#include "json_schema_converter.h" +#include "support/json_parse.h" +#include "support/logging.h" +#include "support/recursion_guard.h" +#include "support/utils.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/grammar.h" + +namespace xgrammar { + +// Short alias for the error type. +using ISTError = InvalidStructuralTagError; + +// Forward declaration for helpers that convert Format to picojson::value. +picojson::value FormatToJSONValue(const Format& format); + +picojson::value StringVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& string : vector) { + array.push_back(picojson::value(string)); + } + return picojson::value(std::move(array)); +} + +picojson::value FormatVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& format : vector) { + array.push_back(xgrammar::FormatToJSONValue(format)); + } + return picojson::value(std::move(array)); +} + +picojson::value TagVectorToJSONArray(const std::vector& vector) { + picojson::array array; + array.reserve(vector.size()); + for (const auto& tag : vector) { + array.push_back(tag.ToJSON()); + } + return picojson::value(std::move(array)); +} + +picojson::value IntOrStringVectorToJSONArray( + const std::vector>& vec +) { + picojson::array array; + array.reserve(vec.size()); + for (const auto& item : vec) { + if (std::holds_alternative(item)) { + array.push_back(picojson::value(static_cast(std::get(item)))); + } else { + array.push_back(picojson::value(std::get(item))); + } + } + return picojson::value(std::move(array)); +} + +/******************** Format To JSON ********************/ +std::string FormatToJSON(const Format& format) { return FormatToJSONValue(format).serialize(); } + +picojson::value FormatToJSONValue(const Format& format) { + return std::visit([&](auto&& arg) { return arg.ToJSON(); }, format); +} + +picojson::value ConstStringFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["value"] = picojson::value(value); + return picojson::value(std::move(obj)); +} + +picojson::value JSONSchemaFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::value schema_val; + if (ParseJSON(schema_val, json_schema).empty()) { + obj["json_schema"] = schema_val; + } else { + obj["json_schema"] = picojson::value(json_schema); + } + obj["style"] = picojson::value(style); + obj["any_order"] = picojson::value(any_order); + if (!excludes.empty()) { + obj["excludes"] = StringVectorToJSONArray(excludes); + } + if (max_whitespace_cnt.has_value()) { + obj["max_whitespace_cnt"] = picojson::value(static_cast(*max_whitespace_cnt)); + } else { + obj["max_whitespace_cnt"] = picojson::value(); + } + return picojson::value(std::move(obj)); +} + +picojson::value GrammarFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["grammar"] = picojson::value(grammar); + return picojson::value(std::move(obj)); +} + +picojson::value RegexFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["pattern"] = picojson::value(pattern); + return picojson::value(std::move(obj)); +} + +picojson::value AnyTextFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["excludes"] = StringVectorToJSONArray(excludes); + obj["detected_end_strs"] = StringVectorToJSONArray(detected_end_strs_); + if (max_tokens >= 0) { + obj["max_tokens"] = picojson::value(static_cast(max_tokens)); + } + if (max_chars >= 0) { + obj["max_chars"] = picojson::value(static_cast(max_chars)); + } + return picojson::value(std::move(obj)); +} + +// These two constructors are defined here rather than inline because instantiating +// vector against the still-incomplete Format variant is ill-formed under C++20. +SequenceFormat::SequenceFormat(std::vector elements) : elements(std::move(elements)) {} + +OrFormat::OrFormat(std::vector elements) : elements(std::move(elements)) {} + +picojson::value SequenceFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["elements"] = FormatVectorToJSONArray(elements); + return picojson::value(std::move(obj)); +} + +picojson::value OrFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["elements"] = FormatVectorToJSONArray(elements); + return picojson::value(std::move(obj)); +} + +picojson::value TagFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + if (std::holds_alternative(begin)) { + obj["begin"] = picojson::value(std::get(begin)); + } else { + obj["begin"] = std::get(begin).ToJSON(); + } + if (content) { + obj["content"] = FormatToJSONValue(*content); + } else { + obj["content"] = picojson::value(); + } + if (std::holds_alternative(end)) { + obj["end"] = std::get(end).ToJSON(); + } else { + const auto& end_strs = std::get>(end); + if (end_strs.size() == 1) { + obj["end"] = picojson::value(end_strs[0]); + } else { + obj["end"] = StringVectorToJSONArray(end_strs); + } + } + return picojson::value(std::move(obj)); +} + +picojson::value TokenFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + if (std::holds_alternative(token)) { + obj["token"] = picojson::value(static_cast(std::get(token))); + } else { + obj["token"] = picojson::value(std::get(token)); + } + return picojson::value(std::move(obj)); +} + +picojson::value ExcludeTokenFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + return picojson::value(std::move(obj)); +} + +picojson::value AnyTokensFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + if (max_tokens >= 0) { + obj["max_tokens"] = picojson::value(static_cast(max_tokens)); + } + return picojson::value(std::move(obj)); +} + +picojson::value TokenTriggeredTagsFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["trigger_tokens"] = IntOrStringVectorToJSONArray(trigger_tokens); + obj["tags"] = TagVectorToJSONArray(tags); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + return picojson::value(std::move(obj)); +} + +picojson::value TriggeredTagsFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["triggers"] = StringVectorToJSONArray(triggers); + obj["tags"] = TagVectorToJSONArray(tags); + obj["excludes"] = StringVectorToJSONArray(excludes); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + obj["detected_end_strs"] = StringVectorToJSONArray(detected_end_strs_); + return picojson::value(std::move(obj)); +} + +picojson::value TagsWithSeparatorFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["tags"] = TagVectorToJSONArray(tags); + obj["separator"] = picojson::value(separator); + obj["at_least_one"] = picojson::value(at_least_one); + obj["stop_after_first"] = picojson::value(stop_after_first); + return picojson::value(std::move(obj)); +} + +picojson::value OptionalFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value PlusFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value StarFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value RepeatFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + obj["min"] = picojson::value(static_cast(min)); + obj["max"] = picojson::value(static_cast(max)); + obj["content"] = FormatToJSONValue(*content); + return picojson::value(std::move(obj)); +} + +picojson::value DispatchFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::array rules_arr; + rules_arr.reserve(rules.size()); + for (const auto& pair : rules) { + picojson::array pair_arr; + pair_arr.push_back(picojson::value(pair.first)); + if (pair.second) { + pair_arr.push_back(FormatToJSONValue(*pair.second)); + } else { + pair_arr.push_back(picojson::value()); + } + rules_arr.push_back(picojson::value(std::move(pair_arr))); + } + obj["rules"] = picojson::value(std::move(rules_arr)); + obj["loop"] = picojson::value(loop); + obj["excludes"] = StringVectorToJSONArray(excludes); + return picojson::value(std::move(obj)); +} + +picojson::value TokenDispatchFormat::ToJSON() const { + picojson::object obj; + obj["type"] = picojson::value(type); + picojson::array rules_arr; + rules_arr.reserve(rules.size()); + for (const auto& pair : rules) { + picojson::array pair_arr; + if (std::holds_alternative(pair.first)) { + pair_arr.push_back(picojson::value(static_cast(std::get(pair.first)))); + } else { + pair_arr.push_back(picojson::value(std::get(pair.first))); + } + if (pair.second) { + pair_arr.push_back(FormatToJSONValue(*pair.second)); + } else { + pair_arr.push_back(picojson::value()); + } + rules_arr.push_back(picojson::value(std::move(pair_arr))); + } + obj["rules"] = picojson::value(std::move(rules_arr)); + obj["loop"] = picojson::value(loop); + obj["exclude_tokens"] = IntOrStringVectorToJSONArray(exclude_tokens); + return picojson::value(std::move(obj)); +} + +/************** StructuralTag Parser **************/ + +class StructuralTagParser { + public: + static Result FromJSON(const std::string& json); + + private: + Result ParseStructuralTag(const picojson::value& value); + + /*! + * \brief Parse a Format object from a JSON value. + * \param value The JSON value to parse. + * \return A Format object if the JSON is valid, otherwise an error message in std::runtime_error. + * \note The "type" field is checked in this function, and not checked in the Parse*Format + * functions. + */ + Result ParseFormat(const picojson::value& value); + Result ParseConstStringFormat(const picojson::object& value); + Result ParseJSONSchemaFormat( + const picojson::object& value, std::optional style_override = std::nullopt + ); + Result ParseAnyTextFormat(const picojson::object& value); + Result ParseGrammarFormat(const picojson::object& value); + Result ParseRegexFormat(const picojson::object& value); + Result ParseSequenceFormat(const picojson::object& value); + Result ParseOrFormat(const picojson::object& value); + /*! \brief ParseTagFormat with extra check for object and the type field. */ + Result ParseTagFormat(const picojson::value& value); + Result ParseTagFormat(const picojson::object& value); + Result ParseTriggeredTagsFormat(const picojson::object& value); + Result ParseTagsWithSeparatorFormat( + const picojson::object& value + ); + Result ParseOptionalFormat(const picojson::object& value); + Result ParsePlusFormat(const picojson::object& value); + Result ParseStarFormat(const picojson::object& value); + Result ParseRepeatFormat(const picojson::object& value); + Result ParseTokenFormat(const picojson::object& value); + Result ParseExcludeTokenFormat(const picojson::object& value); + Result ParseAnyTokensFormat(const picojson::object& value); + Result ParseTokenTriggeredTagsFormat( + const picojson::object& value + ); + Result ParseDispatchFormat(const picojson::object& value); + Result ParseTokenDispatchFormat(const picojson::object& value); + + int parse_format_recursion_depth_ = 0; +}; + +Result StructuralTagParser::FromJSON(const std::string& json) { + picojson::value value; + std::string err = ParseJSON(value, json); + if (!err.empty()) { + return ResultErr("Failed to parse JSON: " + err); + } + return Result::Convert( + StructuralTagParser().ParseStructuralTag(value) + ); +} + +Result StructuralTagParser::ParseStructuralTag(const picojson::value& value +) { + if (!value.is()) { + return ResultErr("Structural tag must be an object"); + } + const auto& obj = value.get(); + // The type field is optional but must be "structural_tag" if present. + if (obj.find("type") != obj.end()) { + if (!obj["type"].is() || obj["type"].get() != "structural_tag") { + return ResultErr("Structural tag's type must be a string \"structural_tag\""); + } + } + // The format field is required. + if (obj.find("format") == obj.end()) { + return ResultErr("Structural tag must have a format field"); + } + auto format = ParseFormat(obj["format"]); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + return ResultOk(std::move(format).Unwrap()); +} + +Result StructuralTagParser::ParseFormat(const picojson::value& value) { + RecursionGuard guard(&parse_format_recursion_depth_); + // The global recursion limit is far deeper than the native stack allows for the recursive passes + // over the format tree (a few KB per level, and Windows has a 1 MB main-thread stack), so cap + // the nesting of formats explicitly. + static constexpr int kMaxFormatDepth = 100; + if (parse_format_recursion_depth_ > kMaxFormatDepth) { + return ResultErr( + "Formats are nested deeper than " + std::to_string(kMaxFormatDepth) + " levels" + ); + } + if (!value.is()) { + return ResultErr("Format must be an object"); + } + const auto& obj = value.get(); + // If type is present, use it to determine the format. + if (obj.find("type") != obj.end()) { + if (!obj["type"].is()) { + return ResultErr("Format's type must be a string"); + } + auto type = obj["type"].get(); + if (type == "const_string") { + return Result::Convert(ParseConstStringFormat(obj)); + } else if (type == "json_schema") { + return Result::Convert(ParseJSONSchemaFormat(obj)); + } else if (type == "any_text") { + return Result::Convert(ParseAnyTextFormat(obj)); + } else if (type == "sequence") { + return Result::Convert(ParseSequenceFormat(obj)); + } else if (type == "or") { + return Result::Convert(ParseOrFormat(obj)); + } else if (type == "tag") { + return Result::Convert(ParseTagFormat(obj)); + } else if (type == "triggered_tags") { + return Result::Convert(ParseTriggeredTagsFormat(obj)); + } else if (type == "tags_with_separator") { + return Result::Convert(ParseTagsWithSeparatorFormat(obj)); + } else if (type == "optional") { + return Result::Convert(ParseOptionalFormat(obj)); + } else if (type == "plus") { + return Result::Convert(ParsePlusFormat(obj)); + } else if (type == "star") { + return Result::Convert(ParseStarFormat(obj)); + } else if (type == "repeat") { + return Result::Convert(ParseRepeatFormat(obj)); + } else if (type == "qwen_xml_parameter") { + return Result::Convert(ParseJSONSchemaFormat(obj, "qwen_xml")); + } else if (type == "grammar") { + return Result::Convert(ParseGrammarFormat(obj)); + } else if (type == "regex") { + return Result::Convert(ParseRegexFormat(obj)); + } else if (type == "token") { + return Result::Convert(ParseTokenFormat(obj)); + } else if (type == "exclude_token") { + return Result::Convert(ParseExcludeTokenFormat(obj)); + } else if (type == "any_tokens") { + return Result::Convert(ParseAnyTokensFormat(obj)); + } else if (type == "token_triggered_tags") { + return Result::Convert(ParseTokenTriggeredTagsFormat(obj)); + } else if (type == "dispatch") { + return Result::Convert(ParseDispatchFormat(obj)); + } else if (type == "token_dispatch") { + return Result::Convert(ParseTokenDispatchFormat(obj)); + } else { + return ResultErr("Format type not recognized: " + type); + } + } + + // If type is not present, try every format type one by one. Tag is prioritized. + auto tag_format = ParseTagFormat(obj); + if (!tag_format.IsErr()) { + return ResultOk(std::move(tag_format).Unwrap()); + } + auto const_string_format = ParseConstStringFormat(obj); + if (!const_string_format.IsErr()) { + return ResultOk(std::move(const_string_format).Unwrap()); + } + auto json_schema_format = ParseJSONSchemaFormat(obj); + if (!json_schema_format.IsErr()) { + return ResultOk(std::move(json_schema_format).Unwrap()); + } + auto any_text_format = ParseAnyTextFormat(obj); + if (!any_text_format.IsErr()) { + return ResultOk(std::move(any_text_format).Unwrap()); + } + auto sequence_format = ParseSequenceFormat(obj); + if (!sequence_format.IsErr()) { + return ResultOk(std::move(sequence_format).Unwrap()); + } + auto or_format = ParseOrFormat(obj); + if (!or_format.IsErr()) { + return ResultOk(std::move(or_format).Unwrap()); + } + auto triggered_tags_format = ParseTriggeredTagsFormat(obj); + if (!triggered_tags_format.IsErr()) { + return ResultOk(std::move(triggered_tags_format).Unwrap()); + } + auto tags_with_separator_format = ParseTagsWithSeparatorFormat(obj); + if (!tags_with_separator_format.IsErr()) { + return ResultOk(std::move(tags_with_separator_format).Unwrap()); + } + auto optional_format = ParseOptionalFormat(obj); + if (!optional_format.IsErr()) { + return ResultOk(std::move(optional_format).Unwrap()); + } + auto plus_format = ParsePlusFormat(obj); + if (!plus_format.IsErr()) { + return ResultOk(std::move(plus_format).Unwrap()); + } + auto star_format = ParseStarFormat(obj); + if (!star_format.IsErr()) { + return ResultOk(std::move(star_format).Unwrap()); + } + auto repeat_format = ParseRepeatFormat(obj); + if (!repeat_format.IsErr()) { + return ResultOk(std::move(repeat_format).Unwrap()); + } + auto tag_dispatch_format = ParseDispatchFormat(obj); + if (!tag_dispatch_format.IsErr()) { + return ResultOk(std::move(tag_dispatch_format).Unwrap()); + } + auto token_tag_dispatch_format = ParseTokenDispatchFormat(obj); + if (!token_tag_dispatch_format.IsErr()) { + return ResultOk(std::move(token_tag_dispatch_format).Unwrap()); + } + return ResultErr("Invalid format: " + value.serialize(false)); +} + +Result StructuralTagParser::ParseConstStringFormat( + const picojson::object& obj +) { + // value is required. + auto value_it = obj.find("value"); + if (value_it == obj.end() || !value_it->second.is()) { + return ResultErr("ConstString format must have a value field with a string"); + } + return ResultOk(value_it->second.get()); +} + +Result StructuralTagParser::ParseJSONSchemaFormat( + const picojson::object& obj, std::optional style_override +) { + // json_schema is required. + auto json_schema_it = obj.find("json_schema"); + if (json_schema_it == obj.end() || + !(json_schema_it->second.is() || json_schema_it->second.is())) { + return ResultErr( + "JSON schema format must have a json_schema field with a object or boolean value" + ); + } + std::string style = "json"; + if (style_override.has_value()) { + style = *style_override; + } else { + auto it = obj.find("style"); + if (it != obj.end() && it->second.is()) { + style = it->second.get(); + if (style != "json" && style != "qwen_xml" && style != "minimax_xml" && + style != "minimax_m3_xml" && style != "deepseek_xml" && style != "glm_xml" && + style != "cohere_xml" && style != "kimi_k3_xml" && style != "deepseek_v4_1_xml") { + return ResultErr( + "style must be \"json\", \"qwen_xml\", \"minimax_xml\", \"minimax_m3_xml\", " + "\"deepseek_xml\", \"glm_xml\", \"cohere_xml\", \"kimi_k3_xml\", or " + "\"deepseek_v4_1_xml\"" + ); + } + } + } + bool any_order = false; + auto any_order_it = obj.find("any_order"); + if (any_order_it != obj.end()) { + if (!any_order_it->second.is()) { + return ResultErr("any_order must be a boolean"); + } + any_order = any_order_it->second.get(); + } + std::optional max_whitespace_cnt = std::nullopt; + auto max_whitespace_cnt_it = obj.find("max_whitespace_cnt"); + if (max_whitespace_cnt_it != obj.end() && !max_whitespace_cnt_it->second.is()) { + if (!max_whitespace_cnt_it->second.is()) { + return ResultErr("max_whitespace_cnt must be an integer or null"); + } + max_whitespace_cnt = static_cast(max_whitespace_cnt_it->second.get()); + } + std::vector excludes; + if (auto it = obj.find("excludes"); it != obj.end()) { + if (!it->second.is()) { + return ResultErr( + "JSONSchema format's excludes must be an array of non-empty strings" + ); + } + for (const auto& value : it->second.get()) { + if (!value.is() || value.get().empty()) { + return ResultErr( + "JSONSchema format's excludes must be an array of non-empty strings" + ); + } + excludes.push_back(value.get()); + } + std::sort(excludes.begin(), excludes.end()); + excludes.erase(std::unique(excludes.begin(), excludes.end()), excludes.end()); + } + // here introduces a serialization/deserialization overhead; try to avoid it in the future. + return ResultOk( + json_schema_it->second.serialize(false), + style, + any_order, + max_whitespace_cnt, + std::move(excludes) + ); +} + +Result ParseOptionalBudget( + const picojson::object& obj, const std::string& field_name, const std::string& format_name +) { + auto it = obj.find(field_name); + if (it == obj.end() || it->second.is()) { + return ResultOk(-1); + } + if (!it->second.is()) { + return ResultErr( + field_name + " in " + format_name + " must be a non-negative 32-bit integer or null" + ); + } + double value = it->second.get(); + if (value < 0 || value > std::numeric_limits::max() || value != std::floor(value)) { + return ResultErr( + field_name + " in " + format_name + " must be a non-negative 32-bit integer or null" + ); + } + return ResultOk(static_cast(value)); +} + +Result StructuralTagParser::ParseAnyTextFormat(const picojson::object& obj +) { + auto max_tokens_result = ParseOptionalBudget(obj, "max_tokens", "any_text"); + if (max_tokens_result.IsErr()) { + return ResultErr(std::move(max_tokens_result).UnwrapErr()); + } + auto max_chars_result = ParseOptionalBudget(obj, "max_chars", "any_text"); + if (max_chars_result.IsErr()) { + return ResultErr(std::move(max_chars_result).UnwrapErr()); + } + int32_t max_tokens = std::move(max_tokens_result).Unwrap(); + int32_t max_chars = std::move(max_chars_result).Unwrap(); + + auto excluded_strs_it = obj.find("excludes"); + if (excluded_strs_it == obj.end()) { + if ((obj.find("type") == obj.end())) { + return ResultErr("Any text format should not have any fields other than type"); + } + return ResultOk(std::vector{}, max_tokens, max_chars); + } + if (!excluded_strs_it->second.is()) { + return ResultErr("AnyText format's excluded_strs field must be an array"); + } + const auto& excluded_strs_array = excluded_strs_it->second.get(); + std::vector excluded_strs; + excluded_strs.reserve(excluded_strs_array.size()); + for (const auto& excluded_str : excluded_strs_array) { + if (!excluded_str.is()) { + return ResultErr("AnyText format's excluded_strs array must contain strings"); + } + excluded_strs.push_back(excluded_str.get()); + } + return ResultOk(std::move(excluded_strs), max_tokens, max_chars); +} + +Result StructuralTagParser::ParseGrammarFormat(const picojson::object& obj +) { + // grammar is required. + auto grammar_it = obj.find("grammar"); + if (grammar_it == obj.end() || !grammar_it->second.is() || + grammar_it->second.get().empty()) { + return ResultErr("Grammar format must have a grammar field with a non-empty string"); + } + return ResultOk(grammar_it->second.get()); +} + +Result StructuralTagParser::ParseRegexFormat(const picojson::object& obj) { + // pattern is required. + auto pattern_it = obj.find("pattern"); + if (pattern_it == obj.end() || !pattern_it->second.is() || + pattern_it->second.get().empty()) { + return ResultErr("Regex format must have a pattern field with a non-empty string"); + } + return ResultOk(pattern_it->second.get()); +} + +Result StructuralTagParser::ParseSequenceFormat( + const picojson::object& obj +) { + // elements is required. + auto elements_it = obj.find("elements"); + if (elements_it == obj.end() || !elements_it->second.is()) { + return ResultErr("Sequence format must have an elements field with an array"); + } + const auto& elements_array = elements_it->second.get(); + std::vector elements; + elements.reserve(elements_array.size()); + for (const auto& element : elements_array) { + auto format = ParseFormat(element); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + elements.push_back(std::move(format).Unwrap()); + } + if (elements.size() == 0) { + return ResultErr("Sequence format must have at least one element"); + } + return ResultOk(std::move(elements)); +} + +Result StructuralTagParser::ParseOrFormat(const picojson::object& obj) { + // elements is required. + auto elements_it = obj.find("elements"); + if (elements_it == obj.end() || !elements_it->second.is()) { + return ResultErr("Or format must have an elements field with an array"); + } + const auto& elements_array = elements_it->second.get(); + std::vector elements; + elements.reserve(elements_array.size()); + for (const auto& element : elements_array) { + auto format = ParseFormat(element); + if (format.IsErr()) { + return ResultErr(std::move(format).UnwrapErr()); + } + elements.push_back(std::move(format).Unwrap()); + } + if (elements.size() == 0) { + return ResultErr("Or format must have at least one element"); + } + return ResultOk(std::move(elements)); +} + +Result StructuralTagParser::ParseTagFormat(const picojson::value& value) { + if (!value.is()) { + return ResultErr("Tag format must be an object"); + } + const auto& obj = value.get(); + if (obj.find("type") != obj.end() && + (!obj["type"].is() || obj["type"].get() != "tag")) { + return ResultErr("Tag format's type must be a string \"tag\""); + } + return ParseTagFormat(obj); +} + +Result StructuralTagParser::ParseTagFormat(const picojson::object& obj) { + // begin is required: string or TokenFormat object + auto begin_it = obj.find("begin"); + if (begin_it == obj.end()) { + return ResultErr("Tag format's begin field must be a string"); + } + std::variant begin; + if (begin_it->second.is()) { + begin = begin_it->second.get(); + } else if (begin_it->second.is()) { + auto tf = ParseTokenFormat(begin_it->second.get()); + if (tf.IsErr()) { + return ResultErr(std::move(tf).UnwrapErr()); + } + begin = std::move(tf).Unwrap(); + } else { + return ResultErr("Tag format's begin field must be a string"); + } + + // content is required. + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Tag format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + + // end is required: string, array of strings, or TokenFormat object + auto end_it = obj.find("end"); + if (end_it == obj.end()) { + return ResultErr("Tag format must have an end field"); + } + + std::variant, TokenFormat> end; + if (end_it->second.is()) { + end = std::vector{end_it->second.get()}; + } else if (end_it->second.is()) { + const auto& end_array = end_it->second.get(); + if (end_array.empty()) { + return ResultErr("Tag format's end array cannot be empty"); + } + std::vector end_strings; + for (const auto& item : end_array) { + if (!item.is()) { + return ResultErr("Tag format's end array must contain only strings"); + } + end_strings.push_back(item.get()); + } + end = std::move(end_strings); + } else if (end_it->second.is()) { + auto tf = ParseTokenFormat(end_it->second.get()); + if (tf.IsErr()) { + return ResultErr(std::move(tf).UnwrapErr()); + } + end = std::move(tf).Unwrap(); + } else { + return ResultErr("Tag format's end field must be a string or array of strings"); + } + + return ResultOk( + std::move(begin), std::make_shared(std::move(content).Unwrap()), std::move(end) + ); +} + +Result StructuralTagParser::ParseTriggeredTagsFormat( + const picojson::object& obj +) { + // triggers is required. + auto triggers_it = obj.find("triggers"); + if (triggers_it == obj.end() || !triggers_it->second.is()) { + return ResultErr("Triggered tags format must have a triggers field with an array"); + } + const auto& triggers_array = triggers_it->second.get(); + std::vector excluded_strs; + std::vector triggers; + triggers.reserve(triggers_array.size()); + for (const auto& trigger : triggers_array) { + if (!trigger.is() || trigger.get().empty()) { + return ResultErr("Triggered tags format's triggers must be non-empty strings"); + } + triggers.push_back(trigger.get()); + } + if (triggers.size() == 0) { + return ResultErr("Triggered tags format's triggers must be non-empty"); + } + // tags is required. + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("Triggered tags format must have a tags field with an array"); + } + const auto& tags_array = tags_it->second.get(); + std::vector tags; + tags.reserve(tags_array.size()); + for (const auto& tag : tags_array) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.size() == 0) { + return ResultErr("Triggered tags format's tags must be non-empty"); + } + // excludes is optional. + auto excludes_it = obj.find("excludes"); + if (excludes_it != obj.end()) { + if (!excludes_it->second.is()) { + return ResultErr("Triggered tags format should have a excludes field with an array" + ); + } + const auto& excludes_array = excludes_it->second.get(); + excluded_strs.reserve(excludes_array.size()); + for (const auto& excluded_str : excludes_array) { + if (!excluded_str.is() || excluded_str.get().empty()) { + return ResultErr("Triggered tags format's excluded_strs must be non-empty strings" + ); + } + excluded_strs.push_back(excluded_str.get()); + } + } + + // at_least_one is optional. + bool at_least_one = false; + auto at_least_one_it = obj.find("at_least_one"); + if (at_least_one_it != obj.end()) { + if (!at_least_one_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = at_least_one_it->second.get(); + } + // stop_after_first is optional. + bool stop_after_first = false; + auto stop_after_first_it = obj.find("stop_after_first"); + if (stop_after_first_it != obj.end()) { + if (!stop_after_first_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = stop_after_first_it->second.get(); + } + return ResultOk( + std::move(triggers), std::move(tags), std::move(excluded_strs), at_least_one, stop_after_first + ); +} + +Result StructuralTagParser::ParseTagsWithSeparatorFormat( + const picojson::object& obj +) { + // tags is required. + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("Tags with separator format must have a tags field with an array"); + } + const auto& tags_array = tags_it->second.get(); + std::vector tags; + tags.reserve(tags_array.size()); + for (const auto& tag : tags_array) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.size() == 0) { + return ResultErr("Tags with separator format's tags must be non-empty"); + } + // separator is required (can be empty string). + auto separator_it = obj.find("separator"); + if (separator_it == obj.end() || !separator_it->second.is()) { + return ResultErr("Tags with separator format's separator field must be a string"); + } + // at_least_one is optional. + bool at_least_one = false; + auto at_least_one_it = obj.find("at_least_one"); + if (at_least_one_it != obj.end()) { + if (!at_least_one_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = at_least_one_it->second.get(); + } + // stop_after_first is optional. + bool stop_after_first = false; + auto stop_after_first_it = obj.find("stop_after_first"); + if (stop_after_first_it != obj.end()) { + if (!stop_after_first_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = stop_after_first_it->second.get(); + } + return ResultOk( + std::move(tags), separator_it->second.get(), at_least_one, stop_after_first + ); +} + +Result StructuralTagParser::ParseOptionalFormat( + const picojson::object& obj +) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Optional format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParsePlusFormat(const picojson::object& obj) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Plus format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParseStarFormat(const picojson::object& obj) { + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Star format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk(std::make_shared(std::move(content).Unwrap())); +} + +Result StructuralTagParser::ParseTokenFormat(const picojson::object& obj) { + auto token_it = obj.find("token"); + if (token_it == obj.end()) { + return ResultErr("TokenFormat must have a token field"); + } + if (token_it->second.is()) { + double d = token_it->second.get(); + if (d != static_cast(static_cast(d))) { + return ResultErr("Token ID must be an integer"); + } + int32_t id = static_cast(d); + if (id < 0) { + return ResultErr("Token ID must be non-negative"); + } + return ResultOk(std::variant(id)); + } else if (token_it->second.is()) { + auto s = token_it->second.get(); + if (s.empty()) { + return ResultErr("Token string must be non-empty"); + } + return ResultOk(std::variant(std::move(s))); + } + return ResultErr("TokenFormat's token must be an integer or string"); +} + +Result>, ISTError> ParseIntOrStringArray( + const picojson::value& val, const std::string& field_name +) { + std::vector> result; + if (!val.is()) { + return ResultErr(field_name + " must be an array"); + } + for (const auto& v : val.get()) { + if (v.is()) { + double d = v.get(); + if (d != static_cast(static_cast(d))) { + return ResultErr(field_name + " elements must be integers, not floats"); + } + int32_t id = static_cast(d); + if (id < 0) { + return ResultErr( + field_name + " elements must be non-negative integers or strings" + ); + } + result.emplace_back(std::in_place_type, id); + } else if (v.is()) { + auto s = v.get(); + if (s.empty()) { + return ResultErr(field_name + " string elements must be non-empty"); + } + result.push_back(std::move(s)); + } else { + return ResultErr(field_name + " elements must be integers or strings"); + } + } + return ResultOk(std::move(result)); +} + +Result StructuralTagParser::ParseExcludeTokenFormat( + const picojson::object& obj +) { + std::vector> exclude_tokens; + auto it = obj.find("exclude_tokens"); + if (it != obj.end()) { + auto parsed = ParseIntOrStringArray(it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + return ResultOk(std::move(exclude_tokens)); +} + +Result StructuralTagParser::ParseAnyTokensFormat( + const picojson::object& obj +) { + std::vector> exclude_tokens; + auto it = obj.find("exclude_tokens"); + if (it != obj.end()) { + auto parsed = ParseIntOrStringArray(it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + auto max_tokens_result = ParseOptionalBudget(obj, "max_tokens", "any_tokens"); + if (max_tokens_result.IsErr()) { + return ResultErr(std::move(max_tokens_result).UnwrapErr()); + } + return ResultOk( + std::move(exclude_tokens), std::move(max_tokens_result).Unwrap() + ); +} + +Result StructuralTagParser::ParseTokenTriggeredTagsFormat( + const picojson::object& obj +) { + // trigger_tokens is required + auto triggers_it = obj.find("trigger_tokens"); + if (triggers_it == obj.end()) { + return ResultErr("TokenTriggeredTagsFormat must have a trigger_tokens field"); + } + auto triggers = ParseIntOrStringArray(triggers_it->second, "trigger_tokens"); + if (triggers.IsErr()) { + return ResultErr(std::move(triggers).UnwrapErr()); + } + auto trigger_tokens = std::move(triggers).Unwrap(); + if (trigger_tokens.empty()) { + return ResultErr("trigger_tokens must be non-empty"); + } + + // tags is required + auto tags_it = obj.find("tags"); + if (tags_it == obj.end() || !tags_it->second.is()) { + return ResultErr("TokenTriggeredTagsFormat must have a tags field with an array"); + } + std::vector tags; + for (const auto& tag : tags_it->second.get()) { + auto tag_format = ParseTagFormat(tag); + if (tag_format.IsErr()) { + return ResultErr(std::move(tag_format).UnwrapErr()); + } + tags.push_back(std::move(tag_format).Unwrap()); + } + if (tags.empty()) { + return ResultErr("TokenTriggeredTagsFormat tags must be non-empty"); + } + + // exclude_tokens is optional + std::vector> exclude_tokens; + auto excludes_it = obj.find("exclude_tokens"); + if (excludes_it != obj.end()) { + auto parsed = ParseIntOrStringArray(excludes_it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + + bool at_least_one = false; + auto alo_it = obj.find("at_least_one"); + if (alo_it != obj.end()) { + if (!alo_it->second.is()) { + return ResultErr("at_least_one must be a boolean"); + } + at_least_one = alo_it->second.get(); + } + + bool stop_after_first = false; + auto saf_it = obj.find("stop_after_first"); + if (saf_it != obj.end()) { + if (!saf_it->second.is()) { + return ResultErr("stop_after_first must be a boolean"); + } + stop_after_first = saf_it->second.get(); + } + + return ResultOk( + std::move(trigger_tokens), + std::move(tags), + std::move(exclude_tokens), + at_least_one, + stop_after_first + ); +} + +Result StructuralTagParser::ParseDispatchFormat( + const picojson::object& obj +) { + auto rules_it = obj.find("rules"); + if (rules_it == obj.end() || !rules_it->second.is()) { + return ResultErr("TagDispatch format must have a rules field with an array"); + } + const auto& rules_array = rules_it->second.get(); + if (rules_array.empty()) { + return ResultErr("TagDispatch format rules must be non-empty"); + } + std::vector>> rules; + rules.reserve(rules_array.size()); + for (const auto& item : rules_array) { + if (!item.is()) { + return ResultErr("TagDispatch pair must be a 2-element array"); + } + const auto& pair_arr = item.get(); + if (pair_arr.size() != 2) { + return ResultErr("TagDispatch pair must be a 2-element array"); + } + if (!pair_arr[0].is()) { + return ResultErr("TagDispatch pair first element must be a string"); + } + std::string trigger = pair_arr[0].get(); + auto content = ParseFormat(pair_arr[1]); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + rules.push_back({std::move(trigger), std::make_shared(std::move(content).Unwrap())}); + } + + bool loop = true; + auto loop_it = obj.find("loop"); + if (loop_it != obj.end()) { + if (!loop_it->second.is()) { + return ResultErr("loop must be a boolean"); + } + loop = loop_it->second.get(); + } + + std::vector excludes; + auto excludes_it = obj.find("excludes"); + if (excludes_it != obj.end()) { + if (!excludes_it->second.is()) { + return ResultErr("excludes must be an array"); + } + for (const auto& e : excludes_it->second.get()) { + if (!e.is() || e.get().empty()) { + return ResultErr("excludes must contain non-empty strings"); + } + excludes.push_back(e.get()); + } + } + + return ResultOk(std::move(rules), loop, std::move(excludes)); +} + +Result StructuralTagParser::ParseTokenDispatchFormat( + const picojson::object& obj +) { + auto rules_it = obj.find("rules"); + if (rules_it == obj.end() || !rules_it->second.is()) { + return ResultErr("TokenTagDispatch format must have a rules field with an array"); + } + const auto& rules_array = rules_it->second.get(); + if (rules_array.empty()) { + return ResultErr("TokenTagDispatch format rules must be non-empty"); + } + std::vector, std::shared_ptr>> rules; + rules.reserve(rules_array.size()); + for (const auto& item : rules_array) { + if (!item.is()) { + return ResultErr("TokenTagDispatch pair must be a 2-element array"); + } + const auto& pair_arr = item.get(); + if (pair_arr.size() != 2) { + return ResultErr("TokenTagDispatch pair must be a 2-element array"); + } + std::variant trigger; + if (pair_arr[0].is()) { + double d = pair_arr[0].get(); + if (d != static_cast(static_cast(d))) { + return ResultErr("Token ID must be an integer"); + } + trigger = static_cast(d); + } else if (pair_arr[0].is()) { + trigger = pair_arr[0].get(); + } else { + return ResultErr("TokenTagDispatch pair first element must be an integer or string" + ); + } + auto content = ParseFormat(pair_arr[1]); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + rules.push_back({std::move(trigger), std::make_shared(std::move(content).Unwrap())}); + } + + bool loop = true; + auto loop_it = obj.find("loop"); + if (loop_it != obj.end()) { + if (!loop_it->second.is()) { + return ResultErr("loop must be a boolean"); + } + loop = loop_it->second.get(); + } + + std::vector> exclude_tokens; + auto excludes_it = obj.find("exclude_tokens"); + if (excludes_it != obj.end()) { + auto parsed = ParseIntOrStringArray(excludes_it->second, "exclude_tokens"); + if (parsed.IsErr()) { + return ResultErr(std::move(parsed).UnwrapErr()); + } + exclude_tokens = std::move(parsed).Unwrap(); + } + + return ResultOk(std::move(rules), loop, std::move(exclude_tokens)); +} + +/************** StructuralTagTokenResolver **************/ + +class StructuralTagTokenResolver { + public: + static std::optional Resolve( + StructuralTag* structural_tag, const std::optional& tokenizer_info + ); + + private: + explicit StructuralTagTokenResolver(const std::optional& tokenizer_info) + : tokenizer_info_(tokenizer_info) {} + + std::optional ResolveFormat(Format* format); + std::optional ResolveTagFormat(TagFormat* tag); + std::optional ResolveTokenFormat(TokenFormat* tf); + std::optional ResolveIntOrStringVec( + const std::vector>& input, std::vector* output + ); + + const std::optional& tokenizer_info_; +}; + +std::optional StructuralTagTokenResolver::Resolve( + StructuralTag* structural_tag, const std::optional& tokenizer_info +) { + return StructuralTagTokenResolver(tokenizer_info).ResolveFormat(&structural_tag->format); +} + +std::optional StructuralTagTokenResolver::ResolveTokenFormat(TokenFormat* tf) { + if (tf->resolved_token_id_ >= 0) return std::nullopt; + if (!std::holds_alternative(tf->token)) return std::nullopt; + if (!tokenizer_info_) { + return ISTError("Token string resolution requires tokenizer_info"); + } + const auto& token_str = std::get(tf->token); + const auto& vocab = tokenizer_info_->GetDecodedVocab(); + for (int32_t i = 0; i < static_cast(vocab.size()); ++i) { + if (vocab[i] == token_str) { + tf->resolved_token_id_ = i; + return std::nullopt; + } + } + return ISTError("Token string \"" + token_str + "\" not found in vocabulary"); +} + +std::optional StructuralTagTokenResolver::ResolveIntOrStringVec( + const std::vector>& input, std::vector* output +) { + output->clear(); + output->reserve(input.size()); + for (const auto& item : input) { + if (std::holds_alternative(item)) { + output->push_back(std::get(item)); + } else { + if (!tokenizer_info_) { + return ISTError("Token string resolution requires tokenizer_info"); + } + const auto& s = std::get(item); + const auto& vocab = tokenizer_info_->GetDecodedVocab(); + bool found = false; + for (int32_t i = 0; i < static_cast(vocab.size()); ++i) { + if (vocab[i] == s) { + output->push_back(i); + found = true; + break; + } + } + if (!found) { + return ISTError("Token string \"" + s + "\" not found in vocabulary"); + } + } + } + return std::nullopt; +} + +std::optional StructuralTagTokenResolver::ResolveTagFormat(TagFormat* tag) { + if (std::holds_alternative(tag->begin)) { + auto err = ResolveTokenFormat(&std::get(tag->begin)); + if (err) return err; + } + if (std::holds_alternative(tag->end)) { + auto err = ResolveTokenFormat(&std::get(tag->end)); + if (err) return err; + } + return ResolveFormat(tag->content.get()); +} + +std::optional StructuralTagTokenResolver::ResolveFormat(Format* format) { + return std::visit( + [&](auto&& arg) -> std::optional { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return ResolveTokenFormat(&arg); + } else if constexpr (std::is_same_v) { + return ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_token_ids_); + } else if constexpr (std::is_same_v) { + return ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + } else if constexpr (std::is_same_v) { + auto err = ResolveIntOrStringVec(arg.trigger_tokens, &arg.resolved_trigger_token_ids_); + if (err) return err; + err = ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + if (err) return err; + for (auto& tag : arg.tags) { + err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + std::vector> trigger_tokens; + trigger_tokens.reserve(arg.rules.size()); + for (const auto& p : arg.rules) { + trigger_tokens.push_back(p.first); + } + auto err = ResolveIntOrStringVec(trigger_tokens, &arg.resolved_trigger_token_ids_); + if (err) return err; + err = ResolveIntOrStringVec(arg.exclude_tokens, &arg.resolved_exclude_token_ids_); + if (err) return err; + for (auto& p : arg.rules) { + if (p.second) { + auto e = ResolveFormat(p.second.get()); + if (e) return e; + } + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& p : arg.rules) { + if (p.second) { + auto err = ResolveFormat(p.second.get()); + if (err) return err; + } + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + return ResolveTagFormat(&arg); + } else if constexpr (std::is_same_v) { + for (auto& elem : arg.elements) { + auto err = ResolveFormat(&elem); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& elem : arg.elements) { + auto err = ResolveFormat(&elem); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& tag : arg.tags) { + auto err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v) { + for (auto& tag : arg.tags) { + auto err = ResolveTagFormat(&tag); + if (err) return err; + } + return std::nullopt; + } else if constexpr (std::is_same_v || std::is_same_v || + std::is_same_v) { + return ResolveFormat(arg.content.get()); + } else { + return std::nullopt; + } + }, + *format + ); +} + +Result StructuralTagParser::ParseRepeatFormat(const picojson::object& obj) { + auto min_it = obj.find("min"); + if (min_it == obj.end() || !min_it->second.is()) { + return ResultErr("Repeat format must have a min field (number)"); + } + auto max_it = obj.find("max"); + if (max_it == obj.end() || !max_it->second.is()) { + return ResultErr("Repeat format must have a max field (number)"); + } + int64_t min = min_it->second.get(); + int64_t max = max_it->second.get(); + int32_t max_value_int32 = std::numeric_limits::max(); + if (max >= 0 && min > max) { + return ResultErr("Repeat min must be <= max"); + } + if (min < 0) { + return ResultErr("Repeat min must be >= 0"); + } + if (max < -1) { + return ResultErr("Repeat max must be -1 (unbounded) or >= 0"); + } + if (max > static_cast(max_value_int32)) { + XGRAMMAR_LOG(WARNING) << "Repeat max is too large, will be set as not limited"; + max = -1; // -1 means unlimited + } + if (min > static_cast(max_value_int32)) { + return ResultErr( + "Repeat min is too large, must be <= " + std::to_string(max_value_int32) + ); + } + auto content_it = obj.find("content"); + if (content_it == obj.end()) { + return ResultErr("Repeat format must have a content field"); + } + auto content = ParseFormat(content_it->second); + if (content.IsErr()) { + return ResultErr(std::move(content).UnwrapErr()); + } + return ResultOk( + static_cast(min), + static_cast(max), + std::make_shared(std::move(content).Unwrap()) + ); +} + +/************** StructuralTag Analyzer **************/ + +/*! + * \brief Analyze a StructuralTag and extract useful information for conversion to Grammar. + */ +class StructuralTagAnalyzer { + public: + static std::optional Analyze(StructuralTag* structural_tag); + + private: + /*! \brief A variant that can hold the pointer of any Format types. */ + using FormatPtrVariant = std::variant< + ConstStringFormat*, + JSONSchemaFormat*, + AnyTextFormat*, + GrammarFormat*, + RegexFormat*, + SequenceFormat*, + OrFormat*, + TagFormat*, + TriggeredTagsFormat*, + TagsWithSeparatorFormat*, + OptionalFormat*, + PlusFormat*, + StarFormat*, + RepeatFormat*, + TokenFormat*, + ExcludeTokenFormat*, + AnyTokensFormat*, + TokenTriggeredTagsFormat*, + DispatchFormat*, + TokenDispatchFormat*>; + + // Call this if we have a pointer to a Format. + std::optional Visit(Format* format); + // Call this if we have a pointer to a variant of Format. + std::optional Visit(FormatPtrVariant format); + + // The following is dispatched from Visit. Don't call them directly because they don't handle + // stack logics. + std::optional VisitSub(ConstStringFormat* format); + std::optional VisitSub(JSONSchemaFormat* format); + std::optional VisitSub(AnyTextFormat* format); + std::optional VisitSub(GrammarFormat* format); + std::optional VisitSub(RegexFormat* format); + std::optional VisitSub(SequenceFormat* format); + std::optional VisitSub(OrFormat* format); + std::optional VisitSub(TagFormat* format); + std::optional VisitSub(TriggeredTagsFormat* format); + std::optional VisitSub(TagsWithSeparatorFormat* format); + std::optional VisitSub(OptionalFormat* format); + std::optional VisitSub(PlusFormat* format); + std::optional VisitSub(StarFormat* format); + std::optional VisitSub(TokenFormat* format); + std::optional VisitSub(ExcludeTokenFormat* format); + std::optional VisitSub(AnyTokensFormat* format); + std::optional VisitSub(TokenTriggeredTagsFormat* format); + std::optional VisitSub(RepeatFormat* format); + std::optional VisitSub(DispatchFormat* format); + std::optional VisitSub(TokenDispatchFormat* format); + + std::vector DetectEndStrings(); + std::vector DetectEndTokenIds(); + bool IsUnlimited(const Format& format); + bool IsExcluded(const Format& format); + + int visit_format_recursion_depth_ = 0; + std::vector stack_; +}; + +std::optional StructuralTagAnalyzer::Analyze(StructuralTag* structural_tag) { + return StructuralTagAnalyzer().Visit(&structural_tag->format); +} + +std::vector StructuralTagAnalyzer::DetectEndStrings() { + for (int i = static_cast(stack_.size()) - 1; i >= 0; --i) { + auto& format = stack_[i]; + if (std::holds_alternative(format)) { + auto* tag = std::get(format); + if (std::holds_alternative>(tag->end)) { + return std::get>(tag->end); + } + return {}; // TokenFormat end — propagated via DetectEndTokenIds + } + } + return {}; +} + +std::vector StructuralTagAnalyzer::DetectEndTokenIds() { + for (int i = static_cast(stack_.size()) - 1; i >= 0; --i) { + auto& format = stack_[i]; + if (std::holds_alternative(format)) { + auto* tag = std::get(format); + if (std::holds_alternative(tag->end)) { + auto& tf = std::get(tag->end); + return {tf.resolved_token_id_}; + } + return {}; + } + } + return {}; +} + +bool StructuralTagAnalyzer::IsUnlimited(const Format& format) { + return std::visit( + [&](auto&& arg) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return arg.is_unlimited_; + } else if constexpr (std::is_same_v) { + return arg.is_unlimited_; + } else if constexpr (std::is_same_v) { + return IsUnlimited(*arg.content); + } else if constexpr (std::is_same_v || std::is_same_v) { + return true; + } else if constexpr (std::is_same_v) { + return arg.max == -1 || (arg.max != 0 && IsUnlimited(*arg.content)); + } else { + return false; + } + }, + format + ); +} + +bool StructuralTagAnalyzer::IsExcluded(const Format& format) { + return std::visit( + [&](auto&& arg) -> bool { + using T = std::decay_t; + if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else if constexpr (std::is_same_v) { + return !arg.excludes.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else if constexpr (std::is_same_v) { + return !arg.exclude_tokens.empty(); + } else { + return false; + } + }, + format + ); +} + +std::optional StructuralTagAnalyzer::Visit(Format* format) { + FormatPtrVariant format_ptr_variant = + std::visit([&](auto&& arg) -> FormatPtrVariant { return &arg; }, *format); + return Visit(format_ptr_variant); +} + +std::optional StructuralTagAnalyzer::Visit(FormatPtrVariant format) { + RecursionGuard guard(&visit_format_recursion_depth_); + + // Push format to stack + stack_.push_back(format); + + // Dispatch to the corresponding visit function + auto result = + std::visit([&](auto&& arg) -> std::optional { return VisitSub(arg); }, format); + + // Pop format from stack + stack_.pop_back(); + + return result; +} + +std::optional StructuralTagAnalyzer::VisitSub(ConstStringFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(JSONSchemaFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(AnyTextFormat* format) { + format->detected_end_strs_ = DetectEndStrings(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(GrammarFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(RegexFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(SequenceFormat* format) { + bool is_any_unlimited = false; + for (auto& element : format->elements) { + auto err = Visit(&element); + if (err.has_value()) { + return err; + } + is_any_unlimited |= IsUnlimited(element) && !IsExcluded(element); + } + format->is_unlimited_ = is_any_unlimited; + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(OrFormat* format) { + bool is_any_unlimited = false; + for (auto& element : format->elements) { + auto err = Visit(&element); + if (err.has_value()) { + return err; + } + is_any_unlimited |= IsUnlimited(element) && !IsExcluded(element); + } + format->is_unlimited_ = is_any_unlimited; + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TagFormat* format) { + auto err = Visit(format->content.get()); + if (err.has_value()) { + return err; + } + auto is_content_unlimited = IsUnlimited(*(format->content)); + if (is_content_unlimited) { + if (std::holds_alternative>(format->end)) { + const auto& ends = std::get>(format->end); + bool has_non_empty_end = false; + for (const auto& end_str : ends) { + if (!end_str.empty()) { + has_non_empty_end = true; + break; + } + } + if (!has_non_empty_end && !IsExcluded(*format->content)) { + return ISTError("When the content is unlimited, at least one end string must be non-empty"); + } + } + // TokenFormat end is always non-empty → no error needed + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TriggeredTagsFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + format->detected_end_strs_ = DetectEndStrings(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TagsWithSeparatorFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(OptionalFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(PlusFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(StarFormat* format) { + return Visit(format->content.get()); +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenFormat* format) { + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(ExcludeTokenFormat* format) { + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(AnyTokensFormat* format) { + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenTriggeredTagsFormat* format) { + for (auto& tag : format->tags) { + auto err = Visit(&tag); + if (err.has_value()) { + return err; + } + } + format->detected_end_token_ids_ = DetectEndTokenIds(); + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(DispatchFormat* format) { + for (auto& pair : format->rules) { + if (pair.second) { + auto err = Visit(pair.second.get()); + if (err.has_value()) return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(TokenDispatchFormat* format) { + for (auto& pair : format->rules) { + if (pair.second) { + auto err = Visit(pair.second.get()); + if (err.has_value()) return err; + } + } + return std::nullopt; +} + +std::optional StructuralTagAnalyzer::VisitSub(RepeatFormat* format) { + return Visit(format->content.get()); +} + +/************** StructuralTag to Grammar Converter **************/ + +class StructuralTagGrammarConverter { + public: + static Result Convert(const StructuralTag& structural_tag); + + private: + StructuralTagGrammarConverter() = default; + + /*! + * \brief Visit a Format and return the rule id of the added rule. + * \param format The Format to visit. + * \return The rule id of the added rule. If the visit fails, the error is returned. + * \note This method uses serialization to deduplicate identical formats. + */ + Result Visit(const Format& format); + Result VisitSub(const ConstStringFormat& format); + Result VisitSub(const JSONSchemaFormat& format); + Result VisitSub(const AnyTextFormat& format); + Result VisitSub(const GrammarFormat& format); + Result VisitSub(const RegexFormat& format); + Result VisitSub(const SequenceFormat& format); + Result VisitSub(const OrFormat& format); + Result VisitSub(const TagFormat& format); + Result VisitSub(const TriggeredTagsFormat& format); + Result VisitSub(const TagsWithSeparatorFormat& format); + Result VisitSub(const OptionalFormat& format); + Result VisitSub(const PlusFormat& format); + Result VisitSub(const StarFormat& format); + Result VisitSub(const TokenFormat& format); + Result VisitSub(const ExcludeTokenFormat& format); + Result VisitSub(const AnyTokensFormat& format); + Result VisitSub(const TokenTriggeredTagsFormat& format); + Result VisitSub(const RepeatFormat& format); + Result VisitSub(const DispatchFormat& format); + Result VisitSub(const TokenDispatchFormat& format); + Grammar AddRootRuleAndGetGrammar(int ref_rule_id); + void SetRuleBodyAndBudgets( + int32_t rule_id, int body_expr_id, int32_t max_tokens, int32_t max_chars = -1 + ); + + bool IsPrefix(const std::string& prefix, const std::string& full_str); + int BuildBeginExpr(const TagFormat& tag); + int BuildEndExpr(const TagFormat& tag); + + GrammarBuilder grammar_builder_; + + /*! + * \brief Cache from format serialization to rule id. + * This enables deduplication of identical formats to reduce grammar size. + */ + std::unordered_map serialization_to_rule_id_; +}; + +bool StructuralTagGrammarConverter::IsPrefix( + const std::string& prefix, const std::string& full_str +) { + return prefix.size() <= full_str.size() && + std::string_view(full_str).substr(0, prefix.size()) == prefix; +} + +Result StructuralTagGrammarConverter::Convert(const StructuralTag& structural_tag +) { + StructuralTagGrammarConverter converter; + auto result = converter.Visit(structural_tag.format); + if (result.IsErr()) { + return ResultErr(std::move(result).UnwrapErr()); + } + // Add a root rule + auto root_rule_id = std::move(result).Unwrap(); + return ResultOk(converter.AddRootRuleAndGetGrammar(root_rule_id)); +} + +Grammar StructuralTagGrammarConverter::AddRootRuleAndGetGrammar(int ref_rule_id) { + auto expr = grammar_builder_.AddRuleRef(ref_rule_id); + auto sequence_expr = grammar_builder_.AddSequence({expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr}); + auto root_rule_id = grammar_builder_.AddRuleWithHint("root", choices_expr); + return grammar_builder_.Get(root_rule_id); +} + +void StructuralTagGrammarConverter::SetRuleBodyAndBudgets( + int32_t rule_id, int body_expr_id, int32_t max_tokens, int32_t max_chars +) { + // A zero-token region is empty by definition. Materialize that in the grammar because -1 is + // also the parser's pre-first-token deadline sentinel. + if (max_tokens == 0) { + body_expr_id = grammar_builder_.AddChoices({grammar_builder_.AddEmptyStr()}); + } + grammar_builder_.UpdateRuleBody(rule_id, body_expr_id); + if (max_tokens >= 0) { + grammar_builder_.UpdateMaxTokens(rule_id, max_tokens); + } + if (max_chars >= 0) { + grammar_builder_.UpdateMaxChars(rule_id, max_chars); + } +} + +Result StructuralTagGrammarConverter::Visit(const Format& format) { + std::string fingerprint = FormatToJSONValue(format).serialize(); + + // Check if we've already processed an identical format + auto it = serialization_to_rule_id_.find(fingerprint); + if (it != serialization_to_rule_id_.end()) { + return ResultOk(it->second); + } + + // Process the format and cache the result + auto result = + std::visit([&](auto&& arg) -> Result { return VisitSub(arg); }, format); + if (result.IsOk()) { + int rule_id = std::move(result).Unwrap(); + serialization_to_rule_id_[fingerprint] = rule_id; + return ResultOk(rule_id); + } + return result; +} + +Result StructuralTagGrammarConverter::VisitSub(const ConstStringFormat& format) { + auto expr = format.value.empty() ? grammar_builder_.AddEmptyStr() + : grammar_builder_.AddByteString(format.value); + auto sequence_expr = grammar_builder_.AddSequence({expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr}); + return ResultOk(grammar_builder_.AddRuleWithHint("const_string", choices_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const JSONSchemaFormat& format) { + auto json_format = JSONFormatFromString(format.style); + if (!json_format.has_value()) { + return ResultErr("Unsupported parsing type: " + format.style); + } + // The whitespace cap comes from the JSONSchemaFormat node (per-tag). + auto sub_grammar = GrammarNormalizer::Apply(JSONSchemaToGrammar( + format.json_schema, + /*any_whitespace=*/true, + /*indent=*/std::nullopt, + /*separators=*/std::nullopt, + /*strict_mode=*/true, + /*max_whitespace_cnt=*/format.max_whitespace_cnt, + /*any_order=*/format.any_order, + /*json_format=*/*json_format, + /*excludes=*/format.excludes + )); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const GrammarFormat& format) { + auto sub_grammar = Grammar::FromEBNF(format.grammar); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const RegexFormat& format) { + auto sub_grammar = Grammar::FromRegex(format.pattern); + auto added_root_rule_id = SubGrammarAdder().Apply(&grammar_builder_, sub_grammar); + return ResultOk(added_root_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const AnyTextFormat& format) { + std::vector all_excludes = format.excludes; + for (const auto& s : format.detected_end_strs_) { + if (!s.empty()) { + all_excludes.push_back(s); + } + } + int body_expr_id; + if (!all_excludes.empty()) { + body_expr_id = + grammar_builder_.AddTagDispatch(Grammar::Impl::TagDispatch{{}, false, all_excludes}); + } else { + auto any_text_expr = grammar_builder_.AddCharacterClassStar({{0, 0x10FFFF}}, false); + auto sequence_expr = grammar_builder_.AddSequence({any_text_expr}); + body_expr_id = grammar_builder_.AddChoices({sequence_expr}); + } + int rule_id = grammar_builder_.AddEmptyRuleWithHint("any_text"); + SetRuleBodyAndBudgets(rule_id, body_expr_id, format.max_tokens, format.max_chars); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const SequenceFormat& format) { + std::vector rule_ref_ids; + rule_ref_ids.reserve(format.elements.size()); + for (const auto& element : format.elements) { + auto result = Visit(element); + if (result.IsErr()) { + return result; + } + int sub_rule_id = std::move(result).Unwrap(); + rule_ref_ids.push_back(grammar_builder_.AddRuleRef(sub_rule_id)); + } + auto expr = grammar_builder_.AddChoices({grammar_builder_.AddSequence(rule_ref_ids)}); + return ResultOk(grammar_builder_.AddRuleWithHint("sequence", expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const OrFormat& format) { + std::vector sequence_ids; + sequence_ids.reserve(format.elements.size()); + for (const auto& element : format.elements) { + auto result = Visit(element); + if (result.IsErr()) { + return result; + } + int sub_rule_id = std::move(result).Unwrap(); + auto rule_ref_expr = grammar_builder_.AddRuleRef(sub_rule_id); + sequence_ids.push_back(grammar_builder_.AddSequence({rule_ref_expr})); + } + auto expr = grammar_builder_.AddChoices(sequence_ids); + return ResultOk(grammar_builder_.AddRuleWithHint("or", expr)); +} + +int StructuralTagGrammarConverter::BuildBeginExpr(const TagFormat& tag) { + if (std::holds_alternative(tag.begin)) { + return grammar_builder_.AddByteString(std::get(tag.begin)); + } + return grammar_builder_.AddTokenSet({std::get(tag.begin).resolved_token_id_}); +} + +int StructuralTagGrammarConverter::BuildEndExpr(const TagFormat& tag) { + if (std::holds_alternative(tag.end)) { + return grammar_builder_.AddTokenSet({std::get(tag.end).resolved_token_id_}); + } + const auto& ends = std::get>(tag.end); + if (ends.size() == 1) { + return ends[0].empty() ? grammar_builder_.AddEmptyStr() + : grammar_builder_.AddByteString(ends[0]); + } + std::vector end_seq_ids; + for (const auto& s : ends) { + auto e = s.empty() ? grammar_builder_.AddEmptyStr() : grammar_builder_.AddByteString(s); + end_seq_ids.push_back(grammar_builder_.AddSequence({e})); + } + auto choice = grammar_builder_.AddChoices(end_seq_ids); + auto rule = grammar_builder_.AddRuleWithHint("tag_end", choice); + return grammar_builder_.AddRuleRef(rule); +} + +Result StructuralTagGrammarConverter::VisitSub(const TagFormat& format) { + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + auto sub_rule_id = std::move(result).Unwrap(); + auto begin_expr = BuildBeginExpr(format); + auto rule_ref_expr = grammar_builder_.AddRuleRef(sub_rule_id); + auto end_expr = BuildEndExpr(format); + + auto sequence_expr_id = grammar_builder_.AddSequence({begin_expr, rule_ref_expr, end_expr}); + auto choices_expr = grammar_builder_.AddChoices({sequence_expr_id}); + return ResultOk(grammar_builder_.AddRuleWithHint("tag", choices_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TriggeredTagsFormat& format) { + // Step 1. Visit all tags and add to grammar + std::vector> trigger_to_tag_ids(format.triggers.size()); + std::vector tag_content_rule_ids; + tag_content_rule_ids.reserve(format.tags.size()); + + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + if (!std::holds_alternative(tag.begin)) { + return ResultErr( + "Tags in triggered_tags must have a string begin, not a token format" + ); + } + const auto& tag_begin = std::get(tag.begin); + int matched_trigger_id = -1; + for (int it_trigger = 0; it_trigger < static_cast(format.triggers.size()); ++it_trigger) { + const auto& trigger = format.triggers[it_trigger]; + if (IsPrefix(trigger, tag_begin)) { + if (matched_trigger_id != -1) { + return ResultErr("One tag matches multiple triggers in a triggered tags format" + ); + } + matched_trigger_id = it_trigger; + } + } + if (matched_trigger_id == -1) { + return ResultErr("One tag does not match any trigger in a triggered tags format"); + } + trigger_to_tag_ids[matched_trigger_id].push_back(it_tag); + + auto result = Visit(*tag.content); + if (result.IsErr()) { + return result; + } + tag_content_rule_ids.push_back(std::move(result).Unwrap()); + } + + // Step 2. Special Case: at_least_one && stop_after_first. + if (format.at_least_one && format.stop_after_first) { + std::vector choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr_id = BuildBeginExpr(tag); + auto rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr_id = BuildEndExpr(tag); + choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_elements); + return ResultOk(grammar_builder_.AddRuleWithHint("triggered_tags", choice_expr_id)); + } + + // Step 3. Normal Case. + // Step 3.1 Get tag_rule_pairs. + std::vector> tag_rule_pairs; + for (int it_trigger = 0; it_trigger < static_cast(format.triggers.size()); ++it_trigger) { + const auto& trigger = format.triggers[it_trigger]; + std::vector choice_elements; + for (const auto& tag_id : trigger_to_tag_ids[it_trigger]) { + const auto& tag = format.tags[tag_id]; + const auto& tag_begin = std::get(tag.begin); + int begin_expr_id = grammar_builder_.AddByteString(tag_begin.substr(trigger.size())); + int rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[tag_id]); + int end_expr_id = BuildEndExpr(tag); + choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_elements); + auto sub_rule_id = grammar_builder_.AddRuleWithHint("triggered_tags_group", choice_expr_id); + tag_rule_pairs.push_back(std::make_pair(trigger, sub_rule_id)); + } + + // Step 3.2 Add TagDispatch. + int32_t rule_expr_id; + bool loop_after_dispatch = !format.stop_after_first; + std::vector all_excludes = format.excludes; + for (const auto& s : format.detected_end_strs_) { + if (!s.empty()) { + all_excludes.push_back(s); + } + } + rule_expr_id = grammar_builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{tag_rule_pairs, loop_after_dispatch, all_excludes} + ); + + // Step 3.3 Consider at_least_one + if (format.at_least_one) { + std::vector first_choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr_id = BuildBeginExpr(tag); + auto rule_ref_expr_id = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr_id = BuildEndExpr(tag); + first_choice_elements.push_back( + grammar_builder_.AddSequence({begin_expr_id, rule_ref_expr_id, end_expr_id}) + ); + } + auto first_choice_expr_id = grammar_builder_.AddChoices(first_choice_elements); + auto first_rule_id = + grammar_builder_.AddRuleWithHint("triggered_tags_first", first_choice_expr_id); + + auto tag_dispatch_rule_id = + grammar_builder_.AddRuleWithHint("triggered_tags_sub", rule_expr_id); + auto ref_first_rule_expr_id = grammar_builder_.AddRuleRef(first_rule_id); + auto ref_tag_dispatch_rule_expr_id = grammar_builder_.AddRuleRef(tag_dispatch_rule_id); + auto sequence_expr_id = + grammar_builder_.AddSequence({ref_first_rule_expr_id, ref_tag_dispatch_rule_expr_id}); + rule_expr_id = grammar_builder_.AddChoices({sequence_expr_id}); + } + + auto rule_id = grammar_builder_.AddRuleWithHint("triggered_tags", rule_expr_id); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const TagsWithSeparatorFormat& format +) { + // The grammar: + // Step 1. tags_rule: call tags + // tags_rule ::= tag1 | tag2 | ... | tagN + // Step 2. Special handling (stop_after_first is true): + // if at_least_one is false: + // root ::= tags_rule | "" + // if at_least_one is true: + // root ::= tags_rule + // Step 3. Normal handling (stop_after_first is false): + // if at_least_one is false: + // root ::= tags_rule tags_rule_sub | "" + // if at_least_one is true: + // root ::= tags_rule tags_rule_sub + // tags_rule_sub ::= sep tags_rule tags_rule_sub | "" + + // Step 1. Construct a rule representing any tag + std::vector choice_ids; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + auto tag_rule_id = Visit(format.tags[it_tag]); + if (tag_rule_id.IsErr()) { + return tag_rule_id; + } + auto tag_rule_ref_id = grammar_builder_.AddRuleRef(std::move(tag_rule_id).Unwrap()); + auto sequence_expr_id = grammar_builder_.AddSequence({tag_rule_ref_id}); + choice_ids.push_back(sequence_expr_id); + } + auto choice_expr_id = grammar_builder_.AddChoices(choice_ids); + auto all_tags_rule_id = + grammar_builder_.AddRuleWithHint("tags_with_separator_tags", choice_expr_id); + + auto all_tags_rule_ref_id = grammar_builder_.AddRuleRef(all_tags_rule_id); + + // Step 2. Special case (stop_after_first is true): + if (format.stop_after_first) { + int32_t rule_body_expr_id; + if (format.at_least_one) { + // root ::= tags_rule + rule_body_expr_id = + grammar_builder_.AddChoices({grammar_builder_.AddSequence({all_tags_rule_ref_id})}); + } else { + // root ::= tags_rule | "" + rule_body_expr_id = grammar_builder_.AddChoices( + {grammar_builder_.AddSequence({all_tags_rule_ref_id}), grammar_builder_.AddEmptyStr()} + ); + } + + auto rule_id = grammar_builder_.AddRuleWithHint("tags_with_separator", rule_body_expr_id); + return ResultOk(rule_id); + } + + // Step 3. Normal handling (stop_after_first is false): + // Step 3.1 Construct sub rule: sub ::= sep tags sub | "" + auto sub_rule_id = grammar_builder_.AddEmptyRuleWithHint("tags_with_separator_sub"); + + auto end_str_sequence_id = grammar_builder_.AddEmptyStr(); + + std::vector sub_sequence_elements; + if (!format.separator.empty()) { + sub_sequence_elements.push_back(grammar_builder_.AddByteString(format.separator)); + } + sub_sequence_elements.push_back(all_tags_rule_ref_id); + sub_sequence_elements.push_back(grammar_builder_.AddRuleRef(sub_rule_id)); + + auto sub_rule_body_id = grammar_builder_.AddChoices( + {grammar_builder_.AddSequence(sub_sequence_elements), end_str_sequence_id} + ); + grammar_builder_.UpdateRuleBody(sub_rule_id, sub_rule_body_id); + + // Step 3.2 Construct root rule + std::vector choices = { + grammar_builder_.AddSequence({all_tags_rule_ref_id, grammar_builder_.AddRuleRef(sub_rule_id)} + ), + }; + if (!format.at_least_one) { + choices.push_back(end_str_sequence_id); + } + auto rule_body_expr_id = grammar_builder_.AddChoices(choices); + auto rule_id = grammar_builder_.AddRuleWithHint("tags_with_separator", rule_body_expr_id); + return ResultOk(rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const OptionalFormat& format) { + // optional: 0 or 1 occurrence -> Choice(content, "") + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto expr = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref})} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("optional", expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const PlusFormat& format) { + // plus: 1 or more occurrences -> content content_star, where content_star = content content_star + // | "" + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("plus_star"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref, star_ref})} + ); + grammar_builder_.UpdateRuleBody(star_rule_id, star_body); + auto plus_expr = grammar_builder_.AddSequence({content_ref, star_ref}); + return ResultOk(grammar_builder_.AddRuleWithHint("plus", plus_expr)); +} + +Result StructuralTagGrammarConverter::VisitSub(const StarFormat& format) { + // star: 0 or more occurrences -> content_star, where content_star = content content_star | "" + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + auto content_ref = grammar_builder_.AddRuleRef(content_rule_id); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("star"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({content_ref, star_ref})} + ); + grammar_builder_.UpdateRuleBody(star_rule_id, star_body); + return ResultOk(grammar_builder_.AddRuleWithHint("star", star_ref)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenFormat& format) { + XGRAMMAR_DCHECK(format.resolved_token_id_ >= 0) + << "TokenFormat must be resolved before conversion"; + auto token_set_expr = grammar_builder_.AddTokenSet({format.resolved_token_id_}); + auto seq = grammar_builder_.AddSequence({token_set_expr}); + auto choices = grammar_builder_.AddChoices({seq}); + return ResultOk(grammar_builder_.AddRuleWithHint("token", choices)); +} + +Result StructuralTagGrammarConverter::VisitSub(const ExcludeTokenFormat& format) { + std::vector all_excludes = format.resolved_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + int expr = grammar_builder_.AddExcludeTokenSet(all_excludes); + auto seq = grammar_builder_.AddSequence({expr}); + auto choices = grammar_builder_.AddChoices({seq}); + return ResultOk(grammar_builder_.AddRuleWithHint("exclude_token", choices)); +} + +Result StructuralTagGrammarConverter::VisitSub(const AnyTokensFormat& format) { + std::vector all_excludes = format.resolved_exclude_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + int exclude_expr = grammar_builder_.AddExcludeTokenSet(all_excludes); + int exclude_seq = grammar_builder_.AddSequence({exclude_expr}); + int exclude_choices = grammar_builder_.AddChoices({exclude_seq}); + int inner_rule = grammar_builder_.AddRuleWithHint("any_tokens_inner", exclude_choices); + auto inner_ref = grammar_builder_.AddRuleRef(inner_rule); + auto star_rule_id = grammar_builder_.AddEmptyRuleWithHint("any_tokens"); + auto star_ref = grammar_builder_.AddRuleRef(star_rule_id); + auto star_body = grammar_builder_.AddChoices( + {grammar_builder_.AddEmptyStr(), grammar_builder_.AddSequence({inner_ref, star_ref})} + ); + SetRuleBodyAndBudgets(star_rule_id, star_body, format.max_tokens); + return ResultOk(star_rule_id); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenTriggeredTagsFormat& format +) { + // Step 1. Visit all tags, map trigger → tag IDs + std::vector> trigger_to_tag_ids(format.trigger_tokens.size()); + std::vector tag_content_rule_ids; + tag_content_rule_ids.reserve(format.tags.size()); + + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + if (!std::holds_alternative(tag.begin)) { + return ResultErr( + "Tags in token_triggered_tags must have a token format begin, not a string" + ); + } + auto begin_token_id = std::get(tag.begin).resolved_token_id_; + + int matched = -1; + for (int it_t = 0; it_t < static_cast(format.resolved_trigger_token_ids_.size()); ++it_t) { + if (format.resolved_trigger_token_ids_[it_t] == begin_token_id) { + if (matched != -1) { + return ResultErr("Tag matches multiple triggers"); + } + matched = it_t; + } + } + if (matched == -1) { + return ResultErr("Tag does not match any trigger"); + } + trigger_to_tag_ids[matched].push_back(it_tag); + + auto result = Visit(*tag.content); + if (result.IsErr()) return result; + tag_content_rule_ids.push_back(std::move(result).Unwrap()); + } + + // Step 2. Special case: at_least_one && stop_after_first + if (format.at_least_one && format.stop_after_first) { + std::vector choice_elements; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr = BuildBeginExpr(tag); + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr = BuildEndExpr(tag); + choice_elements.push_back(grammar_builder_.AddSequence({begin_expr, ref, end_expr})); + } + auto choice = grammar_builder_.AddChoices(choice_elements); + return ResultOk(grammar_builder_.AddRuleWithHint("token_triggered_tags", choice)); + } + + // Step 3. Normal case — TokenTagDispatch + std::vector> trigger_rule_pairs; + for (int it_t = 0; it_t < static_cast(format.trigger_tokens.size()); ++it_t) { + std::vector choice_elements; + for (auto tag_id : trigger_to_tag_ids[it_t]) { + const auto& tag = format.tags[tag_id]; + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[tag_id]); + auto end_expr = BuildEndExpr(tag); + choice_elements.push_back(grammar_builder_.AddSequence({ref, end_expr})); + } + auto choice = grammar_builder_.AddChoices(choice_elements); + auto sub_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_group", choice); + trigger_rule_pairs.push_back({format.resolved_trigger_token_ids_[it_t], sub_rule}); + } + + bool loop = !format.stop_after_first; + std::vector all_excludes = format.resolved_exclude_token_ids_; + for (auto tid : format.detected_end_token_ids_) { + all_excludes.push_back(tid); + } + auto ttd_expr = grammar_builder_.AddTokenTagDispatch( + Grammar::Impl::TokenTagDispatch{trigger_rule_pairs, loop, all_excludes} + ); + int32_t rule_expr_id = ttd_expr; + + if (format.at_least_one) { + std::vector first_choices; + for (int it_tag = 0; it_tag < static_cast(format.tags.size()); ++it_tag) { + const auto& tag = format.tags[it_tag]; + auto begin_expr = BuildBeginExpr(tag); + auto ref = grammar_builder_.AddRuleRef(tag_content_rule_ids[it_tag]); + auto end_expr = BuildEndExpr(tag); + first_choices.push_back(grammar_builder_.AddSequence({begin_expr, ref, end_expr})); + } + auto first_choice = grammar_builder_.AddChoices(first_choices); + auto first_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_first", first_choice); + auto dispatch_rule = grammar_builder_.AddRuleWithHint("token_triggered_tags_sub", rule_expr_id); + auto seq = grammar_builder_.AddSequence( + {grammar_builder_.AddRuleRef(first_rule), grammar_builder_.AddRuleRef(dispatch_rule)} + ); + rule_expr_id = grammar_builder_.AddChoices({seq}); + } + + return ResultOk(grammar_builder_.AddRuleWithHint("token_triggered_tags", rule_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const RepeatFormat& format) { + auto result = Visit(*format.content); + if (result.IsErr()) { + return result; + } + int content_rule_id = std::move(result).Unwrap(); + int repeat_expr_id = grammar_builder_.AddRepeat(content_rule_id, format.min, format.max); + return ResultOk(grammar_builder_.AddRuleWithHint("repeat", repeat_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const DispatchFormat& format) { + std::vector> tag_rule_pairs; + tag_rule_pairs.reserve(format.rules.size()); + for (const auto& pair : format.rules) { + if (!pair.second) { + return ResultErr("TagDispatch pair must have content"); + } + auto result = Visit(*pair.second); + if (result.IsErr()) { + return result; + } + tag_rule_pairs.push_back({pair.first, std::move(result).Unwrap()}); + } + auto rule_expr_id = grammar_builder_.AddTagDispatch( + Grammar::Impl::TagDispatch{std::move(tag_rule_pairs), format.loop, format.excludes} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("tag_dispatch", rule_expr_id)); +} + +Result StructuralTagGrammarConverter::VisitSub(const TokenDispatchFormat& format) { + XGRAMMAR_DCHECK(format.resolved_trigger_token_ids_.size() == format.rules.size()) + << "TokenDispatchFormat must be resolved before conversion"; + std::vector> trigger_rule_pairs; + trigger_rule_pairs.reserve(format.rules.size()); + for (size_t i = 0; i < format.rules.size(); ++i) { + const auto& pair = format.rules[i]; + if (!pair.second) { + return ResultErr("TokenTagDispatch pair must have content"); + } + auto result = Visit(*pair.second); + if (result.IsErr()) { + return result; + } + trigger_rule_pairs.push_back({format.resolved_trigger_token_ids_[i], std::move(result).Unwrap()} + ); + } + std::vector all_excludes = format.resolved_exclude_token_ids_; + auto rule_expr_id = grammar_builder_.AddTokenTagDispatch( + Grammar::Impl::TokenTagDispatch{trigger_rule_pairs, format.loop, all_excludes} + ); + return ResultOk(grammar_builder_.AddRuleWithHint("token_tag_dispatch", rule_expr_id)); +} + +/************** StructuralTag Conversion Public API **************/ + +Result StructuralTagToGrammar( + const std::string& structural_tag_json, const std::optional& tokenizer_info +) { + auto structural_tag_result = StructuralTagParser::FromJSON(structural_tag_json); + if (structural_tag_result.IsErr()) { + return ResultErr(std::move(structural_tag_result).UnwrapErr()); + } + auto structural_tag = std::move(structural_tag_result).Unwrap(); + + auto resolve_err = StructuralTagTokenResolver::Resolve(&structural_tag, tokenizer_info); + if (resolve_err.has_value()) { + return ResultErr(std::move(resolve_err).value()); + } + + auto err = StructuralTagAnalyzer().Analyze(&structural_tag); + if (err.has_value()) { + return ResultErr(std::move(err).value()); + } + + auto result = StructuralTagGrammarConverter::Convert(structural_tag); + if (result.IsErr()) { + return ResultErr(std::move(result).UnwrapErr()); + } + return ResultOk(GrammarNormalizer::Apply(std::move(result).Unwrap())); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/structural_tag.h b/zig/vendor/xgrammar/cpp/structural_tag.h new file mode 100644 index 0000000000..5913d93135 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/structural_tag.h @@ -0,0 +1,415 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/structural_tag_impl.h + * \brief The implementation header for the structural tag. + */ + +#ifndef XGRAMMAR_STRUCTURAL_TAG_H_ +#define XGRAMMAR_STRUCTURAL_TAG_H_ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "support/utils.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +/******************** Structural Tag Definition ********************/ + +// TODO(yixin): Consider moving the definition to Public API. + +struct ConstStringFormat; +struct JSONSchemaFormat; +struct AnyTextFormat; +struct GrammarFormat; +struct RegexFormat; +struct SequenceFormat; +struct OrFormat; +struct TagFormat; +struct TriggeredTagsFormat; +struct TagsWithSeparatorFormat; +struct OptionalFormat; +struct PlusFormat; +struct StarFormat; +struct TokenFormat; +struct ExcludeTokenFormat; +struct AnyTokensFormat; +struct TokenTriggeredTagsFormat; +struct RepeatFormat; +struct DispatchFormat; +struct TokenDispatchFormat; + +using Format = std::variant< + ConstStringFormat, + JSONSchemaFormat, + AnyTextFormat, + GrammarFormat, + RegexFormat, + SequenceFormat, + OrFormat, + TagFormat, + TriggeredTagsFormat, + TagsWithSeparatorFormat, + OptionalFormat, + PlusFormat, + StarFormat, + TokenFormat, + ExcludeTokenFormat, + AnyTokensFormat, + TokenTriggeredTagsFormat, + RepeatFormat, + DispatchFormat, + TokenDispatchFormat>; + +/******************** Basic Formats ********************/ + +struct ConstStringFormat { + static constexpr const char* type = "const_string"; + std::string value; + ConstStringFormat(std::string value) : value(std::move(value)) {} + picojson::value ToJSON() const; +}; + +struct JSONSchemaFormat { + static constexpr const char* type = "json_schema"; + std::string json_schema; + // "json", "qwen_xml", "minimax_xml", "minimax_m3_xml", "deepseek_xml", "glm_xml", + // "cohere_xml", "kimi_k3_xml", "deepseek_v4_1_xml" + std::string style = "json"; + // Whether to allow object properties to appear in any order. See + // Grammar::FromJSONSchema / JSONSchemaToEBNF for the semantics. + bool any_order = false; + // Per-tag cap on consecutive whitespace characters in the JSON-schema content. + std::optional max_whitespace_cnt = std::nullopt; + // Substrings forbidden within string values and property names. + std::vector excludes; + JSONSchemaFormat( + std::string json_schema, + std::string style = "json", + bool any_order = false, + std::optional max_whitespace_cnt = std::nullopt, + std::vector excludes = {} + ) + : json_schema(std::move(json_schema)), + style(std::move(style)), + any_order(any_order), + max_whitespace_cnt(max_whitespace_cnt), + excludes(std::move(excludes)) {} + picojson::value ToJSON() const; +}; + +struct GrammarFormat { + static constexpr const char* type = "grammar"; + std::string grammar; + GrammarFormat(std::string grammar) : grammar(std::move(grammar)) {} + picojson::value ToJSON() const; +}; + +struct RegexFormat { + static constexpr const char* type = "regex"; + std::string pattern; + RegexFormat(std::string pattern) : pattern(std::move(pattern)) {} + picojson::value ToJSON() const; +}; + +struct AnyTextFormat { + static constexpr const char* type = "any_text"; + std::vector excludes; + int32_t max_tokens = -1; + int32_t max_chars = -1; + AnyTextFormat( + std::vector excluded_strs, int32_t max_tokens = -1, int32_t max_chars = -1 + ) + : excludes(std::move(excluded_strs)), max_tokens(max_tokens), max_chars(max_chars) {} + picojson::value ToJSON() const; + + private: + std::vector detected_end_strs_; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TokenFormat { + static constexpr const char* type = "token"; + std::variant token; + TokenFormat(std::variant token) : token(std::move(token)) { + if (std::holds_alternative(this->token)) { + resolved_token_id_ = std::get(this->token); + } + } + picojson::value ToJSON() const; + + private: + int32_t resolved_token_id_ = -1; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct ExcludeTokenFormat { + static constexpr const char* type = "exclude_token"; + std::vector> exclude_tokens; + ExcludeTokenFormat(std::vector> exclude_tokens) + : exclude_tokens(std::move(exclude_tokens)) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct AnyTokensFormat { + static constexpr const char* type = "any_tokens"; + std::vector> exclude_tokens; + int32_t max_tokens = -1; + AnyTokensFormat( + std::vector> exclude_tokens, int32_t max_tokens = -1 + ) + : exclude_tokens(std::move(exclude_tokens)), max_tokens(max_tokens) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_exclude_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +/******************** Combinatorial Formats ********************/ + +struct SequenceFormat { + static constexpr const char* type = "sequence"; + std::vector elements; + SequenceFormat(std::vector elements); + picojson::value ToJSON() const; + + private: + // Detected in StructuralTagAnalyzer + bool is_unlimited_ = false; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct OrFormat { + static constexpr const char* type = "or"; + std::vector elements; + OrFormat(std::vector elements); + picojson::value ToJSON() const; + + private: + // Detected in StructuralTagAnalyzer + bool is_unlimited_ = false; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TagFormat { + static constexpr const char* type = "tag"; + std::variant begin; + std::shared_ptr content; + std::variant, TokenFormat> end; + + TagFormat( + std::variant begin, + std::shared_ptr content, + std::variant, TokenFormat> end + ) + : begin(std::move(begin)), content(std::move(content)), end(std::move(end)) {} + picojson::value ToJSON() const; +}; + +struct TriggeredTagsFormat { + static constexpr const char* type = "triggered_tags"; + std::vector triggers; + std::vector tags; + std::vector excludes; + bool at_least_one = false; + bool stop_after_first = false; + + TriggeredTagsFormat( + std::vector triggers, + std::vector tags, + std::vector excludes, + bool at_least_one, + bool stop_after_first + ) + : triggers(std::move(triggers)), + tags(std::move(tags)), + excludes(std::move(excludes)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + std::vector detected_end_strs_; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TagsWithSeparatorFormat { + static constexpr const char* type = "tags_with_separator"; + std::vector tags; + std::string separator; + bool at_least_one = false; + bool stop_after_first = false; + + TagsWithSeparatorFormat( + std::vector tags, std::string separator, bool at_least_one, bool stop_after_first + ) + : tags(std::move(tags)), + separator(std::move(separator)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct TokenTriggeredTagsFormat { + static constexpr const char* type = "token_triggered_tags"; + std::vector> trigger_tokens; + std::vector tags; + std::vector> exclude_tokens; + bool at_least_one = false; + bool stop_after_first = false; + + TokenTriggeredTagsFormat( + std::vector> trigger_tokens, + std::vector tags, + std::vector> exclude_tokens, + bool at_least_one, + bool stop_after_first + ) + : trigger_tokens(std::move(trigger_tokens)), + tags(std::move(tags)), + exclude_tokens(std::move(exclude_tokens)), + at_least_one(at_least_one), + stop_after_first(stop_after_first) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_trigger_token_ids_; + std::vector resolved_exclude_token_ids_; + std::vector detected_end_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagAnalyzer; + friend class StructuralTagGrammarConverter; +}; + +struct OptionalFormat { + static constexpr const char* type = "optional"; + std::shared_ptr content; + OptionalFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct PlusFormat { + static constexpr const char* type = "plus"; + std::shared_ptr content; + PlusFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct StarFormat { + static constexpr const char* type = "star"; + std::shared_ptr content; + StarFormat(std::shared_ptr content) : content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +struct RepeatFormat { + static constexpr const char* type = "repeat"; + int32_t min; + int32_t max; + std::shared_ptr content; + RepeatFormat(int32_t min, int32_t max, std::shared_ptr content) + : min(min), max(max), content(std::move(content)) {} + picojson::value ToJSON() const; +}; + +/*! + * \brief A format that maps directly to a TagDispatch grammar. + * Accepts ``[trigger string, content format]`` pairs in JSON; each content is converted to a rule + * and the result is a single TagDispatch(loop, excludes). + */ +struct DispatchFormat { + static constexpr const char* type = "dispatch"; + std::vector>> rules; + bool loop = true; + std::vector excludes; + + DispatchFormat( + std::vector>> rules, + bool loop = true, + std::vector excludes = {} + ) + : rules(std::move(rules)), loop(loop), excludes(std::move(excludes)) {} + picojson::value ToJSON() const; +}; + +/*! + * \brief A format that maps directly to a TokenTagDispatch grammar. + * Accepts ``[trigger token, content format]`` pairs in JSON; trigger can be token ID or token + * string (resolved via tokenizer_info). Each content is converted to a rule. + */ +struct TokenDispatchFormat { + static constexpr const char* type = "token_dispatch"; + std::vector, std::shared_ptr>> rules; + bool loop = true; + std::vector> exclude_tokens; + + TokenDispatchFormat( + std::vector, std::shared_ptr>> rules, + bool loop = true, + std::vector> exclude_tokens = {} + ) + : rules(std::move(rules)), loop(loop), exclude_tokens(std::move(exclude_tokens)) {} + picojson::value ToJSON() const; + + private: + std::vector resolved_trigger_token_ids_; + std::vector resolved_exclude_token_ids_; + friend class StructuralTagTokenResolver; + friend class StructuralTagGrammarConverter; +}; + +/******************** Top Level ********************/ + +struct StructuralTag { + static constexpr const char* type = "structural_tag"; + Format format; + + StructuralTag(Format format) : format(std::move(format)) {} +}; + +/******************** Conversion API ********************/ + +/*! + * \brief Convert a structural tag JSON string to a grammar. + * \param structural_tag_json The JSON string of the structural tag. + * \return A grammar if the JSON is valid, otherwise an error message in std::string. + */ +Result StructuralTagToGrammar( + const std::string& structural_tag_json, + const std::optional& tokenizer_info = std::nullopt +); + +} // namespace xgrammar + +#endif // XGRAMMAR_STRUCTURAL_TAG_H_ diff --git a/zig/vendor/xgrammar/cpp/suffix_automata.cc b/zig/vendor/xgrammar/cpp/suffix_automata.cc new file mode 100644 index 0000000000..a3b4c85bac --- /dev/null +++ b/zig/vendor/xgrammar/cpp/suffix_automata.cc @@ -0,0 +1,88 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/suffix_automata.cc + */ + +#include "suffix_automata.h" + +#include +#include +#include +#include + +namespace xgrammar { + +FSMWithStartEnd SuffixAutomata::Build(const std::vector& chunks) { + // Step 1. Build the suffix automaton over the chunk sequence with the standard online + // construction. Each chunk is treated as one symbol of the alphabet. + struct State { + int32_t length = 0; + int32_t suffix_link = -1; + std::map transitions; + }; + + std::vector states(1); + int32_t last = 0; + for (const std::string& chunk : chunks) { + int32_t current = static_cast(states.size()); + states.push_back({states[last].length + 1, -1, {}}); + + int32_t parent = last; + while (parent != -1 && !states[parent].transitions.count(chunk)) { + states[parent].transitions[chunk] = current; + parent = states[parent].suffix_link; + } + if (parent == -1) { + states[current].suffix_link = 0; + } else { + int32_t target = states[parent].transitions.at(chunk); + if (states[parent].length + 1 == states[target].length) { + states[current].suffix_link = target; + } else { + int32_t clone = static_cast(states.size()); + states.push_back(states[target]); + states[clone].length = states[parent].length + 1; + while (parent != -1) { + auto transition = states[parent].transitions.find(chunk); + if (transition == states[parent].transitions.end() || transition->second != target) { + break; + } + transition->second = clone; + parent = states[parent].suffix_link; + } + states[target].suffix_link = clone; + states[current].suffix_link = clone; + } + } + last = current; + } + + // Step 2. Expand the chunk-level automaton into a byte-level FSM. Automaton state i maps to + // FSM state i; every automaton state is accepting. A chunk-labeled transition becomes a chain + // of byte transitions through fresh intermediate states; an empty chunk becomes an epsilon + // transition. + FSM fsm(static_cast(states.size())); + std::vector end_states; + end_states.reserve(states.size()); + for (int32_t index = 0; index < static_cast(states.size()); ++index) { + end_states.push_back(index); + } + for (int32_t index = 0; index < static_cast(states.size()); ++index) { + for (const auto& [chunk, target] : states[index].transitions) { + if (chunk.empty()) { + fsm.AddEpsilonEdge(index, target); + continue; + } + int current_state = index; + for (size_t offset = 0; offset < chunk.size(); ++offset) { + int next_state = offset + 1 == chunk.size() ? static_cast(target) : fsm.AddState(); + uint8_t byte = static_cast(chunk[offset]); + fsm.AddEdge(current_state, next_state, byte, byte); + current_state = next_state; + } + } + } + return FSMWithStartEnd(fsm, 0, std::move(end_states)); +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/suffix_automata.h b/zig/vendor/xgrammar/cpp/suffix_automata.h new file mode 100644 index 0000000000..d4d9dac536 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/suffix_automata.h @@ -0,0 +1,36 @@ +/*! + * Copyright (c) 2026 by Contributors + * \file xgrammar/suffix_automata.h + * \brief Suffix automaton construction for substring expressions. + */ +#ifndef XGRAMMAR_SUFFIX_AUTOMATA_H_ +#define XGRAMMAR_SUFFIX_AUTOMATA_H_ + +#include +#include + +#include "fsm.h" + +namespace xgrammar { + +/*! + * \brief Builds the automaton of a substring expression via a chunk-level suffix automaton. + */ +class SuffixAutomata { + public: + /*! + * \brief Build an FSM that accepts exactly the contiguous subsequences of the chunk list, + * including the empty one. + * \details A suffix automaton is built over the chunk sequence (each chunk is one symbol), so + * the number of automaton states grows linearly with the number of chunks. Every automaton + * state is accepting. Each chunk-labeled transition is then expanded into a chain of byte + * transitions; an empty chunk becomes an epsilon transition. + * \param chunks The list of byte string chunks. Chunks may be empty or repeated. + * \return The FSM with start and end states. + */ + static FSMWithStartEnd Build(const std::vector& chunks); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUFFIX_AUTOMATA_H_ diff --git a/zig/vendor/xgrammar/cpp/support/compact_2d_array.h b/zig/vendor/xgrammar/cpp/support/compact_2d_array.h new file mode 100644 index 0000000000..f95c20e9c1 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/compact_2d_array.h @@ -0,0 +1,397 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/compact_2d_array.h + */ +#ifndef XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ +#define XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" +#include "memory_size.h" +#include "reflection.h" + +namespace xgrammar { + +/*! + * \brief This class implements a Compressed Sparse Row (CSR) array data structure. It stores + * a 2D array in a compressed format, where each row can have a variable number of elements, and + * all rows are stored contiguously in memory. The inserted row is immutable. + * + * \note Inserting new rows into the Compact2DArray will invalidate the existing Row objects. + * + * \tparam DataType The type of elements stored in the Compact2DArray. + * + * \details + * The Compact2DArray stores elements of type DataType in a compressed format, + * where each row can have a variable number of elements. It uses two vectors: + * - data_: stores all elements contiguously + * - indptr_: stores the starting index of each row in data_. Its last element is the size of data_ + * representing the ending index. + * + * This structure allows efficient storage and access for sparse data. + */ +template +class Compact2DArray { + public: + /*! + * \brief The struct representing a row in the Compact2DArray. + */ + struct Row { + /*! \brief The value type is DataType. */ + using value_type = DataType; + + /*! \brief Pointer to the data of the row. */ + const DataType* data; + /*! \brief Length of the row data. */ + int32_t data_len; + + /*! + * \brief Access an element in the row. + * \param i Index of the element to access. + * \return Reference to the element at index i. + */ + const DataType& operator[](int32_t i) const { + XGRAMMAR_DCHECK(i >= 0 && i < data_len) + << "Index " << i << " of the Compact2DArray Row is out of bound"; + return data[i]; + } + + /*! \brief Get the beginning iterator of the row. */ + const DataType* begin() const { return data; } + /*! \brief Get the end iterator of the row. */ + const DataType* end() const { return data + data_len; } + /*! \brief Get the size of the row. */ + int32_t size() const { return data_len; } + + /*! \brief Get a sub-row in [begin, end). */ + Row Slice(int32_t begin, int32_t end) const { + XGRAMMAR_DCHECK(begin >= 0 && begin <= end && end <= data_len) + << "Compact2DArray Row slice is out of bound"; + return {data + begin, end - begin}; + } + + friend std::ostream& operator<<(std::ostream& os, const Row& row) { + os << "["; + for (auto i = 0; i < row.data_len; ++i) { + if (i > 0) { + os << ", "; + } + os << row[i]; + } + os << "]"; + return os; + } + }; + + /*! + * \brief The mutable struct representing a row in the Compact2DArray. + */ + struct MutableRow { + /*! \brief The value type is DataType. */ + using value_type = DataType; + + /*! \brief Pointer to the data of the row. */ + DataType* data; + /*! \brief Length of the row data. */ + int32_t data_len; + + /*! + * \brief Access an element in the row. + * \param i Index of the element to access. + * \return Reference to the element at index i. + */ + DataType& operator[](int32_t i) const { + XGRAMMAR_DCHECK(i >= 0 && i < data_len) + << "Index " << i << " of the Compact2DArray MutableRow is out of bound"; + return data[i]; + } + + /*! \brief Get the beginning iterator of the row. */ + DataType* begin() const { return data; } + /*! \brief Get the end iterator of the row. */ + DataType* end() const { return data + data_len; } + /*! \brief Get the size of the row. */ + int32_t size() const { return data_len; } + }; + + /*! \brief The value type is Row. */ + using value_type = Row; + + /*! \brief Default constructor. */ + Compact2DArray() = default; + + /*! + * \brief Construct a Compact2DArray from an existing CSR representation. + * \param data All row elements stored contiguously. + * \param indptr Row start offsets. Must start with 0, be non-decreasing, and end with + * data.size(). + * \return The constructed Compact2DArray. + */ + static Compact2DArray FromDataAndIndptr(std::vector data, std::vector indptr); + + /*! + * \brief Construct a Compact2DArray from row sizes with default-constructed data. + * \param row_sizes The size of each row. + * \return The constructed Compact2DArray. + */ + static Compact2DArray FromRowSizes(const std::vector& row_sizes); + + /*! + * \brief Reset the Compact2DArray from row sizes with default-constructed data. + * \param row_sizes The size of each row. + */ + void ResetWithRowSizes(const std::vector& row_sizes); + + /****************** Accessors ******************/ + + /*! \brief Get the number of rows in the Compact2DArray. */ + int32_t size() const { return static_cast(indptr_.size()) - 1; } + + /*! + * \brief Check the CSR invariants after deserialization: indptr starts with 0, is non-decreasing + * and ends with the data size. Otherwise row accesses would read out of bounds. + * \return An error message if the invariants are violated. + */ + std::optional Validate() const { + if (indptr_.empty() || indptr_.front() != 0 || + indptr_.back() != static_cast(data_.size()) || + !std::is_sorted(indptr_.begin(), indptr_.end())) { + return "Invalid indptr: it must start with 0, be non-decreasing and end with the data size"; + } + return std::nullopt; + } + + friend std::size_t MemorySize(const Compact2DArray& arr) { + return MemorySize(arr.data_) + MemorySize(arr.indptr_); + } + + /*! + * \brief Access a row in the Compact2DArray. + * \param i Index of the row to access. + * \return Row struct representing the i-th row. + */ + Row operator[](int32_t i) const; + + /*! + * \brief Access a mutable row in the Compact2DArray. + * \param i Index of the row to access. + * \return MutableRow struct representing the i-th row. + */ + MutableRow MutableRowAt(int32_t i); + + /****************** Modifiers ******************/ + + /*! + * \brief Insert a new row of data into the Compact2DArray. + * \param data Pointer to the data to be inserted. + * \param data_len Length of the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const DataType* new_data, int32_t new_data_len); + + /*! + * \brief Insert a new row of data into the Compact2DArray from a vector. + * \param data Vector containing the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const std::vector& new_data); + + /*! + * \brief Insert a new row of data into the Compact2DArray from a Row struct. + * \param row The Row struct containing the data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBack(const Row& row) { return PushBack(row.data, row.data_len); } + + /*! + * \brief Push back a new element in the latest row. + * \param new_data the element to be pushed. + */ + void PushBackInLatestRow(const DataType& new_data) { + XGRAMMAR_DCHECK(!indptr_.empty()) << "Cannot push back in an empty Compact2DArray"; + CheckCanAppendData(1); + data_.push_back(new_data); + indptr_.back()++; + } + + Row Back() { return (*this)[size() - 1]; } + + /*! + * \brief Insert a new row of non-contiguous data into the Compact2DArray. This method inserts a + * single element followed by a sequence of elements. This is useful in the GrammarExpr data + * structure. + * \param data_1 The first element to be inserted. + * \param data_2 Pointer to the remaining data to be inserted. + * \param data_2_len Length of the remaining data to be inserted. + * \return The index of the newly inserted row. + */ + int32_t PushBackNonContiguous(DataType data_1, const DataType* data_2, int32_t data_2_len); + + /*! + * \brief Pop back the last one or multiple rows of the Compact2DArray. + * \param cnt The number of rows to be popped. + */ + void PopBack(const int32_t& cnt) { + indptr_.erase(indptr_.end() - cnt, indptr_.end()); + data_.erase(data_.begin() + indptr_.back(), data_.end()); + return; + } + + /****************** Internal Accessors ******************/ + + /*! \brief Get a pointer to the underlying data array. */ + const DataType* data() const { return data_.data(); } + /*! \brief Get a pointer to the underlying index pointer array. */ + const int32_t* indptr() const { return indptr_.data(); } + + /****************** Printing ******************/ + + friend std::ostream& operator<<(std::ostream& os, const Compact2DArray& compact_2d_array) { + os << "Compact2DArray(["; + for (auto i = 0; i < compact_2d_array.size(); ++i) { + if (i > 0) { + os << ", "; + } + os << compact_2d_array[i]; + } + os << "])"; + return os; + } + + private: + static constexpr size_t kMaxRepresentableSize = static_cast(INT32_MAX); + + inline void CheckCanAppendData(size_t additional_size) const { + XGRAMMAR_ICHECK(additional_size <= kMaxRepresentableSize - data_.size()) + << "Compact2DArray data size would exceed the int32_t limit, likely due to an unbounded " + "grammar pattern causing state explosion"; + } + + /*! \brief Vector storing all elements contiguously. */ + std::vector data_; + /*! \brief Vector storing the starting index of each row in data_. */ + std::vector indptr_{0}; + friend struct member_trait>; +}; + +template +inline typename Compact2DArray::Row Compact2DArray::operator[](int32_t i +) const { + XGRAMMAR_DCHECK(i >= 0 && i < size()) << "Compact2DArray index " << i << " is out of bound"; + int32_t start = indptr_[i]; + int32_t end = indptr_[i + 1]; + return {data_.data() + start, end - start}; +} + +template +inline typename Compact2DArray::MutableRow Compact2DArray::MutableRowAt( + int32_t i +) { + XGRAMMAR_DCHECK(i >= 0 && i < size()) << "Compact2DArray index " << i << " is out of bound"; + int32_t start = indptr_[i]; + int32_t end = indptr_[i + 1]; + return {data_.data() + start, end - start}; +} + +template +inline Compact2DArray Compact2DArray::FromDataAndIndptr( + std::vector data, std::vector indptr +) { + XGRAMMAR_CHECK(!indptr.empty()) << "Compact2DArray indptr cannot be empty"; + XGRAMMAR_CHECK(indptr.front() == 0) << "Compact2DArray indptr must start with 0"; + for (int32_t i = 1; i < static_cast(indptr.size()); ++i) { + XGRAMMAR_CHECK(indptr[i - 1] <= indptr[i]) << "Compact2DArray indptr must be non-decreasing"; + } + XGRAMMAR_CHECK(indptr.back() == static_cast(data.size())) + << "Compact2DArray indptr must end with data.size()"; + + Compact2DArray result; + result.data_ = std::move(data); + result.indptr_ = std::move(indptr); + return result; +} + +template +inline Compact2DArray Compact2DArray::FromRowSizes( + const std::vector& row_sizes +) { + Compact2DArray result; + result.ResetWithRowSizes(row_sizes); + return result; +} + +template +inline void Compact2DArray::ResetWithRowSizes(const std::vector& row_sizes) { + indptr_.resize(row_sizes.size() + 1); + indptr_[0] = 0; + for (int32_t i = 0; i < static_cast(row_sizes.size()); ++i) { + XGRAMMAR_CHECK(row_sizes[i] >= 0) << "Compact2DArray row size cannot be negative"; + XGRAMMAR_CHECK(static_cast(row_sizes[i]) <= kMaxRepresentableSize - indptr_[i]) + << "Compact2DArray data size exceeds the int32_t limit"; + indptr_[i + 1] = indptr_[i] + row_sizes[i]; + } + data_.resize(indptr_.back()); +} + +template +inline int32_t Compact2DArray::PushBack(const DataType* new_data, int32_t new_data_len) { + CheckCanAppendData(static_cast(new_data_len)); + // If the new data is already in the Compact2DArray, we need to copy it to the new memory + // location. + if (new_data >= data_.data() && new_data < data_.data() + data_.size()) { + std::vector new_data_copied(new_data, new_data + new_data_len); + data_.insert(data_.end(), new_data_copied.begin(), new_data_copied.end()); + } else { + data_.insert(data_.end(), new_data, new_data + new_data_len); + } + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +inline int32_t Compact2DArray::PushBack(const std::vector& new_data) { + CheckCanAppendData(new_data.size()); + data_.insert(data_.end(), new_data.begin(), new_data.end()); + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +inline int32_t Compact2DArray::PushBackNonContiguous( + DataType data_1, const DataType* data_2, int32_t data_2_len +) { + CheckCanAppendData(static_cast(data_2_len) + 1); + if (data_2 >= data_.data() && data_2 < data_.data() + data_.size()) { + std::vector new_data_copied(data_2, data_2 + data_2_len); + data_.push_back(data_1); + data_.insert(data_.end(), new_data_copied.begin(), new_data_copied.end()); + } else { + data_.push_back(data_1); + data_.insert(data_.end(), data_2, data_2 + data_2_len); + } + indptr_.push_back(static_cast(data_.size())); + return static_cast(indptr_.size()) - 2; +} + +template +XGRAMMAR_MEMBER_TABLE_TEMPLATE( + Compact2DArray, + "data_", + &Compact2DArray::data_, + "indptr_", + &Compact2DArray::indptr_ +); + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_COMPACT_2D_ARRAY_H_ diff --git a/zig/vendor/xgrammar/cpp/support/container.h b/zig/vendor/xgrammar/cpp/support/container.h new file mode 100644 index 0000000000..2ca660e90b --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/container.h @@ -0,0 +1,166 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/container.h + * \brief The header for container. + */ +#ifndef XGRAMMAR_SUPPORT_CONTAINER_H_ +#define XGRAMMAR_SUPPORT_CONTAINER_H_ +#include + +#include "logging.h" + +namespace xgrammar { + +namespace details { + +template +class NodePool { + public: + NodePool() = default; + + void Reserve(int n) { node_pool_.reserve(n); } + + [[nodiscard]] + int Allocate() { + if (free_list_.empty()) { + int node = Size(); + node_pool_.emplace_back(); + return node; + } else { + int node = free_list_.back(); + free_list_.pop_back(); + return node; + } + } + + void Deallocate(int node) { free_list_.push_back(node); } + + void Clear() { + node_pool_.clear(); + free_list_.clear(); + } + + Node& operator[](int node) { + XGRAMMAR_DCHECK(0 <= node && node < Size()); + return node_pool_[node]; + } + + int Size() const { return static_cast(node_pool_.size()); } + + private: + std::vector node_pool_; + std::vector free_list_; +}; + +} // namespace details + +template +class List { + private: + struct Node { + int prev; + int next; + Value value; + }; + + public: + struct iterator { + public: + iterator(int n, List& c) : node_(n), list_(&c) { + XGRAMMAR_DCHECK(0 <= node_ && node_ < list_->node_pool_.Size()); + } + iterator& operator++() { + node_ = GetNode().next; + return *this; + } + iterator operator++(int) { + iterator tmp = *this; + ++*this; + return tmp; + } + Value& operator*() const { return GetNode().value; } + Value* operator->() const { return &GetNode().value; } + bool operator==(const iterator& rhs) const { + XGRAMMAR_DCHECK(list_ == rhs.list_) << "compare different container is UB"; + return node_ == rhs.node_; // compare different container is UB + } + bool operator!=(const iterator& rhs) const { + XGRAMMAR_DCHECK(list_ == rhs.list_) << "compare different container is UB"; + return node_ != rhs.node_; // compare different container is UB + } + + int Index() const { return node_; } + + private: + friend class List; + Node& GetNode() const { return list_->node_pool_[node_]; } + + int node_; + List* list_; + }; + + List(int reserved = 0) { + node_pool_.Reserve(reserved); + InitGuard(); + } + + iterator PushBack(const Value& value) { + int node = node_pool_.Allocate(); + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + node_pool_[node].value = value; + LinkBefore(node, 0); + return iterator(node, *this); + } + + void MoveBack(int node) { + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + Unlink(node); + LinkBefore(node, 0); + } + + iterator Erase(iterator it) { + int node = it.Index(); + XGRAMMAR_DCHECK(0 < node && node < node_pool_.Size()); + int next = node_pool_[node].next; + Unlink(node); + node_pool_.Deallocate(node); + return iterator(next, *this); + } + + void Clear() { + node_pool_.Clear(); + InitGuard(); + } + + iterator begin() { return iterator(node_pool_[0].next, *this); } + iterator end() { return iterator(0, *this); } + + private: + void InitGuard() { + int node_id = node_pool_.Allocate(); + XGRAMMAR_DCHECK(node_id == 0) << "node 0 should be reserved as guard node"; + node_pool_[0].prev = 0; + node_pool_[0].next = 0; + } + + void LinkBefore(int node, int next) { + int prev = node_pool_[next].prev; + node_pool_[node].prev = prev; + node_pool_[node].next = next; + node_pool_[prev].next = node; + node_pool_[next].prev = node; + } + + void Unlink(int node) { + int prev = node_pool_[node].prev; + int next = node_pool_[node].next; + node_pool_[prev].next = next; + node_pool_[next].prev = prev; + } + + details::NodePool node_pool_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_CONTAINER_H_ diff --git a/zig/vendor/xgrammar/cpp/support/cpptrace.h b/zig/vendor/xgrammar/cpp/support/cpptrace.h new file mode 100644 index 0000000000..c1d7b7a415 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/cpptrace.h @@ -0,0 +1,39 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/cpptrace.h + * \details This file is an encapsulation of the cpptrace library. It helps debugging. This file + * takes effect only when XGRAMMAR_ENABLE_CPPTRACE is set to 1, and only support Linux and + * RelWithDebugInfo or Debug build. + */ +#ifndef XGRAMMAR_SUPPORT_CPPTRACE_H_ +#define XGRAMMAR_SUPPORT_CPPTRACE_H_ + +#if XGRAMMAR_ENABLE_CPPTRACE == 1 +#include +#endif + +#include + +namespace xgrammar { + +#if XGRAMMAR_ENABLE_CPPTRACE == 1 + +// Flag to check if cpptrace feature is enabled +static constexpr bool CPPTRACE_ENABLED = true; + +inline void PrintTrace() { cpptrace::generate_trace().print(); } +inline std::string GetTraceString() { return cpptrace::generate_trace().to_string(true); } + +#else + +static constexpr bool CPPTRACE_ENABLED = false; + +// Provide empty implementation when cpptrace is disabled +inline void PrintTrace() {} +inline std::string GetTraceString() { return ""; } + +#endif + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_CPPTRACE_H_ diff --git a/zig/vendor/xgrammar/cpp/support/dynamic_bitset.h b/zig/vendor/xgrammar/cpp/support/dynamic_bitset.h new file mode 100644 index 0000000000..592a612398 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/dynamic_bitset.h @@ -0,0 +1,363 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/dynamic_bitset.h + * \brief The header for utilities used in grammar-guided generation. + */ +#ifndef XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ +#define XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ + +#include + +#include +#include +#include +#include +#include + +// For __popcnt +#ifdef _MSC_VER +#include +#endif + +#include "json_serializer.h" +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief A bitset whose length is specified at runtime. Note the size cannot be changed after + * construction. + * \details The buffer of the bitset is a uint32_t array. There are two uses for this class: + * - When passing nullptr to data, it maintains an internal buffer for the bitset. + * - When passing a pointer to a buffer with enough size, it uses the external buffer for the + * bitset. + * \details Part of the implementation is adopted from Boost::dynamic_bitset. + */ +class DynamicBitset { + public: + /*! + * \brief Calculate the minimal size of the uint32_t buffer for the bitset with the given size. + * \param element_size The size of the bitset. + * \return The minimal buffer size. + */ + static int GetBufferSize(int element_size) { return (element_size + 31) / 32; } + + /*! + * \brief Construct a empty bitset. This object should be assigned to a valid bitset before using. + */ + DynamicBitset() : size_(0), buffer_size_(0), data_(nullptr), is_internal_(true) {} + + /*! + * \brief Construct a bitset with the given size. + * \param size The size of the bitset. + * \param data The buffer for the bitset. If nullptr, the bitset will maintain an internal buffer. + */ + DynamicBitset(int size, uint32_t* data = nullptr) + : size_(size), buffer_size_(GetBufferSize(size)) { + if (data == nullptr) { + internal_buffer_.resize(buffer_size_, 0); + data_ = internal_buffer_.data(); + is_internal_ = true; + } else { + data_ = data; + is_internal_ = false; + } + } + + /*! \brief Copy constructor. Copy the buffer and manage the memory internally. */ + DynamicBitset(const DynamicBitset& other) + : size_(other.size_), + buffer_size_(other.buffer_size_), + data_(), + internal_buffer_(), + is_internal_(other.is_internal_) { + if (other.is_internal_) { + // copy the internal buffer + internal_buffer_ = other.internal_buffer_; + data_ = internal_buffer_.data(); + } else { + // simply point to the same external buffer + data_ = other.data_; + } + } + + /*! \brief Move constructor. Reset other and take ownership of its buffer. */ + DynamicBitset(DynamicBitset&& other) noexcept + : size_(std::exchange(other.size_, 0)), + buffer_size_(std::exchange(other.buffer_size_, 0)), + data_(std::exchange(other.data_, nullptr)), + internal_buffer_(std::move(other.internal_buffer_)), + is_internal_(std::exchange(other.is_internal_, true)) {} + + /*! \brief Copy assignment. */ + DynamicBitset& operator=(const DynamicBitset& other) { + XGRAMMAR_DCHECK(is_internal_ || size_ >= other.size_) + << "Expanding bitset size is not allowed when the " + "memory of the bitset is externally managed"; + size_ = other.size_; + buffer_size_ = other.buffer_size_; + if (is_internal_) { + internal_buffer_.resize(buffer_size_); + data_ = internal_buffer_.data(); + } + if (data_ != other.data_) { + std::memcpy(data_, other.data_, buffer_size_ * sizeof(uint32_t)); + } + return *this; + } + + /*! \brief Move assignment. */ + DynamicBitset& operator=(DynamicBitset&& other) noexcept { + size_ = other.size_; + buffer_size_ = other.buffer_size_; + is_internal_ = other.is_internal_; + if (is_internal_) { + internal_buffer_ = std::move(other.internal_buffer_); + data_ = internal_buffer_.data(); + } else { + data_ = other.data_; + } + return *this; + } + + /*! \brief Get the value of the bit at the given index. */ + bool operator[](int index) const { + XGRAMMAR_DCHECK(data_ && index >= 0 && index < size_); + return (data_[index / 32] >> (index % 32)) & 1; + } + + /*! \brief Get the size of the bitset. */ + int Size() const { return size_; } + + /*! \brief Set the whole bitset to true. */ + void Set() { + XGRAMMAR_DCHECK(data_); + std::memset(data_, 0xFF, buffer_size_ * sizeof(uint32_t)); + } + + /*! \brief Set the bit at the given index to the given value. */ + void Set(int index, bool value = true) { + XGRAMMAR_DCHECK(data_ && index >= 0 && index < size_); + if (value) { + data_[index / 32] |= 1 << (index % 32); + } else { + data_[index / 32] &= ~(1 << (index % 32)); + } + } + + /*! \brief Set the whole bitset to false. */ + void Reset() { + XGRAMMAR_DCHECK(data_); + std::memset(data_, 0, buffer_size_ * sizeof(uint32_t)); + } + + /*! \brief Set the bit at the given index to false. */ + void Reset(int index) { Set(index, false); } + + /*! \brief Perform a bitwise OR operation between the current bitset and another bitset. */ + DynamicBitset& operator|=(const DynamicBitset& other) { + XGRAMMAR_DCHECK(buffer_size_ <= other.buffer_size_); + for (int i = 0; i < buffer_size_; ++i) { + data_[i] |= other.data_[i]; + } + return *this; + } + + int FindFirstOne() const { return DoFindOneFrom(0); } + + int FindNextOne(int pos) const { + if (pos >= size_ - 1 || size_ == 0) return -1; + ++pos; + int blk = pos / BITS_PER_BLOCK; + int ind = pos % BITS_PER_BLOCK; + uint32_t fore = data_[blk] >> ind; + int result = fore ? pos + LowestBit(fore) : DoFindOneFrom(blk + 1); + return result < size_ ? result : -1; + } + + int FindFirstZero() const { return DoFindZeroFrom(0); } + + int FindNextZero(int pos) const { + if (pos >= size_ - 1 || size_ == 0) return -1; + ++pos; + int blk = pos / BITS_PER_BLOCK; + int ind = pos % BITS_PER_BLOCK; + uint32_t fore = (~data_[blk]) >> ind; + int result = fore ? pos + LowestBit(fore) : DoFindZeroFrom(blk + 1); + return result < size_ ? result : -1; + } + + int Count() const { + int count = 0; + for (int i = 0; i < buffer_size_; ++i) { + count += PopCount(data_[i]); + } + return count; + } + + bool All() const { + if (size_ == 0) return true; + // Check all complete blocks except the last one + for (int i = 0; i < buffer_size_ - 1; ++i) { + if (data_[i] != ~static_cast(0)) { + return false; + } + } + // For the last block, create a mask for valid bits only + int remaining_bits = size_ % BITS_PER_BLOCK; + uint32_t last_block_mask = remaining_bits ? (static_cast(1) << remaining_bits) - 1 + : ~static_cast(0); + return (data_[buffer_size_ - 1] & last_block_mask) == last_block_mask; + } + + bool Any() const { + if (size_ == 0) return false; + // Check all complete blocks except the last one + for (int i = 0; i < buffer_size_ - 1; ++i) { + if (data_[i] != 0) { + return true; + } + } + // For the last block, only consider the valid bits + int remaining_bits = size_ % BITS_PER_BLOCK; + uint32_t last_block_mask = remaining_bits ? (static_cast(1) << remaining_bits) - 1 + : ~static_cast(0); + return (data_[buffer_size_ - 1] & last_block_mask) != 0; + } + + static constexpr int BITS_PER_BLOCK = 32; + + friend std::size_t MemorySize(const DynamicBitset& bitset) { + return bitset.buffer_size_ * sizeof(bitset.data_[0]); + } + + friend picojson::value SerializeJSONValue(const DynamicBitset& bitset) { + XGRAMMAR_DCHECK(bitset.buffer_size_ == GetBufferSize(bitset.size_)); + picojson::array result; + result.reserve(2 + bitset.buffer_size_); + result.emplace_back(picojson::value(static_cast(bitset.size_))); + result.emplace_back(picojson::value(static_cast(bitset.buffer_size_))); + for (int i = 0; i < bitset.buffer_size_; ++i) { + result.emplace_back(picojson::value(static_cast(bitset.data_[i]))); + } + return picojson::value(std::move(result)); + } + + friend std::optional DeserializeJSONValue( + DynamicBitset* bitset, const picojson::value& value, const std::string& type_name + ) { + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + if (arr.size() < 2) { + return ConstructDeserializeError("Except at least 2 elements in the array", type_name); + } + if (!arr[0].is()) { + return ConstructDeserializeError("Expect an integer for size", type_name); + } + int size = static_cast(arr[0].get()); + if (!arr[1].is()) { + return ConstructDeserializeError("Expect an integer for buffer_size", type_name); + } + int buffer_size = static_cast(arr[1].get()); + if (size < 0 || buffer_size != GetBufferSize(size)) { + return ConstructDeserializeError( + "Invalid buffer_size. Buffer size should be ceil(size / 32)", type_name + ); + } + if (static_cast(arr.size()) != static_cast(buffer_size) + 2) { + return ConstructDeserializeError( + "Expect exactly buffer_size + 2 elements in the array", type_name + ); + } + + DynamicBitset result(size); + for (int i = 0; i < buffer_size; ++i) { + if (!arr[i + 2].is()) { + return ConstructDeserializeError("Expect an integer in the array", type_name); + } + int64_t value = arr[i + 2].get(); + if (value < 0 || value > std::numeric_limits::max()) { + return ConstructDeserializeError( + "Integer in the array is " + std::to_string(value) + " and out of the uint32_t range", + type_name + ); + } + result.data_[i] = static_cast(value); + } + *bitset = std::move(result); + return std::nullopt; + } + + bool operator==(const DynamicBitset& other) const { + if (size_ != other.size_) return false; + if (buffer_size_ != other.buffer_size_) return false; + for (int i = 0; i < buffer_size_; ++i) { + if (data_[i] != other.data_[i]) return false; + } + return true; + } + + private: + static int LowestBit(uint32_t value) { +#ifdef __GNUC__ + return __builtin_ctz(value); +#else // __GNUC__ + // From https://stackoverflow.com/a/757266 + static const int MultiplyDeBruijnBitPosition[32] = {0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, + 15, 25, 17, 4, 8, 31, 27, 13, 23, 21, 19, + 16, 7, 26, 12, 18, 6, 11, 5, 10, 9}; + return MultiplyDeBruijnBitPosition[((uint32_t)((value & -value) * 0x077CB531U)) >> 27]; +#endif // __GNUC__ + } + + static int PopCount(uint32_t value) { +#ifdef __GNUC__ + return __builtin_popcount(value); +#elif defined(_MSC_VER) + return __popcnt(value); +#else + XGRAMMAR_LOG(FATAL) << "PopCount is not supported on this platform"; +#endif + } + + int DoFindZeroFrom(int first_block) const { + int position = -1; + for (int i = first_block; i < buffer_size_; ++i) { + if (data_[i] != ~static_cast(0)) { + position = i; + break; + } + } + if (position == -1) return -1; + return position * BITS_PER_BLOCK + LowestBit(~data_[position]); + } + + int DoFindOneFrom(int first_block) const { + int position = -1; + for (int i = first_block; i < buffer_size_; ++i) { + if (data_[i] != 0) { + position = i; + break; + } + } + if (position == -1) return -1; + return position * BITS_PER_BLOCK + LowestBit(data_[position]); + } + + // The size of the bitset. + int size_; + // The size of the buffer. + int buffer_size_; + // The buffer for the bitset. + uint32_t* data_; + // The internal buffer. It is empty if not needed. + std::vector internal_buffer_; + // Whether the buffer is internally managed. + bool is_internal_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_DYNAMIC_BITSET_H_ diff --git a/zig/vendor/xgrammar/cpp/support/encoding.h b/zig/vendor/xgrammar/cpp/support/encoding.h new file mode 100644 index 0000000000..195fde5228 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/encoding.h @@ -0,0 +1,459 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/encoding.h + * \brief Encoding and decoding from/to UTF-8 and escape sequence to/from codepoints. + */ +#ifndef XGRAMMAR_SUPPORT_ENCODING_H_ +#define XGRAMMAR_SUPPORT_ENCODING_H_ +// TODO(yixin): enhance performance + +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! \brief Represents a unicode codepoint. */ +using TCodepoint = int32_t; + +/*! + * \brief Represents an error when handling characters. Will be returned as a special TCodepoint + * value. + */ +enum CharHandlingError : TCodepoint { + /*! \brief The UTF-8 string is invalid. */ + kInvalidUTF8 = -10, + /*! \brief The escape sequence is invalid. */ + kInvalidEscape = -11, + /*! \brief The Latin-1 string is invalid. */ + kInvalidLatin1 = -12, +}; + +/******************** UTF-8 Handling ********************/ + +/*! + * \brief Print a codepoint to a UTF-8 string. + * \param codepoint The codepoint. + * \return The UTF-8 string. + */ +std::string CharToUTF8(TCodepoint codepoint); + +/*! + * \brief Handle the utf-8 first byte. + * \returns (is_valid, total_number_of_bytes, initial_codepoint). + */ +std::tuple HandleUTF8FirstByte(uint8_t byte); + +/*! + * \brief Parse all codepoints in a UTF-8 string. + * \param utf8 The UTF-8 string. + * \param perserve_invalid_bytes If the invalid UTF8 bytes will be preserved in the result. + * \return All codepoints. If the UTF-8 string is invalid, when perserve_invalid_bytes is false, + * the invalid bytes will be added to the result as a TCodepoint. Otherwise, the function will + * return {CharHandlingError::kInvalidUTF8}. + */ +std::vector ParseUTF8(const char* utf8, bool perserve_invalid_bytes = false); + +/*! + * \brief Parse the first codepoint in a UTF-8 string. + * \param utf8 The UTF-8 string. + * \return The codepoint and the number of bytes consumed. If the UTF-8 string is invalid, return + * {CharHandlingError::kInvalidUTF8, 0}. + */ +std::pair ParseNextUTF8(const char* utf8); + +/*! + * \brief Convert a Latin-1 string to a byte sequence. + * \param latin1 The Latin-1 string. + * \return The byte sequence. + */ +std::optional Latin1ToBytes(const std::string& latin1, std::string* result); + +/******************** Escape Handling ********************/ + +/*! + * \brief Convert a codepoint to a escaped string. If the codepoint is not printable, it will be + * escaped. By default the function support escape sequences in C ("\n", "\t", "\u0123"). User + * can specify more escape sequences using additional_escape_map. + * \param codepoint The codepoint. + * \param additional_escape_map A map from codepoint to escape sequence. If the codepoint is in + * the map, it will be escaped using the corresponding escape sequence. e.g. {{'-', "\\-"}}. + * \return The printable string. + */ +std::string EscapeString( + TCodepoint codepoint, + const std::unordered_map& additional_escape_map = {} +); + +/*! + * \brief Convert the given char to a escaped string that can be printed. + * \return The escaped string. + */ +std::string EscapeString(uint8_t raw_char); + +/*! + * \brief Convert the given string to a escaped string that can be printed. + * \return The escaped string. + */ +std::string EscapeString(std::string raw_str); + +/*! + * \brief Convert a hex character to an integer. + * \param c The hex character: 0-9, a-f, A-F. + * \return The integer value of the hex character. If the character is not a valid hex character, + * return -1. + */ +int HexCharToInt(char c); + +/*! + * \brief Parse the first escaped codepoint from a escaped string. data must start with a '\' + * character. + * \param data The escaped string. Can be TCodepoint* (e.g. string decoded from UTF-8) or char*. + * \param additional_escape_map A map from escape sequence to codepoint. If the escape sequence is + * in the map, it will be converted to the corresponding codepoint. e.g. {{"\\-", '-'}}. + * \return The codepoint and the number of bytes consumed. + */ +template +std::pair ParseNextEscaped( + const CharType* data, const std::unordered_map& additional_escape_map = {} +); + +/*! + * \brief Parse the first codepoint from a UTF-8 string. Also checks escape sequences and converts + * the escaped char to its original value. + * \param utf8 The UTF-8 string or the escape sequence. + * \param additional_escape_map A map from escape sequence to codepoint. If the escape sequence is + * in the map, it will be converted to the corresponding codepoint. e.g. {{"\\-", '-'}}. + * \return The codepoint and the number of bytes consumed. If the UTF-8 string is invalid, the + * function returns (CharHandlingError::kInvalidUTF8, 0). If the escape sequence is invalid, the + * function returns (CharHandlingError::kInvalidEscape, 0). + */ +std::pair ParseNextUTF8OrEscaped( + const char* utf8, const std::unordered_map& additional_escape_map = {} +); + +/******************** Implementation ********************/ + +inline std::string CharToUTF8(TCodepoint codepoint) { + XGRAMMAR_DCHECK(codepoint <= 0x10FFFF) << "Invalid codepoint: " << codepoint; + std::string utf8; + if (codepoint <= 0x7F) { + // 1-byte sequence + utf8 += static_cast(codepoint); + } else if (codepoint <= 0x7FF) { + // 2-byte sequence + utf8 += static_cast(0xC0 | ((codepoint >> 6) & 0x1F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } else if (codepoint <= 0xFFFF) { + // 3-byte sequence + utf8 += static_cast(0xE0 | ((codepoint >> 12) & 0x0F)); + utf8 += static_cast(0x80 | ((codepoint >> 6) & 0x3F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } else { + // 4-byte sequence + utf8 += static_cast(0xF0 | ((codepoint >> 18) & 0x07)); + utf8 += static_cast(0x80 | ((codepoint >> 12) & 0x3F)); + utf8 += static_cast(0x80 | ((codepoint >> 6) & 0x3F)); + utf8 += static_cast(0x80 | (codepoint & 0x3F)); + } + return utf8; +} + +inline std::tuple HandleUTF8FirstByte(uint8_t byte) { + static const std::array kFirstByteMask = {0x00, 0x7F, 0x1F, 0x0F, 0x07}; + // clang-format off + static const std::array kUtf8Bytes = { + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + 4, 4, 4, 4, 4, 4, 4, 4, -1, -1, -1, -1, -1, -1, -1, -1, + }; + // clang-format on + auto num_bytes = kUtf8Bytes[static_cast(byte)]; + if (num_bytes == -1) { + return {false, 0, 0}; + } + return {true, num_bytes, byte & kFirstByteMask[num_bytes]}; +} + +inline std::pair ParseNextUTF8(const char* utf8) { + auto [accepted, num_bytes, res] = HandleUTF8FirstByte(utf8[0]); + if (accepted) { + for (int i = 1; i < num_bytes; ++i) { + if (utf8[i] == 0 || (static_cast(utf8[i]) & 0xC0) != 0x80) { + // invalid utf8 + accepted = false; + break; + } + res = (res << 6) | (static_cast(utf8[i]) & 0x3F); + } + } + + if (!accepted) { + // invalid utf8 + return {CharHandlingError::kInvalidUTF8, 0}; + } + + return {res, num_bytes}; +} + +inline std::vector ParseUTF8(const char* utf8, bool perserve_invalid_bytes) { + std::vector codepoints; + while (*utf8 != 0) { + auto [codepoint, num_bytes] = ParseNextUTF8(utf8); + if (codepoint == CharHandlingError::kInvalidUTF8) { + if (perserve_invalid_bytes) { + codepoints.push_back(static_cast(static_cast(utf8[0]))); + utf8 += 1; + continue; + } else { + return {CharHandlingError::kInvalidUTF8}; + } + } + codepoints.push_back(codepoint); + utf8 += num_bytes; + } + return codepoints; +} + +/*! + \brief Convert a Latin-1 string to a byte sequence. + \param latin1 The Latin-1 string. + \param result The output byte sequence. + The function will convert each Latin-1 character to its corresponding byte(s). + For characters in the range [0x00, 0x7F], the corresponding byte is the same as the character. + Otherwise, the character should be encoded in two bytes in UTF-8: + - First byte: 110xxxxx (0xC0 | (char >> 6)) + - Second byte: 10xxxxxx (0x80 | (char & 0x3F)) + Example: + 0xC3 0xBF -> 0xFF + 'A' -> 'A' + \return std::nullopt if the conversion is successful. Otherwise, return + CharHandlingError::kInvalidLatin1 if the Latin-1 string is invalid. +*/ +inline std::optional Latin1ToBytes( + const std::string& latin1, std::string* result +) { + result->clear(); + result->reserve(latin1.size()); + + const size_t len = latin1.size(); + for (size_t i = 0; i < len; ++i) { + unsigned char c1 = static_cast(latin1[i]); + if (c1 < 0x80) { + result->push_back(static_cast(c1)); + } else { + if (i + 1 >= len) { + return CharHandlingError::kInvalidLatin1; + } + + unsigned char c2 = static_cast(latin1[i + 1]); + if ((c2 & 0xC0) != 0x80) { + return CharHandlingError::kInvalidLatin1; + } + + int code = ((c1 & 0x1F) << 6) | (c2 & 0x3F); + if (code < 0x80 || code > 0xFF) { + return CharHandlingError::kInvalidLatin1; + } + + result->push_back(static_cast(code)); + ++i; + } + } + + return std::nullopt; +} + +/*! + \brief Convert a byte sequence to a Latin-1 string. + \param Bytes The input byte sequence. + \param result The output Latin-1 string. + The function will convert each byte in the input to a Latin-1 character. + For bytes in the range [0x00, 0x7F], the corresponding Latin-1 character is the same as the byte. + For bytes in the range [0x80, 0xFF], the corresponding Latin-1 character is represented by two + bytes in UTF-8: + - First byte: 110xxxxx (0xC0 | (byte >> 6)) + - Second byte: 10xxxxxx (0x80 | (byte & 0x3F)) + Example: + 0xFF -> 0xC3 0xBF + 'A' -> 'A' +*/ +inline void ByteToLatin1(const std::string& bytes, std::string* result) { + result->clear(); + result->reserve(bytes.size()); + + for (unsigned char current_char : bytes) { + // Ascii character, directly add to result. + if (current_char <= 0x7F) { + result->push_back(static_cast(current_char)); + continue; + } + + // not Ascii character, convert to Latin-1. + unsigned char latin1_first_byte = 0; + unsigned char latin1_second_byte = 0; + + latin1_first_byte = 0xC0 | (current_char >> 6); + latin1_second_byte = 0x80 | (current_char & 0x3F); + result->push_back(static_cast(latin1_first_byte)); + result->push_back(static_cast(latin1_second_byte)); + } +} + +inline int HexCharToInt(char c) { + if (c >= '0' && c <= '9') { + return c - '0'; + } else if (c >= 'a' && c <= 'f') { + return c - 'a' + 10; + } else if (c >= 'A' && c <= 'F') { + return c - 'A' + 10; + } else { + return -1; + } +} + +inline std::string EscapeString( + TCodepoint codepoint, const std::unordered_map& additional_escape_map +) { + static const std::unordered_map kCodepointToEscape = { + {'\'', "\\\'"}, + {'\"', "\\\""}, + {'\?', "\\?"}, + {'\\', "\\\\"}, + {'\a', "\\a"}, + {'\b', "\\b"}, + {'\f', "\\f"}, + {'\n', "\\n"}, + {'\r', "\\r"}, + {'\t', "\\t"}, + {'\v', "\\v"}, + {'\0', "\\0"}, + {'\x1B', "\\e"} + }; + + if (auto it = additional_escape_map.find(codepoint); it != additional_escape_map.end()) { + return it->second; + } + + if (auto it = kCodepointToEscape.find(codepoint); it != kCodepointToEscape.end()) { + return it->second; + } + + if (codepoint >= 0x20 && codepoint <= 0x7E) { + return std::string({static_cast(codepoint)}); + } + + // convert codepoint to hex + char prefix = codepoint <= 0xFF ? 'x' : codepoint <= 0xFFFF ? 'u' : 'U'; + int width = codepoint <= 0xFF ? 2 : codepoint <= 0xFFFF ? 4 : 8; + std::stringstream ss; + ss << std::setfill('0') << std::setw(width) << std::hex << codepoint; + auto hex = ss.str(); + return std::string("\\") + prefix + hex; +} + +inline std::string EscapeString(uint8_t raw_char) { + return EscapeString(static_cast(raw_char)); +} + +inline std::string EscapeString(std::string raw_str) { + std::string res; + auto codepoints = ParseUTF8(raw_str.c_str(), true); + for (auto c : codepoints) { + res += EscapeString(c); + } + return res; +} + +template +std::pair ParseNextEscaped( + const CharType* data, const std::unordered_map& additional_escape_map +) { + // C escape characters + static const std::unordered_map kEscapeToCodepoint = { + // clang-format off + {'\'', '\''}, {'\"', '\"'}, {'?', '\?'}, {'\\', '\\'}, {'/', '/'}, {'a', '\a'}, + {'b', '\b'}, {'f', '\f'}, {'n', '\n'}, {'r', '\r'}, {'t', '\t'}, {'v', '\v'}, + {'0', '\0'}, {'e', '\x1B'} // clang-format on + }; + if (data[0] != '\\') { + return {CharHandlingError::kInvalidEscape, 0}; + } + + bool escape_char_in_escape_range = + static_cast(static_cast(data[1])) <= 128; + if (!escape_char_in_escape_range) { + return {CharHandlingError::kInvalidEscape, 0}; + } + + if (auto it = additional_escape_map.find(static_cast(data[1])); + it != additional_escape_map.end()) { + return {it->second, 2}; + } + if (auto it = kEscapeToCodepoint.find(static_cast(data[1])); + it != kEscapeToCodepoint.end()) { + return {it->second, 2}; + } + + if (data[1] == 'x') { + // arbitrary length hex + int len = 0; + TCodepoint codepoint = 0; + int32_t digit; + while ((digit = HexCharToInt(data[2 + len])) != -1) { + codepoint = codepoint * 16 + digit; + ++len; + } + if (len == 0) { + return {CharHandlingError::kInvalidEscape, 0}; + } + return {codepoint, len + 2}; + } else if (data[1] == 'u' || data[1] == 'U') { + // 4- or 8-digit hex + int len = data[1] == 'u' ? 4 : 8; + TCodepoint codepoint = 0; + + for (int i = 0; i < len; ++i) { + auto digit = HexCharToInt(data[i + 2]); + if (digit == -1) { + return {CharHandlingError::kInvalidEscape, 0}; + } + codepoint = codepoint * 16 + digit; + } + return {codepoint, len + 2}; + } else { + return {CharHandlingError::kInvalidEscape, 0}; + } +} + +inline std::pair ParseNextUTF8OrEscaped( + const char* utf8, const std::unordered_map& additional_escape_map +) { + if (utf8[0] != '\\') { + return ParseNextUTF8(utf8); + } + return ParseNextEscaped(utf8, additional_escape_map); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_ENCODING_H_ diff --git a/zig/vendor/xgrammar/cpp/support/int_set.h b/zig/vendor/xgrammar/cpp/support/int_set.h new file mode 100644 index 0000000000..f37f9a83ab --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/int_set.h @@ -0,0 +1,132 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/int_set.h + * \brief The header for utilities used in grammar-guided generation. + */ +#ifndef XGRAMMAR_SUPPORT_INT_SET_H_ +#define XGRAMMAR_SUPPORT_INT_SET_H_ + +#include +#include +#include +#include + +namespace xgrammar { + +/*! + * \brief Let lhs be the union of lhs and rhs. Suppose that both sets are sorted. + * \note No additional vectors are allocated, and the time complexity is O(n) + */ +inline void IntsetUnion(std::vector* lhs, const std::vector& rhs) { + int original_lhs_size = lhs->size(); + int rhs_size = rhs.size(); + + lhs->resize(original_lhs_size + rhs_size); + + auto it_lhs = lhs->rbegin() + rhs_size; + auto it_rhs = rhs.rbegin(); + auto it_result = lhs->rbegin(); + + while (it_lhs != lhs->rend() && it_rhs != rhs.rend()) { + if (*it_lhs > *it_rhs) { + *it_result = *it_lhs; + ++it_lhs; + } else if (*it_lhs < *it_rhs) { + *it_result = *it_rhs; + ++it_rhs; + } else { + *it_result = *it_lhs; + ++it_lhs; + ++it_rhs; + } + ++it_result; + } + + while (it_rhs != rhs.rend()) { + *it_result = *it_rhs; + ++it_result; + ++it_rhs; + } + + auto last = std::unique(lhs->begin(), lhs->end()); + lhs->erase(last, lhs->end()); +} + +/*! + * \brief Let lhs be the intersection of lhs and rhs. Suppose that both sets are sorted. + * \note No additional vector is allocated, and the time complexity is O(n). + * \note Support the case where lhs is the universal set by setting lhs to {-1}. The result will be + * rhs then. + */ +inline void IntsetIntersection(std::vector* lhs, const std::vector& rhs) { + if (lhs->size() == 1 && (*lhs)[0] == -1) { + *lhs = rhs; + return; + } + + auto it_lhs = lhs->begin(); + auto it_rhs = rhs.begin(); + auto it_result = lhs->begin(); + + while (it_lhs != lhs->end() && it_rhs != rhs.end()) { + if (*it_lhs < *it_rhs) { + ++it_lhs; + } else if (*it_lhs > *it_rhs) { + ++it_rhs; + } else { + *it_result = *it_lhs; + ++it_lhs; + ++it_rhs; + ++it_result; + } + } + lhs->erase(it_result, lhs->end()); +} + +/*! + * \brief Let lhs = lhs - rhs. Both sets must be sorted. + * \note In-place, no additional vector allocated, O(n) time. + */ +inline void IntsetDifference(std::vector* lhs, const std::vector& rhs) { + auto it_lhs = lhs->begin(); + auto it_rhs = rhs.begin(); + auto it_result = lhs->begin(); + + while (it_lhs != lhs->end() && it_rhs != rhs.end()) { + if (*it_lhs < *it_rhs) { + *it_result++ = *it_lhs++; + } else if (*it_lhs > *it_rhs) { + ++it_rhs; + } else { + ++it_lhs; + ++it_rhs; + } + } + while (it_lhs != lhs->end()) { + *it_result++ = *it_lhs++; + } + lhs->erase(it_result, lhs->end()); +} + +/*! + * \brief Compute result = [0, n) - excluded. excluded must be sorted with values in [0, n). + * \note O(n) time. + */ +inline void IntsetComplement( + std::vector* result, int32_t n, const std::vector& excluded +) { + result->clear(); + result->reserve(n - static_cast(excluded.size())); + auto it = excluded.begin(); + for (int32_t i = 0; i < n; ++i) { + if (it != excluded.end() && *it == i) { + ++it; + } else { + result->push_back(i); + } + } +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_INT_SET_H_ diff --git a/zig/vendor/xgrammar/cpp/support/json_parse.h b/zig/vendor/xgrammar/cpp/support/json_parse.h new file mode 100644 index 0000000000..3deba6916d --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/json_parse.h @@ -0,0 +1,118 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/json_parse.h + * \brief picojson::parse wrappers whose nesting depth is bounded by RecursionGuard. + * picojson recurses once per nesting level and has no depth limit of its own, so a deeply nested + * input would overflow the stack before any of the recursion guards in xgrammar run. + */ +#ifndef XGRAMMAR_SUPPORT_JSON_PARSE_H_ +#define XGRAMMAR_SUPPORT_JSON_PARSE_H_ + +#include + +#include +#include + +#include "logging.h" +#include "recursion_guard.h" + +namespace xgrammar { + +namespace detail { + +/*! + * \brief A picojson parse context that behaves like picojson::default_parse_context, but counts + * every nested array or object in a RecursionGuard. The maximum recursion depth of xgrammar thus + * applies while parsing, and exceeding it throws before the stack overflows. + */ +class DepthGuardedParseContext { + public: + DepthGuardedParseContext(picojson::value* out, int* depth) : out_(out), depth_(depth) {} + DepthGuardedParseContext(const DepthGuardedParseContext&) = delete; + DepthGuardedParseContext& operator=(const DepthGuardedParseContext&) = delete; + + bool set_null() { + *out_ = picojson::value(); + return true; + } + bool set_bool(bool b) { + *out_ = picojson::value(b); + return true; + } + bool set_int64(int64_t i) { + *out_ = picojson::value(i); + return true; + } + bool set_number(double f) { + *out_ = picojson::value(f); + return true; + } + template + bool parse_string(picojson::input& in) { + *out_ = picojson::value(picojson::string_type, false); + return picojson::_parse_string(out_->get(), in); + } + bool parse_array_start() { + *out_ = picojson::value(picojson::array_type, false); + return true; + } + template + bool parse_array_item(picojson::input& in, size_t) { + picojson::array& a = out_->get(); + a.push_back(picojson::value()); + RecursionGuard guard(depth_); + DepthGuardedParseContext child(&a.back(), depth_); + return picojson::_parse(child, in); + } + bool parse_array_stop(size_t) { return true; } + bool parse_object_start() { + *out_ = picojson::value(picojson::object_type, false); + return true; + } + template + bool parse_object_item(picojson::input& in, const std::string& key) { + picojson::object& o = out_->get(); + RecursionGuard guard(depth_); + DepthGuardedParseContext child(&o[key], depth_); + return picojson::_parse(child, in); + } + + private: + picojson::value* out_; + int* depth_; +}; + +} // namespace detail + +/*! + * \brief Parse the JSON value at the start of [begin, end). Same as + * picojson::parse(out, begin, end, err), but the nesting depth is bounded by the maximum + * recursion depth (see RecursionGuard). + * \return The iterator past the parsed value; begin if the depth limit was exceeded. + */ +template +inline Iter ParseJSON(picojson::value& out, Iter begin, Iter end, std::string* err) { + int depth = 0; + detail::DepthGuardedParseContext ctx(&out, &depth); + try { + return picojson::_parse(ctx, begin, end, err); + } catch (const LogFatalError& error) { + *err = error.what(); + return begin; + } +} + +/*! + * \brief Parse a JSON string. Same as picojson::parse(out, json), but the nesting depth is bounded + * by the maximum recursion depth (see RecursionGuard). + * \return The error message, empty on success. + */ +inline std::string ParseJSON(picojson::value& out, const std::string& json) { + std::string err; + ParseJSON(out, json.begin(), json.end(), &err); + return err; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_JSON_PARSE_H_ diff --git a/zig/vendor/xgrammar/cpp/support/json_serializer.h b/zig/vendor/xgrammar/cpp/support/json_serializer.h new file mode 100644 index 0000000000..aaff03cccf --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/json_serializer.h @@ -0,0 +1,640 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/json_serializer.h + * \brief A JSON-based serializer. Automatically generates serialization and deserialization logic + * from reflection. + */ +#ifndef XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ +#define XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "encoding.h" +#include "json_parse.h" +#include "logging.h" +#include "reflection.h" +#include "utils.h" +#include "xgrammar/exception.h" +#include "xgrammar/object.h" + +namespace xgrammar { + +/******************** Interfaces ********************/ + +/*! + * \brief Manages the version of the serialized object. The version will be added to the serialized + * object, and during deserialization, the object's version must match the current serialization + * version in xgrammar. + */ +class SerializeVersion { + public: + /*! + * \brief Returns the current serialization version. + */ + static std::string_view GetVersion() { return kXGrammarSerializeVersion; } + + /*! + * \brief Adds the version info to the serialized object. + */ + static void Apply(picojson::object* object); + + /*! + * \brief Checks if the serialized object's version matches the current serialization version. + * \return An error if the version does not exist or does not match. + */ + static std::optional Check(const picojson::object& object); + + private: + /*! + * \brief The key of the version info in the serialized object. + */ + static constexpr const char kXGrammarSerializeVersionKey[] = "__VERSION__"; + + /*! + * \brief The current serialization version. When the serialization result of any object in + * XGrammar is changed, this version should be bumped. + */ + static constexpr const char kXGrammarSerializeVersion[] = "v16"; +}; + +/*! + * \brief Serializes a value to a JSON value. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * SerializeJSONValue function. For reflection-based types, the serialization logic is automatically + * generated from the defined members. + * \param value The value to be serialized. + * \return The serialized JSON value. + */ +template +picojson::value AutoSerializeJSONValue(const T& value); + +/*! + * \brief Serializes a value to a JSON string. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * SerializeJSONValue function. For reflection-based types, the serialization logic is automatically + * generated from the defined members. + * \param value The value to be serialized. + * \param add_version Whether to add the version info to the serialized object. The addition is + * valid only when the serialized result is an object. + * \return The serialized JSON string. + */ +template +std::string AutoSerializeJSON(const T& value, bool add_version = false); + +/*! + * \brief Deserializes a value from a JSON value. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * DeserializeJSONValue function. For reflection-based types, the deserialization logic is + * automatically generated from the defined members. If a reflection-based type also defines + * `std::optional Validate() const`, it is called after all members are deserialized, + * and a returned message is reported as a deserialization error. This is where a type checks the + * invariants (index ranges, CSR layout, ...) that the constructors normally guarantee. + * \param result The pointer to the result to be deserialized. + * \param value The JSON value to be deserialized. + * \param type_name The name of the type to be deserialized. Used for error message. + * \return The deserialization error if any. + */ +template +std::optional AutoDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name = "" +); + +/*! + * \brief Deserializes a value from a JSON string. + * \details It supports STL types, PImpl types, reflection-based types (whose members are defined + * through XGRAMMAR_MEMBER_TABLE or XGRAMMAR_MEMBER_ARRAY), and types who have defined a global + * DeserializeJSONValue function. For reflection-based types, the deserialization logic is + * automatically generated from the defined members. + * \param result The pointer to the result to be deserialized. + * \param json_string The JSON string to be deserialized. + * \param check_version Whether to check the version info in the serialized object. The check is + * valid only when the serialized object is an object. + * \param type_name The name of the type to be deserialized. Used for error message. + * \return The deserialization error if any. + */ +template +std::optional AutoDeserializeJSON( + T* result, + const std::string& json_string, + bool check_version = false, + const std::string& type_name = "" +); + +/*! + * \brief Constructs a deserialize error with the given error message and type name. + * \param error_message The error message. + * \param type_name The name of the type. + * \return The constructed runtime error. + */ +inline SerializationError ConstructDeserializeError( + const std::string& error_message, const std::string& type_name +); + +/******************** Implementations ********************/ + +inline void SerializeVersion::Apply(picojson::object* object) { + XGRAMMAR_DCHECK(object != nullptr); + XGRAMMAR_DCHECK(object->find(kXGrammarSerializeVersionKey) == object->end()); + (*object)[kXGrammarSerializeVersionKey] = picojson::value(std::string(GetVersion())); +} + +inline std::optional SerializeVersion::Check(const picojson::object& object) { + if (object.find(kXGrammarSerializeVersionKey) == object.end()) { + return DeserializeVersionError( + std::string("Missing version in serialized object: ") + kXGrammarSerializeVersionKey + ); + } + if (object.at(kXGrammarSerializeVersionKey).get() != GetVersion()) { + return DeserializeVersionError( + std::string("Wrong version in serialized object: Got ") + + object.at(kXGrammarSerializeVersionKey).get() + ", expected " + + std::string(GetVersion()) + ); + } + return std::nullopt; +} + +/******************** Template Implementations ********************/ + +namespace detail::json_serializer { + +template +struct has_serialize_json_global : std::false_type {}; + +template +struct has_serialize_json_global< + T, + std::void_t()))>> : std::true_type { + static_assert( + std::is_same_v())), picojson::value>, + "SerializeJSONValue must be a global function returning picojson::value" + ); +}; + +template +struct has_deserialize_json_global : std::false_type {}; + +template +struct has_deserialize_json_global< + T, + std::void_t(), picojson::value{}, std::string{}) + )>> : std::true_type { + static_assert( + std::is_same_v< + decltype(DeserializeJSONValue(std::declval(), picojson::value{}, std::string{})), + std::optional>, + "DeserializeJSONValue must be a global function returning std::optional" + ); + static_assert( + std::is_default_constructible_v, + "global deserializer can only apply to a default constructible type" + ); +}; + +template +struct has_validate : std::false_type {}; + +template +struct has_validate().Validate())>> + : std::true_type { + static_assert( + std::is_same_v().Validate()), std::optional>, + "Validate must be a const member function returning std::optional" + ); +}; + +/*! \brief Runs T::Validate() on a deserialized value if the type defines one. */ +template +inline std::optional ValidateDeserialized( + const T& value, const std::string& type_name +) { + if constexpr (has_validate::value) { + if (auto error = value.Validate()) { + return ConstructDeserializeError(*error, type_name); + } + } + return std::nullopt; +} + +template +inline constexpr bool false_v = false; + +template +inline picojson::value TraitSerializeJSONValue(const T& value) { + using Functor = member_functor; + if constexpr (Functor::value == member_type::kConfig) { + if constexpr (Functor::has_names) { + // normal named struct + picojson::object obj; + obj.reserve(Functor::member_count); + visit_config([&](auto ptr, const char* name, std::size_t) { + XGRAMMAR_DCHECK(obj.find(name) == obj.end()); + obj[name] = AutoSerializeJSONValue(value.*ptr); + }); + return picojson::value(std::move(obj)); + } else if constexpr (Functor::member_count == 1) { + // optimize for single member unnamed structs + constexpr auto member_ptr = std::get<0>(Functor::members); + return AutoSerializeJSONValue(value.*member_ptr); + } else { + // normal unnamed struct + picojson::array arr; + arr.resize(Functor::member_count); + visit_config([&](auto ptr, const char*, std::size_t idx) { + arr[idx] = AutoSerializeJSONValue(value.*ptr); + }); + return picojson::value(std::move(arr)); + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Invalid trait type"); + return picojson::value{}; + } +} + +template +inline std::optional TraitDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name +) { + using Functor = member_functor; + if constexpr (Functor::value == member_type::kConfig) { + if constexpr (Functor::has_names) { + // normal named struct + if (!value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& obj = value.get(); + std::optional err = std::nullopt; + visit_config([&](auto ptr, const char* name, std::size_t idx) { + if (err) { + return; + } else if (obj.find(name) == obj.end()) { + err = ConstructDeserializeError("Missing member " + std::string(name), type_name); + } else if (auto e = AutoDeserializeJSONValue(&(result->*ptr), obj.at(name), type_name)) { + err = e; + } + }); + return err; + } else if constexpr (Functor::member_count == 1) { + // optimize for single member unnamed structs + constexpr auto member_ptr = std::get<0>(Functor::members); + return AutoDeserializeJSONValue(&(result->*member_ptr), value, type_name); + } else { + // normal unnamed struct + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + if (arr.size() != Functor::member_count) { + return ConstructDeserializeError( + "Wrong number of elements in array: Expected " + std::to_string(Functor::member_count) + + ", but got " + std::to_string(arr.size()), + type_name + ); + } + std::optional err = std::nullopt; + visit_config([&](auto ptr, const char*, std::size_t idx) { + if (err) { + return; + } else if (auto e = AutoDeserializeJSONValue(&(result->*ptr), arr[idx], type_name)) { + err = e; + } + }); + return err; + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Invalid trait type"); + XGRAMMAR_UNREACHABLE(); + } +} + +/******************** Customized Serialization ********************/ + +template > +inline picojson::value AutoSerializeJSONValuePImpl(const T& value) { + if (value.IsNull()) return picojson::value{}; + return AutoSerializeJSONValue(*value.ImplPtr()); +} + +template > +inline std::optional AutoDeserializeJSONValuePImpl( + T* result, const picojson::value& value, const std::string& type_name +) { + XGRAMMAR_DCHECK(result->IsNull()); + if (value.is()) { + *result = T{NullObj{}}; + return std::nullopt; + } + auto ptr = std::make_shared(); + if (auto error = AutoDeserializeJSONValue(ptr.get(), value, type_name)) { + return error; + } + *result = T(std::move(ptr)); + return std::nullopt; +} + +} // namespace detail::json_serializer + +inline SerializationError ConstructDeserializeError( + const std::string& error_message, const std::string& type_name +) { + if (type_name.empty()) { + return DeserializeFormatError("Deserialize error: " + error_message); + } else { + return DeserializeFormatError("Deserialize error for type " + type_name + ": " + error_message); + } +} + +template +inline picojson::value AutoSerializeJSONValue(const T& value) { + if constexpr (detail::json_serializer::has_serialize_json_global::value) { + // User-defined SerializeJSONValue (highest priority) + return SerializeJSONValue(value); + } else if constexpr (is_pimpl_class::value) { + // Library-customized serialization methods + return detail::json_serializer::AutoSerializeJSONValuePImpl(value); + } else if constexpr (member_trait::value != member_type::kNone) { + // Trait serialization methods + return detail::json_serializer::TraitSerializeJSONValue(value); + } else if constexpr (std::is_same_v) { + // Below is primitive types + return picojson::value(value); + } else if constexpr (std::is_integral_v || std::is_enum_v) { + return picojson::value(static_cast(value)); + } else if constexpr (std::is_floating_point_v) { + return picojson::value(static_cast(value)); + } else if constexpr (std::is_same_v) { + std::string result; + ByteToLatin1(value, &result); + return picojson::value(result); + } else if constexpr (is_std_optional::value) { + if (value.has_value()) { + return AutoSerializeJSONValue(*value); + } else { + return picojson::value{}; + } + } else if constexpr (is_std_pair::value) { + // std::pair: serialize as an array of size 2 + picojson::array arr; + arr.resize(2); + arr[0] = AutoSerializeJSONValue(value.first); + arr[1] = AutoSerializeJSONValue(value.second); + return picojson::value(std::move(arr)); + } else if constexpr (is_std_vector::value) { + picojson::array arr; + arr.reserve(value.size()); + for (const auto& item : value) { + arr.push_back(AutoSerializeJSONValue(item)); + } + return picojson::value(std::move(arr)); + } else if constexpr (is_std_unordered_set::value) { + std::vector ptr_vec; + ptr_vec.reserve(value.size()); + for (const auto& item : value) { + ptr_vec.push_back(&item); + } + std::sort(ptr_vec.begin(), ptr_vec.end(), [](const auto* a, const auto* b) { return *a < *b; }); + picojson::array arr; + arr.reserve(value.size()); + for (const auto* ptr : ptr_vec) { + arr.push_back(AutoSerializeJSONValue(*ptr)); + } + return picojson::value(std::move(arr)); + } else if constexpr (is_std_unordered_map::value) { + if constexpr (std::is_same_v) { + // unordered_map: map to json object + picojson::object obj; + obj.reserve(value.size()); + for (const auto& item : value) { + obj[item.first] = AutoSerializeJSONValue(item.second); + } + return picojson::value(std::move(obj)); + } else { + // unordered_map (T1 is not string): map to json array of array of size 2 + std::vector ptr_vec; + ptr_vec.reserve(value.size()); + for (const auto& item : value) { + ptr_vec.push_back(&item); + } + std::sort(ptr_vec.begin(), ptr_vec.end(), [](const auto* a, const auto* b) { + return a->first < b->first; + }); + picojson::array arr; + arr.reserve(value.size()); + for (const auto* ptr : ptr_vec) { + const auto& [key, item] = *ptr; + picojson::array sub_arr{AutoSerializeJSONValue(key), AutoSerializeJSONValue(item)}; + arr.push_back(picojson::value(std::move(sub_arr))); + } + return picojson::value(std::move(arr)); + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Cannot serialize this type"); + XGRAMMAR_UNREACHABLE(); + } +} + +template +inline std::string AutoSerializeJSON(const T& value, bool add_version) { + picojson::value json_value = AutoSerializeJSONValue(value); + if (add_version) { + XGRAMMAR_DCHECK(json_value.is()); + SerializeVersion::Apply(&json_value.get()); + } + return picojson::value(json_value).serialize(); +} + +template +inline std::optional AutoDeserializeJSONValue( + T* result, const picojson::value& value, const std::string& type_name +) { + static_assert(!std::is_const_v, "Cannot deserialize into a const type"); + if constexpr (detail::json_serializer::has_deserialize_json_global::value) { + return DeserializeJSONValue(result, value, type_name); + } else if constexpr (is_pimpl_class::value) { + return detail::json_serializer::AutoDeserializeJSONValuePImpl(result, value, type_name); + } else if constexpr (member_trait::value != member_type::kNone) { + if (auto error = detail::json_serializer::TraitDeserializeJSONValue(result, value, type_name)) { + return error; + } + return detail::json_serializer::ValidateDeserialized(*result, type_name); + } else if constexpr (std::is_same_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a boolean", type_name); + } + *result = value.get(); + return std::nullopt; + } else if constexpr (std::is_integral_v || std::is_enum_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect an integer", type_name); + } + *result = static_cast(value.get()); + return std::nullopt; + } else if constexpr (std::is_floating_point_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a floating point number", type_name); + } + *result = static_cast(value.get()); + return std::nullopt; + } else if constexpr (std::is_same_v) { + if (!value.is()) { + return ConstructDeserializeError("Expect a string", type_name); + } + // Now PicoJSON will convert byte sequence to latin-1 string. Convert it back to byte sequence. + auto error = Latin1ToBytes(value.get(), result); + if (error) { + return ConstructDeserializeError( + "XGramamr serializer will serialize byte sequence as latin-1 string, but got invalid " + "latin-1 string", + type_name + ); + } + return std::nullopt; + } else if constexpr (is_std_optional::value) { + // for the following container, T must be default constructible + if (value.is()) { + result->reset(); + return std::nullopt; + } else { + return AutoDeserializeJSONValue(&(result->emplace()), value, type_name); + } + } else if constexpr (is_std_pair::value) { + // std::pair: deserialize from an array of size 2 + if (!value.is()) { + return ConstructDeserializeError("Expect an array for deserializing pair", type_name); + } + const auto& arr = value.get(); + if (arr.size() != 2) { + return ConstructDeserializeError( + "Expect an array of size 2 for deserializing pair", type_name + ); + } + if (auto error = AutoDeserializeJSONValue(&(result->first), arr[0], type_name)) { + return error; + } + if (auto error = AutoDeserializeJSONValue(&(result->second), arr[1], type_name)) { + return error; + } + return std::nullopt; + } else if constexpr (is_std_vector::value) { + if (!value.is()) { + return ConstructDeserializeError("Expect an array", type_name); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + if (auto error = AutoDeserializeJSONValue(&(result->emplace_back()), item, type_name)) { + return error; + } + } + return std::nullopt; + } else if constexpr (is_std_unordered_set::value) { + if (!value.is()) { + return ConstructDeserializeError( + "Expect an array for deserializing unordered set", type_name + ); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + typename T::value_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, item, type_name)) { + return error; + } + result->emplace(std::move(item_value)); + } + return std::nullopt; + } else if constexpr (is_std_unordered_map::value) { + if constexpr (std::is_same_v) { + // unordered_map: convert from json object + if (!value.is()) { + return ConstructDeserializeError("Expect an object", type_name); + } + const auto& obj = value.get(); + result->clear(); + result->reserve(obj.size()); + for (const auto& [key, item] : obj) { + typename T::mapped_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, item, type_name)) { + return error; + } + result->try_emplace(key, std::move(item_value)); + } + return std::nullopt; + } else { + // unordered_map (T1 is not string): convert from json array of array of size 2 + if (!value.is()) { + return ConstructDeserializeError( + "Expect an array for deserializing unordered map", type_name + ); + } + const auto& arr = value.get(); + result->clear(); + result->reserve(arr.size()); + for (const auto& item : arr) { + if (!item.is()) { + return ConstructDeserializeError( + "Expect an array of array of size 2 for deserializing unordered map", type_name + ); + } + const auto& sub_arr = item.get(); + if (sub_arr.size() != 2) { + return ConstructDeserializeError( + "Expect an array of array of size 2 for deserializing unordered map", type_name + ); + } + typename T::key_type key_value{}; + if (auto error = AutoDeserializeJSONValue(&key_value, sub_arr[0], type_name)) { + return error; + } + typename T::mapped_type item_value{}; + if (auto error = AutoDeserializeJSONValue(&item_value, sub_arr[1], type_name)) { + return error; + } + result->emplace(std::move(key_value), std::move(item_value)); + } + return std::nullopt; + } + } else { + // should give an error in this case + static_assert(detail::json_serializer::false_v, "Cannot deserialize this type"); + XGRAMMAR_UNREACHABLE(); + } +} + +template +inline std::optional AutoDeserializeJSON( + T* result, const std::string& json_string, bool check_version, const std::string& type_name +) { + picojson::value json_value; + if (auto error = ParseJSON(json_value, json_string); !error.empty()) { + return InvalidJSONError(error); + } + if (check_version) { + XGRAMMAR_DCHECK(json_value.is()); + if (auto error = SerializeVersion::Check(json_value.get())) { + return error; + } + } + return AutoDeserializeJSONValue(result, json_value, type_name); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_JSON_SERIALIZER_H_ diff --git a/zig/vendor/xgrammar/cpp/support/logging.cc b/zig/vendor/xgrammar/cpp/support/logging.cc new file mode 100644 index 0000000000..f14cf8a2c1 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/logging.cc @@ -0,0 +1,24 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/logging.cc + */ +#include "logging.h" + +namespace xgrammar { + +#if XGRAMMAR_LOG_CUSTOMIZE == 0 + +LogFatal::Entry& LogFatal::GetEntry() { + static thread_local LogFatal::Entry result; + return result; +} + +const char* LogMessage::level_strings_[] = { + ": ", // XGRAMMAR_LOG_LEVEL_INFO + ": Debug: ", // XGRAMMAR_LOG_LEVEL_DEBUG + ": Warning: ", // XGRAMMAR_LOG_LEVEL_WARNING +}; + +#endif // XGRAMMAR_LOG_CUSTOMIZE + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/support/logging.h b/zig/vendor/xgrammar/cpp/support/logging.h new file mode 100644 index 0000000000..27bd30018d --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/logging.h @@ -0,0 +1,236 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/logging.h + * \brief A logging library that supports logging at different levels. + */ +#ifndef XGRAMMAR_SUPPORT_LOGGING_H_ +#define XGRAMMAR_SUPPORT_LOGGING_H_ + +#include +#include +#include +#include +#include + +#include "cpptrace.h" // IWYU pragma: keep + +/*! + * \brief Whether or not customize the logging output. + * If log customize is enabled, the user must implement + * xgrammar::LogFatalImpl and xgrammar::LogMessageImpl. + */ +#ifndef XGRAMMAR_LOG_CUSTOMIZE +#define XGRAMMAR_LOG_CUSTOMIZE 0 +#endif + +namespace xgrammar { + +/*! + * \brief Error type for errors from XGRAMMAR_CHECK, XGRAMMAR_ICHECK, and XGRAMMAR_LOG(FATAL). This + * error contains a backtrace of where it occurred. + */ +class LogFatalError : public std::runtime_error { + public: + /*! \brief Construct an error. Not recommended to use directly. Instead use XGRAMMAR_LOG(FATAL). + * + * \param file The file where the error occurred. + * \param lineno The line number where the error occurred. + * \param message The error message to display. + * \param time The time at which the error occurred. This should be in local time. + */ + LogFatalError( + const std::string& file, + int lineno, + const std::string& message, + std::time_t time = std::time(nullptr) + ) + : std::runtime_error(message), file_(file), lineno_(lineno), time_(time) { + std::ostringstream s; + s << "[" << std::put_time(std::localtime(&time), "%H:%M:%S") << "] " << file << ":" << lineno + << ": " << message << "\n"; + full_message_ = s.str(); + } + + /*! \return The file in which the error occurred. */ + const std::string& file() const { return file_; } + /*! \return The time at which this error occurred. */ + const std::time_t& time() const { return time_; } + /*! \return The line number at which this error occurred. */ + int lineno() const { return lineno_; } + /*! \return The error message. */ + const char* what() const noexcept override { return full_message_.c_str(); } + + private: + std::string file_; + int lineno_; + std::time_t time_; + std::string full_message_; +}; + +// Provide support for customized logging. +#if XGRAMMAR_LOG_CUSTOMIZE +/*! + * \brief Custom implementations of LogFatal. + * + * \sa XGRAMMAR_LOG_CUSTOMIZE + */ +[[noreturn]] void LogFatalImpl(const std::string& file, int lineno, const std::string& message); + +/*! + * \brief Custom implementations of LogMessage. + * + * \sa XGRAMMAR_LOG_CUSTOMIZE + */ +void LogMessageImpl(const std::string& file, int lineno, int level, const std::string& message); + +/*! + * \brief Class to accumulate an error message and throw it. Do not use + * directly, instead use LOG(FATAL). + */ +class LogFatal { + public: + LogFatal(const std::string& file, int lineno) : file_(file), lineno_(lineno) {} +#ifdef _MSC_VER +#pragma disagnostic push +#pragma warning(disable : 4722) +#endif + [[noreturn]] ~LogFatal() noexcept(false) { LogFatalImpl(file_, lineno_, stream_.str()); } +#ifdef _MSC_VER +#pragma disagnostic pop +#endif + std::ostringstream& stream() { return stream_; } + + private: + std::ostringstream stream_; + std::string file_; + int lineno_; +}; + +/*! + * \brief Class to accumulate an log message. Do not use directly, instead use + * LOG(INFO), LOG(WARNING), LOG(ERROR). + */ +class LogMessage { + public: + LogMessage(const std::string& file, int lineno, int level) + : file_(file), lineno_(lineno), level_(level) {} + ~LogMessage() { LogMessageImpl(file_, lineno_, level_, stream_.str()); } + std::ostringstream& stream() { return stream_; } + + private: + std::string file_; + int lineno_; + int level_; + std::ostringstream stream_; +}; + +#else // if XGRAMMAR_LOG_CUSTOMIZE + +/*! + * \brief Class to accumulate an error message and throw it. Do not use + * directly, instead use XGRAMMAR_LOG(FATAL). + * \note The `LogFatal` class is designed to be an empty class to reduce stack size usage. + * To play this trick, we use the thread-local storage to store its internal data. + */ +class LogFatal { + public: + LogFatal(const std::string& file, int lineno) { GetEntry().Init(file, lineno); } +#ifdef _MSC_VER +#pragma disagnostic push +#pragma warning(disable : 4722) +#endif + [[noreturn]] ~LogFatal() noexcept(false) { + GetEntry().Finalize(); + throw; + } +#ifdef _MSC_VER +#pragma disagnostic pop +#endif + std::ostringstream& stream() { return GetEntry().stream_; } + + private: + struct Entry { + void Init(const std::string& file, int lineno) { + this->stream_.str(""); + this->file_ = file; + this->lineno_ = lineno; + } + [[noreturn]] LogFatalError Finalize() noexcept(false) { + LogFatalError error(file_, lineno_, stream_.str()); + throw error; + } + std::ostringstream stream_; + std::string file_; + int lineno_; + }; + + static Entry& GetEntry(); +}; + +/*! + * \brief Class to accumulate an log message. Do not use directly, instead use + * XGRAMMAR_LOG(INFO), XGRAMMAR_LOG(WARNING), XGRAMMAR_LOG(ERROR). + */ +class LogMessage { + public: + LogMessage(const std::string& file, int lineno, int level) { + std::time_t t = std::time(nullptr); + stream_ << "[" << std::put_time(std::localtime(&t), "%H:%M:%S") << "] " << file << ":" << lineno + << level_strings_[level]; + } + ~LogMessage() { std::cerr << (stream_.str() + "\n"); } + std::ostringstream& stream() { return stream_; } + + private: + std::ostringstream stream_; + static const char* level_strings_[]; +}; + +#endif // XGRAMMAR_LOG_CUSTOMIZE + +#define XGRAMMAR_LOG_LEVEL_INFO 0 +#define XGRAMMAR_LOG_LEVEL_DEBUG 1 +#define XGRAMMAR_LOG_LEVEL_WARNING 2 +#define XGRAMMAR_LOG_LEVEL_FATAL 3 + +#define XGRAMMAR_LOG_INFO LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_INFO).stream() +#define XGRAMMAR_LOG_DEBUG LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_DEBUG).stream() +#define XGRAMMAR_LOG_WARNING LogMessage(__FILE__, __LINE__, XGRAMMAR_LOG_LEVEL_WARNING).stream() +#define XGRAMMAR_LOG_FATAL LogFatal(__FILE__, __LINE__).stream() + +/*! + * \brief Log a message at the given level. + * \param level The level of the message. Can be INFO, DEBUG, WARNING, FATAL. + */ +#define XGRAMMAR_LOG(level) XGRAMMAR_LOG_##level + +/*! + * \brief Check if the condition is true. Used for checking the correctness of user inputs. + * \param x The condition to check. + */ +#define XGRAMMAR_CHECK(x) \ + if (!(x)) LogFatal(__FILE__, __LINE__).stream() << "Check failed: (" #x << ") is false: " + +/*! + * \brief Check if the condition is true. Used to guarantee some internal conditions in the code. + * \param x The condition to check. + */ +#define XGRAMMAR_ICHECK(x) \ + if (!(x)) LogFatal(__FILE__, __LINE__).stream() << "Internal check failed: (" #x << ") is false: " + +/*! + * \brief Check if the condition is true. Used to guarantee some internal conditions in the code. + * \note This check is only enabled in debug mode. In release mode, it will be disabled for + * efficiency. This should be used in preference to XGRAMMAR_ICHECK. + * \param x The condition to check. + */ +#if XGRAMMAR_ENABLE_INTERNAL_CHECK +#define XGRAMMAR_DCHECK(x) XGRAMMAR_ICHECK(x) +#else +#define XGRAMMAR_DCHECK(x) \ + while (false) XGRAMMAR_ICHECK(x) +#endif // XGRAMMAR_ENABLE_DCHECK + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_LOGGING_H_ diff --git a/zig/vendor/xgrammar/cpp/support/memory_size.h b/zig/vendor/xgrammar/cpp/support/memory_size.h new file mode 100644 index 0000000000..d5ad99529f --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/memory_size.h @@ -0,0 +1,129 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/memory_size.h + * \brief Compute the memory consumption of a container in heap memory. + */ + +#ifndef XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ +#define XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ + +#include +#include +#include +#include +#include +#include + +#include "reflection.h" + +namespace xgrammar { + +/******************* MemorySize Procotol *******************/ + +template +inline constexpr std::size_t MemorySize(const T& value); + +template +inline constexpr std::size_t MemorySize(const std::pair& pair); + +template +inline constexpr std::size_t MemorySize(const std::tuple& tpl); + +template +inline constexpr std::size_t MemorySize(const std::optional& optional_value); + +inline std::size_t MemorySize(const std::vector& value); + +/******************* MemorySize Implementations *******************/ + +namespace detail::memory_size { + +/*! + * \brief Get the element type of a container. + */ +template +using ElementType = std::decay_t; + +/*! + * \brief A false value for static_assert. + */ +template +inline constexpr bool false_v = false; + +} // namespace detail::memory_size + +/*! + * \brief Compute the memory consumption of a value. + * \tparam T The type of the value. + * \param value The value. + * \return The memory consumption in heap memory of the value in bytes. + */ +template +inline constexpr std::size_t MemorySize(const T& value) { + if constexpr (is_pimpl_class::value) { + // Customized MemorySize + return MemorySize(*value.ImplPtr()); + } else if constexpr (std::is_trivially_copyable_v) { + // Primitive type + return 0; + } else if constexpr (std::is_trivially_copyable_v>) { + // Container of primitive type + return sizeof(detail::memory_size::ElementType) * std::size(value); + } else if constexpr (!std::is_trivially_copyable_v>) { + // Container of non-primitive type: sum up the memory size of all elements + std::size_t size = sizeof(detail::memory_size::ElementType) * std::size(value); + for (const auto& element : value) { + size += MemorySize(element); + } + return size; + } else { + static_assert(detail::memory_size::false_v, "MemorySize is not implemented for this type"); + } +} + +/*! + * \brief Compute the memory consumption of a pair. + * \tparam T1 The type of the first element. + * \tparam T2 The type of the second element. + * \param pair The pair. + * \return The memory consumption in heap memory of the pair. + */ +template +inline constexpr std::size_t MemorySize(const std::pair& pair) { + return MemorySize(pair.first) + MemorySize(pair.second); +} + +/*! + * \brief Compute the memory consumption of a tuple. + * \tparam Ts The types of the tuple. + * \param tpl The tuple. + * \return The memory consumption in heap memory of the tuple. + */ +template +inline constexpr std::size_t MemorySize(const std::tuple& tpl) { + return std::apply([](auto&&... elems) { return (MemorySize(elems) + ... + 0); }, tpl); +} + +/*! + * \brief Compute the memory consumption in heap memory. This function is specialized for + * std::optional. + * \tparam Tp The type of the optional. + * \param range The optional. + * \return The memory consumption in heap memory of the optional. + */ +template +inline constexpr std::size_t MemorySize(const std::optional& optional_value) { + return optional_value.has_value() ? MemorySize(*optional_value) : 0; +} + +/*! + * \brief Compute the memory consumption of a std::vector, which is bit-packed: it stores + * one bit (not one byte) per element. The generic container overload would over-count it by 8x. + * \param value The vector. + * \return The memory consumption in heap memory of the vector in bytes. + */ +inline std::size_t MemorySize(const std::vector& value) { return (value.size() + 7) / 8; } + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_MEMORY_SIZE_H_ diff --git a/zig/vendor/xgrammar/cpp/support/recursion_guard.cc b/zig/vendor/xgrammar/cpp/support/recursion_guard.cc new file mode 100644 index 0000000000..0a182863b4 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/recursion_guard.cc @@ -0,0 +1,50 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/recursion_guard.cc + */ + +#include "recursion_guard.h" + +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +int RecursionGuard::LoadMaxRecursionDepthFromEnv() { + const char* env_value = std::getenv(kMaxRecursionDepthEnvVar); + if (env_value == nullptr) { + return kDefaultMaxRecursionDepth; + } + + int value = 0; + std::string_view sv(env_value); + + // Convert the string to an integer + auto result = std::from_chars(sv.data(), sv.data() + sv.size(), value); + + // Check if the conversion is successful + if (result.ec == std::errc::invalid_argument || result.ec == std::errc::result_out_of_range || + result.ptr != sv.data() + sv.size() || value <= 0) { + XGRAMMAR_LOG(WARNING) << "Env variable XGRAMMAR_MAX_RECURSION_DEPTH is not a valid " + "integer or out of range: '" + << env_value << "', using default " << kDefaultMaxRecursionDepth; + return kDefaultMaxRecursionDepth; + } + + // Check if the value is too large + if (value > kMaxReasonableDepth) { + XGRAMMAR_LOG(WARNING) << "Env variable XGRAMMAR_MAX_RECURSION_DEPTH too large: " << value + << ", clamping to " << kMaxReasonableDepth; + return kMaxReasonableDepth; + } + + return value; +} + +std::atomic RecursionGuard::max_recursion_depth_{LoadMaxRecursionDepthFromEnv()}; + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/support/recursion_guard.h b/zig/vendor/xgrammar/cpp/support/recursion_guard.h new file mode 100644 index 0000000000..3c1ac4ba05 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/recursion_guard.h @@ -0,0 +1,127 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/recursion_guard.h + * \brief The header for recursion depth guard. + */ + +#ifndef XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ +#define XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ + +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief Thread-safe recursion guard to prevent stack overflow + * + * This class provides a RAII-style guard that tracks recursion depth + * and prevents excessive recursion that could lead to stack overflow. + * It uses atomic operations for thread safety and supports configurable + * maximum recursion depth. + */ +class RecursionGuard { + public: + /*! + * \brief Constructor that increments recursion depth + * \param current_recursion_depth Pointer to the current recursion depth counter + * \throws Logs fatal error if max recursion depth is exceeded + */ + explicit RecursionGuard(int* current_recursion_depth) + : current_depth_ptr_(current_recursion_depth) { + auto error = AddRecursionDepth(current_depth_ptr_); + XGRAMMAR_CHECK(error == std::nullopt) << error.value().what(); + } + + /*! + * \brief Reset the recursion depth to 0 + * \param current_recursion_depth Pointer to the current recursion depth counter + */ + static void ResetRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr); + *current_recursion_depth = 0; + } + + /*! + * \brief Destructor that decrements recursion depth + */ + ~RecursionGuard() { SubtractRecursionDepth(current_depth_ptr_); } + + /*! + * \brief Get the maximum allowed recursion depth + * \return Current maximum recursion depth limit + */ + static int GetMaxRecursionDepth() { return max_recursion_depth_.load(std::memory_order_relaxed); } + + /*! + * \brief Set the maximum allowed recursion depth + * \param max_depth New maximum recursion depth limit (must be positive) + */ + static void SetMaxRecursionDepth(int max_depth) { + if (max_depth <= 0 || max_depth > kMaxReasonableDepth) { + XGRAMMAR_LOG(FATAL + ) << "RecursionGuard: Maximum recursion depth must be positive and less than " + << kMaxReasonableDepth << ", got: " << max_depth; + } + max_recursion_depth_.store(max_depth, std::memory_order_relaxed); + } + + static std::optional AddRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr); + int current_depth = ++(*current_recursion_depth); + int max_depth = max_recursion_depth_.load(std::memory_order_relaxed); + if (current_depth > max_depth) { + return std::runtime_error( + "RecursionGuard: Maximum recursion depth exceeded. " + "Current depth: " + + std::to_string(current_depth) + ", Max allowed: " + std::to_string(max_depth) + ); + } + return std::nullopt; + } + + static void SubtractRecursionDepth(int* current_recursion_depth) { + XGRAMMAR_DCHECK(current_recursion_depth != nullptr && *current_recursion_depth > 0); + --(*current_recursion_depth); + } + + private: + /*! + * \brief Get the maximum allowed recursion depth from the environment variable. Used to + * initialize max_recursion_depth_. + * \return Current maximum recursion depth limit + */ + static int LoadMaxRecursionDepthFromEnv(); + + /*! + * \brief Pointer to the recursion depth counter + */ + int* current_depth_ptr_; + + /*! + * \brief Thread-safe global configuration + */ + static std::atomic max_recursion_depth_; + + /*! + * \brief Environment variable name for the maximum recursion depth + */ + inline constexpr static char kMaxRecursionDepthEnvVar[] = "XGRAMMAR_MAX_RECURSION_DEPTH"; + + /*! + * \brief Default maximum recursion depth + */ + inline constexpr static int kDefaultMaxRecursionDepth = 10000; + + /*! + * \brief Maximum reasonable recursion depth + */ + inline constexpr static int kMaxReasonableDepth = 1000000; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_RECURSION_GUARD_H_ diff --git a/zig/vendor/xgrammar/cpp/support/reflection.h b/zig/vendor/xgrammar/cpp/support/reflection.h new file mode 100644 index 0000000000..4703d9a746 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/reflection.h @@ -0,0 +1,300 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/reflection.h + * \brief The header for compile-time reflection. + */ + +#ifndef XGRAMMAR_SUPPORT_REFLECTION_H_ +#define XGRAMMAR_SUPPORT_REFLECTION_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace xgrammar { + +/******************** Core Reflection Types ********************/ + +/*! + * \brief The type of the member trait. + */ +enum class member_type { + kNone = 0, // this is default, which has no member trait + kConfig = 1, // this is a config with member pointers +}; + +/** + * \brief Base trait for member traits. + * + * \tparam T the type whose members are being reflected + * \details Provides a default trait indicating no members. + */ +template +struct member_trait { + static constexpr auto value = member_type::kNone; +}; + +/******************** STL and Custom Type Traits ********************/ + +template +struct is_std_array : std::false_type {}; + +template +struct is_std_array> : std::true_type {}; + +template +struct is_std_pair : std::false_type {}; + +template +struct is_std_pair> : std::true_type {}; + +template +struct is_std_tuple : std::false_type {}; + +template +struct is_std_tuple> : std::true_type {}; + +template +struct is_std_optional : std::false_type {}; + +template +struct is_std_optional> : std::true_type {}; + +template +struct is_std_vector : std::false_type {}; + +template +struct is_std_vector> : std::true_type {}; + +template +struct is_std_unordered_map : std::false_type {}; + +template +struct is_std_unordered_map> : std::true_type {}; + +template +struct is_std_unordered_set : std::false_type {}; + +template +struct is_std_unordered_set> : std::true_type {}; + +/*! + * \brief XGrammar specific: Check if a class is a PImpl class. + */ +template +struct is_pimpl_class : std::false_type {}; + +/*! + * \brief XGrammar specific: Check if a class is a PImpl class. It's true iff the class has a + * member `Impl` and the class is not the same as the `Impl` type. + */ +template +struct is_pimpl_class< + T, + std::void_t, void>>> + : std::true_type {}; + +/*! + * \brief A helper class to print the value when the condition is false. + */ +template +struct DebugAssert { + static_assert(condition); +}; + +/******************** Implementation Details ********************/ + +namespace detail::reflection { + +// We cannot use `static_assert(false)` even in unreachable code in `if constexpr`. +// See https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2023/p2593r1.html +// for more details. +// TL;DR: We use the following `false_v` as a workaround. +template +inline constexpr bool false_v = false; + +// Note that we don't allow empty tables now (that's uncommon). +template +inline constexpr auto make_member_table(X, Y second, Args... args) { + static_assert(sizeof...(args) % 2 == 0, "member table must be even"); + static_assert(std::is_same_v, "first member must be a c-string"); + static_assert(std::is_member_pointer_v, "second member must be a member pointer"); + if constexpr (sizeof...(args) == 0) { + return std::make_tuple(second); + } else { + return std::tuple_cat(std::make_tuple(second), make_member_table(args...)); + } +} + +template +inline constexpr auto make_name_table_aux(std::index_sequence, Tuple tuple) { + return std::array{std::get(tuple)...}; +} + +template +inline constexpr auto make_name_table(Args... args) { + constexpr auto N = sizeof...(args); + static_assert(N % 2 == 0, "name table must be even"); + return make_name_table_aux(std::make_index_sequence{}, std::make_tuple(args...)); +} + +template +inline void visit_config_impl(Fn&& fn, std::index_sequence) { + // This is a helper function to visit each member of the config. + // It uses fold expression to apply the function to each member. + static_assert(Ftor::value == member_type::kConfig, "T must be a config type"); + static constexpr auto get_name = [](std::size_t idx) { + if constexpr (Ftor::has_names) { + return Ftor::names[idx]; + } else { + return ""; + } + }; + return (fn(std::get(Ftor::members), get_name(Idx), Idx), ...); +} + +} // namespace detail::reflection + +/******************** Member Functors and Visitors ********************/ + +/*! + * \brief A functor that provides access to the members of a config type. + * It extracts the members from the `member_trait` specialization for the type `T`. + * A valid `member_trait` specialization must meet the following requirements: + * - It must have a static member `value` of type `member_type`, + * which must be either `kNone` or `kConfig`. + */ +template ::value> +struct member_functor { + static_assert(detail::reflection::false_v, "This specialization should never be used"); +}; + +/*! + * \brief A specialization of `member_functor` for config types. + * A valid `member_trait` specialization for a config type must meet the following: + * - It must have a static member `value` of type `member_type::kConfig`. + * - It must have a static tuple `members` that contains the member pointers. + * - It must have a static array `names` that contains the names of the members. + * - The size of `names` must be either 0 or equal to the number of members in `members`. + * - In the first case, `names` will be empty. + * - In the second case, `names` represent the printed name of each member. + */ +template +struct member_functor { + private: + using _trait_t = member_trait; + using _members_t = std::decay_t; + using _names_t = std::decay_t; + + public: + static constexpr auto value = member_type::kConfig; + static constexpr auto members = _trait_t::members; + static constexpr auto names = _trait_t::names; + static constexpr auto member_count = std::tuple_size_v<_members_t>; + static constexpr auto has_names = names.size() == member_count; + + // some static_asserts to check the member list and name list + static_assert(is_std_tuple<_members_t>::value, "Member list must be a tuple"); + static_assert(is_std_array<_names_t>::value, "Name list must be an array"); + static_assert(member_count > 0, "Member list must not be empty"); + static_assert( + names.size() == member_count || names.size() == 0, + "Name list must be empty or have the same size as member list" + ); +}; + +/*! + * \brief Visit the members of a config type. + * \tparam T The type of the config. + * \tparam Fn The type of the function to visit the members. + * \param fn The function to visit the members. fn's signature should be: + * \code{.cpp} + * (auto ptr, const char* name, size_t idx) -> void + * \endcode + * where `ptr` is the pointer to the member, `name` is the name of the member, and `idx` is the + * index of the member. + */ +template +inline void visit_config(Fn&& fn) { + using Ftor = member_functor; + return detail::reflection::visit_config_impl( + fn, std::make_index_sequence{} + ); +} + +/******************** Registration Macros ********************/ + +/** + * \brief Macros to define member traits for types. + * \details These macros are used to define the structural information of types + * for serialization and reflection purposes. + * + * Macros: + * - \c XGRAMMAR_MEMBER_TABLE: Defines a type with a table of (name, member pointer) pairs. + * - \c XGRAMMAR_MEMBER_ARRAY: Defines a type with an array of member pointers. + * + * Use the `_TEMPLATE` variants for template types. + * + * \example + * \code{.cpp} + * // Example of using XGRAMMAR_MEMBER_TABLE to register (name, member pointer) pairs + * struct SimpleClass { + * int a; + * double b; + * }; + * XGRAMMAR_MEMBER_TABLE(SimpleClass, "name_a", &SimpleClass::a, "name_b", &SimpleClass::b); + * + * // Or register members as an array with XGRAMMAR_MEMBER_ARRAY + * XGRAMMAR_MEMBER_ARRAY(SimpleClass, &SimpleClass::a, &SimpleClass::b); + * + * // Example of using XGRAMMAR_MEMBER_ARRAY to register members from a derived class + * struct Derived : SimpleClass { + * std::string c; + * }; + * XGRAMMAR_MEMBER_TABLE(Derived, "name_a", &Derived::a, "name_b", &Derived::b, "name_c", + * &Derived::c); + * + * // Example of using XGRAMMAR_MEMBER_ARRAY_TEMPLATE for a template type + * // If the default constructor/member is private, you need to declare a friend for member_trait. + * template + * struct TemplateClass { + * private: + * T value; + * TemplateClass() = default; + * friend struct member_trait; + * }; + * template + * XGRAMMAR_MEMBER_ARRAY_TEMPLATE(TemplateClass, &TemplateClass::value); + * \endcode + */ +#define XGRAMMAR_MEMBER_TABLE_TEMPLATE(Type, ...) \ + struct member_trait { \ + static constexpr auto value = member_type::kConfig; \ + static constexpr auto members = detail::reflection::make_member_table(__VA_ARGS__); \ + static constexpr auto names = detail::reflection::make_name_table(__VA_ARGS__); \ + } + +#define XGRAMMAR_MEMBER_ARRAY_TEMPLATE(Type, ...) \ + struct member_trait { \ + static constexpr auto value = member_type::kConfig; \ + static constexpr auto members = std::make_tuple(__VA_ARGS__); \ + static constexpr auto names = std::array{}; \ + } + +#define XGRAMMAR_MEMBER_TABLE(Type, ...) \ + template <> \ + XGRAMMAR_MEMBER_TABLE_TEMPLATE(Type, __VA_ARGS__) + +#define XGRAMMAR_MEMBER_ARRAY(Type, ...) \ + template <> \ + XGRAMMAR_MEMBER_ARRAY_TEMPLATE(Type, __VA_ARGS__) + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_REFLECTION_H_ diff --git a/zig/vendor/xgrammar/cpp/support/thread_pool.h b/zig/vendor/xgrammar/cpp/support/thread_pool.h new file mode 100644 index 0000000000..6bbec3c2d1 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/thread_pool.h @@ -0,0 +1,239 @@ +/*! + * Copyright (c) 2023 by Contributors + * \file xgrammar/support/thread_pool.h + * \brief Thread pool. + */ +#ifndef XGRAMMAR_SUPPORT_THREAD_POOL_H_ +#define XGRAMMAR_SUPPORT_THREAD_POOL_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +/*! + * \brief A thread pool implementation for parallel task execution. + * + * ThreadPool manages a pool of worker threads that can execute tasks asynchronously. + * Tasks are submitted to a queue and executed by available threads from the pool. + * The pool automatically handles thread synchronization and task distribution. + */ +class ThreadPool { + public: + /*! + * \brief Construct a new thread pool with the specified number of threads. + * \param num_threads Number of worker threads to create. Defaults to hardware concurrency. + * \note The pool starts the worker threads immediately upon construction. + */ + ThreadPool(size_t num_threads = std::thread::hardware_concurrency()) { + // Initialize thread pool with num_threads threads + for (size_t i = 0; i < num_threads; ++i) { + workers_.emplace_back([this] { + while (true) { + std::function task; + { + // Lock queue while waiting for new task + std::unique_lock lock(queue_mutex_); + queue_condition_.wait(lock, [this] { return shutdown_ || !task_queue_.empty(); }); + + // Exit thread if shutdown and queue is empty + if (shutdown_ && task_queue_.empty()) return; + + // Get task from queue + task = std::move(task_queue_.front()); + task_queue_.pop(); + } + try { + task(); + } catch (...) { + std::unique_lock lock(queue_mutex_); + if (!first_exception_) { + first_exception_ = std::current_exception(); + } + } + TaskComplete(); + } + }); + } + } + + /*! + * \brief Add a new task to be executed by the thread pool. + * \tparam F Type of the function to execute + * \tparam Args Types of the arguments to pass to the function + * \param f Function to execute + * \param args Arguments to pass to the function + * \return std::shared_future containing the result of the function call + * \note Tasks are executed in FIFO order but may complete in any order. + */ + template + auto Submit(F&& f, Args&&... args) -> std::shared_future> { + using return_type = std::invoke_result_t; + + // Package the task with its arguments into a shared pointer + auto task = std::make_shared>( + std::bind(std::forward(f), std::forward(args)...) + ); + + std::shared_future res = task->get_future().share(); + + { + std::unique_lock lock(queue_mutex_); + XGRAMMAR_CHECK(!shutdown_) << "Cannot submit task to stopped ThreadPool"; + ++unfinished_task_count_; // Increment task count + + // Directly add the task without wrapping + task_queue_.emplace([task]() { (*task)(); }); + } + queue_condition_.notify_one(); + return res; + } + + /*! + * \brief Add a new task to be executed by the thread pool without returning a future. + * \tparam F Type of the function to execute + * \tparam Args Types of the arguments to pass to the function + * \param f Function to execute + * \param args Arguments to pass to the function + * \note Tasks are executed asynchronously by the worker threads. + */ + template + void Execute(F&& f, Args&&... args) { + { + std::unique_lock lock(queue_mutex_); + XGRAMMAR_CHECK(!shutdown_) << "Cannot execute task in stopped ThreadPool"; + ++unfinished_task_count_; // Increment task count + + // Directly add the task without wrapping + task_queue_.emplace(std::bind(std::forward(f), std::forward(args)...)); + } + queue_condition_.notify_one(); + } + + /*! + * \brief Wait until all submitted tasks have finished. + * \note If a task submitted with Execute threw, the first such exception is rethrown here on the + * calling thread. An exception escaping a worker thread would otherwise terminate the process. + */ + void Wait() { + { + std::unique_lock lock(queue_mutex_); + tasks_done_condition_.wait(lock, [this] { return unfinished_task_count_ == 0; }); + } + RethrowTaskException(); + } + + /*! + * \brief Join all threads in the pool. + * + * Sets shutdown flag and waits for all threads to complete their current tasks + * before destroying the pool. Any remaining tasks in the queue will be executed + * before shutdown completes. + * \note If a task submitted with Execute threw, the first such exception is rethrown here on the + * calling thread. An exception escaping a worker thread would otherwise terminate the process. + */ + void Join() { + Shutdown(); + RethrowTaskException(); + } + + /*! + * \brief Destructor that ensures graceful shutdown of the thread pool. + */ + ~ThreadPool() { Shutdown(); } + + // Prevent copying or moving of the thread pool + ThreadPool(const ThreadPool&) = delete; + ThreadPool(ThreadPool&&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + ThreadPool& operator=(ThreadPool&&) = delete; + + private: + void Shutdown() { + { + std::unique_lock lock(queue_mutex_); + if (shutdown_) return; // Already shut down + shutdown_ = true; + } + + queue_condition_.notify_all(); // Wake up all threads so they can exit + for (std::thread& worker : workers_) { + if (worker.joinable()) worker.join(); // Wait for thread to finish + } + } + + void RethrowTaskException() { + std::exception_ptr exception; + { + std::unique_lock lock(queue_mutex_); + std::swap(exception, first_exception_); + } + if (exception) { + std::rethrow_exception(exception); + } + } + + void TaskComplete() { + std::unique_lock lock(queue_mutex_); + --unfinished_task_count_; + if (unfinished_task_count_ == 0) { + tasks_done_condition_.notify_all(); // Notify waiting threads + } + } + + /*! \brief Thread container */ + std::vector workers_; + /*! \brief Task queue */ + std::queue> task_queue_; + /*! \brief Mutex to protect task queue */ + std::mutex queue_mutex_; + /*! \brief Condition variable for thread synchronization */ + std::condition_variable queue_condition_; + /*! \brief Condition variable for task completion */ + std::condition_variable tasks_done_condition_; + /*! \brief Flag to indicate thread pool shutdown */ + bool shutdown_ = false; + /*! \brief Number of unfinished tasks */ + int unfinished_task_count_ = 0; + /*! \brief The first exception thrown by a task, rethrown by Wait() or Join() */ + std::exception_ptr first_exception_ = nullptr; +}; + +inline void ParallelFor(int low, int high, int num_threads, std::function f) { + if (high - low == 1) { + f(low); + return; + } + + ThreadPool pool(num_threads); + + int total = high - low; + int chunk_size = (total + num_threads - 1) / num_threads; + + for (int t = 0; t < num_threads; ++t) { + int start = low + t * chunk_size; + int end = std::min(start + chunk_size, high); + + if (start >= end) break; // No more iterations to process + + pool.Execute([f, start, end]() { + for (int i = start; i < end; ++i) { + f(i); + } + }); + } + pool.Join(); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_THREAD_POOL_H_ diff --git a/zig/vendor/xgrammar/cpp/support/thread_safe_cache.h b/zig/vendor/xgrammar/cpp/support/thread_safe_cache.h new file mode 100644 index 0000000000..28859db66c --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/thread_safe_cache.h @@ -0,0 +1,404 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/thread_safe_cache.h + * \brief The header for thread-safe caching functionality. + */ +#ifndef XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ +#define XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ + +#include +#include // IWYU pragma: keep +#include +#include +#include +#include +#include +#include +#include +#include + +#include "container.h" + +namespace xgrammar { + +/*! + * \brief Primary template for ThreadSafeCache + * \details This class provides thread-safe caching functionality in two forms: + * 1. Single value cache when only Value template parameter is provided + * 2. Key-value cache when both Key and Value template parameters are provided + */ +template +class ThreadSafeCache; + +/*! + * \brief Thread-safe cache for a single computed value + * \tparam Value The type of value being cached + * \details Specialization that provides: + * - Thread-safe access to a single cached value + * - Lazy computation on first access + * - Reader-writer locking for concurrent reads + */ +template +class ThreadSafeCache { + public: + /*! + * \brief Constructs a new single-value cache + * \param compute The function that computes the cached value + */ + explicit ThreadSafeCache(std::function compute) : compute_(std::move(compute)) {} + + /*! + * \brief Gets or computes the cached value + * \return The cached or newly computed value + */ + Value Get() { + // First try reading from cache with shared lock + { + std::shared_lock cache_lock(cache_mutex_); + if (cache_.has_value()) { + return cache_.value(); // Cache hit + } + } + + // Acquire exclusive lock to compute value + std::unique_lock cache_lock(cache_mutex_); + + // Double-check to prevent redundant computation + if (cache_.has_value()) { + return cache_.value(); + } + + Value value = compute_(); + XGRAMMAR_DCHECK(!cache_.has_value()); + cache_ = value; + return value; + } + + /*! + * \brief Clears the cached value + * This function removes the cached value, so the next call to Get() will recompute it. + */ + void Clear() { + std::unique_lock cache_lock(cache_mutex_); + cache_.reset(); + } + + private: + /*! \brief Optional container holding the cached value */ + std::optional cache_; + /*! \brief Function used to compute the value when not cached */ + std::function compute_; + /*! \brief Reader-writer lock protecting access to cache_ */ + std::shared_mutex cache_mutex_; +}; + +/*! + * \brief A thread-safe key-value cache with on-demand computation + * \tparam Key The type of keys used to lookup values. Should be hashable. + * \tparam Value The type of values stored in the cache + * \details This cache provides thread-safe access to computed values with the following features: + * - Lazy computation: Values are only computed when first requested + * - Thread safety: Uses reader-writer locks for concurrent reads + * - Parallel computation: Different keys can be computed simultaneously + * - Double-checked locking: Prevents redundant computation + */ +template +class ThreadSafeCache { + public: + /*! + * \brief Constructs a new thread-safe cache + * \param compute The function that computes values for uncached keys + */ + explicit ThreadSafeCache(std::function compute) + : compute_(std::move(compute)) {} + + /*! + * \brief Gets or computes the value for a key + * \param key The key to lookup + * \return The cached or newly computed value of the key + */ + Value Get(const Key& key) { + // Why we need this: + // - When adding new elements to a unordered_map, the map may be rehashed, + // - which means all the iterators may be invalidated. + // - However, cppreference says: + // - "References and pointers to either key or data stored in the container are only invalidated + // - by erasing that element, even when the corresponding iterator is invalidated." + // - (See https://en.cppreference.com/w/cpp/container/unordered_map) + // - Therefore, we should maintain 2 locks. + // - When we add something to the cache, we should hold the cache_mutex_. + // - When we erase something from the cache, we should hold the clear_mutex_. + + auto erase_lock = std::shared_lock(erase_mutex_); + + // First attempt to read from cache_ + { + auto cache_lock = std::shared_lock(cache_mutex_); + auto it = cache_.find(key); + if (it != cache_.end()) { // Cache hit + auto& entry = it->second; // The iterator is invalidated after releasing the lock + cache_lock.unlock(); // Therefore, we should hold the entry by reference first + + // We should not hold lock here, since this function may be blocking. + return entry.get(compute_, key); + } + } + + // Acquire exclusive lock to compute value + { + auto cache_lock = std::unique_lock(cache_mutex_); + auto& entry = cache_[key]; // Create a new entry + cache_lock.unlock(); // Release the lock before blocking + + // We should not hold lock here, since this function may be blocking. + return entry.get(compute_, key); + } + } + + /*! + * \brief Clears all cached values and associated per-key mutexes + * This function removes all cached key-value pairs, so subsequent calls to Get() will recompute + * them. + */ + void Clear() { + auto erase_lock = std::unique_lock(erase_mutex_); + cache_.clear(); + } + + private: + struct Entry { + Value value; + std::once_flag flag; + const Value& get(const std::function& f, const Key& key) { + // block in this lambda until the value is computed + std::call_once(flag, [&] { value = f(key); }); + return value; + } + }; + + /*! \brief The cache mapping keys to computed values */ + std::unordered_map cache_; + /*! \brief The function used to compute values for uncached keys */ + std::function compute_; + /*! \brief Reader-writer lock protecting access to cache_ */ + std::shared_mutex cache_mutex_; + /*! \brief Mutex protecting removing elements */ + std::shared_mutex erase_mutex_; +}; + +namespace details { + +template +class LRUCacheImpl { + public: + struct Entry { + Value value; // value of the node + int index; // node index + }; + + /*! \brief Visits the node and moves it to the back of the LRU list. Return its value. */ + const Value& LRUVisit(const std::pair& pair) { + const auto& entry = pair.second; + lru_list_.MoveBack(entry.index); + return entry.value; + } + + /*! \brief Initializes the node with the given value and moves it to the back of the LRU list. */ + void LRUInit(std::pair& pair, const Value& init) { + auto& entry = pair.second; + entry.value = init; + entry.index = lru_list_.PushBack(&pair).Index(); + } + + /*! + * \brief Evicts the least recently used nodes until the predicate returns false. + * \param predicate The function that returns true if eviction should continue. + * \param evict The function takes a value and returns true if the value can be evicted. + * This will be only called when the predicate returns true. + * If this function returns true, it should update the size information before return. + * \details This function will evict the least recently used nodes until the predicate returns + * false. The evict function will be called for each node to determine if it should be evicted. + */ + template + void LRUEvict(const Predicate& predicate, const Evict& evict) { + if (!predicate()) return; + + auto iter = lru_list_.begin(); + if (iter == lru_list_.end()) return; + + do { + auto& [key, entry] = **iter; + if (evict(entry.value)) { + iter = lru_list_.Erase(iter); + map_.erase(key); + } else { + ++iter; // simply skip those waiting for computation + } + } while (predicate() && iter != lru_list_.end()); + } + + std::unordered_map& GetMap() { return map_; } + + private: + std::unordered_map map_; + List*> lru_list_; +}; + +} // namespace details + +/** + * \brief A thread-safe key-value cache with on-demand computation and LRU eviction + * \tparam Key The type of keys used to lookup values. Should be hashable. + * \tparam Value The type of values stored in the cache + * \tparam Computer The functor that computes values for uncached keys + * \tparam SizeEstimator The functor that estimates the size of a value in bytes + * \details This cache provides thread-safe access to computed values with the following features: + * - Lazy computation: Values are only computed when first requested + * - LRU eviction: When the cache is full, the least recently used value is evicted + * - Thread safety: Uses reader-writer locks for concurrent reads + * \attention User should guarantee the following: + * 1. The policy class should provide a compute method that takes a key and returns a value. + * 2. The value type should have a MemorySize method that returns the size of the value in bytes. + */ +template +class ThreadSafeLRUCache { + private: + struct SizedValue { + Value value; + std::size_t size; + }; + + public: + inline static constexpr std::size_t kUnlimitedSize = static_cast(-1); + + explicit ThreadSafeLRUCache( + std::size_t max_size = kUnlimitedSize, + const Computer& computer = Computer{}, + const SizeEstimator& size_estimator = SizeEstimator{} + ) + : max_size_(max_size), computer_(computer), size_estimator_(size_estimator), cache_() {} + + std::size_t MaxMemorySize() const { return max_size_; } + std::size_t MemorySize() const { return current_size_; } + + Value Get(const Key& key) { + auto future = GetFuture(key); + return future.get().value; + } + + void Clear() { + // Remove all the ready entries. + const auto lock_map = std::lock_guard{map_mutex_}; + if (this->max_size_ == kUnlimitedSize) + cache_.GetMap().clear(); + else + cache_.LRUEvict( + [] { return true; }, + [&](const std::shared_future& value) { + // always evict and block until the value is ready + try { + current_size_ -= value.get().size; + } catch (...) { + // fine, just ignore the exception, size is not updated + } + return true; + } + ); + } + + private: + std::shared_future GetFuture(const Key& key) { + if (this->max_size_ == kUnlimitedSize) return GetFutureUnlimited(key); + auto& map = cache_.GetMap(); + + { + auto lock_map = std::shared_lock{map_mutex_}; + auto it = map.find(key); + if (it != map.end()) { + // We only need to hold LRU lock when shared lock is held here. + // When unique lock of map_mutex_ is held, only 1 thread can access the + // LRU list at the same time, so we do not need to hold the LRU lock then. + const auto lock_lru = std::lock_guard{lru_mutex_}; + return cache_.LRUVisit(*it); + } + } + + auto task = std::packaged_task{[this, &key] { + auto value = computer_(key); + auto result = SizedValue{value, size_estimator_(value)}; + current_size_ += result.size; + return result; + }}; + + auto lock_map = std::unique_lock{map_mutex_}; + auto [it, success] = map.try_emplace(key); + if (!success) return cache_.LRUVisit(*it); + + // in this case, we insert the task, and we need to compute the value + auto future = task.get_future().share(); + + // perform eviction if the cache is full + cache_.LRUInit(*it, future); + cache_.LRUEvict( + [&] { return current_size_ > max_size_; }, + [&](const std::shared_future& value) { + using namespace std::chrono_literals; + // if not ready, then do not wait and block here + if (value.wait_for(0s) != std::future_status::ready) return false; + try { + current_size_ -= value.get().size; + } catch (...) { + // fine, just ignore the exception, size is not updated + } + return true; + } + ); + + // perform the costly computation outside all locks + lock_map.unlock(); + task(); + return future; + } + + std::shared_future GetFutureUnlimited(const Key& key) { + auto& map = cache_.GetMap(); + + { + auto lock_map = std::shared_lock{map_mutex_}; + auto it = map.find(key); + if (it != map.end()) return it->second.value; + } + + auto task = std::packaged_task{[this, &key] { + auto value = computer_(key); + auto result = SizedValue{value, size_estimator_(value)}; + current_size_ += result.size; + return result; + }}; + + auto lock_map = std::unique_lock{map_mutex_}; + auto [it, success] = map.try_emplace(key); + if (!success) return it->second.value; + + auto future = task.get_future().share(); + it->second.value = future; + + // perform the costly computation outside all locks + lock_map.unlock(); + task(); + return future; + } + + private: + const std::size_t max_size_; + const Computer computer_; + const SizeEstimator size_estimator_; + details::LRUCacheImpl> cache_; + std::atomic_size_t current_size_{0}; + std::shared_mutex map_mutex_; + std::mutex lru_mutex_; +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_THREAD_SAFE_CACHE_H_ diff --git a/zig/vendor/xgrammar/cpp/support/union_find_set.h b/zig/vendor/xgrammar/cpp/support/union_find_set.h new file mode 100644 index 0000000000..be82a51674 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/union_find_set.h @@ -0,0 +1,105 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/support/union_find_set.h + */ +#ifndef XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ +#define XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ + +#include +#include +#include +#include + +#include "logging.h" + +namespace xgrammar { + +template +class UnionFindSet { + private: + std::unordered_map> element_to_parent_and_size_; + + public: + UnionFindSet() = default; + + /*! + * \brief Add a new element to the union-find set. + * \param element The element to add. + * \return True if the element was added successfully, false if it already exists. + */ + bool Add(const T& element) { + if (element_to_parent_and_size_.find(element) != element_to_parent_and_size_.end()) { + return false; // Element already exists. + } + element_to_parent_and_size_[element] = {element, 1}; + return true; + } + + /*! \brief Clear the union find set.*/ + void Clear() { element_to_parent_and_size_.clear(); } + + /*! + * \brief Find the representative of the set containing the element. + * \param element The element to find. + * \return The representative of the set containing the element. + */ + T Find(const T& element) { + XGRAMMAR_CHECK(element_to_parent_and_size_.find(element) != element_to_parent_and_size_.end()) + << "Element not found in union-find set."; + if (element_to_parent_and_size_[element].first != element) { + // Path compression. + element_to_parent_and_size_[element].first = Find(element_to_parent_and_size_[element].first); + } + return element_to_parent_and_size_[element].first; + } + + /*! + * \brief Union two elements into the same set. + * \param a The first element. + * \param b The second element. + */ + void Union(const T& a, const T& b) { + XGRAMMAR_CHECK(element_to_parent_and_size_.find(a) != element_to_parent_and_size_.end()) + << "Element " << a << " not found in union-find set."; + XGRAMMAR_CHECK(element_to_parent_and_size_.find(b) != element_to_parent_and_size_.end()) + << "Element " << b << " not found in union-find set."; + T root_a = Find(a); + T root_b = Find(b); + if (root_a == root_b) { + return; + } + if (element_to_parent_and_size_[root_a].second < element_to_parent_and_size_[root_b].second) { + std::swap(root_a, root_b); + // Make sure root_a is the larger set. + } + element_to_parent_and_size_[root_b].first = root_a; + element_to_parent_and_size_[root_a].second += element_to_parent_and_size_[root_b].second; + } + + int Count(const T& element) const { return element_to_parent_and_size_.count(element); } + + std::vector> GetAllSets() { + std::vector> result; + std::unordered_map root_to_set; + for (const auto& [value, _] : element_to_parent_and_size_) { + auto root = Find(value); + if (root_to_set.find(root) == root_to_set.end()) { + result.emplace_back(); + root_to_set[root] = result.size() - 1; + } + result[root_to_set[root]].push_back(value); + } + // Sort result to make it deterministic + for (auto& vec : result) { + std::sort(vec.begin(), vec.end()); + } + std::sort(result.begin(), result.end(), [](const std::vector& v1, const std::vector& v2) { + return v1.front() < v2.front(); + }); + return result; + } +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_UNION_FIND_SET_H_ diff --git a/zig/vendor/xgrammar/cpp/support/utils.h b/zig/vendor/xgrammar/cpp/support/utils.h new file mode 100644 index 0000000000..d8c974adee --- /dev/null +++ b/zig/vendor/xgrammar/cpp/support/utils.h @@ -0,0 +1,451 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/support/utils.h + * \brief Utility functions. + */ +#ifndef XGRAMMAR_SUPPORT_UTILS_H_ +#define XGRAMMAR_SUPPORT_UTILS_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "logging.h" + +/****************** Hash Library ******************/ + +namespace xgrammar { + +/*! + * \brief Hash and combine value into seed. + * \ref https://www.boost.org/doc/libs/1_84_0/boost/intrusive/detail/hash_combine.hpp + */ +inline void HashCombineBinary(uint64_t& seed, uint64_t value) { + seed ^= value + 0x9e3779b97f4a7c15ull + (seed << 6) + (seed >> 2); +} + +/*! + * \brief Find the hash sum of several size_t args. + */ +template +inline uint64_t HashCombine(Args... args) { + uint64_t seed = 0; + (..., HashCombineBinary(seed, args)); + return seed; +} + +/*! + * \brief Helper class to define the hash function for a struct by its members. + */ +template +struct HashByMembers { + std::size_t operator()(T const& x) const noexcept { + return HashCombine(std::hash>{}(x.*Members)...); + } +}; + +} // namespace xgrammar + +/*! + * \brief Define a hash function for a struct by its members in namespace std. Should be used + * outside of namespace xgrammar. + * \param Type The type of the struct. + * \param ... The member pointers of the struct. + * \example + * \code + * // In the global namespace + * XGRAMMAR_HASH_BY_MEMBERS(Type, &Type::member1, &Type::member2, &Type::member3); + * \endcode + */ +#define XGRAMMAR_HASH_BY_MEMBERS(Type, ...) \ + namespace std { \ + template <> \ + struct hash : public xgrammar::HashByMembers {}; \ + } + +/*! + * \brief Empty specialization of XGRAMMAR_HASH_BY_MEMBERS. + */ +#define XGRAMMAR_HASH_BY_MEMBERS_EMPTY(Type) \ + namespace std { \ + template <> \ + struct hash : public xgrammar::HashByMembers {}; \ + } + +namespace std { + +/*! + * \brief Define the hash function for std::pair. + */ +template +struct hash> { + size_t operator()(const std::pair& pair) const noexcept { + return xgrammar::HashCombine(std::hash{}(pair.first), std::hash{}(pair.second)); + } +}; + +/*! + * \brief Define the hash function for std::tuple. + */ +template +struct hash> { + size_t operator()(const std::tuple& tuple) const noexcept { + return std::apply( + [](const Args&... args) { return xgrammar::HashCombine(std::hash{}(args)...); }, tuple + ); + } +}; + +/*! + * \brief Define the hash function for std::vector. + */ +template +struct hash> { + size_t operator()(const std::vector& vec) const { + uint32_t seed = 0; + for (const auto& item : vec) { + xgrammar::HashCombineBinary(seed, std::hash{}(item)); + } + return seed; + } +}; + +} // namespace std + +namespace xgrammar { + +/****************** Result Library ******************/ + +/*! + * \brief A partial result type that can be used to construct a Result. Holds a result value or an + * error value. + * \tparam T The type of the value + * \tparam IsOk Whether the result is ok + */ +template +struct PartialResult { + template + PartialResult(Args&&... args) : value(std::forward(args)...) {} + T value; +}; + +/*! + * \brief Construct a success result with the arguments to construct a T. + * \tparam T The type of the success value + * \tparam Args The types of the arguments to construct a T + * \param args The arguments to construct a T + * \return A PartialResult with the arguments to construct a T + * \example + * \code + * // Call the constructor of T with the arguments + * return ResultOk(1, 2, 3); + * \endcode + */ +template +inline PartialResult ResultOk(Args&&... args) { + return PartialResult{std::forward(args)...}; +} + +/*! + * \brief Construct a success result with a universal reference (both lvalue and rvalue) + * \tparam T The type of the success value + * \param value The universal reference to the success value + * \return A PartialResult with the universal reference to the success value + * \example + * \code + * T value = T(1, 2, 3); + * // Move the value to the PartialResult + * return ResultOk(std::move(value)); + * \endcode + */ +template +inline PartialResult ResultOk(T&& value) { + return PartialResult{std::forward(value)}; +} + +/*! + * \brief Construct a error result with the arguments to construct a E. + * \tparam E The type of the error value. Default to std::runtime_error. + * \tparam Args The types of the arguments to construct a E + * \param args The arguments to construct a E + * \return A PartialResult with the arguments to construct a E + * \example + * \code + * // Construct a std::runtime_error with a error + * std::runtime_error error("Message"); + * return ResultErr(std::move(error)); + * \endcode + * \code + * // Construct a std::runtime_error with its argument + * return ResultErr("Error"); + * \endcode + * \code + * // Construct an E error with its argument + * return ResultErr("Error"); + * \endcode + */ +template +inline PartialResult ResultErr(Args&&... args) { + return PartialResult{std::forward(args)...}; +} + +/*! + * \brief Construct a error result with a universal reference (both lvalue and rvalue) + * \tparam E The type of the error value + * \param err The universal reference to the error value + * \return A PartialResult with the universal reference to the error value + * \example + * \code + * E err = E("Error"); + * // Move the err to the PartialResult + * return ResultErr(std::move(err)); + * \endcode + */ +template +inline PartialResult ResultErr(E&& err) { + return PartialResult{std::forward(err)}; +} + +/*! + * \brief An always-move Result type similar to Rust's Result, representing either success (Ok) or + * failure (Err). It always uses move semantics for the success and error values. + * \tparam T The type of the success value + * \tparam E The type of the error value + * + * \note The Ok and Err constructor, and all methods of this class (except for ValueRef and ErrRef) + * accept only rvalue references as parameters for performance reasons. You should use std::move to + * convert a Result to an rvalue reference before invoking these methods. Examples for move + * semantics are shown below. + * + * \example Construct a success result with a rvalue reference + * \code + * T value; + * return Result::Ok(std::move(value)); + * \endcode + * \example Construct a error result with a rvalue reference of std::runtime_error + * \code + * std::runtime_error error_msg = std::runtime_error("Error"); + * return Result::Err(std::move(error_msg)); + * \endcode + * \example Construct a error result with a std::runtime_error object constructed with a string + * \code + * std::string error_msg = "Error"; + * return Result::Err(std::move(error_msg)); + * \endcode + * \example Unwrap the rvalue reference of the result + * \code + * Result result = func(); + * if (result.IsOk()) { + * T result_val = std::move(result).Unwrap(); + * } else { + * std::runtime_error error_msg = std::move(result).UnwrapErr(); + * } + * \endcode + */ +template +class Result { + private: + static_assert(!std::is_same_v, "T and E cannot be the same type"); + + public: + /*! \brief Default constructor is deleted to avoid accidental use */ + Result() = delete; + + /*! \brief Construct from Result::Ok */ + template >>> + Result(PartialResult&& partial_result) + : data_(std::in_place_type, std::forward(partial_result.value)) {} + + /*! \brief Construct from Result::Err */ + template >>> + Result(PartialResult&& partial_result) + : data_(std::in_place_type, std::forward(partial_result.value)) {} + + /*! \brief Check if Result contains success value */ + bool IsOk() const { return std::holds_alternative(data_); } + + /*! \brief Check if Result contains error */ + bool IsErr() const { return std::holds_alternative(data_); } + + /*! \brief Get the success value. It assumes (or checks if in debug mode) the result is ok. */ + T Unwrap() && { + XGRAMMAR_DCHECK(IsOk()) << "Called Unwrap() on an Err value"; + return std::get(std::move(data_)); + } + + /*! \brief Get the error value. It assumes (or checks if in debug mode) the result is an error. */ + E UnwrapErr() && { + XGRAMMAR_DCHECK(IsErr()) << "Called UnwrapErr() on an Ok value"; + return std::get(std::move(data_)); + } + + /*! \brief Get the success value if present, otherwise return the provided default */ + T UnwrapOr(T default_value) && { + return IsOk() ? std::get(std::move(data_)) : std::move(default_value); + } + + /*! \brief Map success value to new type using provided function */ + template >> + Result Map(F&& f) && { + if (IsOk()) { + return ResultOk(f(std::get(std::move(data_)))); + } + return ResultErr(std::get(std::move(data_))); + } + + /*! \brief Map error value to new type using provided function */ + template >> + Result MapErr(F&& f) && { + if (IsErr()) { + return ResultErr(f(std::get(std::move(data_)))); + } + return ResultOk(std::get(std::move(data_))); + } + + /*! + * \brief Convert a Result to a Result. U should be convertible to T, and V should be + * convertible to E. + */ + template + static Result Convert(Result&& result) { + if (result.IsOk()) { + return ResultOk(std::move(result).Unwrap()); + } + return ResultErr(std::move(result).UnwrapErr()); + } + + /*! \brief Get a std::variant from the result. */ + std::variant ToVariant() && { return std::move(data_); } + + /*! + * \brief Get a reference to the success value. It assumes (or checks if in debug mode) the + * result is ok. + */ + T& ValueRef() & { + XGRAMMAR_DCHECK(IsOk()) << "Called ValueRef() on an Err value"; + return std::get(data_); + } + + /*! + * \brief Get a reference to the error value. It assumes (or checks if in debug mode) the + * result is an error. + */ + E& ErrRef() & { + XGRAMMAR_DCHECK(IsErr()) << "Called ErrRef() on an Ok value"; + return std::get(data_); + } + + private: + // in-place construct T in variant + template + explicit Result(std::in_place_type_t, Args&&... args) + : data_(std::in_place_type, std::forward(args)...) {} + + // in-place construct E in variant + template + explicit Result(std::in_place_type_t, Args&&... args) + : data_(std::in_place_type, std::forward(args)...) {} + + std::variant data_; +}; + +/****************** Misc ******************/ + +// Sometimes GCC fails to detect some branches will not return, such as when we use LOG(FATAL) +// to raise an error. This macro manually mark them as unreachable to avoid warnings. +#ifdef __GNUC__ +#define XGRAMMAR_UNREACHABLE() __builtin_unreachable() +#else +#define XGRAMMAR_UNREACHABLE() +#endif + +/*! + * \brief An error class that contains a type. The type can be an enum. + */ +template +class TypedError : public std::runtime_error { + public: + explicit TypedError(T type, const std::string& msg) : std::runtime_error(msg), type_(type) {} + const T& Type() const noexcept { return type_; } + + private: + T type_; +}; + +/** + * \brief Helper function to compare two objects by their members. + */ +template +constexpr bool EqualByMembers(const T& lhs, const T& rhs) noexcept { + return std::tie(lhs.*Ms...) == std::tie(rhs.*Ms...); +} + +/** + * \brief Define == and != operator for a struct by its members. + * \param Type The type of the struct. Must be under namespace xgrammar. + * \param ... The member pointers of the struct. + * \example + * \code + * struct Type { + * int member1; + * std::string member2; + * double member3; + * + * XGRAMMAR_EQUAL_BY_MEMBERS(Type, &Type::member1, &Type::member2, &Type::member3); + * }; + * \endcode + */ +#define XGRAMMAR_EQUAL_BY_MEMBERS(Type, ...) \ + friend bool operator==(const Type& lhs, const Type& rhs) noexcept { \ + return EqualByMembers(lhs, rhs); \ + } \ + friend bool operator!=(const Type& lhs, const Type& rhs) noexcept { return !(lhs == rhs); } + +/*! + * \brief Empty specialization of XGRAMMAR_EQUAL_BY_MEMBERS. + */ +#define XGRAMMAR_EQUAL_BY_MEMBERS_EMPTY(Type) \ + friend bool operator==(const Type& lhs, const Type& rhs) noexcept { return true; } \ + friend bool operator!=(const Type& lhs, const Type& rhs) noexcept { return false; } + +/*! + * \brief Throw an error from a variant of multiple error types. + * \param error_variant The variant of multiple error types. + * \tparam Args The types of the error types. Each type should inherit from std::runtime_error. + */ +template +[[noreturn]] void ThrowVariantError(const std::variant& error_variant) { + std::visit([](const auto& e) { throw e; }, error_variant); + XGRAMMAR_UNREACHABLE(); +} + +/*! + * \brief Get the message from a variant of multiple error types. + * \param error_variant The variant of multiple error types. + * \return The message from the error variant. + * \tparam Args The types of the error types. Each type should inherit from std::runtime_error. + */ +template +std::string GetMessageFromVariantError(const std::variant& error_variant) { + return std::visit([](const auto& e) { return e.what(); }, error_variant); +} + +/*! + * \brief Get the type name from a variant of XGrammarError types (each has GetType()). + * \param error_variant The variant of multiple error types. + * \return The type string from the error variant (e.g. "DeserializeVersionError"). + * \tparam Args The types of the error types. Each type should have GetType() const. + */ +template +std::string GetTypeFromVariantError(const std::variant& error_variant) { + return std::visit([](const auto& e) { return e.GetType(); }, error_variant); +} + +} // namespace xgrammar + +#endif // XGRAMMAR_SUPPORT_UTILS_H_ diff --git a/zig/vendor/xgrammar/cpp/testing.cc b/zig/vendor/xgrammar/cpp/testing.cc new file mode 100644 index 0000000000..b7b5bdf156 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/testing.cc @@ -0,0 +1,63 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/testing.cc + */ +#include "testing.h" + +#include + +#include +#include +#include +#include +#include + +#include "grammar_impl.h" +#include "grammar_parser.h" +#include "support/encoding.h" + +namespace xgrammar { + +std::string PrintTokenByIds( + const std::vector& token_ids, const TokenizerInfo& tokenizer_info, int max_print_num +) { + std::stringstream ss; + const auto& sorted_decoded_vocab = tokenizer_info.GetDecodedVocab(); + ss << "["; + int print_num = std::min(static_cast(token_ids.size()), max_print_num); + for (int i = 0; i < print_num; ++i) { + ss << "#" << token_ids[i] << " <" << EscapeString(sorted_decoded_vocab[token_ids[i]]) << ">"; + if (i < print_num - 1) { + ss << ", "; + } + } + if (static_cast(token_ids.size()) > max_print_num) { + ss << ", ..."; + } + ss << "]"; + return ss.str(); +} + +Grammar _EBNFToGrammarNoNormalization( + const std::string& ebnf_string, const std::string& root_rule_name +) { + return ParseEBNF(ebnf_string, root_rule_name); +} + +std::string _PrintGrammarFSMs(const Grammar& grammar) { + XGRAMMAR_CHECK(static_cast(grammar->per_rule_fsms.size()) == grammar->NumRules()) + << "The grammar has no per-rule FSMs; build them first"; + std::string result; + for (int i = 0; i < grammar->NumRules(); i++) { + result += "Rule " + std::to_string(i) + ": " + grammar->GetRule(i).name + ", FSM: "; + if (grammar->per_rule_fsms[i].has_value()) { + result += grammar->per_rule_fsms[i]->GetFsm().ToString(); + } else { + result += "None"; + } + result += "\n"; + } + return result; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/testing.h b/zig/vendor/xgrammar/cpp/testing.h new file mode 100644 index 0000000000..71b3c6331a --- /dev/null +++ b/zig/vendor/xgrammar/cpp/testing.h @@ -0,0 +1,30 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/testing.h + * \brief The header testing utilities. + */ +#ifndef XGRAMMAR_TESTING_H_ +#define XGRAMMAR_TESTING_H_ + +#include +#include + +#include +#include +#include + +namespace xgrammar { + +std::string PrintTokenByIds( + const std::vector& token_ids, const TokenizerInfo& tokenizer_info, int max_print_num +); + +Grammar _EBNFToGrammarNoNormalization( + const std::string& ebnf_string, const std::string& root_rule_name +); + +std::string _PrintGrammarFSMs(const Grammar& grammar); + +} // namespace xgrammar + +#endif // XGRAMMAR_TESTING_H_ diff --git a/zig/vendor/xgrammar/cpp/tokenizer_info.cc b/zig/vendor/xgrammar/cpp/tokenizer_info.cc new file mode 100644 index 0000000000..b580d3c30f --- /dev/null +++ b/zig/vendor/xgrammar/cpp/tokenizer_info.cc @@ -0,0 +1,562 @@ +/*! + * Copyright (c) 2023 by Contributors + * \file xgrammar/tokenizer_info.cc + */ + +#include "xgrammar/tokenizer_info.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "support/encoding.h" +#include "support/json_parse.h" +#include "support/json_serializer.h" +#include "support/logging.h" +#include "tokenizer_info_impl.h" +#include "xgrammar/exception.h" + +namespace xgrammar { + +/************* Token decoders: ByteFallback and ByteLevel *************/ + +class TokenDecoder { + public: + /*! + * \brief Post-process a raw token to the actual token with the given post-processing method. + */ + static std::string DecodeToken(const std::string& token, VocabType vocab_type) { + // TODO(yixin): Avoid allocating new string in decoder calls + if (vocab_type == VocabType::BYTE_FALLBACK) { + return SpaceReplacerDecoder(ByteFallbackDecoder(token)); + } else if (vocab_type == VocabType::BYTE_LEVEL) { + return ByteLevelDecoder(token); + } else { + return token; + } + } + + private: + /*! \brief ByteFallback decoder: transform tokens like <0x1B> to hex char byte 1B */ + static std::string ByteFallbackDecoder(const std::string& token) { + if (token.length() == 6 && token.substr(0, 3) == "<0x" && token.back() == '>') { + int byte = 0; + for (int i = 0; i < 2; ++i) { + byte *= 16; + byte += token[3 + i] >= '0' && token[3 + i] <= '9' ? token[3 + i] - '0' + : token[3 + i] - 'A' + 10; + } + XGRAMMAR_CHECK(byte >= 0 && byte < 256); + return std::string(/*n=*/1, static_cast(byte)); + } + return token; + } + + /*! \brief SpaceReplacer decoder: transform "\u2581" back to space */ + static std::string SpaceReplacerDecoder(const std::string& token) { + // \u2581 is the unicode for "lower one eighth block" + // UTF8 encoding for \u2581 is 0xE2 0x96 0x81 + std::string result; + for (int i = 0; i < static_cast(token.size()); ++i) { + if (i + 2 < static_cast(token.size()) && token[i] == char(0xE2) && + token[i + 1] == char(0x96) && token[i + 2] == char(0x81)) { + result += ' '; + i += 2; + } else { + result += token[i]; + } + } + return result; + } + + /*! + * \brief ByteLevel decoder: inverses the bytes-to-unicode transformation in the encoding + * process as in + * https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59 + */ + static std::string ByteLevelDecoder(const std::string& token) { + // The inverse map of bytes_to_unicode. -1 means there is no mapping to this unicode. + static const std::array char_to_byte_map = { + // clang-format off + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, + 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, + 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, + 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, + 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, -1, + 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, + 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, + 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, + 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, + 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, + 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128, + 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, + 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 173 + // clang-format on + }; + + auto unicode_codepoints = ParseUTF8(token.c_str(), false); + if (unicode_codepoints.size() == 1 && unicode_codepoints[0] == kInvalidUTF8) { + return token; + } + + std::string decoded; + decoded.reserve(unicode_codepoints.size()); + + for (auto unicode_codepoint : unicode_codepoints) { + XGRAMMAR_CHECK(unicode_codepoint >= 0); + if (unicode_codepoint >= static_cast(char_to_byte_map.size()) || + char_to_byte_map[unicode_codepoint] == -1) { + // If there is no mapping, return the original token + return token; + } + decoded += static_cast(char_to_byte_map[unicode_codepoint]); + } + return decoded; + } +}; + +/************* Metadata detection from huggingface tokenizer.json *************/ + +class HFTokenizerAnalyzer { + public: + /*! + * \brief Detect the vocabulary type from tokenizer.json. + * \details Find {"type": "ByteFallback"} or {"type": "ByteLevel"} in "decoder" field of the + * tokenizer. + */ + static VocabType DetectVocabType(const picojson::object& hf_tokenizer_obj) { +#define CHECK_AND_WARNING(condition, message) \ + if (!(condition)) { \ + XGRAMMAR_LOG(WARNING) << "Vocab type detection failed: (" #condition \ + << ") is false: " << (message) << " Using RAW VocabType by default."; \ + return VocabType::RAW; \ + } + + CHECK_AND_WARNING( + hf_tokenizer_obj.count("decoder") && hf_tokenizer_obj.at("decoder").is(), + "Decoder field is not found in tokenizer.json." + ); + + auto decoder_obj = hf_tokenizer_obj.at("decoder").get(); + CHECK_AND_WARNING( + decoder_obj.count("type") && decoder_obj.at("type").is(), + "Type field is not found in decoder field" + ); + auto type = decoder_obj.at("type").get(); + + std::vector decoders; + if (type == "Sequence") { + CHECK_AND_WARNING( + decoder_obj.count("decoders") && decoder_obj.at("decoders").is(), + "Decoders field is not found in a Sequence decoder" + ); + decoders = decoder_obj.at("decoders").get(); + } else { + decoders.emplace_back(hf_tokenizer_obj.at("decoder")); + } + + for (const auto& decoder : decoders) { + CHECK_AND_WARNING(decoder.is(), "Decoder is not an object"); + auto decoder_obj = decoder.get(); + CHECK_AND_WARNING( + decoder_obj.count("type") && decoder_obj.at("type").is(), + "Type field is not found in decoder field" + ); + auto type = decoder_obj.at("type").get(); + if (type == "ByteLevel") { + return VocabType::BYTE_LEVEL; + } else if (type == "ByteFallback") { + return VocabType::BYTE_FALLBACK; + } + } + + // If neither byte_level nor byte_fallback decoder is detected, return RAW. + return VocabType::RAW; + +#undef CHECK_AND_WARNING + } + + static bool DetectPrependNormalizer(const picojson::object& hf_tokenizer_obj) { + if (!hf_tokenizer_obj.count("normalizer") || + !hf_tokenizer_obj.at("normalizer").is()) { + return false; + } + + const picojson::value& normalizer_value = hf_tokenizer_obj.at("normalizer"); + if (!normalizer_value.is()) { + return false; + } + const picojson::object& normalizer_obj = normalizer_value.get(); + if (!normalizer_obj.count("type") || !normalizer_obj.at("type").is()) { + return false; + } + auto type = normalizer_obj.at("type").get(); + + std::vector normalizers; + if (type == "Sequence") { + if (!normalizer_obj.count("normalizers") || + !normalizer_obj.at("normalizers").is()) { + return false; + } + normalizers = normalizer_obj.at("normalizers").get(); + } else { + normalizers.emplace_back(normalizer_value); + } + + for (const auto& normalizer : normalizers) { + if (!normalizer.is()) { + continue; + } + auto normalizer_obj = normalizer.get(); + if (!normalizer_obj.count("type") || !normalizer_obj.at("type").is()) { + continue; + } + auto type = normalizer_obj.at("type").get(); + if (type == "Prepend" && normalizer_obj.count("prepend") && + normalizer_obj.at("prepend").is() && + normalizer_obj.at("prepend").get() == "▁") { + return true; + } + } + return false; + } + + static bool DetectMetaspacePreTokenizer(const picojson::object& hf_tokenizer_obj) { + if (!hf_tokenizer_obj.count("pre_tokenizer") || + !hf_tokenizer_obj.at("pre_tokenizer").is()) { + return false; + } + auto pre_tokenizer_obj = hf_tokenizer_obj.at("pre_tokenizer").get(); + if (!pre_tokenizer_obj.count("type") || !pre_tokenizer_obj.at("type").is()) { + return false; + } + auto type = pre_tokenizer_obj.at("type").get(); + if (!pre_tokenizer_obj.count("prepend_scheme") || + !pre_tokenizer_obj.at("prepend_scheme").is()) { + return false; + } + auto prepend_scheme = pre_tokenizer_obj.at("prepend_scheme").get(); + return type == "Metaspace" && (prepend_scheme == "always" || prepend_scheme == "first"); + } + + /*! + * \brief Detect whether add prefix space from tokenizer.json. + * \details Find {"type": "Prepend", "prepend": "▁"} in "normalizer" field of the tokenizer, or + * "pre_tokenizer": {"type": "Metaspace", "prepend_scheme": "always" | "first"} in the tokenizer. + */ + static bool DetectAddPrefixSpace(const picojson::object& hf_tokenizer_obj) { + return DetectPrependNormalizer(hf_tokenizer_obj) || + DetectMetaspacePreTokenizer(hf_tokenizer_obj); + } +}; + +/************* TokenizerInfo::Impl *************/ + +bool TokenizerInfo::Impl::IsSpecialToken(const std::string& token) { return token == ""; } + +TokenizerInfo::Impl::Impl( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space +) + : vocab_type_(vocab_type), + vocab_size_(vocab_size.value_or(encoded_vocab.size())), + add_prefix_space_(add_prefix_space) { + // vocab_size only ever pads the encoded vocab with special ids; a value below the number of + // real tokens would leave token ids without a slot and corrupt the id -> index table below. + XGRAMMAR_CHECK(vocab_size_ >= static_cast(encoded_vocab.size())) + << "vocab_size (" << vocab_size_ << ") must be at least the number of tokens in the vocab (" + << encoded_vocab.size() << ")."; + decoded_vocab_.reserve(encoded_vocab.size()); + sorted_decoded_vocab_.reserve(encoded_vocab.size()); + for (int i = 0; i < static_cast(encoded_vocab.size()); ++i) { + const std::string& token = TokenDecoder::DecodeToken(encoded_vocab[i], vocab_type_); + decoded_vocab_.push_back(token); + if ((!stop_token_ids && DETECTION_STOP_TOKENS.count(token)) || + (stop_token_ids && + std::find(stop_token_ids->begin(), stop_token_ids->end(), i) != stop_token_ids->end())) { + stop_token_ids_.push_back(i); + } else if (IsSpecialToken(token)) { + special_token_ids_.push_back(i); + } else { + sorted_decoded_vocab_.push_back({i, token}); + } + } + for (int i = encoded_vocab.size(); i < vocab_size_; ++i) { + special_token_ids_.push_back(i); + } + + auto f_compare_token = [](const std::pair& a, + const std::pair& b) { + return a.second < b.second; + }; + std::sort(sorted_decoded_vocab_.begin(), sorted_decoded_vocab_.end(), f_compare_token); + + BuildTokenIdToSortedVocabIndex(); + + // The value means: the subtree is [i, trie_subtree_nodes_range[i]). + trie_subtree_nodes_range_.resize(sorted_decoded_vocab_.size(), 0); + std::stack> prefix_stack; + for (size_t i = 0; i < sorted_decoded_vocab_.size(); ++i) { + const auto& token = sorted_decoded_vocab_[i].second; + while ((!prefix_stack.empty()) && (token.find(prefix_stack.top().first) == std::string::npos)) { + const auto& top_pair = prefix_stack.top(); + trie_subtree_nodes_range_[top_pair.second] = i; + prefix_stack.pop(); + } + prefix_stack.push({token, i}); + } + while (!prefix_stack.empty()) { + const auto& top_pair = prefix_stack.top(); + trie_subtree_nodes_range_[top_pair.second] = sorted_decoded_vocab_.size(); + prefix_stack.pop(); + } + BuildTokenCharData(); +} + +std::optional TokenizerInfo::Impl::Validate() const { + const int64_t num_tokens = decoded_vocab_.size(); + if (vocab_size_ < num_tokens) { + return "vocab_size " + std::to_string(vocab_size_) + " is smaller than the number of tokens " + + std::to_string(num_tokens); + } + for (const auto& [token_id, token] : sorted_decoded_vocab_) { + if (token_id < 0 || token_id >= num_tokens) { + return "sorted_decoded_vocab contains token id " + std::to_string(token_id) + " out of range"; + } + } + if (trie_subtree_nodes_range_.size() != sorted_decoded_vocab_.size()) { + return "trie_subtree_nodes_range must have one entry per sorted token"; + } + auto id_ok = [&](int32_t token_id) { return token_id >= 0 && token_id < vocab_size_; }; + if (!std::all_of(stop_token_ids_.begin(), stop_token_ids_.end(), id_ok) || + !std::all_of(special_token_ids_.begin(), special_token_ids_.end(), id_ok)) { + return "stop_token_ids or special_token_ids contains a token id out of range"; + } + return std::nullopt; +} + +void TokenizerInfo::Impl::BuildTokenIdToSortedVocabIndex() { + token_id_to_sorted_vocab_index_.assign(vocab_size_, -1); + for (int32_t i = 0; i < static_cast(sorted_decoded_vocab_.size()); ++i) { + token_id_to_sorted_vocab_index_[sorted_decoded_vocab_[i].first] = i; + } +} + +void TokenizerInfo::Impl::BuildTokenCharData() { + token_char_counts_.assign(sorted_decoded_vocab_.size(), 0); + int32_t max_chars = 0; + for (int32_t index = 0; index < static_cast(sorted_decoded_vocab_.size()); ++index) { + int32_t count = 0; + for (uint8_t byte : sorted_decoded_vocab_[index].second) { + count += (byte & 0xC0) != 0x80; + } + token_char_counts_[index] = count; + max_chars = std::max(max_chars, count); + } + max_token_chars_ = max_chars; +} + +const std::vector& TokenizerInfo::Impl::GetTokenCharCounts() const { + return token_char_counts_; +} + +int32_t TokenizerInfo::Impl::GetMaxTokenChars() const { return max_token_chars_; } + +std::string TokenizerInfo::Impl::DumpMetadata() const { + return DumpMetadataValue().serialize(false); +} + +picojson::value TokenizerInfo::Impl::DumpMetadataValue() const { + picojson::object obj; + obj["vocab_type"] = picojson::value(static_cast(vocab_type_)); + obj["vocab_size"] = picojson::value(static_cast(vocab_size_)); + obj["add_prefix_space"] = picojson::value(add_prefix_space_); + picojson::array stop_token_ids_array; + for (auto id : stop_token_ids_) { + stop_token_ids_array.push_back(picojson::value(static_cast(id))); + } + obj["stop_token_ids"] = picojson::value(std::move(stop_token_ids_array)); + + return picojson::value(std::move(obj)); +} + +std::optional TokenizerInfo::Impl::CheckMetadataMatch( + const picojson::value& metadata +) const { + if (!metadata.is()) { + return std::runtime_error("Expect an object"); + } + const auto& object = metadata.get(); + if (object.find("vocab_type") == object.end()) { + return std::runtime_error("Missing 'vocab_type' in metadata"); + } + auto vocab_type = object.at("vocab_type").get(); + if (vocab_type != static_cast(vocab_type_)) { + return std::runtime_error( + "Vocab type mismatch: " + std::to_string(vocab_type) + + " != " + std::to_string(static_cast(vocab_type_)) + ); + } + if (object.find("vocab_size") == object.end()) { + return std::runtime_error("Missing 'vocab_size' in metadata"); + } + auto vocab_size = object.at("vocab_size").get(); + if (vocab_size != vocab_size_) { + return std::runtime_error( + "Vocab size mismatch: " + std::to_string(vocab_size) + " != " + std::to_string(vocab_size_) + ); + } + if (object.find("add_prefix_space") == object.end()) { + return std::runtime_error("Missing 'add_prefix_space' in metadata"); + } + auto add_prefix_space = object.at("add_prefix_space").get(); + if (add_prefix_space != add_prefix_space_) { + return std::runtime_error( + "Add prefix space mismatch: " + std::to_string(add_prefix_space) + + " != " + std::to_string(add_prefix_space_) + ); + } + if (object.find("stop_token_ids") == object.end()) { + return std::runtime_error("Missing 'stop_token_ids' in metadata"); + } + auto stop_token_ids = object.at("stop_token_ids").get(); + std::vector stop_token_ids_vec; + stop_token_ids_vec.reserve(stop_token_ids.size()); + for (const auto& id : stop_token_ids) { + if (!id.is()) { + return std::runtime_error("Stop token id is not an integer"); + } + stop_token_ids_vec.push_back(static_cast(id.get())); + } + if (stop_token_ids_vec != stop_token_ids_) { + return std::runtime_error("Stop token ids mismatch"); + } + return std::nullopt; +} + +std::shared_ptr TokenizerInfo::Impl::FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata +) { + picojson::value v; + std::string err = ParseJSON(v, metadata); + XGRAMMAR_CHECK(err.empty()) << "Failed to parse metadata: " << err; + + const picojson::object& obj = v.get(); + + XGRAMMAR_CHECK(obj.count("vocab_type") && obj["vocab_type"].is()) + << "Missing or invalid 'vocab_type' in metadata"; + int vocab_type_int = static_cast(obj["vocab_type"].get()); + XGRAMMAR_CHECK(vocab_type_int == 0 || vocab_type_int == 1 || vocab_type_int == 2) + << "Invalid vocab_type in metadata: " << vocab_type_int; + VocabType vocab_type = static_cast(vocab_type_int); + + XGRAMMAR_CHECK(obj.count("vocab_size") && obj["vocab_size"].is()) + << "Missing or invalid 'vocab_size' in metadata"; + int vocab_size = static_cast(obj["vocab_size"].get()); + + XGRAMMAR_CHECK(obj.count("add_prefix_space") && obj["add_prefix_space"].is()) + << "Missing or invalid 'add_prefix_space' in metadata"; + bool add_prefix_space = obj["add_prefix_space"].get(); + + std::vector stop_token_ids; + XGRAMMAR_CHECK(obj.count("stop_token_ids") && obj["stop_token_ids"].is()) + << "Missing or invalid 'stop_token_ids' in metadata"; + for (const auto& id : obj["stop_token_ids"].get()) { + XGRAMMAR_CHECK(id.is()) << "Stop token id is not an integer"; + stop_token_ids.push_back(static_cast(id.get())); + } + return std::make_shared( + encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space + ); +} + +std::string TokenizerInfo::Impl::DetectMetadataFromHF(const std::string& backend_str) { + picojson::value v; + std::string err = ParseJSON(v, backend_str); + XGRAMMAR_CHECK(err.empty() && v.is()) << "Failed to parse JSON object: " << err; + const picojson::object& obj = v.get(); + VocabType vocab_type = HFTokenizerAnalyzer::DetectVocabType(obj); + bool add_prefix_space = HFTokenizerAnalyzer::DetectAddPrefixSpace(obj); + + // Serialize the metadata + picojson::object metadata_obj; + metadata_obj["vocab_type"] = picojson::value(static_cast(vocab_type)); + metadata_obj["add_prefix_space"] = picojson::value(add_prefix_space); + return picojson::value(metadata_obj).serialize(false); +} + +/************* TokenizerInfo *************/ + +TokenizerInfo::TokenizerInfo( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space +) + : pimpl_(std::make_shared( + encoded_vocab, vocab_type, vocab_size, stop_token_ids, add_prefix_space + )) {} + +int TokenizerInfo::GetVocabSize() const { return pimpl_->GetVocabSize(); } +VocabType TokenizerInfo::GetVocabType() const { return pimpl_->GetVocabType(); } +bool TokenizerInfo::GetAddPrefixSpace() const { return pimpl_->GetAddPrefixSpace(); } +const std::vector& TokenizerInfo::GetDecodedVocab() const { + return pimpl_->GetDecodedVocab(); +} +const std::vector& TokenizerInfo::GetStopTokenIds() const { + return pimpl_->GetStopTokenIds(); +} +const std::vector& TokenizerInfo::GetSpecialTokenIds() const { + return pimpl_->GetSpecialTokenIds(); +} +const std::vector>& TokenizerInfo::GetSortedDecodedVocab() const { + return pimpl_->GetSortedDecodedVocab(); +} + +const std::vector& TokenizerInfo::GetTrieSubtreeNodesRange() const { + return pimpl_->GetTrieSubtreeNodesRange(); +} + +std::string TokenizerInfo::DumpMetadata() const { return pimpl_->DumpMetadata(); } + +TokenizerInfo TokenizerInfo::FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata +) { + return TokenizerInfo(Impl::FromVocabAndMetadata(encoded_vocab, metadata)); +} + +std::string TokenizerInfo::DetectMetadataFromHF(const std::string& backend_str) { + return Impl::DetectMetadataFromHF(backend_str); +} + +std::string TokenizerInfo::SerializeJSON() const { return AutoSerializeJSON(*this, true); } + +std::variant TokenizerInfo::DeserializeJSON( + const std::string& json_string +) { + TokenizerInfo tokenizer_info{NullObj()}; + if (auto err = AutoDeserializeJSON(&tokenizer_info, json_string, true, "TokenizerInfo")) { + return err.value(); + } + // Derived per-token arrays are not serialized; rebuild them like the constructor does. + tokenizer_info->BuildTokenIdToSortedVocabIndex(); + tokenizer_info->BuildTokenCharData(); + return tokenizer_info; +} + +} // namespace xgrammar diff --git a/zig/vendor/xgrammar/cpp/tokenizer_info_impl.h b/zig/vendor/xgrammar/cpp/tokenizer_info_impl.h new file mode 100644 index 0000000000..863254aff4 --- /dev/null +++ b/zig/vendor/xgrammar/cpp/tokenizer_info_impl.h @@ -0,0 +1,143 @@ +#ifndef XGRAMMAR_TOKENIZER_INFO_IMPL_H_ +#define XGRAMMAR_TOKENIZER_INFO_IMPL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "support/reflection.h" +#include "xgrammar/tokenizer_info.h" + +namespace xgrammar { + +class TokenizerInfo::Impl { + public: + explicit Impl() = default; + + Impl( + const std::vector& encoded_vocab, + VocabType vocab_type, + std::optional vocab_size, + std::optional> stop_token_ids, + bool add_prefix_space + ); + + VocabType GetVocabType() const { return vocab_type_; } + bool GetAddPrefixSpace() const { return add_prefix_space_; } + int GetVocabSize() const { return vocab_size_; } + const std::vector& GetDecodedVocab() { return decoded_vocab_; } + const std::vector& GetStopTokenIds() const { return stop_token_ids_; } + const std::vector& GetSpecialTokenIds() const { return special_token_ids_; } + const std::vector>& GetSortedDecodedVocab() const { + return sorted_decoded_vocab_; + } + const std::vector& GetTrieSubtreeNodesRange() const { return trie_subtree_nodes_range_; } + const std::vector& GetTokenIdToSortedVocabIndex() const { + return token_id_to_sorted_vocab_index_; + } + const std::vector& GetTokenCharCounts() const; + int32_t GetMaxTokenChars() const; + void BuildTokenIdToSortedVocabIndex(); + void BuildTokenCharData(); + + /*! + * \brief Check that the token ids and per-token arrays are consistent with the vocabulary. Used + * after deserialization, where the fields are restored verbatim. + * \return An error message if the tokenizer info is malformed. + */ + std::optional Validate() const; + + std::string DumpMetadata() const; + picojson::value DumpMetadataValue() const; + + static std::shared_ptr FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata + ); + + std::optional CheckMetadataMatch(const picojson::value& metadata) const; + + static std::string DetectMetadataFromHF(const std::string& backend_str); + + bool operator==(const Impl& other) const; + + private: + static bool IsSpecialToken(const std::string& decoded_token); + + /*! \brief The vocabulary type. */ + VocabType vocab_type_; + /*! \brief The size of the vocabulary. */ + int vocab_size_; + /*! \brief Whether to add prefix space. */ + bool add_prefix_space_; + + /*! \brief The vocabulary. Special tokens are included. */ + std::vector decoded_vocab_; + /*! \brief All (id, token) pairs sorted in lexicographic order. This sorting is done to + * maximize prefix reuse during matching. Special tokens and stop tokens are not included. */ + std::vector> sorted_decoded_vocab_; + /*! \brief A pesudo-trie. trie_subtree_nodes_range[i] stores how many nodes there are in the + * subtree. */ + std::vector trie_subtree_nodes_range_; + /*! \brief The stop tokens. When the GrammarMatcher can reach the end of the grammar, + * stop tokens can be accepted. */ + std::vector stop_token_ids_; + /*! \brief The special tokens. These tokens are ignored (masked out) during the grammar-guided + * generation. */ + std::vector special_token_ids_; + /*! \brief Reverse mapping: token_id -> index in sorted_decoded_vocab_. -1 if not present. */ + std::vector token_id_to_sorted_vocab_index_; + /*! \brief Unicode codepoint counts for the sorted decoded vocabulary. */ + int32_t max_token_chars_ = 0; + std::vector token_char_counts_; + + /*! + * \brief The tokens used to detect stop tokens from the vocabulary. + * + * LLaMA2: + * LLaMA3: <|end_of_text|>, <|eot_id|> + * Phi-2: <|endoftext|> + * Gemma: , + * DeepSeek-V2: <|end▁of▁sentence|> + */ + inline static const std::unordered_set DETECTION_STOP_TOKENS = { + "", + "<|end_of_text|>", + "<|eot_id|>", + "<|endoftext|>", + "", + "<|eos|>", + "", + "<|end▁of▁sentence|>" + }; + + friend struct member_trait; +}; + +XGRAMMAR_MEMBER_TABLE( + TokenizerInfo::Impl, + "vocab_type", + &TokenizerInfo::Impl::vocab_type_, + "vocab_size", + &TokenizerInfo::Impl::vocab_size_, + "add_prefix_space", + &TokenizerInfo::Impl::add_prefix_space_, + "stop_token_ids", + &TokenizerInfo::Impl::stop_token_ids_, + "special_token_ids", + &TokenizerInfo::Impl::special_token_ids_, + "decoded_vocab", + &TokenizerInfo::Impl::decoded_vocab_, + "sorted_decoded_vocab", + &TokenizerInfo::Impl::sorted_decoded_vocab_, + "trie_subtree_nodes_range", + &TokenizerInfo::Impl::trie_subtree_nodes_range_ +); + +} // namespace xgrammar + +#endif // XGRAMMAR_TOKENIZER_INFO_IMPL_H_ diff --git a/zig/vendor/xgrammar/include/module.modulemap b/zig/vendor/xgrammar/include/module.modulemap new file mode 100644 index 0000000000..619525a81a --- /dev/null +++ b/zig/vendor/xgrammar/include/module.modulemap @@ -0,0 +1,4 @@ +module XGrammar { + umbrella header "xgrammar/xgrammar.h" + export * +} diff --git a/zig/vendor/xgrammar/include/xgrammar/compiler.h b/zig/vendor/xgrammar/include/xgrammar/compiler.h new file mode 100644 index 0000000000..2609e4b4d0 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/compiler.h @@ -0,0 +1,125 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/compiler.h + * \brief The header for the compiler. + */ + +#ifndef XGRAMMAR_COMPILER_H_ +#define XGRAMMAR_COMPILER_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +/*! + * \brief The compiled grammar of a GrammarMatcher. It contains the preprocessing results of the + * grammar and tokenizer. + */ +class CompiledGrammar { + public: + /*! \brief Get the associated grammar. */ + Grammar GetGrammar() const; + + /*! \brief Get the associated tokenizer info. */ + TokenizerInfo GetTokenizerInfo() const; + + /*! \brief Return the approximate memory usage of the grammar in bytes. */ + std::size_t MemorySizeBytes() const; + + /*! \brief Return the serialized JSON string of the compiled grammar. */ + std::string SerializeJSON() const; + + /*! \brief Deserialize a compiled grammar from a JSON string and tokenizer info. */ + static std::variant DeserializeJSON( + const std::string& json_string, const TokenizerInfo& tokenizer_info + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(CompiledGrammar); +}; + +/*! + * \brief A cache to get the compiled grammar for grammar or schema. This class avoids + * redundant preprocessing of the grammar or schema when constructing a CompiledGrammar. + * \note This class is associated with a vocabulary when constructed. The vocabulary is used to + * create every compiled grammar. If multiple toke tables are used to create init + * contexts, an instance of this class for each vocabulary should be created. + */ +class GrammarCompiler { + public: + /*! + * \brief Construct a GrammarCompiler with a vocabulary. This class will always + * create compiled grammars with this vocabulary. + * \param tokenizer_info The tokenizer info. + * \param max_threads The maximum number of threads to use for compiling grammars. + * \param cache_enabled Whether to enable the cache. + * \param max_memory_bytes The maximum memory usage in bytes. + */ + GrammarCompiler( + const TokenizerInfo& tokenizer_info, + int max_threads = 8, + bool cache_enabled = true, + int64_t max_memory_bytes = -1 // unlimited + ); + + /*! \brief Get the compiled grammar for a JSON schema string. */ + CompiledGrammar CompileJSONSchema( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool any_order = false + ); + + /*! \brief Get the compiled grammar for pure JSON. */ + CompiledGrammar CompileBuiltinJSONGrammar(); + + /*! + * \brief Get the compiled grammar for a Lark grammar string. + * \param lark_string The Lark grammar. The root rule must be named "start". + * \param named_grammars Grammar objects or Lark sources that can be referenced with `@name`. + */ + CompiledGrammar CompileLark( + const std::string& lark_string, const std::vector& named_grammars = {} + ); + + /*! \brief Get the compiled grammar for a grammar. */ + CompiledGrammar CompileGrammar(const Grammar& grammar); + + /*! \brief Get the compiled grammar for a grammar. */ + CompiledGrammar CompileGrammar( + const std::string& ebnf_str, const std::string& root_rule_name = "root" + ); + + /*! \brief Get the compiled grammar for a structural tag. */ + CompiledGrammar CompileStructuralTag(const std::string& structural_tag_json); + + /*! \brief Get the compiled grammar for a regex. */ + CompiledGrammar CompileRegex(const std::string& regex); + + /*! \brief Clear the internal cache of compiled grammars. */ + void ClearCache(); + + /*! \brief Return the approximate memory usage of the compiler in bytes. */ + int64_t GetCacheSizeBytes() const; + + /*! \brief Return the approximate memory usage of the compiler in bytes. -1 means unlimited. */ + int64_t CacheLimitBytes() const; + + XGRAMMAR_DEFINE_PIMPL_METHODS(GrammarCompiler); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_COMPILER_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/config.h b/zig/vendor/xgrammar/include/xgrammar/config.h new file mode 100644 index 0000000000..4e5de5e84b --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/config.h @@ -0,0 +1,35 @@ +/*! + * Copyright (c) 2025 by Contributors + * \file xgrammar/config.h + * \brief Global configuration for XGrammar. + */ + +#ifndef XGRAMMAR_CONFIG_H_ +#define XGRAMMAR_CONFIG_H_ + +#include + +namespace xgrammar { + +/*! + * \brief Set the maximum recursion depth for the grammar. + * \param max_recursion_depth The maximum recursion depth. + */ +void SetMaxRecursionDepth(int max_recursion_depth); + +/*! + * \brief Get the maximum recursion depth for the grammar. + * \return The maximum recursion depth. + */ +int GetMaxRecursionDepth(); + +/*! + * \brief Get the serialization version for the grammar. + * \return The serialization version. + * \note This is used to check the compatibility of the serialized grammar. + */ +std::string GetSerializationVersion(); + +} // namespace xgrammar + +#endif // XGRAMMAR_CONFIG_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/exception.h b/zig/vendor/xgrammar/include/xgrammar/exception.h new file mode 100644 index 0000000000..7efb3b0007 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/exception.h @@ -0,0 +1,80 @@ +#ifndef XGRAMMAR_EXCEPTION_H +#define XGRAMMAR_EXCEPTION_H + +#include +#include +#include + +namespace xgrammar { + +/************** Exception Definitions **************/ + +/*! + * \brief Exception thrown when the version in the serialized data does not follow the current + * serialization version. + */ + +struct XGrammarError : std::runtime_error { + XGrammarError(const std::string& message) : std::runtime_error(message) {} + virtual std::string GetType() const { return "XGrammarError"; } +}; + +struct DeserializeVersionError : XGrammarError { + DeserializeVersionError(const std::string& message) + : XGrammarError(std::string("Deserialize version error: ") + message) {} + std::string GetType() const override { return "DeserializeVersionError"; } +}; + +/*! + * \brief Exception thrown when the JSON is invalid. + */ +struct InvalidJSONError : XGrammarError { + InvalidJSONError(const std::string& message) + : XGrammarError(std::string("Invalid JSON error: ") + message) {} + std::string GetType() const override { return "InvalidJSONError"; } +}; + +/*! + * \brief Exception thrown when the serialized data does not follow the expected format. + */ +struct DeserializeFormatError : XGrammarError { + DeserializeFormatError(const std::string& message) + : XGrammarError(std::string("Deserialize format error: ") + message) {} + std::string GetType() const override { return "DeserializeFormatError"; } +}; + +/*! + * \brief Exception thrown when the JSON schema is invalid or not satisfiable. + */ +struct InvalidJSONSchemaError : XGrammarError { + InvalidJSONSchemaError(const std::string& message) + : XGrammarError(std::string("Invalid JSON schema error: ") + message) {} + std::string GetType() const override { return "InvalidJSONSchemaError"; } +}; + +/*! + * \brief Exception thrown when the structural tag is invalid. + */ +struct InvalidStructuralTagError : XGrammarError { + InvalidStructuralTagError(const std::string& message) + : XGrammarError(std::string("Invalid structural tag error: ") + message) {} + std::string GetType() const override { return "InvalidStructuralTagError"; } +}; + +/************** Union Exceptions **************/ + +/*! + * \brief Represents a serialization error. + */ +using SerializationError = + std::variant; + +/*! + * \brief Represents an error from the structural tag conversion. + */ +using StructuralTagError = + std::variant; + +} // namespace xgrammar + +#endif // XGRAMMAR_EXCEPTION_H diff --git a/zig/vendor/xgrammar/include/xgrammar/grammar.h b/zig/vendor/xgrammar/include/xgrammar/grammar.h new file mode 100644 index 0000000000..ffb54672a8 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/grammar.h @@ -0,0 +1,228 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/grammar.h + * \brief The header for the definition and construction of BNF grammar. + */ + +#ifndef XGRAMMAR_GRAMMAR_H_ +#define XGRAMMAR_GRAMMAR_H_ + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +struct StructuralTagItem { + std::string begin; + std::string schema; + std::string end; + + bool operator==(const StructuralTagItem& other) const { + return begin == other.begin && schema == other.schema && end == other.end; + } +}; + +struct NamedGrammar; + +/*! + * \brief This class stores the abstract syntax tree (AST) of the Backus-Naur Form (BNF) grammar. + * The BNF definition here is standard BNF, and the characters are represented using regex-style + * character classes (e.g. [a-z], [^a-z]). + * + * \details + * ### Rules + * The BNF grammar AST consists of a set of rules. Each rule contains a name and a definition, and + * corresponds to a production in the grammar. The definition of a rule is a GrammarExpr. Each rule + * has a rule_id for reference. + * + * ### GrammarExprs + * GrammarExpr is the definition of a rule or part of the definition of a rule. It can contain + * elements, empty string, reference to other GrammarExprs, or reference to other rules. Each + * GrammarExpr corresponds to a grammar_expr_id for reference. + * + * For example, in the following rule: rule ::= ("a" "b") | "c" + * ("a" "b"), "c", ("a" "b") | "c" are all GrammarExprs. + * + * #### Types of GrammarExprs + * Every GrammarExpr is represented by a type as well as a variable-length array containing its + * data. GrammarExpr has several types: + * - Byte string: a string of bytes (0~255). Supports UTF-8 strings. + * - Character class: a range of characters (each character is a unicode codepoint), e.g. [a-z], + * [ac-z]. Can be negated: [^a-z], [^ac-z]. Now only ascii chars is allowed in [], but this + * expression can accept/reject unicode chars. + * - Character class star: a star quantifier of a character class. e.g. [a-z]*, [^a-z]*. + * - EmptyStr: an empty string, i.e. "" + * - Rule reference: a reference to another rule + * - Sequence: a sequence of grammar_exprs, e.g. ("a" "b"). These grammar_exprs are concatenated + * together. + * - Choices: a choice of grammar_exprs, e.g. ("a" "b") | "c". Each grammar_expr can be matched. + * + * #### Storage of GrammarExprs + * Each type of GrammarExpr has a different data format. For the format of each type of GrammarExpr, + * see docs in Grammar::Impl::GrammarExprType. + * + * We store all GrammarExprs in csr_matrix style. That is, they are stored consecutively in one + * vector (data vector) and the starting position of each GrammarExpr is recorded in the indptr + * vector. + * + * \remark The character class star GrammarExpr is for the special support for elements like [a-z]* + * in the grammar. We add it to make the matching more efficient, as we can avoid recursion into + * rules when matching a sequence of characters. It should be used like: + * rule1 ::= ((element1 element2 rule2 ...) | ...) + * rule2 ::= character_class_star_grammar_expr(id_of_a_character_class_grammar_expr) + */ +class Grammar { + static const std::vector& EmptyNamedGrammars(); + + public: + /*! + * \brief Get the EBNF string of the grammar. + */ + std::string ToString() const; + + /*! + * \brief Construct a BNF grammar with a EBNF-formatted string. The grammar will be normalized + * (simplified) by default. + * \param ebnf_string The EBNF-formatted string. + * \param root_rule_name The name of the root rule. + */ + static Grammar FromEBNF( + const std::string& ebnf_string, const std::string& root_rule_name = "root" + ); + + /*! + * \brief Construct a BNF grammar from the json schema string. The schema string should be in the + * format of the schema of a JSON file. We will parse the schema and generate a BNF grammar. + * \param schema The schema string. + * \param indent The number of spaces for indentation. If set to std::nullopt, the output will be + * in one line. Default: 2. + * \param separators Two separators used in the schema: comma and colon. Examples: {",", ":"}, + * {", ", ": "}. If std::nullopt, the default separators will be used: {",", ": "} when the + * indent is not nullopt, and {", ", ": "} otherwise. This follows the convention in python + * json.dumps(). Default: std::nullopt. + * \param strict_mode Whether to use strict mode. In strict mode, the generated grammar will not + * allow properties and items that is not specified in the schema. This is equivalent to + * setting unevaluatedProperties and unevaluatedItems to false. + * + * This helps LLM to generate accurate output in the grammar-guided generation with JSON + * schema. Default: true. + */ + static Grammar FromJSONSchema( + const std::string& schema, + bool any_whitespace = true, + std::optional indent = std::nullopt, + std::optional> separators = std::nullopt, + bool strict_mode = true, + std::optional max_whitespace_cnt = std::nullopt, + bool print_converted_ebnf = false, + bool any_order = false + ); + + /*! + * \brief Construct a grammar from a regular expression string. + * \param regex The regular expression string. + * \param print_converted_ebnf This method will convert the regex to EBNF first. If this is true, + * the converted EBNF string will be printed. For debugging purpose. Default: false. + */ + static Grammar FromRegex(const std::string& regex, bool print_converted_ebnf = false); + + /*! + * \brief Construct a grammar from Lark syntax. + * \param lark_string The Lark grammar. The root rule must be named "start". + * \param tokenizer_info Optional tokenizer metadata used to resolve named special tokens. + * \param named_grammars Grammar objects or Lark sources that can be referenced with `@name`. + */ + static Grammar FromLark( + const std::string& lark_string, + const std::optional& tokenizer_info = std::nullopt, + const std::vector& named_grammars = EmptyNamedGrammars() + ); + + /*! + * \brief Construct a grammar from a structural tag string. + * \param structural_tag_json The structural tag string. + * \param tokenizer_info Optional tokenizer info for resolving string token references. + */ + static std::variant FromStructuralTag( + const std::string& structural_tag_json, + const std::optional& tokenizer_info = std::nullopt + ); + + /*! + * \brief Get the grammar of standard JSON format. We have built-in support for JSON. + * \return The grammar of standard JSON format. + */ + static Grammar BuiltinJSONGrammar(); + + /*! + * \brief Create a grammar that matches any of the grammars in the list. That is equivalent to + * using the `|` operator to concatenate the grammars in the list. + * \param grammars The grammars to create the union of. + * \returns The union of the grammars. + */ + static Grammar Union(const std::vector& grammars); + + /*! + * \brief Create a grammar that matches the concatenation of the grammars in the list. That is + * equivalent to using the `+` operator to concatenate the grammars in the list. + * \param grammars The grammars to create the concatenation of. + * \returns The concatenation of the grammars. + */ + static Grammar Concat(const std::vector& grammars); + + /*! + * \brief Print a BNF grammar. + * \param os The output stream. + * \param grammar The grammar to print. + * \return The output stream. + */ + friend std::ostream& operator<<(std::ostream& os, const Grammar& grammar); + + /*! + * \brief Return the serialized JSON string of the grammar. + * \return The serialized JSON string. + */ + std::string SerializeJSON() const; + + /*! + * \brief Deserialize a grammar from a JSON string. + * \param json_string The JSON string to deserialize. + * \return If the deserialization is successful, return the grammar. Otherwise, return a runtime + * error with the error message. + */ + static std::variant DeserializeJSON(const std::string& json_string); + + // Allow the builder to bind to this grammar's impl for in-place editing (FromMutableGrammar). + friend class GrammarBuilder; + + XGRAMMAR_DEFINE_PIMPL_METHODS(Grammar); +}; + +/*! + * \brief A grammar or Lark source that can be referenced by name from another Lark grammar. + */ +struct NamedGrammar { + /*! \brief The reference name, without the leading `@`. */ + std::string name; + + /*! \brief An existing grammar or a Lark source with its own `start` rule. */ + std::variant grammar; +}; + +inline const std::vector& Grammar::EmptyNamedGrammars() { + static const std::vector empty; + return empty; +} + +} // namespace xgrammar + +#endif // XGRAMMAR_GRAMMAR_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/matcher.h b/zig/vendor/xgrammar/include/xgrammar/matcher.h new file mode 100644 index 0000000000..d3a12b02f0 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/matcher.h @@ -0,0 +1,300 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/matcher.h + * \brief The header for the matcher. + */ + +#ifndef XGRAMMAR_MATCHER_H_ +#define XGRAMMAR_MATCHER_H_ + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace xgrammar { + +int32_t GetBitmaskSize(int vocab_size); + +DLDataType GetBitmaskDLType(); + +void _DebugGetMaskedTokensFromBitmask( + std::vector* rejected_tokens, const DLTensor& token_bitmask, int vocab_size, int index = 0 +); + +std::pair _IsSingleTokenBitmask(const DLTensor& bitmask, int vocab_size, int index); + +void ApplyTokenBitmaskInplaceCPU( + DLTensor* logits, + const DLTensor& bitmask, + int vocab_size = -1, + std::optional> indices = std::nullopt +); + +/*! + * \brief A stateful matcher to match tokens to the specified BNF grammar. This class is the core + * logic of the grammar-guided generation. + * + * \details This class implements the non-deterministic pushdown automaton (NPDA) matching algorithm + * to match characters to a BNF grammar. It keep track of the current state of the matching process + * by maintaining several stacks internally as possible paths in the NPDA. It also supports + * backtracking. + * + * It is particularly capable of finding the set of tokens that are acceptable for the next step + * and storing them in a bitmask. This aids in grammar-guided generation. + * + * \example + * \code + * Tokenizer tokenizer = ...; + * auto compiled_grammar = GrammarMatcher::CreateCompiledGrammar(grammar, + * tokenizer->PostProcessedVocab()); + * GrammarMatcher matcher(compiled_grammar, 10); + * matcher->AcceptToken(67); + * + * // Construct a DLTensor with shape (tokenizer.GetVocabSize() + 31) / 32, and dtype int32. + * DLTensor next_token_bitmask = ...; + * matcher->FillNextTokenBitmask(&next_token_bitmask); + * + * // Rollback is supported + * matcher->Rollback(1); + * \endcode + */ +class GrammarMatcher { + public: + /*! + * \brief Construct a GrammarMatcher from the preprocessing result of type + * CompiledGrammar. + * \param compiled_grammar The compiled grammar. It is obtained through + * CreateCompiledGrammar as a result of preprocessing the grammar and tokenizer. + * \param override_stop_tokens Optional stop token ids that replace those from the tokenizer. + * \param terminate_without_stop_token Whether to terminate when the root rule is complete. + * \param max_rollback_tokens Deprecated and unused. + * \param default_temperature The temperature used when no active rule specifies one. + */ + GrammarMatcher( + const CompiledGrammar& compiled_grammar, + std::optional> override_stop_tokens = std::nullopt, + bool terminate_without_stop_token = false, + int max_rollback_tokens = -1, + std::optional default_temperature = std::nullopt + ); + + /*! + * \brief Accept one token and update the state of the matcher. + * \param token_id The id of the token to accept. + * \return Whether the token is accepted. + * \note Termination state. + * When the end of the root rule is reached, the matcher can only accept the stop token. + * The matcher is terminated after accepting the stop token, i.e. no AcceptToken or + * FindNextTokenMask operations can be performed. The termination state can be canceled + * using Rollback(). + */ + bool AcceptToken(int32_t token_id, bool debug_print = false); + + /*! + * \brief Accept a string and update the state of the matcher. The whole string is considered + * as one step in rollback. It is used to complement the functionality of AcceptToken, and + * AcceptToken should always be used to accept tokens. + * \param input_str The string to be accepted. + * \param debug_print Whether to print information about the internal state of the matcher. + * \return Whether the string is accepted. + */ + bool AcceptString(const std::string& input_str, bool debug_print = false); + + /*! + * \brief Get the set of tokens that are acceptable for the next step and store them in a + * bitmask. + * \param next_token_bitmask The bitmask to store the result. The bitmask must be pre-allocated + * and with shape (GetBitmaskSize(),) and dtype int32. + * \return Whether the bitmask need to be applied (not all-true). + */ + bool FillNextTokenBitmask(DLTensor* next_token_bitmask, int index = 0, bool debug_print = false); + + /*! + * \brief Traverse a draft token tree and fill the token bitmask for each node. + * + * This function performs a DFS traversal of the speculative decoding tree and fills + * the token bitmask for each node based on grammar constraints. + * + * \param retrieve_next_token DLTensor where retrieve_next_token[i] gives the index of + * the child node of node i, or -1 if no child exists. + * \param retrieve_next_sibling DLTensor where retrieve_next_sibling[i] gives the index of + * the sibling node of node i, or -1 if no sibling exists. + * \param draft_tokens DLTensor of draft token ids at each node. + * \param token_bitmask DLTensor to store the bitmask (2D: num_nodes x bitmask_size). + * \param time_threshold Maximum allowed time in seconds for the DFS traversal. + * If the traversal exceeds this threshold, it returns false. + * A value <= 0 disables the timeout (default: -1.0). + * \param temperatures Optional DLTensor to store the effective temperature for each node + * (1D float32 with num_nodes elements). -1 represents no effective temperature; it is + * also written for nodes that were not visited or were rejected. + * \return true if the traversal completed successfully, false if it timed out. + */ + bool TraverseDraftTree( + const DLTensor* retrieve_next_token, + const DLTensor* retrieve_next_sibling, + const DLTensor* draft_tokens, + DLTensor* token_bitmask, + double time_threshold = -1.0, + DLTensor* temperatures = nullptr + ); + + /*! + * \brief Find the jump-forward string for jump-forward decoding. This is the longest string that + will be valid according to the current syntax. + * \note This method does not change the grammar state. + */ + std::string FindJumpForwardString(); + + /*! + * \brief Rollback the matcher to a previous state. + * \param num_tokens The number of tokens to rollback. It cannot exceed the current number of + * steps, nor can it exceed the specified maximum number of rollback tokens. + */ + void Rollback(int num_tokens = 1); + + /*! + * \brief Get the capture groups recorded so far, ordered by completion position. + * \param deduplicate The Earley parser explores parse hypotheses in parallel, so one + * occurrence of a captured rule may complete at several candidate end positions (e.g. a + * /[0-9]+/ body completes after every digit). If true (default), only the last (longest) + * completion of each occurrence is kept. Distinct occurrences — repeated matches of the same + * rule at different positions — are always kept. If false, the raw completion events are + * returned. + * \return A list of (capture_name, matched_bytes) pairs. Rules gain a capture name via the + * Lark attribute rule[capture] / rule[capture="name"] or the EBNF form rule[capture="name"] + * ::= ... The bytes are the input span that the rule matched. + * \note Captures are recorded when a rule is completed during AcceptToken / AcceptString, and + * are rolled back together with Rollback. Mask computation never records captures. Completions + * on parse paths that are later abandoned may still be recorded; for unambiguous grammars + * whose captured rules have unambiguous boundaries, the deduplicated result is exact. + */ + std::vector> GetCaptures(bool deduplicate = true) const; + + /*! + * \brief Check if the matcher has accepted the stop token and terminated. + * \sa AcceptToken + */ + bool IsTerminated() const; + + /*! + * \brief Check if the grammar's root rule has been fully matched by the input accepted so far. + * Unlike IsTerminated(), this does not require the stop token to have been accepted. + * \sa IsTerminated, AcceptToken + */ + bool IsCompleted() const; + + /*! \brief Reset the matcher to the initial state. */ + void Reset(); + + /*! + * \brief Fork the matcher. Returns a new GrammarMatcher with a deep copy of all state except + * compiled_grammar and tokenizer_info, which are shared with this matcher. + */ + GrammarMatcher Fork() const; + + /*! \brief Get the maximum number of rollback tokens allowed. */ + int GetMaxRollbackTokens() const; + + /*! \brief Get the effective sampling temperature for the next token. */ + std::optional GetTemperature() const; + + const std::vector& GetStopTokenIds() const; + + /*! \brief Print the internal state of the matcher. This is only used for debugging. The + * representation of the internal state is subject to change. + */ + std::string _DebugPrintInternalState() const; + + XGRAMMAR_DEFINE_PIMPL_METHODS(GrammarMatcher); +}; + +/*! + * \brief A batched version of GrammarMatcher for better efficiency. It supports batch processing + * of multiple GrammarMatcher objects in parallel. + * + * \details This class provides batched versions of the core methods of GrammarMatcher, including + * FillNextTokenBitmask, AcceptString, and AcceptToken. It utilizes multi-threading to process + * multiple GrammarMatcher objects simultaneously, significantly improving efficiency when dealing + * with a large number of matchers. + */ +class BatchGrammarMatcher { + public: + BatchGrammarMatcher(std::variant max_threads = "auto"); + + /*! + \brief A batched version of FillNextTokenBitmask for better efficiency. + \param matchers The array of GrammarMatcher objects. + \param next_token_bitmask The pre-allocated DLTensor to store the result bitmasks. + \param indices The optional array of indices to specify which matcher corresponds to which slice + of the bitmask tensor. If not provided, all matchers will write to the corresponding + indices(matchers[i] to next_token_bitmask[i]). + \param debug_print Whether to print debug information. Default is false. + */ + void BatchFillNextTokenBitmask( + std::vector* matchers, + DLTensor* next_token_bitmask, + const std::optional>& indices = std::nullopt, + bool debug_print = false + ); + + /*! + * \brief Fill the effective sampling temperature of each matcher into a tensor. + * \param matchers The array of GrammarMatcher objects. + * \param temperatures The pre-allocated 1D float32 CPU tensor to store the result. The entry of + * a matcher without an effective temperature is set to -1. + * \param indices The optional array of indices to specify which element each matcher writes to. + * If not provided, matchers[i] writes to temperatures[i]. + */ + static void BatchFillTemperature( + const std::vector& matchers, + DLTensor* temperatures, + const std::optional>& indices = std::nullopt + ); + + /*! + * \brief A batched version of AcceptString for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param input_strs The array of input strings to be accepted. + * \param debug_print Whether to print debug information. Default is false. + * \return A vector of bytes indicating whether each string is accepted. + */ + static std::vector BatchAcceptString( + std::vector* matchers, + const std::vector& input_strs, + bool debug_print = false + ); + + /*! + * \brief A batched version of AcceptToken for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param token_ids The array of token ids to be accepted. + * \param debug_print Whether to print debug information. Default is false. + * \return A vector of bytes indicating whether each token is accepted. + */ + static std::vector BatchAcceptToken( + std::vector* matchers, + const std::vector& token_ids, + bool debug_print = false + ); + + /*! + * \brief A batched version of Rollback for better efficiency. + * \param matchers The array of GrammarMatcher objects. + * \param num_tokens The array of the number of tokens to rollback for each matcher. + */ + static void BatchRollback( + std::vector* matchers, const std::vector& num_tokens + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(BatchGrammarMatcher); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_MATCHER_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/object.h b/zig/vendor/xgrammar/include/xgrammar/object.h new file mode 100644 index 0000000000..641b5bc526 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/object.h @@ -0,0 +1,51 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/object.h + * \brief Utilities for creating objects. + */ + +#ifndef XGRAMMAR_OBJECT_H_ +#define XGRAMMAR_OBJECT_H_ + +#include // IWYU pragma: keep +#include // IWYU pragma: keep + +namespace xgrammar { + +/*! + * \brief A tag type for creating a null object. + */ +struct NullObj {}; + +/*! + * \brief This macro defines the methods for the PImpl classes. + * \details Many classes in xgrammar are PImpl classes. PImpl classes only stores a shared pointer + * to the implementation. This allows reference-counter-based memory management and efficient + * object copy and passing. We always expose PImpl classes to Python to control over object sharing + * and memory management. Note simple and critical classes should not be defined as PImpl classes, + * but as normal classes for better efficiency. + */ +#define XGRAMMAR_DEFINE_PIMPL_METHODS(TypeName) \ + public: \ + class Impl; \ + /* Construct a null object. Note operating on a null object will fail. */ \ + explicit TypeName(NullObj) : pimpl_(nullptr) {} \ + /* Construct object with a shared pointer to impl. */ \ + explicit TypeName(std::shared_ptr pimpl) : pimpl_(std::move(pimpl)) {} \ + TypeName(const TypeName& other) = default; \ + TypeName(TypeName&& other) noexcept = default; \ + TypeName& operator=(const TypeName& other) = default; \ + TypeName& operator=(TypeName&& other) noexcept = default; \ + bool IsNull() const { return pimpl_ == nullptr; } \ + /* Access the impl pointer. Useful in implementation. */ \ + Impl* ImplPtr() { return pimpl_.get(); } \ + const Impl* ImplPtr() const { return pimpl_.get(); } \ + Impl* operator->() { return pimpl_.get(); } \ + const Impl* operator->() const { return pimpl_.get(); } \ + \ + private: \ + std::shared_ptr pimpl_ + +} // namespace xgrammar + +#endif // XGRAMMAR_OBJECT_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/tokenizer_info.h b/zig/vendor/xgrammar/include/xgrammar/tokenizer_info.h new file mode 100644 index 0000000000..6a8e67444b --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/tokenizer_info.h @@ -0,0 +1,86 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/tokenizer_info.h + * \brief The header for the tokenizer info. + */ + +#ifndef XGRAMMAR_TOKENIZER_INFO_H_ +#define XGRAMMAR_TOKENIZER_INFO_H_ + +#include + +#include +#include +#include +#include +#include + +#include "xgrammar/exception.h" + +namespace xgrammar { + +enum class VocabType : int { + RAW = 0, + BYTE_FALLBACK = 1, + BYTE_LEVEL = 2, +}; + +class TokenizerInfo { + public: + TokenizerInfo( + const std::vector& encoded_vocab, + VocabType vocab_type = VocabType::RAW, + std::optional vocab_size = std::nullopt, + std::optional> stop_token_ids = std::nullopt, + bool add_prefix_space = false + ); + + VocabType GetVocabType() const; + bool GetAddPrefixSpace() const; + int GetVocabSize() const; + const std::vector& GetDecodedVocab() const; + const std::vector& GetStopTokenIds() const; + const std::vector& GetSpecialTokenIds() const; + const std::vector>& GetSortedDecodedVocab() const; + const std::vector& GetTrieSubtreeNodesRange() const; + std::string DumpMetadata() const; + + /*! + * \brief Create a tokenizer info from a vocabulary and metadata. + * \param encoded_vocab The encoded vocabulary. + * \param metadata The metadata. + * \return The tokenizer info. + */ + static TokenizerInfo FromVocabAndMetadata( + const std::vector& encoded_vocab, const std::string& metadata + ); + + /*! + * \brief Detect the metadata from a Hugging Face backend string. + * \param backend_str The Hugging Face backend string. + * \return The metadata. + */ + static std::string DetectMetadataFromHF(const std::string& backend_str); + + /*! + * \brief Return the serialized JSON string of the tokenizer info. + * \return The serialized JSON string. + */ + std::string SerializeJSON() const; + + /*! + * \brief Deserialize a tokenizer info from a JSON string. + * \param json_string The JSON string to deserialize. + * \return If the deserialization is successful, return the tokenizer info. Otherwise, return a + * runtime error with the error message. + */ + static std::variant DeserializeJSON( + const std::string& json_string + ); + + XGRAMMAR_DEFINE_PIMPL_METHODS(TokenizerInfo); +}; + +} // namespace xgrammar + +#endif // XGRAMMAR_TOKENIZER_INFO_H_ diff --git a/zig/vendor/xgrammar/include/xgrammar/xgrammar.h b/zig/vendor/xgrammar/include/xgrammar/xgrammar.h new file mode 100644 index 0000000000..8513b3a9c5 --- /dev/null +++ b/zig/vendor/xgrammar/include/xgrammar/xgrammar.h @@ -0,0 +1,17 @@ +/*! + * Copyright (c) 2024 by Contributors + * \file xgrammar/xgrammar.h + * \brief The header for the support of grammar-guided generation. + */ + +#ifndef XGRAMMAR_XGRAMMAR_H_ +#define XGRAMMAR_XGRAMMAR_H_ + +#include +#include +#include +#include +#include +#include + +#endif // XGRAMMAR_XGRAMMAR_H_