onedrop_hlsl/

types.rs

//! Type-aware post-translator passes for the HLSL→WGSL pipeline.
//!
//! After the regex rewrites in `lib.rs` produce something WGSL-shaped, the
//! remaining failures cluster around HLSL semantics WGSL doesn't honour:
//!
//! 1. **Implicit scalar→vector broadcasts** in calls like
//!    `clamp(<vec3>, 0, 1)`. HLSL broadcasts the scalars `0` and `1` to
//!    `vec3(0)` and `vec3(1)` automatically; WGSL refuses with
//!    `inconsistent type passed as argument #2 to clamp`.
//! 2. **Implicit scalar←vector truncation** in declarations like
//!    `float lum = GetPixel(uv) * c.x + …;`. HLSL takes the first
//!    component of the vec result; WGSL refuses with `the type of lum is
//!    expected to be f32; but got vec3<f32>`.
//!
//! This module fixes both with two passes that share a small symbol
//! table built from the translated source's `var NAME: TYPE` and
//! `let NAME: TYPE` declarations. Type inference for arbitrary
//! expressions is intentionally heuristic: it returns a confident answer
//! for the dominant MD2 patterns (named locals, numeric literals, vec/
//! mat constructors, calls to a small set of known helpers) and
//! `Unknown` otherwise. `Unknown` short-circuits both passes — the
//! translator only injects fixes it's sure about.

use std::collections::HashMap;

/// WGSL types we care about for MD2 user shaders. Anything outside this
/// set falls through as [`WgslType::Unknown`] and skips both passes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WgslType {
    F32,
    I32,
    Vec2F,
    Vec3F,
    Vec4F,
    Mat2F,
    Mat3F,
    Mat4F,
    Bool,
    Unknown,
}

impl WgslType {
    /// `true` for the vector types we know how to broadcast to.
    pub fn is_vec(self) -> bool {
        matches!(self, Self::Vec2F | Self::Vec3F | Self::Vec4F)
    }

    /// `true` for scalar types — broadcast targets (the value being
    /// wrapped, not the dominant arg type).
    pub fn is_scalar(self) -> bool {
        matches!(self, Self::F32 | Self::I32 | Self::Bool)
    }

    /// WGSL spelling, for emission.
    pub fn wgsl_name(self) -> &'static str {
        match self {
            Self::F32 => "f32",
            Self::I32 => "i32",
            Self::Vec2F => "vec2<f32>",
            Self::Vec3F => "vec3<f32>",
            Self::Vec4F => "vec4<f32>",
            Self::Mat2F => "mat2x2<f32>",
            Self::Mat3F => "mat3x3<f32>",
            Self::Mat4F => "mat4x4<f32>",
            Self::Bool => "bool",
            Self::Unknown => "/* unknown */",
        }
    }

    /// Parse the type fragment that appears between `: ` and `=`/`;` in
    /// a `var/let` declaration. Whitespace-trimmed; case-sensitive.
    fn from_decl_str(s: &str) -> Self {
        match s.trim() {
            "f32" => Self::F32,
            "i32" => Self::I32,
            "u32" => Self::I32, // close enough for our purposes
            "vec2<f32>" => Self::Vec2F,
            "vec3<f32>" => Self::Vec3F,
            "vec4<f32>" => Self::Vec4F,
            "mat2x2<f32>" => Self::Mat2F,
            "mat3x3<f32>" => Self::Mat3F,
            "mat4x4<f32>" => Self::Mat4F,
            "bool" => Self::Bool,
            _ => Self::Unknown,
        }
    }
}

/// Maps locally-declared identifier → its WGSL type. Pre-seeded with the
/// uniforms and helper-function bindings the codegen wrapper exposes
/// inside `fs_main` (so that `texsize`, `q1..q32`, `aspect`, etc. resolve
/// to known types).
pub struct SymbolTable {
    pub locals: HashMap<String, WgslType>,
}

impl SymbolTable {
    /// Build a symbol table by scanning `var NAME: TYPE` / `let NAME: TYPE`
    /// declarations in the translated source. Pre-seeded with the wrapper
    /// preamble's bindings so user code can be analysed in isolation.
    pub fn from_source(src: &str) -> Self {
        let mut locals = HashMap::new();

        // Wrapper preamble (kept in sync with `wrap_user_comp_shader` in
        // `onedrop-codegen`). These are `let` aliases inside fs_main, so
        // the user shader sees them unprefixed.
        for (n, t) in WRAPPER_PRELUDE_LOCALS {
            locals.insert((*n).to_string(), *t);
        }

        // Walk the source and pick up every `var X: T` / `let X: T`.
        let bytes = src.as_bytes();
        let mut i = 0;
        while i < bytes.len() {
            // Skip line comments and block comments so `var` etc. inside
            // them don't pollute the symbol table.
            if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'/' {
                while i < bytes.len() && bytes[i] != b'\n' {
                    i += 1;
                }
                continue;
            }
            if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' {
                i += 2;
                while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
                    i += 1;
                }
                i += 2;
                continue;
            }

            let kw_len = match keyword_at(bytes, i, &["var", "let"]) {
                Some(n) => n,
                None => {
                    i += 1;
                    continue;
                }
            };
            let mut j = i + kw_len;
            while j < bytes.len() && bytes[j].is_ascii_whitespace() {
                j += 1;
            }
            let name_start = j;
            while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
                j += 1;
            }
            if j == name_start {
                i = j + 1;
                continue;
            }
            let name = &src[name_start..j];
            while j < bytes.len() && bytes[j].is_ascii_whitespace() {
                j += 1;
            }
            if j >= bytes.len() || bytes[j] != b':' {
                // No type annotation — skip.
                i = j + 1;
                continue;
            }
            j += 1;
            // Read the type fragment until `=` or `;` at depth 0.
            let ty_start = j;
            let mut depth_angle = 0i32;
            while j < bytes.len() {
                match bytes[j] {
                    b'<' => depth_angle += 1,
                    b'>' => depth_angle -= 1,
                    b'=' | b';' if depth_angle == 0 => break,
                    _ => {}
                }
                j += 1;
            }
            let ty = WgslType::from_decl_str(&src[ty_start..j]);
            if !matches!(ty, WgslType::Unknown) {
                // The scanner is flat — it ignores function scopes. When
                // a name is declared with mutually-incompatible types in
                // separate functions (`var tmp: vec2<f32>` in GetDist
                // vs `var tmp: f32` in MinDist on the whoah/martin
                // shared bodies), the downstream coercion passes pick
                // whichever the scanner happened to see last and
                // mis-type every assignment to that name in the other
                // function. Demote to `Unknown` only when the conflict
                // mixes scalar with vec — the assignment passes treat
                // Unknown as "leave alone" and naga gets to validate
                // each scope on its own. Same-shape conflicts (vec3 vs
                // vec4) keep last-wins so the existing swizzle-assign
                // rewrite still has a vec width to work with — those
                // are common when a global `float3 col;` shadows a
                // local `float4 col;` and the user's `col.rgb *= …`
                // only needs the LHS to *be* a vec to lower cleanly.
                use std::collections::hash_map::Entry;
                match locals.entry(name.to_string()) {
                    Entry::Vacant(e) => {
                        e.insert(ty);
                    }
                    Entry::Occupied(mut e) => {
                        let prev = *e.get();
                        if prev != ty
                            && (prev.is_scalar() != ty.is_scalar() || prev.is_vec() != ty.is_vec())
                        {
                            e.insert(WgslType::Unknown);
                        } else {
                            // Same kind (both vec / both scalar) but
                            // different exact type — keep latest.
                            e.insert(ty);
                        }
                    }
                }
            }
            i = j;
        }

        Self { locals }
    }

    pub fn lookup(&self, name: &str) -> Option<WgslType> {
        self.locals.get(name).copied()
    }

    /// Best-effort type inference for an expression text. Returns a
    /// confident type for the dominant patterns, [`WgslType::Unknown`]
    /// otherwise — callers must treat `Unknown` as "do not modify".
    ///
    /// Recognised:
    /// - A numeric literal (with or without `f`/`i`/`u` suffix) → `F32`.
    /// - A bare identifier in the symbol table → that type.
    /// - A constructor call `vec3<f32>(...)` etc. → that type.
    /// - A call to a known helper (`GetPixel`, `GetBlur1..3`, `lum`,
    ///   `textureSample`, etc.) → its return type.
    /// - A swizzle access `<expr>.<comp>` where the prefix is a vec
    ///   type — narrowed to `F32` for `.x/.y/.z/.w/.r/.g/.b/.a` of length 1,
    ///   `Vec2F` for 2, `Vec3F` for 3, `Vec4F` for 4.
    /// - A binary op (`+`, `-`, `*`, `/`) where one side is vec → the
    ///   widest vec type. Falls back to `F32` when both sides are scalar.
    pub fn infer_expr_type(&self, expr: &str) -> WgslType {
        // Strip C-style comments first — `/*was: foo*/` markers from the
        // tex2D rewrite contain `/` and `*`, which the binop splitter
        // would otherwise mistake for top-level division/multiplication
        // and slice the expression in two. We only strip for inference;
        // the visible output keeps the comments for debuggability.
        let stripped = strip_comments(expr);
        let expr = strip_outer_parens(stripped.trim());

        if expr.is_empty() {
            return WgslType::Unknown;
        }

        // Numeric literal.
        if is_numeric_literal(expr) {
            return WgslType::F32;
        }

        // Bare identifier.
        if is_identifier(expr) {
            if let Some(t) = self.lookup(expr) {
                return t;
            }
            // A stray unknown identifier is just unknown — could be a
            // user variable we missed.
            return WgslType::Unknown;
        }

        // Constructor calls and known helper calls.
        if let Some((head, args)) = split_call(expr) {
            if let Some(t) = constructor_type(head) {
                return t;
            }
            if let Some(t) = known_call_return_type(head) {
                return t;
            }
            // Polymorphic builtins (`clamp`, `min`, `max`, …) take the
            // *element type* of their args. We pick the smallest vec
            // among the args (matching the broadcast pass's "narrow to
            // smallest" rule); if no vec is present, fall back to F32.
            if POLY_BUILTINS.contains(&head) {
                let mut smallest: Option<WgslType> = None;
                for a in split_top_level_commas(args) {
                    let t = self.infer_expr_type(a.trim());
                    if t.is_vec() {
                        smallest = Some(match smallest {
                            None => t,
                            Some(s) => narrower(s, t),
                        });
                    }
                }
                return smallest.unwrap_or(WgslType::F32);
            }
        }

        // Binary op at the top level FIRST — `a + b.x` is `a + (b.x)`,
        // so a top-level `+` outranks the swizzle's `.`. Split on
        // `+ - * /` at depth 0 and pick the widest type among operands.
        if let Some(operands) = split_binop_operands(expr) {
            let mut widest = WgslType::Unknown;
            for op in &operands {
                let t = self.infer_expr_type(op);
                widest = widen(widest, t);
            }
            if !matches!(widest, WgslType::Unknown) {
                return widest;
            }
        }

        // Swizzle: split on the *last* `.` at depth 0. Only reached for
        // single-term expressions like `c.xyz` — chained binops were
        // handled above.
        if let Some((prefix, comp)) = split_last_swizzle(expr)
            && is_swizzle_components(comp)
        {
            let prefix_ty = self.infer_expr_type(prefix);
            if prefix_ty.is_vec() {
                return swizzle_target_type(comp.len());
            }
        }

        // Unary prefix (`-x`, `+x`, `!x`) — same type as operand. Placed
        // AFTER the top-level binop split so `a - b` still parses as a
        // subtraction (operands `a` and `b`); only single-term primaries
        // like `-r` or `-(rad-.5)` fall through to here. Without this,
        // `clamp(v, -r, r)` saw arg #2 as Unknown and the broadcast pass
        // wrapped only arg #3 — surfacing as `inconsistent type passed
        // as argument #2 to clamp` (the single largest residual cluster
        // on the post-f83a9cc 2000-sample, ≈24 presets).
        if matches!(expr.as_bytes().first(), Some(b'-' | b'+' | b'!')) {
            return self.infer_expr_type(&expr[1..]);
        }

        WgslType::Unknown
    }
}

/// Wrapper preamble locals exposed to user shaders. Mirrors the `let`
/// aliases in `onedrop-codegen::wrap_user_comp_shader::USER_COMP_FRAGMENT_PREFIX`.
const WRAPPER_PRELUDE_LOCALS: &[(&str, WgslType)] = &[
    ("uv", WgslType::Vec2F),
    ("uv_orig", WgslType::Vec2F),
    ("rad", WgslType::F32),
    ("ang", WgslType::F32),
    ("ret", WgslType::Vec3F),
    ("color", WgslType::Vec3F),
    ("texsize", WgslType::Vec4F),
    ("aspect", WgslType::Vec4F),
    ("time", WgslType::F32),
    ("fps", WgslType::F32),
    ("frame", WgslType::F32),
    ("progress", WgslType::F32),
    ("bass", WgslType::F32),
    ("mid", WgslType::F32),
    ("treb", WgslType::F32),
    ("vol", WgslType::F32),
    ("bass_att", WgslType::F32),
    ("mid_att", WgslType::F32),
    ("treb_att", WgslType::F32),
    ("vol_att", WgslType::F32),
    ("rand_preset", WgslType::Vec4F),
    ("rand_frame", WgslType::Vec4F),
    ("slow_roam_cos", WgslType::Vec4F),
    ("slow_roam_sin", WgslType::Vec4F),
    ("roam_cos", WgslType::Vec4F),
    ("roam_sin", WgslType::Vec4F),
    ("blur1_min", WgslType::F32),
    ("blur1_max", WgslType::F32),
    ("blur2_min", WgslType::F32),
    ("blur2_max", WgslType::F32),
    ("blur3_min", WgslType::F32),
    ("blur3_max", WgslType::F32),
    ("hue_shader", WgslType::Vec3F),
    ("g_fTexSize", WgslType::Vec4F),
    ("texsize_noise_lq", WgslType::Vec4F),
    ("texsize_noise_lq_lite", WgslType::Vec4F),
    ("texsize_noise_mq", WgslType::Vec4F),
    ("texsize_noise_hq", WgslType::Vec4F),
    ("texsize_noisevol_lq", WgslType::Vec4F),
    ("texsize_noisevol_hq", WgslType::Vec4F),
    // q1..q32 and the math constants are scalar f32.
    ("q1", WgslType::F32),
    ("q2", WgslType::F32),
    ("q3", WgslType::F32),
    ("q4", WgslType::F32),
    ("q5", WgslType::F32),
    ("q6", WgslType::F32),
    ("q7", WgslType::F32),
    ("q8", WgslType::F32),
    ("q9", WgslType::F32),
    ("q10", WgslType::F32),
    ("q11", WgslType::F32),
    ("q12", WgslType::F32),
    ("q13", WgslType::F32),
    ("q14", WgslType::F32),
    ("q15", WgslType::F32),
    ("q16", WgslType::F32),
    ("q17", WgslType::F32),
    ("q18", WgslType::F32),
    ("q19", WgslType::F32),
    ("q20", WgslType::F32),
    ("q21", WgslType::F32),
    ("q22", WgslType::F32),
    ("q23", WgslType::F32),
    ("q24", WgslType::F32),
    ("q25", WgslType::F32),
    ("q26", WgslType::F32),
    ("q27", WgslType::F32),
    ("q28", WgslType::F32),
    ("q29", WgslType::F32),
    ("q30", WgslType::F32),
    ("q31", WgslType::F32),
    ("q32", WgslType::F32),
    ("M_PI", WgslType::F32),
    ("M_PI_2", WgslType::F32),
    ("M_INV_PI", WgslType::F32),
    ("M_INV_PI_2", WgslType::F32),
];

fn keyword_at(bytes: &[u8], i: usize, kws: &[&str]) -> Option<usize> {
    if i > 0 {
        let prev = bytes[i - 1];
        if prev.is_ascii_alphanumeric() || prev == b'_' {
            return None;
        }
    }
    for kw in kws {
        let len = kw.len();
        if i + len <= bytes.len() && &bytes[i..i + len] == kw.as_bytes() {
            let next = bytes.get(i + len).copied();
            let is_id = next.is_some_and(|c| c.is_ascii_alphanumeric() || c == b'_');
            if !is_id {
                return Some(len);
            }
        }
    }
    None
}

/// Replace `/* ... */` and `// ...` comments with spaces in-place so the
/// resulting string has the same length and column positions as the
/// original — crucial because [`SymbolTable::infer_expr_type`] doesn't
/// remap offsets when reasoning about an arg's text.
fn strip_comments(src: &str) -> String {
    let bytes = src.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;
    while i < bytes.len() {
        if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'/' {
            while i < bytes.len() && bytes[i] != b'\n' {
                out.push(b' ');
                i += 1;
            }
        } else if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' {
            // Walk to the matching `*/`. Replace every byte (including the
            // markers) with spaces; preserve newlines so line numbers
            // stay aligned.
            while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
                out.push(if bytes[i] == b'\n' { b'\n' } else { b' ' });
                i += 1;
            }
            // Replace the closing `*/` if found.
            if i + 1 < bytes.len() {
                out.push(b' ');
                out.push(b' ');
                i += 2;
            }
        } else {
            out.push(bytes[i]);
            i += 1;
        }
    }
    // SAFETY: input was UTF-8 and we only replaced bytes inside ASCII
    // comment runs with ASCII spaces / newlines, so the result is still
    // valid UTF-8.
    String::from_utf8(out).expect("comment-stripping preserved UTF-8")
}

fn strip_outer_parens(expr: &str) -> &str {
    let mut e = expr;
    loop {
        let trimmed = e.trim();
        if !trimmed.starts_with('(') || !trimmed.ends_with(')') {
            return trimmed;
        }
        // Only strip if the outermost parens enclose the whole expression.
        let bytes = trimmed.as_bytes();
        let mut depth = 0i32;
        let mut closed_at = None;
        for (i, &b) in bytes.iter().enumerate() {
            match b {
                b'(' => depth += 1,
                b')' => {
                    depth -= 1;
                    if depth == 0 {
                        closed_at = Some(i);
                        break;
                    }
                }
                _ => {}
            }
        }
        if closed_at == Some(bytes.len() - 1) {
            e = &trimmed[1..bytes.len() - 1];
        } else {
            return trimmed;
        }
    }
}

fn is_numeric_literal(s: &str) -> bool {
    let s = s.trim();
    if s.is_empty() {
        return false;
    }
    // Strip a trailing type suffix.
    let s = s.trim_end_matches(['f', 'i', 'u', 'h']);
    // A leading `-` is fine — we may see `-1.0`.
    let core = s.strip_prefix('-').unwrap_or(s);
    !core.is_empty()
        && core
            .chars()
            .all(|c| c.is_ascii_digit() || c == '.' || c == 'e' || c == 'E' || c == '+' || c == '-')
        && core.chars().any(|c| c.is_ascii_digit())
}

fn is_identifier(s: &str) -> bool {
    let mut chars = s.chars();
    match chars.next() {
        Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
        _ => return false,
    }
    chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}

/// Split `<head>(<args>)` into (head, args). Returns `None` if the
/// expression isn't a bare call (e.g. has a trailing `.foo`).
fn split_call(expr: &str) -> Option<(&str, &str)> {
    let bytes = expr.as_bytes();
    let open = expr.find('(')?;
    if open == 0 {
        return None;
    }
    // The head must be a single identifier or `vecN<f32>` etc.
    // Check that closing paren is at the very end.
    if bytes.last() != Some(&b')') {
        return None;
    }
    // Walk to confirm balance.
    let mut depth = 0i32;
    for (i, &b) in bytes.iter().enumerate().skip(open) {
        match b {
            b'(' => depth += 1,
            b')' => {
                depth -= 1;
                if depth == 0 {
                    if i != bytes.len() - 1 {
                        return None;
                    }
                    return Some((expr[..open].trim(), &expr[open + 1..i]));
                }
            }
            _ => {}
        }
    }
    None
}

fn constructor_type(head: &str) -> Option<WgslType> {
    Some(match head {
        "vec2<f32>" => WgslType::Vec2F,
        "vec3<f32>" => WgslType::Vec3F,
        "vec4<f32>" => WgslType::Vec4F,
        "mat2x2<f32>" => WgslType::Mat2F,
        "mat3x3<f32>" => WgslType::Mat3F,
        "mat4x4<f32>" => WgslType::Mat4F,
        "f32" => WgslType::F32,
        "i32" => WgslType::I32,
        _ => return None,
    })
}

fn known_call_return_type(head: &str) -> Option<WgslType> {
    Some(match head {
        // MD2 wrapper helpers (in `onedrop-codegen::USER_COMP_HELPERS`).
        "GetPixel" | "GetBlur1" | "GetBlur2" | "GetBlur3" => WgslType::Vec3F,
        "lum" => WgslType::F32,
        // WGSL builtins that return a known type regardless of args.
        "length" | "distance" | "dot" => WgslType::F32,
        "cross" => WgslType::Vec3F,
        "textureSample" => WgslType::Vec4F,
        // Same-type-as-arg builtins — best heuristic is "preserve vec3"
        // because the dominant MD2 body manipulates vec3 colour. A
        // conservative `Unknown` would block the broadcast pass; pick
        // Vec3F since it's the dominant case and a wrong guess only
        // means we fail to inject a needed broadcast (graceful: same
        // behaviour as before this pass).
        // (Intentionally not enumerated — return Unknown so the pass
        // falls back on its own per-arg analysis.)
        _ => return None,
    })
}

/// Return `(prefix, components)` where `<expr>.<components>` is the swizzle.
/// Splits on the last `.` at depth 0.
fn split_last_swizzle(expr: &str) -> Option<(&str, &str)> {
    let bytes = expr.as_bytes();
    let mut depth_paren = 0i32;
    let mut depth_angle = 0i32;
    let mut last_dot = None;
    for (i, &b) in bytes.iter().enumerate() {
        match b {
            b'(' => depth_paren += 1,
            b')' => depth_paren -= 1,
            b'<' => depth_angle += 1,
            b'>' => depth_angle -= 1,
            b'.' if depth_paren == 0 && depth_angle == 0 => last_dot = Some(i),
            _ => {}
        }
    }
    let dot = last_dot?;
    // Don't split numeric literals.
    let pre = expr[..dot].trim();
    if pre.is_empty() || is_numeric_literal(pre) {
        return None;
    }
    Some((pre, &expr[dot + 1..]))
}

fn is_swizzle_components(s: &str) -> bool {
    !s.is_empty()
        && s.len() <= 4
        && s.chars()
            .all(|c| matches!(c, 'x' | 'y' | 'z' | 'w' | 'r' | 'g' | 'b' | 'a'))
}

/// Component count for vec types, 0 for everything else. Used by the
/// truncation pass to size the trailing `.xy`/`.xyz` swizzle when
/// narrowing a wider vec into a smaller one.
pub(crate) fn vec_size(t: WgslType) -> usize {
    match t {
        WgslType::Vec2F => 2,
        WgslType::Vec3F => 3,
        WgslType::Vec4F => 4,
        _ => 0,
    }
}

/// Returns the WgslType for a vec of the given component count. `1` maps
/// to scalar f32, `2..4` to the matching vec type, anything else to
/// `Unknown`. Convenience for AST-driven rewrites that need to construct
/// a target type from an inferred component count.
pub(crate) fn vec_of_size(n: usize) -> WgslType {
    match n {
        1 => WgslType::F32,
        2 => WgslType::Vec2F,
        3 => WgslType::Vec3F,
        4 => WgslType::Vec4F,
        _ => WgslType::Unknown,
    }
}

fn swizzle_target_type(len: usize) -> WgslType {
    match len {
        1 => WgslType::F32,
        2 => WgslType::Vec2F,
        3 => WgslType::Vec3F,
        4 => WgslType::Vec4F,
        _ => WgslType::Unknown,
    }
}

/// Split an expression on top-level binary ops `+ - * /` and return the
/// operands. Returns `None` if no top-level op found (i.e. the expression
/// is a single term).
fn split_binop_operands(expr: &str) -> Option<Vec<&str>> {
    let bytes = expr.as_bytes();
    let mut depth_paren = 0i32;
    let mut depth_angle = 0i32;
    let mut depth_bracket = 0i32;
    let mut splits = Vec::new();
    let mut prev_was_op_or_start = true;
    for (i, &b) in bytes.iter().enumerate() {
        match b {
            b'(' => {
                depth_paren += 1;
                prev_was_op_or_start = false;
            }
            b')' => {
                depth_paren -= 1;
                prev_was_op_or_start = false;
            }
            b'<' => {
                depth_angle += 1;
                prev_was_op_or_start = false;
            }
            b'>' => {
                depth_angle -= 1;
                prev_was_op_or_start = false;
            }
            b'[' => {
                depth_bracket += 1;
                prev_was_op_or_start = false;
            }
            b']' => {
                depth_bracket -= 1;
                prev_was_op_or_start = false;
            }
            b'+' | b'-' | b'*' | b'/'
                if depth_paren == 0
                    && depth_angle == 0
                    && depth_bracket == 0
                    && !prev_was_op_or_start =>
            {
                splits.push(i);
                prev_was_op_or_start = true;
            }
            c if c.is_ascii_whitespace() => {}
            _ => prev_was_op_or_start = false,
        }
    }
    if splits.is_empty() {
        return None;
    }
    let mut out = Vec::with_capacity(splits.len() + 1);
    let mut start = 0;
    for &s in &splits {
        out.push(&expr[start..s]);
        start = s + 1;
    }
    out.push(&expr[start..]);
    Some(out)
}

/// Pick the "widest" of two types — vec wins over scalar; among vecs
/// keep the larger; on conflict default to the first non-unknown.
fn widen(a: WgslType, b: WgslType) -> WgslType {
    match (a, b) {
        (WgslType::Unknown, x) | (x, WgslType::Unknown) => x,
        (x, y) if x == y => x,
        (WgslType::Vec4F, _) | (_, WgslType::Vec4F) => WgslType::Vec4F,
        (WgslType::Vec3F, _) | (_, WgslType::Vec3F) => WgslType::Vec3F,
        (WgslType::Vec2F, _) | (_, WgslType::Vec2F) => WgslType::Vec2F,
        _ => a,
    }
}

// ---------------------------------------------------------------------------
// Pass: scalar→vector broadcast injection in known multi-arg builtins
// ---------------------------------------------------------------------------

/// Built-in functions that require all arguments to share a type and
/// don't accept HLSL-style scalar↔vector broadcasts in WGSL. For each
/// call to one of these, if any arg has a known vec type and another
/// is a scalar, the scalar is wrapped in the matching vec constructor.
const BROADCAST_BUILTINS: &[&str] = &[
    "clamp",
    "min",
    "max",
    "mix",
    "step",
    "smoothstep",
    "pow",
    // `dot(a, b)` requires both args to share a vec size; corpus has
    // `dot(vec2, vec3)` shapes that HLSL silently truncates. Adding it
    // here narrows the larger arg to the smaller vec.
    "dot",
    // `cross(a, b)` is strictly vec3×vec3 in WGSL; corpus pattern is
    // `cross(vec3, textureSample(...))` — arg #2 is vec4 and naga
    // rejects with `wrong type passed as argument #2 to cross`. The
    // narrower-vec rule truncates to vec3.
    "cross",
];

/// Builtins whose return type matches their arg type (HLSL "polymorphic"
/// element-typed functions). Used by [`SymbolTable::infer_expr_type`] —
/// the same set as [`BROADCAST_BUILTINS`] plus a few more pure-passthrough
/// math functions we don't need to broadcast-rewrite but whose return
/// type is "same as arg".
const POLY_BUILTINS: &[&str] = &[
    "clamp",
    "min",
    "max",
    "mix",
    "step",
    "smoothstep",
    "pow",
    "abs",
    "sign",
    "floor",
    "ceil",
    "fract",
    "exp",
    "log",
    "sin",
    "cos",
    "tan",
    "sqrt",
    "normalize",
];

/// Walk `src` looking for calls to broadcast-prone builtins; rewrite
/// scalar arguments to vec constructors where the dominant arg is a vec.
///
/// **Recursion strategy**: for every matched call, we first recursively
/// apply the broadcast rewrite to each argument, *then* decide whether
/// to broadcast at the current level. This ensures that
/// `pow(clamp(<vec>, 0, 1), 0.5)` both rewrites the inner `clamp` and
/// the outer `pow` in a single sweep, instead of needing a fixed-point
/// loop with the walker skipping past unrewritten inner calls.
pub fn inject_broadcasts(src: &str, table: &SymbolTable) -> String {
    let bytes = src.as_bytes();
    let mut out = String::with_capacity(src.len() + 64);
    let mut i = 0;

    while i < bytes.len() {
        // Try to match a builtin name on a word boundary.
        let mut matched = None;
        for name in BROADCAST_BUILTINS {
            let len = name.len();
            if i + len < bytes.len()
                && &bytes[i..i + len] == name.as_bytes()
                && bytes[i + len] == b'('
                && (i == 0 || !(bytes[i - 1].is_ascii_alphanumeric() || bytes[i - 1] == b'_'))
            {
                matched = Some(*name);
                break;
            }
        }

        let Some(name) = matched else {
            out.push(bytes[i] as char);
            i += 1;
            continue;
        };

        // Balance parens to find the call's args.
        let arg_start = i + name.len() + 1;
        let mut j = arg_start;
        let mut depth = 1i32;
        while j < bytes.len() {
            match bytes[j] {
                b'(' => depth += 1,
                b')' => {
                    depth -= 1;
                    if depth == 0 {
                        break;
                    }
                }
                _ => {}
            }
            j += 1;
        }
        if j >= bytes.len() {
            // Unbalanced — emit verbatim and bail.
            out.push_str(&src[i..]);
            return out;
        }

        let args_text = &src[arg_start..j];
        let raw_args = split_top_level_commas(args_text);

        // Recurse first so any nested broadcast-prone call gets rewritten
        // before we infer the type at this level. The recursed text is
        // the source of truth for both the inferred-type lookup and the
        // emit step below.
        let rewritten_args: Vec<String> = raw_args
            .iter()
            .map(|a| inject_broadcasts(a.trim(), table))
            .collect();
        let arg_types: Vec<WgslType> = rewritten_args
            .iter()
            .map(|a| table.infer_expr_type(a))
            .collect();

        // Pick the *smallest* vec among args as the target. HLSL allows
        // `max(vec3, vec4)` (truncating the vec4 to vec3); WGSL refuses
        // and demands all args share a type. The smallest vec wins so we
        // truncate larger args down — matches HLSL's "first arg
        // constrains" rule for cases like `ret = max(ret_vec3, tex2D())`.
        let target = arg_types
            .iter()
            .copied()
            .filter(|t| t.is_vec())
            .reduce(narrower);

        out.push_str(name);
        out.push('(');
        match target {
            Some(vec_ty) => {
                for (k, (arg, ty)) in rewritten_args.iter().zip(arg_types.iter()).enumerate() {
                    if k > 0 {
                        out.push_str(", ");
                    }
                    let wrap = arg_wrap(*ty, vec_ty, arg);
                    match wrap {
                        ArgWrap::Broadcast => {
                            out.push_str(vec_ty.wgsl_name());
                            out.push('(');
                            out.push_str(arg);
                            out.push(')');
                        }
                        ArgWrap::Truncate(swizzle) => {
                            out.push('(');
                            out.push_str(arg);
                            out.push(')');
                            out.push_str(swizzle);
                        }
                        ArgWrap::None => {
                            out.push_str(arg);
                        }
                    }
                }
            }
            None => {
                // No vec dominance — emit args (post-recursion) verbatim.
                for (k, arg) in rewritten_args.iter().enumerate() {
                    if k > 0 {
                        out.push_str(", ");
                    }
                    out.push_str(arg);
                }
            }
        }
        out.push(')');
        i = j + 1;
    }

    out
}

enum ArgWrap {
    None,
    Broadcast,
    Truncate(&'static str),
}

/// Decide how to coerce an arg of `arg_ty` to `target` (the smallest vec
/// among the call's args). Scalars get broadcast; larger vecs get
/// truncated; equal vecs and unknowns mostly pass through (except for
/// literal scalars, which we still broadcast even when type inference
/// returned Unknown).
fn arg_wrap(arg_ty: WgslType, target: WgslType, arg: &str) -> ArgWrap {
    if arg_ty == target {
        return ArgWrap::None;
    }
    if arg_ty.is_scalar() {
        return ArgWrap::Broadcast;
    }
    if matches!(arg_ty, WgslType::Unknown) && is_numeric_literal(arg.trim()) {
        return ArgWrap::Broadcast;
    }
    if arg_ty.is_vec() && target.is_vec() && vec_size(arg_ty) > vec_size(target) {
        return ArgWrap::Truncate(match vec_size(target) {
            2 => ".xy",
            3 => ".xyz",
            _ => "",
        });
    }
    ArgWrap::None
}

/// Reverse of [`widen`] — pick the smaller vec when both args are vecs.
/// Used by the broadcast pass so a `max(vec3, vec4)` call truncates the
/// vec4 to vec3 (HLSL semantics), not the other way round.
fn narrower(a: WgslType, b: WgslType) -> WgslType {
    match (a, b) {
        (WgslType::Unknown, x) | (x, WgslType::Unknown) => x,
        (x, y) if x == y => x,
        (WgslType::Vec2F, _) | (_, WgslType::Vec2F) => WgslType::Vec2F,
        (WgslType::Vec3F, _) | (_, WgslType::Vec3F) => WgslType::Vec3F,
        (WgslType::Vec4F, _) | (_, WgslType::Vec4F) => WgslType::Vec4F,
        _ => a,
    }
}

fn split_top_level_commas(s: &str) -> Vec<&str> {
    let bytes = s.as_bytes();
    let mut out = Vec::new();
    let mut depth_paren = 0i32;
    let mut depth_angle = 0i32;
    let mut start = 0usize;
    for (i, &b) in bytes.iter().enumerate() {
        match b {
            b'(' => depth_paren += 1,
            b')' => depth_paren -= 1,
            b'<' => depth_angle += 1,
            b'>' => depth_angle -= 1,
            b',' if depth_paren == 0 && depth_angle == 0 => {
                out.push(&s[start..i]);
                start = i + 1;
            }
            _ => {}
        }
    }
    out.push(&s[start..]);
    out
}

// ---------------------------------------------------------------------------
// Pass: assignment-target type coercion
// ---------------------------------------------------------------------------

/// Walk the source for `<bare_ident> <op>= <expr>;` statements (where
/// `<op>` is empty for plain assignment or one of `+ - * /` for compound
/// assignment). When the inferred type of `<expr>` doesn't match the
/// declared type of `<bare_ident>`, inject the HLSL implicit conversion
/// (broadcast or truncate).
///
/// This is the assignment-statement analogue of [`inject_truncations`],
/// which handles `var X: T = ...` declarations. The dominant motivating
/// case is `ret += tex3D(...)*0.5;` where `ret` is vec3 in the wrapper
/// preamble and the tex3D fallback returns vec4 — WGSL refuses with
/// `InvalidBinaryOperandTypes`.
///
/// Conservative scope:
/// - LHS must be a bare identifier (no swizzle, no array index).
/// - LHS must resolve to a known vec or scalar type in `table`.
/// - RHS must infer to a known type that differs from LHS.
/// - Compound `*=` / `/=` only fire when RHS is scalar→vec (broadcast)
///   or larger vec→smaller (truncate) — leaving matrix-mul ambiguity
///   well alone.
pub fn inject_assignment_coercions(src: &str, table: &SymbolTable) -> String {
    let bytes = src.as_bytes();
    let mut out = String::with_capacity(src.len() + 64);
    let mut i = 0usize;
    let mut at_stmt_start = true;
    let mut paren = 0i32;

    while i < bytes.len() {
        // Track comments — never start a statement inside a comment.
        if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'/' {
            while i < bytes.len() && bytes[i] != b'\n' {
                out.push(bytes[i] as char);
                i += 1;
            }
            continue;
        }
        if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' {
            let s = i;
            i += 2;
            while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
                i += 1;
            }
            if i + 1 < bytes.len() {
                i += 2;
            }
            out.push_str(&src[s..i]);
            continue;
        }

        // Track paren depth — only attempt assignment matching at depth 0.
        match bytes[i] {
            b'(' => paren += 1,
            b')' => paren -= 1,
            _ => {}
        }

        if !at_stmt_start || paren != 0 || !bytes[i].is_ascii_alphabetic() && bytes[i] != b'_' {
            // Update statement-start tracking and continue. Whitespace
            // leaves the state untouched; statement boundaries set it
            // back to true; everything else falsifies it.
            if !bytes[i].is_ascii_whitespace() {
                at_stmt_start = matches!(bytes[i], b';' | b'{' | b'}');
            }
            out.push(bytes[i] as char);
            i += 1;
            continue;
        }

        // Try to parse `<ident>(.<swizzle>)? <op>= <expr>;`.
        let id_start = i;
        let mut p = i;
        while p < bytes.len() && (bytes[p].is_ascii_alphanumeric() || bytes[p] == b'_') {
            p += 1;
        }
        let id_end = p;
        let name = &src[id_start..id_end];
        // Optionally consume a swizzle suffix (`.x`, `.xy`, `.xyz`,
        // `.xyzw`). Lets us coerce `uv2.x = tex2D(...)` (a real MD2 pattern
        // in midgitstraights / suksma-neck) by treating the LHS as f32
        // when the swizzle is single-component, vecN otherwise. The
        // accessor characters mirror WGSL's set (`xyzw` + the legacy
        // `rgba` aliases normalised away earlier in the pipeline — kept
        // here as defence in depth in case a future pass reintroduces them).
        let mut sw_len = 0usize;
        if p < bytes.len() && bytes[p] == b'.' && p + 1 < bytes.len() {
            let sw_start = p + 1;
            let mut q = sw_start;
            while q < bytes.len()
                && matches!(
                    bytes[q],
                    b'x' | b'y' | b'z' | b'w' | b'r' | b'g' | b'b' | b'a'
                )
            {
                q += 1;
            }
            if q > sw_start && q - sw_start <= 4 {
                sw_len = q - sw_start;
                p = q;
            }
        }
        // Skip whitespace.
        while p < bytes.len() && bytes[p].is_ascii_whitespace() {
            p += 1;
        }
        if p >= bytes.len() {
            out.push_str(&src[id_start..p]);
            i = p;
            at_stmt_start = false;
            continue;
        }
        // Look for assignment op: `=`, `+=`, `-=`, `*=`, `/=`.
        let (op_byte, op_len) = match bytes[p] {
            b'=' if bytes.get(p + 1) != Some(&b'=') => (None, 1),
            b'+' if bytes.get(p + 1) == Some(&b'=') => (Some(b'+'), 2),
            b'-' if bytes.get(p + 1) == Some(&b'=') => (Some(b'-'), 2),
            b'*' if bytes.get(p + 1) == Some(&b'=') => (Some(b'*'), 2),
            b'/' if bytes.get(p + 1) == Some(&b'=') => (Some(b'/'), 2),
            _ => {
                // Not an assignment statement — pass through.
                out.push_str(&src[id_start..p]);
                i = p;
                at_stmt_start = false;
                continue;
            }
        };
        let op_end = p + op_len;
        // Read RHS until `;` at depth 0.
        let rhs_start = op_end;
        let mut q = rhs_start;
        let mut dpar = 0i32;
        let mut dbr = 0i32;
        while q < bytes.len() {
            match bytes[q] {
                b'(' => dpar += 1,
                b')' => dpar -= 1,
                b'[' => dbr += 1,
                b']' => dbr -= 1,
                b';' if dpar == 0 && dbr == 0 => break,
                _ => {}
            }
            q += 1;
        }
        if q >= bytes.len() {
            out.push_str(&src[id_start..q]);
            i = q;
            continue;
        }
        let rhs_raw = &src[rhs_start..q];
        let rhs_trimmed = rhs_raw.trim();

        // Look up LHS type. If unknown, pass through.
        let Some(base_ty) = table.lookup(name) else {
            out.push_str(&src[id_start..q]);
            i = q;
            at_stmt_start = false;
            continue;
        };

        // If we consumed a swizzle suffix, the effective LHS type is
        // derived from the swizzle length (single-component = scalar,
        // N-component = vec_N). The underlying ident still needs to be a
        // vec — bail otherwise (writing to `f32.x` makes no sense).
        let lhs_ty = if sw_len == 0 {
            base_ty
        } else if base_ty.is_vec() {
            match sw_len {
                1 => WgslType::F32,
                2 => WgslType::Vec2F,
                3 => WgslType::Vec3F,
                4 => WgslType::Vec4F,
                _ => {
                    out.push_str(&src[id_start..q]);
                    i = q;
                    at_stmt_start = false;
                    continue;
                }
            }
        } else {
            out.push_str(&src[id_start..q]);
            i = q;
            at_stmt_start = false;
            continue;
        };

        // Don't bother for non-scalar/non-vec types (matrix assignment is
        // outside our scope).
        if !lhs_ty.is_scalar() && !lhs_ty.is_vec() {
            out.push_str(&src[id_start..q]);
            i = q;
            at_stmt_start = false;
            continue;
        }

        let rhs_ty = table.infer_expr_type(rhs_trimmed);

        // Emit `<ident>(.swizzle)? <op>= ` verbatim.
        out.push_str(&src[id_start..op_end]);
        out.push(' ');

        // Coercion handles compound ops on scalar LHS and swizzle-LHS
        // scalars. The AdamFX CollaborationFX preset writes
        // `dx += GetPixel(...) - GetPixel(...);` with `dx: f32` — without
        // this, the `op_byte.is_none()` guard let the vec3 RHS slide
        // through to naga, which rejected the InvalidStoreTypes. The
        // midgitstraights / suksma-neck preset writes `uv2.x = 1 -
        // tex2D(...);` with `uv2: vec2<f32>` — the swizzle LHS path
        // narrows the LHS to f32 and truncates the vec4 RHS to `.x`.
        let coerced = match (lhs_ty, rhs_ty) {
            // vec = scalar: broadcast. Only for plain `=` and `+=`/`-=`
            // (broadcast doesn't change `*=`/`/=` semantics on vec, but
            // we avoid them out of paranoia).
            (l, WgslType::F32) | (l, WgslType::I32)
                if l.is_vec() && matches!(op_byte, None | Some(b'+') | Some(b'-')) =>
            {
                Some(format!("{}({})", l.wgsl_name(), rhs_trimmed))
            }
            // vec_n = vec_m where m > n: truncate.
            (l, r) if l.is_vec() && r.is_vec() && vec_size(r) > vec_size(l) => {
                let sw = match vec_size(l) {
                    2 => ".xy",
                    3 => ".xyz",
                    _ => "",
                };
                Some(format!("({}){}", rhs_trimmed, sw))
            }
            // f32 = vec: truncate with `.x`. Now fires for compound ops
            // too — `dx += vec3` becomes `dx += (vec3).x`, matching the
            // HLSL implicit-truncation semantics.
            (l, r) if l.is_scalar() && r.is_vec() => Some(format!("({}).x", rhs_trimmed)),
            _ => None,
        };

        if let Some(new_rhs) = coerced {
            out.push_str(&new_rhs);
        } else {
            out.push_str(rhs_raw);
        }
        out.push(';');
        i = q + 1;
        at_stmt_start = true;
    }

    out
}

// ---------------------------------------------------------------------------
// Pass: scalar→vec implicit broadcast in declarations + f32←vec truncation
// ---------------------------------------------------------------------------

/// Walk `var X: TYPE = <expr>;` (or `let`) declarations and fix HLSL-
/// implicit-conversion mismatches between LHS type and RHS type:
///
/// - **`f32 = vec`**: rewrite RHS to `(<expr>).x` (HLSL truncation).
/// - **`vec = scalar`**: rewrite RHS to `vecN<f32>(<expr>)` (HLSL broadcast).
///
/// The pass is intentionally conservative: only fires when the LHS type
/// is known and the inferred RHS type is the *other* of the two cases
/// above. Anything ambiguous (`Unknown` RHS) passes through.
///
/// Why this exists: MD2 user shaders frequently write
/// `float3 bg = pow(length(dz), 0.7)*2 + GetBlur1(uv).y*2;` — RHS is
/// scalar, LHS is vec3; HLSL silently broadcasts, WGSL rejects with
/// `the type of bg is expected to be vec3<f32> but got f32`. Mirror
/// case: `float lum = GetPixel(uv) * c.x;` — RHS vec3, LHS f32.
pub fn inject_truncations(src: &str, table: &SymbolTable) -> String {
    let bytes = src.as_bytes();
    let mut out = String::with_capacity(src.len() + 64);
    let mut i = 0;

    while i < bytes.len() {
        // Skip comments to avoid false positives in disabled code blocks.
        if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'/' {
            while i < bytes.len() && bytes[i] != b'\n' {
                out.push(bytes[i] as char);
                i += 1;
            }
            continue;
        }

        let kw_len = match keyword_at(bytes, i, &["var", "let"]) {
            Some(n) => n,
            None => {
                out.push(bytes[i] as char);
                i += 1;
                continue;
            }
        };

        // Try to parse `var/let NAME: f32 = <expr>;`. Anything that
        // doesn't fit this exact shape passes through verbatim.
        let kw_end = i + kw_len;
        let mut p = kw_end;
        while p < bytes.len() && bytes[p].is_ascii_whitespace() {
            p += 1;
        }
        let name_start = p;
        while p < bytes.len() && (bytes[p].is_ascii_alphanumeric() || bytes[p] == b'_') {
            p += 1;
        }
        if p == name_start {
            out.push_str(&src[i..kw_end]);
            i = kw_end;
            continue;
        }
        let after_name = p;
        while p < bytes.len() && bytes[p].is_ascii_whitespace() {
            p += 1;
        }
        if p >= bytes.len() || bytes[p] != b':' {
            out.push_str(&src[i..after_name]);
            i = after_name;
            continue;
        }
        p += 1;
        // Read type until `=` or `;`.
        let ty_start = p;
        while p < bytes.len() && bytes[p] != b'=' && bytes[p] != b';' {
            p += 1;
        }
        let ty_str = src[ty_start..p].trim();
        let lhs_ty = WgslType::from_decl_str(ty_str);
        // Only act on scalar or vec LHS — matrix/bool/unknown pass through.
        if !matches!(
            lhs_ty,
            WgslType::F32 | WgslType::I32 | WgslType::Vec2F | WgslType::Vec3F | WgslType::Vec4F
        ) {
            out.push_str(&src[i..p]);
            i = p;
            continue;
        }
        if p >= bytes.len() || bytes[p] != b'=' {
            // No initialiser — passthrough.
            out.push_str(&src[i..p]);
            i = p;
            continue;
        }
        let eq_at = p;
        p += 1;
        // RHS until `;` at depth 0.
        let rhs_start = p;
        let mut depth_paren = 0i32;
        let mut depth_bracket = 0i32;
        while p < bytes.len() {
            match bytes[p] {
                b'(' => depth_paren += 1,
                b')' => depth_paren -= 1,
                b'[' => depth_bracket += 1,
                b']' => depth_bracket -= 1,
                b';' if depth_paren == 0 && depth_bracket == 0 => break,
                _ => {}
            }
            p += 1;
        }
        if p >= bytes.len() {
            // Unterminated — bail.
            out.push_str(&src[i..]);
            return out;
        }
        let rhs = &src[rhs_start..p];
        let rhs_trimmed = rhs.trim();
        let rhs_ty = table.infer_expr_type(rhs_trimmed);

        // Emit the prefix (`var NAME: TYPE`) verbatim.
        out.push_str(&src[i..=eq_at]);

        match (lhs_ty, rhs_ty) {
            // f32/i32 = vec: truncate with `.x`.
            (l, r) if l.is_scalar() && r.is_vec() => {
                out.push_str(" (");
                out.push_str(rhs_trimmed);
                out.push_str(").x");
            }
            // vec = scalar: broadcast via constructor.
            (l, WgslType::F32) | (l, WgslType::I32) if l.is_vec() => {
                out.push(' ');
                out.push_str(l.wgsl_name());
                out.push('(');
                out.push_str(rhs_trimmed);
                out.push(')');
            }
            // vecN = vecM where N < M: truncate via swizzle.
            (l, r) if l.is_vec() && r.is_vec() && vec_size(r) > vec_size(l) => {
                let swizzle = match vec_size(l) {
                    2 => ".xy",
                    3 => ".xyz",
                    _ => "",
                };
                out.push_str(" (");
                out.push_str(rhs_trimmed);
                out.push(')');
                out.push_str(swizzle);
            }
            _ => {
                out.push_str(rhs);
            }
        }
        out.push(';');
        i = p + 1;
    }

    out
}

// ---------------------------------------------------------------------------
// Pass: swizzle-LHS assignment rewrite
// ---------------------------------------------------------------------------

/// Rewrite `target.<swizzle> <op>= <rhs>;` into a full-vector
/// reconstruction. WGSL refuses multi-component swizzles on the LHS of an
/// assignment (`invalid left-hand side of assignment`); HLSL allows them
/// freely. The dominant cases in `test-presets-200/`:
///
/// - `ret.xy *= diff;`
/// - `ret.xyz = tex2D(...).xyz;`
/// - `ret.zy /= diff2;`
///
/// Rewrite shape (for `target: vec3<f32>`):
///
/// `ret.xy = expr;`  →  `ret = vec3<f32>((expr).x, (expr).y, ret.z);`
/// `ret.xy *= expr;` →  `ret = vec3<f32>(ret.x * (expr).x, ret.y * (expr).y, ret.z);`
///
/// Constraints:
/// - `target` must resolve to a known vec3/vec4 in `table`. Skip otherwise.
/// - Swizzle length must be 2..=target_size, all distinct, all in `{x,y,z,w}`.
/// - Reordering swizzles (`yx`, `zy`) are handled — each target component
///   picks the matching swizzle slot.
/// - Scalar RHS is broadcast (each lane uses the same `(expr)`).
/// - RHS with unknown type is treated as a vec matching the swizzle width
///   (the dominant case in real shaders); the resulting WGSL still
///   validates because positional swizzle access on a vec is well-typed.
pub fn inject_swizzle_assignments(src: &str, table: &SymbolTable) -> String {
    use regex::Regex;
    use std::sync::LazyLock;

    static SWZ_RE: LazyLock<Regex> = LazyLock::new(|| {
        // Line-anchored: only fires at statement start (after whitespace
        // or the wrapper's body prefix). Captures:
        //   1 = indent, 2 = target ident, 3 = swizzle, 4 = op, 5 = rhs
        Regex::new(
            r"(?m)^([\t ]*)([A-Za-z_][A-Za-z0-9_]*)\.([xyzwrgba]{2,4})\s*([+\-*/]?=)\s*([^;]+);",
        )
        .unwrap()
    });

    SWZ_RE
        .replace_all(src, |caps: &regex::Captures| {
            let indent = &caps[1];
            let target = &caps[2];
            let swizzle = &caps[3];
            let op = &caps[4];
            let rhs = caps[5].trim();

            // Normalise rgba/xyzw — accept either MD2 colour spelling but
            // emit canonical xyzw inside fs_main (the wrapper uses xyzw).
            let swizzle_xyzw = normalise_swizzle(swizzle);

            // All components must be unique — duplicate components on the
            // LHS are an HLSL error too; bail rather than fabricate.
            if !all_unique(&swizzle_xyzw) {
                return caps[0].to_string();
            }

            // Lookup target type. Only handle vec3/vec4; bail otherwise.
            let target_ty = match table.lookup(target) {
                Some(WgslType::Vec3F) => WgslType::Vec3F,
                Some(WgslType::Vec4F) => WgslType::Vec4F,
                _ => return caps[0].to_string(),
            };
            let target_size = vec_size(target_ty);

            // Swizzle width must not exceed target's width.
            if swizzle_xyzw.len() > target_size {
                return caps[0].to_string();
            }

            // Compute new value for each component of the target vec.
            let comps = match target_size {
                3 => &['x', 'y', 'z'][..],
                4 => &['x', 'y', 'z', 'w'][..],
                _ => return caps[0].to_string(),
            };
            let rhs_is_scalar = table.infer_expr_type(rhs).is_scalar();
            let mut lane_exprs: Vec<String> = Vec::with_capacity(target_size);
            for &c in comps {
                if let Some(pos) = swizzle_xyzw.iter().position(|&s| s == c) {
                    // Lane in RHS at this swizzle position. The RHS is
                    // indexed positionally (lane 0 = .x, lane 1 = .y,
                    // …) — NOT by the swizzle letter at this position,
                    // which equals the target component `c` itself.
                    let lane_letter = match pos {
                        0 => 'x',
                        1 => 'y',
                        2 => 'z',
                        3 => 'w',
                        _ => return caps[0].to_string(),
                    };
                    let rhs_lane = if rhs_is_scalar || swizzle_xyzw.len() == 1 {
                        // Scalar RHS — broadcast unchanged value.
                        format!("({rhs})")
                    } else {
                        format!("({rhs}).{lane_letter}")
                    };
                    let new_val = match op {
                        "=" => rhs_lane,
                        "+=" => format!("{target}.{c} + {rhs_lane}"),
                        "-=" => format!("{target}.{c} - {rhs_lane}"),
                        "*=" => format!("{target}.{c} * {rhs_lane}"),
                        "/=" => format!("{target}.{c} / {rhs_lane}"),
                        _ => return caps[0].to_string(),
                    };
                    lane_exprs.push(new_val);
                } else {
                    // Component not in swizzle — keep unchanged.
                    lane_exprs.push(format!("{target}.{c}"));
                }
            }

            format!(
                "{indent}{target} = {ty}({args});",
                ty = target_ty.wgsl_name(),
                args = lane_exprs.join(", ")
            )
        })
        .to_string()
}

/// Map `r/g/b/a` to `x/y/z/w` while preserving identity for `x/y/z/w`.
/// MD2 user shaders mix the two conventions freely; the codegen wrapper
/// always uses `xyzw` for its locals, so canonicalising on emit keeps
/// downstream lane lookups stable.
fn normalise_swizzle(s: &str) -> Vec<char> {
    s.chars()
        .map(|c| match c {
            'r' => 'x',
            'g' => 'y',
            'b' => 'z',
            'a' => 'w',
            other => other,
        })
        .collect()
}

fn all_unique(letters: &[char]) -> bool {
    let mut seen = [false; 4];
    for &c in letters {
        let idx = match c {
            'x' => 0,
            'y' => 1,
            'z' => 2,
            'w' => 3,
            _ => return false,
        };
        if seen[idx] {
            return false;
        }
        seen[idx] = true;
    }
    true
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn symbol_table_picks_up_var_decls() {
        let src = "var foo: vec3<f32> = vec3<f32>(0.0); let bar: f32 = 1.0;";
        let t = SymbolTable::from_source(src);
        assert_eq!(t.lookup("foo"), Some(WgslType::Vec3F));
        assert_eq!(t.lookup("bar"), Some(WgslType::F32));
    }

    #[test]
    fn symbol_table_seeds_wrapper_locals() {
        // `uv` is in the wrapper preamble — it must resolve even with no
        // user-side `var` for it.
        let t = SymbolTable::from_source("// nothing here");
        assert_eq!(t.lookup("uv"), Some(WgslType::Vec2F));
        assert_eq!(t.lookup("texsize"), Some(WgslType::Vec4F));
        assert_eq!(t.lookup("q1"), Some(WgslType::F32));
        assert_eq!(t.lookup("M_PI_2"), Some(WgslType::F32));
    }

    #[test]
    fn infer_expr_numeric_literal() {
        let t = SymbolTable::from_source("");
        assert_eq!(t.infer_expr_type("1.0"), WgslType::F32);
        assert_eq!(t.infer_expr_type("-3.14"), WgslType::F32);
        assert_eq!(t.infer_expr_type("0"), WgslType::F32);
    }

    #[test]
    fn infer_expr_known_helpers() {
        let t = SymbolTable::from_source("");
        assert_eq!(t.infer_expr_type("GetPixel(uv)"), WgslType::Vec3F);
        assert_eq!(t.infer_expr_type("GetBlur1(uv)"), WgslType::Vec3F);
        assert_eq!(t.infer_expr_type("lum(ret)"), WgslType::F32);
        assert_eq!(
            t.infer_expr_type("textureSample(sampler_main_texture, sampler_main, uv)"),
            WgslType::Vec4F
        );
    }

    #[test]
    fn infer_expr_constructor() {
        let t = SymbolTable::from_source("");
        assert_eq!(t.infer_expr_type("vec3<f32>(1, 0, 0)"), WgslType::Vec3F);
        assert_eq!(t.infer_expr_type("vec4<f32>(c)"), WgslType::Vec4F);
    }

    #[test]
    fn infer_expr_swizzle_narrows_to_scalar() {
        let t = SymbolTable::from_source("var c: vec4<f32> = vec4<f32>(1);");
        assert_eq!(t.infer_expr_type("c.x"), WgslType::F32);
        assert_eq!(t.infer_expr_type("c.xy"), WgslType::Vec2F);
        assert_eq!(t.infer_expr_type("c.xyz"), WgslType::Vec3F);
    }

    #[test]
    fn infer_expr_binop_widens_to_vec() {
        let t = SymbolTable::from_source("var c: vec4<f32> = vec4<f32>(1);");
        // `GetPixel(uv) * c.x` → vec3 * f32 = vec3.
        assert_eq!(t.infer_expr_type("GetPixel(uv) * c.x"), WgslType::Vec3F);
    }

    #[test]
    fn broadcast_clamp_with_scalar_bounds() {
        let t = SymbolTable::from_source("");
        let src = "ret = clamp(GetBlur1(uv), 0.0, 1.0);";
        let out = inject_broadcasts(src, &t);
        assert!(
            out.contains("clamp(GetBlur1(uv), vec3<f32>(0.0), vec3<f32>(1.0))"),
            "got: {out}"
        );
    }

    #[test]
    fn broadcast_pow_with_scalar_exponent() {
        let t = SymbolTable::from_source("");
        let src = "ret = pow(GetPixel(uv), 0.5);";
        let out = inject_broadcasts(src, &t);
        assert!(
            out.contains("pow(GetPixel(uv), vec3<f32>(0.5))"),
            "got: {out}"
        );
    }

    #[test]
    fn broadcast_mix_with_scalar_lerp_factor() {
        let t = SymbolTable::from_source("");
        let src = "ret = mix(a, b, 0.3);";
        // Without context (a, b unknown), no rewrite — the pass is
        // conservative.
        let out = inject_broadcasts(src, &t);
        assert_eq!(out, src);

        // Now with vec3 a and b in scope:
        let t = SymbolTable::from_source(
            "var a: vec3<f32> = vec3<f32>(0); var b: vec3<f32> = vec3<f32>(1);",
        );
        let out = inject_broadcasts(src, &t);
        assert!(out.contains("mix(a, b, vec3<f32>(0.3))"), "got: {out}");
    }

    #[test]
    fn broadcast_truncates_larger_vec_to_smaller() {
        // Real preset pattern: `ret = max(ret_vec3, tex2D(...))`. HLSL
        // truncates the vec4 result to vec3; WGSL needs an explicit `.xyz`.
        let t = SymbolTable::from_source("var ret: vec3<f32> = vec3<f32>(0);");
        let src = "ret = max(ret, textureSample(t, s, uv));";
        let out = inject_broadcasts(src, &t);
        assert!(
            out.contains("max(ret, (textureSample(t, s, uv)).xyz)"),
            "got: {out}"
        );
    }

    #[test]
    fn broadcast_skipped_when_all_args_scalar() {
        let t = SymbolTable::from_source("");
        let src = "var x: f32 = clamp(0.5, 0.0, 1.0);";
        let out = inject_broadcasts(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn truncation_f32_eq_vec3_inserts_dot_x() {
        let t = SymbolTable::from_source("");
        let src = "var gx1: f32 = GetPixel(uv) + GetBlur1(uv);";
        let out = inject_truncations(src, &t);
        assert!(
            out.contains("var gx1: f32 = (GetPixel(uv) + GetBlur1(uv)).x;"),
            "got: {out}"
        );
    }

    #[test]
    fn truncation_skipped_when_rhs_already_scalar() {
        let t = SymbolTable::from_source("");
        let src = "var x: f32 = 1.0 + 2.0;";
        let out = inject_truncations(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn truncation_skipped_when_lhs_is_vec() {
        let t = SymbolTable::from_source("");
        let src = "var v: vec3<f32> = GetPixel(uv);";
        let out = inject_truncations(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn broadcast_vec3_eq_scalar_wraps_in_constructor() {
        // Real preset pattern: `float3 bg = pow(length(dz), 0.7)*2 + GetBlur1(uv).y*2;`
        // RHS is f32, LHS is vec3 — must inject `vec3<f32>(<rhs>)`.
        let t = SymbolTable::from_source("var dz: vec3<f32> = vec3<f32>(0);");
        let src = "var bg: vec3<f32> = pow(length(dz), 0.7)*2 + GetBlur1(uv).y*2;";
        let out = inject_truncations(src, &t);
        assert!(
            out.contains(
                "var bg: vec3<f32> = vec3<f32>(pow(length(dz), 0.7)*2 + GetBlur1(uv).y*2);"
            ),
            "got: {out}"
        );
    }

    #[test]
    fn truncation_vec3_eq_vec4_appends_xyz() {
        // Real preset pattern: `float3 ret2 = tex2D(sampler_main, uv);`
        // where tex2D returns vec4 — must inject `.xyz`.
        let t = SymbolTable::from_source("");
        let src = "var ret2: vec3<f32> = textureSample(t, s, uv);";
        let out = inject_truncations(src, &t);
        assert!(
            out.contains("var ret2: vec3<f32> = (textureSample(t, s, uv)).xyz;"),
            "got: {out}"
        );
    }

    #[test]
    fn truncation_vec2_eq_vec4_appends_xy() {
        let t = SymbolTable::from_source("");
        let src = "var sam: vec2<f32> = textureSample(t, s, uv);";
        let out = inject_truncations(src, &t);
        assert!(
            out.contains("var sam: vec2<f32> = (textureSample(t, s, uv)).xy;"),
            "got: {out}"
        );
    }

    #[test]
    fn broadcast_skipped_when_rhs_already_vec() {
        let t = SymbolTable::from_source("");
        let src = "var v: vec3<f32> = GetPixel(uv);";
        let out = inject_truncations(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn truncation_skipped_on_unknown_rhs_type() {
        let t = SymbolTable::from_source("");
        // RHS references an unknown identifier — we can't be sure of the
        // type, so don't touch it.
        let src = "var x: f32 = some_user_function(uv);";
        let out = inject_truncations(src, &t);
        assert_eq!(out, src);
    }

    // ----------------------------------------------------------------
    // Swizzle-LHS assignment rewrite
    // ----------------------------------------------------------------

    #[test]
    fn swizzle_xy_assignment_on_vec3_target() {
        // `ret` is vec3<f32> per the wrapper preamble. `ret.xy = expr;`
        // rebuilds the full vec3 with the .x/.y lanes from expr and the
        // .z lane unchanged.
        let t = SymbolTable::from_source("");
        let src = "ret.xy = diff;";
        let out = inject_swizzle_assignments(src, &t);
        assert_eq!(
            out, "ret = vec3<f32>((diff).x, (diff).y, ret.z);",
            "got: {out}"
        );
    }

    #[test]
    fn swizzle_xyz_full_replace_on_vec3_target() {
        let t = SymbolTable::from_source("");
        let src = "ret.xyz = tex2D(s, uv).xyz;";
        let out = inject_swizzle_assignments(src, &t);
        // Whole vec replaced — `.z` lane sourced from `(rhs).z`, not the
        // unchanged target.
        assert_eq!(
            out,
            "ret = vec3<f32>((tex2D(s, uv).xyz).x, (tex2D(s, uv).xyz).y, (tex2D(s, uv).xyz).z);",
        );
    }

    #[test]
    fn swizzle_compound_mul_assign() {
        let t = SymbolTable::from_source("");
        let src = "ret.xy *= diff;";
        let out = inject_swizzle_assignments(src, &t);
        // Compound: each lane = target.<comp> * rhs.<lane>.
        assert_eq!(
            out,
            "ret = vec3<f32>(ret.x * (diff).x, ret.y * (diff).y, ret.z);",
        );
    }

    #[test]
    fn swizzle_reordered_zy() {
        // `ret.zy = expr;` puts expr.x into ret.z and expr.y into ret.y.
        let t = SymbolTable::from_source("");
        let src = "ret.zy = pair;";
        let out = inject_swizzle_assignments(src, &t);
        // x unchanged, y = pair.y, z = pair.x.
        assert_eq!(out, "ret = vec3<f32>(ret.x, (pair).y, (pair).x);",);
    }

    #[test]
    fn swizzle_rgba_normalised_to_xyzw() {
        let t = SymbolTable::from_source("");
        let src = "ret.rg = pair;";
        let out = inject_swizzle_assignments(src, &t);
        // `rg` maps to `xy`. Same emit shape as `ret.xy`.
        assert_eq!(out, "ret = vec3<f32>((pair).x, (pair).y, ret.z);",);
    }

    #[test]
    fn swizzle_skipped_when_target_unknown() {
        // Unknown identifier — must NOT rewrite.
        let t = SymbolTable::from_source("");
        let src = "user_var.xy = stuff;";
        let out = inject_swizzle_assignments(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn swizzle_skipped_on_single_component() {
        // Single-component swizzles like `ret.x = foo;` are valid WGSL —
        // leave them alone.
        let t = SymbolTable::from_source("");
        let src = "ret.z = stuff;";
        let out = inject_swizzle_assignments(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn swizzle_skipped_when_duplicate_components() {
        // `ret.xx = foo` is meaningless on the LHS — bail.
        let t = SymbolTable::from_source("");
        let src = "ret.xx = pair;";
        let out = inject_swizzle_assignments(src, &t);
        assert_eq!(out, src);
    }

    #[test]
    fn swizzle_div_compound_on_vec3() {
        let t = SymbolTable::from_source("");
        let src = "ret.zy /= diff2;";
        let out = inject_swizzle_assignments(src, &t);
        // x unchanged, y = ret.y / diff2.y, z = ret.z / diff2.x.
        assert_eq!(
            out,
            "ret = vec3<f32>(ret.x, ret.y / (diff2).y, ret.z / (diff2).x);",
        );
    }

    #[test]
    fn swizzle_xy_on_vec4_target() {
        // vec4 case: `ret4: vec4<f32>`, swizzle `xy`, only x/y change.
        let t = SymbolTable::from_source("var ret4: vec4<f32> = vec4<f32>(0.0);");
        let src = "ret4.xy = pair;";
        let out = inject_swizzle_assignments(src, &t);
        assert_eq!(out, "ret4 = vec4<f32>((pair).x, (pair).y, ret4.z, ret4.w);",);
    }
}