onedrop_hlsl/rewrite/

qa_stub.rs

//! Pass: `_qa..._qh` stub injection.

use super::*;

// ---------------------------------------------------------------------------
// Pass: `_qa..._qh` stub injection
// ---------------------------------------------------------------------------

/// `_qa`/`_qb`/.../`_qh` are MD2 runtime-defined aliases that pack the 32
/// `q1..q32` evaluator channels into 8 × `float4` (`_qa = float4(q1, q2, q3, q4)`,
/// `_qb = float4(q5, …)`, etc.). User comp shaders read them directly,
/// typically via `mul(uv, float2x2(_qb))` for kaleidoscope rotation.
///
/// The standalone HLSL parser never sees their declarations because they
/// live in the host-provided prelude — naga rejects with
/// `no definition in scope for identifier: '_qb'`. We stub the full
/// `_qa..._qh` family — many presets reference `_qb` and a long tail of
/// the others.
///
/// For each referenced identifier we:
///
/// - prepend `float4 _qX = float4(0.0, 0.0, 0.0, 0.0);` at the start of
///   `shader_body { ... }` so any later read picks up a defined value, and
/// - expand every `float2x2(_qX)` call into the four-scalar constructor
///   form `float2x2(_qX.x, _qX.y, _qX.z, _qX.w)` so it survives
///   `replace_types` into a valid `mat2x2<f32>(...)` call (WGSL refuses
///   the single-vec4 form).
///
/// If no `_qX` identifier is referenced, returns the input unchanged.
pub(crate) fn inject_qa_stub(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut uses = QaUses::default();
    if let Some(body) = &tu.shader_body {
        scan_block_for_qa(body, &mut uses);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            scan_block_for_qa(&f.body, &mut uses);
        }
    }
    if uses.referenced_mask == 0 {
        return src.to_string();
    }
    let Some(body) = &tu.shader_body else {
        return src.to_string();
    };

    let mut edits: Vec<TextEdit> = uses
        .mat_call_spans
        .iter()
        .map(|(sp, letter)| TextEdit {
            start: sp.start,
            end: sp.end,
            replacement: format!(
                "float2x2(_q{letter}.x, _q{letter}.y, _q{letter}.z, _q{letter}.w)"
            ),
        })
        .collect();

    if let Some(open) = find_brace_after(src, body.span.start) {
        let mut stub = String::new();
        for (idx, letter) in QA_LETTERS.iter().enumerate() {
            if uses.referenced_mask & (1u8 << idx) != 0 {
                stub.push_str(&format!(" float4 _q{letter} = float4(0.0, 0.0, 0.0, 0.0);"));
            }
        }
        edits.push(TextEdit {
            start: open,
            end: open,
            replacement: stub,
        });
    }
    apply_edits(src, &mut edits)
}

/// Letters `a..h` map to the 8 four-channel groups packing `q1..q32`.
const QA_LETTERS: [char; 8] = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'];

/// Returns `Some(bit_index)` if `name` is one of `_qa..._qh`.
fn qa_letter_index(name: &str) -> Option<usize> {
    let bytes = name.as_bytes();
    if bytes.len() != 3 || bytes[0] != b'_' || bytes[1] != b'q' {
        return None;
    }
    let c = bytes[2];
    if (b'a'..=b'h').contains(&c) {
        Some((c - b'a') as usize)
    } else {
        None
    }
}

#[derive(Default)]
struct QaUses {
    /// Bit-set of referenced `_qa..._qh` identifiers (bit 0 = `_qa`).
    referenced_mask: u8,
    /// Each entry is `(call_span, ident_letter)` for a `float2x2(_qX)` call
    /// to expand into the four-scalar form.
    mat_call_spans: Vec<(Span, char)>,
}

fn scan_block_for_qa(b: &Block, out: &mut QaUses) {
    for s in &b.stmts {
        scan_stmt_for_qa(s, out);
    }
}

fn scan_stmt_for_qa(s: &Stmt, out: &mut QaUses) {
    match s {
        Stmt::LocalDecl(d) => {
            if let Some(init) = &d.init {
                scan_expr_for_qa(init, out);
            }
        }
        Stmt::Assign(a) => {
            scan_expr_for_qa(&a.target, out);
            scan_expr_for_qa(&a.value, out);
        }
        Stmt::Expr(e) => scan_expr_for_qa(e, out),
        Stmt::If(i) => {
            scan_expr_for_qa(&i.cond, out);
            scan_stmt_for_qa(&i.then_branch, out);
            if let Some(e) = &i.else_branch {
                scan_stmt_for_qa(e, out);
            }
        }
        Stmt::While(w) => {
            scan_expr_for_qa(&w.cond, out);
            scan_stmt_for_qa(&w.body, out);
        }
        Stmt::For(f) => {
            if let Some(init) = &f.init {
                scan_stmt_for_qa(init, out);
            }
            if let Some(c) = &f.cond {
                scan_expr_for_qa(c, out);
            }
            if let Some(st) = &f.step {
                scan_expr_for_qa(st, out);
            }
            scan_stmt_for_qa(&f.body, out);
        }
        Stmt::Return(Some(e)) => scan_expr_for_qa(e, out),
        Stmt::Block(b) => scan_block_for_qa(b, out),
        Stmt::Return(None) | Stmt::Break | Stmt::Continue => {}
    }
}

fn scan_expr_for_qa(e: &Expr, out: &mut QaUses) {
    match e {
        Expr::Ident(name, _) => {
            if let Some(idx) = qa_letter_index(name) {
                out.referenced_mask |= 1u8 << idx;
            }
        }
        Expr::Lit(_) => {}
        Expr::Binary(b) => {
            scan_expr_for_qa(&b.lhs, out);
            scan_expr_for_qa(&b.rhs, out);
        }
        Expr::Unary(u) => scan_expr_for_qa(&u.operand, out),
        Expr::Ternary(t) => {
            scan_expr_for_qa(&t.cond, out);
            scan_expr_for_qa(&t.then_expr, out);
            scan_expr_for_qa(&t.else_expr, out);
        }
        Expr::Call(c) => {
            if c.callee == "float2x2"
                && c.args.len() == 1
                && let Expr::Ident(n, _) = &c.args[0]
                && let Some(idx) = qa_letter_index(n)
            {
                let letter = QA_LETTERS[idx];
                out.mat_call_spans.push((c.span, letter));
                out.referenced_mask |= 1u8 << idx;
                return;
            }
            for a in &c.args {
                scan_expr_for_qa(a, out);
            }
        }
        Expr::Member(m) => scan_expr_for_qa(&m.base, out),
        Expr::Swizzle(s) => scan_expr_for_qa(&s.base, out),
        Expr::Index(i) => {
            scan_expr_for_qa(&i.base, out);
            scan_expr_for_qa(&i.index, out);
        }
        Expr::InitList(l) => {
            for e in &l.elems {
                scan_expr_for_qa(e, out);
            }
        }
        Expr::Assign(a) => {
            scan_expr_for_qa(&a.target, out);
            scan_expr_for_qa(&a.value, out);
        }
    }
}

/// Locate the byte position immediately after the first `{` at or past
/// `start`. `shader_body { … }`'s parser hands us a span starting at the
/// keyword, not the brace, so this helper bridges the gap. Returns `None`
/// when no `{` is found (malformed source — we skip injection in that case).
fn find_brace_after(src: &str, start: u32) -> Option<u32> {
    let bytes = src.as_bytes();
    let mut i = start as usize;
    while i < bytes.len() && bytes[i] != b'{' {
        i += 1;
    }
    if i < bytes.len() {
        Some((i + 1) as u32)
    } else {
        None
    }
}