onedrop_hlsl/rewrite/

chained_init.rs

//! Pass: lower chained assignment in local-decl init position.
//!
//! HLSL accepts `T x = y = expr;` (and deeper chains like
//! `T x = y = z = expr;`), evaluating right-to-left. WGSL has no
//! assignment expression — `var x: T = (y = expr);` is rejected at parse.
//! The corpus shape that hits this is a Martin/Rovastar idiom found in
//! ~8 / 2 000 random presets:
//!
//! ```hlsl
//! float2 ruv = uv = 0.5 + (uv-0.5)*(1+(rc.y*0.05));
//! ```
//!
//! Rewrite: lift the inner assignments out as statements that run
//! *before* the decl, then initialise the decl from the leftmost LHS:
//!
//! ```hlsl
//! uv = 0.5 + (uv-0.5)*(1+(rc.y*0.05)); float2 ruv = uv;
//! ```
//!
//! Detection runs on the AST produced by [`crate::parse::parse_hlsl`].
//! The parser already lowers the chain into a sequence of synthetic
//! [`Stmt::Assign`]s followed by the [`Stmt::LocalDecl`]; here we just
//! recognise that the synthetic assigns' spans sit *after* the decl's
//! span in source order (whereas a hand-written
//! `uv = 0.5; float2 ruv = uv;` has the assign's span *before* the
//! decl's). One edit replaces the original chained source with the
//! lowered statements; the existing trailing `;` is preserved.

use super::*;

pub(crate) fn rewrite_chained_assign_inits(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>) {
    let stmts = &b.stmts;
    let mut i = 0;
    while i < stmts.len() {
        if let Stmt::LocalDecl(d) = &stmts[i] {
            let chain_len = count_preceding_synthetic(stmts, i, d.span.start);
            if chain_len > 0 {
                emit_chain_edit(d, &stmts[i - chain_len..i], src, edits);
            }
        }
        descend(&stmts[i], src, edits);
        i += 1;
    }
}

fn descend(s: &Stmt, src: &str, edits: &mut Vec<TextEdit>) {
    match s {
        Stmt::If(i) => {
            descend(&i.then_branch, src, edits);
            if let Some(e) = &i.else_branch {
                descend(e, src, edits);
            }
        }
        Stmt::While(w) => descend(&w.body, src, edits),
        Stmt::For(f) => descend(&f.body, src, edits),
        Stmt::Block(b) => walk_block(b, src, edits),
        _ => {}
    }
}

/// How many `Stmt::Assign`s immediately before index `i` are synthetic
/// chained-init lowerings — identified by a span that starts *after* the
/// LocalDecl's span starts (i.e. the assignment text sits inside the
/// chained `T x = y = …;` block, but the AST lifted it ahead of the decl).
fn count_preceding_synthetic(stmts: &[Stmt], i: usize, decl_start: u32) -> usize {
    let mut k = 0;
    while k < i {
        let prev = &stmts[i - k - 1];
        let Stmt::Assign(a) = prev else { break };
        if a.span.start <= decl_start {
            break;
        }
        k += 1;
    }
    k
}

fn emit_chain_edit(d: &LocalDecl, synth: &[Stmt], src: &str, edits: &mut Vec<TextEdit>) {
    // synth is in AST order: synth[0] is the innermost assign (RHS = the
    // user's final expr); synth[len-1] is the outermost (RHS = the next
    // chain link). The source text for the whole chained init spans from
    // the decl's start to the innermost assign's value end.
    let last_value_end = match &synth[0] {
        Stmt::Assign(a) => a.value.span().end,
        _ => return,
    };
    let start = d.span.start;
    if last_value_end <= start {
        return; // sanity: spans must point forward
    }

    let mut text = String::new();
    // Emit assignments in evaluation order (innermost first).
    for stmt in synth {
        let Stmt::Assign(a) = stmt else {
            return;
        };
        text.push_str(slice(src, a.target.span()));
        text.push_str(" = ");
        text.push_str(slice(src, a.value.span()));
        text.push_str("; ");
    }
    // Emit the decl: original "T name" text, then `= <leftmost LHS>`.
    text.push_str(slice(src, d.span));
    if let Some(init) = &d.init {
        text.push_str(" = ");
        text.push_str(slice(src, init.span()));
    }

    edits.push(TextEdit {
        start,
        end: last_value_end,
        replacement: text,
    });
}

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

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

    #[test]
    fn two_deep_chain_lowered() {
        let src = "shader_body { float2 ruv = uv = 0.5; }";
        let out = rewrite_chained_assign_inits(src);
        assert!(out.contains("uv = 0.5; float2 ruv = uv"), "got: {out}");
        // Trailing `;` from the original source is preserved.
        assert!(out.trim_end_matches('}').trim_end().ends_with(';'));
    }

    #[test]
    fn three_deep_chain_lowered() {
        let src = "shader_body { float2 ruv = uv = uv2 = 0.5; }";
        let out = rewrite_chained_assign_inits(src);
        assert!(
            out.contains("uv2 = 0.5; uv = uv2; float2 ruv = uv"),
            "got: {out}"
        );
    }

    #[test]
    fn hand_written_assign_then_decl_untouched() {
        let src = "shader_body { uv = 0.5; float2 ruv = uv; }";
        let out = rewrite_chained_assign_inits(src);
        assert_eq!(out, src);
    }

    #[test]
    fn realistic_corpus_shape() {
        // The exact idiom from MilkDrop2077.11337.milk and 7 sibling
        // presets in the seed-42 2000-sample.
        let src = "shader_body { float3 rc = GetBlur1(uv); \
                   float2 ruv = uv = 0.5 + (uv-0.5)*(1+(rc.y*0.05)); \
                   ret = tex2D(sampler_main, ruv).xyz; }";
        let out = rewrite_chained_assign_inits(src);
        assert!(
            out.contains("uv = 0.5 + (uv-0.5)*(1+(rc.y*0.05)); float2 ruv = uv"),
            "got: {out}"
        );
        // The non-chained statements around it must survive untouched.
        assert!(out.contains("float3 rc = GetBlur1(uv);"));
        assert!(out.contains("ret = tex2D(sampler_main, ruv).xyz;"));
    }
}