inferencelayer 0.2.1

//! Flat, GPU-uploadable encoding of the G1 tables + the contract the G2 WGSL interpreter fills in.
//!
//! This module is the pure-Rust half of Workstream G2: it turns [`JsonSchemaTables`] and
//! [`TokenByteTable`] into the flat `u32` buffers the WGSL kernel uploads and indexes — mirroring the
//! `chars`/`offs` style of the OSFQL `GrammarTables`/`VocabTable` — and it documents (below) the
//! kernel + harness surface the fresh G2 session implements against these buffers. No WGSL and no
//! `wgpu` objects live here yet; keeping the encoding pure and unit-tested de-risks G2 (the GPU reads
//! exactly what the CPU FSM reads).
//!
//! ## What G2 adds (the `GpuGrammarHarness` contract)
//!
//! A device-side harness owning three storage buffers and three compute entry points, mirroring the
//! solo OSFQL `GpuGrammar` surface (`state_buf`/`mask_buf`/`chosen_buf`, `mask_workgroups`,
//! `make_mask_bg`/`make_advance_bg`) and producing an [`inferencelayer::SpecGrammar`] for the speculative
//! seam. Intended Rust API (the RED drive gate `tests/json_grammar_gpu.rs` is written against it):
//!
//! ```ignore
//! let h = GpuGrammarHarness::new(&ctx, &GpuSchema::from_tables(&t), &GpuVocab::from_table(&vocab))?;
//! let mask: Vec<u32> = h.mask_main(&ctx, state.as_words());       // per-token accept bits over the vocab
//! let next: [u32; 64] = h.advance_main(&ctx, state.as_words(), id); // advance the FSM by a chosen token
//! h.snap_advance_main(&ctx, slot, draft_id);                       // speculative snapshot+advance
//! ```
//!
//! WGSL entry points (names fixed by the task, matching the OSFQL kernels):
//! - `mask_main` — `@workgroup_size(256)`, grid `(mask_workgroups, 1)`. Cooperative-load the 64-word
//!   [`FsmState`] into `var<workgroup> sh_state: array<u32,64>` ONCE per workgroup (`if lid.x < 64u`,
//!   then `workgroupBarrier()`), then each thread walks its token's byte run — `vocab_bytes[off[id]
//!   .. off[id+1]]` — through the FSM `step_byte` transitions and writes `mask[id] = accept`. This is
//!   the [`FsmState`]/table port of the OSFQL `mask_main` at gpu_grammar.rs:711-751.
//! - `advance_main` — one thread; fold `step_byte` over the chosen token's bytes into `state_buf`.
//! - `snap_advance_main` — snapshot `state_buf` into slot `p` of a snapshot buffer, then advance by
//!   `drafts[p]`; drives [`inferencelayer::SpecGrammar`]'s `snap_adv_bgs[p]` (state stride = 64·4 bytes).
//!
//! WGSL gotchas carried from the codebase: `meta` is reserved (use `gm`); do NOT dynamically index
//! private arrays (Naga spills to DRAM — measured 4.2×); keep the 64-word state in `var<workgroup>`
//! and cooperative-load it; the module stays wasm32-clean (no fs/env).

use super::json_schema::JsonSchemaTables;
use super::token_bytes::TokenByteTable;

/// The `words`-buffer section bases (u32 INDICES into [`GpuSchema::words`], not byte offsets), plus
/// counts and the root node — the values the WGSL interpreter reads from its uniform to locate each
/// table. Every base is where that table's flat array begins inside the single `words` buffer.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SchemaMeta {
    /// Root node index (into the nodes table).
    pub root: u32,
    /// Base of the nodes table (always 0). Node `n`'s record = `words[nodes_base + n*NODE_WORDS ..]`.
    pub nodes_base: u32,
    /// Base of the props table. Property `p` = `words[props_base + p*2 ..]` = `[key_span, value_node]`.
    pub props_base: u32,
    /// Base of the choice-literals table (one span index per Choice member).
    pub choice_lits_base: u32,
    /// Base of the literal-spans table. Span `s` = `words[spans_base + s*2 ..]` = `[byte_off, len]`.
    pub spans_base: u32,
    /// Node count.
    pub n_nodes: u32,
    /// Property count (the props table holds `2 * n_props` words).
    pub n_props: u32,
    /// Choice-literal count.
    pub n_choice_lits: u32,
    /// Literal-span count (the spans table holds `2 * n_spans` words).
    pub n_spans: u32,
    /// Literal byte count (the `bytes` buffer length).
    pub n_bytes: u32,
}

impl SchemaMeta {
    /// The header as a flat `u32` run for upload as a uniform (11 words; pad to a 16-byte multiple on
    /// the device side as WGSL uniforms require).
    pub fn as_words(&self) -> [u32; 11] {
        [
            self.root,
            self.nodes_base,
            self.props_base,
            self.choice_lits_base,
            self.spans_base,
            self.n_nodes,
            self.n_props,
            self.n_choice_lits,
            self.n_spans,
            self.n_bytes,
            0,
        ]
    }
}

/// Flat, GPU-uploadable schema tables: the four `u32` sub-tables of [`JsonSchemaTables`] concatenated
/// into ONE `words` buffer (indexed via [`SchemaMeta`] bases), plus the literal bytes exploded to one
/// byte per `u32` in `bytes` for direct WGSL indexing (the OSFQL `chars`-buffer convention).
#[derive(Clone, Debug)]
pub struct GpuSchema {
    /// `[nodes.. | props.. | choice_lits.. | literal_spans..]` — one contiguous storage buffer.
    pub words: Vec<u32>,
    /// Literal bytes (property keys + Choice literals), one byte per `u32`, zero-extended.
    pub bytes: Vec<u32>,
    /// Section bases + counts + root for the WGSL uniform.
    pub meta: SchemaMeta,
}

impl GpuSchema {
    /// Pack compiled tables into the flat GPU layout. Lossless; see [`Self::nodes`] et al. for the
    /// read-back accessors the round-trip test checks against the CPU tables.
    pub fn from_tables(t: &JsonSchemaTables) -> Self {
        let nodes = t.raw_nodes();
        let props = t.raw_props();
        let choice_lits = t.raw_choice_lits();
        let spans = t.raw_literal_spans();
        let lit_bytes = t.raw_literal_bytes();

        let mut words =
            Vec::with_capacity(nodes.len() + props.len() + choice_lits.len() + spans.len());
        let nodes_base = words.len() as u32;
        words.extend_from_slice(nodes);
        let props_base = words.len() as u32;
        words.extend_from_slice(props);
        let choice_lits_base = words.len() as u32;
        words.extend_from_slice(choice_lits);
        let spans_base = words.len() as u32;
        words.extend_from_slice(spans);

        let bytes: Vec<u32> = lit_bytes.iter().map(|&b| b as u32).collect();
        let meta = SchemaMeta {
            root: t.root(),
            nodes_base,
            props_base,
            choice_lits_base,
            spans_base,
            n_nodes: (nodes.len() / JsonSchemaTables::NODE_WORDS) as u32,
            n_props: (props.len() / 2) as u32,
            n_choice_lits: choice_lits.len() as u32,
            n_spans: (spans.len() / 2) as u32,
            n_bytes: lit_bytes.len() as u32,
        };
        Self { words, bytes, meta }
    }

    /// Read-back of the nodes sub-table (for the round-trip check; the WGSL indexes `words` the same
    /// way via `meta.nodes_base`).
    pub fn nodes(&self) -> &[u32] {
        let a = self.meta.nodes_base as usize;
        let b = a + self.meta.n_nodes as usize * JsonSchemaTables::NODE_WORDS;
        &self.words[a..b]
    }
    /// Read-back of the props sub-table (`2 * n_props` words).
    pub fn props(&self) -> &[u32] {
        let a = self.meta.props_base as usize;
        &self.words[a..a + self.meta.n_props as usize * 2]
    }
    /// Read-back of the choice-literals sub-table.
    pub fn choice_lits(&self) -> &[u32] {
        let a = self.meta.choice_lits_base as usize;
        &self.words[a..a + self.meta.n_choice_lits as usize]
    }
    /// Read-back of the literal-spans sub-table (`2 * n_spans` words).
    pub fn spans(&self) -> &[u32] {
        let a = self.meta.spans_base as usize;
        &self.words[a..a + self.meta.n_spans as usize * 2]
    }
    /// The literal bytes for span `s` (`[byte_off, len]`), reconstructed from the flat buffers — the
    /// exact sequence the WGSL matches property keys / Choice literals against.
    pub fn literal(&self, span: u32) -> Vec<u8> {
        let base = self.meta.spans_base as usize + span as usize * 2;
        let off = self.words[base] as usize;
        let len = self.words[base + 1] as usize;
        self.bytes[off..off + len]
            .iter()
            .map(|&w| w as u8)
            .collect()
    }
}

/// Flat, GPU-uploadable token vocabulary: token `id`'s bytes = `bytes[off[id] .. off[id+1]]`, one
/// byte per `u32` (the OSFQL `VocabTable` `chars`/`off` convention). The `mask_main` kernel walks
/// each token's run in parallel.
#[derive(Clone, Debug)]
pub struct GpuVocab {
    /// All token bytes concatenated, one byte per `u32`.
    pub bytes: Vec<u32>,
    /// Per-token start offsets into `bytes`; length = token count + 1.
    pub off: Vec<u32>,
}

impl GpuVocab {
    /// Flatten a [`TokenByteTable`] into the GPU layout.
    pub fn from_table(table: &TokenByteTable) -> Self {
        let mut bytes = Vec::new();
        let mut off = vec![0u32];
        for id in 0..table.len() as u32 {
            let piece = table.bytes(id).unwrap_or(&[]);
            bytes.extend(piece.iter().map(|&b| b as u32));
            off.push(bytes.len() as u32);
        }
        Self { bytes, off }
    }

    /// Token count.
    pub fn len(&self) -> usize {
        self.off.len() - 1
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Token `id`'s raw bytes, reconstructed from the flat buffers (round-trip helper).
    pub fn piece(&self, id: u32) -> Vec<u8> {
        let a = self.off[id as usize] as usize;
        let b = self.off[id as usize + 1] as usize;
        self.bytes[a..b].iter().map(|&w| w as u8).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::grammar::json_schema::NodeKind;
    use serde_json::json;

    fn schemas() -> Vec<serde_json::Value> {
        vec![
            json!({ "type": "number" }),
            json!({ "type": "boolean" }),
            json!({ "enum": ["red", "green", "blue"] }),
            json!({ "type": "array", "items": { "type": "integer" }, "minItems": 1, "maxItems": 3 }),
            json!({
                "type": "object",
                "properties": { "name": { "type": "string" }, "age": { "type": "integer" } },
                "required": ["name", "age"],
                "additionalProperties": false
            }),
        ]
    }

    #[test]
    fn should_round_trip_flat_schema_encoding() {
        for schema in schemas() {
            let t = JsonSchemaTables::compile(&schema).expect("compile");
            let g = GpuSchema::from_tables(&t);
            // Lossless: each packed sub-table equals the source flat array.
            assert_eq!(g.nodes(), t.raw_nodes(), "nodes sub-table");
            assert_eq!(g.props(), t.raw_props(), "props sub-table");
            assert_eq!(
                g.choice_lits(),
                t.raw_choice_lits(),
                "choice_lits sub-table"
            );
            assert_eq!(g.spans(), t.raw_literal_spans(), "spans sub-table");
            assert_eq!(g.bytes.len(), t.raw_literal_bytes().len(), "bytes length");
            // Meta base arithmetic: reading node kinds through the flat layout matches the CPU tables.
            for n in 0..g.meta.n_nodes {
                let flat_kind =
                    g.words[g.meta.nodes_base as usize + n as usize * JsonSchemaTables::NODE_WORDS];
                assert_eq!(flat_kind, t.kind(n) as u32, "flat kind of node {n}");
            }
            // Every literal (property keys + Choice members) reconstructs byte-for-byte.
            for s in 0..g.meta.n_spans {
                assert_eq!(g.literal(s), t.literal_for_test(s), "literal span {s}");
            }
        }
    }

    #[test]
    fn should_encode_object_keys_in_required_order() {
        // The flat props preserve the required-array order that the FSM enforces.
        let t = JsonSchemaTables::compile(&json!({
            "type": "object",
            "properties": { "name": { "type": "string" }, "age": { "type": "integer" } },
            "required": ["name", "age"],
            "additionalProperties": false
        }))
        .unwrap();
        let g = GpuSchema::from_tables(&t);
        let root = g.meta.root;
        let kind =
            g.words[g.meta.nodes_base as usize + root as usize * JsonSchemaTables::NODE_WORDS];
        assert_eq!(kind, NodeKind::Object as u32);
        // property 0's key span → "name", property 1's → "age"
        let prop_start =
            g.words[g.meta.nodes_base as usize + root as usize * JsonSchemaTables::NODE_WORDS + 1];
        let key0_span = g.words[g.meta.props_base as usize + prop_start as usize * 2];
        let key1_span = g.words[g.meta.props_base as usize + (prop_start as usize + 1) * 2];
        assert_eq!(g.literal(key0_span), b"\"name\"");
        assert_eq!(g.literal(key1_span), b"\"age\"");
    }

    #[test]
    fn should_round_trip_flat_vocab_encoding() {
        let tok = r#"{ "model": { "vocab": { "ĠParis": 0, "Hello": 1, "Ġworld": 2 } } }"#;
        let table = TokenByteTable::from_tokenizer_json(tok.as_bytes()).expect("byte table");
        let g = GpuVocab::from_table(&table);
        assert_eq!(g.len(), table.len());
        for id in 0..table.len() as u32 {
            assert_eq!(g.piece(id), table.bytes(id).unwrap(), "vocab piece {id}");
        }
    }
}