inferencelayer 0.2.13

//! GPU-resident JSON-schema grammar: the WGSL transcription of [`flat_step_byte`] + a wgpu harness.
//!
//! The kernel [`GRAMMAR_WGSL`]'s `mask_main` is a language port of the proven flat-table interpreter
//! ([`super::flat_ref::flat_step_byte`]): identical 64-word state, identical `(kind, phase)`
//! dispatch, identical `words`/`bytes` index arithmetic. One thread per vocabulary token walks that
//! token's byte run through the automaton and writes an accept bit — the per-token grammar mask.
//! The state is cooperative-loaded into `var<workgroup>` once per workgroup (the OSFQL pattern), and
//! each thread threads its own `var<private>` working copy.
//!
//! Byte-exactness against the CPU FSM is the whole contract, gated in `tests/json_grammar_gpu.rs`.
//! Perf note: the private working state is dynamically indexed (the stack frame base is
//! `4 + 4·(depth-1)`), which Naga spills to scratch — correct but not yet fast; a bucketed /
//! first-char-prefiltered pass is the follow-up optimization (mirrors the OSFQL correctness-first,
//! then-optimize arc). `meta` is a reserved WGSL word, so the uniform is named `gm`.

use crate::forward::ShaderModuleTuned as _;
use super::harness::{GpuSchema, GpuVocab};
use crate::GpuCtx;
use crate::forward::uni;
use anyhow::Result;

/// Build a bind group from EXPLICIT `(binding, buffer)` pairs. Each grammar entry point uses a
/// different subset of the module's bindings, so the engine's positional `make_bg` (bindings 0..n
/// then the uniform) does not apply.
fn bg_from(
    ctx: &GpuCtx,
    pl: &wgpu::ComputePipeline,
    entries: &[(u32, &wgpu::Buffer)],
) -> wgpu::BindGroup {
    let e: Vec<wgpu::BindGroupEntry> = entries
        .iter()
        .map(|(b, buf)| wgpu::BindGroupEntry {
            binding: *b,
            resource: buf.as_entire_binding(),
        })
        .collect();
    ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &pl.get_bind_group_layout(0),
        entries: &e,
    })
}

/// Compile one entry point of the shared grammar module. The engine's `forward::pipeline` helper
/// hardcodes `main`; the grammar module carries several entries (`mask_main`, and later
/// `advance_main`/`snap_advance_main`), so it names the entry explicitly.
fn pipeline_entry(ctx: &GpuCtx, label: &str, src: &str, entry: &str) -> wgpu::ComputePipeline {
    let module = ctx
        .device
        .shader_module_tuned(wgpu::ShaderModuleDescriptor {
            label: Some(label),
            source: wgpu::ShaderSource::Wgsl(src.into()),
        });
    ctx.device
        .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(label),
            layout: None,
            module: &module,
            entry_point: Some(entry),
            compilation_options: Default::default(),
            cache: None,
        })
}

/// The fixed WGSL grammar interpreter. `mask_main` fills the per-token accept mask from a schema +
/// vocab + the current 64-word FSM state; the transitions mirror `flat_ref::flat_step_byte`.
pub const GRAMMAR_WGSL: &str = r#"
struct Meta {
  root: u32, nodes_base: u32, props_base: u32, choice_lits_base: u32, spans_base: u32,
  n_nodes: u32, n_props: u32, n_choice_lits: u32, n_spans: u32, n_bytes: u32,
  n_vocab: u32, pad: u32,
};
@group(0) @binding(0) var<storage, read>       words:    array<u32>;
@group(0) @binding(1) var<storage, read>       lbytes:   array<u32>;   // literal bytes, 1/word
@group(0) @binding(2) var<storage, read>       vbytes:   array<u32>;   // vocab bytes, 1/word
@group(0) @binding(3) var<storage, read>       voff:     array<u32>;   // per-token start offsets
@group(0) @binding(4) var<storage, read_write> state:    array<u32>;   // 64-word FSM state
@group(0) @binding(5) var<storage, read_write> mask:     array<u32>;   // per-token accept bit
@group(0) @binding(6) var<uniform>             gm:       Meta;
@group(0) @binding(7) var<storage, read>       params:   array<u32>;   // [token_id, snapshot slot]
@group(0) @binding(8) var<storage, read_write> status:   array<u32>;   // [accepted]
@group(0) @binding(9) var<storage, read_write> snap:     array<u32>;   // [slots × 64] snapshots
@group(0) @binding(10) var<storage, read_write> firstb:  array<u32>;   // [256] first-byte accept

var<workgroup> sh_state: array<u32, 64>;
var<private>   st:       array<u32, 64>;

// -- flat table reads (mirror flat_ref::Flat) --
fn kind(node: u32) -> u32 { return words[gm.nodes_base + node * 5u]; }
fn field(node: u32, k: u32) -> u32 { return words[gm.nodes_base + node * 5u + 1u + k]; }
fn span_len(span: u32) -> u32 { return words[gm.spans_base + span * 2u + 1u]; }
fn span_byte(span: u32, j: u32) -> u32 {
  let off = words[gm.spans_base + span * 2u];
  return lbytes[off + j];
}
fn obj_count(node: u32) -> u32 { return field(node, 1u); }
fn obj_value(node: u32, i: u32) -> u32 {
  let base = gm.props_base + (field(node, 0u) + i) * 2u;
  return words[base + 1u];
}
fn obj_key_span(node: u32, i: u32) -> u32 {
  return words[gm.props_base + (field(node, 0u) + i) * 2u];
}
fn arr_item(node: u32) -> u32 { return field(node, 0u); }
fn arr_min(node: u32) -> u32 { return field(node, 1u); }
fn arr_max(node: u32) -> u32 { return field(node, 2u); }
fn str_min(node: u32) -> u32 { return field(node, 0u); }
fn str_max(node: u32) -> u32 { return field(node, 1u); }
fn choice_count(node: u32) -> u32 { return field(node, 1u); }
fn choice_span(node: u32, c: u32) -> u32 {
  return words[gm.choice_lits_base + field(node, 0u) + c];
}

// -- state helpers (st) --
fn push(node: u32, phase: u32, a: u32, b: u32) -> bool {
  let d = st[0];
  if (d >= 15u) { return false; }
  let base = 4u + d * 4u;
  st[base] = node; st[base + 1u] = phase; st[base + 2u] = a; st[base + 3u] = b;
  st[0] = d + 1u;
  return true;
}
fn pop() { let d = st[0]; if (d > 0u) { st[0] = d - 1u; } }

fn is_hex(c: u32) -> bool {
  return (c >= 0x30u && c <= 0x39u) || (c >= 0x41u && c <= 0x46u) || (c >= 0x61u && c <= 0x66u);
}
fn hexval(c: u32) -> u32 {
  if (c >= 0x30u && c <= 0x39u) { return c - 0x30u; }
  if (c >= 0x41u && c <= 0x46u) { return c - 0x41u + 10u; }
  return c - 0x61u + 10u;
}
fn is_digit(c: u32) -> bool { return c >= 0x30u && c <= 0x39u; }
fn is_num_terminal(phase: u32) -> bool {
  return phase == 52u || phase == 53u || phase == 55u || phase == 58u; // ZERO/INT_MORE/FRAC_MORE/EXP_MORE
}
fn number_extends(is_int: bool, phase: u32, a: u32, byte: u32) -> bool {
  if (a >= 24u) { return false; }
  if (phase == 52u) { return (!is_int) && (byte == 0x2Eu || byte == 0x65u || byte == 0x45u); }
  if (phase == 53u) { return is_digit(byte) || ((!is_int) && (byte == 0x2Eu || byte == 0x65u || byte == 0x45u)); }
  if (phase == 55u) { return is_digit(byte) || byte == 0x65u || byte == 0x45u; }
  if (phase == 58u) { return is_digit(byte); }
  return false;
}

// fresh (phase,a,b) for a new frame; returns phase in .x, a in .y, b in .z.
fn fresh(node: u32) -> vec3<u32> {
  let k = kind(node);
  if (k == 0u) { return vec3<u32>(0u, 0u, 0u); }   // OBJ_OPEN
  if (k == 1u) { return vec3<u32>(10u, 0u, 0u); }  // ARR_OPEN
  if (k == 2u) { return vec3<u32>(20u, 0u, 0u); }  // STR_OPEN
  if (k == 3u || k == 4u) { return vec3<u32>(50u, 0u, 0u); } // N_START
  // Choice: CH_MATCH, a=0, b=all-viable bitmask
  let count = choice_count(node);
  var mask_b = 0u;
  if (count >= 32u) { mask_b = 0xFFFFFFFFu; } else { mask_b = (1u << count) - 1u; }
  return vec3<u32>(60u, 0u, mask_b);
}

// WGSL has no forward declarations: these two must precede step_string, their only caller.
fn cont(base: u32, byte: u32, lo: u32, hi: u32, next: u32, count: bool, a: u32) -> bool {
  if (byte >= lo && byte <= hi) { st[base + 1u] = next; if (count) { st[base + 2u] = a + 1u; } return true; }
  return false;
}
fn set_phase_if(base: u32, ok: bool, next: u32) -> bool {
  if (ok) { st[base + 1u] = next; }
  return ok;
}

// step_string: mutate st for frame i; return accept. Mirrors flat_ref::step_string.
fn step_string(i: u32, node: u32, phase: u32, a: u32, b: u32, byte: u32) -> bool {
  let base = 4u + i * 4u;
  if (phase == 20u) { // STR_OPEN
    if (byte == 0x22u) { st[base + 1u] = 21u; st[base + 2u] = 0u; return true; }
    return false;
  }
  if (phase == 21u) { // STR_BODY
    let mn = str_min(node); let mx = str_max(node);
    if (byte == 0x22u) { if (a >= mn) { pop(); return true; } return false; }
    if (byte == 0x5Cu) { if (a < mx) { st[base + 1u] = 29u; return true; } return false; }
    if (a >= mx) { return false; }
    if (byte >= 0x20u && byte <= 0x7Fu && byte != 0x22u && byte != 0x5Cu) { st[base + 2u] = a + 1u; return true; }
    var nx = 0u;
    if (byte >= 0xC2u && byte <= 0xDFu) { nx = 22u; }
    else if (byte == 0xE0u) { nx = 24u; }
    else if (byte == 0xEDu) { nx = 25u; }
    else if ((byte >= 0xE1u && byte <= 0xECu) || (byte >= 0xEEu && byte <= 0xEFu)) { nx = 23u; }
    else if (byte == 0xF0u) { nx = 27u; }
    else if (byte == 0xF4u) { nx = 28u; }
    else if (byte >= 0xF1u && byte <= 0xF3u) { nx = 26u; }
    else { return false; }
    st[base + 1u] = nx;
    return true;
  }
  // UTF-8 continuation phases: (lo,hi,next,count)
  if (phase == 22u) { return cont(base, byte, 0x80u, 0xBFu, 21u, true, a); }  // C1 -> BODY, count
  if (phase == 23u) { return cont(base, byte, 0x80u, 0xBFu, 22u, false, a); } // C2 -> C1
  if (phase == 24u) { return cont(base, byte, 0xA0u, 0xBFu, 22u, false, a); } // C2_E0
  if (phase == 25u) { return cont(base, byte, 0x80u, 0x9Fu, 22u, false, a); } // C2_ED
  if (phase == 26u) { return cont(base, byte, 0x80u, 0xBFu, 23u, false, a); } // C3 -> C2
  if (phase == 27u) { return cont(base, byte, 0x90u, 0xBFu, 23u, false, a); } // C3_F0
  if (phase == 28u) { return cont(base, byte, 0x80u, 0x8Fu, 23u, false, a); } // C3_F4
  if (phase == 29u) { // STR_ESC
    if (byte == 0x22u || byte == 0x5Cu || byte == 0x2Fu || byte == 0x62u || byte == 0x66u || byte == 0x6Eu || byte == 0x72u || byte == 0x74u) {
      st[base + 1u] = 21u; st[base + 2u] = a + 1u; return true;
    }
    if (byte == 0x75u) { st[base + 1u] = 30u; st[base + 3u] = 0u; return true; }
    return false;
  }
  if (phase == 30u || phase == 31u || phase == 32u) { // U0/U1/U2
    if (is_hex(byte)) {
      var shift = 4u;
      if (phase == 30u) { shift = 12u; } else if (phase == 31u) { shift = 8u; }
      st[base + 3u] = b | (hexval(byte) << shift);
      st[base + 1u] = phase + 1u;
      return true;
    }
    return false;
  }
  if (phase == 33u) { // U3
    if (is_hex(byte)) {
      let cu = b | hexval(byte);
      if (cu >= 0xD800u && cu <= 0xDBFFu) { st[base + 1u] = 34u; return true; }
      if (cu >= 0xDC00u && cu <= 0xDFFFu) { return false; }
      st[base + 1u] = 21u; st[base + 2u] = a + 1u; st[base + 3u] = 0u; return true;
    }
    return false;
  }
  if (phase == 34u) { return set_phase_if(base, byte == 0x5Cu, 35u); } // SB -> SU on '\'
  if (phase == 35u) { return set_phase_if(base, byte == 0x75u, 36u); } // SU -> SL0 on 'u'
  if (phase == 36u) { return set_phase_if(base, byte == 0x64u || byte == 0x44u, 37u); } // SL0: d/D
  if (phase == 37u) { return set_phase_if(base, (byte >= 0x63u && byte <= 0x66u) || (byte >= 0x43u && byte <= 0x46u), 38u); } // SL1: c..f/C..F
  if (phase == 38u) { return set_phase_if(base, is_hex(byte), 39u); } // SL2
  if (phase == 39u) { // SL3
    if (is_hex(byte)) { st[base + 1u] = 21u; st[base + 2u] = a + 1u; st[base + 3u] = 0u; return true; }
    return false;
  }
  return false;
}

fn step_number(i: u32, is_int: bool, phase: u32, a: u32, byte: u32) -> bool {
  let base = 4u + i * 4u;
  let na = a + 1u;
  let cap = a < 24u;
  // go(next): set phase+a, return true
  if (phase == 50u) { // N_START
    if (byte == 0x2Du) { st[base+1u]=51u; st[base+2u]=na; return true; }
    if (byte == 0x30u) { st[base+1u]=52u; st[base+2u]=na; return true; }
    if (byte >= 0x31u && byte <= 0x39u) { st[base+1u]=53u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 51u) { // N_INT_FIRST
    if (byte == 0x30u) { st[base+1u]=52u; st[base+2u]=na; return true; }
    if (byte >= 0x31u && byte <= 0x39u) { st[base+1u]=53u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 52u) { // N_INT_ZERO
    if ((!is_int) && cap && byte == 0x2Eu) { st[base+1u]=54u; st[base+2u]=na; return true; }
    if ((!is_int) && cap && (byte == 0x65u || byte == 0x45u)) { st[base+1u]=56u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 53u) { // N_INT_MORE
    if (cap && is_digit(byte)) { st[base+1u]=53u; st[base+2u]=na; return true; }
    if ((!is_int) && cap && byte == 0x2Eu) { st[base+1u]=54u; st[base+2u]=na; return true; }
    if ((!is_int) && cap && (byte == 0x65u || byte == 0x45u)) { st[base+1u]=56u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 54u) { // N_FRAC_FIRST
    if (cap && is_digit(byte)) { st[base+1u]=55u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 55u) { // N_FRAC_MORE
    if (cap && is_digit(byte)) { st[base+1u]=55u; st[base+2u]=na; return true; }
    if (cap && (byte == 0x65u || byte == 0x45u)) { st[base+1u]=56u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 56u) { // N_EXP_SIGN
    if (cap && (byte == 0x2Bu || byte == 0x2Du)) { st[base+1u]=57u; st[base+2u]=na; return true; }
    if (cap && is_digit(byte)) { st[base+1u]=58u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 57u) { // N_EXP_FIRST
    if (cap && is_digit(byte)) { st[base+1u]=58u; st[base+2u]=na; return true; }
    return false;
  }
  if (phase == 58u) { // N_EXP_MORE
    if (cap && is_digit(byte)) { st[base+1u]=58u; st[base+2u]=na; return true; }
    return false;
  }
  return false;
}

fn step_choice(i: u32, node: u32, a: u32, mask_in: u32, byte: u32) -> bool {
  let base = 4u + i * 4u;
  let count = choice_count(node);
  var newmask = 0u;
  var done = false;
  for (var c = 0u; c < count; c = c + 1u) {
    if (((mask_in >> c) & 1u) == 0u) { continue; }
    let span = choice_span(node, c);
    let len = span_len(span);
    if (a < len && span_byte(span, a) == byte) {
      if (a + 1u == len) { done = true; } else { newmask = newmask | (1u << c); }
    }
  }
  if (done) { pop(); return true; }
  if (newmask != 0u) { st[base + 2u] = a + 1u; st[base + 3u] = newmask; return true; }
  return false;
}

// One byte transition over the current state; mirrors flat_ref::flat_step_byte's loop.
fn gstep(byte: u32) -> bool {
  for (var iter = 0u; iter < 128u; iter = iter + 1u) {
    let d = st[0];
    if (d == 0u) { return false; }
    let i = d - 1u;
    let base = 4u + i * 4u;
    let node = st[base];
    let phase = st[base + 1u];
    let a = st[base + 2u];
    let b = st[base + 3u];
    let k = kind(node);
    if (k == 0u) { // Object
      if (phase == 0u) { if (byte == 0x7Bu) { st[base+1u]=1u; return true; } return false; } // OBJ_OPEN expect {
      if (phase == 1u) { // OBJ_FIRST (structural)
        if (obj_count(node) == 0u) { st[base+1u]=6u; } else { st[base+1u]=2u; st[base+2u]=0u; st[base+3u]=0u; }
        continue;
      }
      if (phase == 4u) { // OBJ_VALUE (structural: push child)
        let vnode = obj_value(node, a);
        st[base+1u]=5u;
        let fr = fresh(vnode);
        if (!push(vnode, fr.x, fr.y, fr.z)) { return false; }
        continue;
      }
      if (phase == 2u) { // OBJ_KEY
        let span = obj_key_span(node, a);
        let klen = span_len(span);
        if (b < klen && span_byte(span, b) == byte) {
          if (b + 1u == klen) { st[base+1u]=3u; st[base+3u]=0u; } else { st[base+3u]=b+1u; }
          return true;
        }
        return false;
      }
      if (phase == 3u) { if (byte == 0x3Au) { st[base+1u]=4u; return true; } return false; } // OBJ_COLON expect :
      if (phase == 5u) { // OBJ_AFTER
        let count = obj_count(node);
        if (a + 1u < count) {
          if (byte == 0x2Cu) { st[base+2u]=a+1u; st[base+1u]=2u; st[base+3u]=0u; return true; }
        } else if (byte == 0x7Du) { pop(); return true; }
        return false;
      }
      if (phase == 6u) { if (byte == 0x7Du) { pop(); return true; } return false; } // OBJ_CLOSE_EMPTY
      return false;
    } else if (k == 1u) { // Array
      if (phase == 10u) { if (byte == 0x5Bu) { st[base+1u]=11u; st[base+2u]=0u; return true; } return false; } // ARR_OPEN
      if (phase == 11u) { // ARR_FIRST
        let item = arr_item(node); let mn = arr_min(node); let mx = arr_max(node);
        if (byte == 0x5Du) { if (mn == 0u) { pop(); return true; } return false; }
        if (mx > 0u) {
          st[base+1u]=12u; st[base+2u]=1u;
          let fr = fresh(item);
          if (!push(item, fr.x, fr.y, fr.z)) { return false; }
          continue;
        }
        return false;
      }
      if (phase == 12u) { // ARR_AFTER
        let mn = arr_min(node); let mx = arr_max(node);
        if (byte == 0x5Du) { if (a >= mn) { pop(); return true; } return false; }
        if (byte == 0x2Cu && a < mx) { st[base+1u]=13u; return true; }
        return false;
      }
      if (phase == 13u) { // ARR_ITEM
        let item = arr_item(node);
        st[base+1u]=12u; st[base+2u]=a+1u;
        let fr = fresh(item);
        if (!push(item, fr.x, fr.y, fr.z)) { return false; }
        continue;
      }
      return false;
    } else if (k == 2u) {
      return step_string(i, node, phase, a, b, byte);
    } else if (k == 3u || k == 4u) {
      let is_int = (k == 4u);
      if (is_num_terminal(phase) && !number_extends(is_int, phase, a, byte)) { pop(); continue; }
      return step_number(i, is_int, phase, a, byte);
    } else if (k == 5u) {
      return step_choice(i, node, a, b, byte);
    }
    return false;
  }
  return false;
}

// FIRST-BYTE PREFILTER: one thread per byte value tests whether the CURRENT state accepts it.
// A token whose first byte is rejected here is rejected, period (step_token folds step_byte), so
// mask_main can discard it after a single load — no 64-word state copy, no walk. This is exact,
// not an approximation: it is precisely step_byte's decision on the token's first byte.
@compute @workgroup_size(256)
fn firstb_main(@builtin(local_invocation_id) lid: vec3<u32>) {
  let byte = lid.x;
  for (var w = 0u; w < 64u; w = w + 1u) { st[w] = state[w]; }
  firstb[byte] = select(0u, 1u, gstep(byte));
}

@compute @workgroup_size(256)
fn mask_main(@builtin(global_invocation_id) gid: vec3<u32>,
             @builtin(local_invocation_id) lid: vec3<u32>) {
  if (lid.x < 64u) { sh_state[lid.x] = state[lid.x]; }
  workgroupBarrier();
  let id = gid.x;
  if (id >= gm.n_vocab) { return; }
  let lo = voff[id];
  let hi = voff[id + 1u];
  // Empty piece: no bytes to walk, so the state is unchanged and the token is trivially allowed
  // (matches step_token over an empty slice).
  if (lo >= hi) { mask[id] = 1u; return; }
  // Prefilter: reject on the first byte without touching the (scratch-spilling) private state.
  if (firstb[vbytes[lo]] == 0u) { mask[id] = 0u; return; }
  for (var w = 0u; w < 64u; w = w + 1u) { st[w] = sh_state[w]; }
  var ok = true;
  for (var p = lo; p < hi; p = p + 1u) {
    if (!gstep(vbytes[p])) { ok = false; break; }
  }
  mask[id] = select(0u, 1u, ok);
}

// Walk `params[0]`'s bytes from the current state into `st`; report acceptance in status[0] and
// commit back to `state` only on success (a rejected token must never corrupt the live state).
fn walk_and_commit() {
  for (var w = 0u; w < 64u; w = w + 1u) { st[w] = state[w]; }
  let id = params[0];
  let lo = voff[id];
  let hi = voff[id + 1u];
  var ok = true;
  for (var p = lo; p < hi; p = p + 1u) {
    if (!gstep(vbytes[p])) { ok = false; break; }
  }
  status[0] = select(0u, 1u, ok);
  if (ok) {
    for (var w = 0u; w < 64u; w = w + 1u) { state[w] = st[w]; }
  }
}

// Advance the live state by the chosen token (params[0]).
@compute @workgroup_size(1)
fn advance_main() {
  walk_and_commit();
}

// Speculative seam: snapshot the live state into slot params[1], then advance by the draft token
// params[0]. Restoring a snapshot is a host-side buffer copy (no kernel needed).
@compute @workgroup_size(1)
fn snap_advance_main() {
  let slot = params[1];
  for (var w = 0u; w < 64u; w = w + 1u) { snap[slot * 64u + w] = state[w]; }
  walk_and_commit();
}
"#;

/// The three grammar pipelines, compiled ONCE per device. Shader compilation costs 150-450 ms, so
/// a per-request compile would be a non-starter for serving: build this once, then stamp out a
/// [`GpuGrammar`] per request (buffers + bind groups only).
pub struct GrammarKernels {
    firstb_pl: wgpu::ComputePipeline,
    mask_pl: wgpu::ComputePipeline,
    adv_pl: wgpu::ComputePipeline,
    snap_pl: wgpu::ComputePipeline,
}

impl GrammarKernels {
    /// Compile the grammar module's entry points.
    pub fn new(ctx: &GpuCtx) -> Self {
        Self {
            firstb_pl: pipeline_entry(ctx, "json_firstb", GRAMMAR_WGSL, "firstb_main"),
            mask_pl: pipeline_entry(ctx, "json_mask_main", GRAMMAR_WGSL, "mask_main"),
            adv_pl: pipeline_entry(ctx, "json_advance_main", GRAMMAR_WGSL, "advance_main"),
            snap_pl: pipeline_entry(ctx, "json_snap_advance", GRAMMAR_WGSL, "snap_advance_main"),
        }
    }
}

/// The tokenizer's byte table on-device. Fixed per model, so it is uploaded ONCE and shared by every
/// request's grammar (a 150k-token blob per request would dwarf the schema itself).
pub struct GrammarVocab {
    n_vocab: u32,
    vbytes: wgpu::Buffer,
    voff: wgpu::Buffer,
}

impl GrammarVocab {
    /// Upload a flat vocabulary.
    pub fn new(ctx: &GpuCtx, vocab: &GpuVocab) -> Result<Self> {
        let n_vocab = vocab.len() as u32;
        anyhow::ensure!(n_vocab > 0, "empty vocabulary");
        Ok(Self {
            n_vocab,
            vbytes: ctx.storage_bytes(bytemuck::cast_slice(&pad1(&vocab.bytes))),
            voff: ctx.storage_bytes(bytemuck::cast_slice(&vocab.off)),
        })
    }

    /// Token count.
    pub fn len(&self) -> u32 {
        self.n_vocab
    }

    /// Whether the vocabulary is empty (never true — the constructor refuses it).
    pub fn is_empty(&self) -> bool {
        self.n_vocab == 0
    }
}

/// One request's device-side JSON-schema grammar: the compiled schema, the live 64-word FSM state,
/// and bind groups over the SHARED [`GrammarKernels`] + [`GrammarVocab`]. Exposes [`Self::mask`]
/// (per-token accept bits), [`Self::advance`] (commit a chosen token), and
/// [`Self::snap_advance`]/[`Self::restore`] (the speculative seam) — all byte-exact with the CPU FSM.
pub struct GpuGrammar {
    n_vocab: u32,
    slots: u32,
    state: wgpu::Buffer,
    mask: wgpu::Buffer,
    params: wgpu::Buffer,
    status: wgpu::Buffer,
    snap: wgpu::Buffer,
    firstb_pl: wgpu::ComputePipeline,
    firstb_bg: wgpu::BindGroup,
    mask_pl: wgpu::ComputePipeline,
    mask_bg: wgpu::BindGroup,
    adv_pl: wgpu::ComputePipeline,
    adv_bg: wgpu::BindGroup,
    snap_pl: wgpu::ComputePipeline,
    snap_bg: wgpu::BindGroup,
}

impl GpuGrammar {
    /// Default snapshot slots (the speculative draft depth the seam supports).
    pub const DEFAULT_SLOTS: u32 = 16;

    /// Build a per-request grammar over the shared kernels + vocab (uploads only the schema).
    pub fn new(
        ctx: &GpuCtx,
        kernels: &GrammarKernels,
        vocab: &GrammarVocab,
        schema: &GpuSchema,
    ) -> Result<Self> {
        Self::new_with_slots(ctx, kernels, vocab, schema, Self::DEFAULT_SLOTS)
    }

    /// [`Self::new`] with an explicit snapshot-slot count (speculative draft depth).
    pub fn new_with_slots(
        ctx: &GpuCtx,
        kernels: &GrammarKernels,
        vocab: &GrammarVocab,
        schema: &GpuSchema,
        slots: u32,
    ) -> Result<Self> {
        anyhow::ensure!(slots > 0, "at least one snapshot slot");
        let n_vocab = vocab.n_vocab;
        let words = ctx.storage_bytes(bytemuck::cast_slice(&pad1(&schema.words)));
        let lbytes = ctx.storage_bytes(bytemuck::cast_slice(&pad1(&schema.bytes)));
        let state = ctx.storage_bytes(bytemuck::cast_slice(&[0u32; 64]));
        let mask = ctx.empty(n_vocab as usize);
        let params = ctx.storage_bytes(bytemuck::cast_slice(&[0u32; 2]));
        let status = ctx.storage_bytes(bytemuck::cast_slice(&[0u32; 1]));
        let snap = ctx.empty(slots as usize * 64);
        let firstb = ctx.empty(256);
        let m = &schema.meta;
        let gm = uni(
            ctx,
            bytemuck::cast_slice(&[
                m.root,
                m.nodes_base,
                m.props_base,
                m.choice_lits_base,
                m.spans_base,
                m.n_nodes,
                m.n_props,
                m.n_choice_lits,
                m.n_spans,
                m.n_bytes,
                n_vocab,
                0u32,
            ]),
        );
        let (vbytes, voff) = (&vocab.vbytes, &vocab.voff);

        // wgpu resources are refcounted, so the bind groups keep the shared vocab buffers alive.
        let firstb_bg = bg_from(
            ctx,
            &kernels.firstb_pl,
            &[
                (0, &words),
                (1, &lbytes),
                (4, &state),
                (6, &gm),
                (10, &firstb),
            ],
        );
        let mask_bg = bg_from(
            ctx,
            &kernels.mask_pl,
            &[
                (0, &words),
                (1, &lbytes),
                (2, vbytes),
                (3, voff),
                (4, &state),
                (5, &mask),
                (6, &gm),
                (10, &firstb),
            ],
        );
        let adv_bg = bg_from(
            ctx,
            &kernels.adv_pl,
            &[
                (0, &words),
                (1, &lbytes),
                (2, vbytes),
                (3, voff),
                (4, &state),
                (6, &gm),
                (7, &params),
                (8, &status),
            ],
        );
        let snap_bg = bg_from(
            ctx,
            &kernels.snap_pl,
            &[
                (0, &words),
                (1, &lbytes),
                (2, vbytes),
                (3, voff),
                (4, &state),
                (6, &gm),
                (7, &params),
                (8, &status),
                (9, &snap),
            ],
        );
        Ok(Self {
            n_vocab,
            slots,
            state,
            mask,
            params,
            status,
            snap,
            firstb_pl: kernels.firstb_pl.clone(),
            firstb_bg,
            mask_pl: kernels.mask_pl.clone(),
            mask_bg,
            adv_pl: kernels.adv_pl.clone(),
            adv_bg,
            snap_pl: kernels.snap_pl.clone(),
            snap_bg,
        })
    }

    /// Overwrite the device FSM state (the 64 words of [`super::FsmState::as_words`]).
    pub fn set_state(&self, ctx: &GpuCtx, words: &[u32; 64]) {
        ctx.queue
            .write_buffer(&self.state, 0, bytemuck::cast_slice(words));
    }

    /// Read the live device FSM state back (the same 64 words the CPU threads).
    pub fn read_state(&self, ctx: &GpuCtx) -> Result<[u32; 64]> {
        let v = ctx.read_u32(&self.state, 64)?;
        let mut w = [0u32; 64];
        w.copy_from_slice(&v);
        Ok(w)
    }

    /// Run `mask_main` and read back the per-token accept bits (`1` = the grammar accepts stepping
    /// this token from the current state, `0` = rejected).
    /// Compute the mask on-device WITHOUT reading it back — the GPU-resident path, where the mask
    /// feeds a mask-apply kernel instead of the host. Blocks until the work completes so callers
    /// can time the kernel alone.
    pub fn mask_dispatch(&self, ctx: &GpuCtx) -> Result<()> {
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        {
            let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            p.set_pipeline(&self.firstb_pl);
            p.set_bind_group(0, &self.firstb_bg, &[]);
            p.dispatch_workgroups(1, 1, 1);
            p.set_pipeline(&self.mask_pl);
            p.set_bind_group(0, &self.mask_bg, &[]);
            p.dispatch_workgroups(self.n_vocab.div_ceil(256), 1, 1);
        }
        ctx.queue.submit([enc.finish()]);
        ctx.device
            .poll(wgpu::PollType::wait_indefinitely())
            .map_err(|e| anyhow::anyhow!("poll: {e:?}"))?;
        Ok(())
    }

    pub fn mask(&self, ctx: &GpuCtx) -> Result<Vec<u32>> {
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        {
            // One pass: the 256-thread first-byte prefilter, then the per-token walk that reads
            // it. Dispatches within a pass see each other's writes in order.
            let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            p.set_pipeline(&self.firstb_pl);
            p.set_bind_group(0, &self.firstb_bg, &[]);
            p.dispatch_workgroups(1, 1, 1);
            p.set_pipeline(&self.mask_pl);
            p.set_bind_group(0, &self.mask_bg, &[]);
            p.dispatch_workgroups(self.n_vocab.div_ceil(256), 1, 1);
        }
        ctx.queue.submit([enc.finish()]);
        ctx.read_u32(&self.mask, self.n_vocab as usize)
    }

    /// Advance the live state by `token`. Returns whether the grammar accepted it; on rejection the
    /// device state is left untouched (a bad token can never corrupt the live automaton).
    pub fn advance(&self, ctx: &GpuCtx, token: u32) -> Result<bool> {
        self.run_one(ctx, &self.adv_pl, &self.adv_bg, token, 0)
    }

    /// Snapshot the live state into `slot`, then advance by the draft `token` — the speculative
    /// seam. Pair with [`Self::restore`] to roll back a rejected draft.
    pub fn snap_advance(&self, ctx: &GpuCtx, slot: u32, token: u32) -> Result<bool> {
        anyhow::ensure!(slot < self.slots, "snapshot slot {slot} out of range");
        self.run_one(ctx, &self.snap_pl, &self.snap_bg, token, slot)
    }

    /// Restore the live state from snapshot `slot` (a device-side buffer copy; no kernel).
    pub fn restore(&self, ctx: &GpuCtx, slot: u32) -> Result<()> {
        anyhow::ensure!(slot < self.slots, "snapshot slot {slot} out of range");
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        enc.copy_buffer_to_buffer(&self.snap, slot as u64 * 64 * 4, &self.state, 0, 64 * 4);
        ctx.queue.submit([enc.finish()]);
        Ok(())
    }

    /// Shared driver for the two single-thread kernels: upload `[token, slot]`, dispatch, read the
    /// accept flag.
    fn run_one(
        &self,
        ctx: &GpuCtx,
        pl: &wgpu::ComputePipeline,
        bg: &wgpu::BindGroup,
        token: u32,
        slot: u32,
    ) -> Result<bool> {
        anyhow::ensure!(token < self.n_vocab, "token id {token} out of range");
        ctx.queue
            .write_buffer(&self.params, 0, bytemuck::cast_slice(&[token, slot]));
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        {
            let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            p.set_pipeline(pl);
            p.set_bind_group(0, bg, &[]);
            p.dispatch_workgroups(1, 1, 1);
        }
        ctx.queue.submit([enc.finish()]);
        Ok(ctx.read_u32(&self.status, 1)?[0] == 1)
    }
}

/// Ensure a `u32` slice is non-empty for buffer upload (an empty storage buffer is invalid); a
/// zero pad is never indexed (`voff`/counts bound every read).
fn pad1(v: &[u32]) -> Vec<u32> {
    if v.is_empty() { vec![0u32] } else { v.to_vec() }
}