onedrop_hlsl/rewrite/

bool_arith.rs

//! Pass: bool RHS in f32 context → `select(0.0, 1.0, cond)`.

use super::*;

// ---------------------------------------------------------------------------
// Pass: bool RHS in f32 context → `select(0.0, 1.0, cond)`
// ---------------------------------------------------------------------------

/// HLSL coerces bool → float silently (`float mask = uv.x > 0.5;` yields
/// 0.0/1.0). WGSL refuses with `the type of 'mask' is expected to be
/// 'f32'; but got 'bool'` (or `cannot convert elements of vec3<bool> to
/// f32` for vector comparisons). One of the largest comp failure
/// clusters in the corpus.
///
/// Strategy: walk `LocalDecl`/`Assign` whose target type resolves to
/// `f32` / `vec_n<f32>`; when the RHS is a comparison or short-circuit
/// boolean, wrap it in `select(0.0, 1.0, <cond>)` (or `select(vec3<f32>(0.0),
/// vec3<f32>(1.0), <cond>)` for vector). Stays conservative: only fires when
/// both the target type and the RHS top-level operator are unambiguous.
pub(crate) fn rewrite_bool_to_float(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut ctx = WalkCtx::new(src);
    ctx.seed_globals(&tu);
    if let Some(body) = &tu.shader_body {
        walk_block_for_bool(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, type_from_typeref(&p.ty));
            }
            walk_block_for_bool(&f.body, &mut ctx);
            ctx.scope_pop();
        }
    }
    apply_edits(src, &mut ctx.edits)
}

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

fn walk_stmt_for_bool(s: &Stmt, ctx: &mut WalkCtx) {
    match s {
        Stmt::LocalDecl(d) => {
            // The LHS type used for the bool-coercion decision must
            // reflect what the *downstream* WGSL var ends up as, not the
            // literal HLSL token. `int mask = (c.y > 0);` is a dominant
            // failure; `replace_types` later rewrites `int` → `f32`, but
            // `type_from_typeref` returns `I32`, so the bool RHS wrap was
            // skipped. Pretend `int` is already `f32` (same widening rule
            // that `replace_types` applies to `int`/`half*`/`double*`).
            let target_ty = decl_target_float_kind(&d.ty);
            ctx.declare(&d.name, target_ty);
            if let Some(init) = &d.init {
                try_wrap_bool_rhs(init, target_ty, ctx);
                walk_expr_for_bool_arith(init, ctx);
            }
        }
        Stmt::Assign(a) => {
            // For Assign, infer the target type from the LHS expression.
            let target_ty = infer_type_static(&a.target, ctx);
            try_wrap_bool_rhs(&a.value, target_ty, ctx);
            walk_expr_for_bool(&a.target, ctx);
            walk_expr_for_bool_arith(&a.value, ctx);
        }
        Stmt::Expr(e) => walk_expr_for_bool_arith(e, ctx),
        Stmt::If(i) => {
            walk_expr_for_bool(&i.cond, ctx);
            walk_stmt_for_bool(&i.then_branch, ctx);
            if let Some(e) = &i.else_branch {
                walk_stmt_for_bool(e, ctx);
            }
        }
        Stmt::While(w) => {
            walk_expr_for_bool(&w.cond, ctx);
            walk_stmt_for_bool(&w.body, ctx);
        }
        Stmt::For(f) => {
            ctx.scope_push();
            if let Some(init) = &f.init {
                walk_stmt_for_bool(init, ctx);
            }
            if let Some(c) = &f.cond {
                walk_expr_for_bool(c, ctx);
            }
            if let Some(st) = &f.step {
                walk_expr_for_bool(st, ctx);
            }
            walk_stmt_for_bool(&f.body, ctx);
            ctx.scope_pop();
        }
        // `return` expressions are arithmetic context — they feed into
        // the caller's bool-or-float expectations. The corpus shape
        // `return !mask*domain + mask*refrac_uv` (Flexi's
        // `uv_lens_half_sphere`) needs the bool-typed `mask` operands
        // wrapped just like a `LocalDecl` RHS would.
        Stmt::Return(Some(e)) => walk_expr_for_bool_arith(e, ctx),
        Stmt::Block(b) => walk_block_for_bool(b, ctx),
        Stmt::Return(None) | Stmt::Break | Stmt::Continue => {}
    }
}

/// Top-level walk for nested sub-statements; the bool-wrap heuristic only
/// fires on the RHS of an assignment / decl, so this just recurses without
/// emitting edits of its own.
#[allow(clippy::only_used_in_recursion)]
fn walk_expr_for_bool(e: &Expr, ctx: &mut WalkCtx) {
    match e {
        Expr::Binary(b) => {
            walk_expr_for_bool(&b.lhs, ctx);
            walk_expr_for_bool(&b.rhs, ctx);
        }
        Expr::Unary(u) => walk_expr_for_bool(&u.operand, ctx),
        Expr::Ternary(t) => {
            walk_expr_for_bool(&t.cond, ctx);
            walk_expr_for_bool(&t.then_expr, ctx);
            walk_expr_for_bool(&t.else_expr, ctx);
        }
        Expr::Call(c) => {
            for a in &c.args {
                walk_expr_for_bool(a, ctx);
            }
        }
        Expr::Member(m) => walk_expr_for_bool(&m.base, ctx),
        Expr::Swizzle(s) => walk_expr_for_bool(&s.base, ctx),
        Expr::Index(i) => {
            walk_expr_for_bool(&i.base, ctx);
            walk_expr_for_bool(&i.index, ctx);
        }
        Expr::InitList(l) => {
            for e in &l.elems {
                walk_expr_for_bool(e, ctx);
            }
        }
        Expr::Assign(a) => {
            walk_expr_for_bool(&a.target, ctx);
            walk_expr_for_bool(&a.value, ctx);
        }
        Expr::Lit(_) | Expr::Ident(_, _) => {}
    }
}

/// Walk an expression that lives in arithmetic context (the RHS of a
/// decl/assign, an `*=` RHS, a `*`/`+`/`-`/`/` sub-tree). For every
/// `Binary` whose operator expects numeric operands, wrap any
/// bool-producing child (comparison or short-circuit) in `select(0.0,
/// 1.0, …)`. Targets a dominant validate-stage failure
/// (`InvalidBinaryOperandTypes`), exemplified by the MD2 gating idiom
/// `value * (rs.z > 0) * (rs.z < hlim)`.
///
/// We don't try to peek through `Member`/`Swizzle`/`Index`/Call args
/// here — those go through `try_wrap_bool_rhs` and the call-site coerce
/// pass at function granularity. Conservative wrapping keeps the edits
/// strictly zero-length insertions so existing inner edits (from
/// `rewrite_binary_vec_mismatches`, etc.) don't fight us.
fn walk_expr_for_bool_arith(e: &Expr, ctx: &mut WalkCtx) {
    if let Expr::Binary(b) = e
        && is_arith_op(b.op)
    {
        try_wrap_bool_arith_operand(&b.lhs, ctx);
        try_wrap_bool_arith_operand(&b.rhs, ctx);
    }
    walk_expr_for_bool(e, ctx);
    // Recurse into common arithmetic structures so multi-level
    // `a + (b > c) * d` chains catch every bool operand.
    match e {
        Expr::Binary(b) => {
            walk_expr_for_bool_arith(&b.lhs, ctx);
            walk_expr_for_bool_arith(&b.rhs, ctx);
        }
        Expr::Unary(u) => {
            // `-<bool>` and `+<bool>` are invalid in WGSL (naga's
            // `InvalidUnaryOperandType(Negate, …)` — the dominant fs_main
            // validation sub-bucket on the post-bb4754a 2000-sample,
            // ≈18 presets). Treat the operand as arithmetic context so
            // a comparison inside (`-(lum(ret)<sin(time*90))*…`) gets
            // wrapped in `select(0.0, 1.0, …)` *before* the unary
            // operator applies — yielding the valid `-(select(…))*…`.
            // `!<bool>` stays handled by the parent's
            // `try_wrap_bool_arith_operand` since `is_boolean_producing`
            // accepts `Unary(Not, bool)` as a whole.
            if matches!(u.op, UnaryOp::Neg | UnaryOp::Pos) {
                try_wrap_bool_arith_operand(&u.operand, ctx);
            }
            walk_expr_for_bool_arith(&u.operand, ctx);
        }
        Expr::Ternary(t) => {
            walk_expr_for_bool_arith(&t.then_expr, ctx);
            walk_expr_for_bool_arith(&t.else_expr, ctx);
        }
        Expr::Call(c) => {
            for a in &c.args {
                walk_expr_for_bool_arith(a, ctx);
            }
        }
        _ => {}
    }
}

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

/// Emit a `select(0.0, 1.0, (<expr>))` zero-length-insertion wrap around
/// `e` if its top-level operator is a boolean producer. Mirrors
/// [`try_wrap_bool_rhs`] but is called from arithmetic context, and also
/// emits the vector-shape form
/// `select(vec3<f32>(0.0), vec3<f32>(1.0), (vec_cmp))` when the bool's
/// operands are vector-typed, since WGSL `select(scalar, scalar,
/// vecN<bool>)` is invalid.
fn try_wrap_bool_arith_operand(e: &Expr, ctx: &mut WalkCtx) {
    if !is_boolean_producing(e, ctx) {
        return;
    }
    let text = &ctx.src[e.span().start as usize..e.span().end as usize];
    if looks_like_wgsl_generic(text) {
        return;
    }
    if is_chained_comparison(e) {
        return;
    }
    // Skip if already wrapped in `select(…)` — re-running the pass on its
    // own output should be idempotent. We detect by checking the byte
    // immediately before the start span: a preceding `(` of a `select(`
    // call form is the marker.
    let start = e.span().start as usize;
    if start >= 8 {
        let preceding = &ctx.src[start.saturating_sub(8)..start];
        if preceding.contains("1.0, ") || preceding.contains("1.0,(") {
            return;
        }
    }
    let (zero, one) = bool_shape_select_literals(e, ctx);
    let span = e.span();
    ctx.edits.push(TextEdit {
        start: span.start,
        end: span.start,
        replacement: format!("select({zero}, {one}, ("),
    });
    ctx.edits.push(TextEdit {
        start: span.end,
        end: span.end,
        replacement: "))".to_string(),
    });
}

/// Pick the `(zero, one)` literals for a `select(0.0/vecN, …)` wrap based
/// on the *shape* of the bool-producing expression. A scalar comparison
/// (e.g. `rs.z > 0`) keeps the scalar `0.0`/`1.0`; a vector comparison
/// (e.g. `noise >= 0.1`) needs `vecN<f32>(0.0)`/`vecN<f32>(1.0)` so the
/// resulting `select(...)` typechecks under WGSL's strict-shape rules.
fn bool_shape_select_literals(e: &Expr, ctx: &mut WalkCtx) -> (&'static str, &'static str) {
    let shape = infer_bool_shape(e, ctx);
    match shape {
        WgslType::Vec2F => ("vec2<f32>(0.0)", "vec2<f32>(1.0)"),
        WgslType::Vec3F => ("vec3<f32>(0.0)", "vec3<f32>(1.0)"),
        WgslType::Vec4F => ("vec4<f32>(0.0)", "vec4<f32>(1.0)"),
        _ => ("0.0", "1.0"),
    }
}

/// Walks the bool-producing subtree and returns the inferred *vector*
/// shape of the underlying values being compared. Returns `F32` for
/// scalar bool (so the caller emits scalar literals), `VecN` for vector
/// bool, `Unknown` when nothing can be inferred (caller falls back to
/// scalar — which matches the conservative pre-AA behaviour).
fn infer_bool_shape(e: &Expr, ctx: &mut WalkCtx) -> WgslType {
    match e {
        Expr::Binary(b) => match b.op {
            BinaryOp::Eq
            | BinaryOp::Ne
            | BinaryOp::Lt
            | BinaryOp::Le
            | BinaryOp::Gt
            | BinaryOp::Ge => widen_type(
                infer_type_static(&b.lhs, ctx),
                infer_type_static(&b.rhs, ctx),
            ),
            BinaryOp::And | BinaryOp::Or => {
                widen_type(infer_bool_shape(&b.lhs, ctx), infer_bool_shape(&b.rhs, ctx))
            }
            _ => WgslType::Unknown,
        },
        Expr::Unary(u) if matches!(u.op, UnaryOp::Not) => infer_bool_shape(&u.operand, ctx),
        Expr::Swizzle(s) => match s.components.len() {
            2 => WgslType::Vec2F,
            3 => WgslType::Vec3F,
            4 => WgslType::Vec4F,
            _ => WgslType::F32,
        },
        _ => WgslType::Unknown,
    }
}

/// Maps an HLSL declaration's textual type to the WGSL type the
/// downstream `replace_types` pass will produce — for the purpose of
/// deciding whether the RHS should be bool-coerced. Treats `int`, `half`,
/// `double` and their vector siblings as float-kind, mirroring the
/// substitution rules in [`crate::replace_types`].
pub(super) fn decl_target_float_kind(t: &TypeRef) -> WgslType {
    match t.name.as_str() {
        "int" | "half" | "double" | "half1" => WgslType::F32,
        "half2" | "double2" => WgslType::Vec2F,
        "half3" | "double3" => WgslType::Vec3F,
        "half4" | "double4" => WgslType::Vec4F,
        _ => type_from_typeref(t),
    }
}

/// If `rhs` is a comparison / boolean operator and `target_ty` is a
/// numeric float type, emit two zero-length edits to wrap the RHS in
/// `select(<zero>, <one>, <rhs>)`. Zero-length insertions don't collide
/// with inner edits emitted by other passes.
fn try_wrap_bool_rhs(rhs: &Expr, target_ty: WgslType, ctx: &mut WalkCtx) {
    if !is_float_kind(target_ty) {
        return;
    }
    if !is_boolean_producing(rhs, ctx) {
        return;
    }
    // Conservative guard against false positives where the HLSL parser
    // misinterprets a WGSL-style generic (`vec3<f32>(…)`, `array<f32, 4>`)
    // as a chained comparison `vec3 < f32 > (…)`. The `<…>` brackets show
    // up textually; if we see one in the RHS slice, this is not a real
    // boolean expression — bail.
    let rhs_text = &ctx.src[rhs.span().start as usize..rhs.span().end as usize];
    if looks_like_wgsl_generic(rhs_text) {
        return;
    }
    // Skip chained comparisons (`a < b > c`) — almost always parser noise
    // from generic-type lookalikes, never a real preset idiom.
    if is_chained_comparison(rhs) {
        return;
    }
    let (zero, one) = match target_ty {
        WgslType::F32 => ("0.0", "1.0"),
        WgslType::Vec2F => ("vec2<f32>(0.0)", "vec2<f32>(1.0)"),
        WgslType::Vec3F => ("vec3<f32>(0.0)", "vec3<f32>(1.0)"),
        WgslType::Vec4F => ("vec4<f32>(0.0)", "vec4<f32>(1.0)"),
        _ => return,
    };

    // WGSL `select` rejects `select(scalar, scalar, vec<bool>)` — the
    // condition's shape must match the value shape, or the values must
    // be scalar AND the condition scalar. Detect the "scalar target /
    // vector cmp" mismatch (HLSL silently extracts component 0 when
    // assigning `float x = vec_cmp;`) and wrap the cmp in `(cmp).x` to
    // recover scalar bool.
    let cmp_shape = infer_bool_shape(rhs, ctx);
    let needs_scalar_extract = matches!(target_ty, WgslType::F32)
        && matches!(
            cmp_shape,
            WgslType::Vec2F | WgslType::Vec3F | WgslType::Vec4F
        );

    let span = rhs.span();
    if needs_scalar_extract {
        // Result form: `select(0.0, 1.0, ((cmp)).x)`.
        ctx.edits.push(TextEdit {
            start: span.start,
            end: span.start,
            replacement: format!("select({zero}, {one}, (("),
        });
        ctx.edits.push(TextEdit {
            start: span.end,
            end: span.end,
            replacement: ")).x)".to_string(),
        });
    } else {
        ctx.edits.push(TextEdit {
            start: span.start,
            end: span.start,
            replacement: format!("select({zero}, {one}, ("),
        });
        ctx.edits.push(TextEdit {
            start: span.end,
            end: span.end,
            replacement: "))".to_string(),
        });
    }
}

fn looks_like_wgsl_generic(text: &str) -> bool {
    text.contains("<f32>")
        || text.contains("<i32>")
        || text.contains("<u32>")
        || text.contains("<bool>")
}

fn is_chained_comparison(e: &Expr) -> bool {
    let Expr::Binary(b) = e else { return false };
    let is_cmp = |op: BinaryOp| {
        matches!(
            op,
            BinaryOp::Eq | BinaryOp::Ne | BinaryOp::Lt | BinaryOp::Le | BinaryOp::Gt | BinaryOp::Ge
        )
    };
    if !is_cmp(b.op) {
        return false;
    }
    let inner = |x: &Expr| matches!(x, Expr::Binary(c) if is_cmp(c.op));
    inner(&b.lhs) || inner(&b.rhs)
}

fn is_float_kind(t: WgslType) -> bool {
    matches!(
        t,
        WgslType::F32 | WgslType::Vec2F | WgslType::Vec3F | WgslType::Vec4F
    )
}

/// True when the top-level operator of `e` is one that produces a bool
/// value in WGSL (comparison, short-circuit boolean, or a swizzle over
/// one). We follow `Swizzle` so that `(vec >= 0).xyz` reaches this
/// function during a `vec3` assignment and the bool-vec3 result still
/// gets the `select(...)` wrap.
///
/// The `ctx` argument lets us also catch the HLSL idiom of using a
/// bool-typed *variable* in arithmetic context (corpus shape:
/// `bool mask = …; return !mask*domain + mask*refrac_uv;`). HLSL
/// silently coerces those bool operands to 0.0 / 1.0; WGSL refuses
/// the multiplication outright, so the wrap has to apply at the
/// identifier level too.
fn is_boolean_producing(e: &Expr, ctx: &mut WalkCtx) -> bool {
    match e {
        Expr::Binary(b) => matches!(
            b.op,
            BinaryOp::Eq
                | BinaryOp::Ne
                | BinaryOp::Lt
                | BinaryOp::Le
                | BinaryOp::Gt
                | BinaryOp::Ge
                | BinaryOp::And
                | BinaryOp::Or
        ),
        Expr::Unary(u) => matches!(u.op, UnaryOp::Not) && is_boolean_producing(&u.operand, ctx),
        Expr::Swizzle(s) => is_boolean_producing(&s.base, ctx),
        Expr::Ident(name, _) => matches!(ctx.lookup(name), WgslType::Bool),
        _ => false,
    }
}