onedrop_hlsl/rewrite/

modf_arity.rs

//! Pass: lower HLSL two-arg `modf(x, &out)` to WGSL's one-arg struct form.
//!
//! HLSL spec: `modf(x, &out)` returns the fractional part of `x` and
//! writes the integer part (truncated toward zero) into `out`.
//!
//! WGSL spec: `modf(x)` returns `__modf_result_f32 { fract: f32, whole:
//! f32 }`. The two-arg HLSL form is rejected at parse with
//! `too many arguments passed to 'modf'`.
//!
//! Corpus shape (5 / 2 000 warp presets, all in the same `PutDist`
//! helper inherited from the whoah / martin / amandio packs):
//!
//! ```hlsl
//! float fg, fb;
//! fg = modf((1-x)*255.0, fb);
//! ```
//!
//! Rewrite injects a unique temp at the call site so the input is only
//! evaluated once, then replaces the original assign with the two
//! component reads:
//!
//! ```wgsl
//! let _md2_modf_K = modf((1-x)*255.0);
//! fb = _md2_modf_K.whole;
//! fg = _md2_modf_K.fract;
//! ```
//!
//! Only fires on `Stmt::Assign(target, =, Call("modf", [arg, out_ident]))`
//! where `out_ident` is a bare `Expr::Ident`. Compound assigns (`+=` etc.)
//! and out-arg shapes the user wrote as `modf(x, &out)` (with HLSL's `&`
//! prefix) are out-of-scope — neither appears in the corpus.

use super::*;
use crate::lex::Span;
use std::sync::atomic::{AtomicUsize, Ordering};

static MODF_TEMP_COUNTER: AtomicUsize = AtomicUsize::new(0);

pub(crate) fn rewrite_modf_arity(src: &str) -> String {
    let Ok(tu) = parse_hlsl(src) else {
        return src.to_string();
    };
    let mut edits = Vec::new();
    if let Some(body) = &tu.shader_body {
        walk_block(body, src, &mut edits);
    }
    for item in &tu.items {
        if let Item::Function(f) = item {
            walk_block(&f.body, src, &mut edits);
        }
    }
    apply_edits(src, &mut edits)
}

fn walk_block(b: &Block, src: &str, edits: &mut Vec<TextEdit>) {
    for s in &b.stmts {
        walk_stmt(s, src, edits);
    }
}

fn walk_stmt(s: &Stmt, src: &str, edits: &mut Vec<TextEdit>) {
    match s {
        Stmt::Assign(a) if matches!(a.op, AssignOp::Set) => {
            try_emit_modf_rewrite(a, src, edits);
        }
        Stmt::LocalDecl(d) => {
            if let Some(init) = &d.init {
                try_emit_modf_in_init(d, init, src, edits);
            }
        }
        Stmt::Block(b) => walk_block(b, src, edits),
        Stmt::If(i) => {
            walk_stmt(&i.then_branch, src, edits);
            if let Some(e) = &i.else_branch {
                walk_stmt(e, src, edits);
            }
        }
        Stmt::While(w) => walk_stmt(&w.body, src, edits),
        Stmt::For(f) => walk_stmt(&f.body, src, edits),
        _ => {}
    }
}

fn try_emit_modf_rewrite(a: &AssignStmt, src: &str, edits: &mut Vec<TextEdit>) {
    let Expr::Call(c) = &a.value else { return };
    if !c.callee.eq_ignore_ascii_case("modf") || c.args.len() != 2 {
        return;
    }
    let Expr::Ident(_, _) = &c.args[1] else {
        return;
    };
    let target_text = slice(src, a.target.span());
    let out_text = slice(src, c.args[1].span());
    let arg_text = slice(src, c.args[0].span());
    let tmp = unique_temp();
    let stmt_text = format!(
        "let {tmp} = modf({arg_text}); {out_text} = {tmp}.whole; {target_text} = {tmp}.fract;"
    );
    edits.push(TextEdit {
        start: a.span.start,
        end: span_includes_semi(src, a.span).end,
        replacement: stmt_text,
    });
}

fn try_emit_modf_in_init(d: &LocalDecl, init: &Expr, src: &str, edits: &mut Vec<TextEdit>) {
    // `float fg = modf(x, fb);` — same shape but inside a local-decl init.
    // The LocalDecl's `span` only covers `T name` (not the init), so the
    // replacement range has to extend to the `;` after the init expr.
    let Expr::Call(c) = init else { return };
    if !c.callee.eq_ignore_ascii_case("modf") || c.args.len() != 2 {
        return;
    }
    let Expr::Ident(_, _) = &c.args[1] else {
        return;
    };
    let ty_text = slice(src, d.ty.span);
    let out_text = slice(src, c.args[1].span());
    let arg_text = slice(src, c.args[0].span());
    let tmp = unique_temp();
    let replacement = format!(
        "let {tmp} = modf({arg_text}); {out_text} = {tmp}.whole; {ty} {name} = {tmp}.fract;",
        ty = ty_text,
        name = d.name,
    );
    let init_end = init.span().end;
    let stmt_end = span_includes_semi(
        src,
        Span {
            start: d.span.start,
            end: init_end,
            line: d.span.line,
            col: d.span.col,
        },
    )
    .end;
    edits.push(TextEdit {
        start: d.span.start,
        end: stmt_end,
        replacement,
    });
}

fn unique_temp() -> String {
    let n = MODF_TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
    format!("_md2_modf_{n}")
}

fn slice(src: &str, span: Span) -> &str {
    &src[span.start as usize..span.end as usize]
}

fn span_includes_semi(src: &str, span: Span) -> Span {
    let end = span.end as usize;
    let bytes = src.as_bytes();
    let mut i = end;
    while i < bytes.len() && bytes[i].is_ascii_whitespace() {
        i += 1;
    }
    if i < bytes.len() && bytes[i] == b';' {
        Span {
            end: (i + 1) as u32,
            ..span
        }
    } else {
        span
    }
}

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

    #[test]
    fn two_arg_modf_in_assignment_gets_lowered() {
        let src = r#"shader_body {
    float fg;
    float fb;
    fg = modf((1-fg)*255.0, fb);
}"#;
        let out = rewrite_modf_arity(src);
        assert!(
            out.contains("modf((1-fg)*255.0)"),
            "single-arg call missing: {out}"
        );
        assert!(out.contains(".whole"), "whole field missing: {out}");
        assert!(out.contains(".fract"), "fract field missing: {out}");
    }

    #[test]
    fn one_arg_modf_left_alone() {
        // The WGSL-shape call survives — no two-arg shape to detect.
        let src = r#"shader_body { float r = modf(0.5); }"#;
        let out = rewrite_modf_arity(src);
        assert_eq!(out, src);
    }

    #[test]
    fn out_arg_must_be_ident() {
        // `modf(x, foo.bar)` is unusual and we don't try to lower it —
        // the existing source survives, naga will report the same error.
        let src = r#"shader_body {
    float fg;
    float2 v;
    fg = modf(0.5, v.x);
}"#;
        let out = rewrite_modf_arity(src);
        assert_eq!(out, src);
    }

    #[test]
    fn translate_roundtrip_putdist_pattern() {
        let hlsl = r#"
float2 PutDist(float x) {
    float fg = 0.0;
    float fb = 0.0;
    fg = modf((1-x)*255.0, fb);
    return float2(fg, fb);
}
shader_body { ret = float3(PutDist(0.5), 0); }
"#;
        let wgsl = translate_shader(hlsl).expect("translates");
        assert!(
            wgsl.contains("modf((1-x)*255.0)"),
            "expected single-arg modf, got:\n{wgsl}"
        );
        assert!(
            wgsl.contains(".whole") && wgsl.contains(".fract"),
            "struct field access missing:\n{wgsl}"
        );
        // The original two-arg form is gone.
        assert!(
            !wgsl.contains("modf((1-x)*255.0,fb)"),
            "stale two-arg call: {wgsl}"
        );
    }

    #[test]
    fn compound_assign_with_modf_left_alone() {
        // `fg += modf(x, fb);` — out-of-scope; the corpus uses `=` only.
        let src = r#"shader_body {
    float fg;
    float fb;
    fg += modf(0.5, fb);
}"#;
        let out = rewrite_modf_arity(src);
        assert_eq!(out, src);
    }

    #[test]
    fn modf_in_local_decl_init() {
        let src = r#"shader_body {
    float fb;
    float fg = modf(0.5, fb);
}"#;
        let out = rewrite_modf_arity(src);
        assert!(
            out.contains(".whole") && out.contains(".fract"),
            "got: {out}"
        );
        assert!(!out.contains("modf(0.5, fb)"), "stale two-arg: {out}");
    }
}