onedrop_hlsl/rewrite/

mod.rs

//! AST-driven HLSL→HLSL rewrites.
//!
//! Operates on the HLSL source *before* the regex pipeline in
//! [`crate::translate_shader`]. Each pass parses the source with
//! [`crate::parse::parse_hlsl`], walks the AST to spot a specific HLSL
//! idiom that the downstream regex passes can't recover from, and emits a
//! list of [`TextEdit`]s applied back to the source.
//!
//! Falls back silently when the parser can't consume the input — the
//! original source is returned unchanged so the regex pipeline still has a
//! chance. The corpus parse rate is 100 % on `test-presets-200/`, so the
//! fallback is exercised only by malformed test inputs in practice.
//!
//! Each pass lives in its own submodule and is composed in order by
//! [`apply_all`]. The shared [`WalkCtx`] and type-inference helpers live
//! here in `mod.rs`; private items in this file are accessible to the
//! pass submodules under Rust's descendant-visibility rules.

use crate::ast::*;
use crate::lex::Span;
use crate::parse::parse_hlsl;
use crate::types::{SymbolTable, WgslType, vec_of_size, vec_size};

mod array_lower;
mod bare_expr;
mod binop_vec;
mod bool_arith;
mod brace_init;
mod chained_init;
mod comma_paren;
mod embedded_assign;
mod for_init;
mod modf_arity;
mod qa_stub;
mod scalar_swizzle;
mod swizzle_assign;
mod ternary;
mod texture_uv;
mod user_fn;
mod vec_cmp;

#[cfg(test)]
mod tests;

// Re-export the entry point of each pass so callers (and the tests
// module) see them under the `rewrite::` path the original monolithic
// file exposed.
#[cfg(test)]
pub(super) use array_lower::{lower_array_globals, lower_local_arrays};
pub(super) use binop_vec::rewrite_binary_vec_mismatches;
pub(super) use bool_arith::rewrite_bool_to_float;
pub(super) use brace_init::rewrite_brace_init;
pub(super) use chained_init::rewrite_chained_assign_inits;
pub(super) use comma_paren::rewrite_comma_paren;
pub(super) use embedded_assign::rewrite_embedded_assigns;
#[cfg(test)]
pub(super) use for_init::rewrite_for_int_init;
pub(super) use modf_arity::rewrite_modf_arity;
pub(super) use qa_stub::inject_qa_stub;
pub(super) use scalar_swizzle::rewrite_scalar_swizzle;
pub(super) use swizzle_assign::rewrite_swizzle_assigns;
pub(super) use ternary::rewrite_ternary_to_select;
pub(super) use texture_uv::coerce_texture_uv_args;
pub(super) use user_fn::coerce_user_fn_args;
pub(super) use vec_cmp::rewrite_vec_scalar_compare;

use array_lower::{collect_global_array_edits, collect_local_array_edits};
use bare_expr::collect_bare_expr_stmt_edits;
use for_init::collect_for_int_edits;

/// Apply every AST-driven rewrite pass in series. Returns the rewritten
/// source. If parsing fails at any step, returns the source unchanged from
/// that step onward — the pipeline is best-effort.
///
/// Pass order matters: the HLSL-preserving passes (stub injection, binop
/// truncation, UV coercion) run first so each one can re-parse the
/// in-flight source. The "lowering" passes (array globals, local arrays,
/// `for (int …)` init) emit WGSL-shaped syntax that the parser can't
/// re-consume, so they all share **one** parse via [`apply_lowerings`] and
/// commit their edits together at the end.
pub fn apply_all(src: &str) -> String {
    // `(a, b, c)` (C-style comma operator in paren-expression position) —
    // HLSL evaluates each and yields the rightmost; WGSL has no comma op
    // and refuses the syntax. Token-driven strip of `a, b,` so the parens
    // hold only the rightmost. Runs first because every later pass sees
    // a cleaner source — no comma-op surprises in expression positions.
    let s = rewrite_comma_paren(src);
    // `float2x2 rot = { a, b, c, d };` — HLSL brace-init of vec/mat
    // locals. WGSL has no brace-init form for these types and the
    // downstream local-decl regex skips any decl containing `{` (to avoid
    // greedy-swallowing function bodies), so the rewrite must run as an
    // AST pass. Lowers to `T x = T(a, b, c, d);`.
    let s = rewrite_brace_init(&s);
    // HLSL `modf(x, out)` is the two-arg form (writes int part to out,
    // returns fract). WGSL `modf(x)` returns a struct; the two-arg form
    // is a parse error. Lift to a `let` temp + two assigns. Runs after
    // brace_init / chained_init / comma_paren so the rewritten source
    // still parses for the later AST passes.
    let s = rewrite_modf_arity(&s);
    // `T x = y = expr;` (chained-assign in init position) — HLSL accepts,
    // WGSL has no assignment expression. Lower to `y = expr; T x = y;`
    // first so every later pass and the downstream regex pipeline see
    // valid-shape declarations.
    let s = rewrite_chained_assign_inits(&s);
    // `lerp(a, tmp = expr, b)` and similar — HLSL assignment-as-arg shapes.
    // Lift the side effect to a preceding statement so the downstream
    // regex pipeline doesn't see `tmp = expr` in a position WGSL refuses.
    let s = rewrite_embedded_assigns(&s);
    let s = inject_qa_stub(&s);
    // Ternary `cond ? a : b` doesn't exist in WGSL. Lower to
    // `select(b, a, cond)` early so subsequent passes don't trip on the `?`
    // token (the regex pipeline doesn't model it; type-aware passes give
    // up). Emit replacements as HLSL-shaped `select(...)` calls so later
    // passes can still re-parse the source.
    let s = rewrite_ternary_to_select(&s);
    // Multi-component swizzle assigns (`uv.xy += rhs`) are illegal in WGSL.
    // Lower them to a per-component sequence before any pass that re-parses
    // (no later pass tries to re-recognise the original `lhs.SWIZ = rhs`
    // shape — it's purely about getting valid LHSs to naga).
    let s = rewrite_swizzle_assigns(&s);
    let s = rewrite_binary_vec_mismatches(&s);
    let s = coerce_texture_uv_args(&s);
    let s = coerce_user_fn_args(&s);
    // `vec3 >= 0.1` and friends. HLSL broadcasts the scalar for
    // comparisons; WGSL refuses. Must run before `rewrite_bool_to_float`
    // so the wrap below sees a well-typed `vec3<bool>` comparison
    // (otherwise its select() ends up wrapping a still-invalid mixed-width
    // operand).
    let s = rewrite_vec_scalar_compare(&s);
    // `float mask = a > b;` and `float3 x = (vec_a > vec_b);` need the
    // comparison wrapped in `select(0.0, 1.0, …)` since WGSL doesn't coerce
    // bool↔f32. Runs late so it sees the post-typed source.
    let s = rewrite_bool_to_float(&s);
    // `<scalar>.xxx` is HLSL broadcast syntax; WGSL doesn't accept it.
    // Rewrite to a constructor `float3(<scalar>)` (regex pipeline later
    // maps `float3` → `vec3<f32>`).
    let s = rewrite_scalar_swizzle(&s);
    apply_lowerings(&s)
}

/// Collect every "WGSL-output" lowering edit against a **single** parse of
/// `src`, then apply them in one pass. Splitting these across separate
/// pipeline steps would corrupt the second step: once the first emits
/// `var name: array<…>` (or `var i: i32`), the source is no longer valid
/// HLSL and a fresh `parse_hlsl` call returns `Err`, silently skipping the
/// later passes. The for-int init rewrite once turned that latent hazard
/// into a real regression on the Stahlregen + suksma `dotes` pair (the
/// for-int rewrite ran first, then `lower_array_globals` couldn't
/// re-parse). The combined pass below avoids the order trap.
fn apply_lowerings(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut edits = Vec::new();
    collect_global_array_edits(&tu, src, &mut edits);
    if let Some(body) = &tu.shader_body {
        collect_local_array_edits(body, src, &mut edits);
        collect_for_int_edits(body, &mut edits);
        collect_bare_expr_stmt_edits(body, src, &mut edits);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            collect_local_array_edits(&f.body, src, &mut edits);
            collect_for_int_edits(&f.body, &mut edits);
            collect_bare_expr_stmt_edits(&f.body, src, &mut edits);
        }
    }
    apply_edits(src, &mut edits)
}

// ---------------------------------------------------------------------------
// Text edit machinery
// ---------------------------------------------------------------------------

/// One textual replacement on the original source. `start..end` are byte
/// offsets; `replacement` is the new text.
#[derive(Debug, Clone)]
struct TextEdit {
    start: u32,
    end: u32,
    replacement: String,
}

/// Apply a list of non-overlapping edits in order. Sorts by start position
/// descending so each application doesn't invalidate the remaining offsets.
/// Edits that overlap are dropped silently — the AST walker is responsible
/// for ensuring non-overlap by emitting one edit per syntactic position.
fn apply_edits(src: &str, edits: &mut [TextEdit]) -> String {
    if edits.is_empty() {
        return src.to_string();
    }
    edits.sort_by(|a, b| b.start.cmp(&a.start));
    let mut out = src.to_string();
    let mut last_start = u32::MAX;
    for e in edits.iter() {
        if e.end > last_start {
            continue; // overlap — skip
        }
        let s = e.start as usize;
        let n = e.end as usize;
        out.replace_range(s..n, &e.replacement);
        last_start = e.start;
    }
    out
}

/// AST walker state: a small stacked symbol table plus a sink for text
/// edits. Each `walk_*` consults and updates the table; `walk_expr` returns
/// the inferred type of the expression so callers can decide whether to
/// emit a fix-up edit.
struct WalkCtx<'a> {
    src: &'a str,
    scopes: Vec<std::collections::HashMap<String, WgslType>>,
    edits: Vec<TextEdit>,
}

impl<'a> WalkCtx<'a> {
    fn new(src: &'a str) -> Self {
        // Bootstrap from the existing SymbolTable so the wrapper preludes
        // (`uv`, `ret`, `roam_sin`, `q1`, …) are recognised. We then drop
        // the SymbolTable and copy its locals into the bottom scope.
        let seed = SymbolTable::from_source("");
        let mut ctx = Self {
            src,
            scopes: vec![Default::default()],
            edits: Vec::new(),
        };
        for (k, v) in seed.locals {
            ctx.scopes[0].insert(k, v);
        }
        ctx
    }

    fn seed_globals(&mut self, tu: &TranslationUnit) {
        for item in &tu.items {
            if let Item::GlobalVar(g) = item {
                let ty = if g.array_len.is_some() {
                    // Arrays don't broadcast as vecs; mark Unknown so
                    // the binop pass doesn't try to coerce them.
                    WgslType::Unknown
                } else {
                    type_from_typeref(&g.ty)
                };
                self.declare(&g.name, ty);
            }
        }
    }

    fn scope_push(&mut self) {
        self.scopes.push(Default::default());
    }
    fn scope_pop(&mut self) {
        self.scopes.pop();
    }
    fn declare(&mut self, name: &str, ty: WgslType) {
        let last = self.scopes.last_mut().expect("at least one scope");
        last.insert(name.to_string(), ty);
    }
    fn lookup(&self, name: &str) -> WgslType {
        for s in self.scopes.iter().rev() {
            if let Some(t) = s.get(name) {
                return *t;
            }
        }
        WgslType::Unknown
    }

    /// Slice of original source between span bounds.
    fn slice(&self, span: Span) -> &'a str {
        &self.src[span.start as usize..span.end as usize]
    }

    fn emit_truncation(&mut self, span: Span, target_size: usize) {
        // `1` collapses a vec down to scalar via `.x` — used by the
        // user-fn arg coercion when a preset feeds a vec3 to a helper
        // declared `float lavcol(float t)`.
        let swizzle = match target_size {
            1 => ".x",
            2 => ".xy",
            3 => ".xyz",
            _ => return,
        };
        let text = self.slice(span);
        // Avoid double-wrapping if the text already ends with the target
        // swizzle (idempotency on repeat runs).
        if text.ends_with(swizzle) {
            return;
        }
        // Split wrap into two zero-length insertions so any inner edits
        // inside `span` can coexist. An earlier implementation used a
        // single `replace(start..end, "(text).xyz")` edit, which got
        // dropped by [`apply_edits`]'s overlap guard whenever a nested
        // binop or UV truncation already emitted within the same range
        // (real example: `.15*tex2D(noise, uv2*ret)` — the inner
        // `uv2*ret → uv2*ret.xy` edit shadowed the outer wrap to `.xyz`
        // of the whole `.15*tex2D(...)` subexpression, leaving the
        // goody+martin preset with a vec4×vec3 multiply naga rejected).
        // Two zero-length inserts at the operand's start and end don't
        // overlap any range-based inner edit, so the wrap survives.
        if needs_parens_for_swizzle(text) {
            self.edits.push(TextEdit {
                start: span.start,
                end: span.start,
                replacement: "(".to_string(),
            });
            self.edits.push(TextEdit {
                start: span.end,
                end: span.end,
                replacement: format!("){swizzle}"),
            });
        } else {
            // Bare ident / parenthesised whole — no extra parens needed.
            // Single zero-length append at end is equivalent to the old
            // replace-range edit (the original text is preserved by the
            // range we don't touch).
            self.edits.push(TextEdit {
                start: span.end,
                end: span.end,
                replacement: swizzle.to_string(),
            });
        }
    }

    /// Wrap a vec2 expression in `vec3<f32>(<expr>, 0.0)` so it satisfies
    /// a helper that demands a vec3 (real example: MD2 presets calling
    /// `lum(roam_sin.yx)` where the runtime padded implicitly). Emitted as
    /// two zero-length insertions instead of one replace-range edit so any
    /// inner edits inside the arg (e.g. a swizzle rewrite) don't get
    /// dropped by [`apply_edits`]'s overlap guard.
    fn emit_pad_vec2_to_vec3(&mut self, span: Span) {
        self.edits.push(TextEdit {
            start: span.start,
            end: span.start,
            replacement: "vec3<f32>(".to_string(),
        });
        self.edits.push(TextEdit {
            start: span.end,
            end: span.end,
            replacement: ", 0.0)".to_string(),
        });
    }

    /// Wrap a scalar expression in `vec2<f32>(<expr>)` so it satisfies a
    /// 2D-texture coord parameter. MD2 presets occasionally pass a single
    /// channel (`tex2D(sampler_noise_hq, uv.x)`) where naga demands a
    /// vec2; HLSL silently broadcast in that situation. Emitted as two
    /// zero-length insertions, same overlap-safety rationale as
    /// [`Self::emit_pad_vec2_to_vec3`].
    fn emit_scalar_to_vec2(&mut self, span: Span) {
        self.edits.push(TextEdit {
            start: span.start,
            end: span.start,
            replacement: "vec2<f32>(".to_string(),
        });
        self.edits.push(TextEdit {
            start: span.end,
            end: span.end,
            replacement: ")".to_string(),
        });
    }
}

fn needs_parens_for_swizzle(s: &str) -> bool {
    // A bare identifier or a parenthesised whole doesn't need extra
    // parens. Any operator at the top level does. We conservatively
    // assume any non-ident / non-call shape needs parens.
    let t = s.trim();
    if t.is_empty() {
        return false;
    }
    // Already parenthesised whole?
    if t.starts_with('(') && t.ends_with(')') && balanced(t) {
        return false;
    }
    // Bare identifier?
    if t.chars().all(|c| c.is_alphanumeric() || c == '_') {
        return false;
    }
    // A member/swizzle chain `foo.bar.baz` is OK to append `.xyz` to.
    if is_ident_chain(t) {
        return false;
    }
    true
}

fn is_ident_chain(t: &str) -> bool {
    let mut started = false;
    for c in t.chars() {
        if c == '.' {
            if !started {
                return false;
            }
            started = false;
        } else if c.is_alphanumeric() || c == '_' {
            started = true;
        } else {
            return false;
        }
    }
    started
}

fn balanced(t: &str) -> bool {
    // Quick check: parens balance such that the whole text is wrapped by
    // the outermost pair (not just two separate paren groups).
    let mut depth = 0i32;
    for (i, c) in t.char_indices() {
        match c {
            '(' => depth += 1,
            ')' => {
                depth -= 1;
                if depth == 0 && i != t.len() - 1 {
                    return false;
                }
            }
            _ => {}
        }
    }
    depth == 0
}

fn type_from_typeref(t: &TypeRef) -> WgslType {
    match t.name.as_str() {
        "float" => WgslType::F32,
        "int" => WgslType::I32,
        "bool" => WgslType::Bool,
        "float2" | "vec2" => WgslType::Vec2F,
        "float3" | "vec3" => WgslType::Vec3F,
        "float4" | "vec4" => WgslType::Vec4F,
        "float2x2" | "mat2x2" => WgslType::Mat2F,
        "float3x3" | "mat3x3" => WgslType::Mat3F,
        "float4x4" | "mat4x4" => WgslType::Mat4F,
        "float1" => WgslType::F32, // HLSL 1-vec, rare
        _ => WgslType::Unknown,
    }
}

fn widen_type(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,
    }
}

fn constructor_return(callee: &str) -> Option<WgslType> {
    Some(match callee {
        "float" => WgslType::F32,
        "int" => WgslType::I32,
        "float2" | "vec2" => WgslType::Vec2F,
        "float3" | "vec3" => WgslType::Vec3F,
        "float4" | "vec4" => WgslType::Vec4F,
        _ => return None,
    })
}

fn builtin_return(callee: &str, args: &[Expr], ctx: &mut WalkCtx) -> WgslType {
    match callee {
        // HLSL is case-insensitive on builtins; the corpus mixes `tex2D`
        // and `tex2d` freely (the regex pass normalises to `tex2D` later
        // but that runs AFTER these AST passes). Without the lowercase
        // forms here `builtin_return` returned Unknown for `tex2d(...)`,
        // which disabled the binop walker's vec4↔vec3 truncation on
        // `ret = tex2d(...) - ret;` — by far the dominant residual
        // InvalidBinaryOperandTypes cluster on the 2000-sample.
        "tex2D" | "tex2d" | "tex3D" | "tex3d" | "textureSample" => WgslType::Vec4F,
        "GetPixel" => WgslType::Vec3F,
        "GetBlur1" | "GetBlur2" | "GetBlur3" => WgslType::Vec3F,
        "lum" => WgslType::F32,
        "length" | "dot" | "distance" => WgslType::F32,
        "sin" | "cos" | "tan" | "asin" | "acos" | "atan" | "atan2" | "sqrt" | "abs" | "sign"
        | "floor" | "ceil" | "fract" | "exp" | "log" | "saturate" | "frac" => {
            // Same type as arg.
            if let Some(a) = args.first() {
                return infer_type(a, ctx);
            }
            WgslType::F32
        }
        "clamp" | "min" | "max" | "mix" | "smoothstep" | "step" | "pow" | "lerp" => {
            // Widen across all args.
            let mut t = WgslType::Unknown;
            for a in args {
                t = widen_type(t, infer_type(a, ctx));
            }
            t
        }
        "normalize" => {
            if let Some(a) = args.first() {
                return infer_type(a, ctx);
            }
            WgslType::Unknown
        }
        // HLSL's `mul(vec, mat)` and `mul(mat, vec)` returns a vector of
        // the input vec's shape. Without this, corpus expressions like
        // `mul(uv, mat2x2) + ret1*.1` infer the LHS as `Unknown`, which
        // disables `binop_vec`'s size-mismatch truncation and surfaces as
        // `InvalidBinaryOperandTypes Add(vec2, vec3)` at validate time.
        // Matrix constructors are `mat2x2`/`mat3x3`/`mat4x4` (the
        // translator's downstream type form).
        "mul" => {
            // Return the vector-shaped argument's type. If both are vecs,
            // pick the first; if neither is, widen them.
            let mut found_vec: WgslType = WgslType::Unknown;
            for a in args {
                let t = infer_type(a, ctx);
                if t.is_vec() {
                    found_vec = t;
                    break;
                }
            }
            if matches!(found_vec, WgslType::Unknown) {
                let mut t = WgslType::Unknown;
                for a in args {
                    t = widen_type(t, infer_type(a, ctx));
                }
                t
            } else {
                found_vec
            }
        }
        // Matrix constructors emit a matrix; downstream `binop_vec`
        // only triggers when both operands are vecs, so returning a
        // matrix type here keeps `vec * mat` quiet (it will not match
        // the both-is-vec truncation guard). Currently we don't carry
        // distinct matrix dimensions through `WgslType`; treat the call
        // as opaque (Unknown) — the binop walker already returns
        // `widen_type(Vec2F, Unknown) = Vec2F` which is the actual
        // result of `vec2 * mat2x2`, so the downstream truncation
        // proceeds correctly.
        "mat2x2" | "mat3x3" | "mat4x4" | "float2x2" | "float3x3" | "float4x4" => WgslType::Unknown,
        _ => WgslType::Unknown,
    }
}

/// Type-only version of `walk_expr` — doesn't emit edits. Used by
/// builtin-return inference where we want the type of an arg without
/// double-walking it.
fn infer_type(e: &Expr, ctx: &mut WalkCtx) -> WgslType {
    match e {
        Expr::Lit(l) => match l.value {
            LitValue::Int(_) | LitValue::Float(_) => WgslType::F32,
            LitValue::Bool(_) => WgslType::Bool,
        },
        Expr::Ident(name, _) => ctx.lookup(name),
        Expr::Swizzle(s) => {
            let base = infer_type(&s.base, ctx);
            if base.is_vec() {
                vec_of_size(s.components.len())
            } else {
                WgslType::Unknown
            }
        }
        Expr::Call(c) => {
            if let Some(t) = constructor_return(&c.callee) {
                return t;
            }
            builtin_return(&c.callee, &c.args, ctx)
        }
        Expr::Unary(u) => infer_type(&u.operand, ctx),
        Expr::Binary(b) => widen_type(infer_type(&b.lhs, ctx), infer_type(&b.rhs, ctx)),
        _ => WgslType::Unknown,
    }
}

/// Stand-alone type inference that doesn't emit edits — used by the bool
/// pass to peek at an Assign's LHS type before deciding whether to wrap.
fn infer_type_static(e: &Expr, ctx: &mut WalkCtx) -> WgslType {
    match e {
        Expr::Ident(name, _) => ctx.lookup(name),
        Expr::Lit(l) => match l.value {
            LitValue::Int(_) | LitValue::Float(_) => WgslType::F32,
            LitValue::Bool(_) => WgslType::Bool,
        },
        Expr::Swizzle(s) => {
            let base = infer_type_static(&s.base, ctx);
            if base.is_vec() || matches!(base, WgslType::F32 | WgslType::I32) {
                vec_of_size(s.components.len())
            } else {
                WgslType::Unknown
            }
        }
        Expr::Call(c) => {
            if let Some(t) = constructor_return(&c.callee) {
                return t;
            }
            builtin_return(&c.callee, &c.args, ctx)
        }
        Expr::Binary(b) => widen_type(
            infer_type_static(&b.lhs, ctx),
            infer_type_static(&b.rhs, ctx),
        ),
        _ => WgslType::Unknown,
    }
}