inferencelayer 0.2.1

//! The **flat-table reference interpreter** — the executable spec the G2 WGSL kernel transcribes.
//!
//! [`JsonFsm`](super::json_fsm) reads the structured [`JsonSchemaTables`](super::json_schema); the
//! GPU cannot. This module re-expresses the identical automaton over ONLY the flat `u32` buffers of
//! a [`GpuSchema`] + [`GpuVocab`] — the same `words`/`bytes`/`off` arrays the WGSL uploads and the
//! same index arithmetic it will use — and the test below gates it byte-for-byte against `JsonFsm`
//! over the G1 corpus. Once green, the WGSL is a mechanical language port of [`flat_step_byte`]:
//! same 64-word state, same `(kind, phase)` dispatch, same buffer reads. No `wgpu`, no WGSL, pure
//! Rust — so it costs zero GPU build cycles to prove the encoding is semantically sufficient.
//!
//! The state is the identical `[u32; 64]` layout of [`FsmState`](super::json_fsm::FsmState): word 0
//! = depth, words 4.. = 15 frames of `[node, phase, a, b]`. Every phase constant and transition
//! below mirrors `json_fsm.rs` arm-for-arm; keep the two in lockstep (the gate enforces it).
//!
//! The `pub(crate)` interpreter fns have no non-test caller yet — the G2 WGSL harness is their
//! consumer — so the whole module allows dead code until that lands.
#![allow(dead_code)]

use super::harness::GpuSchema;

const HEADER_WORDS: usize = 4;
const FRAME_WORDS: usize = 4;
const MAX_STACK: usize = 15;
const NUMBER_BYTE_CAP: u32 = 24;
const NODE_WORDS: usize = 5;

// Node kind discriminants (mirror `json_schema::NodeKind`).
const K_OBJECT: u32 = 0;
const K_ARRAY: u32 = 1;
const K_STRING: u32 = 2;
const K_NUMBER: u32 = 3;
const K_INTEGER: u32 = 4;
const K_CHOICE: u32 = 5;

// Phase constants (mirror `json_fsm.rs`).
const OBJ_OPEN: u32 = 0;
const OBJ_FIRST: u32 = 1;
const OBJ_KEY: u32 = 2;
const OBJ_COLON: u32 = 3;
const OBJ_VALUE: u32 = 4;
const OBJ_AFTER: u32 = 5;
const OBJ_CLOSE_EMPTY: u32 = 6;
const ARR_OPEN: u32 = 10;
const ARR_FIRST: u32 = 11;
const ARR_AFTER: u32 = 12;
const ARR_ITEM: u32 = 13;
const STR_OPEN: u32 = 20;
const STR_BODY: u32 = 21;
const STR_C1: u32 = 22;
const STR_C2: u32 = 23;
const STR_C2_E0: u32 = 24;
const STR_C2_ED: u32 = 25;
const STR_C3: u32 = 26;
const STR_C3_F0: u32 = 27;
const STR_C3_F4: u32 = 28;
const STR_ESC: u32 = 29;
const STR_U0: u32 = 30;
const STR_U1: u32 = 31;
const STR_U2: u32 = 32;
const STR_U3: u32 = 33;
const STR_SB: u32 = 34;
const STR_SU: u32 = 35;
const STR_SL0: u32 = 36;
const STR_SL1: u32 = 37;
const STR_SL2: u32 = 38;
const STR_SL3: u32 = 39;
const N_START: u32 = 50;
const N_INT_FIRST: u32 = 51;
const N_INT_ZERO: u32 = 52;
const N_INT_MORE: u32 = 53;
const N_FRAC_FIRST: u32 = 54;
const N_FRAC_MORE: u32 = 55;
const N_EXP_SIGN: u32 = 56;
const N_EXP_FIRST: u32 = 57;
const N_EXP_MORE: u32 = 58;
const CH_MATCH: u32 = 60;

// -- flat table reads (the exact index arithmetic the WGSL performs against `words`/`bytes`) --------

/// A read-only view over the flat schema buffers, resolving table bases from the header.
pub(crate) struct Flat<'a> {
    words: &'a [u32],
    bytes: &'a [u32],
    nodes_base: usize,
    props_base: usize,
    choice_lits_base: usize,
    spans_base: usize,
    root: u32,
}

impl<'a> Flat<'a> {
    pub(crate) fn new(s: &'a GpuSchema) -> Self {
        Self {
            words: &s.words,
            bytes: &s.bytes,
            nodes_base: s.meta.nodes_base as usize,
            props_base: s.meta.props_base as usize,
            choice_lits_base: s.meta.choice_lits_base as usize,
            spans_base: s.meta.spans_base as usize,
            root: s.meta.root,
        }
    }
    fn kind(&self, node: u32) -> u32 {
        self.words[self.nodes_base + node as usize * NODE_WORDS]
    }
    fn field(&self, node: u32, k: usize) -> u32 {
        self.words[self.nodes_base + node as usize * NODE_WORDS + 1 + k]
    }
    fn obj_count(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    fn obj_value(&self, node: u32, i: u32) -> u32 {
        let base = self.props_base + (self.field(node, 0) + i) as usize * 2;
        self.words[base + 1]
    }
    /// Byte `j` of object `node`'s key `i` (via the props → span → bytes chain).
    fn obj_key_byte(&self, node: u32, i: u32, j: u32) -> u32 {
        let base = self.props_base + (self.field(node, 0) + i) as usize * 2;
        let span = self.words[base];
        self.span_byte(span, j)
    }
    fn obj_key_len(&self, node: u32, i: u32) -> u32 {
        let base = self.props_base + (self.field(node, 0) + i) as usize * 2;
        self.span_len(self.words[base])
    }
    fn arr_item(&self, node: u32) -> u32 {
        self.field(node, 0)
    }
    fn arr_min(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    fn arr_max(&self, node: u32) -> u32 {
        self.field(node, 2)
    }
    fn str_min(&self, node: u32) -> u32 {
        self.field(node, 0)
    }
    fn str_max(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    fn choice_count(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    fn choice_lit_byte(&self, node: u32, c: u32, j: u32) -> u32 {
        let span = self.words[self.choice_lits_base + (self.field(node, 0) + c) as usize];
        self.span_byte(span, j)
    }
    fn choice_lit_len(&self, node: u32, c: u32) -> u32 {
        let span = self.words[self.choice_lits_base + (self.field(node, 0) + c) as usize];
        self.span_len(span)
    }
    fn span_len(&self, span: u32) -> u32 {
        self.words[self.spans_base + span as usize * 2 + 1]
    }
    fn span_byte(&self, span: u32, j: u32) -> u32 {
        let off = self.words[self.spans_base + span as usize * 2];
        self.bytes[(off + j) as usize]
    }
}

// -- state helpers (the [u32;64] layout, mirroring FsmState) ----------------------------------------

type St = [u32; 64];

fn depth(s: &St) -> usize {
    s[0] as usize
}
fn set_depth(s: &mut St, d: usize) {
    s[0] = d as u32;
}
fn fbase(i: usize) -> usize {
    HEADER_WORDS + i * FRAME_WORDS
}
fn node_of(s: &St, i: usize) -> u32 {
    s[fbase(i)]
}
fn phase_of(s: &St, i: usize) -> u32 {
    s[fbase(i) + 1]
}
fn a_of(s: &St, i: usize) -> u32 {
    s[fbase(i) + 2]
}
fn b_of(s: &St, i: usize) -> u32 {
    s[fbase(i) + 3]
}
fn set_phase(s: &mut St, i: usize, p: u32) {
    s[fbase(i) + 1] = p;
}
fn set_a(s: &mut St, i: usize, a: u32) {
    s[fbase(i) + 2] = a;
}
fn set_b(s: &mut St, i: usize, b: u32) {
    s[fbase(i) + 3] = b;
}
fn push(s: &mut St, node: u32, phase: u32, a: u32, b: u32) -> bool {
    let d = depth(s);
    if d >= MAX_STACK {
        return false;
    }
    let base = fbase(d);
    s[base] = node;
    s[base + 1] = phase;
    s[base + 2] = a;
    s[base + 3] = b;
    set_depth(s, d + 1);
    true
}
fn pop(s: &mut St) {
    let d = depth(s);
    if d > 0 {
        set_depth(s, d - 1);
    }
}

/// The fresh `(phase, a, b)` for a new frame of `node` (mirrors `json_fsm::fresh`).
fn fresh(f: &Flat, node: u32) -> (u32, u32, u32) {
    match f.kind(node) {
        K_OBJECT => (OBJ_OPEN, 0, 0),
        K_ARRAY => (ARR_OPEN, 0, 0),
        K_STRING => (STR_OPEN, 0, 0),
        K_NUMBER | K_INTEGER => (N_START, 0, 0),
        K_CHOICE => {
            let count = f.choice_count(node);
            let mask = (((1u64 << count) - 1) & 0xFFFF_FFFF) as u32;
            (CH_MATCH, 0, mask)
        }
        _ => (0, 0, 0),
    }
}

/// The starting state: one frame for the root node.
pub(crate) fn initial(f: &Flat) -> St {
    let mut s = [0u32; 64];
    let (p, a, b) = fresh(f, f.root);
    push(&mut s, f.root, p, a, b);
    s
}

fn is_hex(byte: u32) -> bool {
    matches!(byte, 0x30..=0x39 | 0x41..=0x46 | 0x61..=0x66)
}
fn hexval(byte: u32) -> u32 {
    match byte {
        0x30..=0x39 => byte - 0x30,
        0x41..=0x46 => byte - 0x41 + 10,
        0x61..=0x66 => byte - 0x61 + 10,
        _ => 0,
    }
}
fn is_digit(byte: u32) -> bool {
    matches!(byte, 0x30..=0x39)
}
fn is_number_terminal(phase: u32) -> bool {
    matches!(phase, N_INT_ZERO | N_INT_MORE | N_FRAC_MORE | N_EXP_MORE)
}
fn number_extends(is_int: bool, phase: u32, a: u32, byte: u32) -> bool {
    if a >= NUMBER_BYTE_CAP {
        return false;
    }
    match phase {
        N_INT_ZERO => !is_int && matches!(byte, 0x2E | 0x65 | 0x45),
        N_INT_MORE => is_digit(byte) || (!is_int && matches!(byte, 0x2E | 0x65 | 0x45)),
        N_FRAC_MORE => is_digit(byte) || matches!(byte, 0x65 | 0x45),
        N_EXP_MORE => is_digit(byte),
        _ => false,
    }
}

/// Structural, byte-free transitions (OBJ_FIRST, OBJ_VALUE) — mirrors `json_fsm::normalize`.
fn normalize(f: &Flat, s: &mut St) {
    loop {
        let d = depth(s);
        if d == 0 {
            return;
        }
        let i = d - 1;
        let node = node_of(s, i);
        let kind = f.kind(node);
        let phase = phase_of(s, i);
        if kind == K_OBJECT && phase == OBJ_FIRST {
            if f.obj_count(node) == 0 {
                set_phase(s, i, OBJ_CLOSE_EMPTY);
            } else {
                set_phase(s, i, OBJ_KEY);
                set_a(s, i, 0);
                set_b(s, i, 0);
            }
        } else if kind == K_OBJECT && phase == OBJ_VALUE {
            let vnode = f.obj_value(node, a_of(s, i));
            set_phase(s, i, OBJ_AFTER);
            let (fp, fa, fb) = fresh(f, vnode);
            if !push(s, vnode, fp, fa, fb) {
                return;
            }
        } else {
            return;
        }
    }
}

/// Advance one byte; `Some(new_state)` iff accepted. Mirrors `JsonFsm::step_byte` arm-for-arm.
pub(crate) fn flat_step_byte(f: &Flat, state: &St, byte: u32) -> Option<St> {
    let mut s = *state;
    loop {
        let d = depth(&s);
        if d == 0 {
            return None;
        }
        let i = d - 1;
        let node = node_of(&s, i);
        let phase = phase_of(&s, i);
        let a = a_of(&s, i);
        let b = b_of(&s, i);
        match f.kind(node) {
            K_OBJECT => match phase {
                OBJ_OPEN => return expect(&mut s, i, byte, 0x7B, OBJ_FIRST),
                OBJ_FIRST | OBJ_VALUE => {
                    normalize(f, &mut s);
                    continue;
                }
                OBJ_KEY => {
                    let klen = f.obj_key_len(node, a);
                    if b < klen && f.obj_key_byte(node, a, b) == byte {
                        if b + 1 == klen {
                            set_phase(&mut s, i, OBJ_COLON);
                            set_b(&mut s, i, 0);
                        } else {
                            set_b(&mut s, i, b + 1);
                        }
                        return Some(s);
                    }
                    return None;
                }
                OBJ_COLON => return expect(&mut s, i, byte, 0x3A, OBJ_VALUE),
                OBJ_AFTER => {
                    let count = f.obj_count(node);
                    if a + 1 < count {
                        if byte == 0x2C {
                            set_a(&mut s, i, a + 1);
                            set_phase(&mut s, i, OBJ_KEY);
                            set_b(&mut s, i, 0);
                            return Some(s);
                        }
                    } else if byte == 0x7D {
                        pop(&mut s);
                        return Some(s);
                    }
                    return None;
                }
                OBJ_CLOSE_EMPTY => {
                    if byte == 0x7D {
                        pop(&mut s);
                        return Some(s);
                    }
                    return None;
                }
                _ => return None,
            },
            K_ARRAY => match phase {
                ARR_OPEN => {
                    if byte == 0x5B {
                        set_phase(&mut s, i, ARR_FIRST);
                        set_a(&mut s, i, 0);
                        return Some(s);
                    }
                    return None;
                }
                ARR_FIRST => {
                    let (item, min, max) = (f.arr_item(node), f.arr_min(node), f.arr_max(node));
                    if byte == 0x5D {
                        if min == 0 {
                            pop(&mut s);
                            return Some(s);
                        }
                        return None;
                    }
                    if max > 0 {
                        set_phase(&mut s, i, ARR_AFTER);
                        set_a(&mut s, i, 1);
                        let (fp, fa, fb) = fresh(f, item);
                        if !push(&mut s, item, fp, fa, fb) {
                            return None;
                        }
                        continue;
                    }
                    return None;
                }
                ARR_AFTER => {
                    let (min, max) = (f.arr_min(node), f.arr_max(node));
                    if byte == 0x5D {
                        if a >= min {
                            pop(&mut s);
                            return Some(s);
                        }
                        return None;
                    }
                    if byte == 0x2C && a < max {
                        set_phase(&mut s, i, ARR_ITEM);
                        return Some(s);
                    }
                    return None;
                }
                ARR_ITEM => {
                    let item = f.arr_item(node);
                    set_phase(&mut s, i, ARR_AFTER);
                    set_a(&mut s, i, a + 1);
                    let (fp, fa, fb) = fresh(f, item);
                    if !push(&mut s, item, fp, fa, fb) {
                        return None;
                    }
                    continue;
                }
                _ => return None,
            },
            K_STRING => return step_string(f, &mut s, i, node, phase, a, b, byte).then_some(s),
            K_NUMBER | K_INTEGER => {
                let is_int = f.kind(node) == K_INTEGER;
                if is_number_terminal(phase) && !number_extends(is_int, phase, a, byte) {
                    pop(&mut s);
                    continue;
                }
                return step_number(&mut s, i, is_int, phase, a, byte).then_some(s);
            }
            K_CHOICE => return step_choice(f, &mut s, i, node, a, b, byte).then_some(s),
            _ => return None,
        }
    }
}

fn expect(s: &mut St, i: usize, byte: u32, want: u32, next: u32) -> Option<St> {
    if byte == want {
        set_phase(s, i, next);
        Some(*s)
    } else {
        None
    }
}

#[allow(clippy::too_many_arguments)]
fn step_string(
    f: &Flat,
    s: &mut St,
    i: usize,
    node: u32,
    phase: u32,
    a: u32,
    b: u32,
    byte: u32,
) -> bool {
    match phase {
        STR_OPEN => {
            if byte == 0x22 {
                set_phase(s, i, STR_BODY);
                set_a(s, i, 0);
                true
            } else {
                false
            }
        }
        STR_BODY => {
            let (min, max) = (f.str_min(node), f.str_max(node));
            if byte == 0x22 {
                if a >= min {
                    pop(s);
                    return true;
                }
                return false;
            }
            if byte == 0x5C {
                if a < max {
                    set_phase(s, i, STR_ESC);
                    return true;
                }
                return false;
            }
            if a >= max {
                return false;
            }
            if (0x20..=0x7F).contains(&byte) && byte != 0x22 && byte != 0x5C {
                set_a(s, i, a + 1);
                return true;
            }
            let next = match byte {
                0xC2..=0xDF => STR_C1,
                0xE0 => STR_C2_E0,
                0xED => STR_C2_ED,
                0xE1..=0xEC | 0xEE..=0xEF => STR_C2,
                0xF0 => STR_C3_F0,
                0xF4 => STR_C3_F4,
                0xF1..=0xF3 => STR_C3,
                _ => return false,
            };
            set_phase(s, i, next);
            true
        }
        STR_C1 => cont(s, i, byte, 0x80, 0xBF, STR_BODY, true, a),
        STR_C2 => cont(s, i, byte, 0x80, 0xBF, STR_C1, false, a),
        STR_C2_E0 => cont(s, i, byte, 0xA0, 0xBF, STR_C1, false, a),
        STR_C2_ED => cont(s, i, byte, 0x80, 0x9F, STR_C1, false, a),
        STR_C3 => cont(s, i, byte, 0x80, 0xBF, STR_C2, false, a),
        STR_C3_F0 => cont(s, i, byte, 0x90, 0xBF, STR_C2, false, a),
        STR_C3_F4 => cont(s, i, byte, 0x80, 0x8F, STR_C2, false, a),
        STR_ESC => match byte {
            0x22 | 0x5C | 0x2F | 0x62 | 0x66 | 0x6E | 0x72 | 0x74 => {
                set_phase(s, i, STR_BODY);
                set_a(s, i, a + 1);
                true
            }
            0x75 => {
                set_phase(s, i, STR_U0);
                set_b(s, i, 0);
                true
            }
            _ => false,
        },
        STR_U0 | STR_U1 | STR_U2 => {
            if is_hex(byte) {
                let shift = match phase {
                    STR_U0 => 12,
                    STR_U1 => 8,
                    _ => 4,
                };
                set_b(s, i, b | (hexval(byte) << shift));
                set_phase(s, i, phase + 1);
                true
            } else {
                false
            }
        }
        STR_U3 => {
            if is_hex(byte) {
                let cu = b | hexval(byte);
                if (0xD800..=0xDBFF).contains(&cu) {
                    set_phase(s, i, STR_SB);
                    true
                } else if (0xDC00..=0xDFFF).contains(&cu) {
                    false
                } else {
                    set_phase(s, i, STR_BODY);
                    set_a(s, i, a + 1);
                    set_b(s, i, 0);
                    true
                }
            } else {
                false
            }
        }
        STR_SB => set_phase_if(s, i, byte == 0x5C, STR_SU),
        STR_SU => set_phase_if(s, i, byte == 0x75, STR_SL0),
        STR_SL0 => set_phase_if(s, i, byte == 0x64 || byte == 0x44, STR_SL1),
        STR_SL1 => set_phase_if(s, i, matches!(byte, 0x63..=0x66 | 0x43..=0x46), STR_SL2),
        STR_SL2 => set_phase_if(s, i, is_hex(byte), STR_SL3),
        STR_SL3 if is_hex(byte) => {
            set_phase(s, i, STR_BODY);
            set_a(s, i, a + 1);
            set_b(s, i, 0);
            true
        }
        _ => false,
    }
}

#[allow(clippy::too_many_arguments)]
fn cont(s: &mut St, i: usize, byte: u32, lo: u32, hi: u32, next: u32, count: bool, a: u32) -> bool {
    if (lo..=hi).contains(&byte) {
        set_phase(s, i, next);
        if count {
            set_a(s, i, a + 1);
        }
        true
    } else {
        false
    }
}

fn set_phase_if(s: &mut St, i: usize, ok: bool, next: u32) -> bool {
    if ok {
        set_phase(s, i, next);
    }
    ok
}

fn step_number(s: &mut St, i: usize, is_int: bool, phase: u32, a: u32, byte: u32) -> bool {
    let na = a + 1;
    let cap = a < NUMBER_BYTE_CAP;
    let go = |s: &mut St, next: u32| {
        set_phase(s, i, next);
        set_a(s, i, na);
        true
    };
    match phase {
        N_START => match byte {
            0x2D => go(s, N_INT_FIRST),
            0x30 => go(s, N_INT_ZERO),
            0x31..=0x39 => go(s, N_INT_MORE),
            _ => false,
        },
        N_INT_FIRST => match byte {
            0x30 => go(s, N_INT_ZERO),
            0x31..=0x39 => go(s, N_INT_MORE),
            _ => false,
        },
        N_INT_ZERO => {
            if !is_int && cap && byte == 0x2E {
                go(s, N_FRAC_FIRST)
            } else if !is_int && cap && matches!(byte, 0x65 | 0x45) {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_INT_MORE => {
            if cap && is_digit(byte) {
                go(s, N_INT_MORE)
            } else if !is_int && cap && byte == 0x2E {
                go(s, N_FRAC_FIRST)
            } else if !is_int && cap && matches!(byte, 0x65 | 0x45) {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_FRAC_FIRST => {
            if cap && is_digit(byte) {
                go(s, N_FRAC_MORE)
            } else {
                false
            }
        }
        N_FRAC_MORE => {
            if cap && is_digit(byte) {
                go(s, N_FRAC_MORE)
            } else if cap && matches!(byte, 0x65 | 0x45) {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_EXP_SIGN => {
            if cap && matches!(byte, 0x2B | 0x2D) {
                go(s, N_EXP_FIRST)
            } else if cap && is_digit(byte) {
                go(s, N_EXP_MORE)
            } else {
                false
            }
        }
        N_EXP_FIRST => {
            if cap && is_digit(byte) {
                go(s, N_EXP_MORE)
            } else {
                false
            }
        }
        N_EXP_MORE if cap && is_digit(byte) => go(s, N_EXP_MORE),
        _ => false,
    }
}

fn step_choice(f: &Flat, s: &mut St, i: usize, node: u32, a: u32, mask: u32, byte: u32) -> bool {
    let count = f.choice_count(node);
    let pos = a;
    let mut newmask = 0u32;
    let mut done = false;
    for c in 0..count {
        if (mask >> c) & 1 == 0 {
            continue;
        }
        let len = f.choice_lit_len(node, c);
        if pos < len && f.choice_lit_byte(node, c, pos) == byte {
            if pos + 1 == len {
                done = true;
            } else {
                newmask |= 1 << c;
            }
        }
    }
    if done {
        pop(s);
        true
    } else if newmask != 0 {
        set_a(s, i, a + 1);
        set_b(s, i, newmask);
        true
    } else {
        false
    }
}

/// Walk a token's raw bytes; `Some(new_state)` iff every byte is accepted (mirrors `step_token`).
pub(crate) fn flat_step_token(f: &Flat, state: &St, bytes: &[u8]) -> Option<St> {
    let mut s = *state;
    for &byte in bytes {
        s = flat_step_byte(f, &s, byte as u32)?;
    }
    Some(s)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::grammar::harness::GpuSchema;
    use crate::grammar::json_fsm::{FsmState, JsonFsm};
    use crate::grammar::json_schema::JsonSchemaTables;
    use serde_json::json;

    /// The corpus: one representative schema per node kind + nesting, each with a set of valid
    /// instances whose byte prefixes exercise every transition. Mirrors the G1 corpus shape.
    fn corpus() -> Vec<(serde_json::Value, Vec<&'static str>)> {
        vec![
            // Flat object of strings (Object + StringFree + key matching + stack push/pop).
            (
                json!({
                    "type": "object",
                    "properties": {"name": {"type": "string"}, "city": {"type": "string"}},
                    "required": ["name", "city"],
                    "additionalProperties": false
                }),
                vec![
                    r#"{"name":"Ada","city":"Paris"}"#,
                    r#"{"name":"","city":"x"}"#,
                ],
            ),
            // Enums + bool + null + number + integer (Choice + Number + Integer).
            (
                json!({
                    "type": "object",
                    "properties": {
                        "role": {"enum": ["admin", "user"]},
                        "active": {"type": "boolean"},
                        "note": {"type": "null"},
                        "score": {"type": "number"},
                        "age": {"type": "integer"}
                    },
                    "required": ["role", "active", "note", "score", "age"],
                    "additionalProperties": false
                }),
                vec![
                    r#"{"role":"admin","active":true,"note":null,"score":-0.5e10,"age":-7}"#,
                    r#"{"role":"user","active":false,"note":null,"score":0,"age":42}"#,
                ],
            ),
            // Array with bounds + nested object (Array + nesting).
            (
                json!({
                    "type": "object",
                    "properties": {
                        "tags": {"type": "array", "items": {"type": "string"}, "minItems": 0, "maxItems": 3},
                        "pt": {
                            "type": "object",
                            "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}},
                            "required": ["x", "y"],
                            "additionalProperties": false
                        }
                    },
                    "required": ["tags", "pt"],
                    "additionalProperties": false
                }),
                vec![
                    r#"{"tags":[],"pt":{"x":1,"y":2}}"#,
                    r#"{"tags":["a","bb"],"pt":{"x":-3,"y":0}}"#,
                ],
            ),
            // Unicode string content: multi-byte + escapes + surrogate pair.
            (
                json!({
                    "type": "object",
                    "properties": {"s": {"type": "string"}},
                    "required": ["s"],
                    "additionalProperties": false
                }),
                vec![r#"{"s":"café-日本語"}"#, r#"{"s":"tab\there \"q\" A 😀"}"#],
            ),
        ]
    }

    /// The flat interpreter must agree with the reference [`JsonFsm`] on EVERY byte of every valid
    /// instance: same accept/reject decision AND same resulting 64-word state — proving the flat
    /// buffers carry enough to reproduce the automaton (what the WGSL will do).
    #[test]
    fn should_match_reference_fsm_over_flat_tables() {
        let mut checked = 0usize;
        for (schema, instances) in corpus() {
            let tables = JsonSchemaTables::compile(&schema).expect("compile");
            let gpu = GpuSchema::from_tables(&tables);
            let f = Flat::new(&gpu);

            // Initial states must match exactly.
            let ref_init = JsonFsm::initial(&tables);
            let flat_init = initial(&f);
            assert_eq!(
                &flat_init,
                ref_init.as_words(),
                "initial state mismatch for {schema}"
            );

            for inst in instances {
                let mut rs = JsonFsm::initial(&tables);
                let mut fs = initial(&f);
                for (bi, &byte) in inst.as_bytes().iter().enumerate() {
                    let r = JsonFsm::step_byte(&tables, &rs, byte);
                    let g = flat_step_byte(&f, &fs, byte as u32);
                    assert_eq!(
                        r.is_some(),
                        g.is_some(),
                        "accept/reject diverged at byte {bi} (0x{byte:02x}) of {inst:?}"
                    );
                    if let (Some(rn), Some(gn)) = (r, g) {
                        assert_eq!(
                            rn.as_words(),
                            &gn,
                            "state diverged after byte {bi} (0x{byte:02x}) of {inst:?}"
                        );
                        rs = rn;
                        fs = gn;
                    } else {
                        break;
                    }
                }
                // A complete valid instance must leave BOTH interpreters complete (depth 0).
                assert!(
                    JsonFsm::is_complete(&tables, &rs),
                    "reference not complete after {inst:?}"
                );
                assert_eq!(depth(&fs), 0, "flat interp not complete after {inst:?}");
                checked += 1;
            }
        }
        assert!(checked >= 8, "corpus too small ({checked})");
        eprintln!("flat-table interpreter OK: {checked} instances byte-exact vs JsonFsm");
    }

    /// The per-token mask the WGSL `mask_main` produces must equal "does `step_token` accept this
    /// token from the current state" — verified here at every token boundary of a real instance
    /// using a small synthetic byte vocabulary (the exact GPU semantics, CPU-side).
    #[test]
    fn should_reproduce_step_token_mask_via_flat_tables() {
        // A byte-level vocab: single-byte tokens for the JSON structural/content bytes we use,
        // plus a few multi-byte pieces (to exercise mid-token walks) and a never-legal one.
        let pieces: Vec<Vec<u8>> = vec![
            b"{".to_vec(),
            b"}".to_vec(),
            b"\"".to_vec(),
            b":".to_vec(),
            b",".to_vec(),
            b"name".to_vec(),
            b"city".to_vec(),
            b"Ada".to_vec(),
            b"Paris".to_vec(),
            b"xyz".to_vec(),
            b"na".to_vec(),  // prefix of "name" — legal partial
            b"nax".to_vec(), // wrong continuation — illegal
            b"\"}".to_vec(), // multi-byte structural
        ];
        let schema = json!({
            "type": "object",
            "properties": {"name": {"type": "string"}, "city": {"type": "string"}},
            "required": ["name", "city"],
            "additionalProperties": false
        });
        let tables = JsonSchemaTables::compile(&schema).expect("compile");
        let gpu = GpuSchema::from_tables(&tables);
        let f = Flat::new(&gpu);

        // Drive a valid instance token-by-token; at each boundary, the flat mask over the vocab
        // must equal the reference step_token accept set.
        let plan: &[&[u8]] = &[b"{", b"\"", b"name", b"\"", b":", b"\"", b"Ada", b"\""];
        let mut rs = JsonFsm::initial(&tables);
        let mut fs = initial(&f);
        let mut boundaries = 0usize;
        for step in plan {
            for (id, piece) in pieces.iter().enumerate() {
                let refn = JsonFsm::step_token(&tables, &rs, piece).is_some();
                let flatn = flat_step_token(&f, &fs, piece).is_some();
                assert_eq!(
                    refn, flatn,
                    "token {id} ({piece:?}) mask diverged at boundary {boundaries}"
                );
            }
            // Advance both by the planned real token.
            rs = JsonFsm::step_token(&tables, &rs, step).expect("plan token valid (ref)");
            fs = flat_step_token(&f, &fs, step).expect("plan token valid (flat)");
            boundaries += 1;
        }
        // Round-trips through the public FsmState words API (the exact GPU state_buf contents).
        assert_eq!(FsmState::from_words(fs).as_words(), &fs);
        eprintln!(
            "flat-table token mask OK: {boundaries} boundaries × {} tokens",
            pieces.len()
        );
    }
}