inferencelayer 0.2.4

//! JSON-Schema → flat table compiler (the "program" the byte-FSM interprets).
//!
//! A schema is compiled ONCE per request (pure Rust, microseconds) into a set of flat `u32`/`u8`
//! arrays — the exact buffers the future WGSL interpreter (G2) will upload and read. The CPU FSM in
//! [`super::json_fsm`] and the GPU kernel therefore execute the SAME tables; anything that cannot be
//! encoded in these tables surfaces HERE as a compile-time [`anyhow::Error`] naming the construct,
//! which is the whole de-risking purpose of G1.
//!
//! Supported subset = **OpenAI strict mode**: `object`/`array`/`string`/`number`/`integer`/
//! `boolean`/`null`; `properties`+`required` (every property required, matched in the `required`
//! array's order); `additionalProperties:false` mandatory; `enum`/`const` over scalars; bounded
//! arrays (`minItems`/`maxItems` ≤ 255) and strings (`minLength`/`maxLength`); nesting depth ≤ 10.
//! `pattern`/`anyOf`/`oneOf`/`allOf`/`$ref`/`format`/`minimum`/… are refused loudly.
//!
//! ## Table format (read this for G2)
//!
//! - `nodes: Vec<u32>` — stride [`Self::NODE_WORDS`] = 5: `[kind, f1, f2, f3, f4]`. `kind` is a
//!   [`NodeKind`] discriminant. Field meaning by kind:
//!   - `Object`  → `f1 = prop_start` (property index into `props`), `f2 = prop_count`.
//!   - `Array`   → `f1 = item_node`, `f2 = min_items`, `f3 = max_items` ([`Self::UNBOUNDED`] = none).
//!   - `String`  → `f1 = min_len`, `f2 = max_len` ([`Self::UNBOUNDED`] = none).
//!   - `Number`/`Integer` → no fields.
//!   - `Choice`  → `f1 = lit_start` (index into `choice_lits`), `f2 = lit_count` (≤ 32).
//! - `props: Vec<u32>` — stride 2: `[key_span, value_node]` per property, in `required` order.
//! - `choice_lits: Vec<u32>` — one span index per Choice member.
//! - `literal_spans: Vec<u32>` — stride 2: `[offset, len]` into `literal_bytes`.
//! - `literal_bytes: Vec<u8>` — the pooled bytes of every property key and Choice literal (keys and
//!   scalar members are stored in their CANONICAL JSON serialization, quotes included).
//! - `root: u32` — the root node index.

use anyhow::{Result, bail};
use serde_json::Value;

/// Node kinds present in a compiled table. `boolean`/`null`/`const`/`enum` all lower to
/// [`NodeKind::Choice`] (match one of N byte literals), so the interpreter has exactly six shapes.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum NodeKind {
    /// `{ … }` — a fixed, ordered, fully-required property list.
    Object = 0,
    /// `[ … ]` — a homogeneous item schema with optional `minItems`/`maxItems`.
    Array = 1,
    /// `"…"` — the JSON string grammar with `minLength`/`maxLength` (code-point counted).
    StringFree = 2,
    /// The full JSON number grammar, capped at [`super::json_fsm::NUMBER_BYTE_CAP`] bytes.
    Number = 3,
    /// The JSON number grammar restricted to `-?(0|[1-9][0-9]*)` (no fraction/exponent).
    Integer = 4,
    /// Match exactly one of N canonical byte literals (`boolean`/`null`/`const`/`enum`).
    Choice = 5,
}

impl NodeKind {
    /// The reserved discriminant for a future regex-DFA node (`pattern`); never emitted today, so a
    /// compiled table never carries it. Kept so G2's WGSL `switch` can reserve the arm now.
    pub const RESERVED_PATTERN: u32 = 6;
    /// Reserved discriminants for the union/reference constructs G1 refuses (documented for G2).
    pub const RESERVED_ANYOF: u32 = 7;
    pub const RESERVED_ONEOF: u32 = 8;
    pub const RESERVED_ALLOF: u32 = 9;
    pub const RESERVED_REF: u32 = 10;

    fn from_u32(v: u32) -> Self {
        match v {
            0 => NodeKind::Object,
            1 => NodeKind::Array,
            2 => NodeKind::StringFree,
            3 => NodeKind::Number,
            4 => NodeKind::Integer,
            5 => NodeKind::Choice,
            _ => unreachable!("compiled tables only ever hold node kinds 0..=5, got {v}"),
        }
    }
}

/// The compiled, flat, GPU-uploadable representation of one JSON schema. See the module docs for the
/// exact array layout the WGSL interpreter (G2) mirrors.
#[derive(Clone, Debug)]
pub struct JsonSchemaTables {
    nodes: Vec<u32>,
    props: Vec<u32>,
    choice_lits: Vec<u32>,
    literal_spans: Vec<u32>,
    literal_bytes: Vec<u8>,
    root: u32,
}

impl JsonSchemaTables {
    /// Words per node record in [`Self::nodes`].
    pub const NODE_WORDS: usize = 5;
    /// `max_items`/`max_len` sentinel meaning "no upper bound declared".
    pub const UNBOUNDED: u32 = u32::MAX;
    /// The maximum container nesting depth accepted (OpenAI strict-mode ceiling).
    pub const MAX_DEPTH: usize = 10;
    /// The maximum number of members a `Choice` (enum) may carry — bounded by the viable-candidate
    /// bitmask that lives in one state word.
    pub const MAX_CHOICE: usize = 32;

    /// Compile a JSON schema into interpreter tables, or return an error naming the first construct
    /// outside the supported subset. Pure and allocation-light — safe on any per-request hot path.
    pub fn compile(schema: &Value) -> Result<Self> {
        let mut b = Builder::default();
        let root = b.compile_node(schema, 1)?;
        Ok(Self {
            nodes: b.nodes,
            props: b.props,
            choice_lits: b.choice_lits,
            literal_spans: b.literal_spans,
            literal_bytes: b.literal_bytes,
            root,
        })
    }

    // -- accessors used by the interpreter (crate-internal; G2 reads the raw arrays instead) -------

    pub(crate) fn root(&self) -> u32 {
        self.root
    }

    // Raw flat arrays — the exact data G2 packs into GPU buffers (see `grammar::harness::GpuSchema`).
    // The WGSL interpreter indexes these identically to the CPU accessors below.
    // Staged for the G2 GPU-resident grammar integration; no crate-internal caller yet.
    #[allow(dead_code)]
    pub(crate) fn raw_nodes(&self) -> &[u32] {
        &self.nodes
    }
    #[allow(dead_code)]
    pub(crate) fn raw_props(&self) -> &[u32] {
        &self.props
    }
    #[allow(dead_code)]
    pub(crate) fn raw_choice_lits(&self) -> &[u32] {
        &self.choice_lits
    }
    #[allow(dead_code)]
    pub(crate) fn raw_literal_spans(&self) -> &[u32] {
        &self.literal_spans
    }
    #[allow(dead_code)]
    pub(crate) fn raw_literal_bytes(&self) -> &[u8] {
        &self.literal_bytes
    }

    /// The literal bytes for span `s`, reconstructed from the CPU tables — the oracle the flat
    /// [`crate::grammar::GpuSchema::literal`] round-trip check compares against.
    #[cfg(test)]
    pub(crate) fn literal_for_test(&self, s: u32) -> Vec<u8> {
        let base = s as usize * 2;
        let off = self.literal_spans[base] as usize;
        let len = self.literal_spans[base + 1] as usize;
        self.literal_bytes[off..off + len].to_vec()
    }

    pub(crate) fn kind(&self, node: u32) -> NodeKind {
        NodeKind::from_u32(self.nodes[node as usize * Self::NODE_WORDS])
    }

    fn field(&self, node: u32, k: usize) -> u32 {
        self.nodes[node as usize * Self::NODE_WORDS + 1 + k]
    }

    pub(crate) fn obj_count(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    pub(crate) fn obj_key(&self, node: u32, i: usize) -> &[u8] {
        let base = (self.field(node, 0) as usize + i) * 2;
        self.literal(self.props[base])
    }
    pub(crate) fn obj_value(&self, node: u32, i: usize) -> u32 {
        let base = (self.field(node, 0) as usize + i) * 2;
        self.props[base + 1]
    }

    pub(crate) fn arr_item(&self, node: u32) -> u32 {
        self.field(node, 0)
    }
    pub(crate) fn arr_min(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    pub(crate) fn arr_max(&self, node: u32) -> u32 {
        self.field(node, 2)
    }

    pub(crate) fn str_min(&self, node: u32) -> u32 {
        self.field(node, 0)
    }
    pub(crate) fn str_max(&self, node: u32) -> u32 {
        self.field(node, 1)
    }

    pub(crate) fn choice_count(&self, node: u32) -> u32 {
        self.field(node, 1)
    }
    pub(crate) fn choice_lit(&self, node: u32, c: usize) -> &[u8] {
        let span = self.choice_lits[self.field(node, 0) as usize + c];
        self.literal(span)
    }

    pub(crate) fn literal(&self, span: u32) -> &[u8] {
        let off = self.literal_spans[span as usize * 2] as usize;
        let len = self.literal_spans[span as usize * 2 + 1] as usize;
        &self.literal_bytes[off..off + len]
    }

    /// The set of bytes a FRESH frame for `node` will accept as its opening byte — i.e. the first
    /// byte of any value of that node's type. Used by the interpreter to enumerate array-boundary
    /// choices without pushing a child.
    pub(crate) fn open_set(&self, node: u32, set: &mut [bool; 256]) {
        match self.kind(node) {
            NodeKind::Object => set[b'{' as usize] = true,
            NodeKind::Array => set[b'[' as usize] = true,
            NodeKind::StringFree => set[b'"' as usize] = true,
            NodeKind::Number | NodeKind::Integer => {
                set[b'-' as usize] = true;
                for d in b'0'..=b'9' {
                    set[d as usize] = true;
                }
            }
            NodeKind::Choice => {
                let count = self.choice_count(node) as usize;
                for c in 0..count {
                    let lit = self.choice_lit(node, c);
                    if let Some(&first) = lit.first() {
                        set[first as usize] = true;
                    }
                }
            }
        }
    }
}

#[derive(Default)]
struct Builder {
    nodes: Vec<u32>,
    props: Vec<u32>,
    choice_lits: Vec<u32>,
    literal_spans: Vec<u32>,
    literal_bytes: Vec<u8>,
}

/// JSON-Schema keywords G1 understands (structural). Anything else that is not a pure annotation is
/// refused by name.
const SUPPORTED_KEYWORDS: &[&str] = &[
    "type",
    "properties",
    "required",
    "additionalProperties",
    "items",
    "minItems",
    "maxItems",
    "minLength",
    "maxLength",
    "enum",
    "const",
];

/// Annotation-only keywords that do not constrain the instance and are safely ignored.
const ANNOTATION_KEYWORDS: &[&str] = &[
    "title",
    "description",
    "$comment",
    "$schema",
    "$id",
    "$defs",
    "definitions",
    "default",
    "examples",
    "readOnly",
    "writeOnly",
    "deprecated",
];

impl Builder {
    fn intern_literal(&mut self, bytes: &[u8]) -> u32 {
        let off = self.literal_bytes.len() as u32;
        self.literal_bytes.extend_from_slice(bytes);
        let span = (self.literal_spans.len() / 2) as u32;
        self.literal_spans.push(off);
        self.literal_spans.push(bytes.len() as u32);
        span
    }

    fn push_node(&mut self, kind: NodeKind, fields: [u32; 4]) -> u32 {
        let idx = (self.nodes.len() / JsonSchemaTables::NODE_WORDS) as u32;
        self.nodes.push(kind as u32);
        self.nodes.extend_from_slice(&fields);
        idx
    }

    fn compile_node(&mut self, schema: &Value, depth: usize) -> Result<u32> {
        let obj = schema
            .as_object()
            .ok_or_else(|| anyhow::anyhow!("each schema must be a JSON object, got: {schema}"))?;
        // Refuse any non-annotation keyword outside the subset, naming it.
        for key in obj.keys() {
            if !SUPPORTED_KEYWORDS.contains(&key.as_str())
                && !ANNOTATION_KEYWORDS.contains(&key.as_str())
            {
                bail!(
                    "unsupported JSON Schema keyword `{key}`: the G1 subset is OpenAI strict mode \
                     (object/array/string/number/integer/boolean/null, properties, required, \
                     additionalProperties:false, enum, const, min/maxItems, min/maxLength)"
                );
            }
        }

        if obj.contains_key("enum") {
            let members = obj["enum"]
                .as_array()
                .ok_or_else(|| anyhow::anyhow!("`enum` must be an array"))?;
            return self.compile_choice(members);
        }
        if obj.contains_key("const") {
            return self.compile_choice(std::slice::from_ref(&obj["const"]));
        }
        match obj.get("type") {
            Some(Value::String(t)) => match t.as_str() {
                "object" => self.compile_object(obj, depth),
                "array" => self.compile_array(obj, depth),
                "string" => self.compile_string(obj),
                "number" => Ok(self.push_node(NodeKind::Number, [0; 4])),
                "integer" => Ok(self.push_node(NodeKind::Integer, [0; 4])),
                "boolean" => self.compile_choice(&[Value::Bool(true), Value::Bool(false)]),
                "null" => self.compile_choice(&[Value::Null]),
                other => bail!("unsupported `type`: {other:?}"),
            },
            Some(Value::Array(_)) => bail!(
                "unsupported: `type` as a union array behaves like `anyOf` and is outside the subset"
            ),
            Some(other) => bail!("`type` must be a string, got: {other}"),
            None => bail!("schema must declare `type`, `enum`, or `const`"),
        }
    }

    fn compile_object(
        &mut self,
        obj: &serde_json::Map<String, Value>,
        depth: usize,
    ) -> Result<u32> {
        if depth > JsonSchemaTables::MAX_DEPTH {
            bail!(
                "nesting depth exceeds the {} level cap",
                JsonSchemaTables::MAX_DEPTH
            );
        }
        // `additionalProperties: false` is mandatory in strict mode.
        match obj.get("additionalProperties") {
            Some(Value::Bool(false)) => {}
            Some(_) => bail!("`additionalProperties` must be `false` (strict mode)"),
            None => bail!("object schema must set `additionalProperties: false` (strict mode)"),
        }
        let empty = serde_json::Map::new();
        let properties = match obj.get("properties") {
            Some(Value::Object(m)) => m,
            Some(_) => bail!("`properties` must be an object"),
            None => &empty,
        };
        let required: Vec<&str> = match obj.get("required") {
            Some(Value::Array(a)) => a
                .iter()
                .map(|v| {
                    v.as_str()
                        .ok_or_else(|| anyhow::anyhow!("`required` entries must be strings"))
                })
                .collect::<Result<_>>()?,
            Some(_) => bail!("`required` must be an array"),
            None => Vec::new(),
        };
        // Strict mode: `required` must list every property exactly once (canonical order source).
        if required.len() != properties.len() {
            bail!(
                "strict mode requires every property in `required` (have {} required, {} properties)",
                required.len(),
                properties.len()
            );
        }
        // JSON object key ORDER is semantically irrelevant to VALIDATION (the jsonschema oracle is
        // order-independent) — the FSM enforces an order purely so canonical generation is
        // DETERMINISTIC. Because the check above forces required-set == properties-set, the key
        // CONTENT always matches the schema regardless; only the SEQUENCE contract is taken from the
        // `required` array (a parsed `properties` object cannot preserve declaration order here —
        // this workspace does not enable serde_json's `preserve_order`, and flipping it is a
        // workspace-global semantic change, not ours to make).
        let mut pairs = Vec::with_capacity(required.len());
        for key in &required {
            let child = properties.get(*key).ok_or_else(|| {
                anyhow::anyhow!("`required` names `{key}`, absent from `properties`")
            })?;
            let value_node = self.compile_node(child, depth + 1)?;
            let key_lit =
                serde_json::to_vec(&Value::String((*key).to_string())).expect("string serializes");
            let key_span = self.intern_literal(&key_lit);
            pairs.push((key_span, value_node));
        }
        let prop_start = (self.props.len() / 2) as u32;
        for (ks, vn) in &pairs {
            self.props.push(*ks);
            self.props.push(*vn);
        }
        Ok(self.push_node(NodeKind::Object, [prop_start, pairs.len() as u32, 0, 0]))
    }

    fn compile_array(&mut self, obj: &serde_json::Map<String, Value>, depth: usize) -> Result<u32> {
        if depth > JsonSchemaTables::MAX_DEPTH {
            bail!(
                "nesting depth exceeds the {} level cap",
                JsonSchemaTables::MAX_DEPTH
            );
        }
        let items = obj
            .get("items")
            .ok_or_else(|| anyhow::anyhow!("array schema requires an `items` schema"))?;
        if !items.is_object() {
            bail!("`items` must be a single schema object (tuple `items` arrays are unsupported)");
        }
        let item_node = self.compile_node(items, depth + 1)?;
        let min = bound(obj, "minItems", 0)?;
        let max = bound(obj, "maxItems", JsonSchemaTables::UNBOUNDED)?;
        if max != JsonSchemaTables::UNBOUNDED && min > max {
            bail!("minItems ({min}) exceeds maxItems ({max})");
        }
        Ok(self.push_node(NodeKind::Array, [item_node, min, max, 0]))
    }

    fn compile_string(&mut self, obj: &serde_json::Map<String, Value>) -> Result<u32> {
        let min = bound(obj, "minLength", 0)?;
        let max = bound(obj, "maxLength", JsonSchemaTables::UNBOUNDED)?;
        if max != JsonSchemaTables::UNBOUNDED && min > max {
            bail!("minLength ({min}) exceeds maxLength ({max})");
        }
        Ok(self.push_node(NodeKind::StringFree, [min, max, 0, 0]))
    }

    fn compile_choice(&mut self, members: &[Value]) -> Result<u32> {
        if members.is_empty() {
            bail!("`enum`/`const` must have at least one member");
        }
        // Canonicalize each scalar member to its JSON serialization; dedupe.
        let mut lits: Vec<Vec<u8>> = Vec::new();
        for m in members {
            if m.is_object() || m.is_array() {
                bail!("`enum`/`const` members must be scalars (string/number/boolean/null)");
            }
            let lit = serde_json::to_vec(m).expect("scalar serializes");
            if !lits.contains(&lit) {
                lits.push(lit);
            }
        }
        if lits.len() > JsonSchemaTables::MAX_CHOICE {
            bail!(
                "`enum` has {} members but the limit is {}: the viable-candidate set is tracked as \
                 a bitmask in one 32-bit state word, so a larger enum is refused rather than \
                 silently narrowed (the covers()-pattern — refuse, never mis-mask)",
                lits.len(),
                JsonSchemaTables::MAX_CHOICE
            );
        }
        // No member's canonical form may be a byte-prefix of another's, or termination would be
        // ambiguous (the shorter could end while the longer wants to continue). STRING members can
        // never trip this — every string literal ends in `"`, which is not a prefix of a longer
        // string's opening — so in practice this only ever fires for numeric enums like [1, 12].
        for i in 0..lits.len() {
            for j in 0..lits.len() {
                if i != j && lits[j].starts_with(&lits[i]) {
                    bail!(
                        "`enum`/`const` member {:?} is a byte-prefix of {:?}, which makes \
                         termination ambiguous; unsupported. (String members never trip this — the \
                         closing quote disambiguates — so this only affects numeric enums such as \
                         [1, 12].)",
                        String::from_utf8_lossy(&lits[i]),
                        String::from_utf8_lossy(&lits[j])
                    );
                }
            }
        }
        let lit_start = self.choice_lits.len() as u32;
        for lit in &lits {
            let span = self.intern_literal(lit);
            self.choice_lits.push(span);
        }
        Ok(self.push_node(NodeKind::Choice, [lit_start, lits.len() as u32, 0, 0]))
    }
}

/// Read a non-negative integer bound (`minItems`/`maxItems`/`minLength`/`maxLength`), enforcing the
/// ≤ 255 ceiling, or return `default` when absent.
fn bound(obj: &serde_json::Map<String, Value>, key: &str, default: u32) -> Result<u32> {
    match obj.get(key) {
        None => Ok(default),
        Some(v) => {
            let n = v
                .as_u64()
                .ok_or_else(|| anyhow::anyhow!("`{key}` must be a non-negative integer"))?;
            if n > 255 {
                bail!("`{key}` = {n} exceeds the 255 ceiling");
            }
            Ok(n as u32)
        }
    }
}