onedrop_hlsl/rewrite/

vec_cmp.rs

//! Pass: broadcast scalar operand of vec/scalar comparisons.

use super::bool_arith::decl_target_float_kind;
use super::*;

// ---------------------------------------------------------------------------
// Pass: broadcast scalar operand of vec/scalar comparisons
// ---------------------------------------------------------------------------

/// HLSL auto-broadcasts the scalar operand of a comparison
/// (`vec3 >= 0.1` → element-wise `vec3 >= float3(0.1)`); WGSL rejects
/// `vec3<f32> >= f32` with `InvalidBinaryOperandTypes` (a dominant
/// validate-stage cluster, op=Ge/Le/Gt/Lt/Eq/Ne with
/// `lhs_type: Vector`). Wraps the scalar operand in a `vecN<f32>(…)`
/// constructor so the comparison evaluates as a vector-bool.
///
/// Stays conservative: only fires when one side is a *known* vector
/// and the other is a *known* scalar f32. `(a < b) <= c` chained
/// comparisons (where one side is already bool) are left alone — they
/// belong to the eval long-tail, not this cluster.
pub(crate) fn rewrite_vec_scalar_compare(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut ctx = WalkCtx::new(src);
    // Seed globals with the downstream-widened type (int → f32) rather
    // than the literal HLSL type, mirroring `walk_stmt_for_vsc`'s
    // body-decl path. `static const int anz = 35;` parses as a global, and
    // after `replace_types` widens to `f32`; without this, the body's
    // `n <= anz` lookup would see anz as I32 and the cast wrap wouldn't
    // fire.
    for item in &tu.items {
        if let Item::GlobalVar(g) = item {
            let ty = if g.array_len.is_some() {
                WgslType::Unknown
            } else {
                decl_target_float_kind(&g.ty)
            };
            ctx.declare(&g.name, ty);
        }
    }
    if let Some(body) = &tu.shader_body {
        walk_block_for_vsc(body, &mut ctx);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            ctx.scope_push();
            for p in &f.params {
                ctx.declare(&p.name, decl_target_float_kind(&p.ty));
            }
            walk_block_for_vsc(&f.body, &mut ctx);
            ctx.scope_pop();
        }
    }
    apply_edits(src, &mut ctx.edits)
}

fn walk_block_for_vsc(b: &Block, ctx: &mut WalkCtx) {
    ctx.scope_push();
    for s in &b.stmts {
        walk_stmt_for_vsc(s, ctx);
    }
    ctx.scope_pop();
}

fn walk_stmt_for_vsc(s: &Stmt, ctx: &mut WalkCtx) {
    match s {
        Stmt::LocalDecl(d) => {
            // Most `int x = …;` decls become `var x: f32 = …;` after
            // `replace_types` widens int → f32, so the broadcast decision
            // must reflect the *downstream* type. The exception is
            // `for (int n = …; …)` loop counters, which
            // `collect_for_int_edits` keeps as `i32`. The `For` arm below
            // handles those specially.
            ctx.declare(&d.name, decl_target_float_kind(&d.ty));
            if let Some(init) = &d.init {
                walk_expr_for_vsc(init, ctx);
            }
        }
        Stmt::Assign(a) => {
            walk_expr_for_vsc(&a.target, ctx);
            walk_expr_for_vsc(&a.value, ctx);
        }
        Stmt::Expr(e) => walk_expr_for_vsc(e, ctx),
        Stmt::If(i) => {
            walk_expr_for_vsc(&i.cond, ctx);
            walk_stmt_for_vsc(&i.then_branch, ctx);
            if let Some(e) = &i.else_branch {
                walk_stmt_for_vsc(e, ctx);
            }
        }
        Stmt::While(w) => {
            walk_expr_for_vsc(&w.cond, ctx);
            walk_stmt_for_vsc(&w.body, ctx);
        }
        Stmt::For(f) => {
            ctx.scope_push();
            if let Some(init) = &f.init {
                // For-loop counters keep their declared HLSL type (i32 for
                // `int`, f32 for `float`), unlike the body-decl `int → f32`
                // widening. Declare with `type_from_typeref` so a `for (int n
                // = 1; n <= anz; …)` comparison sees lhs=I32 against
                // rhs=F32 and the cast wrap fires.
                if let Stmt::LocalDecl(d) = init.as_ref() {
                    ctx.declare(&d.name, type_from_typeref(&d.ty));
                    if let Some(e) = &d.init {
                        walk_expr_for_vsc(e, ctx);
                    }
                } else {
                    walk_stmt_for_vsc(init, ctx);
                }
            }
            if let Some(c) = &f.cond {
                walk_expr_for_vsc(c, ctx);
            }
            if let Some(st) = &f.step {
                walk_expr_for_vsc(st, ctx);
            }
            walk_stmt_for_vsc(&f.body, ctx);
            ctx.scope_pop();
        }
        Stmt::Return(Some(e)) => walk_expr_for_vsc(e, ctx),
        Stmt::Block(b) => walk_block_for_vsc(b, ctx),
        Stmt::Return(None) | Stmt::Break | Stmt::Continue => {}
    }
}

fn walk_expr_for_vsc(e: &Expr, ctx: &mut WalkCtx) {
    match e {
        Expr::Binary(b) => {
            walk_expr_for_vsc(&b.lhs, ctx);
            walk_expr_for_vsc(&b.rhs, ctx);
            if is_compare_op(b.op) {
                try_broadcast_scalar_compare(b, ctx);
            } else if is_arith_op(b.op) {
                try_promote_i32_arith(b, ctx);
            }
        }
        Expr::Unary(u) => walk_expr_for_vsc(&u.operand, ctx),
        Expr::Ternary(t) => {
            walk_expr_for_vsc(&t.cond, ctx);
            walk_expr_for_vsc(&t.then_expr, ctx);
            walk_expr_for_vsc(&t.else_expr, ctx);
        }
        Expr::Call(c) => {
            for a in &c.args {
                walk_expr_for_vsc(a, ctx);
            }
        }
        Expr::Member(m) => walk_expr_for_vsc(&m.base, ctx),
        Expr::Swizzle(s) => walk_expr_for_vsc(&s.base, ctx),
        Expr::Index(i) => {
            walk_expr_for_vsc(&i.base, ctx);
            walk_expr_for_vsc(&i.index, ctx);
        }
        Expr::InitList(l) => {
            for e in &l.elems {
                walk_expr_for_vsc(e, ctx);
            }
        }
        Expr::Assign(a) => {
            walk_expr_for_vsc(&a.target, ctx);
            walk_expr_for_vsc(&a.value, ctx);
        }
        Expr::Lit(_) | Expr::Ident(_, _) => {}
    }
}

fn is_compare_op(op: BinaryOp) -> bool {
    matches!(
        op,
        BinaryOp::Eq | BinaryOp::Ne | BinaryOp::Lt | BinaryOp::Le | BinaryOp::Gt | BinaryOp::Ge
    )
}

fn is_arith_op(op: BinaryOp) -> bool {
    matches!(
        op,
        BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div | BinaryOp::Rem
    )
}

/// Promote an `i32` operand of an `f32 ⟦arith⟧ i32` binop to `f32` by
/// wrapping it in `float(…)`. WGSL refuses mixed `f32 * i32`
/// arithmetic; corpus shape (`martin - pixies party 2 random
/// mosscity.milk`):
///
/// ```text
/// for (int n=1; n<=anz; n++) {
///     z = 1 - fract(1.0*n/anz - 1.0*fract(-t_rel/anz));
/// }
/// ```
///
/// After `for_init`, `n: i32`. The condition's `n <= anz` already gets
/// a `f32(n)` cast from [`try_broadcast_scalar_compare`]; the body's
/// `1.0*n` needs the same treatment.
fn try_promote_i32_arith(b: &crate::ast::BinaryExpr, ctx: &mut WalkCtx) {
    let lt = infer_type_static(&b.lhs, ctx);
    let rt = infer_type_static(&b.rhs, ctx);
    let int_side: &Expr = match (lt, rt) {
        (WgslType::I32, WgslType::F32) => &b.lhs,
        (WgslType::F32, WgslType::I32) => &b.rhs,
        // Both i32: leave alone (pure-int arithmetic is valid).
        // Vector mixes are handled by the comparison path / binop_vec.
        _ => return,
    };
    // Skip if already a `float(…)` / `f32(…)` cast — keeps the pass
    // idempotent.
    if let Expr::Call(c) = int_side
        && matches!(c.callee.as_str(), "float" | "f32")
    {
        return;
    }
    let span = int_side.span();
    ctx.edits.push(TextEdit {
        start: span.start,
        end: span.start,
        replacement: "float(".to_string(),
    });
    ctx.edits.push(TextEdit {
        start: span.end,
        end: span.end,
        replacement: ")".to_string(),
    });
}

/// Inspect a comparison binop's operands; if one is a known vector and
/// the other is a known scalar f32, emit a constructor wrap around the
/// scalar so the comparison evaluates as a vector-bool. Zero-length
/// insertions only, so the wrap survives even when other passes have
/// also edited the operands.
fn try_broadcast_scalar_compare(b: &crate::ast::BinaryExpr, ctx: &mut WalkCtx) {
    let lt = infer_type_static(&b.lhs, ctx);
    let rt = infer_type_static(&b.rhs, ctx);
    enum Wrap<'a> {
        ToVec(WgslType, &'a Expr),
        I32ToF32(&'a Expr),
    }
    let wrap: Wrap = match (lt.is_vec(), rt.is_vec()) {
        // vec ⟦cmp⟧ scalar — broadcast the scalar to vec
        (true, false) if matches!(rt, WgslType::F32 | WgslType::I32) => Wrap::ToVec(lt, &b.rhs),
        (false, true) if matches!(lt, WgslType::F32 | WgslType::I32) => Wrap::ToVec(rt, &b.lhs),
        // i32 ⟦cmp⟧ f32 — promote the i32 side to f32. `for (var n: i32 = …;
        // n <= anz; …)` where `anz: f32` is the dominant idiom (`for (int n =
        // …; n <= floatvar; …)` in HLSL).
        (false, false) => match (lt, rt) {
            (WgslType::I32, WgslType::F32) => Wrap::I32ToF32(&b.lhs),
            (WgslType::F32, WgslType::I32) => Wrap::I32ToF32(&b.rhs),
            _ => return,
        },
        _ => return,
    };
    match wrap {
        Wrap::ToVec(vec_ty, scalar_side) => {
            // Skip if the scalar side is already a constructor of matching
            // width — the wrap would be a no-op and accumulate on idempotent
            // re-runs.
            if let Expr::Call(c) = scalar_side
                && let Some(t) = constructor_return(&c.callee)
                && t == vec_ty
            {
                return;
            }
            let span = scalar_side.span();
            // Emit the HLSL-shape constructor (`float3(x)`) rather than the
            // WGSL shape (`vec3<f32>(x)`). The downstream `replace_types`
            // pass rewrites it to the WGSL form, and avoiding `<f32>` in
            // our inserted text means the later `looks_like_wgsl_generic`
            // guard in `rewrite_bool_to_float` doesn't get tripped — it
            // interprets `<f32>` as a sign of HLSL parser misreading
            // `vec3<f32>(…)` as a chained comparison.
            let ctor = match vec_ty {
                WgslType::Vec2F => "float2",
                WgslType::Vec3F => "float3",
                WgslType::Vec4F => "float4",
                _ => return,
            };
            ctx.edits.push(TextEdit {
                start: span.start,
                end: span.start,
                replacement: format!("{ctor}("),
            });
            ctx.edits.push(TextEdit {
                start: span.end,
                end: span.end,
                replacement: ")".to_string(),
            });
        }
        Wrap::I32ToF32(int_side) => {
            // Skip if already a `float(…)` cast or `f32(…)` cast.
            if let Expr::Call(c) = int_side
                && matches!(c.callee.as_str(), "float" | "f32")
            {
                return;
            }
            let span = int_side.span();
            ctx.edits.push(TextEdit {
                start: span.start,
                end: span.start,
                replacement: "float(".to_string(),
            });
            ctx.edits.push(TextEdit {
                start: span.end,
                end: span.end,
                replacement: ")".to_string(),
            });
        }
    }
}