inferencelayer 0.2.1

//! The fixed table-interpreter: a byte-driven pushdown automaton over [`JsonSchemaTables`].
//!
//! This is the CPU reference the WGSL kernel (G2) will mirror byte-for-byte, so it is written as a
//! TABLE interpreter, not an AST walk: all schema structure lives in the flat tables, and all
//! runtime state lives in one fixed [`STATE_WORDS`]-word array. Every operation is deterministic and
//! has no interior mutability, so it is trivially wasm-safe and GPU-portable.
//!
//! ## State layout (`[u32; 64]`)
//!
//! - word 0: `depth` — number of active frames (0 = complete/empty).
//! - words 1..=3: reserved (kept zero; round the layout to 64 and give G2 spare words without
//!   shifting any frame offset).
//! - words 4..64: [`MAX_STACK`] = 15 frames of [`FRAME_WORDS`] = 4 words each, frame `i` at
//!   `4 + 4*i`: `[node, phase, a, b]`.
//!   - `node`  — the schema node this frame executes.
//!   - `phase` — the sub-state within that node (the `*_PHASE`/`STR_*`/`N_*`/`CH_*` constants).
//!   - `a`     — primary counter: object property index · array item count · string char count ·
//!     number byte count · choice match position.
//!   - `b`     — secondary counter: object matched-key-length · string `\u` accumulator · choice
//!     viable-candidate bitmask.
//!
//! The stack holds a frame per OPEN value (containers AND the innermost scalar). Nesting ≤ 10
//! containers plus one leaf scalar ⇒ ≤ 11 frames; 15 leaves headroom.
//!
//! ## Step contract
//!
//! `step_byte` runs one loop: apply structural transitions that consume no byte (an object entering
//! its key section, or delegating a property value to a fresh child frame); let a completed
//! open-ended scalar (number/integer — the only values with no closing delimiter) fall through to
//! its parent; otherwise apply the single byte transition for the top frame. `allowed_bytes`
//! recomputes the same mask directly; the two paths are cross-checked exhaustively by the gate.

use super::json_schema::{JsonSchemaTables, NodeKind};

/// The fixed state-word count. Locked in lockstep with the frame arithmetic below and mirrored by
/// G2's WGSL `array<u32,64>`.
pub const STATE_WORDS: usize = 64;
const HEADER_WORDS: usize = 4;
const FRAME_WORDS: usize = 4;
/// Maximum simultaneous frames (see the layout note). 15 covers depth-10 nesting with headroom.
pub const MAX_STACK: usize = 15;
/// The JSON number grammar is capped at this many bytes (matches the OpenAI strict-mode budget).
pub const NUMBER_BYTE_CAP: u32 = 24;

// Lockstep guards — the WGSL hardcodes the same 64-word layout; keep these in sync or the GPU
// buffers are mis-sized (mirrors the OSFQL discipline at gpu_grammar.rs).
const _: () = assert!(HEADER_WORDS + FRAME_WORDS * MAX_STACK == STATE_WORDS);
const _: () = assert!(STATE_WORDS == 64, "state layout is a fixed 64 words");

const F_NODE: usize = 0;
const F_PHASE: usize = 1;
const F_A: usize = 2;
const F_B: usize = 3;

// -- phase constants (namespaced by kind; dispatch is always on (kind, phase)) --------------------

// Object.
const OBJ_OPEN: u32 = 0; // expect `{`
const OBJ_FIRST: u32 = 1; // structural: enter key section or empty-close
const OBJ_KEY: u32 = 2; // match property key literal (a = prop index, b = matched len)
const OBJ_COLON: u32 = 3; // expect `:`
const OBJ_VALUE: u32 = 4; // structural: push the property value child
const OBJ_AFTER: u32 = 5; // after a value: expect `,` (more props) or `}`
const OBJ_CLOSE_EMPTY: u32 = 6; // empty object: expect `}`

// Array.
const ARR_OPEN: u32 = 10; // expect `[`
const ARR_FIRST: u32 = 11; // before the first item: `]` or an item
const ARR_AFTER: u32 = 12; // after an item: `]` or `,`
const ARR_ITEM: u32 = 13; // after a `,`: a required item

// String (a = code-point count; b = `\u` accumulator during escapes).
const STR_OPEN: u32 = 20; // expect opening `"`
const STR_BODY: u32 = 21; // string content
const STR_C1: u32 = 22; // expect 1 UTF-8 continuation, then char complete
const STR_C2: u32 = 23; // expect a continuation, then C1
const STR_C2_E0: u32 = 24; // E0 lead: continuation 0xA0..=0xBF, then C1
const STR_C2_ED: u32 = 25; // ED lead: continuation 0x80..=0x9F, then C1
const STR_C3: u32 = 26; // expect a continuation, then C2
const STR_C3_F0: u32 = 27; // F0 lead: continuation 0x90..=0xBF, then C2
const STR_C3_F4: u32 = 28; // F4 lead: continuation 0x80..=0x8F, then C2
const STR_ESC: u32 = 29; // after `\`
const STR_U0: u32 = 30; // \u hex digit 1
const STR_U1: u32 = 31; // \u hex digit 2
const STR_U2: u32 = 32; // \u hex digit 3
const STR_U3: u32 = 33; // \u hex digit 4, then classify BMP/surrogate
const STR_SB: u32 = 34; // high surrogate seen: expect `\`
const STR_SU: u32 = 35; // expect `u`
const STR_SL0: u32 = 36; // low surrogate hex 1 (must be D)
const STR_SL1: u32 = 37; // low surrogate hex 2 (must be C..F)
const STR_SL2: u32 = 38; // low surrogate hex 3
const STR_SL3: u32 = 39; // low surrogate hex 4, then pair complete

// Number/Integer (a = byte count). Terminable phases can legally end the value.
const N_START: u32 = 50; // optional `-`, then first int digit
const N_INT_FIRST: u32 = 51; // first int digit after `-`
const N_INT_ZERO: u32 = 52; // leading 0 (terminable; no more int digits)
const N_INT_MORE: u32 = 53; // int digits after 1-9 (terminable)
const N_FRAC_FIRST: u32 = 54; // first fraction digit after `.`
const N_FRAC_MORE: u32 = 55; // fraction digits (terminable)
const N_EXP_SIGN: u32 = 56; // optional exponent sign
const N_EXP_FIRST: u32 = 57; // first exponent digit
const N_EXP_MORE: u32 = 58; // exponent digits (terminable)

// Choice (a = match position; b = viable-candidate bitmask).
const CH_MATCH: u32 = 60;

/// One [`STATE_WORDS`]-word interpreter state. `Copy` so callers thread it functionally with no
/// interior mutability (each `step_*` returns a fresh state).
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct FsmState {
    w: [u32; STATE_WORDS],
}

impl FsmState {
    /// The fixed state-word count, exposed for the lockstep gate.
    pub const WORDS: usize = STATE_WORDS;

    /// The raw 64 state words — the exact contents of the GPU `state_buf`. G2 uploads these to the
    /// device and reads them back after `advance_main`, so CPU and GPU thread the identical layout.
    pub fn as_words(&self) -> &[u32; STATE_WORDS] {
        &self.w
    }

    /// Reconstruct a state from raw words read back from the GPU `state_buf`.
    pub fn from_words(words: [u32; STATE_WORDS]) -> Self {
        Self { w: words }
    }

    fn empty() -> Self {
        Self {
            w: [0; STATE_WORDS],
        }
    }
    fn depth(&self) -> usize {
        self.w[0] as usize
    }
    fn set_depth(&mut self, d: usize) {
        self.w[0] = d as u32;
    }
    fn base(i: usize) -> usize {
        HEADER_WORDS + i * FRAME_WORDS
    }
    fn node(&self, i: usize) -> u32 {
        self.w[Self::base(i) + F_NODE]
    }
    fn phase(&self, i: usize) -> u32 {
        self.w[Self::base(i) + F_PHASE]
    }
    fn a(&self, i: usize) -> u32 {
        self.w[Self::base(i) + F_A]
    }
    fn b(&self, i: usize) -> u32 {
        self.w[Self::base(i) + F_B]
    }
    fn set_phase(&mut self, i: usize, p: u32) {
        self.w[Self::base(i) + F_PHASE] = p;
    }
    fn set_a(&mut self, i: usize, a: u32) {
        self.w[Self::base(i) + F_A] = a;
    }
    fn set_b(&mut self, i: usize, b: u32) {
        self.w[Self::base(i) + F_B] = b;
    }
    /// Push a fresh frame; returns `false` if the depth budget is exhausted.
    fn push(&mut self, node: u32, phase: u32, a: u32, b: u32) -> bool {
        let d = self.depth();
        if d >= MAX_STACK {
            return false;
        }
        let base = Self::base(d);
        self.w[base + F_NODE] = node;
        self.w[base + F_PHASE] = phase;
        self.w[base + F_A] = a;
        self.w[base + F_B] = b;
        self.set_depth(d + 1);
        true
    }
    fn pop(&mut self) {
        let d = self.depth();
        if d > 0 {
            self.set_depth(d - 1);
        }
    }
}

/// The initial fields of a fresh frame for `node` (phase, a, b).
fn fresh(t: &JsonSchemaTables, node: u32) -> (u32, u32, u32) {
    match t.kind(node) {
        NodeKind::Object => (OBJ_OPEN, 0, 0),
        NodeKind::Array => (ARR_OPEN, 0, 0),
        NodeKind::StringFree => (STR_OPEN, 0, 0),
        NodeKind::Number | NodeKind::Integer => (N_START, 0, 0),
        NodeKind::Choice => {
            let count = t.choice_count(node);
            // Viable-candidate bitmask, all set (count ≤ 32 ⇒ fits 32 bits).
            let mask = (((1u64 << count) - 1) & 0xFFFF_FFFF) as u32;
            (CH_MATCH, 0, mask)
        }
    }
}

fn hexval(byte: u8) -> Option<u32> {
    match byte {
        b'0'..=b'9' => Some((byte - b'0') as u32),
        b'a'..=b'f' => Some((byte - b'a' + 10) as u32),
        b'A'..=b'F' => Some((byte - b'A' + 10) as u32),
        _ => None,
    }
}

fn is_number_terminal(phase: u32) -> bool {
    matches!(phase, N_INT_ZERO | N_INT_MORE | N_FRAC_MORE | N_EXP_MORE)
}

/// Whether `byte` extends an already-complete number from a terminal `phase` (mirrors the terminal
/// arms of [`step_number`] exactly). Used to decide the "scalar falls through to parent" step.
fn number_extends(kind: NodeKind, phase: u32, a: u32, byte: u8) -> bool {
    if a >= NUMBER_BYTE_CAP {
        return false;
    }
    let is_int = matches!(kind, NodeKind::Integer);
    match phase {
        N_INT_ZERO => !is_int && matches!(byte, b'.' | b'e' | b'E'),
        N_INT_MORE => byte.is_ascii_digit() || (!is_int && matches!(byte, b'.' | b'e' | b'E')),
        N_FRAC_MORE => byte.is_ascii_digit() || matches!(byte, b'e' | b'E'),
        N_EXP_MORE => byte.is_ascii_digit(),
        _ => false,
    }
}

/// Apply the structural, byte-free transitions (`OBJ_FIRST`, `OBJ_VALUE`) until the top frame
/// expects a concrete byte (or the stack is empty). Shared by `allowed_bytes`/`is_complete` and
/// the head of `step_byte`.
fn normalize(t: &JsonSchemaTables, s: &mut FsmState) {
    loop {
        let d = s.depth();
        if d == 0 {
            return;
        }
        let i = d - 1;
        let node = s.node(i);
        match (t.kind(node), s.phase(i)) {
            (NodeKind::Object, OBJ_FIRST) => {
                if t.obj_count(node) == 0 {
                    s.set_phase(i, OBJ_CLOSE_EMPTY);
                } else {
                    s.set_phase(i, OBJ_KEY);
                    s.set_a(i, 0);
                    s.set_b(i, 0);
                }
            }
            (NodeKind::Object, OBJ_VALUE) => {
                let vnode = t.obj_value(node, s.a(i) as usize);
                s.set_phase(i, OBJ_AFTER); // parent resumes here when the child pops
                let (fp, fa, fb) = fresh(t, vnode);
                if !s.push(vnode, fp, fa, fb) {
                    return;
                }
            }
            _ => return,
        }
    }
}

/// The reference interpreter. Deterministic, no interior mutability.
pub struct JsonFsm;

impl JsonFsm {
    /// The starting state: a single frame for the root node.
    pub fn initial(t: &JsonSchemaTables) -> FsmState {
        let mut s = FsmState::empty();
        let root = t.root();
        let (p, a, b) = fresh(t, root);
        s.push(root, p, a, b);
        s
    }

    /// Advance by one output byte, or `None` if the byte is not allowed here.
    pub fn step_byte(t: &JsonSchemaTables, state: &FsmState, byte: u8) -> Option<FsmState> {
        let mut s = *state;
        loop {
            let d = s.depth();
            if d == 0 {
                return None; // already complete: nothing may follow
            }
            let i = d - 1;
            let node = s.node(i);
            let phase = s.phase(i);
            let a = s.a(i);
            let b = s.b(i);
            match t.kind(node) {
                NodeKind::Object => match phase {
                    OBJ_OPEN => return step_expect(&mut s, i, byte, b'{', OBJ_FIRST),
                    OBJ_FIRST | OBJ_VALUE => {
                        normalize(t, &mut s);
                        continue;
                    }
                    OBJ_KEY => {
                        let key = t.obj_key(node, a as usize);
                        let ml = b as usize;
                        if ml < key.len() && key[ml] == byte {
                            if ml + 1 == key.len() {
                                s.set_phase(i, OBJ_COLON);
                                s.set_b(i, 0);
                            } else {
                                s.set_b(i, (ml + 1) as u32);
                            }
                            return Some(s);
                        }
                        return None;
                    }
                    OBJ_COLON => return step_expect(&mut s, i, byte, b':', OBJ_VALUE),
                    OBJ_AFTER => {
                        let count = t.obj_count(node);
                        if a + 1 < count {
                            if byte == b',' {
                                s.set_a(i, a + 1);
                                s.set_phase(i, OBJ_KEY);
                                s.set_b(i, 0);
                                return Some(s);
                            }
                        } else if byte == b'}' {
                            s.pop();
                            return Some(s);
                        }
                        return None;
                    }
                    OBJ_CLOSE_EMPTY => {
                        if byte == b'}' {
                            s.pop();
                            return Some(s);
                        }
                        return None;
                    }
                    _ => return None,
                },
                NodeKind::Array => match phase {
                    ARR_OPEN => {
                        if byte == b'[' {
                            s.set_phase(i, ARR_FIRST);
                            s.set_a(i, 0);
                            return Some(s);
                        }
                        return None;
                    }
                    ARR_FIRST => {
                        let (item, min, max) = (t.arr_item(node), t.arr_min(node), t.arr_max(node));
                        if byte == b']' {
                            if min == 0 {
                                s.pop();
                                return Some(s);
                            }
                            return None;
                        }
                        if max > 0 {
                            s.set_phase(i, ARR_AFTER);
                            s.set_a(i, 1);
                            let (fp, fa, fb) = fresh(t, item);
                            if !s.push(item, fp, fa, fb) {
                                return None;
                            }
                            continue; // the child consumes this byte
                        }
                        return None;
                    }
                    ARR_AFTER => {
                        let (min, max) = (t.arr_min(node), t.arr_max(node));
                        if byte == b']' {
                            if a >= min {
                                s.pop();
                                return Some(s);
                            }
                            return None;
                        }
                        if byte == b',' && a < max {
                            s.set_phase(i, ARR_ITEM);
                            return Some(s);
                        }
                        return None;
                    }
                    ARR_ITEM => {
                        let item = t.arr_item(node);
                        s.set_phase(i, ARR_AFTER);
                        s.set_a(i, a + 1);
                        let (fp, fa, fb) = fresh(t, item);
                        if !s.push(item, fp, fa, fb) {
                            return None;
                        }
                        continue; // the child consumes this byte
                    }
                    _ => return None,
                },
                NodeKind::StringFree => {
                    return step_string(t, &mut s, i, node, phase, a, b, byte).then_some(s);
                }
                NodeKind::Number | NodeKind::Integer => {
                    let kind = t.kind(node);
                    if is_number_terminal(phase) && !number_extends(kind, phase, a, byte) {
                        s.pop();
                        continue; // the completed number's parent handles this byte
                    }
                    return step_number(&mut s, i, kind, phase, a, byte).then_some(s);
                }
                NodeKind::Choice => {
                    return step_choice(t, &mut s, i, node, phase, a, b, byte).then_some(s);
                }
            }
        }
    }

    /// The set of bytes that would be accepted next. Computed directly (independent of `step_byte`);
    /// the gate cross-checks the two over the full 0..=255 sweep.
    pub fn allowed_bytes(t: &JsonSchemaTables, state: &FsmState) -> [bool; 256] {
        let mut set = [false; 256];
        let mut s = *state;
        normalize(t, &mut s);
        fill_top(t, &s, &mut set);
        set
    }

    /// Whether the generation may legally STOP here — the root value has been fully emitted.
    pub fn is_complete(t: &JsonSchemaTables, state: &FsmState) -> bool {
        let mut s = *state;
        loop {
            normalize(t, &mut s);
            let d = s.depth();
            if d == 0 {
                return true;
            }
            let i = d - 1;
            let kind = t.kind(s.node(i));
            if matches!(kind, NodeKind::Number | NodeKind::Integer)
                && is_number_terminal(s.phase(i))
            {
                s.pop(); // an open-ended scalar can end; check the parent
                continue;
            }
            return false; // any other open frame needs a closing byte first
        }
    }

    /// Advance by one token's raw bytes (folds [`Self::step_byte`]); `None` if any byte is rejected.
    /// Multi-byte UTF-8 tokens split mid-character walk correctly by construction.
    pub fn step_token(t: &JsonSchemaTables, state: &FsmState, bytes: &[u8]) -> Option<FsmState> {
        let mut s = *state;
        for &byte in bytes {
            s = Self::step_byte(t, &s, byte)?;
        }
        Some(s)
    }
}

/// Consume `byte` iff it equals `want`, moving to `next` phase.
fn step_expect(s: &mut FsmState, i: usize, byte: u8, want: u8, next: u32) -> Option<FsmState> {
    if byte == want {
        s.set_phase(i, next);
        Some(*s)
    } else {
        None
    }
}

#[allow(clippy::too_many_arguments)]
fn step_string(
    t: &JsonSchemaTables,
    s: &mut FsmState,
    i: usize,
    node: u32,
    phase: u32,
    a: u32,
    b: u32,
    byte: u8,
) -> bool {
    match phase {
        STR_OPEN => {
            if byte == b'"' {
                s.set_phase(i, STR_BODY);
                s.set_a(i, 0);
                true
            } else {
                false
            }
        }
        STR_BODY => {
            let (min, max) = (t.str_min(node), t.str_max(node));
            if byte == b'"' {
                if a >= min {
                    s.pop();
                    return true;
                }
                return false;
            }
            if byte == b'\\' {
                if a < max {
                    s.set_phase(i, STR_ESC);
                    return true;
                }
                return false;
            }
            if a >= max {
                return false;
            }
            // 1-byte content: printable ASCII except `"` and `\` (control bytes < 0x20 excluded).
            if (0x20..=0x7F).contains(&byte) && byte != b'"' && byte != b'\\' {
                s.set_a(i, a + 1);
                return true;
            }
            // multi-byte UTF-8 leads (well-formedness enforced by the continuation phases)
            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,
            };
            s.set_phase(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 {
            b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't' => {
                s.set_phase(i, STR_BODY);
                s.set_a(i, a + 1);
                true
            }
            b'u' => {
                s.set_phase(i, STR_U0);
                s.set_b(i, 0);
                true
            }
            _ => false,
        },
        STR_U0 | STR_U1 | STR_U2 => {
            if let Some(h) = hexval(byte) {
                let shift = match phase {
                    STR_U0 => 12,
                    STR_U1 => 8,
                    _ => 4,
                };
                s.set_b(i, b | (h << shift));
                s.set_phase(i, phase + 1);
                true
            } else {
                false
            }
        }
        STR_U3 => {
            if let Some(h) = hexval(byte) {
                let cu = b | h;
                if (0xD800..=0xDBFF).contains(&cu) {
                    s.set_phase(i, STR_SB); // high surrogate: a low surrogate must follow
                    true
                } else if (0xDC00..=0xDFFF).contains(&cu) {
                    false // lone low surrogate is invalid JSON
                } else {
                    s.set_phase(i, STR_BODY);
                    s.set_a(i, a + 1);
                    s.set_b(i, 0);
                    true
                }
            } else {
                false
            }
        }
        STR_SB => step_set_phase(s, i, byte == b'\\', STR_SU),
        STR_SU => step_set_phase(s, i, byte == b'u', STR_SL0),
        STR_SL0 => step_set_phase(s, i, byte == b'd' || byte == b'D', STR_SL1),
        STR_SL1 => step_set_phase(
            s,
            i,
            matches!(byte, b'c' | b'd' | b'e' | b'f' | b'C' | b'D' | b'E' | b'F'),
            STR_SL2,
        ),
        STR_SL2 => step_set_phase(s, i, hexval(byte).is_some(), STR_SL3),
        STR_SL3 if hexval(byte).is_some() => {
            s.set_phase(i, STR_BODY);
            s.set_a(i, a + 1); // the surrogate pair is one code point
            s.set_b(i, 0);
            true
        }
        _ => false,
    }
}

/// A UTF-8 continuation step: accept `byte` in `range`, move to `next`; when `count`, the character
/// is complete so bump the code-point counter.
fn cont(
    s: &mut FsmState,
    i: usize,
    byte: u8,
    range: std::ops::RangeInclusive<u8>,
    next: u32,
    count: bool,
    a: u32,
) -> bool {
    if range.contains(&byte) {
        s.set_phase(i, next);
        if count {
            s.set_a(i, a + 1);
        }
        true
    } else {
        false
    }
}

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

fn step_number(s: &mut FsmState, i: usize, kind: NodeKind, phase: u32, a: u32, byte: u8) -> bool {
    let is_int = matches!(kind, NodeKind::Integer);
    let na = a + 1;
    let cap = a < NUMBER_BYTE_CAP;
    let go = |s: &mut FsmState, next: u32| {
        s.set_phase(i, next);
        s.set_a(i, na);
        true
    };
    match phase {
        N_START => match byte {
            b'-' => go(s, N_INT_FIRST),
            b'0' => go(s, N_INT_ZERO),
            b'1'..=b'9' => go(s, N_INT_MORE),
            _ => false,
        },
        N_INT_FIRST => match byte {
            b'0' => go(s, N_INT_ZERO),
            b'1'..=b'9' => go(s, N_INT_MORE),
            _ => false,
        },
        N_INT_ZERO => {
            if !is_int && cap && byte == b'.' {
                go(s, N_FRAC_FIRST)
            } else if !is_int && cap && matches!(byte, b'e' | b'E') {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_INT_MORE => {
            if cap && byte.is_ascii_digit() {
                go(s, N_INT_MORE)
            } else if !is_int && cap && byte == b'.' {
                go(s, N_FRAC_FIRST)
            } else if !is_int && cap && matches!(byte, b'e' | b'E') {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_FRAC_FIRST => {
            if cap && byte.is_ascii_digit() {
                go(s, N_FRAC_MORE)
            } else {
                false
            }
        }
        N_FRAC_MORE => {
            if cap && byte.is_ascii_digit() {
                go(s, N_FRAC_MORE)
            } else if cap && matches!(byte, b'e' | b'E') {
                go(s, N_EXP_SIGN)
            } else {
                false
            }
        }
        N_EXP_SIGN => {
            if cap && matches!(byte, b'+' | b'-') {
                go(s, N_EXP_FIRST)
            } else if cap && byte.is_ascii_digit() {
                go(s, N_EXP_MORE)
            } else {
                false
            }
        }
        N_EXP_FIRST => {
            if cap && byte.is_ascii_digit() {
                go(s, N_EXP_MORE)
            } else {
                false
            }
        }
        N_EXP_MORE if cap && byte.is_ascii_digit() => go(s, N_EXP_MORE),
        _ => false,
    }
}

#[allow(clippy::too_many_arguments)]
fn step_choice(
    t: &JsonSchemaTables,
    s: &mut FsmState,
    i: usize,
    node: u32,
    _phase: u32,
    a: u32,
    mask: u32,
    byte: u8,
) -> bool {
    let count = t.choice_count(node) as usize;
    let pos = a as usize;
    let mut newmask = 0u32;
    let mut done = false;
    for c in 0..count {
        if (mask >> c) & 1 == 0 {
            continue;
        }
        let lit = t.choice_lit(node, c);
        if pos < lit.len() && lit[pos] == byte {
            if pos + 1 == lit.len() {
                done = true; // fully matched (unique by the no-prefix invariant)
            } else {
                newmask |= 1 << c;
            }
        }
    }
    if done {
        s.pop();
        true
    } else if newmask != 0 {
        s.set_a(i, a + 1);
        s.set_b(i, newmask);
        true
    } else {
        false
    }
}

// -- direct allowed-mask computation (mirrors the step transitions arm-for-arm) --------------------

fn fill_top(t: &JsonSchemaTables, s: &FsmState, set: &mut [bool; 256]) {
    let d = s.depth();
    if d == 0 {
        return;
    }
    let i = d - 1;
    let node = s.node(i);
    let phase = s.phase(i);
    let a = s.a(i);
    let b = s.b(i);
    match t.kind(node) {
        NodeKind::Object => match phase {
            OBJ_OPEN => set[b'{' as usize] = true,
            OBJ_KEY => {
                let key = t.obj_key(node, a as usize);
                if (b as usize) < key.len() {
                    set[key[b as usize] as usize] = true;
                }
            }
            OBJ_COLON => set[b':' as usize] = true,
            OBJ_AFTER => {
                if a + 1 < t.obj_count(node) {
                    set[b',' as usize] = true;
                } else {
                    set[b'}' as usize] = true;
                }
            }
            OBJ_CLOSE_EMPTY => set[b'}' as usize] = true,
            _ => {}
        },
        NodeKind::Array => match phase {
            ARR_OPEN => set[b'[' as usize] = true,
            ARR_FIRST => {
                let (item, min, max) = (t.arr_item(node), t.arr_min(node), t.arr_max(node));
                if min == 0 {
                    set[b']' as usize] = true;
                }
                if max > 0 {
                    t.open_set(item, set);
                }
            }
            ARR_AFTER => {
                let (min, max) = (t.arr_min(node), t.arr_max(node));
                if a >= min {
                    set[b']' as usize] = true;
                }
                if a < max {
                    set[b',' as usize] = true;
                }
            }
            ARR_ITEM => t.open_set(t.arr_item(node), set),
            _ => {}
        },
        NodeKind::StringFree => fill_string(t, node, phase, a, set),
        NodeKind::Number | NodeKind::Integer => {
            fill_number(t.kind(node), phase, a, set);
            if is_number_terminal(phase) {
                // A completed open-ended scalar may hand the next byte to its parent.
                let mut parent = *s;
                parent.pop();
                fill_top(t, &parent, set);
            }
        }
        NodeKind::Choice => {
            let count = t.choice_count(node) as usize;
            let pos = a as usize;
            for c in 0..count {
                if (b >> c) & 1 == 0 {
                    continue;
                }
                let lit = t.choice_lit(node, c);
                if pos < lit.len() {
                    set[lit[pos] as usize] = true;
                }
            }
        }
    }
}

fn fill_string(t: &JsonSchemaTables, node: u32, phase: u32, a: u32, set: &mut [bool; 256]) {
    let mark = |set: &mut [bool; 256], lo: u8, hi: u8| {
        for c in lo..=hi {
            set[c as usize] = true;
        }
    };
    match phase {
        STR_OPEN => set[b'"' as usize] = true,
        STR_BODY => {
            let (min, max) = (t.str_min(node), t.str_max(node));
            if a >= min {
                set[b'"' as usize] = true;
            }
            if a < max {
                set[b'\\' as usize] = true;
                for c in 0x20u8..=0x7F {
                    if c != b'"' && c != b'\\' {
                        set[c as usize] = true;
                    }
                }
                mark(set, 0xC2, 0xDF);
                mark(set, 0xE0, 0xEF);
                mark(set, 0xF0, 0xF4);
            }
        }
        STR_C1 | STR_C2 | STR_C3 => mark(set, 0x80, 0xBF),
        STR_C2_E0 => mark(set, 0xA0, 0xBF),
        STR_C2_ED => mark(set, 0x80, 0x9F),
        STR_C3_F0 => mark(set, 0x90, 0xBF),
        STR_C3_F4 => mark(set, 0x80, 0x8F),
        STR_ESC => {
            for &c in b"\"\\/bfnrtu" {
                set[c as usize] = true;
            }
        }
        STR_U0 | STR_U1 | STR_U2 | STR_U3 | STR_SL2 | STR_SL3 => fill_hex(set),
        STR_SB => set[b'\\' as usize] = true,
        STR_SU => set[b'u' as usize] = true,
        STR_SL0 => {
            set[b'd' as usize] = true;
            set[b'D' as usize] = true;
        }
        STR_SL1 => {
            for &c in b"cdefCDEF" {
                set[c as usize] = true;
            }
        }
        _ => {}
    }
}

fn fill_hex(set: &mut [bool; 256]) {
    for c in b'0'..=b'9' {
        set[c as usize] = true;
    }
    for c in b'a'..=b'f' {
        set[c as usize] = true;
    }
    for c in b'A'..=b'F' {
        set[c as usize] = true;
    }
}

fn fill_number(kind: NodeKind, phase: u32, a: u32, set: &mut [bool; 256]) {
    let is_int = matches!(kind, NodeKind::Integer);
    let cap = a < NUMBER_BYTE_CAP;
    let digits = |set: &mut [bool; 256]| {
        for c in b'0'..=b'9' {
            set[c as usize] = true;
        }
    };
    match phase {
        N_START => {
            set[b'-' as usize] = true;
            digits(set);
        }
        N_INT_FIRST => digits(set),
        N_INT_ZERO => {
            if !is_int && cap {
                set[b'.' as usize] = true;
                set[b'e' as usize] = true;
                set[b'E' as usize] = true;
            }
        }
        N_INT_MORE => {
            if cap {
                digits(set);
                if !is_int {
                    set[b'.' as usize] = true;
                    set[b'e' as usize] = true;
                    set[b'E' as usize] = true;
                }
            }
        }
        N_FRAC_FIRST => {
            if cap {
                digits(set);
            }
        }
        N_FRAC_MORE => {
            if cap {
                digits(set);
                set[b'e' as usize] = true;
                set[b'E' as usize] = true;
            }
        }
        N_EXP_SIGN => {
            if cap {
                set[b'+' as usize] = true;
                set[b'-' as usize] = true;
                digits(set);
            }
        }
        N_EXP_FIRST | N_EXP_MORE if cap => {
            digits(set);
        }
        _ => {}
    }
}